import filecmp
import os
import shutil
import pathlib
import numpy as np
from netCDF4 import Dataset
import xarray as xr
from logging import warning
import copy
from ......utils import path
from ......utils.hdf5 import _hdf5_lock
from .utils import replace_dates
from ......utils.check.errclass import CifError
[docs]
def make_inicond(
self, datastore, runsubdir, mode, datei, datef,
ini_type="inicond"
):
"""Write or symlink the CHIMERE initial concentration file ``INI_CONCS.0.nc``.
For the **first sub-simulation period** only (``datei == self.datei``).
For each active species:
* If no data was modified in the CIF datastore, symlinks the original
file from disk.
* If concentrations were modified (e.g. by a prior transform), copies
the original file and overwrites the relevant species variable.
* For **tangent-linear** mode, also creates ``INI_CONCS.0.increment.nc``
initialised to zero perturbations.
* Handles **ensemble/perturbed** species by reading from the reference
species and writing to the perturbed copy.
After all species are written, calls :func:`replace_dates` to ensure
the NetCDF time axis matches the CHIMERE convention.
Args:
self: CHIMERE model plugin instance (carries ``datei``, ``nho``,
``adj_refdir``, ``chemistry.acspecies``, and ``inicond``).
datastore (dict): tracer-ID-keyed mapping of CIF data-store entries.
runsubdir (str): path to the period run directory.
mode (str): ``'fwd'``, ``'tl'``, or ``'adj'``.
datei (datetime): period start date.
datef (datetime): period end date.
ini_type (str): datastore key prefix; ``'inicond'`` for normal use
or ``'restart_inicond'`` in EnSRF restart mode.
"""
ddi = min(datei, datef)
# Fixed name for INI_CONCS files
fileout = f"{runsubdir}/INI_CONCS.0.nc"
fileoutincr = f"{runsubdir}/INI_CONCS.0.increment.nc"
nho = self.nho
# # Getting the forward initial concentrations (needed even for adjoint)
# # if chained period
# subsimu_dates = self.subsimu_dates
# if mode == "adj" and hasattr(self, "fwd_chain") \
# and datei != self.datei:
# date_index = np.where(subsimu_dates == datei)[0][0]
# refdir = self.adj_refdir
# filein = subsimu_dates[date_index - 1].strftime(
# "{}/chain/end.%Y%m%d%H.{}.nc".format(refdir, nho)
# )
# path.link(filein, fileout)
# Exit if not first period
if datei != self.datei:
return
# If in datastore, take data, otherwise, link to original INI_CONCS
data2dump = {}
data2dump_tl = {}
for spec in self.chemistry.acspecies.attributes:
trid = (ini_type, spec)
ref_trid = (ini_type, spec)
# Update ref_trid if in ensemble mode
perturbed_from_reference = False
if hasattr(self, "perturbed_species"):
if spec in self.perturbed_species:
ref_trcr = self.perturbed_species[spec]
ref_trid = (ini_type, ref_trcr)
# If spec not explicitly defined in datastore,
# fetch general component information if available
if trid in datastore:
pass
# If trid is a perturbed species, needs to reformat name to fit
# CHIMERE naming convention for perturbed species
elif trid in [(k[0], k[1].replace("__sample#", "")) for k in datastore]:
trid_ind = [
(k[0], k[1].replace("__sample#", "")) for k in datastore
].index(trid)
trid = list(datastore.keys())[trid_ind]
# If trid comes from an unperturbed input, but in ensemble mode
# needs to fetch info from the reference species
elif ref_trid in datastore:
trid = ref_trid
perturbed_from_reference = True
elif trid not in datastore and (ini_type, "") in datastore:
trid = (ini_type, "")
else:
continue
tracer = datastore[trid]
tracer_data = tracer["data"][ddi]
# If no data is provided, just copy from original file
if "spec" not in tracer_data:
dirorig = tracer["dirorig"]
fileorig = tracer["fileorig"]
fileini = f"{dirorig}/{fileorig}"
# If does not exist, just link
linked = False
if not os.path.isfile(fileout):
path.link(fileini, fileout)
linked = True
# Simply link if restart_inicond
if ini_type == "restart_inicond":
continue
# Otherwise, check for difference
if not linked or perturbed_from_reference:
if not os.path.isfile(fileini):
warning(f"The initial condition file {fileini} does not exist to initialize the species {spec}. The concentrations will be initialized to zero. Please check your yaml if you expect a different behaviour")
continue
if not filecmp.cmp(fileini, fileout) or perturbed_from_reference:
spec2read = spec if not perturbed_from_reference else ref_trid[1]
with _hdf5_lock:
with Dataset(fileini, "r") as ds:
if spec2read not in ds.variables:
warning(f"{spec} is not accounted for in the initial conditions file {fileini}. Please check your IC whether it should be")
continue
if "Time" in [d.name for d in ds.variables[spec2read].get_dims()]:
raise CifError(
f"Warning! The original file {fileini} is not in the "
f"correct format. More specifically, the dimensions of "
f"the variable {spec2read} include 'Time' whereas it should be"
f" (bottom_top, south_north, west_east)"
)
iniin = ds.variables[spec2read][:]
iniin = xr.DataArray(
iniin[np.newaxis, ...],
coords={"time": [min(datei, datef)]},
dims=("time", "lev", "lat", "lon"),
)
# If ini file is still a link, should be copied
# to be able to modify it locally
if os.path.islink(fileout):
file_orig = pathlib.Path(fileout).resolve()
os.unlink(fileout)
shutil.copy(file_orig, fileout)
# Now saves to buffer variable before saving
data2dump[spec] = copy.deepcopy(iniin)
# Repeat operations for tangent linear
if mode == "tl":
# If does not exist, copy
if not os.path.isfile(fileoutincr):
shutil.copy(fileini, fileoutincr)
with _hdf5_lock:
with Dataset(fileoutincr, "a") as fout:
if spec in fout.variables:
fout.variables[spec][:] = 0.0
continue
with xr.open_dataset(fileout) as ds:
if spec not in ds:
continue
iniin = xr.DataArray(
0. * ds[spec].values[np.newaxis, ...],
coords={"time": [min(datei, datef)]},
dims=("time", "lev", "lat", "lon"),
)
data2dump_tl[spec] = copy.deepcopy(iniin)
else:
# Replace existing link by copy
# of original file to modify it
path.copyfromlink(fileout)
# Write initial conditions
ini_fwd = tracer_data["spec"]
data2dump[spec] = copy.deepcopy(ini_fwd)
if mode == "tl":
path.copyfromlink(fileoutincr)
ini_tl = tracer_data.get("incr", 0.0 * ini_fwd)
data2dump_tl[spec] = copy.deepcopy(ini_tl)
# Now dump buffer data if any
if len(data2dump) > 0:
self.inicond.write(
list(data2dump.keys()), fileout, data2dump, comp_type="inicond"
)
# Now dump buffer data if any
if len(data2dump_tl) > 0:
self.inicond.write(
list(data2dump_tl.keys()), fileoutincr, data2dump_tl, comp_type="inicond"
)
# Check that the dates are consistent with what CHIMERE expects
replace_dates(fileout, [min(datei, datef)], self.ignore_input_dates)
if mode == "tl":
replace_dates(fileoutincr, [min(datei, datef)],
self.ignore_input_dates)
