Interpolate single-variable datasets without copying, coarsen skipping NaN

CHANGELOG.rst

@@ -2,6 +2,29 @@
 Changelog
 =========
 
+Unreleased
+----------
+
+**Changes of results** (compared to 1.1.0, see the entries below for details)
+
+* the coarsening of the source grid by ``interp2grid`` (to a coarser destination grid) skips NaN: a coarsened
+  cell with some NaN values gets the mean of its valid values instead of NaN, only cells with all values NaN stay NaN
+* coarsened float32 values differ from before by float32 rounding, growing with the window size (for 2 m
+  temperature, up to 3e-5 K for a 2 x 2 window and 2e-4 K for 10 x 10)
+
+**Changes**
+
+* changed ``interp2poi`` and ``interp2grid`` to no longer copy a dataset with a single data variable before
+  interpolation (new ``copy`` parameter of ``ds_utils.ds_to_da``, default ``True``). For the operational grids
+  (with cached neighbors), POIs are interpolated up to about 1.5 times faster (e.g. ICON-CH2 ``nearest`` 0.18 s →
+  0.12 s) with 0.1 to 0.2 GB less memory for ICON, and IFS to the seamless 2 km grid with ``nearest_alt`` 0.39 s →
+  0.24 s. A sliced source that is not contiguous in memory is now copied into contiguous rows per grid point, which
+  is slower for large sources: ICON-CH2 in LV95 to the seamless 2 km grid with ``nearest_alt`` or ``temp_fixgrad``
+  0.23 s → 0.29 s and 0.34 s → 0.39 s
+* changed the coarsening of the source grid by ``interp2grid`` to skip NaN (e.g. for the snow limit, NaN below the
+  orography), with bottleneck. ICON-CH1 to the seamless 2 km grid (with cached neighbors) is about 1.4 times
+  faster (e.g. ``nearest_alt`` 0.30 s → 0.21 s), with about 0.05 GB more peak memory (copy of the windows)
+
 1.1.0 (2026-11-12)
 ------------------
 

grid_process/ds_utils.py

@@ -19,6 +19,7 @@
 
 def ds_to_da(
     ds: xr.Dataset,
+    copy: bool = True,
 ) -> tuple[xr.DataArray, dict[Hashable, dict[Hashable, Any]]]:
     """
     Transform a dataset containing model data and one or several data
@@ -30,6 +31,10 @@ def ds_to_da(
     ----------
     ds : xarray.Dataset
         contain one or more variables
+    copy : bool, optional
+        if False and ``ds`` has a single data variable spanning all its dimensions, the dataarray is a view of its
+        data, without a copy: writing into it modifies ``ds``. Several variables are always stacked (copied).
+        Default is True.
 
     Returns
     -------
@@ -45,6 +50,14 @@ def ds_to_da(
     """
 
     vars_attrs = {varname: da.attrs for varname, da in ds.data_vars.items()}
+    if not copy and len(ds.data_vars) == 1:
+        name = next(iter(ds.data_vars))
+        if set(ds[name].dims) == set(ds.dims):  # otherwise to_array keeps coordinates on other dimensions
+            # same as to_array, whose np.stack copies the data
+            da = ds[name].expand_dims(variable=[name])
+            da.attrs = dict(ds.attrs)
+            da.name = None
+            return da, vars_attrs
     da = ds.to_array()
     return da, vars_attrs
 

grid_process/grid_interp.py

@@ -16,6 +16,7 @@
 from collections.abc import Callable, Hashable
 from math import isclose, prod, ceil
 
+import bottleneck as bn
 import numpy as np
 from numpy.typing import NDArray
 import pandas as pd
@@ -1009,7 +1010,7 @@ def _array2ds(data: xr.DataArray) -> xr.Dataset:
         return _interp_regular(data, spatial_dims, poi_info, method)  # type: ignore[arg-type]
 
     data = slice_boundingbox(data, poi_info, method, spatial_dims, **kwargs)
-    data, vars_attrs = ds_to_da(data)  # type: ignore[assignment]
+    data, vars_attrs = ds_to_da(data, copy=False)  # type: ignore[assignment]
     data = data.transpose(*spatial_dims, ...)
     coords_flat, data_flat, pois_coords = _flatten_concat(
         data,  # type: ignore[arg-type]
@@ -1697,7 +1698,7 @@ def _array2ds(data: xr.DataArray) -> xr.Dataset:
     dst_grid = dst_grid.transpose(*dst_spatial_dims)
 
     data = slice_coarsen(data, dst_grid, method, src_spatial_dims, dst_spatial_dims, **kwargs)
-    data, vars_attrs = ds_to_da(data)  # type: ignore[assignment]
+    data, vars_attrs = ds_to_da(data, copy=False)  # type: ignore[assignment]
     data = data.transpose(*src_spatial_dims, ...)
     coords_flat, data_flat, dst_coords_flat = _flatten_concat_grids(
         data,  # type: ignore[arg-type]
@@ -1792,8 +1793,7 @@ def _coarsen_interp() -> xr.Dataset:
                 )
                 if ratio > 1
             }
-            # reduce(np.mean) propagates NaNs like mean(skipna=False), which fails with xarray 2026.9.0
-            interp_data = interp_data.coarsen(coarsen_specs, boundary="trim").reduce(np.mean)
+            interp_data = interp_data.coarsen(coarsen_specs, boundary="trim").reduce(_window_nanmean)
 
         interp_data = _y_increasing(interp_data, src_spatial_dims[0])
         x_targets = x_coords
@@ -2064,12 +2064,23 @@ def extra_buffer(ratio: int) -> int:
                 for dim, ratio in coarsen_specs.items()
             }
         )
-        # reduce(np.mean) propagates NaNs like mean(skipna=False), which fails with xarray 2026.9.0
-        return data.isel(sub_domain).coarsen(coarsen_specs, boundary="trim").reduce(np.mean)
+        return data.isel(sub_domain).coarsen(coarsen_specs, boundary="trim").reduce(_window_nanmean)
 
     return data_sub
 
 
+def _window_nanmean(values: np.ndarray, axis: tuple[int, ...]) -> np.ndarray:
+    """
+    Mean over the coarsening window ``axis`` (for ``coarsen(...).reduce``) ignoring NaN: a block is NaN only if all
+    its values are NaN. bottleneck, with the window moved last and flattened, is faster than ``np.mean`` over the
+    window axes. It sums float32 in float32: the relative error grows with the window size (about 1e-7 for a 2 x 2
+    window, 2e-6 for 40 x 40).
+    """
+
+    values = np.moveaxis(values, axis, tuple(range(-len(axis), 0)))
+    return bn.nanmean(values.reshape(values.shape[: -len(axis)] + (-1,)), axis=-1)
+
+
 def get_buffer(method: str, **kwargs: Any) -> int:
     """
     Buffer specific to the interpolation method to add when slicing the input dataset to the target grid or POIs

test/test_ds_utils.py

@@ -22,6 +22,31 @@ def test_ds_to_da(model_data):
     assert expected_dims == set(da_out.dims)
 
 
+def test_ds_to_da_no_copy(model_data):
+    """
+    With copy=False, a single variable gives the same dataarray as to_array, as a view of its data. Several
+    variables, or a variable not spanning all dimensions, are stacked (copied)
+    """
+
+    single = model_data[[next(iter(model_data.data_vars))]].assign_attrs(source="test")
+    name = next(iter(single.data_vars))
+    da_view, vars_attrs = dsu.ds_to_da(single, copy=False)
+
+    assert da_view.identical(single.to_array())
+    assert vars_attrs == {name: single[name].attrs}
+    assert np.shares_memory(da_view.values, single[name].values)
+    assert not np.shares_memory(dsu.ds_to_da(single)[0].values, single[name].values)  # default: copy
+
+    if len(model_data.data_vars) > 1:
+        da_stack, _ = dsu.ds_to_da(model_data, copy=False)
+        assert da_stack.identical(model_data.to_array())
+        assert not np.shares_memory(da_stack.values, model_data[name].values)
+
+    other_dim = single.assign_coords(extra=("extra", [1, 2]))  # a dimension of a coordinate only
+    da_other, _ = dsu.ds_to_da(other_dim, copy=False)
+    assert da_other.identical(other_dim.to_array())
+
+
 def test_cast_dtype(model_data):
     """Simple testcase for the cast_dtype function from ds_utils module"""
 

test/test_grid_interp.py

@@ -1822,3 +1822,54 @@ def test_get_buffer(method, kwargs, expected):
     from grid_process.grid_interp import get_buffer  # pylint: disable=import-outside-toplevel
 
     assert get_buffer(method, **kwargs) == expected
+
+
+def test_window_nanmean():
+    """
+    Coarsening skips NaN (a block is NaN only if all its values are), keeps float32, and leaves integers to np.mean
+    """
+
+    from grid_process.grid_interp import _window_nanmean  # pylint: disable=import-outside-toplevel
+
+    values = np.arange(32, dtype=np.float32).reshape(2, 4, 4)
+    values[0, 0, 1] = np.nan  # block (0, 0) of the first field: 3 valid values
+    values[1, 2:, 2:] = np.nan  # block (1, 1) of the second field: all NaN
+    da = xr.DataArray(values, dims=("step", "y", "x"))
+
+    out = da.coarsen(y=2, x=2, boundary="trim").reduce(_window_nanmean)
+
+    assert out.dtype == np.float32
+    with warnings.catch_warnings():  # np.nanmean warns for the all-NaN block, bottleneck does not
+        warnings.simplefilter("ignore", RuntimeWarning)
+        expected = np.nanmean(values.reshape(2, 2, 2, 2, 2), axis=(2, 4))
+    np.testing.assert_allclose(out.values, expected, rtol=1e-6)
+    assert out.values[0, 0, 0] == np.float32((0 + 4 + 5) / 3)
+    assert np.isnan(out.values[1, 1, 1])
+    assert np.isnan(out.values).sum() == 1
+
+    ints = xr.DataArray(np.arange(16, dtype=np.int16).reshape(4, 4), dims=("y", "x"))
+    np.testing.assert_array_equal(
+        ints.coarsen(y=2, x=2).reduce(_window_nanmean), ints.coarsen(y=2, x=2).reduce(np.mean)
+    )
+
+
+def test_window_nanmean_large_window():
+    """Large float32 windows with NaN: close to the float64 mean (float32 sums), all-NaN blocks stay NaN"""
+
+    from grid_process.grid_interp import _window_nanmean  # pylint: disable=import-outside-toplevel
+
+    rng = np.random.default_rng(0)
+    values = (rng.standard_normal((3, 2, 20, 30)) * 10 + 280).astype(np.float32)
+    values[rng.random(values.shape) < 0.3] = np.nan
+    values[0, 0, :10, :10] = np.nan  # an all-NaN block
+    da = xr.DataArray(values, dims=("step", "member", "y", "x"))
+
+    out = da.coarsen(y=10, x=10).reduce(_window_nanmean)
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", RuntimeWarning)
+        expected = np.nanmean(values.astype(np.float64).reshape(3, 2, 2, 10, 3, 10), axis=(3, 5))
+    assert out.dtype == np.float32
+    np.testing.assert_allclose(out.values, expected, rtol=1e-6)
+    assert np.isnan(out.values[0, 0, 0, 0])
+    assert np.isnan(out.values).sum() == 1