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.
- grid_process/__init__.py +14 -0
- grid_process/aggreg.py +447 -0
- grid_process/blending.py +206 -0
- grid_process/config/__init__.py +5 -0
- grid_process/config/grid_process.yaml +20 -0
- grid_process/config/library_settings.py +68 -0
- grid_process/da_utils.py +104 -0
- grid_process/data_model.py +531 -0
- grid_process/ds_utils.py +61 -0
- grid_process/ens_stats.py +396 -0
- grid_process/grid_interp.py +2016 -0
- grid_process/grid_interp_utils.py +208 -0
- grid_process/helpers.py +174 -0
- grid_process/optimization.py +50 -0
- grid_process/sun.py +502 -0
- grid_process/time_utils.py +110 -0
- grid_process/timeseries.py +547 -0
- grid_process/units_conv.py +209 -0
- grid_process-1.0.0.dist-info/METADATA +155 -0
- grid_process-1.0.0.dist-info/RECORD +23 -0
- grid_process-1.0.0.dist-info/WHEEL +4 -0
- grid_process-1.0.0.dist-info/licenses/AUTHORS +11 -0
- grid_process-1.0.0.dist-info/licenses/LICENSE +29 -0
|
@@ -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
|
+
}
|