Source code for pycif.plugins.datastreams.fluxes.dalecbethy.write

import os
from logging import info

import numpy as np
import xarray as xr

from .....utils.check.errclass import CifError, CifKeyError
from .....utils.hdf5 import _hdf5_lock


# Units the Fortran reader validates against (see model_sources/dalecbethy/
# src/ncread_forcing.f90::ncread_dynamic_forcing).
_UNITS = {
    "swrad": "W/m2",
    "temperature": "Celsius",
    "precipitation": "mm/h",
    "lwdown": "W/m2",
    "soil_temperature": "Celsius",
}




def write(self, name, flx_file, flx, mode="a", metadata=None, **kwargs):
    """Write D&B meteorological forcing to a native ``dynforcing.nc`` file.

    Builds/appends into the D&B-native forcing format expected by
    ``model_sources/dalecbethy/src/ncread_forcing.f90::
    ncread_dynamic_forcing``: ``ng``/``time`` dimensions and one ``(ng,
    time)`` variable per meteorological component, with the exact
    ``units`` attribute the Fortran reader checks for. The mandatory
    ``yyyymmddhh`` calendar variable is *not* written here (it is not a
    per-component ``("meteo", <component>)`` mapper input): see
    :func:`pycif.plugins.models.dalecbethy.io.inputs.make_forcing.make_meteo`,
    which copies it as-is from the original forcing file.

    Values are scattered back from the tracer's active sample points
    (``self.domain.active``, a subset of the full ``self.domain.ng`` grid
    cells) to their original ``ng`` row; grid cells outside
    ``self.domain.active`` are left as ``NaN``.

    Args:
        self: the fluxes Plugin (with a ``dalecbethy`` ``domain`` attached).
        name (str or list[str]): name(s) of the component(s) to write.
        flx_file (str): the ``dynforcing.nc`` file to write/append to.
        flx (xarray.DataArray or dict[str, xarray.DataArray]): forcing data,
            with dimensions ``(time, lev, lat, lon)`` (``lev``/``lat`` of
            size 1), as returned by :func:`~.read.read`.
        mode (str): ``"w"`` to overwrite, ``"a"`` to append.
        metadata (dict, optional): if given, ``metadata["domain"]`` is used
            instead of ``self.domain``.
        **kwargs: unused, kept for interface compatibility.

    Raises:
        CifKeyError: if no domain information can be found.
        CifError: if the domain has no ``active``/``ng`` sample-point
            selection, or a component has no known expected ``units``.
    """
    if isinstance(name, str):
        flx = {name: flx}
        name = [name]

    if hasattr(self, "domain"):
        domain = self.domain
    elif isinstance(metadata, dict) and "domain" in metadata:
        domain = metadata["domain"]
    else:
        raise CifKeyError("Could not find information about the domain")

    if not hasattr(domain, "active") or not hasattr(domain, "ng"):
        raise CifError(
            "Writing D&B dynamic forcing requires a 'dalecbethy' domain "
            "(with an 'active'/'ng' sample-point selection)."
        )
    active = domain.active
    ng = domain.ng

    time = flx[name[0]]["time"].values

    if os.path.isfile(flx_file) and mode == "a":
        with _hdf5_lock:
            with xr.open_dataset(flx_file) as ds_existing:
                if ds_existing.sizes.get("time") != time.size:
                    raise CifError(
                        f"Trying to append D&B forcing data with "
                        f"{time.size} time steps into '{flx_file}' which "
                        f"already has {ds_existing.sizes.get('time')}."
                    )

    data_vars = {}
    for k in name:
        if k not in _UNITS:
            raise CifError(
                f"Unknown D&B dynamic forcing component '{k}': expected "
                f"units are only known for {sorted(_UNITS)}."
            )

        squeezed = flx[k].squeeze(["lat", "lev"]).transpose("time", "lon")
        full = np.full((time.size, ng), np.nan)
        full[:, active] = squeezed.values

        data_vars[k] = xr.DataArray(
            full.T, dims=("ng", "time"),
            attrs={"units": _UNITS[k], "long_name": f"{k} (pyCIF-generated)"}
        )

    ds = xr.Dataset(data_vars)

    if not os.path.isfile(flx_file) or mode == "w":
        info(f"writing D&B dynamic forcing {name} to '{flx_file}'")
        with _hdf5_lock:
            ds.to_netcdf(flx_file, mode="w")
    else:
        info(f"appending D&B dynamic forcing {name} to '{flx_file}'")
        with _hdf5_lock:
            ds.to_netcdf(flx_file, mode="a")