Source code for pycif.plugins.models.lmdz_acc.chemistry.compute_chemistry

from typing import List, Tuple

import numpy as np
import xarray as xr

from .chemical_scheme import Reaction

BOLTZMAN = 1.380649e-23
DRY_MASS = 28.966




def compute_chemistry_step(
    reaction_list: List[Reaction],
    molar_masses: xr.DataArray,
    dt: float,
    mmr: xr.DataArray,
    mmr_tl: xr.DataArray,
    prescr: xr.DataArray,
    prescr_tl: xr.DataArray,
    pmid: xr.DataArray,
    temp: xr.DataArray,
) -> Tuple[xr.DataArray, xr.DataArray]:
    """Reimplementation of LMDZ's "compute_chem_tl" subroutine.

    Args:
        reaction_list (list of Reaction): list of reactions
        molar_masses (xr.DataArray): molar masses of active species
        dt (float): time step [s]
        mmr (xr.DataArray): mass ratio (forward) [kg/kg]
        mmr_tl (xr.DataArray): mass ratio (tangent) [kg/kg]
        prescr (xr.DataArray): prescribed species concentrations (forward) [molec/cm3]
        prescr_tl (xr.DataArray): prescribed species concentrations (tangent) [molec/cm3]
        pmid (xr.DataArray): pressure field [Pa]
        temp (xr.DataArray): temperature field [K]

    Returns:
        xr.DataArray: losses [kg/kg]
        xr.DataArray: losses_tl [kg/kg]
    """
    convert = DRY_MASS / molar_masses * pmid / (BOLTZMAN * temp) * 1.0e-6
    conc = mmr * convert  # mmr -> molec/cm3
    conc_tl = mmr_tl * convert  # mmr -> molec/cm3
    chem_losses = conc.copy(data=np.zeros(conc.shape))
    chem_losses_tl = conc.copy(data=np.zeros(conc.shape))

    # Loop over all reactions (computing reacting concentrations in molec/cm3)
    for reac in reaction_list:
        rates = reac.rates(temp, pmid)

        # Initialize current reaction concentrations
        reac_conc = rates.copy(deep=True)
        reac_conc_tl = rates.copy(data=np.zeros(rates.shape))

        # --- Computing rates (forward)
        # Loop over active reactants (computing rates)
        for spec in reac.active_reactants:
            reac_conc *= conc[spec.index, ...]

        # Loop over prescribed reactants (computing rates)
        for spec in reac.prescribed_reactants:
            reac_conc *= prescr[spec.index, ...]

        # --- Computing rates (tangent)
        # Loop over active reactants (computing rates, tangent)
        for spec in reac.active_reactants:
            tmp = rates.copy(deep=True)

            # Sub loop over active reactants
            for spec_inner in reac.active_reactants:
                if spec.name != spec_inner.name:
                    tmp *= conc[spec_inner.index, ...]
                else:
                    tmp *= conc_tl[spec_inner.index, ...]

            # Loop over prescribed reactants (computing rates, forward)
            for spec_inner in reac.prescribed_reactants:
                tmp *= prescr[spec_inner.index, ...]

            reac_conc_tl += tmp

        # Loop over prescribed reactants (computing rates, tangent)
        for spec in reac.prescribed_reactants:
            tmp = rates.copy(deep=True)

            # Sub loop over active reactants
            for spec_inner in reac.active_reactants:
                tmp *= conc[spec_inner.index, ...]

            # Loop over prescribed reactants (computing rates, forward)
            for spec_inner in reac.prescribed_reactants:
                if spec.name != spec_inner.name:
                    tmp *= prescr[spec_inner.index, ...]
                else:
                    tmp *= prescr_tl[spec_inner.index, ...]

            reac_conc_tl += tmp

        # --- Applying rates (forward and tangent)
        # Loop over prescribed reactants (applying rates for losses)
        for spec in reac.active_reactants:
            chem_losses[spec.index, ...] -= reac_conc
            chem_losses_tl[spec.index, ...] -= reac_conc_tl

    chem_losses = chem_losses / convert * dt  # molec/cm3 -> mmr
    chem_losses_tl = chem_losses_tl / convert * dt  # molec/cm3 -> mmr

    return chem_losses, chem_losses_tl