Source code for pycif.plugins.models.lmdz_ico.io.native2inputs_adj

from __future__ import annotations

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

import numpy as np
import pandas as pd
import xarray as xr

from .outputs import fetch_end
from .....utils.check.errclass import CifFileNotFoundError, CifValueError
from .....utils.hdf5 import _hdf5_lock




def native2inputs_adj(
    self,
    datastore: dict[tuple[str, str], dict[str | datetime.datetime, Any]],
    input_type: str,
    datei: datetime.datetime,
    datef: datetime.datetime,
    runsubdir: str | PathLike,
    mode: Literal["fwd", "tl", "adj"] = "fwd",
    onlyinit: bool = False,
    do_simu: bool = True,
    check_transforms: bool = False,
    **kwargs,
) -> dict[tuple[str, str], dict[str | datetime.datetime, Any]]:
    """Converts data at the model data resolution to model compatible input
    files.

    Args:
        self: the model Plugin
        input_type (str): one of 'fluxes', 'obs'
        datastore: data to convert
            if input_type == 'fluxes', a dictionary with flux maps
            if input_type == 'obs', a pandas dataframe with the observations
        datei, datef: date interval of the sub-simulation
        mode (str): running mode: one of 'fwd', 'adj' and 'tl'
        runsubdir (str): sub-directory for the current simulation
        workdir (str): the directory of the whole pycif simulation

    Notes:
        - LMDZ expects daily inputs; if the periods in the control vector are
        longer than one day, period values are uniformly de-aggregated to the
        daily scale; this is done with pandas function 'asfreq' and the option
        'ffill' as 'forward-filling'
        See Pandas page for details:
        https://pandas.pydata.org/pandas-docs/stable/generated/pandas
        .DataFrame.asfreq.html

    """
    runsubdir = Path(runsubdir)
    date = pd.to_datetime(runsubdir.name, format="%Y-%m-%d_%H-%M")
    ddi = min(datei, datef)

    # Fetch end concentrations of adjoint for chain simulation
    if input_type == "endconcs":
        datastore = fetch_end(
            self, datastore, datei, runsubdir, mode, onlyinit, check_transforms
        )

    # Reading only output files related to given input_type
    if input_type in ("inicond", "flux", "prescrconcs", "prodloss3d"):
        for trid in datastore:
            (comp, spec) = trid

            if comp == "inicond":
                file_out = runsubdir.parent / "chain" / f"restart_{date:%Y-%m-%d-%H-00}.nc"
            else:
                file_out = runsubdir / f"{comp}_out.nc"

            if not file_out.is_file():
                raise CifFileNotFoundError(
                    f"LMDZ did not produce adjoint output file '{file_out}' for "
                    f"'{comp}' in '{runsubdir}'"
                )

            with _hdf5_lock:
                with xr.open_dataset(file_out) as ds:
                    da = ds[spec]

            if self.grid == "regular":
                da = da.drop_vars(["lat", "lon"])
            elif self.grid == "dynamico":
                da = da.rename({"cell": "lon"}).expand_dims("lat", axis=1)
            else:
                raise CifValueError(f"Unknown grid type '{self.grid}'")

            if input_type == "flux":
                da = da.expand_dims("lev", axis=1)
                expected_time = self.flux_input_dates[ddi][:-1].astype("datetime64")
                if not np.all(da.time.values == expected_time):
                    raise CifValueError(
                        f"Invalid time coordinate in LMDZ adjoint output file '{file_out}'"
                    )

            elif input_type == "inicond":
                da = da.expand_dims({"time": [date]})

            else:
                # TODO: Is this ever reached ??
                da["time"] = (["time"], self.input_dates[ddi][:-1])

            # Check dimensions, wrong dimension order will cause errors on #
            # following transforms
            if da.dims != ("time", "lev", "lat", "lon"):
                raise CifValueError(
                    f"Invalid dimensions in LMDZ adjoint output file '{file_out}'"
                )

            datastore[trid]["data"][ddi]["adj_out"] = da

    return datastore