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

from __future__ import annotations

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

import pandas as pd
import xarray as xr
from pandas import DataFrame

from ......utils.check.errclass import CifFileNotFoundError, CifRuntimeError
from ......utils.hdf5 import _hdf5_lock




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

    cols = []
    for input_type, _ in data2dump:
        if input_type == "concs":
            cols.append("spec")
            if mode == "tl":
                cols.append("incr")
        else:
            cols.append(input_type)

    if sim_file.exists():
        # Read observations, if simulation file is present, observation file should
        # also exits
        with _hdf5_lock:
            with (
                xr.open_dataset(obs_file) as ds_obs,
                xr.open_dataset(sim_file) as ds_sim,
            ):
                # Restore original order (previously sorted by time)
                ds_sim = ds_sim[cols].sortby(ds_obs.index)
                ds_obs = ds_obs.sortby("index")

                tracer_index = ds_obs["itrac"].values  # starts from 1
                sim = ds_sim.to_dataframe()

    elif not obs_file.exists():
        # Empty dataframe
        tracer_index = pd.Series([0], dtype=int)
        sim = pd.DataFrame(columns=cols)

    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 = [spec.lower() for spec in self.chemistry.active_species]

    dataout = {}

    for (input_type, tracer), datastore in data2dump.items():
        tracer = tracer.replace("__sample#", "_")
        itrac = active_species.index(tracer.lower()) + 1

        df = sim.loc[tracer_index == itrac]
        dataloc = datastore[ddi]

        if len(dataloc) == 0:
            continue

        if self.skip_obs_chunking:
            if len(df) != len(dataloc):
                raise CifRuntimeError(
                    "Length mismatch between source and target dataframes. "
                    + "Disable 'skip_obs_chunking' to properly compute alignements"
                )
        else:
            # Redistributing extracting data into correct rows of the dataframe
            chunk_indexes = self.chunk_indexes[ddi][input_type][tracer]
            if chunk_indexes is None:
                raise CifRuntimeError(
                    f"chunk indexes for {input_type}/{tracer} are uninitialized"
                )

            df = df.loc[chunk_indexes]

        if input_type == "concs":
            dataloc.loc[:, ("maindata", "spec")] = df.loc[:, "spec"].to_numpy()

            if mode == "tl":
                dataloc.loc[:, ("maindata", "incr")] = df.loc[:, "incr"].to_numpy()
            else:
                dataloc.loc[:, ("maindata", "incr")] = 0.0

        else:
            dataloc.loc[:, ("maindata", "spec")] = df.loc[:, input_type].to_numpy()
            dataloc.loc[:, ("maindata", "incr")] = 0.0

        dataout[(input_type, tracer)] = dataloc.copy()

    return dataout