Source code for pycif.plugins.models.lmdz_ico.io.inputs.fluxes

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
from .ensemble import ensemble_trid




def make_fluxes(
    self,
    datastore: dict[tuple[str, str], dict[str | datetime.datetime, Any]],
    datei: datetime.datetime,
    runsubdir: str | PathLike,
    mode: Literal["fwd", "tl", "adj"],
) -> None:
    """Write flux input files for one sub-simulation period"""
    for spec in self.chemistry.emitted_species:
        trid = ensemble_trid(self, ("flux", spec), datastore)

        if trid not in datastore:
            continue

        data = datastore[trid]["data"][datei]

        # If not determined by the control vector, read input file with datavect
        # plugin 'read' method
        if "spec" not in data:
            input_tracer = datastore[trid]["tracer"]
            data["spec"] = input_tracer.read(
                spec,
                input_tracer.varname,
                dates=datastore[trid]["input_dates"][datei],
                files=datastore[trid]["input_files"][datei],
                tracer=input_tracer,
            )

        flux_file = Path(runsubdir, "flux.nc")
        self.emissions.write(spec, flux_file, data["spec"])

        if mode == "tl" and "incr" in data:
            self.emissions.write(f"{spec}_tl", flux_file, data["incr"])

        # Force time units to seconds since beginning of month
        with _hdf5_lock:
            with xr.open_dataset(flux_file) as ds:
                ref_datetime = pd.Timestamp(year=datei.year, month=datei.month, day=1)
                seconds = (pd.to_datetime(
                    ds["time"].values) - ref_datetime).total_seconds()
                time = xr.DataArray(
                    data=seconds.values.astype("int32"),
                    dims=["time"],
                    name="time",
                    attrs={
                        "standard_name": "time",
                        "long_name": "time",
                        "units": f"seconds since {ref_datetime:%Y-%m-%d %H:%M:%S}",
                        "calendar": "proleptic_gregorian",
                    },
                )

        with _hdf5_lock:
            time.to_dataset().to_netcdf(flux_file, mode="a")