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

from __future__ import annotations

from typing import Any

from ....utils.classes.chemistries import Chemistry
from ...chemistries.lmdz.utils import format_reaction, parse_reaction

# 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)



# pylint: disable=unused-argument


def perturb_model(self, nsamples: int, transf_mapper):
    # raise NotImplementedError("chemistry attribute access not updated yet")
    self.nsamples = nsamples
    self.perturbed_species = {}
    self.original_species = list(self.chemistry.active_species)
    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_active_species = self.original_species
    restart_id = 1

    # Looping over active species
    for spec in list_active_species:
        spec_plg = self.chemistry.active_species[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
                self.chemistry.active_species[spec_sample] = spec_sample_plg
                append_attribute(self.chemistry.acspecies,
                                 spec_sample, spec_sample_plg)

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

    # Perturb emmited species in the chemical scheme
    emitted_species = []
    # Looping over emmited species
    for spec in self.chemistry.emitted_species:
        if spec in dont_perturb_spec:
            emitted_species.append(spec)

        else:
            for i in range(nsamples):
                emitted_species.append(f"{spec}_{i:03d}")

    # Replacing emmited species from chemical scheme
    self.chemistry.emitted_species = emitted_species

    # Perturb reactions in the chemical scheme
    reactions = []
    # Looping over reaction
    for reac_str in self.chemistry.reactions:
        loss_spec, prod_spec, prod_stoi, rate_type, rate_terms = (
            self.chemistry.parse_reaction(reac_str)
        )

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

        # All active species in reaction be perturbed
        if not any(reac_dont_pertub_spec):
            # 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_active_species 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_active_species 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_reac_str = format_reaction(
                    sample_loss_spec,
                    sample_prod_spec,
                    prod_stoi,
                    rate_type,
                    rate_terms,
                )

                # Adding sample reaction to chemical scheme
                reactions.append(sample_reac_str)

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

        else:
            # Dont perturb reaction
            reactions.append(reac_str)

    self.chemistry.reactions = reactions

    # Dump updated chemical scheme
    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)