from pathlib import Path
import numpy as np
import pandas as pd
import xarray as xr
from ...chemistry import (
BOLTZMAN,
compute_chemistry_step_ad,
parse_chemical_scheme,
read_kinetic,
read_prescr,
)
from .fake_end import read_air_mass
from ......utils.hdf5 import _hdf5_lock
[docs]
def make_adjoint_out(self, runsubdir, ddi, ref_fwd_dir):
runsubdir = Path(runsubdir)
ref_fwd_dir = Path(ref_fwd_dir)
# Dimensions
nspec = len(self.chemistry.active_species)
nlev = self.domain.nlev
nlat = self.domain.nlat
nlon = self.domain.nlon
if not hasattr(self.domain, "areas"):
self.domain.calc_areas()
area = self.domain.areas
# Fetch initial conditions from reference forward
inicond = xr.DataArray(
data=np.zeros((nspec, nlev, nlat, nlon)),
dims=["spec", "lev", "lat", "lon"],
)
inicond_file = ref_fwd_dir / ddi.strftime("%Y-%m-%d_%H-%M") / "start.nc"
with _hdf5_lock:
with xr.open_dataset(inicond_file) as ds:
for index, spec in enumerate(self.chemistry.active_species):
inicond[index, ...] = ds[spec].values
# Read mass of air in the atmosphere and increase precision
mass = read_air_mass(self, ddi, runsubdir).isel(time=0)
# Read adjoint sensitivity at beginning of run
if ddi != self.subsimu_dates[-2]:
sum_mass = float(mass.sum())
# Parsing schemical scheme
reaction_list, molar_masses = parse_chemical_scheme(self)
# Reading kinetic.nc, dims: (time, lev, lat, lon)
kinetic = read_kinetic(self, ddi, runsubdir)
ref_pmid, ref_temp = kinetic.pmid, kinetic.temp
# Reading prescr_*.nc, dims: (spec, time, lev, lat, lon) [molec/cm3]
ref_prescr = read_prescr(self, ddi, runsubdir)
ref_prescr = ref_prescr * ref_pmid / (BOLTZMAN * ref_temp) * 1.0e-6
inicond_ad = xr.DataArray(
data=np.zeros((nspec, nlev, nlat, nlon)),
dims=["spec", "lev", "lat", "lon"],
)
sensit_start_file = runsubdir / "start.nc"
with _hdf5_lock:
with xr.open_dataset(sensit_start_file) as ds_ini:
for index, spec in enumerate(self.chemistry.active_species):
inicond_ad[index, ...] = ds[spec].values
ref_sensit = inicond_ad.sum(["lev", "lat", "lon"]) / sum_mass
mmr_ad = compute_chemistry_step_ad(
reaction_list,
molar_masses,
3600 * 24 * pd.to_datetime(ddi).days_in_month,
inicond,
inicond_ad,
ref_prescr.isel(time=0),
ref_pmid.isel(time=0),
ref_temp.isel(time=0),
)
ini_chem_sensit = mmr_ad.sum(["lev", "lat", "lon"])
else:
ref_sensit = xr.DataArray(data=np.zeros(nspec), dims=["spec"])
ini_chem_sensit = xr.DataArray(data=np.zeros(nspec), dims=["spec"])
# Put adjoint sensitivities to zero or propagate previous values
ds_ini = xr.Dataset(
{
spec: mass * ref_sensit[index] + ini_chem_sensit[index]
for index, spec in enumerate(self.chemistry.active_species)
}
)
with _hdf5_lock:
ds_ini.to_netcdf(runsubdir / "restart.nc")
# Fluxes
ntime = len(self.flx_input_dates[ddi]) - 1
ones = xr.DataArray(
data=np.ones((ntime, nlat, nlon)),
dims=["time", "lat", "lon"],
)
dt = pd.to_timedelta(self.flx_tresol).total_seconds()
ds_flux = xr.Dataset(
{
spec: ones * ref_sensit[index] * dt * area[np.newaxis, ...]
for index, spec in enumerate(self.chemistry.emitted_species)
}
)
with _hdf5_lock:
ds_flux.to_netcdf(runsubdir / "flux_out.nc")
ntime = len(self.input_dates[ddi]) - 1
zeros = xr.DataArray(
data=np.zeros((ntime, nlev, nlat, nlon)),
dims=["time", "lev", "lat", "lon"],
)
# Prescribed concentrations
ds_prescr = xr.Dataset({spec: zeros for spec in self.chemistry.prescribed_species})
with _hdf5_lock:
ds_prescr.to_netcdf(runsubdir / "prescr_out.nc")
# Prodloss
ds_prodloss = xr.Dataset({spec: zeros for spec in self.chemistry.prodloss_species})
with _hdf5_lock:
ds_prodloss.to_netcdf(runsubdir / "prodloss3d_out.nc")
