Skip to content

Interpolate single-variable datasets without copying, coarsen skipping NaN - #11

Merged
matschaer merged 2 commits into
mainfrom
perf/single-var-nan-coarsen
Nov 12, 2026
Merged

matschaer merged 2 commits into
mainfrom
perf/single-var-nan-coarsen

Conversation

@matschaer

Copy link
Copy Markdown
Collaborator

Summary

  • No copy of single-variable datasets (cd24e93): interp2poi and interp2grid call ds_utils.ds_to_da(..., copy=False). When there is one data variable and it spans all dimensions, they get a view via expand_dims instead of the np.stack copy made by to_array. The default stays copy=True.
  • NaN-skipping coarsening (01e36a0): both coarsening sites (_coarsen_interp and slice_coarsen) use coarsen(...).reduce(_window_nanmean). This is a 2-line helper that moves the window axes last and calls bottleneck.nanmean. A coarsened cell with some NaN gets the mean of its valid values; only all-NaN cells stay NaN. This prepares for regridding the snow limit, which is NaN below the orography.

Performance (operational grids, later calls with cached neighbors)

Case main branch
ICON-CH2 → POIs, nearest 0.18 s 0.12 s (0.1–0.2 GB less memory)
IFS → seamless 2 km, nearest_alt 0.39 s 0.24 s
ICON-CH1 → seamless 2 km, nearest_alt 0.32 s 0.22 s (+0.05 GB peak)
gridefix-read (1.1 M cells → 252 × 376, 100 fields) 0.13 s 0.11 s

Known regression: ICON-CH2 in LV95 → seamless 2 km gets slower, from 0.23 to 0.29 s with nearest_alt and from 0.34 to 0.39 s with temp_fixgrad. The source there is a sliced, non-contiguous view, and the reshape now copies it with a gather. Compacting the data as main did is slower in 9 of the 12 operational grid cases.

Changes of results

  • Coarsened cells with partial NaN are no longer NaN (see above).
  • Coarsened float32 values differ by float32 rounding, which grows with the window size. For 2 m temperature that is up to 3e-5 K for a 2 × 2 window and 2.4e-4 K for 10 × 10 (global 1° benchmark); the relative error is about 2e-6 for a 40 × 40 window. ERA5-Land → RhiresY is identical.

Tests

  • New test_ds_to_da_no_copy: checks the view, the default copy, multiple variables, and a dimension that only a coordinate has.
  • New test_window_nanmean and test_window_nanmean_large_window: check partial and all-NaN blocks, float32, integers equal to np.mean, and a 10 × 10 window within rtol=1e-6 of the float64 nanmean.
  • All 262 tests pass; ruff and mypy are clean; pylint gives 10.00.

🤖 Generated with Claude Code

Mathieu Schaer and others added 2 commits November 12, 2026 05:16
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
@matschaer
matschaer merged commit 8a43f7c into main Nov 12, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant