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

from __future__ import annotations

import datetime
from os import PathLike
from pathlib import Path
from typing import Any, Literal

import pandas as pd
import xarray as xr

from ......utils.hdf5 import _hdf5_lock




def fetch_end(
    self,
    data2dump: dict[tuple[str, str], dict[str | datetime.datetime, Any]],
    ddi: datetime.datetime,
    runsubdir: str | PathLike,
    mode: Literal["fwd", "tl", "adj"],
    onlyinit: bool = False,
    check_transforms: bool = False,
) -> dict[tuple[str, str], dict[str | datetime.datetime, Any]]:
    runsubdir = Path(runsubdir)
    date = pd.to_datetime(ddi)
    fileorig = f"chain/restart_{date:%Y-%m-%d-%H-00}.nc"

    if mode in ("fwd", "tl"):
        dataout = {}
        for trid in data2dump:
            _, spec = trid
            dataout[trid] = {"fileorig": fileorig}

            # Read sensitivity only when checking transforms
            if check_transforms and not onlyinit and mode == "tl":
                with _hdf5_lock:
                    with xr.open_dataset(runsubdir.parent / fileorig) as ds:
                        ds = ds.expand_dims("time")
                        dataout[trid]["incr"] = ds[f"{spec}_tl"].values

        return dataout

    else:
        for trid in data2dump:
            _, spec = trid
            data2dump[trid]["data"][ddi]["fileorig"] = fileorig  # TODO: why "data" ???

            # Read sensitivity when checking transforms
            if check_transforms and not onlyinit:
                with _hdf5_lock:
                    with xr.open_dataset(runsubdir.parent / fileorig) as ds:
                        ds = ds.expand_dims("time")
                        data2dump[trid]["data"][ddi]["adj_out"] = ds[spec].values

        return data2dump