grid-process 1.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,2016 @@
1
+ # Copyright (c) 2026 MeteoSwiss, contributors listed in AUTHORS
2
+ #
3
+ # Distributed under the terms of the BSD 3-Clause License.
4
+ #
5
+ # SPDX-License-Identifier: BSD-3-Clause
6
+
7
+ # pylint: disable=too-many-lines
8
+
9
+ """
10
+ Module for spatial interpolation of gridded datasets
11
+ """
12
+
13
+ import logging
14
+ from functools import partial
15
+ from typing import Any, Literal, overload
16
+ from collections.abc import Callable
17
+ from math import isclose, prod, ceil
18
+
19
+ import numpy as np
20
+ from numpy.typing import NDArray
21
+ import xarray as xr
22
+ import geopandas as gpd
23
+ from scipy.interpolate import griddata
24
+ from scipy.spatial import KDTree
25
+ from pyproj import CRS
26
+
27
+ from grid_process import CONFIG
28
+ from grid_process.helpers import make_accessor, reproject
29
+ from grid_process import data_model as dm
30
+ from grid_process.optimization import memoize
31
+ from grid_process.ds_utils import ds_to_da, cast_dtype
32
+ from grid_process.da_utils import da_to_ds
33
+
34
+ logger = logging.getLogger(__name__)
35
+
36
+
37
+ # Register xarray accessor so module functions can also be called as methods
38
+ # on xr.Dataset objects
39
+ @xr.register_dataset_accessor("grid_interp")
40
+ class GridInterp:
41
+ """
42
+ Xarray accessors class for spatial interpolation of gridded data
43
+ """
44
+
45
+ def __init__(self, ds: xr.Dataset):
46
+ self._data = ds
47
+
48
+
49
+ dataset_accessor = partial(make_accessor, accessor_class=GridInterp)
50
+
51
+ POI_MODEL = CONFIG.data_model.poi
52
+ POI_DIM = POI_MODEL.dim
53
+ POI_COORD_COL = POI_MODEL.coord_from
54
+ SRC_ATTRS_KEEP = ("source", "title", "institution", "history", "references", "comment")
55
+ SWISS_CRS = ("EPSG:2056", "EPSG:21781")
56
+
57
+ InterpMethods = (
58
+ Literal[
59
+ "idw",
60
+ "temp",
61
+ "temp_fixgrad",
62
+ "nearest_alt",
63
+ "snowlim",
64
+ "nearest_3D",
65
+ "linear_3D",
66
+ "idw_3D",
67
+ "upscaled",
68
+ "upscaled_3D",
69
+ ]
70
+ | Literal["nearest", "linear", "cubic"]
71
+ )
72
+
73
+
74
+ def interp_temp( # pylint: disable=too-many-locals
75
+ coords_flat: np.ndarray,
76
+ data_flat: np.ndarray,
77
+ pois_coords: np.ndarray,
78
+ z_fact: int = 25,
79
+ k: int = 36,
80
+ weight_power: float | None = None,
81
+ grad_only: bool = False,
82
+ chunk_size: int = 20,
83
+ ) -> np.ndarray:
84
+ """
85
+ Interpolation from grid to POIs for temperature data, based on local
86
+ soundings. The grid does not need to be regular. The sounding curves are
87
+ fitted on the k nearest neighbors considering a distance that more heavily
88
+ penalizes elevation differences for ``z_fact`` > 1.
89
+
90
+ Parameters
91
+ ----------
92
+ coords_flat : numpy.ndarray, shape (n_grids, m)
93
+ array of flattened grid coordinates. n_grids is the number of model
94
+ grid points, m the number of coordinates dimensions (i.e. 3)
95
+ data_flat : numpy.ndarray, shape (n_grids, ...)
96
+ array containing the flattened model data with various trailing
97
+ dimensions
98
+ pois_coords : numpy.ndarray , shape (n_pois, m)
99
+ array of POI coordinates. n_pois is the number of POIs, m the number of
100
+ coordinates dimensions (i.e. 3)
101
+ z_fact : positive int, optional
102
+ factor to multiply the z coordinates to find nearest neighbors based on a
103
+ non-euclidean distance that penalizes more altitude differences. Default
104
+ value is 50. A value of 1 will use regular euclidean distances.
105
+ k : positive int, optional
106
+ number of considered nearest neighbors for the local sounding. Default
107
+ is 36
108
+ weight_power : positive float, optional
109
+ a weight vector is passed to the sounding fit according to
110
+ ``w = d^(-weight_power)``, with d being the (non-euclidean) distances between
111
+ the considered POI and the k nearest grid points. The greater is the
112
+ ``weight_power``, the stronger is the penalization with distance. A
113
+ ``weight_power`` of 0 is equivalent to None (i.e. same weights for all
114
+ points). By default, gives same weight for all points
115
+ grad_only : bool, optional
116
+ if True, only uses the local gradient to correct the nearest gridpoint
117
+ temperature based on the altitude difference with the POI. Default is
118
+ False
119
+ chunk_size : positive int, optional
120
+ number of POIs processed per sounding batch. Smaller values reduce peak
121
+ memory (dominant cost: ``chunk_size × k × n_trail × 12`` bytes) and
122
+ keep the working set in L3 cache, which is typically faster when
123
+ ``n_trail`` is large. Larger values reduce loop overhead and are better
124
+ for small trailing dimensions or when n_pois is small. Default is 20.
125
+
126
+ Returns
127
+ -------
128
+ poi_data : numpy.ndarray, shape (n_pois, ...)
129
+ array containing the interpolated values and the same trailing
130
+ dimensions as the input ``data_flat`` array
131
+
132
+ Raises
133
+ ------
134
+ ValueError : Raised if the orography is flat, as no sounding can be estimated.
135
+
136
+ See also
137
+ --------
138
+ :func:`interp2poi`
139
+
140
+ Note
141
+ ----
142
+ For POIs with altitude higher than 300m above the highest gridpoint or lower
143
+ than 500m below the lowest gridpoint in the local neighborhood used for the fit,
144
+ the fixgrad method is used with default gradient of 0.006 [K / m].
145
+ """
146
+
147
+ # make sure k is smaller or equal to the number of grid points
148
+ k = min(k, len(coords_flat))
149
+
150
+ # rescale z coords
151
+ coords_flat_zfact = coords_flat.copy()
152
+ pois_coords_zfact = pois_coords.copy()
153
+ coords_flat_zfact[:, -1] *= z_fact
154
+ pois_coords_zfact[:, -1] *= z_fact
155
+
156
+ # find indexes and distances of k nearest neighbors in the non-euclidean dist
157
+ dist_nn, ind_nn = nn_lookup_cached(coords_flat_zfact, pois_coords_zfact, k)
158
+
159
+ # preallocate output array
160
+ data_dtype = data_flat.dtype
161
+ poi_data = np.empty((pois_coords.shape[0],) + data_flat.shape[1:], dtype=data_dtype)
162
+
163
+ # gather neighbor altitudes for all POIs at once (small: n_pois × k floats)
164
+ all_alt = coords_flat[ind_nn, -1] # (n_pois, k)
165
+
166
+ if weight_power is not None:
167
+ # weights for the sounding fit according to the distance
168
+ dist_weight: np.ndarray | None = _compute_dist_weight(dist_nn, weight_power, data_dtype)
169
+ else:
170
+ dist_weight = None
171
+
172
+ poi_alts = pois_coords[:, -1]
173
+
174
+ chunk_size = max(1, chunk_size)
175
+
176
+ if grad_only:
177
+ # grad_only always uses a deg=1 fit regardless of altitude position, so
178
+ # the deg3/deg1 split is unnecessary. Only min/max quantiles are needed
179
+ # to classify fixgrad vs sounding POIs.
180
+ alt_minmax = np.quantile(all_alt, [0, 1], axis=1) # shape (2, n_pois)
181
+ fixgrad_mask = (poi_alts < alt_minmax[0] - 600) | (poi_alts > alt_minmax[1] + 200)
182
+ sounding_groups = [(np.nonzero(~fixgrad_mask)[0], 1, True)]
183
+ else:
184
+ # classify each POI into three groups: fixgrad, deg1 sounding, deg3 sounding
185
+ alt_q = np.quantile(all_alt, [0, 0.15, 0.85, 1], axis=1) # shape (4, n_pois)
186
+ fixgrad_mask = (poi_alts < alt_q[0] - 600) | (poi_alts > alt_q[-1] + 200)
187
+ sounding_mask = ~fixgrad_mask
188
+ deg3_mask = sounding_mask & (alt_q[1] < poi_alts) & (poi_alts < alt_q[2])
189
+ deg1_mask = sounding_mask & ~deg3_mask
190
+ sounding_groups = [
191
+ (np.nonzero(deg1_mask)[0], 1, False),
192
+ (np.nonzero(deg3_mask)[0], 3, False),
193
+ ]
194
+
195
+ # fixgrad group: only needs the nearest neighbor per POI (same for both branches)
196
+ fg_idx = np.nonzero(fixgrad_mask)[0]
197
+ if fg_idx.size > 0:
198
+ poi_data[fg_idx] = _fixgrad(
199
+ data_flat[ind_nn[fg_idx, 0]],
200
+ all_alt[fg_idx, 0],
201
+ poi_alts[fg_idx],
202
+ )
203
+
204
+ # sounding groups: process in memory-bounded chunks
205
+ for idx, deg, go in sounding_groups:
206
+ for start in range(0, idx.size, chunk_size):
207
+ chunk_idx = idx[start : start + chunk_size]
208
+ chunk_w = dist_weight[chunk_idx] if dist_weight is not None else None
209
+ poi_data[chunk_idx] = _sounding_batch(
210
+ data_flat[ind_nn[chunk_idx]],
211
+ all_alt[chunk_idx],
212
+ poi_alts[chunk_idx],
213
+ deg=deg,
214
+ w=chunk_w,
215
+ grad_only=go,
216
+ )
217
+
218
+ return poi_data
219
+
220
+
221
+ def _normalize_altitudes(
222
+ altitudes: np.ndarray,
223
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
224
+ """
225
+ Normalize altitudes per POI to improve numerical conditioning of the
226
+ polynomial fit.
227
+
228
+ Parameters
229
+ ----------
230
+ altitudes : numpy.ndarray, shape (n_sub, k)
231
+ neighbor altitudes for each POI
232
+
233
+ Returns
234
+ -------
235
+ alt_norm : numpy.ndarray, shape (n_sub, k)
236
+ normalized altitudes with zero mean and unit standard deviation per POI
237
+ alt_mean : numpy.ndarray, shape (n_sub, 1)
238
+ per-POI altitude mean used for normalization
239
+ alt_std : numpy.ndarray, shape (n_sub, 1)
240
+ per-POI altitude standard deviation used for normalization
241
+ """
242
+
243
+ alt_mean = altitudes.mean(axis=1, keepdims=True)
244
+ alt_std = altitudes.std(axis=1, keepdims=True).clip(1e-10)
245
+ return (altitudes - alt_mean) / alt_std, alt_mean, alt_std
246
+
247
+
248
+ def _build_vandermonde(alt_norm: np.ndarray, deg: int) -> np.ndarray:
249
+ """
250
+ Build a batched Vandermonde matrix from normalized altitudes.
251
+
252
+ Parameters
253
+ ----------
254
+ alt_norm : numpy.ndarray, shape (n_sub, k)
255
+ normalized neighbor altitudes
256
+ deg : positive int
257
+ polynomial degree
258
+
259
+ Returns
260
+ -------
261
+ V : numpy.ndarray, shape (n_sub, k, deg+1)
262
+ Vandermonde matrix with columns [1, x, x^2, ..., x^deg]
263
+ """
264
+
265
+ exponents = np.arange(deg + 1, dtype=np.float64)
266
+ return alt_norm[:, :, None] ** exponents
267
+
268
+
269
+ def _apply_weights(vandermonde: np.ndarray, values: np.ndarray, w: np.ndarray | None) -> tuple[np.ndarray, np.ndarray]:
270
+ """
271
+ Apply per-POI weights to the Vandermonde matrix and right-hand side.
272
+
273
+ Parameters
274
+ ----------
275
+ vandermonde : numpy.ndarray, shape (n_sub, k, deg+1)
276
+ Vandermonde matrix
277
+ values : numpy.ndarray, shape (n_sub, k, n_trail)
278
+ flattened data values
279
+ w : numpy.ndarray, shape (n_sub, k) or None
280
+ weight vector per POI; if None, uniform weights are used
281
+
282
+ Returns
283
+ -------
284
+ vw : numpy.ndarray, shape (n_sub, k, deg+1)
285
+ weight-scaled Vandermonde matrix
286
+ yw : numpy.ndarray, shape (n_sub, k, n_trail)
287
+ weight-scaled right-hand side
288
+ """
289
+
290
+ if w is not None:
291
+ ww = w[:, :, None].astype(np.float64, copy=False)
292
+ return vandermonde * ww, values * ww
293
+ return vandermonde, values
294
+
295
+
296
+ def _qr_polyfit_batch(vw: np.ndarray, yw: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
297
+ """
298
+ Solve a batch of weighted least-squares problems via QR decomposition and
299
+ detect rank-deficient systems.
300
+
301
+ Parameters
302
+ ----------
303
+ vw : numpy.ndarray, shape (n_sub, k, deg+1)
304
+ weighted Vandermonde matrix
305
+ yw : numpy.ndarray, shape (n_sub, k, n_trail)
306
+ weighted right-hand side
307
+
308
+ Returns
309
+ -------
310
+ coeffs : numpy.ndarray, shape (n_sub, deg+1, n_trail)
311
+ polynomial coefficients; zero for rank-deficient systems
312
+ rank_deficient : numpy.ndarray, shape (n_sub,), dtype bool
313
+ True for systems where the R diagonal indicates near-singularity
314
+ """
315
+
316
+ q_mat, r_mat = np.linalg.qr(vw)
317
+ qt_b = q_mat.swapaxes(-2, -1) @ yw # (n_sub, deg+1, n_trail)
318
+ diag_abs = np.abs(np.diagonal(r_mat, axis1=-2, axis2=-1)) # (n_sub, deg+1)
319
+ rank_ok = diag_abs.min(axis=-1) >= 1e-10 * diag_abs.max(axis=-1)
320
+ coeffs = np.zeros_like(qt_b)
321
+ ok_idx = np.nonzero(rank_ok)[0]
322
+ if ok_idx.size > 0:
323
+ coeffs[ok_idx] = np.linalg.solve(r_mat[ok_idx], qt_b[ok_idx])
324
+ return coeffs, ~rank_ok
325
+
326
+
327
+ def _eval_poly_batch(coeffs: np.ndarray, poi_norm: np.ndarray, deg: int) -> np.ndarray:
328
+ """
329
+ Evaluate a batch of polynomials at normalized POI altitudes.
330
+
331
+ Parameters
332
+ ----------
333
+ coeffs : numpy.ndarray, shape (n_sub, deg+1, n_trail)
334
+ polynomial coefficients in the normalized altitude space
335
+ poi_norm : numpy.ndarray, shape (n_sub,)
336
+ normalized target altitudes for each POI
337
+ deg : positive int
338
+ polynomial degree
339
+
340
+ Returns
341
+ -------
342
+ result : numpy.ndarray, shape (n_sub, n_trail)
343
+ evaluated polynomial values
344
+ """
345
+
346
+ exponents = np.arange(deg + 1, dtype=np.float64)
347
+ vander_poi = poi_norm[:, None] ** exponents # (n_sub, deg+1)
348
+ return (vander_poi[:, None, :] @ coeffs)[:, 0, :] # (n_sub, n_trail)
349
+
350
+
351
+ def _sounding_batch( # pylint: disable=too-many-locals
352
+ data_array: np.ndarray,
353
+ altitudes: np.ndarray,
354
+ poi_alts: np.ndarray,
355
+ deg: int = 3,
356
+ w: np.ndarray | None = None,
357
+ grad_only: bool = False,
358
+ ) -> np.ndarray:
359
+ """
360
+ Returns predicted temperature values for a batch of POIs according to the
361
+ local sounding, using a batched QR decomposition.
362
+
363
+ Each POI has its own neighborhood and independent polynomial fit. All fits
364
+ are computed simultaneously via ``numpy.linalg.qr`` and
365
+ ``numpy.linalg.solve`` to avoid per-POI Python-level dispatch.
366
+
367
+ Parameters
368
+ ----------
369
+ data_array : numpy.ndarray, shape (n_sub, k, ...)
370
+ contains the temperature values of the k nearest neighbor grid points
371
+ for each of the n_sub POIs in this batch
372
+ altitudes : numpy.ndarray, shape (n_sub, k)
373
+ contains the altitudes of the k nearest neighbor grid points for each POI
374
+ poi_alts : numpy.ndarray, shape (n_sub,)
375
+ altitude of each POI in this batch
376
+ deg : positive int
377
+ degree of the polynomial fit. Only used if ``grad_only=False``
378
+ w : numpy.ndarray, shape (n_sub, k) or None
379
+ per-POI weight vectors for the polynomial fit. If None, uniform weights
380
+ are used for all points
381
+ grad_only : bool
382
+ if True, only uses the local gradient to correct the temperature of the
383
+ nearest gridpoint based on the altitude difference with the POI.
384
+ In this case, the gradient is bound to [-0.01, 0.03] K m-1. Default is False
385
+
386
+ Returns
387
+ -------
388
+ fitted_val : numpy.ndarray, shape (n_sub, ...)
389
+ array of fitted temperatures for each POI in the batch
390
+
391
+ See also
392
+ --------
393
+ :func:`interp_temp`
394
+ """
395
+
396
+ n_sub = data_array.shape[0]
397
+ trail_shape = data_array.shape[2:]
398
+ n_trail = int(np.prod(trail_shape)) if trail_shape else 1
399
+
400
+ # flatten trailing dims and cast to float64 for numerical stability
401
+ values = data_array.reshape(n_sub, data_array.shape[1], n_trail, order="C").astype(np.float64, copy=False)
402
+ alt_norm, alt_mean, alt_std = _normalize_altitudes(altitudes)
403
+
404
+ if grad_only:
405
+ vandermonde = _build_vandermonde(alt_norm, 1)
406
+ vw, yw = _apply_weights(vandermonde, values, w)
407
+ coeffs, _ = _qr_polyfit_batch(vw, yw)
408
+ # convert slope from normalized space back to original altitude units
409
+ slope = (coeffs[:, 1] / alt_std[:, 0, None]).clip(-0.01, 0.03)
410
+ fitted_val = values[:, 0] + slope * (poi_alts[:, None] - altitudes[:, 0, None])
411
+ else:
412
+ vandermonde = _build_vandermonde(alt_norm, deg)
413
+ vw, yw = _apply_weights(vandermonde, values, w)
414
+ coeffs, rank_deficient = _qr_polyfit_batch(vw, yw)
415
+
416
+ poi_norm = (poi_alts - alt_mean[:, 0]) / alt_std[:, 0]
417
+ fitted_val = _eval_poly_batch(coeffs, poi_norm, deg).reshape(n_sub, *trail_shape)
418
+
419
+ # overwrite rank-deficient results with a deg=1 fallback. If deg is already 1, raise an error instead.
420
+ bad_idx = np.nonzero(rank_deficient)[0]
421
+ if bad_idx.size > 0:
422
+ if deg == 1:
423
+ raise ValueError(
424
+ f"deg=1 sounding is rank-deficient for {bad_idx.size} POI(s): "
425
+ "all k neighbor altitudes are identical. Check the surface altitude coordinate."
426
+ )
427
+ fitted_val[bad_idx] = _sounding_batch(
428
+ data_array[bad_idx],
429
+ altitudes[bad_idx],
430
+ poi_alts[bad_idx],
431
+ deg=1,
432
+ w=w[bad_idx] if w is not None else None,
433
+ grad_only=False,
434
+ )
435
+
436
+ return fitted_val.reshape((n_sub,) + trail_shape)
437
+
438
+
439
+ def interp_temp_fixgrad(
440
+ coords_flat: np.ndarray,
441
+ data_flat: np.ndarray,
442
+ pois_coords: np.ndarray,
443
+ r_search: float = 2,
444
+ z_weight: int = 200,
445
+ grad: float = 0.006,
446
+ ) -> np.ndarray:
447
+ """
448
+ Interpolation from grid to POIs for temperature data by taking the
449
+ surrounding gridpoint with the most similar altitude and further correct
450
+ for altitude difference with a fixed gradient.
451
+
452
+ Parameters
453
+ ----------
454
+ coords_flat : numpy.ndarray, shape (n_grids, m)
455
+ array of flattened grid coordinates. n_grids is the number of model
456
+ grid points, m the number of coordinates dimensions (i.e. 3)
457
+ data_flat : numpy.ndarray, shape (n_grids, ...)
458
+ array containing the flattened model data with various trailing
459
+ dimensions
460
+ pois_coords : numpy.ndarray , shape (n_pois, m)
461
+ array of POI coordinates. n_pois is the number of POIs, m the number of
462
+ coordinates dimensions (i.e. 3)
463
+ r_search : positive float, optional
464
+ radius in model gridpoints that define the number of neighbors included
465
+ in the search. Default is 4
466
+ z_weight : positive int, optional
467
+ weight for the z axis in the distance computation of the "most adequat"
468
+ gridpoint. ``d = horizontal_dist + z_weight * vertical_dist``.
469
+ Default value is 200
470
+ grad : positive float
471
+ gradient in [K / m] to correct for altitude difference between POI and
472
+ nearest grid point. Default is 0.006 [K / m]
473
+
474
+ Returns
475
+ -------
476
+ poi_data : numpy.ndarray, shape (n_pois, ...)
477
+ array containing the interpolated values and the same trailing
478
+ dimensions as the input ``data_flat`` array.
479
+
480
+ See also
481
+ --------
482
+ :func:`interp2poi`
483
+ """
484
+
485
+ ind = _grid_assignation(coords_flat, pois_coords, r_search, z_weight)
486
+
487
+ return _fixgrad(
488
+ data_flat[ind,],
489
+ coords_flat[ind, -1],
490
+ pois_coords[:, -1],
491
+ grad,
492
+ )
493
+
494
+
495
+ def _fixgrad(
496
+ best_gp_data: np.ndarray,
497
+ best_gp_alt: np.ndarray,
498
+ pois_alt: np.ndarray,
499
+ grad: float = 0.006,
500
+ ) -> np.ndarray:
501
+ """
502
+ Correct temperature based on altitudinal differences and a fix gradient.
503
+
504
+ Parameters
505
+ ----------
506
+ best_gp_data : numpy.ndarray, shape (n_pois, ...)
507
+ contains the temperature values of the most appropriate gridpoint for
508
+ each POI
509
+ best_gp_alt : numpy.ndarray, shape (n_pois,)
510
+ contains the altitude of the most appropriate gridpoint for each POI.
511
+ pois_alt : numpy.ndarray, shape (n_pois,)
512
+ contains the altitude of each POI
513
+ grad : positive float
514
+ gradient in [K / m] to correct for altitude difference between POI and
515
+ nearest grid point. Default is 0.006 [K / m]
516
+
517
+ Returns
518
+ -------
519
+ poi_data : numpy.ndarray, shape (n_pois, ...)
520
+ array of altitude corrected temperatures
521
+
522
+ See also
523
+ --------
524
+ :func:`interp_temp_fixgrad`
525
+ """
526
+
527
+ ndim_trail = best_gp_data.ndim - 1
528
+ delta_t = (best_gp_alt - pois_alt) * grad
529
+ delta_t = delta_t.astype(best_gp_data.dtype, copy=False)
530
+
531
+ return best_gp_data + delta_t.reshape(delta_t.shape + (1,) * ndim_trail)
532
+
533
+
534
+ def interp_nearest_alt(
535
+ coords_flat: np.ndarray,
536
+ data_flat: np.ndarray,
537
+ pois_coords: np.ndarray,
538
+ r_search: float = 2,
539
+ z_weight: int = 500,
540
+ ) -> np.ndarray:
541
+ """
542
+ Interpolation from grid to POIs for elevation dependent data. The method finds
543
+ the nearest grid point according to a distance that use a much larger weight
544
+ along the z axis. For wind data for instance, grid points at similar altitudes
545
+ are preferred. The grid does not need to be regular.
546
+
547
+ Parameters
548
+ ----------
549
+ coords_flat : numpy.ndarray, shape (n_grids, m)
550
+ array of flattened grid coordinates. n_grids is the number of model
551
+ grid points, m the number of coordinates dimensions (i.e. 3)
552
+ data_flat : numpy.ndarray, shape (n_grids, ...)
553
+ array containing the flattened model data with various trailing
554
+ dimensions
555
+ pois_coords : numpy.ndarray , shape (n_pois, m)
556
+ array of POI coordinates. n_pois is the number of POIs, m the number of
557
+ coordinates dimensions (i.e. 3)
558
+ r_search : positive float, optional
559
+ radius in model gridpoints that define the number of neighbors included
560
+ in the search. Default is 4
561
+ z_weight : positive int, optional
562
+ weight for the z axis in the distance computation of the "most adequat"
563
+ gridpoint. ``d = horizontal_dist + z_weight * vertical_dist``.
564
+ Default value is 500
565
+
566
+ Returns
567
+ -------
568
+ poi_data : numpy.ndarray, shape (n_pois, ...)
569
+ array containing the interpolated values and the same trailing
570
+ dimensions as the input ``data_flat`` array
571
+
572
+ See also
573
+ --------
574
+ :func:`interp2poi`
575
+ """
576
+
577
+ ind = _grid_assignation(coords_flat, pois_coords, r_search, z_weight)
578
+
579
+ return data_flat[ind,]
580
+
581
+
582
+ def _grid_assignation(
583
+ coords_flat: np.ndarray,
584
+ pois_coords: np.ndarray,
585
+ r_search: float,
586
+ z_weight: int,
587
+ deltah_min: int = 50,
588
+ ) -> np.ndarray:
589
+ """
590
+ Returns the index of the "most appropriate" gridpoint based on a distance
591
+ metric that gives a much greater weight to vertical distances. This is the
592
+ standard used for verification in COSMO.
593
+ """
594
+
595
+ k = round(np.pi * r_search**2)
596
+ # find indexes and distances of nearest neighbor grids for all POIs
597
+ dist_nn_horz, ind_nn = nn_lookup_cached(coords_flat[:, :2], pois_coords[:, :2], k)
598
+
599
+ dist_nn_vert = np.abs(coords_flat[ind_nn, -1] - pois_coords[:, [-1]])
600
+ dist_weighted = dist_nn_horz + z_weight * (dist_nn_vert - deltah_min).clip(0)
601
+
602
+ return ind_nn[np.arange(ind_nn.shape[0]), np.argmin(dist_weighted, axis=1)]
603
+
604
+
605
+ def interp_idw(
606
+ coords_flat: np.ndarray,
607
+ data_flat: np.ndarray,
608
+ pois_coords: np.ndarray,
609
+ k: int = 4,
610
+ weight_power: float = 1.0,
611
+ ) -> np.ndarray:
612
+ """
613
+ Performs interpolations from grid to POIs using inverse distance
614
+ weightning.
615
+
616
+ Parameters
617
+ ----------
618
+ coords_flat : numpy.ndarray, shape (n_grids,m)
619
+ array of flattened grid coordinates. n_grids is the number of model
620
+ grid points, m the number of coordinates dimensions (i.e. 2 or 3)
621
+ data_flat : numpy.ndarray, shape (n_grids, ...)
622
+ array containing the flattened model data with various trailing
623
+ dimensions
624
+ pois_coords : numpy.ndarray , shape (n_pois, m)
625
+ array of POI coordinates. n_pois is the number of POIs, m the number of
626
+ coordinates dimensions (i.e. 2 or 3)
627
+ k : positive int, optional
628
+ number of considered nearest neighbors. Default is 4
629
+ weight_power : positive float, optional
630
+ the interpolated value is the weighted average of the k nearest grid
631
+ points accordi g to ``w = d^(-weight_power)``, with d being the
632
+ distance between the considered POI and the gird points. The greater is
633
+ the ``weight_power``, the stronger is the penalization with distance.
634
+ A ``weight_power`` of 0 produces a normal average of the k nearest grid
635
+ points
636
+
637
+ Returns
638
+ -------
639
+ poi_data : numpy.ndarray, shape (n_pois, ...)
640
+ array containing the interpolated values and the same trailing
641
+ dimensions as the input ``data_flat`` array
642
+
643
+ See also
644
+ --------
645
+ :func:`interp2poi`
646
+ """
647
+
648
+ # find indexes and distances of the k nearest neighbor grids for all POIS
649
+ dist_nn, ind_nn = nn_lookup_cached(coords_flat, pois_coords, k)
650
+
651
+ # compute the weights
652
+ dist_weight = _compute_dist_weight(dist_nn, weight_power, data_flat.dtype)
653
+
654
+ # vectorized weigthed average
655
+ ndim_trail = data_flat.ndim - 1
656
+ poi_data = np.sum(
657
+ data_flat[ind_nn,] * dist_weight.reshape(dist_weight.shape + (1,) * ndim_trail),
658
+ axis=1,
659
+ )
660
+ return poi_data
661
+
662
+
663
+ def interp_upscaled(
664
+ coords_flat: np.ndarray,
665
+ data_flat: np.ndarray,
666
+ pois_coords: np.ndarray,
667
+ k: int = 4,
668
+ agg_func: Callable = np.max,
669
+ ) -> np.ndarray:
670
+ """
671
+ Perform interpolations from grid to POIs by computing a group statistic
672
+ over the k nearest neighbors (e.g. upscaling).
673
+
674
+ Parameters
675
+ ----------
676
+ coords_flat : numpy.ndarray, shape (n_grids, m)
677
+ array of flattened grid coordinates. n_grids is the number of model
678
+ grid points, m the number of coordinates dimensions (i.e. 2 or 3)
679
+ data_flat : numpy.ndarray, shape (n_grids, ...)
680
+ array containing the flattened model data with various trailing
681
+ dimensions
682
+ pois_coords : numpy.ndarray , shape (n_pois, m)
683
+ array of POI coordinates. n_pois is the number of POIs, m the number of
684
+ coordinates dimensions (i.e. 2 or 3).
685
+ k : positive int, optional
686
+ number of considered nearest neighbors. Default is 4
687
+ agg_func : callable, optional
688
+ specify the spatial aggregation function to use (i.e. numpy.mean,
689
+ numpy.max, numpy.min, numpy.std). Default is numpy.max.
690
+
691
+ Returns
692
+ -------
693
+ poi_data : numpy.ndarray, shape (n_pois, ...)
694
+ array containing the interpolated values and the same trailing
695
+ dimensions as the input ``data_flat`` array
696
+
697
+ See also
698
+ --------
699
+ :func:`interp2poi`
700
+ """
701
+
702
+ # find indexes and distances of the k nearest neighbor grids for all POIS
703
+ _, ind_nn = nn_lookup_cached(coords_flat, pois_coords, k)
704
+
705
+ # group aggregation
706
+ return agg_func(
707
+ data_flat[ind_nn,],
708
+ axis=1,
709
+ )
710
+
711
+
712
+ def interp_nearest(coords_flat: np.ndarray, data_flat: np.ndarray, pois_coords: np.ndarray) -> np.ndarray:
713
+ """
714
+ Performs nearest neighbor interpolation from grid to POIs.
715
+
716
+ Parameters
717
+ ----------
718
+ coords_flat : numpy.ndarray, shape (n_grids, m)
719
+ array of flattened grid coordinates. n_grids is the number of model
720
+ grid points, m the number of coordinates dimensions (i.e. 2 or 3).
721
+ data_flat : numpy.ndarray, shape (n_grids, ...)
722
+ array containing the flattened model data with various trailing
723
+ dimensions
724
+ pois_coords : numpy.ndarray , shape (n_pois, m)
725
+ array of POI coordinates. n_pois is the number of POIs, m the number of
726
+ coordinates dimensions (i.e. 2 or 3).
727
+
728
+ Returns
729
+ -------
730
+ poi_data : numpy.ndarray, shape (n_pois, ...)
731
+ array containing the interpolated values and the same trailing
732
+ dimensions as the input ``data_flat`` array
733
+
734
+ See also
735
+ --------
736
+ :func:`interp2poi`
737
+ """
738
+
739
+ # find indexes and distances of the k nearest neighbor grids for all POIS
740
+ _, ind_nn = nn_lookup_cached(coords_flat, pois_coords, 1)
741
+
742
+ return data_flat[ind_nn]
743
+
744
+
745
+ def interp_snowlim(
746
+ coords_flat: np.ndarray,
747
+ data_flat: np.ndarray,
748
+ pois_coords: np.ndarray,
749
+ k: int = 36,
750
+ d_max: float = 40000,
751
+ ) -> np.ndarray:
752
+ """
753
+ Special interpolation routine for snow limit. Snow limits below model
754
+ topography are undefined, resulting in datasets with randomly distributed
755
+ NaNs. This routine checks for the nearest non-NaN grid points. If this
756
+ point is further away than d_max, the snowlimit is set to a value between
757
+ 0 and 200 [m] depending on the POI altitude. Otherwise, the minimum between
758
+ the nearest non-NaN grid point and the minimum altitude of NaN grid points
759
+ among the k nearest neighbors is taken.
760
+
761
+ Parameters
762
+ ----------
763
+ coords_flat : numpy.ndarray, shape (n_grids, m)
764
+ array of flattened grid coordinates. n_grids is the number of model
765
+ grid points, m the number of coordinates dimensions (i.e. 3)
766
+ data_flat : numpy.ndarray, shape (n_grids, ...)
767
+ Array containing the flattened model data with various trailing
768
+ dimension.
769
+ pois_coords : numpy.ndarray , shape (n_pois, m)
770
+ array of POI coordinates. n_pois is the number of POIs, m the number of
771
+ coordinates dimensions (i.e. 3)
772
+ k : positive int
773
+ Number of considered nearest neighbors to check for altitude of NaN
774
+ grid points.
775
+ d_max : positive float
776
+ if all grid points within a radius of d_max meter from the POI are NaNs
777
+ the snow limit is set to a value between 0 and 200 [m] depending on
778
+ the POI altitude
779
+
780
+ Returns
781
+ -------
782
+ poi_data : numpy.ndarray, shape (n_pois, ...)
783
+ array containing the interpolated values and the same trailing
784
+ dimensions as the input ``data_flat`` array
785
+
786
+ See also
787
+ --------
788
+ :func:`interp2poi`
789
+ """
790
+
791
+ data_shape = data_flat.shape
792
+ n_pois = len(pois_coords)
793
+
794
+ poi_dummy = _fillna_alt( # NaNs filled with model topo
795
+ coords_flat, data_flat, pois_coords, min(data_shape[0], k)
796
+ )
797
+
798
+ data_flat = data_flat.reshape((data_shape[0], -1), order="C")
799
+
800
+ # array with default values when there is only NaNs within distance d_max
801
+ poi_data = (pois_coords[:, 2] - 200).clip(0, 200).astype(data_flat.dtype, copy=False)
802
+ poi_data -= poi_data % 50 # floor to multiple of 50
803
+ poi_data = np.broadcast_to(poi_data[:, None], (n_pois, data_flat.shape[-1])).copy()
804
+
805
+ for idx, array in enumerate(data_flat.T):
806
+ ind = ~np.isnan(array)
807
+ if ind.any():
808
+ dist_nn, ind_nn = nn_lookup(coords_flat[ind, :2], pois_coords[:, :2], 1)
809
+ bool_dist = dist_nn < d_max
810
+ poi_data[bool_dist, idx] = array[ind][ind_nn[bool_dist]]
811
+
812
+ poi_data = poi_data.reshape((n_pois,) + data_shape[1:], order="C")
813
+
814
+ return np.minimum(poi_data, poi_dummy)
815
+
816
+
817
+ def _fillna_alt(coords_flat: np.ndarray, data_flat: np.ndarray, pois_coords: np.ndarray, k: int) -> np.ndarray:
818
+ """
819
+ Interp to POIs with nearest neighbor and fill NaNs by simply taking the
820
+ altitude of the lowest grid point among the k nearest neighbors that have
821
+ an undefined snow limit.
822
+
823
+ Parameters
824
+ ----------
825
+ coords_flat : numpy.ndarray, shape (n_grids, m)
826
+ array of flattened grid coordinates. n_grids is the number of model
827
+ grid points, m the number of coordinates dimensions (i.e. 3)
828
+ data_flat : numpy.ndarray, shape (n_grids, ...)
829
+ array containing the flattened model data with various trailing
830
+ dimensions
831
+ pois_coords : numpy.ndarray , shape (n_pois, m)
832
+ array of POI coordinates. n_pois is the number of POIs, m the number of
833
+ coordinates dimensions (i.e. 3)
834
+ k : positive int
835
+ number of considered nearest neighbors to check for altitude of NaN
836
+
837
+ Returns
838
+ -------
839
+ data_poi : numpy.ndarray, shape (n_pois, ...)
840
+ array containing the interpolated values/model topography with same
841
+ trailing dimensions as the input ``data_flat`` array.
842
+
843
+ See also
844
+ -------
845
+ :func:`interp_snowlim`
846
+ """
847
+
848
+ ndim_trail = data_flat.ndim - 1
849
+ _, ind_nn = nn_lookup(coords_flat[:, :2], pois_coords[:, :2], k)
850
+ data_poi = data_flat[ind_nn]
851
+ ind_na = np.isnan(data_poi)
852
+ alt_poi = coords_flat[ind_nn, -1].reshape(ind_nn.shape + (1,) * ndim_trail)
853
+ alt_poi = np.broadcast_to(alt_poi, data_poi.shape)
854
+ alt_poi = np.amin(alt_poi, axis=1, initial=10000, where=ind_na)
855
+ data_poi = data_poi[:, 0]
856
+ data_poi[ind_na[:, 0]] = alt_poi[ind_na[:, 0]]
857
+
858
+ return data_poi
859
+
860
+
861
+ @dataset_accessor
862
+ def interp2poi(
863
+ data: xr.Dataset,
864
+ poi_info: gpd.GeoDataFrame,
865
+ method: InterpMethods,
866
+ try_reg: bool = True,
867
+ **kwargs: Any,
868
+ ) -> xr.Dataset:
869
+ """
870
+ Interpolates gridded spatial data to a set of Points of Interest (POIs).
871
+
872
+ This function provides a unified interface for interpolating values from regular or unstructured spatial grids
873
+ to arbitrary POI locations, supporting a variety of interpolation methods and data models. It replaces the
874
+ original spatial grid dimensions in `data` with a POI dimension, containing interpolated values for each POI.
875
+
876
+ Parameters
877
+ ----------
878
+ data : xarray.Dataset
879
+ The source grid dataset. Supported data models include:
880
+
881
+ - CF-convention grids (see https://cfconventions.org/)
882
+ - internal MCH models
883
+ - Unstructured grids
884
+
885
+ Typical spatial dimension and coordinate names are supported, and can be extended via environment variables:
886
+
887
+ - Y coordinate: export GRID_PROCESS__data_model__grid__y (e.g., '["y", "latitude", "lat", "dummy_y"]')
888
+ - X coordinate: export GRID_PROCESS__data_model__grid__x (e.g., '["x", "longitude", "lon", "dummy_x"]')
889
+ - Unstructured dimension: export GRID_PROCESS__data_model__grid__unstructured_dim (e.g., '["dummy_cell"]')
890
+ - When the spatial coordinates variables have a CF compliant `standard_name` attribute, arbitrary names can be
891
+ used directly.
892
+
893
+ Spatial CRS must be referenced via one of:
894
+
895
+ - Global attribute 'crs' (e.g., "EPSG:2056" or anything accepted by :meth:`pyproj.CRS.from_user_input`)
896
+ - Global attribute 'grid_mapping' (internal grid database convention)
897
+ - CF-convention grid mapping variable
898
+ - If coordinates are named 'latitude'/'longitude' or 'lat'/'lon', WGS84 (EPSG:4326) is assumed.
899
+
900
+ For methods requiring altitude, ensure a surface altitude coordinate is present; extend supported names via
901
+ GRID_PROCESS__data_model__grid__surf_alt.
902
+ poi_info : geopandas.GeoDataFrame
903
+ POI locations and attributes. Must have geometry and a defined CRS.
904
+ For methods requiring altitude, a POI altitude column or z-component in geometry (in units of meters) is needed.
905
+ Extend supported POI altitude column names via:
906
+
907
+ - export GRID_PROCESS__data_model__poi__alt_from (e.g., '["height_masl", "altitude", "dummy_height"]')
908
+
909
+ The column `name` is used for the output POI coordinate by default; override via:
910
+
911
+ - export GRID_PROCESS__data_model__poi__coord_from (e.g., dummy_name)
912
+ method : str
913
+ Interpolation method. Supported options include:
914
+
915
+ {'nearest', 'linear', 'cubic', 'idw', 'temp', 'temp_fixgrad', 'nearest_alt', 'snowlim',
916
+ 'nearest_3D', 'linear_3D', 'idw_3D', 'upscaled', 'upscaled_3D'}.
917
+
918
+ 3D methods include altitude in neighbor selection. For unstructured grids, prefer 'idw' with k=4 and
919
+ weight_power=2 over 'linear' for performence reasons.
920
+ try_reg : bool, default True
921
+ If True, attempts regular grid interpolation for structured grids with 1D coordinates when ``method`` is
922
+ 'nearest' or 'linear'. Regular grid interpolation is typically more efficient for large grids, but may be less
923
+ so for small grids with many trailing dimensions.
924
+ **kwargs : Any
925
+ Additional arguments passed to the specific interpolation sub-function, named `interp_<method>`.
926
+
927
+ Returns
928
+ -------
929
+ xr.Dataset
930
+ Dataset where original spatial dimensions are replaced with a POI dimension containing interpolated values.
931
+ The output POI dimension name can be overridden via GRID_PROCESS__data_model__poi__dim (e.g. station).
932
+
933
+ See also
934
+ --------
935
+ :func:`interp_temp`, :func:`interp_temp_fixgrad`, :func:`interp_nearest_alt`,
936
+ :func:`interp_idw`, :func:`interp_upscaled`, :func:`interp_nearest`,
937
+ :func:`interp_snowlim`, :func:`scipy.interpolate.griddata`, :meth:`xarray.Dataset.interp`
938
+
939
+ Notes
940
+ -----
941
+ - Grid points containing only NaNs across all dimensions are excluded from interpolation, enabling easy masking
942
+ (e.g., excluding water pixels with `ds.where(ds.fr_land > 0.5).interp2poi(..., try_reg=False)`).
943
+ - For regular grid interpolation, all NaN grid points are kept.
944
+ - Environment variables allow flexible support for custom grid and POI data models.
945
+ """
946
+
947
+ def _fullnan_ds(data: xr.DataArray) -> xr.Dataset:
948
+ """
949
+ Returns a dataset full of NaNs with correct dimensions
950
+ """
951
+
952
+ da = (
953
+ data.isel({dim: 0 for dim in spatial_dims}, drop=True)
954
+ .expand_dims({POI_DIM: poi_info[POI_COORD_COL].to_numpy(dtype=str)})
955
+ .copy()
956
+ )
957
+ da.attrs = _add_attrs(data.attrs)
958
+ da[:] = np.nan
959
+ return da_to_ds(da, vars_attrs)
960
+
961
+ def _array2ds(data: xr.DataArray) -> xr.Dataset:
962
+ """
963
+ Converts an array with interpolated values into a dataset with correct
964
+ dimensions, coordinates and attributes
965
+ """
966
+
967
+ da = xr.DataArray(
968
+ poi_data,
969
+ dims=(POI_DIM,) + data.dims[len(spatial_dims) :],
970
+ coords={
971
+ POI_DIM: poi_info[POI_COORD_COL].to_numpy(dtype=str),
972
+ **data.isel({dim: 0 for dim in spatial_dims}, drop=True).coords,
973
+ },
974
+ attrs=_add_attrs(data.attrs),
975
+ )
976
+ return da_to_ds(da, vars_attrs)
977
+
978
+ spatial_dims = dm.get_spatial_dims(data)
979
+ data = dm.set_grid_mapping(data)
980
+ cond_reg = method in ("nearest", "linear") and len(spatial_dims) > 1 and set(spatial_dims).issubset(data.coords)
981
+ if try_reg and cond_reg:
982
+ return _interp_regular(data, spatial_dims, poi_info, method) # type: ignore[arg-type]
983
+
984
+ data = slice_boundingbox(data, poi_info, method, spatial_dims, **kwargs)
985
+ data, vars_attrs = ds_to_da(data) # type: ignore[assignment]
986
+ data = data.transpose(*spatial_dims, ...)
987
+ coords_flat, data_flat, pois_coords = _flatten_concat(
988
+ data, # type: ignore[arg-type]
989
+ spatial_dims,
990
+ poi_info,
991
+ method,
992
+ )
993
+
994
+ if len(coords_flat) == 0:
995
+ return _fullnan_ds(data) # type: ignore[arg-type]
996
+
997
+ if method in INTERP_FUNC:
998
+ poi_data = INTERP_FUNC[method](coords_flat, data_flat, pois_coords, **kwargs) # type: ignore[operator]
999
+
1000
+ elif method in ("linear", "linear_3D", "cubic"):
1001
+ poi_data = griddata(coords_flat, data_flat, pois_coords, method.rsplit("_")[0])
1002
+ poi_data = poi_data.astype(data_flat.dtype, copy=False)
1003
+ else:
1004
+ raise ValueError(f"Unknown interpolation method {method!r}")
1005
+
1006
+ return _array2ds(data) # type: ignore[arg-type]
1007
+
1008
+
1009
+ def _interp_regular( # pylint: disable=too-many-locals
1010
+ data: xr.Dataset,
1011
+ spatial_dims: tuple[str, str],
1012
+ poi_info: gpd.GeoDataFrame,
1013
+ method: Literal["nearest", "linear"] = "linear",
1014
+ ) -> xr.Dataset:
1015
+ """
1016
+ Wrapper for interpolation from a regular grid to POIs.
1017
+ Handle reprojection of the points to the native crs of the input grid.
1018
+ Also assigns coordinate labels to the output poi dimension according to the
1019
+ POI_COORD_COL column of ``poi_info``. Used in :func:`interp2poi`
1020
+ """
1021
+
1022
+ def _interp_cast_dtype(
1023
+ method: Literal["nearest", "linear"],
1024
+ ) -> xr.Dataset:
1025
+ """
1026
+ Wrapper for xarray.Dataset.interp to return data in the same dtype as
1027
+ input.
1028
+ """
1029
+
1030
+ poi_data = data.interp({x_dim: pois_coords[0], y_dim: pois_coords[1]}, method, assume_sorted=True).drop_vars(
1031
+ spatial_dims, errors="ignore"
1032
+ )
1033
+
1034
+ cast_dtype(poi_data, data)
1035
+
1036
+ return poi_data
1037
+
1038
+ data = _drop_coords_reg(data, spatial_dims)
1039
+ y_dim, x_dim = spatial_dims
1040
+ data = _y_increasing(data, y_dim)
1041
+
1042
+ # Reproject POIs coordinates in the grid CRS
1043
+ pois_coords = _convert_poi(poi_info, dm.get_grid_crs(data), data[y_dim].dtype)
1044
+ poi_dim_coord = poi_info[POI_COORD_COL].to_numpy(dtype=str)
1045
+ pois_coords = (
1046
+ xr.DataArray(pois_coords[0], coords=[(POI_DIM, poi_dim_coord)]),
1047
+ xr.DataArray(pois_coords[1], coords=[(POI_DIM, poi_dim_coord)]),
1048
+ )
1049
+
1050
+ if method == "nearest":
1051
+ dx, dy = np.diff(data[x_dim].values), np.diff(data[y_dim].values)
1052
+ cst_spacing = isclose(dx.min(), dx.max(), rel_tol=1e-4) and isclose(dy.min(), dy.max(), rel_tol=1e-4)
1053
+ dtype_bool = not all(data[var].dtype in (np.float32, np.float64) for var in data)
1054
+ chunk_bool = len(data.chunks) > 0
1055
+ if cst_spacing and (chunk_bool or dtype_bool):
1056
+ grid_spacing = float(abs(dx[0])), float(abs(dy[0]))
1057
+ poi_data = _interp_nearest_regular(data, spatial_dims, grid_spacing, pois_coords)
1058
+ else:
1059
+ poi_data = _interp_cast_dtype("nearest")
1060
+ else:
1061
+ poi_data = _interp_cast_dtype(method)
1062
+
1063
+ poi_data.attrs = _add_attrs(data.attrs)
1064
+
1065
+ return poi_data.transpose(POI_DIM, ...)
1066
+
1067
+
1068
+ def _y_increasing(ds: xr.Dataset, y_dim: str) -> xr.Dataset:
1069
+ """
1070
+ Reverse along the y dimension if coordinates are decreasing instead of increasing.
1071
+ This allows to use ``assume_sorted=True`` in :meth:`xarray.Dataset.interp` which is significantly faster.
1072
+ """
1073
+
1074
+ if dm.is_decreasing1d(ds[y_dim]):
1075
+ ds = ds.isel({y_dim: slice(None, None, -1)})
1076
+
1077
+ return ds
1078
+
1079
+
1080
+ def _drop_coords_reg(data: xr.Dataset, spatial_dims: tuple[str, str]) -> xr.Dataset:
1081
+ """
1082
+ Drop non-dimensional coordinate variables that have spatial dimensions.
1083
+ """
1084
+
1085
+ notdim_coords = set(data.coords) - set(data.dims)
1086
+ spd = set(spatial_dims)
1087
+ to_drop = [coord for coord in notdim_coords if spd.issubset(data[coord].dims)]
1088
+
1089
+ return data.drop_vars(to_drop)
1090
+
1091
+
1092
+ def _interp_nearest_regular(
1093
+ data: xr.Dataset,
1094
+ spatial_dims: tuple[str, str],
1095
+ grid_spacing: tuple[float, float],
1096
+ pois_coords: tuple[xr.DataArray, xr.DataArray],
1097
+ ) -> xr.Dataset:
1098
+ """
1099
+ Perform fast nearest neighbor interpolation from a regular grid with
1100
+ constant spacing along each dimension to POIs. Preserve dtype
1101
+ so the memory footprint is limited compared to xarray.Dataset.interp which
1102
+ cast to numpy.float64.
1103
+
1104
+ Note
1105
+ ----
1106
+ If the spacing along a dimension is not constant, results may be wrong.
1107
+ """
1108
+
1109
+ y_dim, x_dim = spatial_dims
1110
+ data = data.sel(
1111
+ {
1112
+ x_dim: slice(
1113
+ pois_coords[0].min() - grid_spacing[0],
1114
+ pois_coords[0].max() + grid_spacing[0],
1115
+ ),
1116
+ y_dim: slice(
1117
+ pois_coords[1].min() - grid_spacing[1],
1118
+ pois_coords[1].max() + grid_spacing[1],
1119
+ ),
1120
+ }
1121
+ )
1122
+ data = data.chunk() # cast to dask
1123
+ grid_coords = (data[x_dim].values, data[y_dim].values)
1124
+ hkey = (
1125
+ tuple(pois_coords[0].poi.values.tolist()),
1126
+ pois_coords[0].values.tobytes(),
1127
+ pois_coords[1].values.tobytes(),
1128
+ grid_coords[0].tobytes(),
1129
+ grid_coords[1].tobytes(),
1130
+ )
1131
+ # pylint: disable-next=unbalanced-tuple-unpacking # returns one index array per spatial dimension
1132
+ indx, indy = nn_lookup_regular(grid_coords, grid_spacing, pois_coords, hkey=hkey) # type: ignore[call-arg]
1133
+ return data.isel({x_dim: indx, y_dim: indy}).drop_vars(spatial_dims, errors="ignore")
1134
+
1135
+
1136
+ @memoize(CONFIG.cache_size.nn_lookup)
1137
+ def nn_lookup(grid_coords: np.ndarray, points: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
1138
+ """
1139
+ Find the k nearest grid points and corresponding distances for a set of
1140
+ POIs.
1141
+
1142
+ Parameters
1143
+ ----------
1144
+ grid_coords : numpy.ndarray, shape (n, m)
1145
+ array of flattened grid coordinates. N is the number of model grid
1146
+ points, m the number of coordinates dimensions (i.e. 2 or 3)
1147
+ points : numpy.ndarray, shape (p, m)
1148
+ array of POIs coordinates. P is the number of POIs, m the number of
1149
+ coordinates dimensions (i.e. 2 or 3)
1150
+ k : positive int
1151
+ number of considered nearest neighbors
1152
+
1153
+ Returns
1154
+ -------
1155
+ dd : numpy.ndarray, shape (p, k)
1156
+ array of nearest neighbor distances.
1157
+ ii : numpy.ndarray, shape (p, k)
1158
+ array of nearest neighbor grid indexes
1159
+
1160
+ See also
1161
+ --------
1162
+ :func:`nn_lookup_cached`
1163
+ """
1164
+
1165
+ tree = KDTree(grid_coords)
1166
+ dd, ii = tree.query(points, k, workers=-1)
1167
+
1168
+ return dd.astype(points.dtype, copy=False), ii
1169
+
1170
+
1171
+ def nn_lookup_cached(grid_coords: np.ndarray, points: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
1172
+ """
1173
+ Wrapper to :func:`nn_lookup` function to use it with a cache
1174
+ """
1175
+
1176
+ hkey = (grid_coords.tobytes(), points.tobytes(), k)
1177
+ return nn_lookup(grid_coords, points, k, hkey=hkey) # type: ignore[call-arg]
1178
+
1179
+
1180
+ @memoize(CONFIG.cache_size.nn_lookup_regular)
1181
+ def nn_lookup_regular(
1182
+ grid_coords: tuple[np.ndarray, np.ndarray],
1183
+ grid_spacing: tuple[float, float],
1184
+ points: tuple[np.ndarray | xr.DataArray, np.ndarray | xr.DataArray],
1185
+ ) -> list[np.ndarray | xr.DataArray]:
1186
+ """
1187
+ Find the nearest grid point for a set of POIs when the source
1188
+ grid is regular and the spacing constant along each dimension.
1189
+
1190
+ Parameters
1191
+ ----------
1192
+ grid_coords : tuple of m 1-D arrays
1193
+ contains flattened grid coordinates arrays for all m dimensions
1194
+ grid_spacing : tuple of m floats
1195
+ contains the grid spacing for all m dimensions
1196
+ points : tuple of m 1-D arrays (or xarray.DataArray with dim poi)
1197
+ contains POIs coordinates for all m dimensions. It assumes the same
1198
+ order of dimensions as in ``grid_coords``
1199
+
1200
+ Returns
1201
+ -------
1202
+ index : list of m 1-D arrays (or xarray DataArray with dimension `poi`)
1203
+ Used to index the dataset in each spatial dimension. The output indexes
1204
+ have length equals to the number of points
1205
+ """
1206
+
1207
+ if len(points) != len(grid_coords):
1208
+ raise ValueError("The number of grid and points dimensions does not match!")
1209
+ grid_index = []
1210
+ for n, dim_coords in enumerate(grid_coords):
1211
+ dim_spacing = grid_spacing[n]
1212
+ dim_origin = np.min(dim_coords) - dim_spacing / 2
1213
+ dim_index = (points[n] - dim_origin) // dim_spacing
1214
+ dim_index = np.clip(dim_index, 0, dim_coords.size - 1).astype(int)
1215
+ grid_index.append(dim_index)
1216
+ return grid_index
1217
+
1218
+
1219
+ def _compute_dist_weight(dist: np.ndarray, power: float, dtype: np.dtype) -> np.ndarray:
1220
+ """
1221
+ Compute weights based on the inverse distance. The weights scales to 1.
1222
+ """
1223
+
1224
+ increm = np.mean(dist, axis=1, keepdims=True) / 50 # avoid division by 0
1225
+ w = 1 / np.power(dist + increm, power)
1226
+ w = w.astype(dtype, copy=False)
1227
+ return w / np.sum(w, axis=1, keepdims=True)
1228
+
1229
+
1230
+ def _add_attrs(data_attrs: dict) -> dict:
1231
+ """
1232
+ Add attrs from source grid
1233
+ """
1234
+
1235
+ return {key: data_attrs[key] for key in SRC_ATTRS_KEEP if key in data_attrs}
1236
+
1237
+
1238
+ def _flatten_concat(
1239
+ data: xr.DataArray,
1240
+ spatial_dims: tuple[str, ...],
1241
+ poi_info: gpd.GeoDataFrame,
1242
+ method: str,
1243
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
1244
+ """
1245
+ Flattens the grid and concatenate columnwise `x` and `y` coordinates into a
1246
+ single 2D array. Also, reprojects the POIs coordinates if they are not in
1247
+ the same CRS than the model coordinates
1248
+ """
1249
+
1250
+ data, grid_y, grid_x, grid_crs = _flatten_coords_crs(data, spatial_dims)
1251
+
1252
+ if method in HEIGHT_METHODS:
1253
+ # include elevation
1254
+ surf_alt = dm.get_surf_alt(data)
1255
+ coords_flat = np.c_[
1256
+ data[grid_x].values,
1257
+ data[grid_y].values,
1258
+ data[surf_alt].values,
1259
+ ]
1260
+
1261
+ if grid_crs not in SWISS_CRS:
1262
+ # project to CH1903+/LV95 to have units in [m]
1263
+ new_crs = CRS.from_string(SWISS_CRS[0])
1264
+ coords_flat[:, 0], coords_flat[:, 1] = _reproject_grid_cached(
1265
+ coords_flat[:, 0], coords_flat[:, 1], grid_crs, new_crs
1266
+ )
1267
+ grid_crs = new_crs
1268
+ else:
1269
+ coords_flat = np.c_[data[grid_x].values, data[grid_y].values]
1270
+
1271
+ # Reproject POIs coordinates in the grid CRS
1272
+ pois_coords = _convert_poi(poi_info, grid_crs, coords_flat.dtype)
1273
+ if method in HEIGHT_METHODS:
1274
+ pois_coords_cat = np.c_[
1275
+ pois_coords[0],
1276
+ pois_coords[1],
1277
+ dm.get_poi_alt(poi_info),
1278
+ ]
1279
+ else:
1280
+ pois_coords_cat = np.c_[pois_coords[0], pois_coords[1]]
1281
+
1282
+ if method == "snowlim":
1283
+ return coords_flat, data.values, pois_coords_cat
1284
+
1285
+ # check for all NaN grid points
1286
+ ind = ~data.isnull().all(data.dims[1:]).values
1287
+
1288
+ return coords_flat[ind], data.values[ind], pois_coords_cat
1289
+
1290
+
1291
+ def _convert_poi(poi_info: gpd.GeoDataFrame, grid_crs: CRS, dtype: np.dtype) -> tuple[np.ndarray, np.ndarray]:
1292
+ """
1293
+ Cached reprojection of POIs coordinates in the grid CRS. Returns a tuple of xy
1294
+ coordinates with the specified dtype.
1295
+ """
1296
+
1297
+ pois_coords = poi_info.get_coordinates().values
1298
+ pois_coords = pois_coords.astype(dtype, copy=False)
1299
+
1300
+ poi_crs = poi_info.crs.to_2d()
1301
+ if grid_crs == poi_crs:
1302
+ return pois_coords[:, 0], pois_coords[:, 1]
1303
+
1304
+ hkey = (pois_coords[:, 0].tobytes(), pois_coords[:, 1].tobytes(), grid_crs)
1305
+ return _reproject_poi(pois_coords, poi_crs, grid_crs, hkey=hkey) # type: ignore[call-arg]
1306
+
1307
+
1308
+ @memoize(CONFIG.cache_size.reproject_poi)
1309
+ def _reproject_poi(coords: np.ndarray, src_crs: CRS, dst_crs: CRS) -> tuple[np.ndarray, np.ndarray]:
1310
+ """
1311
+ Reprojection of POIs coordinates that can be cached
1312
+ """
1313
+
1314
+ x_coords, y_coords = reproject(coords[:, 0], coords[:, 1], src_crs, dst_crs)
1315
+ dtype = coords.dtype
1316
+ x_coords = x_coords.astype(dtype, copy=False)
1317
+ y_coords = y_coords.astype(dtype, copy=False)
1318
+ return x_coords, y_coords
1319
+
1320
+
1321
+ @memoize(CONFIG.cache_size.reproject_grid)
1322
+ def _reproject_grid(
1323
+ x_coords: np.ndarray, y_coords: np.ndarray, src_crs: CRS, dst_crs: CRS
1324
+ ) -> tuple[np.ndarray, np.ndarray]:
1325
+ """
1326
+ Reprojection of grid coordinates that can be cached
1327
+ """
1328
+
1329
+ return reproject(x_coords, y_coords, src_crs, dst_crs)
1330
+
1331
+
1332
+ def _reproject_grid_cached(
1333
+ x_coords: np.ndarray, y_coords: np.ndarray, src_crs: CRS, dst_crs: CRS
1334
+ ) -> tuple[np.ndarray, np.ndarray]:
1335
+ """
1336
+ Cached reprojection of grid coordinates
1337
+ """
1338
+
1339
+ if dst_crs != src_crs:
1340
+ hkey = (x_coords.tobytes(), y_coords.tobytes(), x_coords.shape, dst_crs)
1341
+ x_coords, y_coords = _reproject_grid( # type: ignore[call-arg]
1342
+ x_coords, y_coords, src_crs, dst_crs, hkey=hkey
1343
+ )
1344
+
1345
+ return x_coords.astype(np.float64, copy=False), y_coords.astype(np.float64, copy=False)
1346
+
1347
+
1348
+ @dataset_accessor
1349
+ def interp2grid(
1350
+ data: xr.Dataset,
1351
+ dst_grid: xr.Dataset,
1352
+ method: InterpMethods,
1353
+ try_reg: bool = True,
1354
+ **kwargs: Any,
1355
+ ) -> xr.Dataset:
1356
+ """
1357
+ Interpolates data from a source grid to a destination grid, supporting regular and unstructured grids.
1358
+
1359
+ This function provides a unified interface for grid-to-grid interpolation, supporting a variety of interpolation
1360
+ methods and data models. The output dataset contains the data of the source grid mapped onto the destination grid
1361
+ coordinates.
1362
+
1363
+ Parameters
1364
+ ----------
1365
+ data : xarray.Dataset
1366
+ The source grid dataset. Supported data models include:
1367
+
1368
+ - CF-convention grids (see https://cfconventions.org/)
1369
+ - internal MCH models
1370
+ - Unstructured grids
1371
+
1372
+ Typical spatial dimension and coordinate names are supported, and can be extended via environment variables:
1373
+
1374
+ - Y coordinate: export GRID_PROCESS__data_model__grid__y (e.g., '["y", "latitude", "lat", "dummy_y"]')
1375
+ - X coordinate: export GRID_PROCESS__data_model__grid__x (e.g., '["x", "longitude", "lon", "dummy_x"]')
1376
+ - Unstructured dimension: export GRID_PROCESS__data_model__grid__unstructured_dim (e.g., '["dummy_cell"]')
1377
+ - When the spatial coordinates variables have a CF compliant `standard_name` attribute, arbitrary names can be
1378
+ used directly.
1379
+
1380
+ Spatial CRS must be referenced via one of:
1381
+
1382
+ - Global attribute 'crs' (e.g., "EPSG:2056" or anything accepted by :meth:`pyproj.CRS.from_user_input`)
1383
+ - Global attribute 'grid_mapping' (internal grid database convention)
1384
+ - CF-convention grid mapping variable
1385
+ - If coordinates are named 'latitude'/'longitude' or 'lat'/'lon', WGS84 (EPSG:4326) is assumed.
1386
+
1387
+ For methods requiring altitude, ensure a surface altitude coordinate is present; extend supported names via
1388
+ GRID_PROCESS__data_model__grid__surf_alt.
1389
+ dst_grid : xarray.Dataset
1390
+ The destination grid to interpolate to. Same supported data models as the source grid, including conventions for
1391
+ dimensions and coordinates as well as spatial CRS.
1392
+ For methods requiring altitude, ensure a surface altitude coordinate is also present.
1393
+ The destination dataset may contain only spatial coordinates or also data variables which will be ignored.
1394
+ method : str
1395
+ Interpolation method. Supported options include:
1396
+
1397
+ {'nearest', 'linear', 'cubic', 'idw', 'temp', 'temp_fixgrad', 'nearest_alt', 'snowlim',
1398
+ 'nearest_3D', 'linear_3D', 'idw_3D', 'upscaled', 'upscaled_3D'}.
1399
+
1400
+ 3D methods include altitude in neighbor selection. For unstructured grids, prefer 'idw' with k=4 and
1401
+ weight_power=2 over 'linear' for performence reasons.
1402
+ try_reg : bool, default True
1403
+ If True, attempts regular grid interpolation for structured source grids with 1D coordinates when ``method`` is
1404
+ 'nearest' or 'linear'. Regular grid interpolation is typically more efficient for large grids, but may be less
1405
+ so for small grids with many trailing dimensions.
1406
+ **kwargs : Any
1407
+ Additional arguments passed to the specific interpolation sub-function, named `interp_<method>`.
1408
+
1409
+ Returns
1410
+ -------
1411
+ xr.Dataset
1412
+ Dataset with values from `data` interpolated onto the coordinates of `dst_grid`.
1413
+
1414
+ See also
1415
+ --------
1416
+ :func:`interp_temp`, :func:`interp_temp_fixgrad`, :func:`interp_nearest_alt`,
1417
+ :func:`interp_idw`, :func:`interp_upscaled`, :func:`interp_nearest`,
1418
+ :func:`interp_snowlim`, :func:`scipy.interpolate.griddata`, :meth:`xarray.Dataset.interp`
1419
+
1420
+ Notes
1421
+ -----
1422
+ - Not all interpolation methods scale efficiently with large input datasets or and/or destination grids.
1423
+ - Recommended methods: 'nearest', 'linear' (if the source grid is regular with 1D coordinates), 'idw' (with k=4),
1424
+ 'temp_fixgrad', 'nearest_3D', 'nearest_alt'.
1425
+ - Environment variables allow flexible support for custom grid data models.
1426
+ """
1427
+
1428
+ def _array2ds(data: xr.DataArray) -> xr.Dataset:
1429
+ """
1430
+ Converts an array with interpolated values into a dataset with correct
1431
+ dimensions, coordinates and attributes
1432
+ """
1433
+
1434
+ da = xr.DataArray(
1435
+ interp_data.reshape(
1436
+ (
1437
+ dst_grid.sizes[dst_spatial_dims[0]],
1438
+ dst_grid.sizes[dst_spatial_dims[1]],
1439
+ )
1440
+ + data.shape[len(src_spatial_dims) :]
1441
+ ),
1442
+ dims=dst_spatial_dims + data.dims[len(src_spatial_dims) :],
1443
+ coords={
1444
+ **data.isel({dim: 0 for dim in src_spatial_dims}, drop=True).coords,
1445
+ **dst_grid.coords,
1446
+ },
1447
+ attrs=_combine_attrs(data.attrs, dst_grid.attrs),
1448
+ )
1449
+ return da_to_ds(da, vars_attrs)
1450
+
1451
+ src_spatial_dims = dm.get_spatial_dims(data)
1452
+ dst_spatial_dims = dm.get_spatial_dims(dst_grid)
1453
+ data = dm.set_grid_mapping(data)
1454
+ dst_grid = dm.set_grid_mapping(dst_grid)
1455
+ dst_grid = dst_grid.drop_dims(set(dst_grid.dims) - set(dst_spatial_dims))
1456
+
1457
+ cond_reg = (
1458
+ method in ("nearest", "linear") and len(src_spatial_dims) > 1 and set(src_spatial_dims).issubset(data.coords)
1459
+ )
1460
+ if try_reg and cond_reg:
1461
+ return _interp_grid_regular(
1462
+ data,
1463
+ src_spatial_dims, # type: ignore[arg-type]
1464
+ dst_grid,
1465
+ dst_spatial_dims, # type: ignore[arg-type]
1466
+ method, # type: ignore[arg-type]
1467
+ )
1468
+
1469
+ dst_grid = dst_grid.transpose(*dst_spatial_dims)
1470
+
1471
+ data = slice_coarsen(data, dst_grid, method, src_spatial_dims, dst_spatial_dims, **kwargs)
1472
+ data, vars_attrs = ds_to_da(data) # type: ignore[assignment]
1473
+ data = data.transpose(*src_spatial_dims, ...)
1474
+ coords_flat, data_flat, dst_coords_flat = _flatten_concat_grids(
1475
+ data, # type: ignore[arg-type]
1476
+ src_spatial_dims,
1477
+ dst_grid,
1478
+ dst_spatial_dims, # type: ignore[arg-type]
1479
+ method,
1480
+ )
1481
+
1482
+ if method in INTERP_FUNC:
1483
+ interp_data = INTERP_FUNC[method]( # type: ignore[operator]
1484
+ coords_flat, data_flat, dst_coords_flat, **kwargs
1485
+ )
1486
+
1487
+ elif method in ("linear", "linear_3D", "cubic"):
1488
+ interp_data = griddata(coords_flat, data_flat, dst_coords_flat, method.rsplit("_")[0])
1489
+ interp_data = interp_data.astype(data_flat.dtype, copy=False)
1490
+ else:
1491
+ raise ValueError(f"Unknown interpolation method {method!r}")
1492
+
1493
+ return _array2ds(data) # type: ignore[arg-type]
1494
+
1495
+
1496
+ def _interp_grid_regular(
1497
+ data: xr.Dataset,
1498
+ src_spatial_dims: tuple[str, str],
1499
+ dst_grid: xr.Dataset,
1500
+ dst_spatial_dims: tuple[str, str],
1501
+ method: Literal["nearest", "linear"] = "linear",
1502
+ ) -> xr.Dataset:
1503
+ """
1504
+ Wrapper for interpolation from a regular grid to another one.
1505
+ Handle reprojection of the destination grid to the native crs of the source grid.
1506
+ The destination grid does not need to be regular. Used in :func:`interp2grid`.
1507
+
1508
+ Note
1509
+ ----
1510
+ For very large grids in 'x' and 'y', the interpolation is much more efficient
1511
+ when the input and destination grids are on the same native CRS.
1512
+ """
1513
+
1514
+ def _cond_dst_1d() -> bool:
1515
+ """
1516
+ Whether the destination coordinates should be kept as 1D (when the source and destinations grids are in the
1517
+ same CRS). Depending on the size of the grids, xarray "nearest" interpolation is much faster when destination
1518
+ coordinates are 2D.
1519
+ """
1520
+
1521
+ data_trail_size = prod([size for dim, size in data.sizes.items() if dim not in src_spatial_dims])
1522
+ dst_sizes = dst_grid.sizes
1523
+ dim_ratio = (dst_sizes[dst_spatial_dims[1]] * dst_sizes[dst_spatial_dims[0]]) / data_trail_size
1524
+ bool_dst_1d = dim_ratio > 2000 or method == "linear"
1525
+ return bool_dst_1d and set(dst_spatial_dims).issubset(dst_grid.coords)
1526
+
1527
+ def _reproject_dst_2d() -> tuple[xr.DataArray, xr.DataArray]:
1528
+ """
1529
+ Reproject destination 2D coordinates in the source grid CRS
1530
+ """
1531
+
1532
+ x_coords_val = dst_grid[dst_coords_crs[1]].values
1533
+ y_coords_val = dst_grid[dst_coords_crs[0]].values
1534
+
1535
+ if x_coords_val.ndim == 1:
1536
+ x_coords_val, y_coords_val = np.meshgrid(x_coords_val, y_coords_val)
1537
+
1538
+ # Reproject destination 2D coordinates in the grid CRS
1539
+ x_coords_val, y_coords_val = _reproject_grid_cached(x_coords_val, y_coords_val, dst_coords_crs[2], src_crs)
1540
+
1541
+ return xr.DataArray(x_coords_val, dims=dst_spatial_dims), xr.DataArray(y_coords_val, dims=dst_spatial_dims)
1542
+
1543
+ def _coarsen_interp() -> xr.Dataset:
1544
+ """
1545
+ Coarsen input dataset if needed (downsampling) and perform regular interpolation
1546
+ """
1547
+
1548
+ src_coords_crs = src_spatial_dims + (src_crs,)
1549
+ res_ratio_x, res_ratio_y = _get_coarsen_ratio(
1550
+ data,
1551
+ src_spatial_dims,
1552
+ src_coords_crs,
1553
+ dst_grid,
1554
+ dst_spatial_dims,
1555
+ dst_coords_crs,
1556
+ )
1557
+ interp_data = data
1558
+ if (res_ratio_x > 1) or (res_ratio_y > 1):
1559
+ coarsen_specs = {
1560
+ dim: ratio
1561
+ for dim, ratio in (
1562
+ (src_spatial_dims[0], res_ratio_y),
1563
+ (src_spatial_dims[1], res_ratio_x),
1564
+ )
1565
+ if ratio > 1
1566
+ }
1567
+ # reduce(np.mean) propagates NaNs like mean(skipna=False), which fails with xarray 2026.9.0
1568
+ interp_data = interp_data.coarsen(coarsen_specs, boundary="trim").reduce(np.mean)
1569
+
1570
+ interp_data = _y_increasing(interp_data, src_spatial_dims[0])
1571
+
1572
+ interp_data = interp_data.interp(
1573
+ {src_spatial_dims[1]: x_coords, src_spatial_dims[0]: y_coords},
1574
+ method=method,
1575
+ assume_sorted=True,
1576
+ ).drop_vars(src_spatial_dims)
1577
+ cast_dtype(interp_data, data)
1578
+ interp_data.attrs = _combine_attrs(data.attrs, dst_grid.attrs)
1579
+ interp_data = interp_data.assign_coords(dst_grid.coords)
1580
+ return interp_data
1581
+
1582
+ data = _drop_coords_reg(data, src_spatial_dims)
1583
+ src_crs = dm.get_grid_crs(data)
1584
+
1585
+ if _cond_dst_1d():
1586
+ dst_crs = dm.get_grid_crs(dst_grid)
1587
+ if src_crs == dst_crs:
1588
+ dst_coords_crs = dst_spatial_dims + (dst_crs,)
1589
+ x_coords, y_coords = (
1590
+ dst_grid[dst_spatial_dims[1]].drop_vars(dst_spatial_dims[1]),
1591
+ dst_grid[dst_spatial_dims[0]].drop_vars(dst_spatial_dims[0]),
1592
+ )
1593
+ return _coarsen_interp()
1594
+
1595
+ dst_grid = dst_grid.transpose(*dst_spatial_dims)
1596
+ # chose to use lat/lon if present in dst_grid otherwise native coordinates
1597
+ dst_coords_crs = dm.get_coords_name_crs(dst_grid, dst_spatial_dims)
1598
+
1599
+ x_coords, y_coords = _reproject_dst_2d()
1600
+
1601
+ return _coarsen_interp()
1602
+
1603
+
1604
+ def _flatten_concat_grids( # pylint: disable=too-many-locals
1605
+ data: xr.DataArray,
1606
+ spatial_dims: tuple[str, ...],
1607
+ dst_grid: xr.Dataset,
1608
+ dst_spatial_dims: tuple[str, str],
1609
+ method: str,
1610
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
1611
+ """
1612
+ Flattens the grids and concatenate columnwise `x` and `y` coordinates into a
1613
+ single 2D array. Also, reprojects the destination grid coordinates if they are not
1614
+ in the same CRS than the input model coordinates.
1615
+ """
1616
+
1617
+ def _choose_crs_meters() -> CRS:
1618
+ """
1619
+ Choose the target CRS with units in [m] for reprojection.
1620
+ """
1621
+
1622
+ if reproj_dst:
1623
+ if grid_crs in SWISS_CRS:
1624
+ return grid_crs
1625
+ return CRS.from_string(SWISS_CRS[1]) if dst_grid_crs == SWISS_CRS[1] else CRS.from_string(SWISS_CRS[0])
1626
+
1627
+ if dst_grid_crs in SWISS_CRS:
1628
+ return dst_grid_crs
1629
+ return CRS.from_string(SWISS_CRS[1]) if grid_crs == SWISS_CRS[1] else CRS.from_string(SWISS_CRS[0])
1630
+
1631
+ data, grid_y, grid_x, grid_crs = _flatten_coords_crs(data, spatial_dims)
1632
+ flat_dim = spatial_dims[0] if len(spatial_dims) == 1 else "cell"
1633
+ dst_grid, dst_grid_y, dst_grid_x, dst_grid_crs = _flatten_coords_crs(dst_grid, dst_spatial_dims, flat_dim)
1634
+
1635
+ reproj_dst = dst_grid.sizes[flat_dim] <= data.sizes[flat_dim]
1636
+
1637
+ if method in HEIGHT_METHODS:
1638
+ surf_alt = dm.get_surf_alt(data)
1639
+ dst_surf_alt = dm.get_surf_alt(dst_grid)
1640
+ # include elevation
1641
+ coords_flat = np.c_[
1642
+ data[grid_x].values,
1643
+ data[grid_y].values,
1644
+ data[surf_alt].values,
1645
+ ]
1646
+ dst_coords_flat = np.c_[
1647
+ dst_grid[dst_grid_x].values,
1648
+ dst_grid[dst_grid_y].values,
1649
+ dst_grid[dst_surf_alt].values,
1650
+ ]
1651
+ target_crs = _choose_crs_meters()
1652
+ coords_flat[:, 0], coords_flat[:, 1] = _reproject_grid_cached(
1653
+ coords_flat[:, 0], coords_flat[:, 1], grid_crs, target_crs
1654
+ )
1655
+ dst_coords_flat[:, 0], dst_coords_flat[:, 1] = _reproject_grid_cached(
1656
+ dst_coords_flat[:, 0], dst_coords_flat[:, 1], dst_grid_crs, target_crs
1657
+ )
1658
+ return coords_flat, data.values, dst_coords_flat
1659
+
1660
+ # Non-height methods
1661
+ coords_flat = np.c_[data[grid_x].values, data[grid_y].values]
1662
+ dst_coords_flat = np.c_[dst_grid[dst_grid_x].values, dst_grid[dst_grid_y].values]
1663
+ if reproj_dst: # reproject destination coordinates in the grid CRS
1664
+ dst_coords_flat[:, 0], dst_coords_flat[:, 1] = _reproject_grid_cached(
1665
+ dst_coords_flat[:, 0], dst_coords_flat[:, 1], dst_grid_crs, grid_crs
1666
+ )
1667
+ return coords_flat, data.values, dst_coords_flat
1668
+
1669
+ # reproject grid coordinates in the destination CRS
1670
+ coords_flat[:, 0], coords_flat[:, 1] = _reproject_grid_cached(
1671
+ coords_flat[:, 0], coords_flat[:, 1], grid_crs, dst_grid_crs
1672
+ )
1673
+ return coords_flat, data.values, dst_coords_flat
1674
+
1675
+
1676
+ @overload
1677
+ def _flatten_coords_crs(
1678
+ grid: xr.DataArray, spatial_dims: tuple[str, ...], flat_dim: str = "cell"
1679
+ ) -> tuple[xr.DataArray, str, str, CRS]: ...
1680
+ @overload
1681
+ def _flatten_coords_crs(
1682
+ grid: xr.Dataset, spatial_dims: tuple[str, ...], flat_dim: str = "cell"
1683
+ ) -> tuple[xr.Dataset, str, str, CRS]: ...
1684
+
1685
+
1686
+ def _flatten_coords_crs(
1687
+ grid: xr.DataArray | xr.Dataset,
1688
+ spatial_dims: tuple[str, ...],
1689
+ flat_dim: str = "cell",
1690
+ ) -> tuple[xr.DataArray | xr.Dataset, str, str, CRS]:
1691
+ """
1692
+ Flatten input dataset if not already and extract spatial coordinates and
1693
+ corresponding crs to be used.
1694
+ """
1695
+
1696
+ if len(spatial_dims) > 1:
1697
+ # flatten grid
1698
+ grid = grid.stack({flat_dim: spatial_dims}).transpose(flat_dim, ...)
1699
+ spatial_dims = (flat_dim,)
1700
+
1701
+ grid_y, grid_x, grid_crs = dm.get_coords_name_crs(grid, spatial_dims)
1702
+
1703
+ return grid, grid_y, grid_x, grid_crs
1704
+
1705
+
1706
+ def slice_boundingbox(
1707
+ data: xr.Dataset,
1708
+ poi_info: gpd.GeoDataFrame,
1709
+ method: str,
1710
+ spatial_dims: tuple[str, ...] | None = None,
1711
+ **kwargs: Any,
1712
+ ) -> xr.Dataset:
1713
+ """
1714
+ Slice input dataset spatial dimensions to the bounding box of the target POIs
1715
+ plus a buffer depending on the interpolation method and corresponding kwargs.
1716
+ """
1717
+
1718
+ if spatial_dims is None:
1719
+ spatial_dims = dm.get_spatial_dims(data)
1720
+
1721
+ grid_y, grid_x, grid_crs = dm.get_coords_name_crs(data, spatial_dims)
1722
+ pois_coords = _convert_poi(poi_info, grid_crs, data[grid_y].dtype)
1723
+ yx_bnds = (
1724
+ (pois_coords[1].min().item(), pois_coords[1].max().item()),
1725
+ (pois_coords[0].min().item(), pois_coords[0].max().item()),
1726
+ )
1727
+
1728
+ buffer = _get_buffer(method, **kwargs)
1729
+ sub_domain, _ = _get_sub_domain(data, spatial_dims, (grid_y, grid_x), yx_bnds, buffer)
1730
+ return data.isel(sub_domain)
1731
+
1732
+
1733
+ def slice_coarsen( # pylint: disable=too-many-locals
1734
+ data: xr.Dataset,
1735
+ dst_grid: xr.Dataset,
1736
+ method: str,
1737
+ spatial_dims: tuple[str, ...] | None = None,
1738
+ dst_spatial_dims: tuple[str, ...] | None = None,
1739
+ **kwargs: Any,
1740
+ ) -> xr.Dataset:
1741
+ """
1742
+ Slice input dataset to the extend of the destination grid and coarsen it if the
1743
+ destination grid has lower resolution (downsampling).
1744
+ """
1745
+
1746
+ def extra_buffer(ratio: int) -> int:
1747
+ """
1748
+ Extra buffer to add when input data is coarsened.
1749
+ """
1750
+
1751
+ if buffer == 1:
1752
+ return ceil(ratio / 2)
1753
+
1754
+ return (ratio - 1) * (buffer - 1)
1755
+
1756
+ if spatial_dims is None:
1757
+ spatial_dims = dm.get_spatial_dims(data)
1758
+ if dst_spatial_dims is None:
1759
+ dst_spatial_dims = dm.get_spatial_dims(dst_grid)
1760
+
1761
+ coords_crs = dm.get_coords_name_crs(data, spatial_dims)
1762
+ dst_coords_crs = dm.get_coords_name_crs(dst_grid, dst_spatial_dims)
1763
+ grid_crs, dst_grid_crs = coords_crs[2], dst_coords_crs[2]
1764
+
1765
+ x_coords, y_coords = (
1766
+ dst_grid[dst_coords_crs[1]].values,
1767
+ dst_grid[dst_coords_crs[0]].values,
1768
+ )
1769
+
1770
+ if grid_crs != dst_grid_crs:
1771
+ if x_coords.ndim == 1:
1772
+ x_coords, y_coords = np.meshgrid(x_coords, y_coords)
1773
+
1774
+ # only keep the outer rectangle
1775
+ x_coords = np.append(x_coords[[0, -1], :].ravel(), x_coords[1:-1, [0, -1]].ravel())
1776
+ y_coords = np.append(y_coords[[0, -1], :].ravel(), y_coords[1:-1, [0, -1]].ravel())
1777
+
1778
+ # Reproject destination coordinates in the grid CRS
1779
+ x_coords, y_coords = reproject(x_coords, y_coords, dst_grid_crs, grid_crs)
1780
+
1781
+ yx_bnds = ((y_coords.min(), y_coords.max()), (x_coords.min(), x_coords.max()))
1782
+
1783
+ buffer = _get_buffer(method, **kwargs)
1784
+ sub_domain, res = _get_sub_domain(data, spatial_dims, coords_crs[:2], yx_bnds, buffer)
1785
+ data_sub = data.isel(sub_domain)
1786
+
1787
+ # coarsen input dataset if needed (downsampling)
1788
+ res_ratio_x, res_ratio_y = _get_coarsen_ratio(
1789
+ data_sub,
1790
+ spatial_dims,
1791
+ coords_crs,
1792
+ dst_grid,
1793
+ dst_spatial_dims, # type: ignore[arg-type]
1794
+ dst_coords_crs,
1795
+ res,
1796
+ )
1797
+ if (res_ratio_x > 1) or (res_ratio_y > 1):
1798
+ if len(spatial_dims) == 1:
1799
+ logger.warning(
1800
+ "Unstructured grid cannot be coarsened. Interpolation to a lower"
1801
+ " resolution grid may create aliasing artifacts."
1802
+ )
1803
+ return data_sub
1804
+
1805
+ coarsen_specs = {
1806
+ dim: ratio
1807
+ for dim, ratio in (
1808
+ (spatial_dims[1], res_ratio_x),
1809
+ (spatial_dims[0], res_ratio_y),
1810
+ )
1811
+ if ratio > 1
1812
+ }
1813
+ sub_domain.update(
1814
+ {
1815
+ dim: slice( # type: ignore[misc]
1816
+ max(sub_domain[dim].start - extra_buffer(ratio), 0), # type: ignore[union-attr]
1817
+ sub_domain[dim].stop + extra_buffer(ratio), # type: ignore[union-attr]
1818
+ )
1819
+ for dim, ratio in coarsen_specs.items()
1820
+ }
1821
+ )
1822
+ # reduce(np.mean) propagates NaNs like mean(skipna=False), which fails with xarray 2026.9.0
1823
+ return data.isel(sub_domain).coarsen(coarsen_specs, boundary="trim").reduce(np.mean)
1824
+
1825
+ return data_sub
1826
+
1827
+
1828
+ def _get_buffer(method: str, **kwargs: Any) -> int:
1829
+ """
1830
+ Buffer specific to the interpolation method to add when slicing input dataset
1831
+ to target grid or POIs
1832
+ """
1833
+
1834
+ buffer = 1
1835
+ if (method == "cubic") or ("3D" in method):
1836
+ buffer += 1
1837
+ if "k" in kwargs:
1838
+ buffer += ceil(np.sqrt(kwargs["k"] / np.pi))
1839
+ if "r_search" in kwargs:
1840
+ buffer += ceil(kwargs["r_search"])
1841
+ return buffer
1842
+
1843
+
1844
+ def _get_sub_domain( # pylint: disable=too-many-locals
1845
+ data: xr.Dataset,
1846
+ spatial_dims: tuple[str, ...],
1847
+ spatial_coords: tuple[str, str],
1848
+ yx_bnds: tuple[tuple[float, float], tuple[float, float]],
1849
+ buffer: int,
1850
+ ) -> tuple[dict[str, slice], None] | tuple[dict[str, NDArray], float]:
1851
+ """
1852
+ Get a dict to index ``data`` in ``spatial_dims`` according to ``yx_bnds``
1853
+ plus a ``buffer``. Also return an estimate of the resolution if input grid is
1854
+ unstructured.
1855
+ """
1856
+
1857
+ def get_cond_xy(buffer_crs: float) -> NDArray[np.bool_]:
1858
+ """
1859
+ Get a boolean array to indicate which points are within x_bnds and y_bnds
1860
+ +- ``buffer_crs`` to index data on the unstructured grid dimension.
1861
+ """
1862
+
1863
+ cond_x = (data[grid_x].values > x_bnds[0] - buffer_crs) & (data[grid_x].values < x_bnds[1] + buffer_crs)
1864
+ cond_y = (data[grid_y].values > y_bnds[0] - buffer_crs) & (data[grid_y].values < y_bnds[1] + buffer_crs)
1865
+ return cond_x & cond_y
1866
+
1867
+ def get_ind_startend(coordname: str, dim: str, bnds: tuple[float, float]) -> tuple[int, int]:
1868
+ """
1869
+ Get start and end indices to slice data on ``dim`` based on ``bnds``.
1870
+ """
1871
+
1872
+ coords = data[coordname]
1873
+ dim_size = data.sizes[dim]
1874
+
1875
+ bool_decrease = dm.is_decreasing(coords, dim)
1876
+
1877
+ if bool_decrease:
1878
+ coords = coords.isel({dim: slice(None, None, -1)})
1879
+
1880
+ # Find start index: first value greater than lower bound
1881
+ id_start = (coords > bnds[0]).argmax(dim).values.min().item()
1882
+
1883
+ # Find end index: first value greater than upper bound, or use full size
1884
+ end_mask = coords > bnds[1]
1885
+ id_end = end_mask.argmax(dim).values.max().item() if end_mask.values.any() else dim_size
1886
+
1887
+ if bool_decrease:
1888
+ start_ori = id_start
1889
+ id_start = dim_size - id_end
1890
+ id_end = dim_size - start_ori
1891
+
1892
+ return id_start, id_end
1893
+
1894
+ grid_y, grid_x = spatial_coords
1895
+ y_bnds, x_bnds = yx_bnds
1896
+
1897
+ if len(spatial_dims) == 1:
1898
+ res_start = max(abs(np.diff(data[grid_x][:2])[0]), abs(np.diff(data[grid_y][:2])[0]))
1899
+ res_end = max(abs(np.diff(data[grid_x][-2:])[0]), abs(np.diff(data[grid_y][-2:])[0]))
1900
+ res_approx = max(res_start, res_end)
1901
+ cond_xy = get_cond_xy(2 * res_approx)
1902
+ res = _get_flat_grid_res(data[grid_x][cond_xy], data[grid_y][cond_xy])
1903
+ cond_xy = get_cond_xy(buffer * res * 1.5)
1904
+ return {spatial_dims[0]: cond_xy}, res
1905
+
1906
+ y_dim, x_dim = spatial_dims
1907
+ idx_start, idx_end = get_ind_startend(grid_x, x_dim, x_bnds)
1908
+ idy_start, idy_end = get_ind_startend(grid_y, y_dim, y_bnds)
1909
+ idx_start = max(idx_start - buffer, 0)
1910
+ idx_end += buffer
1911
+ idy_start = max(idy_start - buffer, 0)
1912
+ idy_end += buffer
1913
+ return {x_dim: slice(idx_start, idx_end), y_dim: slice(idy_start, idy_end)}, None
1914
+
1915
+
1916
+ def _get_flat_grid_res(coords_x: xr.DataArray, coords_y: xr.DataArray) -> float:
1917
+ """
1918
+ Get approximative resolution of a flattened grid. Assumes similar resolution
1919
+ in both horizontal dimensions.
1920
+ """
1921
+
1922
+ def dcoords(coords: xr.DataArray) -> float:
1923
+ return np.median(abs(np.diff(coords.values)))
1924
+
1925
+ return max(dcoords(coords_x), dcoords(coords_y))
1926
+
1927
+
1928
+ def _get_coarsen_ratio( # pylint: disable=too-many-locals
1929
+ data: xr.Dataset,
1930
+ spatial_dims: tuple[str, ...],
1931
+ coords_crs: tuple[str, str, CRS],
1932
+ dst_grid: xr.Dataset,
1933
+ dst_spatial_dims: tuple[str, str],
1934
+ dst_coords_crs: tuple[str, str, CRS],
1935
+ res: float | None = None,
1936
+ ) -> tuple[int, int]:
1937
+ """
1938
+ Get the resolution ratios in `x` and `y` between input data and destination grid
1939
+ to eventually coarsen the input data before interpolation.
1940
+ """
1941
+
1942
+ if res is not None:
1943
+ res_src_x = res_src_y = res
1944
+ else:
1945
+ center_src_x, center_src_y = _center_coords(data, spatial_dims, coords_crs[:2]) # type: ignore[arg-type]
1946
+ res_src_x = center_src_x.diff(spatial_dims[1]).values.mean()
1947
+ res_src_y = center_src_y.diff(spatial_dims[0]).values.mean()
1948
+
1949
+ center_dst_x, center_dst_y = _center_coords(dst_grid, dst_spatial_dims, dst_coords_crs[:2])
1950
+ src_crs, dst_crs = coords_crs[2], dst_coords_crs[2]
1951
+ if src_crs != dst_crs:
1952
+ center_dst_x.values, center_dst_y.values = reproject(center_dst_x.values, center_dst_y.values, dst_crs, src_crs)
1953
+
1954
+ res_dst_x = center_dst_x.diff(dst_spatial_dims[1]).values.mean()
1955
+ res_dst_y = center_dst_y.diff(dst_spatial_dims[0]).values.mean()
1956
+ return abs(round(res_dst_x / res_src_x)), abs(round(res_dst_y / res_src_y))
1957
+
1958
+
1959
+ def _center_coords(
1960
+ grid: xr.Dataset, dims: tuple[str, str], coords: tuple[str, str]
1961
+ ) -> tuple[xr.DataArray, xr.DataArray]:
1962
+ """
1963
+ Returns the "x" and "y" coordinates for the 4 center points of input grid domain.
1964
+ Used to compute the resolution of the grid.
1965
+ """
1966
+
1967
+ y_dim, x_dim = dims
1968
+ grid_y, grid_x = coords
1969
+ mid_x = int(grid.sizes[x_dim] / 2)
1970
+ mid_y = int(grid.sizes[y_dim] / 2)
1971
+ sub_x = grid[grid_x].isel({x_dim: slice(mid_x - 1, mid_x + 1)})
1972
+ sub_y = grid[grid_y].isel({y_dim: slice(mid_y - 1, mid_y + 1)})
1973
+ if y_dim in sub_x.dims:
1974
+ sub_x = sub_x.isel({y_dim: slice(mid_y - 1, mid_y + 1)})
1975
+ sub_y = sub_y.isel({x_dim: slice(mid_x - 1, mid_x + 1)})
1976
+ else:
1977
+ sub_y, sub_x = xr.broadcast(sub_y, sub_x)
1978
+
1979
+ return sub_x, sub_y
1980
+
1981
+
1982
+ def _combine_attrs(src_attrs: dict, dst_attrs: dict) -> dict:
1983
+ """
1984
+ Combine attrs from source and destination grids
1985
+ """
1986
+
1987
+ from_dst = ("grid_mapping", "crs", "bounding_box")
1988
+
1989
+ combined_attrs = {key: src_attrs[key] for key in SRC_ATTRS_KEEP if key in src_attrs}
1990
+ combined_attrs.update({key: dst_attrs[key] for key in from_dst if key in dst_attrs})
1991
+ return combined_attrs
1992
+
1993
+
1994
+ HEIGHT_METHODS = (
1995
+ "temp",
1996
+ "temp_fixgrad",
1997
+ "snowlim",
1998
+ "nearest_alt",
1999
+ "idw_3D",
2000
+ "nearest_3D",
2001
+ "linear_3D",
2002
+ "upscaled_3D",
2003
+ )
2004
+
2005
+ INTERP_FUNC = {
2006
+ "temp": interp_temp,
2007
+ "temp_fixgrad": interp_temp_fixgrad,
2008
+ "nearest_alt": interp_nearest_alt,
2009
+ "snowlim": interp_snowlim,
2010
+ "nearest": interp_nearest,
2011
+ "nearest_3D": interp_nearest,
2012
+ "idw": interp_idw,
2013
+ "idw_3D": interp_idw,
2014
+ "upscaled": interp_upscaled,
2015
+ "upscaled_3D": interp_upscaled,
2016
+ }