Source code for pycif.plugins.domains.gridded_NetCDF.utils

from logging import debug
from typing import Union
import numpy as np
import xarray as xr




def get_bounds(
    ds: xr.Dataset,
    coord_name: str,
    flip_coord: bool,
    min_bound: float,
    max_bound: float,
    regular_extent: bool
) -> np.ndarray:
    """Get the bounds of a given coordinates.
    """
    bounds_name = ds[coord_name].attrs.get('bounds')
    # TODO: detect if coord dimensions == (N, 2) (-> in this bounds_name = coord_name)

    bounds = None

    if bounds_name is None:
        shape = ds[coord_name].shape
        if len(shape) == 2 and shape[1] == 2:
            bounds_name = coord_name

    if bounds_name is not None:
        bounds = ds[bounds_name].values
        bounds = np.concatenate([bounds[:, 0], bounds[-1:, 1]])
        bounds = np.flip(bounds) if flip_coord else bounds
        return bounds

    else:
        debug(f"no bounds found for coordinate '{coord_name}', "
              f"assuming '{coord_name}' represents the cell centers")

        coord = ds[coord_name].values
        coord = np.flip(coord) if flip_coord else coord

        delta = np.unique(np.round(np.diff(coord), 10))
        if len(delta) == 1:
            bounds = coord - delta[0] / 2

            return np.concatenate(
                [[max(bounds[0], min_bound)],
                 bounds[1:],
                 [min(bounds[-1] + delta[0], max_bound)]
                 ]
            )

        elif len(delta) != 1 and regular_extent:
            raise ValueError(
                f"'{coord_name}' coordinate is not regular.\n"
                "If the coordinate is not expected to be regular, please specify the corresponding "
                "argument in the yaml: `regular_XXXX: False`\n"
                f"Otherwise, the list of possible deltas is: {delta}.\n"
                "If all delta are very similar apart from rounding difference, "
                "it is possible to force a constant delta by using the parameter "
                "'delta_XXXX' in the yml. Please check the documentation for "
                "the domain plugin 'gridded_netcdf/std'"
            )

        else:
            bounds = np.concatenate([
                coord[[0]] - np.diff(coord)[0] / 2,
                coord[1:] - np.diff(coord) / 2,
                coord[[-1]] + np.diff(coord)[-1] / 2
            ])
            return bounds





def get_centers(coord: np.ndarray, flip_coord: bool) -> np.ndarray:
    coord = np.flip(coord) if flip_coord else coord
    center = coord[:-1] + np.diff(coord) / 2

    return center





def find_coord(
    ds: Union[xr.Dataset, xr.DataArray],
    name: str,
    coord_name: str = None
) -> str:
    """Find a variable or coordinate name in a dataset based on its 'long_name'
    or 'standard_name' attribute.
    """
    all_keys = [*ds, *ds.coords, *ds.dims] if isinstance(ds, xr.Dataset) \
        else [*ds.coords, *ds.dims]

    # Search in coordinates and variables
    for var_name in all_keys:
        if len(ds[var_name].shape) > 1:
            continue

        if name == ds[var_name].attrs.get('standard_name'):
            return var_name

        elif name == ds[var_name].attrs.get('long_name'):
            return var_name

        elif name == var_name:
            return name

    if coord_name in all_keys:
        return coord_name

    raise ValueError("there is no variable or coordinate corresponding to "
                     f"'{name}' or '{coord_name}' in the dataset")