Source code for pycif.plugins.datastreams.fields.lmdz_chemfield_ico.write

from __future__ import annotations

from os import PathLike
from typing import Any

import xarray as xr


# pylint: disable=unused-argument


def write(
    self,
    name: str,
    path: str | PathLike,
    data: xr.DataArray,
    **kwargs,
) -> None:
    """Write prescribed species files for LMDZ

    Args:
        self: this plugin
        path (str): path to the file to write
        data (xarray.Dataset or xarray.DataArray): Data to write
    """
    if not isinstance(data, xr.DataArray):
        plg = self.plugin
        plg_name = f"{plg.name} / {plg.version} / {plg.type}"
        raise TypeError(
            f"Plugin {plg_name} write method 'data' argument should of type "
            f"xarray.DataArray. Passed 'data' argument type: {type(data)}"
        )

    data = data.squeeze("lat")
    data = data.rename({"lon": "cell"})

    # For deposition velocity fields
    if data.sizes["lev"] == 1:
        data = data.squeeze("lev")

    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
    ds = xr.Dataset({name: (data.dims, data.values)}, coords=coords)
    ds.to_netcdf(path, mode="a", encoding={"time": {'dtype': "int32"}})