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

from __future__ import annotations

import copy
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 ......utils.check.errclass import CifFileNotFoundError
from ......utils.hdf5 import _hdf5_lock




def read_sim(
    self,
    data2dump: dict[tuple[str, str], dict[str | datetime.datetime, Any]],
    ddi: datetime.datetime,
    runsubdir: str | PathLike,
    mode: Literal["fwd", "tl", "adj"],
) -> dict[tuple[str, str], Any]:
    """Extract simulated concentrations"""
    runsubdir = Path(runsubdir)
    obs_file = runsubdir / "obs.nc"
    sim_file = runsubdir / "obs_out.nc"

    if sim_file.exists():
        # Read simulation file (should not be empty)
        with _hdf5_lock:
            with xr.open_dataset(sim_file) as ds:
                sim = ds.to_dataframe()

        # Read observations, if simulation file is present, observation file should also
        with _hdf5_lock:
            with xr.open_dataset(obs_file) as ds:
                tracer_index = ds["itrac"].values - 1

    elif not obs_file.exists():
        # Empty dataframe
        sim = pd.DataFrame(
            columns=[
                "spec",
                "incr",
                "pressure",
                "dpressure",
                "hlay",
                "airm",
            ]  # type: ignore
        )
        tracer_index = np.zeros(0, dtype=int)

    else:
        # Simulation file, should be present if observation file is
        raise CifFileNotFoundError(
            f"LMDZ did not produce a obs_out.nc file in '{runsubdir}'."
        )

    active_species = list(self.chemistry.active_species)
    tracer_name = pd.Series(active_species).iloc[tracer_index].values

    dataout = {}
    for trid in data2dump:
        component, spec = trid
        dataloc = data2dump[trid][ddi]

        if len(dataloc) == 0:
            continue

        # Putting values to the local data store
        spec_str = spec.replace("__sample#", "_")
        mask = tracer_name == spec_str
        sim_spec = sim.loc[mask]

        # Redistributing extracting data into correct rows of the dataframe
        chunk_indexes = self.chunk_indexes[ddi][component][spec_str]
        dataavg = sim_spec.loc[chunk_indexes]

        # Put values in dataloc
        dataloc.loc[:, ("maindata", "spec")] = dataavg.loc[:, "spec"].values
        if mode == "tl":
            dataloc.loc[:, ("maindata", "incr")] = dataavg.loc[:, "incr"].values

        # Put simulated value into correct column
        # Different case if concs, or other parameters such as pressure
        # Put pressure and other auxiliary data into spec column for later
        # interpolation
        if trid[0] != "concs":
            # Column name
            col = trid[0]
            dataloc.loc[:, ("maindata", "spec")] = copy.deepcopy(
                dataavg.loc[:, col].values
            )
            dataloc.loc[:, ("maindata", "incr")] = 0.0

        dataout[trid] = copy.deepcopy(dataloc)

    return dataout