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

from logging import debug

import xarray as xr

from .....utils.check.errclass import CifTypeError, CifValueError, IllegalArgumentError
from .....utils.hdf5 import _hdf5_lock




def read(self, name, varnames, dates, files, tracer=None, **kwargs):
    """Read a species from an LMDz restart file into a pyCIF DataArray.

    Resolves the variable name to read in the file, in order of
    precedence: ``varnames`` if given (for a non-perturbed species; raises
    if ``tracer.restart_id`` is also set), otherwise ``tracer.restart_id``
    if set, otherwise the ``restart_id`` of the matching species in the
    chemical scheme (``self.chemistry.acspecies``). A numeric
    ``restart_id`` is turned into the conventional ``qNN`` variable name.

    Args:
        self: The field/tracer Plugin; ``self.chemistry.acspecies`` is used
            when no explicit ``restart_id``/``varnames`` is given.
        name: Name of the component/species.
        varnames: Explicit variable name to read; incompatible with
            ``tracer.restart_id`` unless ``name`` is a perturbed sample
            (contains ``"__sample#"``).
        dates: List of date entries; only the first date is used as the
            output time coordinate.
        files: List with a single restart file path.
        tracer: Tracer/component configuration; may define ``restart_id``.
        **kwargs: Unused.

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

    Raises:
        IllegalArgumentError: If both ``varnames`` and ``tracer.restart_id``
            are given for a non-perturbed species.
        CifTypeError: If ``varnames`` is given but is not a string.
        CifValueError: If more than one file is given, or if the species is
            not found in the chemical scheme.
    """
    is_perturbated = "__sample#" in name

    if varnames and not is_perturbated:
        if hasattr(tracer, 'restart_id'):
            raise IllegalArgumentError(
                "'restart_id' and 'varname' arguments cannot be used together")

        if not isinstance(varnames, str):
            raise CifTypeError(
                "'varnames' argument should be of type str, not list")

        var_name = varnames

    else:
        if hasattr(tracer, 'restart_id'):
            debug("getting 'restart_id' from initial condition plugin's arguments")
            restart_id = tracer.restart_id

        else:
            debug("getting 'restart_id' from chemical scheme")
            if not hasattr(self.chemistry.acspecies, name):
                raise CifValueError(
                    f"active species '{name}' not found in the chemical scheme")

            spec = getattr(self.chemistry.acspecies, name)
            restart_id = spec.restart_id

        if isinstance(restart_id, str):
            var_name = restart_id
        else:
            var_name = f'q{restart_id:02d}'

    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[var_name].values

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

    return xmod