Source code for pycif.plugins.models.lmdz_ico.chemistry.chemical_scheme

from dataclasses import dataclass
from typing import Any, List, Literal, Tuple

import numpy as np
import xarray as xr

from ...lmdz_acc.chemistry.chemical_scheme import P_REF, Species
from .....utils.check.errclass import CifNotImplementedError, CifValueError

Model = Any




@dataclass
class Reaction:
    """A class representing a chemical reaction

    Args:
        active_reactants (list of Species): Active reactants
        prescribed_reactants (list of Species): Prescribed reactants
        active_products (list of Species): Active products
        active_product_stoi (list of int): Active product stoichiometric numbers
        reac_type (int): Reaction type
        rate_constants (list of float) Reaction rate constants
    """

    active_reactants: List[Species]  # Active reactants
    prescribed_reactants: List[Species]  # Prescribed reactants
    active_products: List[Species]  # Active products
    active_product_stoi: List[int]  # Active product stoichiometric numbers
    reac_type: Literal[
        "constant", "simplified_arrhenius", "arrhenius", "pressure"
    ]  # Reaction type
    rate_constants: List[float]  # Reaction rate constants



    def rates(self, temp: xr.DataArray, pmid: xr.DataArray) -> xr.DataArray:
        """Compute reaction rates

        Args:
            temp (xr.DataArray): temperature field [K]
            pmid (xr.DataArray): pressure field [Pa]

        Returns:
            xr.DataArray: rates [molec/cm2/s2]
        """
        # Constant rates
        if self.reac_type == "constant":
            da = self.rate_constants[0] * temp.copy(data=np.ones(temp.shape))

        # Arrhenius simplified rates
        elif self.reac_type == "simplified_arrhenius":
            a, b = self.rate_constants
            da = a * np.exp(-b / temp)

        # Arrhenius complete rates
        elif self.reac_type == "arrhenius":
            a, b, c = self.rate_constants
            da = a * np.exp(-b / temp) * (300.0 / temp) ** c

        # Pressure rates
        elif self.reac_type == "pressure":
            a, b, c = self.rate_constants
            da = a * (b + c * pmid / P_REF)

        # Photolysis rates
        elif self.reac_type == "photolysis":
            raise CifNotImplementedError("Photolysis reaction rates are not supported")

        else:
            raise CifValueError(f"unexpected reaction type '{self.reac_type}'")

        return da  # type: ignore






def parse_chemical_scheme(self: Model) -> Tuple[List[Reaction], xr.DataArray]:
    """Parse the chemical scheme, partial reimplementation of LMDZ's
    "read_chemical_scheme" subroutine

    Args:
        self (Model)

    Returns:
        (list of Reaction, xr.DataArray): list of reactions, molar masses
    """
    species_dict = {}
    for i, spec in enumerate(self.chemistry.active_species):
        species_dict[spec] = Species(name=spec, type="ac", index=i)
    for i, spec in enumerate(self.chemistry.prescribed_species):
        species_dict[spec] = Species(name=spec, type="pr", index=i)

    molar_masses = xr.DataArray(
        data=[spec.molar_mass for spec in self.chemistry.active_species.values()],
        dims=["spec"],
    )

    if self.chemistry.prodloss_species:
        raise CifNotImplementedError("prodloss species are not implemented")

    if self.chemistry.deposition_species:
        raise CifNotImplementedError("deposition species are not implemented")

    reaction_list = []
    for reac_str in self.chemistry.reactions:
        reactants, products, stoi, reac_type, rate_constants = (
            self.chemistry.parse_reaction(self, reac_str)
        )

        reactants = [species_dict[spec] for spec in reactants]
        products = [species_dict[spec] for spec in products]

        ac_reactants = [spec for spec in reactants if spec.is_active()]
        pr_reactants = [spec for spec in reactants if spec.is_prescribed()]
        ac_products = [spec for spec in products if spec.is_active()]

        if ac_products:
            raise CifNotImplementedError(
                "reactions with active products are not implemented"
            )

        reaction_list.append(
            Reaction(
                active_reactants=ac_reactants,
                prescribed_reactants=pr_reactants,
                active_products=[],
                active_product_stoi=[],
                reac_type=reac_type,
                rate_constants=rate_constants,
            )
        )

    return reaction_list, molar_masses