import datetime
import os
from logging import warning
from typing import Optional
import numpy as np
import xarray as xr
from ......utils import path
from ......utils.hdf5 import _hdf5_lock
from .ensemble import ensemble_trid
from ......utils.check.errclass import CifValueError
[docs]
def get_spec_var_name(self, spec: str) -> str:
"""Return the LMDZ restart-file variable name for species *spec*.
Reads ``spec_obj.restart_id`` from the chemistry plugin. If
``restart_id`` is an int it is formatted as ``qNN``; if a string it is
used directly.
Args:
self: LMDZ-old model plugin instance.
spec (str): active species name.
Returns:
str: variable name inside the LMDZ restart NetCDF file.
"""
# Getting restart id
spec_obj = getattr(self.chemistry.acspecies, spec)
restart_id = spec_obj.restart_id
if isinstance(restart_id, str):
var_name = restart_id
else:
var_name = f'q{restart_id:02d}'
return var_name
[docs]
def write_inicond(self, var_name: str, filename: str, data: xr.Dataset) -> None:
"""Write one species' initial-condition array into a LMDZ restart file.
Validates that *data* contains no NaN values before delegating to
``self.inicond.write``.
Args:
self: LMDZ-old model plugin instance.
var_name (str): variable name in the restart file.
filename (str): path to the restart NetCDF file to write.
data (xr.Dataset): concentration data to write.
Raises:
ValueError: if *data* contains NaN values.
"""
if np.any(np.isnan(data)):
raise CifValueError(
f"error while writing inicond file '{filename}' variable "
f"'{var_name}', there is NaN values in the data")
self.inicond.write(var_name, filename, data)
[docs]
def update_controle_var(
self,
ddi: datetime.datetime,
target_path: str,
source_path: Optional[str] = None
) -> None:
"""Update start file with the controle variable from the configuration
file 'model' paragraph or the 'controle' variable in 'source_path' file
Args:
ddi (datetime)
source_path (str)
target_path (str)
"""
controle_var = None
if hasattr(self, 'controle_var'):
controle_var = self.controle_var
elif source_path is not None:
with _hdf5_lock:
with xr.open_dataset(source_path) as ds:
if 'controle' in ds:
controle_var = ds['controle']
if controle_var is None:
warning(
"could not retrieve 'controle' variable form the "
"configuration file 'model' paragraph or 'inicond' file"
)
return
controle_var[4] = ddi.year
controle_var[5] = (ddi - datetime.datetime(ddi.year, 1, 1)).days
with _hdf5_lock:
controle_var.to_netcdf(target_path, mode='a')
[docs]
def make_inicond(
self, datastore, datei, datef, runsubdir, mode, onlyinit,
ini_type="inicond"
):
"""Write or symlink the LMDZ-old initial-condition restart file.
For the first sub-simulation period only (``datei == self.datei``).
If concentrations in *datastore* were not modified by CIF, symlinks the
original restart file. Otherwise copies the file and overwrites the
relevant species variables. Updates the ``controle`` variable to align
LMDZ's internal clock with *datei*.
Args:
self: LMDZ-old model plugin instance.
datastore (dict): tracer-ID-keyed CIF data-store entries.
datei (datetime): period start date.
datef (datetime): period end date.
runsubdir (str): path to the period run directory.
mode (str): ``'fwd'``, ``'tl'``, or ``'adj'``.
onlyinit (bool): skip if ``True``.
ini_type (str): datastore key prefix.
"""
if not hasattr(self, 'chemistry'):
raise CifValueError("no chemical scheme. LMDZ needs a chemical scheme")
ddi = min(datei, datef)
# In adjoint mode (not onlyinit), do nothing as no concentration is needed
if mode == "adj" and not onlyinit:
return
# Generating reference initial conditions if first sub-simulation
for spec in self.chemistry.acspecies.attributes:
trid = ensemble_trid(self, (ini_type, spec), datastore)
if trid not in datastore:
continue
target_path = os.path.join(runsubdir, "start.nc")
target_path_tl = os.path.join(runsubdir, "start_tl.nc")
updated_controle_var = False
var_name = get_spec_var_name(self, spec)
data = datastore[trid]['data'][ddi]
if "spec" not in data:
ref_dir = datastore[trid]["dirorig"]
ref_file = datastore[trid]["fileorig"]
source_path = os.path.join(ref_dir, ddi.strftime(ref_file))
# If restart_inicond restart file is already formatted and contains
# all species variables -> simply link the restart file
if ini_type == "restart_inicond":
path.link(source_path, target_path)
continue
# Reading input file with datavect plugin 'read' method
input_tracer = datastore[trid]['tracer']
da = input_tracer.read(
trid[1],
input_tracer.varname,
dates=datastore[trid]['input_dates'][ddi],
files=datastore[trid]['input_files'][ddi],
tracer=input_tracer
)
# Writing data
write_inicond(self, var_name, target_path, da)
if self.plugin.version == "std" and not updated_controle_var:
update_controle_var(self, ddi, target_path, source_path)
if mode == "tl":
warning("Setting LMDZ tangent-linear initial conditions to "
f"zeros for {spec}.")
da_tl = da.copy(data=np.zeros(da.shape))
write_inicond(self, var_name, target_path_tl, da_tl)
if self.plugin.version == "std" and not updated_controle_var:
update_controle_var(self, ddi, target_path_tl, source_path)
updated_controle_var = True
# Updates data
elif mode in ["fwd", "tl"] and ddi == self.datei:
write_inicond(self, var_name, target_path, data['spec'])
if self.plugin.version == "std" and not updated_controle_var:
update_controle_var(self, ddi, target_path)
if mode == "tl":
if 'incr' in data:
write_inicond(self, var_name, target_path_tl, data['incr'])
else:
warning("Setting LMDZ tangent-linear initial conditions "
f"to zeros for {spec}.")
with _hdf5_lock:
with xr.open_dataset(target_path) as ds_target:
da = ds_target[var_name]
da_tl = da.copy(data=np.zeros(da.shape))
write_inicond(self, var_name, target_path_tl, da_tl)
if self.plugin.version == "std" and not updated_controle_var:
update_controle_var(self, ddi, target_path_tl)
updated_controle_var = True
# The part below is used by the old DISPERSION only
elif mode == "adj" and self.plugin.version == "std":
for spec in self.chemistry.acspecies.attributes:
warning("WARNING: setting initial conditions to zero. "
"LMDZ TL will not pass the test of the adjoint")
da = data['spec']
da_adj = da.copy(data=np.zeros(da.shape))
var_name = get_spec_var_name(self, spec)
write_inicond(self, var_name, target_path, da_adj)
if self.plugin.version == "std" and not updated_controle_var:
update_controle_var(self, ddi, target_path)
updated_controle_var = True
