Merge pull request #835 from DHI/interpolator_dataclass

Interpolation ids and weights always go together

Author: ecomodeller

mikeio/_interpolation.py

@@ -1,31 +1,9 @@
 from __future__ import annotations
-from typing import TYPE_CHECKING, overload
+from dataclasses import dataclass
 import numpy as np
 
-if TYPE_CHECKING:
-    from .dataset import Dataset, DataArray
-
-from .spatial import GeometryUndefined
-
 
 def get_idw_interpolant(distances: np.ndarray, p: float = 2) -> np.ndarray:
-    """IDW interpolant for 2d array of distances.
-
-    https://pro.arcgis.com/en/pro-app/help/analysis/geostatistical-analyst/how-inverse-distance-weighted-interpolation-works.htm
-
-    Parameters
-    ----------
-    distances : array-like
-        distances between interpolation point and grid point
-    p : float, optional
-        power of inverse distance weighting, default=2
-
-    Returns
-    -------
-    np.array
-        weights
-
-    """
     is_1d = distances.ndim == 1
     if is_1d:
         distances = np.atleast_2d(distances)
@@ -45,123 +23,68 @@ def get_idw_interpolant(distances: np.ndarray, p: float = 2) -> np.ndarray:
     return weights
 
 
-@overload
-def interp2d(
-    data: np.ndarray | DataArray,
-    elem_ids: np.ndarray,
-    weights: np.ndarray | None = None,
-    shape: tuple[int, ...] | None = None,
-) -> np.ndarray: ...
-
-
-@overload
-def interp2d(
-    data: Dataset,
-    elem_ids: np.ndarray,
-    weights: np.ndarray | None = None,
-    shape: tuple[int, ...] | None = None,
-) -> Dataset: ...
-
-
-def interp2d(
-    data: Dataset | DataArray | np.ndarray,
-    elem_ids: np.ndarray,
-    weights: np.ndarray | None = None,
-    shape: tuple[int, ...] | None = None,
-) -> Dataset | np.ndarray:
-    """interp spatially in data (2d only).
-
-    Parameters
-    ----------
-    data : mikeio.Dataset, DataArray, or ndarray
-        dfsu data
-    elem_ids : ndarray(int)
-        n sized array of 1 or more element ids used for interpolation
-    weights : ndarray(float), optional
-        weights with same size as elem_ids used for interpolation
-    shape: tuple, optional
-            reshape output
-
-    Returns
-    -------
-    ndarray, Dataset, or DataArray
-        spatially interped data with same type and shape as input
-
-    Examples
-    --------
-    >>> elem_ids, weights = dfs.get_spatial_interpolant(coords)
-    >>> dsi = interp2d(ds, elem_ids, weights)
-
-    """
-    from .dataset import DataArray, Dataset
-
-    if isinstance(data, Dataset):
-        ds = data.copy()
-
-        ni = len(elem_ids)
-
-        interp_data_vars = {}
-
-        for da in ds:
-            key = da.name
-            if "time" not in da.dims:
-                idatitem = _interp_itemstep(da.to_numpy(), elem_ids, weights)
-                if shape:
-                    idatitem = idatitem.reshape(*shape)
-
-            else:
-                nt, _ = da.shape
-                # use dtype of da
-                idatitem = np.empty(shape=(nt, ni), dtype=da.values.dtype)
-                for step in range(nt):
-                    idatitem[step, :] = _interp_itemstep(
-                        da[step].to_numpy(), elem_ids, weights
-                    )
-                if shape:
-                    idatitem = idatitem.reshape((nt, *shape))
-
-            dims = ("time", "element")  # TODO is this the best?
-            interp_data_vars[key] = DataArray(
-                data=idatitem,
-                time=da.time,
-                dims=dims,
-                item=da.item,
-                geometry=GeometryUndefined(),
-            )
-
-        new_ds = Dataset(interp_data_vars, validate=False)
-        return new_ds
-
-    if isinstance(data, DataArray):
-        # TODO why doesn't this return a DataArray?
-        data = data.to_numpy()
-
-    if isinstance(data, np.ndarray):
+@dataclass
+class Interpolant:
+    ids: np.ndarray
+    weights: np.ndarray
+
+    @staticmethod
+    def from_distances(distances: np.ndarray, p: float = 2) -> np.ndarray:
+        """IDW interpolant for 2d array of distances.
+
+        https://pro.arcgis.com/en/pro-app/help/analysis/geostatistical-analyst/how-inverse-distance-weighted-interpolation-works.htm
+
+        Parameters
+        ----------
+        distances : array-like
+            distances between interpolation point and grid point
+        p : float, optional
+            power of inverse distance weighting, default=2
+
+        Returns
+        -------
+        np.array
+            weights
+
+        """
+        return get_idw_interpolant(distances, p)
+
+    def interp1d(self, data: np.ndarray) -> np.ndarray:
+        ids = self.ids
+        weights = self.weights
+        result = np.dot(data[:, ids], weights)
+        assert isinstance(result, np.ndarray)
+        return result
+
+    def interp2d(
+        self,
+        data: np.ndarray,
+    ) -> np.ndarray:
+        """interp spatially in data (2d only).
+
+        Parameters
+        ----------
+        data : ndarray
+            dfsu data
+
+        Returns
+        -------
+        ndarray
+            spatially interpolated data
+
+        """
+        weights = self.weights
+        elem_ids = self.ids
+
         if data.ndim == 1:
-            # data is single item and single time step
-            idatitem = _interp_itemstep(data, elem_ids, weights)
-            if shape:
-                idatitem = idatitem.reshape(*shape)
-            return idatitem
-
-    ni = len(elem_ids)
-    datitem = data
-    nt, _ = datitem.shape
-    idatitem = np.empty(shape=(nt, ni), dtype=datitem.dtype)
-    for step in range(nt):
-        idatitem[step, :] = _interp_itemstep(datitem[step], elem_ids, weights)
-    if shape:
-        idatitem = idatitem.reshape((nt, *shape))
-    return idatitem
-
-
-def _interp_itemstep(
-    data: np.ndarray,
-    elem_ids: np.ndarray,
-    weights: np.ndarray | None = None,
-) -> np.ndarray:
-    if weights is None:
-        return data[elem_ids]
-    else:
-        idat = data[elem_ids] * weights
-        return np.sum(idat, axis=1) if weights.ndim == 2 else idat
+            idat = data[elem_ids] * weights.astype(data.dtype)
+            return np.sum(idat, axis=1) if weights.ndim == 2 else idat
+        elif data.ndim == 2:
+            # data shape: (nt, nelem)
+
+            # data[:, elem_ids]: (nt, ni)
+            # weights: (ni,) or (ni, nweights)
+            idat = data[:, elem_ids] * weights.astype(data.dtype)  # broadcasting
+            return np.sum(idat, axis=-1) if weights.ndim == 2 else idat
+        else:
+            raise ValueError("data must be 1D or 2D array")

mikeio/_track.py

@@ -84,7 +84,7 @@ def _extract_track(
 
     # spatial interpolation
     n_pts = 1 if method == "nearest" else 5
-    elem_ids, weights = geometry.get_2d_interpolant(
+    interpolant = geometry.get_2d_interpolant(
         coords[i_start : (i_end + 1)], n_nearest=n_pts
     )
 
@@ -131,7 +131,9 @@ def is_EOF(step: int) -> bool:
             continue
 
         w = (t_rel[t] - t1) / timestep  # time-weight
-        eid = elem_ids[i_interp]
+        eid = interpolant.ids[i_interp]
+        weights = interpolant.weights
+        # TODO move to interpolation module?
         if np.any(eid > 0):
             dati = (1 - w) * np.dot(d1[:, eid], weights[i_interp])
             dati = dati + w * np.dot(d2[:, eid], weights[i_interp])

mikeio/dataset/_dataarray.py

@@ -25,13 +25,15 @@
 import pandas as pd
 from mikecore.DfsuFile import DfsuFileType
 
+
 from ..eum import EUMType, EUMUnit, ItemInfo
 from .._time import _get_time_idx_list, _n_selected_timesteps
 
 if TYPE_CHECKING:
     from ._dataset import Dataset
     import xarray
     from numpy.typing import ArrayLike
+    from mikeio._interpolation import Interpolant
 
 
 from ..spatial import (
@@ -1014,8 +1016,6 @@ def interp(
         x: float | None = None,
         y: float | None = None,
         z: float | None = None,
-        n_nearest: int = 3,
-        interpolant: tuple[Any, Any] | None = None,
         **kwargs: Any,
     ) -> DataArray:
         """Interpolate data in time and space.
@@ -1045,7 +1045,7 @@ def interp(
             z-coordinate of point to be interpolated to, by default None
         n_nearest : int, optional
             When using IDW interpolation, how many nearest points should
-            be used, by default: 3
+            be used.
         interpolant : tuple, optional
             Precomputed interpolant, by default None
         **kwargs: Any
@@ -1075,15 +1075,12 @@ def interp(
         if z is not None:
             raise NotImplementedError()
 
-        geometry: GeometryPoint2D | GeometryPoint3D | GeometryUndefined = (
-            GeometryUndefined()
-        )
+        geometry: GeometryPoint2D | GeometryUndefined = GeometryUndefined()
+        interpolant = kwargs.get("interpolant")
 
         # interp in space
-        if (x is not None) or (y is not None) or (z is not None):
-            coords = [(x, y)]
-
-            if isinstance(self.geometry, Grid2D):  # TODO DIY bilinear interpolation
+        if (x is not None) or (y is not None):
+            if isinstance(self.geometry, Grid2D):
                 if x is None or y is None:
                     raise ValueError("both x and y must be specified")
 
@@ -1094,31 +1091,27 @@ def interp(
                 )
             elif isinstance(self.geometry, Grid1D):
                 if interpolant is None:
-                    interpolant = self.geometry.get_spatial_interpolant(coords)  # type: ignore
-                dai = self.geometry.interp(self.to_numpy(), *interpolant).flatten()
+                    assert x is not None
+                    interpolant = self.geometry.get_spatial_interpolant(x)
+                dai = interpolant.interp1d(self.to_numpy()).flatten()
                 geometry = GeometryUndefined()
-            elif isinstance(self.geometry, GeometryFM3D):
-                raise NotImplementedError("Interpolation in 3d is not yet implemented")
             elif isinstance(self.geometry, GeometryFM2D):
                 if x is None or y is None:
                     raise ValueError("both x and y must be specified")
 
                 if interpolant is None:
                     interpolant = self.geometry.get_2d_interpolant(
-                        coords,  # type: ignore
-                        n_nearest=n_nearest,
+                        xy=[(x, y)],  # type: ignore
                         **kwargs,  # type: ignore
                     )
-                dai = self.geometry.interp2d(self, *interpolant).flatten()  # type: ignore
-                if z is None:
-                    geometry = GeometryPoint2D(
-                        x=x, y=y, projection=self.geometry.projection
-                    )
-                # this is not supported yet (see above)
-                # else:
-                #    geometry = GeometryPoint3D(
-                #        x=x, y=y, z=z, projection=self.geometry.projection
-                #    )
+                dai = interpolant.interp2d(self.to_numpy()).flatten()
+                geometry = GeometryPoint2D(
+                    x=x, y=y, projection=self.geometry.projection
+                )
+            else:
+                raise NotImplementedError(
+                    f"Interpolation in {self.geometry} is not yet implemented"
+                )
 
             da = DataArray(
                 data=dai,
@@ -1285,7 +1278,7 @@ def interp_na(self, axis: str = "time", **kwargs: Any) -> DataArray:
     def interp_like(
         self,
         other: DataArray | Grid2D | GeometryFM2D | pd.DatetimeIndex,
-        interpolant: tuple[Any, Any] | None = None,
+        interpolant: Interpolant | None = None,
         **kwargs: Any,
     ) -> DataArray:
         """Interpolate in space (and in time) to other geometry (and time axis).
@@ -1297,7 +1290,7 @@ def interp_like(
         ----------
         other: Dataset, DataArray, Grid2D, GeometryFM, pd.DatetimeIndex
             The target geometry (and time axis) to interpolate to
-        interpolant: tuple, optional
+        interpolant: Interpolant, optional
             Reuse pre-calculated index and weights
         **kwargs: Any
             additional kwargs are passed to interpolation method
@@ -1337,16 +1330,19 @@ def interp_like(
             raise NotImplementedError()
 
         if interpolant is None:
-            elem_ids, weights = self.geometry.get_2d_interpolant(xy, **kwargs)
-        else:
-            elem_ids, weights = interpolant
+            interpolant = self.geometry.get_2d_interpolant(xy, **kwargs)
 
         if isinstance(geom, (Grid2D, GeometryFM2D)):
-            shape = (geom.ny, geom.nx) if isinstance(geom, Grid2D) else None
+            ari = interpolant.interp2d(data=self.to_numpy())
+            if isinstance(geom, Grid2D):
+                shape = (
+                    (self.n_timesteps, geom.ny, geom.nx)
+                    if self.dims[0] == "time"
+                    else (geom.ny, geom.nx)
+                )
+                ari = ari.reshape(shape)
 
-            ari = self.geometry.interp2d(
-                data=self.to_numpy(), elem_ids=elem_ids, weights=weights, shape=shape
-            )
+            assert ari.dtype == self.dtype
         else:
             raise NotImplementedError(
                 "Interpolation to other geometry not yet supported"

mikeio/dataset/_dataset.py

@@ -811,7 +811,6 @@ def interp(
         x: float | None = None,
         y: float | None = None,
         z: float | None = None,
-        n_nearest: int = 3,
         **kwargs: Any,
     ) -> Dataset:
         """Interpolate data in time and space.
@@ -878,7 +877,6 @@ def interp(
             ):  # TODO remove this when all geometries implements the same method
                 interpolant = self.geometry.get_2d_interpolant(
                     xy,  # type: ignore
-                    n_nearest=n_nearest,
                     **kwargs,  # type: ignore
                 )
                 das = [da.interp(x=x, y=y, interpolant=interpolant) for da in self]