Source code for pycif.plugins.datastreams.fields.lmdz_inicond_reg.read

from __future__ import annotations

import datetime

import numpy as np
import xarray as xr
from .....utils.check.errclass import CifValueError
from .....utils.hdf5 import _hdf5_lock




def read(
    self,
    name: str,
    varnames: str,
    dates: list[tuple[datetime.datetime, datetime.datetime]],
    files: list[str],
    tracer: object | None = None,
    **kwargs,
) -> xr.DataArray:
    """Read a mixing-ratio initial condition (regular grid) into a pyCIF DataArray.

    For backward compatibility with older LMDZ files, strips a leading
    singleton time dimension if present, and drops a trailing duplicated
    (cyclic-closure) longitude column if the longitude size is
    ``self.domain.nlon + 1``. The resulting shape is validated against
    ``self.domain``.

    Args:
        self: The field/tracer Plugin; ``self.domain`` gives the expected
            ``(nlev, nlat, nlon)`` shape.
        name: Name of the component; used as fallback variable name.
        varnames: Name of the variable to read; if empty, ``name`` is used
            instead.
        dates: List of date entries; only the first date is used as the
            output time coordinate.
        files: List with a single inicond file path.
        tracer: Unused.
        **kwargs: Unused.

    Returns:
        xarray.DataArray: A 4-dimensional ``(time, lev, lat, lon)`` array
        with a single time step.

    Raises:
        CifValueError: If more than one file is given, or if the data
            shape (after backward-compatibility adjustments) does not
            match ``self.domain``.
    """
    varnames = varnames if varnames else name

    if len(files) > 1:
        raise CifValueError("multiple files provided for initial conditions")
    (path,) = files

    with _hdf5_lock:
        with xr.open_dataset(path) as ds:
            data = ds[varnames].values

    # Backward compatibility with old LMDZ files
    if data.shape[0] == 1:
        data = data[0, ...]
    if data.shape[2] == self.domain.nlon + 1:
        data = data[:, :, :-1]

    if data.shape != (self.domain.nlev, self.domain.nlat, self.domain.nlon):
        raise CifValueError(
            f"unexpected initial conditions data shape {data.shape}, "
            f"should be ({self.domain.nlev}, {self.domain.nlat}, {self.domain.nlon})"
        )

    xmod = xr.DataArray(
        data[np.newaxis, :, :, :],
        coords={"time": [dates[0][0]]},
        dims=("time", "lev", "lat", "lon"),
    )

    return xmod