Source code for pycif.plugins.models.lmdz_acc.perturb_model

from typing import Any
from ....utils.classes.chemistries import Chemistry


# Aliases for type hinting
Plugin = Any




def append_attribute(plugin: Plugin, key: str, attr: Plugin) -> None:
    setattr(plugin, key, attr)
    plugin.attributes.append(key)





def remove_attribute(plugin: Plugin, key: str) -> None:
    delattr(plugin, key)
    plugin.attributes.remove(key)





def perturb_model(self, nsamples, transf_mapper):
    self.nsamples = nsamples
    self.perturbed_species = {}
    self.original_species = self.chemistry.acspecies.attributes.copy()
    self.original_restart_ids = {}

    # List of species to not perturb
    dont_perturb_spec = getattr(self, 'dont_perturb_species', [])

    # Perturb active species in the chemical scheme
    list_acspecies = self.original_species
    restart_id = 1

    # Looping over active species
    for spec in list_acspecies:
        spec_plg = getattr(self.chemistry.acspecies, spec)
        self.original_restart_ids[spec] = spec_plg.restart_id

        if spec in dont_perturb_spec:
            spec_plg.restart_id = restart_id
            restart_id += 1

        else:
            # Looping over samples
            for i in range(nsamples):
                spec_sample = f"{spec}_{i:03d}"
                self.perturbed_species[spec_sample] = spec

                # Creating sample species
                spec_sample_plg = Chemistry(plg_orig=spec_plg)
                spec_sample_plg.restart_id = restart_id  # type: ignore
                restart_id += 1

                # Adding sample species to chemical scheme
                append_attribute(self.chemistry.acspecies,
                                 spec_sample, spec_sample_plg)

            # Removing original species from chemical scheme
            remove_attribute(self.chemistry.acspecies, spec)

    if hasattr(self.chemistry, 'emis_species'):
        # Perturb emmited species in the chemical scheme
        list_emispecies = self.chemistry.emis_species.attributes.copy()

        # Looping over emmited species
        for spec in list_emispecies:
            if spec in dont_perturb_spec:
                continue

            spec_plg = getattr(self.chemistry.emis_species, spec)

            # Looping over samples
            for i in range(nsamples):
                spec_sample = f"{spec}_{i:03d}"

                # Creating sample species
                spec_sample_plg = Chemistry(plg_orig=spec_plg)

                # Adding sample species to chemical scheme
                append_attribute(self.chemistry.emis_species,
                                 spec_sample, spec_sample_plg)

            # Removing original species from chemical scheme
            remove_attribute(self.chemistry.emis_species, spec)

    if hasattr(self.chemistry, 'reactions'):
        # Perturb reactions in the chemical scheme
        list_reactions = self.chemistry.reactions.attributes.copy()

        # Looping over reaction
        for key in list_reactions:
            reac_str = getattr(self.chemistry.reactions, key)

            # Parsing reaction string
            formula, rate = reac_str.split(maxsplit=1)

            # Paring species
            losses, products = formula.split("->", maxsplit=1)
            loss_spec = losses.split('+')
            stoi_prod_spec = products.split('+')

            # Parsing stoichiometric numbers for products
            prod_stoi = []
            prod_spec = []
            for val in stoi_prod_spec:
                if "*" in val:
                    stoi, spec = val.split("*", maxsplit=1)
                else:
                    stoi = 1
                    spec = val

                prod_stoi.append(stoi)
                prod_spec.append(spec)

            reac_dont_pertub_spec = [
                spec in dont_perturb_spec for spec in loss_spec + prod_spec
                if spec in list_acspecies
            ]

            if not any(reac_dont_pertub_spec):
                # All active species in reaction be perturbed
                # Looping over samples
                for i in range(nsamples):
                    # Replacing losses species with sample species
                    sample_loss_spec = []
                    for spec in loss_spec:
                        if spec in list_acspecies and spec not in dont_perturb_spec:
                            sample_loss_spec.append(f"{spec}_{i:03d}")
                        else:
                            sample_loss_spec.append(spec)

                    # Replacing product species with sample species
                    sample_prod_spec = []
                    for spec in prod_spec:
                        if spec in list_acspecies and spec not in dont_perturb_spec:
                            sample_prod_spec.append(f"{spec}_{i:03d}")
                        else:
                            sample_prod_spec.append(spec)

                    # Making sample reaction string
                    sample_stoi_prod_spec = [
                        f"{stoi}*{spec}" if stoi > 1 else spec
                        for stoi, spec in zip(prod_stoi, sample_prod_spec)
                    ]
                    sample_losses = '+'.join(sample_loss_spec)
                    sample_products = '+'.join(sample_stoi_prod_spec)
                    sample_reac_str = f"{sample_losses}->{sample_products} {rate}"

                    # Adding sample reaction to chemical scheme
                    append_attribute(self.chemistry.reactions,
                                     f"{key}_{i:03d}", sample_reac_str)

                # Removing original reaction from chemical scheme
                remove_attribute(self.chemistry.reactions, key)

            elif not all(reac_dont_pertub_spec):
                # Mix of species to perturb and to not perturb in reaction
                raise ValueError(
                    f"reaction '{key}: {reac_str}' contain both species to "
                    "perturb and to not perturb"
                )

    # Dump updated chemical scheme as text files for CHIMERE executable
    self.chemistry.create_chemicalscheme()

    if hasattr(self, 'approx_thresholds'):
        # Perturb 'approx_threshold' input argument
        dict_approx = self.approx_thresholds.copy()

        # Looping over emmited species
        for spec, threshold in dict_approx.items():
            if spec in dont_perturb_spec:
                continue

            # Looping over samples
            for i in range(nsamples):
                spec_sample = f"{spec}_{i:03d}"

                # Adding sample species to 'approx_thresholds'
                self.approx_thresholds[spec_sample] = threshold

            # Removing original species
            self.approx_thresholds.pop(spec)