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

from __future__ import annotations

import datetime
from os import PathLike

import pandas as pd
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 | PathLike],
    ddi: datetime.datetime | None = None,
    **kwargs,
):
    """Read LMDz chemical field data (regular grid) into a pyCIF DataArray.

    Opens each distinct file in ``files`` once (validating that the
    requested variable exists), selects the day index (relative to the
    file's month, computed from ``ddi``) for each requested date,
    concatenates the selected time steps along ``time_counter``, expands a
    missing vertical dimension (for deposition-velocity/surface-only
    fields), renames dimensions to pyCIF's ``(time, lev, lat, lon)``
    convention, drops non-time coordinates, and transposes accordingly.

    Args:
        self: The field/tracer Plugin.
        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 ``(start, end)`` date tuples to extract, matching
            ``files``.
        files: List of files matching ``dates``; each distinct file is
            opened only once.
        ddi: Reference date used to compute each date's day-of-month index
            within its file. Mandatory.
        **kwargs: Unused.

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

    Raises:
        CifValueError: If ``ddi`` is not given, or if the requested
            variable is not found in a file.
    """
    varnames = varnames if varnames else name
    if ddi is None:
        raise CifValueError("'ddi' argument is required")

    da_list = []
    file_ref = ""
    da = None

    for (di, _), file_path in zip(dates, files):
        if file_path != file_ref:
            file_ref = file_path

            with _hdf5_lock:
                with xr.open_dataset(file_path) as ds:
                    if varnames not in ds:
                        raise CifValueError(
                            f"Variable '{varnames}' not found in '{file_path}'"
                        )
                    da = ds[varnames]

        if da is None:
            raise CifValueError(f"Error while reading '{file_path}'")

        # Assume monthly files with daily resolution
        day_index = (di - ddi).days + ddi.day - 1
        da_list.append(da.isel(time_counter=[day_index]))

    xmod = xr.concat(da_list, dim="time_counter")

    # Adding vertical dimension for deposition velocity fields
    if "presnivs" not in xmod.dims:
        xmod = xmod.expand_dims("presnivs", axis=1)

    xmod = xmod.rename(time_counter="time", presnivs="lev")

    # Dropping coordinates
    for coord_name in xmod.coords:
        if coord_name != "time":
            xmod = xmod.drop(coord_name)

    # Reordering dimensions
    xmod = xmod.transpose("time", "lev", "lat", "lon")

    return xmod