Source code for pycif.plugins.models.lmdz_old.io.inputs.ensemble

from typing import Any, Dict, Tuple

# Aliases for type hinting
Model = Any




def ensemble_trid(
    self: Model,
    trid: Tuple[str, str],
    datastore: Dict[Tuple[str, str], Any]
) -> Tuple[str, str]:
    """Replace 'spec' in the tracer id 'trid' (component, spec) by the suitable
    species name when in ensemble mode or do nothing when not in ensemble mode

    Args:
        self: Model
        trid (str, str): Tracer id (component, species)
        datastore (dict (str, str) -> Any): Datastore

    Returns:
        (str, str): New tracer id
    """
    comp, spec = trid
    ref_trid = trid

    # Update ref_trid if in ensemble mode
    if hasattr(self, 'perturbed_species'):
        if spec in self.perturbed_species:
            ref_trcr = self.perturbed_species[spec]
            ref_trid = (comp, ref_trcr)

    # If spec not explicitly defined in datastore, fetch general component
    # information if available
    if trid not in datastore:
        # If trid is a perturbed species, needs to reformat name to fit
        # LMDZ naming convention for perturbed species
        for comp_name, tracer_name in datastore.keys():
            if trid == (comp_name, tracer_name.replace("__sample#", "_")):
                # Tracer 'trid' is a perturbated species, reformatting is id
                trid = (comp_name, tracer_name)
                break

        else:
            # Branch here when exiting the for loop without breaking, i.e. if
            # trid is not a perturbated species

            # If trid comes from an unperturbed input in ensemble mode, the
            # info is fetched from the reference species
            if ref_trid in datastore:
                trid = ref_trid

            elif (comp, "") in datastore:
                trid = (comp, "")

    return trid