from __future__ import annotations
import datetime
from os import PathLike
from pathlib import Path
from typing import Literal
import pandas as pd
import xarray as xr
from ......utils import path
from ......utils.check.errclass import CifRuntimeError
from ......utils.hdf5 import _hdf5_lock
from ......utils.namelist import format_fortran_namelist
[docs]
def make_parameters_file(
self,
datei: pd.Timestamp | datetime.datetime,
datef: pd.Timestamp | datetime.datetime,
runsubdir: str | PathLike,
mode: Literal["fwd", "tl", "adj"],
) -> None:
"""Write the LMDZ-ico Fortran parameters namelist for one sub-period.
Generates the ``parameters.nml`` file inside *runsubdir* containing
domain dimensions, time-step counts, and I/O paths needed by the LMDZ
icosahedral-grid executable.
Args:
self: LMDZ-ico model plugin instance.
datei: period start date.
datef: period end date.
runsubdir: path to the period run directory.
mode: ``'fwd'``, ``'tl'``, or ``'adj'``.
"""
datei = pd.to_datetime(datei)
datef = pd.to_datetime(datef)
assert datei < datef
assert (
datei.hour % 3 == 0
and datei.minute == 0
and datei.second == 0
and datei.microsecond == 0
)
assert (
datef.hour % 3 == 0
and datef.minute == 0
and datef.second == 0
and datef.microsecond == 0
)
namelists = {
"dims": {
"ntrac": len(self.chemistry.active_species),
},
"tracers": {
"tracer_species": list(self.chemistry.active_species),
"tracer_flux": [
spec in self.chemistry.emitted_species
for spec in self.chemistry.active_species
],
"tracer_molar_mass": [
spec.molar_mass for spec in self.chemistry.active_species.values()
],
},
"times": {
"datetime_begin": f"{datei:%Y-%m-%d %H:%M:%S}",
"datetime_end": f"{datef:%Y-%m-%d %H:%M:%S}",
"nsplit_dyn": self.nsplit_dyn,
"nsplit_phy": self.nsplit_phy,
},
"flags": {
"forward": mode in ("fwd", "tl"),
"tangent": mode == "tl",
"do_physics": self.do_physics,
"do_chemistry": self.do_chemistry,
"read_start_ad": mode == "adj" and datei != self.subsimu_dates[-2],
"save_trajq": mode == "adj" or not self.no_trajectory_file,
"use_float": self.traj_dtype == "float32",
"save_mass": mode == "fwd" and self.dump_total_mass,
"save_sinks": mode == "fwd" and self.dump_chemical_sink,
"save_lifetime": mode == "fwd" and self.dump_chemical_lifetime,
},
}
with open(Path(runsubdir, "parameters.nml"), "w") as f:
for name, variables in namelists.items():
f.write(format_fortran_namelist(name, variables))
[docs]
def make_auxiliary(
self,
ddi: datetime.datetime,
runsubdir: str | PathLike, # noqa: E501
mode: Literal["fwd", "tl", "adj"] = "fwd",
onlyinit: bool = False,
do_simu: bool = True,
**kwargs,
) -> None:
"""Write auxiliary files for one sub-simulation period"""
# If mode is "fwd" or "tl" but onlyinit is True,
# it means that we are initializing an adjoint by running backward
# the adjoint pipeline
if mode in ["tl", "fwd"] and onlyinit:
mode = "adj"
if not do_simu or onlyinit:
return
runsubdir = Path(runsubdir)
subsimu_ddi, subsimu_ddf = self.subsimu_intervals[ddi]
# Linking executable
path.link(self.model_dir / "dispersion.e", runsubdir / "dispersion.e")
# Linking the domain file (create it if needed)
domain_file = self.model_dir / "domain.nc"
if not domain_file.is_file():
domain_ds = xr.Dataset(coords=self.domain.get_domain_coords(vertical=True))
with _hdf5_lock:
domain_ds.to_netcdf(domain_file)
path.link(domain_file, runsubdir / "domain.nc")
if self.grid == "dynamico":
path.link(self.model_dir / "mesh.nc", runsubdir / "mesh.nc")
# Dump parameters files
make_parameters_file(self, subsimu_ddi, subsimu_ddf, runsubdir, mode)
if self.do_chemistry:
path.link(
self.chemistry.chem_dir / "chemical_scheme.nml",
runsubdir / "chemical_scheme.nml",
)
# Link to trajq from corresponding forward simulation
if mode == "adj":
if not hasattr(self, "adj_refdir"):
raise CifRuntimeError(
"Adjoint LMDZ couldn't be initialized with forward trajq.bin files"
)
else:
for spec in self.chemistry.active_species:
filename = f"trajq_{spec}_{ddi:%Y-%m-%d-%H-00}.bin"
source = Path(self.adj_refdir, "chain", filename)
target = runsubdir / f"trajq_{spec}.bin"
path.link(source, target)
