Source code for pycif.plugins.domains.dalecbethy.read_domain

from pathlib import Path

import numpy as np
import xarray as xr

from ....utils.check.errclass import CifKeyError, CifValueError




def read_grid(domain, **kwargs):
    """Read D&B's sample-point domain from its static forcing NetCDF file.

    Grid-cells are collapsed into a 1D list of ``nsp`` "sample points",
    skipping grid-cells where ``condition_simulate == 0``.

    Sets on *domain*: ``unstructured_domain``, ``nlon`` (=``nsp``),
    ``nlat`` (=1), ``zlon``, ``zlat`` (shape ``(1, nsp)``), plus a
    single-level vertical coordinate. Also sets ``ng`` (the full, unfiltered
    size of the static forcing file's flat grid-cell dimension) and
    ``active``, the indices (into that ``ng`` dimension) of the retained
    sample points, in the same order as ``zlon``/``zlat``; other D&B NetCDF
    files indexed along the same ``ng`` dimension (e.g. D&B's own
    ``dynforcing.nc``, see
    ``pycif.plugins.datastreams.fluxes.dalecbethy``) must be subset/scattered
    with these to line up with this domain.

    Args:
        domain: dalecbethy domain plugin instance with ``dir``/``file`` set.
        **kwargs: unused.

    Raises:
        CifKeyError: if a mandatory dimension/variable is missing.
        CifValueError: if no active sample point is found.
    """
    domain.unstructured_domain = True
    file_path = Path(domain.dir, domain.file)

    with xr.open_dataset(file_path) as ds:
        if "condition_simulate" not in ds:
            raise CifKeyError(
                f"mandatory variable 'condition_simulate' not found in "
                f"'{file_path}'."
            )
        condition_simulate = ds["condition_simulate"].values

        if "lon" not in ds or "lat" not in ds:
            raise CifKeyError(
                f"mandatory 'lon'/'lat' coordinates not found in "
                f"'{file_path}'."
            )
        lon = ds["lon"].values
        lat = ds["lat"].values

        active = np.where(condition_simulate != 0)[0]
        nsp = active.size

        if nsp == 0:
            raise CifValueError(
                f"no active sample point found in D&B static forcing file "
                f"'{file_path}'."
            )

        # D&B is a point cloud, not a regular grid: expose it the same way
        # as pycif.plugins.domains.unstructured_NetCDF (1 x nsp).
        domain.nlon = nsp
        domain.nlat = 1
        domain.ng = condition_simulate.size
        domain.active = active
        domain.zlon = lon[active][np.newaxis, :]
        domain.zlat = lat[active][np.newaxis, :]

        # Complete the corner coordinates
        dlon = np.min(np.diff(np.unique(domain.zlon))) / 2
        dlat = np.min(np.diff(np.unique(domain.zlat))) / 2

        domain.zlonc = np.concatenate(
            [lon[active][np.newaxis, :] - dlon,
             lon[active][np.newaxis, :] - dlon,
             lon[active][np.newaxis, :] + dlon,
             lon[active][np.newaxis, :] + dlon],
            axis=0
        )
        domain.zlatc = np.concatenate(
            [lat[active][np.newaxis, :] - dlat,
             lat[active][np.newaxis, :] + dlat,
             lat[active][np.newaxis, :] + dlat,
             lat[active][np.newaxis, :] - dlat],
            axis=0
        )

        domain.nlev = 1
        domain.sigma_a_mid = [1]
        domain.sigma_b_mid = [0]