from __future__ import annotations
import datetime
import shutil
from logging import debug
from os import PathLike
from pathlib import Path
from typing import Any, Literal
import numpy as np
import pandas as pd
import xarray as xr
from ......utils.check.errclass import CifFileNotFoundError, CifValueError
from ......utils.dates import date_range
from ......utils.hdf5 import _hdf5_lock
from ...chemistry import (
BOLTZMAN,
DRY_MASS,
compute_chemistry_step,
parse_chemical_scheme,
read_kinetic,
read_prescr,
strip_spatial_coords,
)
NANO = np.timedelta64(1, "ns")
# Aliases for type hinting
Model = Any
[docs]
def read_obs(self, runsubdir: str | PathLike) -> tuple[int, np.ndarray, np.ndarray]:
"""Reads time step and species index columns from 'obs.nc"""
obs_file = Path(runsubdir, "obs.nc")
with _hdf5_lock:
with xr.open_dataset(obs_file) as ds:
nobs = ds.sizes["index"]
tstep = (ds["time"].values / self.dt.total_seconds()).astype(int)
spec_index = ds["itarac"].values
return nobs, tstep, spec_index
[docs]
def read_air_mass(
self,
ddi: datetime.datetime,
runsubdir: str | PathLike,
) -> xr.DataArray:
"""Read the LMDZ-ico air-mass field from the run directory"""
di = pd.to_datetime(ddi)
di = di if di.is_month_start else di - pd.offsets.MonthBegin()
freq = self.mass_fluxes_freq
ntime = int(di.days_in_month * 24 * 3600 / freq.total_seconds()) # type: ignore
time = xr.DataArray(data=pd.date_range(di, periods=ntime, freq=freq), dims=["time"])
mass_file = Path(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,
ddi: datetime.datetime,
runsubdir: str | PathLike,
) -> tuple[xr.DataArray, xr.DataArray]:
"""Reads emissions "mod_*.bin" binary files"""
di = pd.to_datetime(ddi)
di = di if di.is_month_start else di - pd.offsets.MonthBegin()
if not hasattr(self.domain, "areas"):
self.domain.calc_areas()
areas = self.domain.areas
dt = pd.to_timedelta(self.flx_tresol).total_seconds()
emis = None
emis_tl = None
for spec in self.chemistry.active_species:
emis_file = Path(runsubdir, f"flux_{spec}.nc")
with _hdf5_lock:
with xr.open_dataset(emis_file) as ds:
emis_spec = ds[spec] * areas[np.newaxis, ...] * dt
emis_tl_spec = ds[f"{spec}_tl"] * areas[np.newaxis, ...] * dt
emis_spec = emis_spec.sum(["lat", "lon"]).expand_dims("spec")
emis_tl_spec = emis_tl_spec.sum(["lat", "lon"]).expand_dims("spec")
if emis is None and emis_tl is None:
emis = emis_spec
emis_tl = emis_tl_spec
else:
emis = xr.concat([emis, emis_spec], dim="spec") # type: ignore
emis_tl = xr.concat([emis_tl, emis_tl_spec], dim="spec") # type: ignore
return emis, emis_tl # type: ignore
[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 perturb_ref_restart(
self,
restart_file: str | PathLike,
ref_restart: str | PathLike,
) -> None:
"""Write a perturbed LMDZ-ico restart file by copying from a reference.
Copies *ref_restart* to *restart_file* and overwrites active-species
mass mixing ratios with CIF-modified concentrations.
Args:
self: LMDZ-ico model plugin instance.
restart_file: output restart path.
ref_restart: reference restart path to copy from.
"""
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:
da = xr.DataArray(data=ds[spec].values, dims=["time", "lev", "lat", "lon"])
if spec in dont_perturb_spec:
self.inicond.write(spec, restart_file, da)
else:
# Looping over perturbated species of species 'spec'
for spec_sample in self.active_species:
if self.perturbed_species.get(spec_sample, "_") != spec:
continue
self.inicond.write(spec_sample, restart_file, da)
[docs]
def make_end(
self,
runsubdir: str | PathLike,
ddi: datetime.datetime,
ref_fwd_dir: str | PathLike,
) -> 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
"""
# Dimensions
nspec = len(self.chemistry.active_species)
nlev = self.domain.nlev
nlat = self.domain.nlat
nlon = self.domain.nlon
runsubdir = Path(runsubdir)
ref_fwd_dir = Path(ref_fwd_dir)
nspec = len(self.chemistry.active_species)
if not ref_fwd_dir.is_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"
)
# Initial conditions, dims: (spec, lev, lat, lon)
mmr = xr.DataArray(
data=np.zeros((nspec, nlev, nlat, nlon)),
dims=["spec", "lev", "lat", "lon"],
)
mmr_tl = xr.DataArray(
data=np.zeros((nspec, nlev, nlat, nlon)),
dims=["spec", "lev", "lat", "lon"],
)
# Reading initial conditions from NetCDF file
with _hdf5_lock:
with xr.open_dataset(runsubdir / "start.nc") as ds:
for index, spec in enumerate(self.chemistry.active_species):
mmr[index, ...] = ds[spec].values
mmr_tl[index, ...] = ds[f"{spec}_tl"].values
# 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)
if self.do_chemistry and reaction_list:
# 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(self, 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() # type: ignore
# 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-%H-00")
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
if self.do_chemistry and reaction_list:
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
if self.do_chemistry and reaction_list:
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"])
else:
delta_chem = xr.zeros_like(emis)
delta_chem_tl = xr.zeros_like(emis)
debug(
f"\ntime step [{start_str} -> {end_str}], dt = {dt} seconds:\n"
f" species: {list(self.chemistry.active_species)}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
# Fetch original end from reference forward
restart_file = Path(runsubdir, "restart.nc")
ref_restart = Path(ref_fwd_dir, "chain", ddi.strftime("restart_%Y-%m-%d.nc"))
if hasattr(self, "perturbed_species"):
# Assuming restart files from reference forward run are not perturbated
perturb_ref_restart(self, restart_file, ref_restart)
else:
shutil.copy(ref_restart, restart_file)
ds_restart = xr.Dataset(
{
f"{spec}_tl": increment.isel(spec=index)
for index, spec in enumerate(self.chemistry.active_species)
}
)
with _hdf5_lock:
ds_restart.to_netcdf(restart_file, mode="a")
# Writing observations output file
ds_obs = xr.Dataset(
{
"spec": (["index"], np.zeros(nobs)),
"incr": (["index"], np.full(nobs, sim_tl)),
"pressure": (["index"], np.zeros(nobs)),
"dpressure": (["index"], np.zeros(nobs)),
"hlay": (["index"], np.zeros(nobs)),
"airm": (["index"], np.zeros(nobs)),
}
)
with _hdf5_lock:
ds_obs.to_netcdf(runsubdir / "obs_out.nc")
