import os
from dataclasses import dataclass
from typing import Any, List, Literal, Tuple
import numpy as np
import pandas as pd
import xarray as xr
from .....utils.check.errclass import CifNotImplementedError, CifValueError
Model = Any
P_REF = 101325.0 # pressure at 1 atm
# 1, 2, 3, 3, 1, 7, 4, 8, 8, 4, 2, 8, 1, 7, 1, 6, 3, 4, 2, 5, 4, 5, 5, 5, 1
N_RATES_CONSTANTS = [1, 2, 3, 3] # ltabrate
[docs]
@dataclass
class Species:
"""A class representing a species
Args:
name (str): species name
type ('ac' or 'pr'): species type, 'ac' (active) or 'pr' (prescribed)
index (int): species index in arrays
restart_id (int): species restart id (active species only)
"""
name: str
type: Literal["ac", "pr"]
index: int = -1
restart_id: int = -1
[docs]
def is_active(self) -> bool:
"""Return ``True`` if this species is an active (transported) tracer."""
return self.type == "ac"
[docs]
def is_prescribed(self) -> bool:
"""Return ``True`` if this species is prescribed (fixed boundary condition)."""
return self.type == "pr"
[docs]
@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[1, 2, 3, 4] # Reaction type
rate_constants: List[float] # Reaction rate constants
[docs]
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 == 1:
da = self.rate_constants[0] * temp.copy(data=np.ones(temp.shape))
# Arrhenius simplified rates
elif self.reac_type == 2:
a, b = self.rate_constants
da = a * np.exp(-b / temp)
# Arrhenius complete rates
elif self.reac_type == 3:
a, b, c = self.rate_constants
da = a * np.exp(-b / temp) * (300.0 / temp)**c
# Pressure rates
elif self.reac_type == 4:
a, b, c = self.rate_constants
da = a * (b + c * pmid / P_REF)
# Photolysis rates
elif self.reac_type == 5:
raise CifNotImplementedError(
"Photolysis reaction rates are not supported")
else:
raise CifValueError(f"unexpected reaction type '{self.reac_type}'")
return da
[docs]
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
"""
scheme_id = self.chemistry.schemeid
chem_dir = os.path.join(self.workdir, "chemical_scheme", scheme_id)
# Parsing species file
all_spec_file = os.path.join(chem_dir, f"ALL_SPECIES.{scheme_id}")
df = pd.read_csv(all_spec_file, sep=' ', header=None)
species_dict = {spec_name: Species(name=spec_name, type=spec_type)
for spec_name, spec_type in zip(df.iloc[:, 0], df.iloc[:, 1])}
# Parsing active species file
ac_spec_file = os.path.join(chem_dir, f"ACTIVE_SPECIES.{scheme_id}")
df = pd.read_csv(ac_spec_file, sep=' ', header=None)
for i, (spec_name, restart_id) in enumerate(zip(df.iloc[:, 1], df.iloc[:, 2])):
assert species_dict[spec_name].is_active()
species_dict[spec_name].index = i
species_dict[spec_name].restart_id = restart_id
molar_masses = xr.DataArray(data=df.iloc[:, 3].values, dims=['spec'])
# Parsing prescribed species file
pr_spec_file = os.path.join(chem_dir, f"PRESCRIBED_SPECIES.{scheme_id}")
with open(pr_spec_file, 'r') as f:
lines = f.readlines()
for i, spec_name in enumerate(lines):
assert species_dict[spec_name.strip()].is_prescribed()
species_dict[spec_name.strip()].index = i
# Parsing prodloss species file
prodloss_file = os.path.join(chem_dir, f"PRODLOSS_SPECIES.{scheme_id}")
if os.path.getsize(prodloss_file) > 0:
raise CifNotImplementedError("prodloss species are not implemented")
# Parsing deposition species file
dep_file = os.path.join(chem_dir, f"DEPO_SPEC.{scheme_id}")
if os.path.getsize(dep_file) > 0:
raise CifNotImplementedError("deposition species are not implemented")
# Parsing reaction files
reac_file = os.path.join(chem_dir, f"CHEMISTRY.{scheme_id}")
rates_file = os.path.join(chem_dir, f"REACTION_RATES.{scheme_id}")
stoi_file = os.path.join(chem_dir, f"STOICHIOMETRY.{scheme_id}")
reaction_list = []
if os.path.isfile(reac_file) and os.path.isfile(rates_file):
with open(reac_file, "r") as f_reac, open(rates_file, "r") as f_rates:
reac_line = f_reac.readline()
rates_line = f_rates.readline()
while reac_line and rates_line:
reac_str = reac_line.strip().split()
rates_str = rates_line.strip().split()
n_reac = int(reac_str[0]) # Number of reactants
reactants = [
species_dict[spec_name] for spec_name in reac_str[1 : n_reac + 1]
]
n_prod = int(reac_str[n_reac + 1]) # Number of products
products = [
species_dict[spec_name]
for spec_name in reac_str[n_reac + 2 : n_prod + 1]
if spec_name in species_dict
]
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"
)
reac_type = int(rates_str[0])
n_rates = N_RATES_CONSTANTS[reac_type - 1]
rates = [float(r) for r in rates_str[1 : n_rates + 1]]
reaction_list.append(
Reaction(
active_reactants=ac_reactants,
prescribed_reactants=pr_reactants,
active_products=ac_products,
active_product_stoi=len(ac_products) * [1],
reac_type=reac_type,
rate_constants=rates,
)
)
reac_line = f_reac.readline()
rates_line = f_rates.readline()
# Parsing stoichiometric numbers files
if os.path.getsize(stoi_file) > 0:
df = pd.read_csv(stoi_file, sep=" ", header=None)
for i in range(len(df)):
spec_name = df.iloc[i, 0]
stoi = df.iloc[i, 1]
ireac = df.iloc[i, 2]
prod_ind = reaction_list[ireac].active_products.index(spec_name)
reaction_list[ireac].active_product_stoi_nums[prod_ind] = stoi
return reaction_list, molar_masses
