Source code for pycif.plugins.models.lmdz_acc.io.outputs.fetch_end

from ......utils.path import link
from ......utils.hdf5 import _hdf5_lock
import shutil
import xarray as xr
import numpy as np




def fetch_end(
        self, data2dump, runsubdir, mode, ddi, ddf,
        check_transforms=False, onlyinit=False
):
    """Register end-concentration restart paths after a LMDZ-ACC run.

    For **forward / TL** mode: records the restart file path
    (``chain/restart_YYYYMMDDHHMM.nc``) in *data2dump*.

    For **adjoint** mode: records the adjoint restart path and optionally
    reads the sensitivity fields back into the data store when
    *check_transforms* is active.

    Args:
        self: LMDZ-ACC model plugin instance.
        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).
        check_transforms (bool): if ``True``, also read sensitivity back.
        onlyinit (bool): skip reading if ``True``.

    Returns:
        dict: updated *data2dump*.
    """

    if mode != "adj":
        dataout = {}
        fileorig = \
            ddi.strftime("chain/restart_%Y%m%d%H%M.nc")
        fileorig_tl = \
            ddi.strftime("chain/restart_tl_%Y%m%d%H%M.bin")
        for trid in data2dump:
            dataout[trid] = {"fileorig": fileorig,
                             "fileorig_tl": fileorig_tl}

            # Read sensitivity only when checking transforms
            if not (check_transforms and not onlyinit and mode == "tl"):
                continue

            spec = trid[1]
            restartID = getattr(self.chemistry.acspecies, spec).restart_id
            var = f"q{restartID:02d}" if type(restartID) != str \
                else restartID
            nlat, nlon = self.domain.zlon.shape
            nlev = self.domain.nlev

            # For the old version, shift the count of bits
            offset = 0
            if self.plugin.version == "std":
                offset = 4

            incr_out = np.fromfile(
                f"{runsubdir}/../{fileorig_tl}",
                offset=offset
            ).reshape((-1, nlev, nlat, nlon))

            ind_species = self.chemistry.acspecies.attributes.index(trid[1])
            dataout[trid]["incr"] = incr_out[ind_species]

        return dataout

    else:
        fileorig = \
            ddi.strftime("chain/restart_%Y%m%d%H%M.nc")
        for trid in data2dump:
            data2dump[trid]["data"][ddi]["fileorig"] = fileorig

            # Read sensitivity when checking transforms
            if check_transforms and not onlyinit:
                spec = trid[1]
                restartID = getattr(self.chemistry.acspecies, spec).restart_id
                var = f"q{restartID:02d}" if type(restartID) != str \
                    else restartID
                with _hdf5_lock:
                    data2dump[trid]["data"][ddi]["adj_out"] = \
                        xr.open_dataset(f"{runsubdir}/../{fileorig}")[var].values

        return data2dump





def make_restart_adj(self, runsubdir, ddi, ref_fwd_dir):
    """Create a zeroed LMDZ-ACC adjoint restart file for the current period.

    Copies the reference forward restart to the adjoint run directory and
    zeros all active-species mass mixing ratios so the LMDZ-ACC adjoint
    can accumulate sensitivities.

    Args:
        self: LMDZ-ACC model plugin instance.
        runsubdir (str): path to the period run directory.
        ddi (datetime): period start date.
        ref_fwd_dir (str): directory of the reference forward run.
    """
    # Fetch original end from reference forward
    end_file = f"{runsubdir}/restart.nc"
    ref_end = ddi.strftime(
        f"{ref_fwd_dir}/chain/restart_%Y%m%d%H%M.nc")
    shutil.copy(ref_end, end_file)

    # Put adjoint sensitivities to zero or propagate previous values
    with _hdf5_lock:
        ds = xr.open_dataset(end_file, mode="a")
        for spec in self.chemistry.acspecies.attributes:
            restartID = getattr(self.chemistry.acspecies, spec).restart_id
            var = f"q{restartID:02d}" if type(restartID) != str \
                else restartID

            # Read adjoint sensitivity at beginning of run
            ref_sensit = 0
            if ddi != self.subsimu_dates[-2]:
                sensit_start_file = f"{runsubdir}/start.nc"
                ds_end = xr.open_dataset(sensit_start_file)
                mass = ds_end["masse"][:]
                sensit = ds_end[var][:]
                ref_sensit = (mass * sensit).sum() / mass.sum()

            ds[var][:] = ref_sensit

    with _hdf5_lock:
        ds.to_netcdf(end_file, mode="a")