import datetime
from logging import debug
import os
from typing import Any, Literal, Tuple, Union
import numpy as np
import pandas as pd
import xarray as xr
from ..inputs.inicond import update_controle_var
from ......utils.dates import date_range
from ......utils.hdf5 import _hdf5_lock
from ......utils.path import link
from ......utils.check.errclass import CifFileNotFoundError, CifValueError
from ...chemistry import (BOLTZMAN, DRY_MASS, compute_chemistry_step,
parse_chemical_scheme, read_inicond_file,
read_kinetic, read_prescr, strip_spatial_coords)
NANO = np.timedelta64(1, 'ns')
# Aliases for type hinting
Model = Any
[docs]
def read_obs(runsubdir: str) -> Tuple[int, np.ndarray, np.ndarray]:
"""Reads 'tstep' and species index columns from 'obs.bin' the binary file
and convert from Fortran to Python indices (i = i - 1)
Args:
runsubdir (str): runsubdir (str): sub simulation run directory
Returns:
(int, 1D array, 1D array): number of observations, tstep, species index
"""
obs_file = os.path.join(runsubdir, "obs.bin")
obs_data = np.fromfile(obs_file, dtype='float').reshape((-1, 6), order="F")
nobs, _ = obs_data.shape
tstep = obs_data[:, 0].astype('int64') - 1 # Fortran to Python index
spec_index = np.floor(obs_data[:, 5]).astype('int64') - 1
return nobs, tstep, spec_index
[docs]
def read_air_mass(
self: Model,
ddi: datetime.datetime,
runsubdir: str
) -> xr.DataArray:
"""Read the LMDZ-old air-mass field from the run directory.
Args:
self: LMDZ-old model plugin instance.
ddi: period start date.
runsubdir: path to the period run directory.
Returns:
xr.DataArray: air-mass field on ``(time, lev, lat, lon)`` grid.
"""
di = pd.to_datetime(ddi)
di = di if di.is_month_start else di - pd.offsets.MonthBegin()
freq = pd.Timedelta(self.meteo.offtstep, 'seconds')
ntime = int(di.days_in_month * 24 * 3600 / freq.seconds)
time = xr.DataArray(data=pd.date_range(di, periods=ntime, freq=freq),
dims=['time'])
mass_file = os.path.join(runsubdir, "fluxstoke.nc")
with _hdf5_lock:
with xr.open_dataset(mass_file) as ds:
da = ds['masse']
da = da.astype('float64')
da = da.rename({'time_counter': 'time', 'sig_s': 'lev'})
da = strip_spatial_coords(da)
da['time'] = time
return da
[docs]
def read_emissions(
self: Model,
ddi: datetime.datetime,
runsubdir: str
) -> Tuple[xr.DataArray, xr.DataArray]:
"""Reads emissions 'mod_*.bin' binary files
Args:
self (Model)
ddi (datetime.datetime): sub simulation start datetime
runsubdir (str): runsubdir (str): sub simulation run directory
Returns:
(xr.DataArray, xr.DataArray): emissions forward, emissions tangent
"""
# Domain infos
nspec = len(self.chemistry.acspecies.attributes)
ntime = self.flx_input_dates[ddi].size - 1
nlat = self.domain.nlat
nlon = self.domain.nlon
di = pd.to_datetime(ddi)
di = di if di.is_month_start else di - pd.offsets.MonthBegin()
time = xr.DataArray(data=pd.date_range(
di, periods=ntime, freq=self.flx_tresol), dims=['time'])
if not hasattr(self.domain, 'areas'):
self.domain.calc_areas()
areas = self.domain.areas
dt = pd.to_timedelta(self.flx_tresol).total_seconds()
emis = xr.DataArray(data=np.zeros((nspec, ntime, nlat, nlon)),
dims=['spec', 'time', 'lat', 'lon'],
coords={'time': time})
emis_tl = xr.DataArray(data=np.zeros((nspec, ntime, nlat, nlon)),
dims=['spec', 'time', 'lat', 'lon'],
coords={'time': time})
for index, spec in enumerate(self.chemistry.acspecies.attributes):
emis_file = os.path.join(runsubdir, f"mod_{spec}.bin")
data = np.fromfile(emis_file)
data = data.reshape((2, nlon, nlat, ntime), order="F")
# kg/m2/s -> kg (by time step)
emis_data = data[0].T * areas[np.newaxis, ...] * dt
emis_data_tl = data[1].T * areas[np.newaxis, ...] * dt
emis[index, ...] = emis_data
emis_tl[index, ...] = emis_data_tl
# Uniform increment
emis = emis.sum(['lat', 'lon'])
emis_tl = emis_tl.sum(['lat', 'lon'])
return emis, emis_tl
[docs]
def time_slice(
da: xr.DataArray,
di: np.datetime64,
df: np.datetime64,
agg_method: Literal["mean", "sum"] = "mean"
) -> xr.DataArray:
"""Pick the values of a DataArray between 'di' and 'df' and aggregate along
the time dimension
Args:
da (xr.DataArray): input DataArray with 'time' dimension and coordinate
di (np.datetime64): slice start
df (np.datetime64): slice end (NOT included)
agg_method ("mean" or "sum", optional): aggregation method. Defaults to "mean".
Returns:
xr.DataArray: sliced and aggregated DataArray
"""
# sliced_da = da.sel(time=slice(di, df), method='ffill')
sliced_da = da.sel(time=slice(di, df - NANO))
if agg_method == "mean":
return sliced_da.mean('time')
elif agg_method == "sum":
return sliced_da.sum('time')
else:
raise CifValueError(f"unexpected agg_method, '{agg_method}'")
[docs]
def restart_id_to_var_name(restart_id: Union[int, str]) -> str:
"""Convert a LMDZ restart ID to a NetCDF variable name.
Args:
restart_id: integer (→ ``'qNN'``) or string (used as-is).
Returns:
str: NetCDF variable name.
"""
if isinstance(restart_id, str):
return restart_id
else:
return f'q{restart_id:02d}'
[docs]
def get_var_name(self: Model, spec: str) -> str:
"""Return the restart NetCDF variable name for species *spec*.
Args:
self: LMDZ-old model plugin instance.
spec: active species name.
Returns:
str: NetCDF variable name.
"""
spec_plg = getattr(self.chemistry.acspecies, spec)
restart_id = spec_plg.restart_id
return restart_id_to_var_name(restart_id)
[docs]
def perturb_ref_restart(
self: Model,
ddi: datetime.datetime,
restart_file: str,
ref_restart: str
) -> None:
"""Write a perturbed LMDZ-old restart file by copying from a reference.
Copies *ref_restart* to *restart_file* and overwrites active-species
variables with CIF-modified concentrations. Species in
``self.dont_perturb_species`` are left unchanged.
Args:
self: LMDZ-old model plugin instance.
ddi: period start date.
restart_file: output restart path.
ref_restart: reference restart path to copy from.
"""
lmdz_version = self.plugin.version
dont_perturb_spec = getattr(self, 'dont_perturb_species', [])
with _hdf5_lock:
with xr.open_dataset(ref_restart) as ds:
# Looping over original species
for spec in self.original_species:
restart_id = self.original_restart_ids[spec]
var_name = restart_id_to_var_name(restart_id)
da = xr.DataArray(data=ds[var_name].values,
dims=['time', 'lev', 'lat', 'lon'])
if spec in dont_perturb_spec:
target_var_name = get_var_name(self, spec)
self.inicond.write(target_var_name, restart_file, da)
else:
# Looping over perturbated species of species 'spec'
for spec_sample in self.chemistry.acspecies.attributes:
if self.perturbed_species.get(spec_sample, "_") != spec:
continue
target_var_name = get_var_name(self, spec_sample)
self.inicond.write(target_var_name, restart_file, da)
if lmdz_version == "std":
update_controle_var(self, ddi, restart_file, ref_restart)
[docs]
def make_end(
self: Model,
runsubdir: str,
ddi: datetime.datetime,
ref_fwd_dir: str
) -> None:
"""TODO
Args:
self (Model)
ddi (datetime.datetime): sub simulation start datetime
runsubdir (str): sub simulation run directory
ref_fwd_dir (str): reference forward run directory
"""
lmdz_version = self.plugin.version
nspec = len(self.chemistry.acspecies.attributes)
if not os.path.isdir(ref_fwd_dir):
raise CifFileNotFoundError(
f"reference forward run directory '{ref_fwd_dir}' not found, "
"a reference forward run of the model is needed to run the "
"approximated operator"
)
# Fetch original end from reference forward
restart_file = os.path.join(runsubdir, "restart.nc")
ref_restart = os.path.join(ref_fwd_dir, "chain",
ddi.strftime("restart_%Y%m%d%H%M.nc"))
if hasattr(self, 'perturbed_species'):
# Assuming restart files from reference forward run are not perturbated
perturb_ref_restart(self, ddi, restart_file, ref_restart)
else:
link(ref_restart, restart_file)
# Initial conditions, dims: (spec, lev, lat, lon)
mmr = read_inicond_file(self, runsubdir, "fwd")
mmr_tl = read_inicond_file(self, runsubdir, "tl")
assert nspec == mmr.shape[0]
# Air mass (with double precision), dims: (time, lev, lat, lon)
ref_mass = read_air_mass(self, ddi, runsubdir)
# The atmosphere is considered as well mixed (i.e. uniform)
ones = mmr.copy(data=np.ones(mmr.shape))
tot_mass = ref_mass.isel(time=0).sum(['lev', 'lat', 'lon'])
mmr_scalar = (ref_mass.isel(time=0) *
mmr).sum(['lev', 'lat', 'lon']) / tot_mass # dims: (spec)
mmr_scalar_tl = (ref_mass.isel(time=0) *
mmr_tl).sum(['lev', 'lat', 'lon']) / tot_mass # dims: (spec)
# Emissions, dims: (time) [kg]
ref_emis, ref_emis_tl = read_emissions(self, ddi, runsubdir)
# 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_tl = ref_prescr.copy(data=np.zeros(ref_prescr.shape))
ref_prescr = ref_prescr * ref_pmid / (BOLTZMAN * ref_temp) * 1.0e-6
# Observations time step (LMDZ time step)
nobs, obs_lmdz_tstep, spec_index = read_obs(runsubdir)
# LMDZ time steps datetimes
tstep_dates = np.array(self.tstep_dates[ddi][:-1], dtype='datetime64[ns]')
# Approximation intergration time step datetimes
ddf = self.input_dates[ddi][-1]
time_periods = date_range(ddi, ddf, period=self.approx_time_step)
time_periods = np.array(time_periods, dtype='datetime64[ns]')
periods_dt = pd.to_timedelta(np.diff(time_periods)).total_seconds()
# LMDZ time step to approximation time step (only using 'tstep','dstep' is
# ignored there)
diff = tstep_dates[:, np.newaxis] - time_periods[np.newaxis, :-1]
diff = diff.astype('float64')
diff[diff < 0] = np.nan
obs_approx_tstep = np.nanargmin(diff, axis=1)
sim_tl = np.zeros(nobs)
ddi_str = ddi.strftime("%Y-%m-%d")
debug(f"Approximating operator for sub-simulation {ddi_str} with "
f"'{self.approx_time_step}' time-step")
# Newton scheme loop
for i, (dt, start, end) in enumerate(zip(periods_dt, time_periods[:-1], time_periods[1:])):
# Filling sim_tl with values at current time step start
ppm_tl = 1e6 * mmr_scalar_tl * DRY_MASS / molar_masses # kg/kg -> ppm
ppm_tl[ppm_tl < 0.0] = 0.0
# Getting observation indices corresponding to the current time step
time_mask = obs_approx_tstep[obs_lmdz_tstep] == i
for index in range(nspec):
# Getting observation indices corresponding to the current time
# step AND species index
obs_mask = time_mask & (spec_index == index)
sim_tl[obs_mask] = ppm_tl.values[index]
# Slicing fields for the current time step
prescr = time_slice(ref_prescr, start, end)
prescr_tl = time_slice(ref_prescr_tl, start, end)
pmid = time_slice(ref_pmid, start, end)
temp = time_slice(ref_temp, start, end)
mass = time_slice(ref_mass, start, end)
emis = time_slice(ref_emis, start, end, "sum")
emis_tl = time_slice(ref_emis_tl, start, end, "sum")
start_str = pd.to_datetime(start).isoformat()
end_str = pd.to_datetime(end).isoformat()
debug(
f"\nbefore time step [{start_str} -> {end_str}], dt = {dt} seconds:\n"
f" before species: {self.chemistry.acspecies.attributes}\n"
f" before mmr (fwd): {mmr_scalar.values.tolist()}\n"
f" before mmr (tl): {mmr_scalar_tl.values.tolist()}\n"
f" before ppm (tl): {ppm_tl.values.tolist()}\n"
)
# Chemistry
delta_chem, delta_chem_tl = compute_chemistry_step(
reaction_list, molar_masses, dt,
mmr_scalar * ones, mmr_scalar_tl * ones,
prescr, prescr_tl, pmid, temp
)
# The atmosphere is considered as well mixed (i.e. uniform)
tot_mass = mass.sum(['lev', 'lat', 'lon'])
delta_chem = (mass * delta_chem).sum(['lev', 'lat', 'lon'])
delta_chem_tl = (mass * delta_chem_tl).sum(['lev', 'lat', 'lon'])
debug(
f"\ntime step [{start_str} -> {end_str}], dt = {dt} seconds:\n"
f" species: {self.chemistry.acspecies.attributes}\n"
f" emis (fwd): {(emis / tot_mass).values.tolist()}\n"
f" emis (tl): {(emis_tl / tot_mass).values.tolist()}\n"
f" chem (fwd): {(delta_chem / tot_mass).values.tolist()}\n"
f" chem (tl): {(delta_chem_tl / tot_mass).values.tolist()}\n"
f" mmr (fwd): {mmr_scalar.values.tolist()}\n"
f" mmr (tl): {mmr_scalar_tl.values.tolist()}\n"
f" ppm (tl): {ppm_tl.values.tolist()}\n"
)
# Apply Newton scheme time step
mmr_scalar += (emis + delta_chem) / tot_mass
mmr_scalar_tl += (emis_tl + delta_chem_tl) / tot_mass
mmr_scalar_tl[mmr_scalar_tl < 0.0] = 0.0
if np.all(mmr_scalar_tl.values == 0.0):
break
# Writing tangent restart file
increment = (mmr_scalar_tl * ones).values.T
restart_tl_file = os.path.join(runsubdir, "restart_tl.bin")
with open(restart_tl_file, "bw") as f:
if lmdz_version == "std":
f.write(b"0000")
increment.flatten(order="F").tofile(f)
# Writing observations output file
data = np.zeros((nobs, 6))
data[:, 1] = sim_tl # Putting sim_tl in the corresponding column
sim_file = os.path.join(runsubdir, "obs_out.bin")
data.flatten(order="F").tofile(sim_file)
