From cd24e93121714ad336a07db65b59fbd86ebb798f Mon Sep 17 00:00:00 2001 From: Mathieu Schaer Date: Thu, 12 Nov 2026 05:16:52 +0100 Subject: [PATCH 1/2] Interpolate a single-variable dataset without copying it interp2poi and interp2grid stacked the data variables with Dataset.to_array, whose np.stack copies the data even for a single variable. ds_to_da gets a copy parameter (default True, unchanged): with copy=False, a single data variable spanning all the dimensions is returned as a view (expand_dims), identical to to_array. Both functions only read the array. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.rst | 13 +++++++++++++ grid_process/ds_utils.py | 13 +++++++++++++ grid_process/grid_interp.py | 4 ++-- test/test_ds_utils.py | 25 +++++++++++++++++++++++++ 4 files changed, 53 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7a5d723..3e8165b 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -2,6 +2,19 @@ Changelog ========= +Unreleased +---------- + +**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 + 1.1.0 (2026-11-12) ------------------ diff --git a/grid_process/ds_utils.py b/grid_process/ds_utils.py index 0bc4b69..1fdd297 100644 --- a/grid_process/ds_utils.py +++ b/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 diff --git a/grid_process/grid_interp.py b/grid_process/grid_interp.py index 9c50b83..a44bbde 100644 --- a/grid_process/grid_interp.py +++ b/grid_process/grid_interp.py @@ -1009,7 +1009,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 +1697,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] diff --git a/test/test_ds_utils.py b/test/test_ds_utils.py index 7125c68..5c015c4 100644 --- a/test/test_ds_utils.py +++ b/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""" From 01e36a0158c39b4714e2aa5a1d7040216d142b80 Mon Sep 17 00:00:00 2001 From: Mathieu Schaer Date: Thu, 12 Nov 2026 05:18:05 +0100 Subject: [PATCH 2/2] Coarsen the source grid skipping NaN, with bottleneck The coarsening of interp2grid (source downsampled to a coarser destination) used coarsen().reduce(np.mean), which gives NaN for a block with any NaN. _window_nanmean averages the valid values of each block with bottleneck.nanmean (only all-NaN blocks stay NaN), e.g. for the snow limit, NaN below the orography. With the window moved last and flattened, it is faster than np.mean (ICON-CH1 to the seamless 2 km grid 0.30 s -> 0.21 s). bottleneck sums float32 in float32: the error grows with the window size (about 1e-7 relative for 2 x 2, 2e-6 for 40 x 40). Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.rst | 10 ++++++++ grid_process/grid_interp.py | 19 +++++++++++--- test/test_grid_interp.py | 51 +++++++++++++++++++++++++++++++++++++ 3 files changed, 76 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 3e8165b..57e5fc2 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -5,6 +5,13 @@ 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 @@ -14,6 +21,9 @@ Unreleased 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) ------------------ diff --git a/grid_process/grid_interp.py b/grid_process/grid_interp.py index a44bbde..58ad3cf 100644 --- a/grid_process/grid_interp.py +++ b/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 @@ -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 diff --git a/test/test_grid_interp.py b/test/test_grid_interp.py index 3263038..0d794ac 100644 --- a/test/test_grid_interp.py +++ b/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