grid-process 1.0.0__py3-none-any.whl

grid_process/grid_interp_utils.py

@@ -0,0 +1,208 @@
+# Copyright (c) 2026 MeteoSwiss, contributors listed in AUTHORS
+#
+# Distributed under the terms of the BSD 3-Clause License.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Module with utilities for spatial interpolation of gridded datasets
+"""
+
+from typing import Any
+
+import numpy as np
+from pyproj import CRS
+import xarray as xr
+
+from grid_process.helpers import reproject
+from grid_process import data_model as dm
+
+
+Y_DIM, X_DIM = dm.YX[0]
+LAT, LON = dm.LATLON[0]
+
+
+def create_dst_grid(  # pylint: disable=too-many-locals
+    dst_grid_crs: Any,
+    bottomleft: tuple[float, float],
+    topright: tuple[float, float] | None = None,
+    nrowcol: tuple[int, int] | None = None,
+    resolution: float | None = None,
+    add_latlon: bool = False,
+    dtype: np.dtype = np.dtype("float32"),
+) -> xr.Dataset:
+    """
+    Create a destination grid on which data can be interpolated. The output
+    can be directly used with :func:`interp2grid`. Two out of the three optional
+    input arguments ``topright``, ``nrowcol``, ``resolution`` must be provided.
+
+    Parameters
+    ----------
+    dst_grid_crs : string, int, dict or a CRS object
+        Anything accepted by pyproj.crs.CRS.from_user_input:
+        https://pyproj4.github.io/pyproj/dev/api/crs/crs.html#pyproj.crs.CRS.from_user_input
+        (e.g. EPSG:4326)
+    bottomleft : tuple of floats
+        (y, x) coordinates of the center of the bottom left corner pixel
+    topright : tuple of floats, optional
+        (y, x) coordinates of the center of the top right corner pixel.
+        If provided with ``nrowcol``, this is the exact coordinates. If provided
+        with ``resolution`` this is a closed upper bound. The exact coordinates of
+        the top right corner pixel will then be a multiple of the resolution with
+        respect to the bottom left corner.
+    nrowcol : tuple of int, optional
+        number of rows (y) and columns (x) respectively
+    resolution : float, optional
+        resolution of the destination grid in the units of the ``grid_mapping``
+    add_latlon : bool
+        If True, add latitude and longitude 2d arrays as coordinates
+    dtype : numpy.dtype
+        The dtype in which spatial coordinates are returned
+
+    Returns
+    -------
+    dst_grid : xarray.Dataset
+        dataset with spatial horizontal coordinates and an associated crs attribute
+
+    See also
+    --------
+    :func:`interp2grid`
+    """
+
+    start_y, start_x = bottomleft
+
+    if topright is not None and nrowcol is not None:
+        end_y, end_x = topright
+        nrows, ncols = nrowcol
+        y_coords = np.linspace(start_y, end_y, nrows, dtype=dtype)
+        x_coords = np.linspace(start_x, end_x, ncols, dtype=dtype)
+    elif topright is not None and resolution is not None:
+        end_y, end_x = topright
+        y_coords = np.linspace(start_y, end_y, num=int((end_y - start_y) / resolution) + 1, dtype=dtype)
+        x_coords = np.linspace(start_x, end_x, num=int((end_x - start_x) / resolution) + 1, dtype=dtype)
+    elif nrowcol is not None and resolution is not None:
+        nrows, ncols = nrowcol
+        end_y = start_y + (nrows - 1) * resolution
+        end_x = start_x + (ncols - 1) * resolution
+        y_coords = np.linspace(start_y, end_y, nrows, dtype=dtype)
+        x_coords = np.linspace(start_x, end_x, ncols, dtype=dtype)
+    else:
+        raise ValueError("At least two of 'topright', 'nrowcol', 'resolution' arguments must be provided")
+
+    dst_grid = xr.Dataset(
+        coords={X_DIM: x_coords, Y_DIM: y_coords},
+        attrs={"crs": dst_grid_crs},
+    )
+
+    if add_latlon:
+        x_2d, y_2d = np.meshgrid(x_coords, y_coords)
+        dst_crs = CRS.from_user_input(dst_grid_crs)
+        longitude, latitude = reproject(x_2d, y_2d, dst_crs, "EPSG:4326")
+        dst_grid = dst_grid.assign_coords(
+            {
+                LON: ((Y_DIM, X_DIM), longitude.astype(dtype, copy=False)),
+                LAT: ((Y_DIM, X_DIM), latitude.astype(dtype, copy=False)),
+            }
+        )
+
+    return dst_grid
+
+
+def create_dst_grid_from_src(  # pylint: disable=too-many-locals
+    src_grid: xr.Dataset,
+    dst_grid_crs: Any,
+    add_latlon: bool = False,
+) -> xr.Dataset:
+    """
+    Creates a destination grid on which data can be interpolated from a source grid
+    in a different Coordinate Reference System. Domain and sizes are preserved.
+    The output can be directly used with :func:`interp2grid`.
+
+    Parameters
+    ----------
+    src_grid : xarray Dataset
+        The source grid from which a new destination grid with same sizes and domain
+        is derived. Supported data models include:
+
+          - CF-convention grids (see https://cfconventions.org/)
+          - internal MCH models
+          - Unstructured grids
+
+        Typical spatial dimension and coordinate names are supported, and can be extended via environment variables:
+
+          - Y coordinate: export GRID_PROCESS__data_model__grid__y (e.g., '["y", "latitude", "lat", "dummy_y"]')
+          - X coordinate: export GRID_PROCESS__data_model__grid__x (e.g., '["x", "longitude", "lon", "dummy_x"]')
+          - Unstructured dimension: export GRID_PROCESS__data_model__grid__unstructured_dim (e.g., '["dummy_cell"]')
+          - When the spatial coordinates variables have a CF compliant `standard_name` attribute, arbitrary names can be
+            used directly.
+
+        Spatial CRS must be referenced via one of:
+
+          - Global attribute 'crs' (e.g., "EPSG:2056" or anything accepted by :meth:`pyproj.CRS.from_user_input`)
+          - Global attribute 'grid_mapping' (internal grid database convention)
+          - CF-convention grid mapping variable
+          - If coordinates are named 'latitude'/'longitude' or 'lat'/'lon', WGS84 (EPSG:4326) is assumed.
+    dst_grid_crs : string, int, dict or a CRS object
+        Anything accepted by pyproj.crs.CRS.from_user_input:
+        https://pyproj4.github.io/pyproj/dev/api/crs/crs.html#pyproj.crs.CRS.from_user_input
+        (e.g. EPSG:4326)
+    add_latlon : bool
+        If True, add latitude and longitude 2d arrays as coordinates
+
+    Returns
+    -------
+    dst_grid : xarray.Dataset
+        dataset with spatial horizontal coordinates and an associated crs attribute
+
+    See also
+    --------
+    :func:`interp2grid`
+    """
+
+    spatial_dims = dm.get_spatial_dims(src_grid)
+    src_grid = dm.set_grid_mapping(src_grid)
+    src_y, src_x, src_crs = dm.get_coords_name_crs(src_grid, spatial_dims)
+    dst_crs = CRS.from_user_input(dst_grid_crs)
+    x_min, x_max = src_grid[src_x].values.min(), src_grid[src_x].values.max()
+    y_min, y_max = src_grid[src_y].values.min(), src_grid[src_y].values.max()
+    x_bnds = (x_min, x_min, x_max, x_max)
+    y_bnds = (y_min, y_max, y_min, y_max)
+    x_bnds, y_bnds = reproject(x_bnds, y_bnds, src_crs, dst_crs)
+    x_bnds, y_bnds = (min(x_bnds), max(x_bnds)), (min(y_bnds), max(y_bnds))
+
+    if len(spatial_dims) == 1:  # need to infer expected sizes for x and y dims
+        xy_ratio = (x_bnds[1] - x_bnds[0]) / (y_bnds[1] - y_bnds[0])
+        dst_x_size = round(np.sqrt(src_grid.sizes[spatial_dims[0]] * xy_ratio))
+        dst_y_size = round(np.sqrt(src_grid.sizes[spatial_dims[0]] / xy_ratio))
+        bool_decrease = False
+    else:
+        dst_x_size, dst_y_size = (
+            src_grid.sizes[spatial_dims[1]],
+            src_grid.sizes[spatial_dims[0]],
+        )
+        bool_decrease = dm.is_decreasing(src_grid[src_y], spatial_dims[0])
+
+    src_x_dtype, src_y_dtype = src_grid[src_x].dtype, src_grid[src_y].dtype
+    dst_x_dtype = src_x_dtype if src_x_dtype == np.float64 else np.float32
+    dst_y_dtype = src_y_dtype if src_y_dtype == np.float64 else np.float32
+    dst_x_coords = np.linspace(x_bnds[0], x_bnds[1], dst_x_size, dtype=dst_x_dtype)
+    dst_y_coords = np.linspace(y_bnds[0], y_bnds[1], dst_y_size, dtype=dst_y_dtype)
+    if bool_decrease:
+        dst_y_coords = dst_y_coords[::-1]
+
+    dst_grid = xr.Dataset(
+        coords={X_DIM: dst_x_coords, Y_DIM: dst_y_coords},
+        attrs={"crs": dst_grid_crs},
+    )
+
+    if add_latlon:
+        x_2d, y_2d = np.meshgrid(dst_x_coords, dst_y_coords)
+        longitude, latitude = reproject(x_2d, y_2d, dst_crs, "EPSG:4326")
+        dst_grid = dst_grid.assign_coords(
+            {
+                LON: ((Y_DIM, X_DIM), longitude.astype(dst_x_dtype, copy=False)),
+                LAT: ((Y_DIM, X_DIM), latitude.astype(dst_y_dtype, copy=False)),
+            }
+        )
+
+    return dst_grid

grid_process/helpers.py

@@ -0,0 +1,174 @@
+# Copyright (c) 2026 MeteoSwiss, contributors listed in AUTHORS
+#
+# Distributed under the terms of the BSD 3-Clause License.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Module for general helper functions and decorators
+"""
+
+import logging
+import time
+from functools import wraps
+from typing import Any, overload
+from collections.abc import Callable, Hashable
+
+import xarray as xr
+import numpy as np
+from numpy.typing import NDArray
+from pyproj import Transformer, CRS
+
+logger = logging.getLogger(__name__)
+
+
+def timefn(func: Callable[..., Any]) -> Callable[..., Any]:
+    """
+    Generic decorator to log function processing times
+    """
+
+    @wraps(func)
+    def measure_time(*args: Any, **kwargs: Any) -> Any:
+        t0 = time.monotonic()
+        result = func(*args, **kwargs)
+        logger.debug("@timefn: %s took %.3f sec", func.__name__, time.monotonic() - t0)
+        return result
+
+    return measure_time
+
+
+@overload
+def make_accessor(func: Callable[..., xr.Dataset], accessor_class: type) -> Callable[..., xr.Dataset]: ...
+@overload
+def make_accessor(func: Callable[..., xr.DataArray], accessor_class: type) -> Callable[..., xr.DataArray]: ...
+
+
+def make_accessor(
+    func: Callable[..., xr.Dataset | xr.DataArray], accessor_class: type
+) -> Callable[..., xr.Dataset | xr.DataArray]:
+    """
+    Generic decorator to turn functions operating on xarray objects into
+    accessors based on a registered dataset or dataarray ``accessor_class``
+    """
+
+    @wraps(func)
+    def method(accessor: xr.Dataset | xr.DataArray, *args: Any, **kwargs: Any) -> xr.Dataset | xr.DataArray:
+        return func(accessor._data, *args, **kwargs)  # pylint: disable=protected-access  # own accessor classes
+
+    setattr(accessor_class, func.__name__, method)
+    return func
+
+
+def keepattrs(func: Callable[..., xr.Dataset]) -> Callable[..., xr.Dataset]:
+    """
+    Decorator to keep dataset and data variables attributes for functions
+    operating on xarray.Datasets
+    """
+
+    @wraps(func)
+    def wrapper(*args: xr.Dataset, **kwargs: Any) -> xr.Dataset:
+        vars_attrs = {name: da.attrs for name, da in args[0].data_vars.items()}
+        ds_attrs = args[0].attrs
+        ds_out = func(*args, **kwargs)
+        update_vars_attrs(ds_out, vars_attrs)
+        ds_out.attrs.update(ds_attrs)
+        return ds_out
+
+    return wrapper
+
+
+def update_vars_attrs(ds: xr.Dataset, vars_attrs: dict[Hashable, dict[Hashable, Any]]) -> None:
+    """
+    Updates data variables attributes in a xarray.Dataset. The operation is
+    done "in-place".
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        can contain one or several data variables
+    vars_attrs : dict
+        contains the attributes of each data variables in the input dataset
+    """
+
+    for varname, attributes in vars_attrs.items():
+        ds[varname].attrs.update(attributes)
+
+
+def replace_vars_attrs(ds: xr.Dataset, vars_attrs: dict[str, dict]) -> None:
+    """
+    Replaces data variables attributes in a xarray.Dataset. The operation is
+    done "in-place".
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        can contain one or several data variables
+    vars_attrs : dict
+        contains the new attributes of each data variables in the input dataset
+        to be replaced
+    """
+
+    for varname, attributes in vars_attrs.items():
+        ds[varname].attrs = attributes
+
+
+@overload
+def reproject(
+    x_coords: NDArray[np.floating],
+    y_coords: NDArray[np.floating],
+    src_crs: CRS | str,
+    dst_crs: CRS | str,
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]: ...
+@overload
+def reproject(
+    x_coords: list[float], y_coords: list[float], src_crs: CRS | str, dst_crs: CRS | str
+) -> tuple[list[float], list[float]]: ...
+@overload
+def reproject(
+    x_coords: tuple[float, ...],
+    y_coords: tuple[float, ...],
+    src_crs: CRS | str,
+    dst_crs: CRS | str,
+) -> tuple[tuple[float, ...], tuple[float, ...]]: ...
+
+
+def reproject(
+    x_coords: NDArray[np.floating] | list[float] | tuple[float, ...],
+    y_coords: NDArray[np.floating] | list[float] | tuple[float, ...],
+    src_crs: CRS | str,
+    dst_crs: CRS | str,
+) -> (
+    tuple[NDArray[np.float64], NDArray[np.float64]]
+    | tuple[list[float], list[float]]
+    | tuple[tuple[float, ...], tuple[float, ...]]
+):
+    """
+    Reprojects `x` and `y` coordinates from a source CRS to a new specified
+    one.
+
+    Parameters
+    ----------
+    x_coords : numpy.ndarray, list or tuple
+        contains `x` coordinates.
+    y_coords : numpy.ndarray, list or tuple
+        contains `y` coordinates
+    src_crs : pyproj.CRS of the source or equivalent string
+        The value can be anything accepted by
+        pyproj.CRS.from_user_input, such as an authority string
+        (e.g. "EPSG:4326") or a WKT string
+    dst_crs : pyproj.CRS of the destination or equivalent string
+        The value can be anything accepted by
+        pyproj.CRS.from_user_input, such as an authority string
+        (eg "EPSG:4326") or a WKT string
+
+    Returns
+    -------
+    x_coords_new : numpy.ndarray, list or tuple of float64
+        same as input ``x_coords`` in the new CRS and cast as float64
+    y_coords_new : numpy.ndarray, list or tuple of float64
+        same as input ``y_coords`` in the new CRS and cast as float64
+    """
+
+    transformer = Transformer.from_crs(src_crs, dst_crs, always_xy=True)
+
+    return transformer.transform(x_coords, y_coords)

grid_process/optimization.py

@@ -0,0 +1,50 @@
+# Copyright (c) 2026 MeteoSwiss, contributors listed in AUTHORS
+#
+# Distributed under the terms of the BSD 3-Clause License.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Module for optimization of computations
+"""
+
+import logging
+from functools import wraps
+from typing import Any, TypeVar, cast
+from collections.abc import Callable, Hashable
+
+logger = logging.getLogger(__name__)
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def memoize(maxsize: int = 10) -> Callable[[F], F]:
+    """
+    LRU cache decorator implementation for any arbitrary input since the
+    caching is purely based on the optional keyword argument ``hkey``.
+    """
+
+    def real_memoize(func: F) -> F:
+        cache: dict[Hashable, Any] = {}
+        hkeys = []
+
+        @wraps(func)
+        def _memoize(*args: Any, **kwargs: Any) -> Any:
+            hkey = kwargs.pop("hkey", None)
+            if hkey in cache:
+                logger.debug("@memoize: %s cached!", func.__name__)
+                return cache[hkey]
+            logger.debug("@memoize: %s computing ...", func.__name__)
+            result = func(*args, **kwargs)
+            if hkey is not None:
+                cache[hkey] = result
+                hkeys.append(hkey)
+                if len(hkeys) > maxsize:
+                    cache.pop(hkeys.pop(0))
+                logger.debug("@memoize: %s cache size is %s}", func.__name__, len(cache))
+
+            return result
+
+        return cast(F, _memoize)
+
+    return real_memoize