import numpy as np
import xarray as xr
import shutil
import pandas as pd
import os
import calendar
from ......utils.hdf5 import _hdf5_lock
BOLTZ = 1.38044e-16
DRY_MASS = 28.966
P_REF = 101325.
[docs]
def findspec(charspec, iallqmax, all_species):
"""Find the index of a species by name in the LMDZ-ACC species list.
Args:
charspec (str): species name to search for.
iallqmax (int): total number of species.
all_species (list[dict]): list of species dicts with a ``'name'`` key.
Returns:
int: 0-based index of *charspec* in *all_species*, or ``-1`` if not found.
"""
nospec = -1
for iq in range(iallqmax):
if charspec == all_species[iq]['name']:
nospec = iq
break
return nospec
[docs]
def comp_rates(nreac, temp, pmid, idtyperate, tabrate):
"""Compute gas-phase reaction rates for the LMDZ-ACC chemistry scheme.
Evaluates reaction rate expressions of multiple types (constant,
Arrhenius, pressure-dependent, …) using lookup table *tabrate* and
the atmospheric state (*temp*, *pmid*).
Args:
nreac (int): number of reactions.
temp (np.ndarray): temperature field (K), shape ``(...)``.
pmid (np.ndarray): mid-level pressure field (Pa), same shape.
idtyperate (np.ndarray): integer reaction-type codes, shape
``(nreac,)``.
tabrate (np.ndarray): rate constant table, shape
``(nparams, nreac)``.
Returns:
np.ndarray: reaction rates, shape ``temp.shape + (nreac,)``.
"""
rate = np.zeros(temp.shape + (nreac,))
# -- Computation of rates
for nr in range(nreac):
ity = idtyperate[nr]
# -- Constant rates
if ity == 1:
rate[..., nr] = tabrate[0, nr]
# -- Arrhenius simplified rates
elif ity == 2:
rate[..., nr] = tabrate[0, nr] * np.exp(-tabrate[1, nr] / temp)
# -- Arrhenius complete rates
elif ity == 3:
rate[..., nr] = tabrate[0, nr] * np.exp(-tabrate[1, nr] / temp) * (300. / temp) ** tabrate[2, nr]
# -- Pressure rates
elif ity == 4:
rate[..., nr] = tabrate[0, nr] * (tabrate[1, nr] * tabrate[2, nr] * pmid / P_REF)
# -- Photolysis rates
# elif ity == 5:
# for ij in range(ijratesmax):
# if jrates[ij]['idreac'] == nr:
# nrj = ij
# rate[..., nr] = refjrates[..., nrj]
return rate
[docs]
def make_chem_modout(self, runsubdir, ddi, inicond, inicond_ad):
"""Compute adjoint sensitivity of initial conditions through the chemistry scheme.
Reads the chemical scheme (reactions, stoichiometry, rate parameters)
from the LMDZ-ACC run directory, integrates chemical production/loss
over the period using :func:`comp_rates`, and propagates the adjoint
sensitivity *inicond_ad* backward through the chemistry operator to
produce sensitivity w.r.t. initial mass mixing ratios.
Args:
self: LMDZ-ACC model plugin instance (carries chemistry and domain).
runsubdir (str): path to the period run directory.
ddi (datetime): period start date.
inicond (np.ndarray): forward initial conditions (mass mixing ratio).
inicond_ad (np.ndarray): adjoint sensitivity at period end.
Returns:
np.ndarray: adjoint sensitivity w.r.t. initial conditions, same
shape as *inicond_ad*.
"""
# Domain infos
domain = self.domain
chemistry = self.chemistry
nlev = domain.nlev
nlat = domain.nlat
nlon = domain.nlon
nspec = len(chemistry.acspecies.attributes)
if not self.do_chemistry:
return np.zeros((nlon, nlat, nlev, nspec))
iprescrmax = 0
iprodmax = 0
idepmax = 0
ijratesmax = 0
ntabmax = 22
# -- Change ltabrate to properly retrieve the right number of coefficients used for the reaction
ltabrate = [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]
# -- Read the chemistry scheme
workdir = chemistry.workdir
dirchem_ref = f"{workdir}/chemical_scheme/{chemistry.schemeid}/"
# -- Reading active/output species names
file_chem = f"{dirchem_ref}/ACTIVE_SPECIES.{chemistry.schemeid}"
df_chem = pd.read_csv(file_chem, header=None, sep=" ")
species = [{'id': id, 'name': spec}
for id, spec in enumerate(df_chem.iloc[:, 0])]
adv_mass = df_chem.iloc[:, 2].values
# -- Reading all species names
file_chem = f"{dirchem_ref}/ALL_SPECIES.{chemistry.schemeid}"
df_chem = pd.read_csv(file_chem, header=None, sep=" ")
all_species = [{'id': id, 'name': spec, 'type': type, 'dep': False, 'iddep': 0, 'prod': False, 'idprod': 0}
for id, (spec, type) in enumerate(zip(df_chem.iloc[:, 0], df_chem.iloc[:, 1]))]
iallqmax = len(all_species)
# -- Reading prescribed species names
file_chem = f"{dirchem_ref}/PRESCRIBED_SPECIES.{chemistry.schemeid}"
if os.path.getsize(file_chem) > 0:
df_chem = pd.read_csv(file_chem, header=None, sep=" ")
prescr_species = [{'id': id, 'name': spec}
for id, spec in enumerate(df_chem.iloc[:, 0])]
iprescrmax = len(prescr_species)
# -- Reading prodloss species names
file_chem = f"{dirchem_ref}/PRODLOSS_SPECIES.{chemistry.schemeid}"
if os.path.getsize(file_chem) > 0:
df_chem = pd.read_csv(file_chem, header=None, sep=" ")
prodloss_species = [{'id': id, 'name': spec}
for id, spec in enumerate(df_chem.iloc[:, 0])]
iprodmax = len(prodloss_species)
# -- Reading deposition species names
file_chem = f"{dirchem_ref}/DEPO_SPEC.{chemistry.schemeid}"
if os.path.getsize(file_chem) > 0:
df_chem = pd.read_csv(file_chem, header=None, sep=" ")
dep_species = [{'id': id, 'name': spec}
for id, spec in enumerate(df_chem.iloc[:, 0])]
idepmax = len(dep_species)
# -- Make dep and prod associations
for iall in range(iallqmax):
for ipl in range(iprodmax):
if all_species[iall]['name'] == prodloss_species[ipl]['name']:
all_species[iall]['prod'] = True
all_species[iall]['idprod'] = ipl
for idp in range(idepmax):
if all_species[iall]['name'] == dep_species[idp]['name']:
all_species[iall]['dep'] = True
all_species[iall]['iddep'] = idp
# -- Reading reaction addressing arrays of chemistry
file_chem = f"{dirchem_ref}/CHEMISTRY.{chemistry.schemeid}"
df_chem = pd.read_csv(file_chem, header=None, sep=" ")
nreac = df_chem.shape[0]
nreactants = np.zeros(nreac, dtype=np.int8)
kreacl = np.zeros(iallqmax, dtype=np.int8)
idreacl = np.zeros((iallqmax, nreac), dtype=np.int8)
idspecl = np.zeros((nreac, 10), dtype=np.int8)
kreacp = np.zeros(iallqmax, dtype=np.int8)
idreacp = np.zeros((iallqmax, nreac), dtype=np.int8)
stoi = np.ones((iallqmax, nreac), dtype=np.int8)
tabrate = np.zeros((ntabmax, nreac))
idtyperate = np.zeros(nreac, dtype=np.int8)
for ireac in range(nreac):
nreactants[ireac] = df_chem.iloc[ireac, 0]
reactants = df_chem.iloc[ireac, 1:nreactants[ireac] + 1].values
nprods = int(df_chem.iloc[ireac, nreactants[ireac] + 1])
prods = df_chem.iloc[ireac, nreactants[ireac] + 2:nprods + nreactants[ireac] + 2].values
# -- Addressing the reactant list
for ire in range(nreactants[ireac]):
idloss = findspec(reactants[ire], iallqmax, all_species)
if idloss >= 0:
kreacl[idloss] = kreacl[idloss] + 1
idreacl[idloss, kreacl[idloss] - 1] = ireac
idspecl[ireac, ire] = idloss
# -- Addressing the product list
for ip in range(nprods):
idprod = findspec(prods[ip], iallqmax, all_species)
if idprod >= 0:
kreacp[idprod] = kreacp[idprod] + 1
idreacp[idprod, kreacp[idprod] - 1] = ireac
# -- Reading and addressing stoichiometric coefficients
file_chem = f"{dirchem_ref}/STOICHIOMETRY.{chemistry.schemeid}"
if os.path.getsize(file_chem) > 0:
df_chem = pd.read_csv(file_chem, header=None, sep=" ")
nlines = df_chem.shape[0]
for i in range(nlines):
reactants[0] = df_chem.iloc[i, 0]
coeff = df_chem.iloc[i, 1]
ireac = df_chem.iloc[i, 2]
idprod = findspec(reactants[0], iallqmax, all_species)
if idprod > 0:
stoi[idprod, ireac] = coeff
# -- Reading the reaction rates constants
file_chem = f"{dirchem_ref}/REACTION_RATES.{chemistry.schemeid}"
df_chem = pd.read_csv(file_chem, header=None, sep=" ")
for ireac in range(nreac):
tr = df_chem.iloc[ireac, 0]
tabrate[:ltabrate[tr], ireac] = df_chem.iloc[ireac, 1:ltabrate[tr] + 1].values
idtyperate[ireac] = tr
# -- Reading the photolysis reaction rates constants
file_chem = f"{dirchem_ref}/PHOTO_RATES.{chemistry.schemeid}"
if os.path.getsize(file_chem) > 0:
df_chem = pd.read_csv(file_chem, header=None, sep=" ")
jrates = [{'idj': id, 'idreac': ireac}
for id, ireac in enumerate(df_chem.iloc[:, 0])]
ijratesmax = len(jrates)
# -- Fetch kinetic variables
kinetic_file = f"{runsubdir}/kinetic.nc"
with _hdf5_lock:
with xr.open_dataset(kinetic_file) as ds:
temp_chem = ds.temp[0, ...].T.astype(np.float64)
pmid_chem = ds.pmid[0, ...].T.astype(np.float64)
# make lon cyclic
temp_chem = np.concatenate([temp_chem, temp_chem[-1:, ...]])
pmid_chem = np.concatenate([pmid_chem, pmid_chem[-1:, ...]])
convert = 10. * pmid_chem[..., np.newaxis] / \
(BOLTZ * temp_chem[..., np.newaxis])
# -- Fetch prescribed species
refprescr = []
for iq in range(iprescrmax):
prescr = prescr_species[iq]['name']
prescr_file = f"{runsubdir}/prescr_{prescr}.nc"
with _hdf5_lock:
ds = xr.open_dataset(prescr_file)
refprescr.append(ds[prescr].mean(axis=0).values)
refprescr = np.array(refprescr).T
# make lon cyclic
refprescr = np.concatenate([refprescr, refprescr[-1:, ...]])
# -- Converting variables from MMR/VMR to molecules/cm3
adv_mass = adv_mass[np.newaxis, np.newaxis, np.newaxis, :]
vmr = inicond * DRY_MASS / adv_mass
molec = vmr * convert
refprescr *= convert
# -- Setting new adjoint variables
delt = 3600 * 24 * calendar.monthrange(ddi.year, ddi.month)[1]
d_chem_ad = inicond_ad * delt * adv_mass / (convert * DRY_MASS)
molec_ad = np.zeros_like(d_chem_ad)
refprescr_ad = np.zeros(d_chem_ad.shape[:-1] + (iprescrmax,))
# -- Loss reactions
loss_ad = -d_chem_ad
rates = comp_rates(nreac, temp_chem, pmid_chem, idtyperate, tabrate)
for ns in range(nspec)[::-1]:
for kr in range(kreacl[ns])[::-1]:
ir = idreacl[ns, kr]
for ire in range(nreactants[ir])[::-1]:
nrat = rates[..., ir]
for irea in range(nreactants[ir])[::-1]:
if irea != ire:
if all_species[idspecl[ir, irea]]['type'] == 'ac':
nrat = nrat * molec[..., idspecl[ir, irea]]
elif all_species[idspecl[ir, irea]]['type'] == 'pr':
nrat = nrat * refprescr[..., idspecl[ir, irea] - nspec]
for irea in range(nreactants[ir])[::-1]:
if irea == ire:
if all_species[idspecl[ir, irea]]['type'] == 'ac':
molec_ad[..., idspecl[ir, irea]] += nrat * loss_ad[..., ns]
elif all_species[idspecl[ir, irea]]['type'] == 'pr':
refprescr_ad[..., idspecl[ir, irea] - nspec] += nrat * loss_ad[..., ns]
vmr_ad = molec_ad * convert
mmr_ad = vmr_ad * DRY_MASS / adv_mass
return mmr_ad
[docs]
def make_mod_out(self, runsubdir, ddi, ref_fwd_dir):
"""Create a zeroed LMDZ-ACC adjoint-output restart file for the current period.
Reads the reference forward restart file from *ref_fwd_dir*, copies it
to the adjoint run directory, and zeros out all active-species mass
mixing ratios so the LMDZ-ACC adjoint executable can accumulate
sensitivities into it.
Args:
self: LMDZ-ACC model plugin instance (carries domain and chemistry).
runsubdir (str): path to the period run directory.
ddi (datetime): period start date.
ref_fwd_dir (str): directory of the reference forward run.
"""
# Domain info
domain = self.domain
nlon = domain.nlon
nlat = domain.nlat
nlev = domain.nlev
if not hasattr(domain, "areas"):
domain.calc_areas()
areas = domain.areas
# Fetch original end from reference forward
end_file = f"{runsubdir}/restart.nc"
ref_end = ddi.strftime(f"{ref_fwd_dir}/chain/restart_%Y%m%d%H%M.nc")
shutil.copy(ref_end, end_file)
# Fetch initial conditions from reference forward
inicond_file = ddi.strftime(f"{ref_fwd_dir}/%Y-%m-%d_%H-%M/start.nc")
if os.path.isfile(inicond_file):
with _hdf5_lock:
ds = xr.open_dataset(inicond_file)
inicond = []
for spec in self.chemistry.acspecies.attributes:
restartID = getattr(self.chemistry.acspecies, spec).restart_id
var = f"q{restartID:02d}" if type(restartID) != str \
else restartID
inicond.append(ds[var].values[0])
inicond = np.array(inicond).T
else:
inicond_file = ddi.strftime(f"{ref_fwd_dir}/%Y-%m-%d_%H-%M/start.bin")
inicond = np.fromfile(
inicond_file, offset=4).reshape((nlon, nlat, nlev, nspec),
order="F")
# Read mass of air in the atmosphere and increase precision
mass_file = f"{runsubdir}/fluxstoke.nc"
with _hdf5_lock:
mass = xr.open_dataset(mass_file)["masse"][0].values[np.newaxis]
mass = mass.astype(np.float64)
# Put adjoint sensitivities to zero or propagate previous values
with _hdf5_lock:
ds = xr.open_dataset(end_file, mode="a")
ref_sensit = {}
ini_chem_sensit = {}
for spec in self.chemistry.acspecies.attributes:
restartID = getattr(self.chemistry.acspecies, spec).restart_id
var = f"q{restartID:02d}" if type(restartID) != str \
else restartID
# Read adjoint sensitivity at beginning of run
ref_sensit[spec] = 0
ini_chem_sensit[spec] = 0
if ddi != self.subsimu_dates[-2]:
sensit_start_file = f"{runsubdir}/start.nc"
with _hdf5_lock:
ds_ini = xr.open_dataset(sensit_start_file)
inicond_ad = ds_ini[var][:].values.T
ref_sensit[spec] = inicond_ad.sum() / mass.sum()
ini_chem_sensit[spec] = \
make_chem_modout(self, runsubdir, ddi, inicond, inicond_ad).T
with _hdf5_lock:
ds[var][:] = mass * ref_sensit[spec] + ini_chem_sensit[spec]
# Put initial condition sensitivities to mod_init
aini_file = f"{runsubdir}/mod_init_{spec}_out.bin"
ds[var][:].values[0].T.flatten(order="F").tofile(aini_file)
with _hdf5_lock:
ds.to_netcdf(end_file, mode="a")
# Fluxes
for spec in self.chemistry.emis_species.attributes:
aemis_file = f"{runsubdir}/mod_fluxes_{spec}_out.bin"
ndates = len(self.flx_input_dates[ddi])
emis_sensit = \
ref_sensit[spec] \
* pd.to_timedelta(self.flx_tresol).total_seconds() \
* areas.T[..., np.newaxis] * np.ones((nlon, nlat, ndates - 1))
emis_sensit.flatten(order="F").tofile(aemis_file)
# Prescribed concentrations
ndates = len(self.input_dates[ddi])
if hasattr(self.chemistry, "prescrconcs"):
for spec in self.chemistry.prescrconcs.attributes:
aprescr_file = f"{runsubdir}/mod_scale_{spec}_out.bin"
np.zeros((nlon, nlat, nlev, ndates - 1)).tofile(aprescr_file)
# Prodloss scaling
if hasattr(self.chemistry, "prodloss3d"):
for spec in self.chemistry.prescrconcs.attributes:
aprescr_file = f"{runsubdir}/mod_prodscale_{spec}_out.bin"
np.zeros((nlon, nlat, nlev, ndates - 1)).tofile(aprescr_file)
