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

from __future__ import annotations

from os import PathLike
from typing import Any

import xarray as xr
from .....utils.check.errclass import CifTypeError
from .....utils.hdf5 import _hdf5_lock




def write(
    self,
    name: str,
    path: str | PathLike,
    data: xr.DataArray,
    metadata: dict[str, Any] | None = None,
    **kwargs,
) -> None:
    """Write prescribed species fluxes to a regular lat-lon LMDZ NetCDF file.

    Builds coordinates from `self.domain.get_domain_coords()` (plus a
    ``time`` coordinate, if present in `data`), squeezes the ``lev``
    dimension of `data`, and appends the resulting dataset to `path`.

    Args:
        self: this plugin
        name (str): name of the flux variable to write
        path (str): path to the file to write (or append to)
        data (xarray.DataArray): Data to write
        metadata (dict, optional): Unused, kept for interface consistency
            with other flux plugins.

    Raises:
        CifTypeError: If `data` is not an `xarray.DataArray`.
    """
    if not isinstance(data, xr.DataArray):
        plg = self.plugin
        plg_name = f"{plg.name} / {plg.version} / {plg.type}"
        raise CifTypeError(
            f"Plugin {plg_name} write method 'data' argument should of type "
            f"xarray.DataArray. Passed 'data' argument type: {type(data)}"
        )

    coords = self.domain.get_domain_coords()

    if "time" in data.coords:
        # fmt: off
        coords["time"] = (["time"], data.time.values, {
            "standard_name": "time",
            "long_name":     "time axis",
            "axis":          "T",
        })
        # fmt: on

    # Formating dataset
    data = data.squeeze("lev")
    ds = xr.Dataset({name: (data.dims, data.values)}, coords=coords)
    with _hdf5_lock:
        ds.to_netcdf(path, mode="a", encoding={"time": {'dtype': "int32"}})