import xarray as xr
import numpy as np
import os
from ......utils.check.errclass import CifError
from ......utils.hdf5 import _hdf5_lock
[docs]
def fetch_end(self, data2dump, runsubdir, mode, ddi, ddf,
check_transforms=False, onlyinit=False):
"""Register end-concentration output paths (and optionally read data) after a CHIMERE run.
For **adjoint** mode: records the ``aend.*.nc`` file path in *data2dump*,
and when *check_transforms* is active, reads the adjoint sensitivity
fields (``_ad_o`` and ``_ad``) from the distributed tile layout back into
the data store.
For **forward/TL** mode: records the ``end.*.nc`` file path, and when
*check_transforms* is active in TL mode, reads the TL increment
(``_o_tl`` and ``_tl``) back into the data store.
Args:
self: CHIMERE model plugin instance (carries ``nho``, ``domain``,
``nzdoms``, ``nmdoms``).
data2dump (dict): tracer-ID-keyed data-store entries to update.
runsubdir (str): path to the period run directory.
mode (str): ``'fwd'``, ``'tl'``, or ``'adj'``.
ddi (datetime): period start date.
ddf (datetime): period end date (unused, kept for API consistency).
check_transforms (bool): if ``True``, also read output fields back
into *data2dump* (used by the transform-checker).
onlyinit (bool): if ``True``, skip reading fields (initialisation
pass).
Returns:
dict: updated *data2dump* with ``fileorig`` paths (and optionally
concentration / sensitivity arrays) set.
"""
if mode == "adj":
fileorig = \
ddi.strftime(
f"chain/aend.%Y%m%d%H.{self.nho}.nc")
for trid in data2dump:
data2dump[trid]["data"][ddi]["fileorig"] = fileorig
# Read sensitivity when checking transforms
if check_transforms and not onlyinit:
# Initialize sub-domains like in master_init_mpi
nmerid, nzonal = self.domain.zlon.shape
nzdoms = self.nzdoms
nmdoms = self.nmdoms
nz_tile = list(range(0, nzonal + 1, int(nzonal / nzdoms)))
if len(nz_tile) == 1:
nz_tile = [nzonal]
else:
nz_tile[-1] = nzonal
nz_tile = np.diff(nz_tile)
redistrib_nz = nz_tile[-1] - nz_tile[0]
nz_tile[-1] -= int(np.ceil(redistrib_nz / 2))
nz_tile[0] += int(np.ceil(redistrib_nz / 2))
nm_tile = list(range(0, nmerid, int(nmerid / nmdoms)))
if len(nm_tile) == 1:
nm_tile = [nmerid]
else:
nm_tile[-1] = nmerid
nm_tile = np.diff(nm_tile)
redistrib_nm = nm_tile[-1] - nm_tile[0]
nm_tile[-1] -= int(np.ceil(redistrib_nm / 2))
nm_tile[0] += int(np.ceil(redistrib_nm / 2))
# Create indexing with 3-cell buffers between tiles
cum_index_i = [0] + list(np.cumsum(nm_tile))
list_i = [
ii for i in range(nmdoms)
for ii in list(range(
i * 6 + cum_index_i[i] + 3,
i * 6 + cum_index_i[i] + 3 + nm_tile[i]
))
]
cum_index_j = [0] + list(np.cumsum(nz_tile))
list_j = [
jj for j in range(nzdoms)
for jj in list(range(
j * 6 + cum_index_j[j] + 3,
j * 6 + cum_index_j[j] + 3 + nz_tile[j]
))
]
with _hdf5_lock:
data2dump[trid]["data"][ddi]["adj_out"] = np.concatenate([
xr.open_dataset(f"{runsubdir}/../{fileorig}")[
f"{trid[1]}_ad_o"].values[:, :-1, :nmerid, :nzonal],
xr.open_dataset(f"{runsubdir}/../{fileorig}")[
f"{trid[1]}_ad"].values[:, :-1, list_i][..., list_j]
], axis=0)
return data2dump
else:
dataout = {}
fileorig = \
ddi.strftime(
f"chain/end.%Y%m%d%H.{self.nho}.nc")
for trid in data2dump:
dataout[trid] = {"fileorig": fileorig}
# Read sensitivity when checking transforms
if check_transforms and not onlyinit and mode == "tl":
with _hdf5_lock:
dataout[trid]["incr"] = np.concatenate([
xr.open_dataset(f"{runsubdir}/../{fileorig}")[
f"{trid[1]}_o_tl"].values,
xr.open_dataset(f"{runsubdir}/../{fileorig}")[
f"{trid[1]}_tl"].values
], axis=0)
return dataout
[docs]
def make_end(self, file_end, ddi, ref_fwd_dir):
"""Create a zeroed TL end-concentration file for tangent-linear initialisation.
Reads the reference forward end file (``chain/end.*.nc`` in
*ref_fwd_dir*), appends zero-initialised TL increment variables for
all active species plus the fixed-species set (``M``, ``O2``, ``N2``,
``H2O``), and writes the result to *file_end*.
Both the forward (``<spec>``) and output-level (``<spec>_o_tl``) TL
variables are initialised to zero.
Args:
self: CHIMERE model plugin instance (carries ``nho``, ``domain``,
``chemistry.acspecies``).
file_end (str): output path for the initialised TL file.
ddi (datetime): period start date (used to build the filename).
ref_fwd_dir (str): directory containing the reference forward chain.
Raises:
Exception: if the reference forward end file does not exist.
"""
# Domain
domain = self.domain
# Active species
acspec = self.chemistry.acspecies.attributes
all_species = acspec + ["M", "O2", "N2", "H2O"]
all_species = all_species + [f"{s}_o" for s in all_species]
all_species = [f"{s}_tl" for s in all_species]
# Fetch original end from reference forward
ref_end = ddi.strftime(f"{ref_fwd_dir}/chain/end.%Y%m%d%H.{self.nho}.nc")
if not os.path.isfile(ref_end):
raise CifError(
f"Could not find the end.nc file from a reference forward: {ref_end}"
)
with _hdf5_lock:
ds_ref = xr.open_dataset(ref_end)
# Add TL variables
for spec in all_species:
ds_ref[spec] = (
("Time", "bottom_top", "south_north", "west_east"),
np.zeros((1, domain.nlev, domain.nlat, domain.nlon)))
# Add unit attribute
ds_ref[spec].attrs = {"units": "molecules/cm3"}
# Dump to end
ds_ref.to_netcdf(file_end, "w", format="NETCDF3_CLASSIC",
encoding={'Times': {'char_dim_name': 'DateStrLen'},
'species': {'char_dim_name': 'SpStrLen'}},
unlimited_dims={'Time': True})
[docs]
def make_aend(self, file_end, ddi):
"""Create a zeroed adjoint end-sensitivity file for adjoint initialisation.
Writes ``{file_end}`` as a NetCDF3 CLASSIC file containing zeroed
sensitivity fields for all active and fixed species (``M``, ``O2``,
``N2``, ``H2O``). Both the full-domain (``<spec>_ad``) and output-
level (``<spec>_ad_o``) arrays are included, accounting for the MPI
tile buffer padding (``6 * nzdoms`` / ``6 * nmdoms`` cells).
Args:
self: CHIMERE model plugin instance (carries ``domain``,
``chemistry.acspecies``, ``nzdoms``, ``nmdoms``, ``nivout``).
file_end (str): output path for the initialised adjoint file.
ddi (datetime): period start date (written into the ``Times``
variable).
"""
# Domain
domain = self.domain
# Active species
acspec = self.chemistry.acspecies.attributes
all_species = acspec + ["M", "O2", "N2", "H2O"]
all_species = [f"{s}_ad" for s in all_species]
all_species = all_species + [f"{s}_o" for s in all_species]
# Initialize zeros adjoint sensitivities
ds = xr.Dataset(
{s: (("Time", "bottom_top", "south_north", "west_east"),
np.zeros((1, self.nivout + 1,
domain.nlat + 6 * self.nmdoms, domain.nlon + 6 * self.nzdoms)))
for s in all_species}
)
# Add unit attribute
for s in all_species:
ds[s].attrs = {"units": "molecules/cm3"}
# Transform times to CHIMERE strings
str_dates = [ddi.strftime("%Y-%m-%d_%H:%M:00")]
date_dtype = np.dtype(('S', 19))
ds["Times"] = xr.DataArray(str_dates, dims=["Time"]).astype(date_dtype)
# list of species
spec_dtype = np.dtype(('S', 23))
ds["species"] = xr.DataArray(
[s.ljust(23) for s in self.chemistry.acspecies.attributes +
["M", "O2", "N2", "H2O"]]
+ (len(self.chemistry.acspecies.attributes) + 4) * [""],
dims=["Species"]).astype(spec_dtype)
# Dump to end
with _hdf5_lock:
ds.to_netcdf(file_end, "w", format="NETCDF3_CLASSIC",
encoding={'Times': {'char_dim_name': 'DateStrLen'},
'species': {'char_dim_name': 'SpStrLen'}},
unlimited_dims={'Time': True})
