pycontrails 0.54.0__cp312-cp312-macosx_10_13_x86_64.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.

Potentially problematic release.


This version of pycontrails might be problematic. Click here for more details.

Files changed (109) hide show
  1. pycontrails/__init__.py +70 -0
  2. pycontrails/_version.py +16 -0
  3. pycontrails/core/__init__.py +30 -0
  4. pycontrails/core/aircraft_performance.py +641 -0
  5. pycontrails/core/airports.py +226 -0
  6. pycontrails/core/cache.py +881 -0
  7. pycontrails/core/coordinates.py +174 -0
  8. pycontrails/core/fleet.py +470 -0
  9. pycontrails/core/flight.py +2314 -0
  10. pycontrails/core/flightplan.py +220 -0
  11. pycontrails/core/fuel.py +140 -0
  12. pycontrails/core/interpolation.py +721 -0
  13. pycontrails/core/met.py +2833 -0
  14. pycontrails/core/met_var.py +307 -0
  15. pycontrails/core/models.py +1181 -0
  16. pycontrails/core/polygon.py +549 -0
  17. pycontrails/core/rgi_cython.cpython-312-darwin.so +0 -0
  18. pycontrails/core/vector.py +2190 -0
  19. pycontrails/datalib/__init__.py +12 -0
  20. pycontrails/datalib/_leo_utils/search.py +250 -0
  21. pycontrails/datalib/_leo_utils/static/bq_roi_query.sql +6 -0
  22. pycontrails/datalib/_leo_utils/vis.py +59 -0
  23. pycontrails/datalib/_met_utils/metsource.py +746 -0
  24. pycontrails/datalib/ecmwf/__init__.py +73 -0
  25. pycontrails/datalib/ecmwf/arco_era5.py +340 -0
  26. pycontrails/datalib/ecmwf/common.py +109 -0
  27. pycontrails/datalib/ecmwf/era5.py +550 -0
  28. pycontrails/datalib/ecmwf/era5_model_level.py +487 -0
  29. pycontrails/datalib/ecmwf/hres.py +782 -0
  30. pycontrails/datalib/ecmwf/hres_model_level.py +459 -0
  31. pycontrails/datalib/ecmwf/ifs.py +284 -0
  32. pycontrails/datalib/ecmwf/model_levels.py +434 -0
  33. pycontrails/datalib/ecmwf/static/model_level_dataframe_v20240418.csv +139 -0
  34. pycontrails/datalib/ecmwf/variables.py +267 -0
  35. pycontrails/datalib/gfs/__init__.py +28 -0
  36. pycontrails/datalib/gfs/gfs.py +646 -0
  37. pycontrails/datalib/gfs/variables.py +100 -0
  38. pycontrails/datalib/goes.py +772 -0
  39. pycontrails/datalib/landsat.py +569 -0
  40. pycontrails/datalib/sentinel.py +511 -0
  41. pycontrails/datalib/spire.py +739 -0
  42. pycontrails/ext/bada.py +41 -0
  43. pycontrails/ext/cirium.py +14 -0
  44. pycontrails/ext/empirical_grid.py +140 -0
  45. pycontrails/ext/synthetic_flight.py +430 -0
  46. pycontrails/models/__init__.py +1 -0
  47. pycontrails/models/accf.py +406 -0
  48. pycontrails/models/apcemm/__init__.py +8 -0
  49. pycontrails/models/apcemm/apcemm.py +982 -0
  50. pycontrails/models/apcemm/inputs.py +226 -0
  51. pycontrails/models/apcemm/static/apcemm_yaml_template.yaml +183 -0
  52. pycontrails/models/apcemm/utils.py +437 -0
  53. pycontrails/models/cocip/__init__.py +29 -0
  54. pycontrails/models/cocip/cocip.py +2616 -0
  55. pycontrails/models/cocip/cocip_params.py +299 -0
  56. pycontrails/models/cocip/cocip_uncertainty.py +285 -0
  57. pycontrails/models/cocip/contrail_properties.py +1517 -0
  58. pycontrails/models/cocip/output_formats.py +2261 -0
  59. pycontrails/models/cocip/radiative_forcing.py +1262 -0
  60. pycontrails/models/cocip/radiative_heating.py +520 -0
  61. pycontrails/models/cocip/unterstrasser_wake_vortex.py +403 -0
  62. pycontrails/models/cocip/wake_vortex.py +396 -0
  63. pycontrails/models/cocip/wind_shear.py +120 -0
  64. pycontrails/models/cocipgrid/__init__.py +9 -0
  65. pycontrails/models/cocipgrid/cocip_grid.py +2573 -0
  66. pycontrails/models/cocipgrid/cocip_grid_params.py +138 -0
  67. pycontrails/models/dry_advection.py +494 -0
  68. pycontrails/models/emissions/__init__.py +21 -0
  69. pycontrails/models/emissions/black_carbon.py +594 -0
  70. pycontrails/models/emissions/emissions.py +1353 -0
  71. pycontrails/models/emissions/ffm2.py +336 -0
  72. pycontrails/models/emissions/static/default-engine-uids.csv +239 -0
  73. pycontrails/models/emissions/static/edb-gaseous-v29b-engines.csv +596 -0
  74. pycontrails/models/emissions/static/edb-nvpm-v29b-engines.csv +215 -0
  75. pycontrails/models/humidity_scaling/__init__.py +37 -0
  76. pycontrails/models/humidity_scaling/humidity_scaling.py +1025 -0
  77. pycontrails/models/humidity_scaling/quantiles/era5-model-level-quantiles.pq +0 -0
  78. pycontrails/models/humidity_scaling/quantiles/era5-pressure-level-quantiles.pq +0 -0
  79. pycontrails/models/issr.py +210 -0
  80. pycontrails/models/pcc.py +327 -0
  81. pycontrails/models/pcr.py +154 -0
  82. pycontrails/models/ps_model/__init__.py +17 -0
  83. pycontrails/models/ps_model/ps_aircraft_params.py +376 -0
  84. pycontrails/models/ps_model/ps_grid.py +505 -0
  85. pycontrails/models/ps_model/ps_model.py +1017 -0
  86. pycontrails/models/ps_model/ps_operational_limits.py +540 -0
  87. pycontrails/models/ps_model/static/ps-aircraft-params-20240524.csv +68 -0
  88. pycontrails/models/ps_model/static/ps-synonym-list-20240524.csv +103 -0
  89. pycontrails/models/sac.py +459 -0
  90. pycontrails/models/tau_cirrus.py +168 -0
  91. pycontrails/physics/__init__.py +1 -0
  92. pycontrails/physics/constants.py +116 -0
  93. pycontrails/physics/geo.py +989 -0
  94. pycontrails/physics/jet.py +837 -0
  95. pycontrails/physics/thermo.py +451 -0
  96. pycontrails/physics/units.py +472 -0
  97. pycontrails/py.typed +0 -0
  98. pycontrails/utils/__init__.py +1 -0
  99. pycontrails/utils/dependencies.py +66 -0
  100. pycontrails/utils/iteration.py +13 -0
  101. pycontrails/utils/json.py +188 -0
  102. pycontrails/utils/temp.py +50 -0
  103. pycontrails/utils/types.py +165 -0
  104. pycontrails-0.54.0.dist-info/LICENSE +178 -0
  105. pycontrails-0.54.0.dist-info/METADATA +179 -0
  106. pycontrails-0.54.0.dist-info/NOTICE +43 -0
  107. pycontrails-0.54.0.dist-info/RECORD +109 -0
  108. pycontrails-0.54.0.dist-info/WHEEL +5 -0
  109. pycontrails-0.54.0.dist-info/top_level.txt +3 -0
@@ -0,0 +1,721 @@
1
+ """Interpolation utilities."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import dataclasses
6
+ import logging
7
+ from typing import Any, Literal, overload
8
+
9
+ import numpy as np
10
+ import numpy.typing as npt
11
+ import scipy.interpolate
12
+ import xarray as xr
13
+
14
+ from pycontrails.core import rgi_cython # type: ignore[attr-defined]
15
+
16
+ # ------------------------------------------------------------------------------
17
+ # Multidimensional interpolation
18
+ # ------------------------------------------------------------------------------
19
+
20
+
21
+ logger = logging.getLogger(__name__)
22
+
23
+
24
+ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolator):
25
+ """Support for performant interpolation over a regular grid.
26
+
27
+ This class is a thin wrapper around the
28
+ :class:`scipy.interpolate.RegularGridInterpolator` in order to make typical
29
+ ``pycontrails`` use-cases more efficient.
30
+
31
+ #. Avoid ``RegularGridInterpolator`` constructor validation. In :func:`interp`,
32
+ parameters are carefully crafted to fit into the intended form, thereby making
33
+ validation unnecessary.
34
+ #. Override the :meth:`_evaluate_linear` method with a faster implementation. See
35
+ the :meth:`_evaluate_linear` docstring for more information.
36
+
37
+ This class should not be used directly. Instead, use the :func:`interp` function.
38
+
39
+ .. versionchanged:: 0.40.0
40
+
41
+ The :meth:`_evaluate_linear` method now uses a Cython implementation. The dtype
42
+ of the output is now consistent with the dtype of the underlying :attr:`values`
43
+
44
+ Parameters
45
+ ----------
46
+ points : tuple[npt.NDArray[np.float64], ...]
47
+ Coordinates of the grid points.
48
+ values : npt.NDArray[np.float64]
49
+ Grid values. The shape of this array must be compatible with the
50
+ coordinates. An error is raised if the dtype is not ``np.float32``
51
+ or ``np.float64``.
52
+ method : str
53
+ Passed into :class:`scipy.interpolate.RegularGridInterpolator`
54
+ bounds_error : bool
55
+ Passed into :class:`scipy.interpolate.RegularGridInterpolator`
56
+ fill_value : float | np.float64 | None
57
+ Passed into :class:`scipy.interpolate.RegularGridInterpolator`
58
+ """
59
+
60
+ def __init__(
61
+ self,
62
+ points: tuple[npt.NDArray[np.float64], ...],
63
+ values: npt.NDArray[np.float64],
64
+ method: str,
65
+ bounds_error: bool,
66
+ fill_value: float | np.float64 | None,
67
+ ):
68
+ if values.dtype not in (np.float32, np.float64):
69
+ msg = f"values must be a float array, not {values.dtype}"
70
+ raise ValueError(msg)
71
+
72
+ self.grid = points
73
+ self.values = values
74
+ # TODO: consider supporting updated tensor-product spline methods
75
+ # see https://github.com/scipy/scipy/releases/tag/v1.13.0
76
+ self.method = _pick_method(scipy.__version__, method)
77
+ self.bounds_error = bounds_error
78
+ self.fill_value = fill_value
79
+ self._spline = None
80
+
81
+ def _prepare_xi_simple(self, xi: npt.NDArray[np.float64]) -> npt.NDArray[np.bool_]:
82
+ """Run looser version of :meth:`_prepare_xi`.
83
+
84
+ Parameters
85
+ ----------
86
+ xi : npt.NDArray[np.float64]
87
+ Points at which to interpolate.
88
+
89
+ Returns
90
+ -------
91
+ npt.NDArray[np.bool_]
92
+ A 1-dimensional Boolean array indicating which points are out of bounds.
93
+ If ``bounds_error`` is ``True``, this will be all ``False``.
94
+ """
95
+
96
+ if self.bounds_error:
97
+ for i, p in enumerate(xi.T):
98
+ g0 = self.grid[i][0]
99
+ g1 = self.grid[i][-1]
100
+ if not (np.all(p >= g0) and np.all(p <= g1)):
101
+ msg = f"One of the requested xi is out of bounds in dimension {i}"
102
+ raise ValueError(msg)
103
+
104
+ return np.zeros(xi.shape[0], dtype=bool)
105
+
106
+ return self._find_out_of_bounds(xi.T)
107
+
108
+ def __call__(
109
+ self, xi: npt.NDArray[np.float64], method: str | None = None
110
+ ) -> npt.NDArray[np.float64]:
111
+ """Evaluate the interpolator at the given points.
112
+
113
+ Parameters
114
+ ----------
115
+ xi : npt.NDArray[np.float64]
116
+ Points at which to interpolate. Must have shape ``(n, ndim)``, where
117
+ ``ndim`` is the number of dimensions of the interpolator.
118
+ method : str | None
119
+ Override the :attr:`method` to keep parity with the base class.
120
+
121
+ Returns
122
+ -------
123
+ npt.NDArray[np.float64]
124
+ Interpolated values. Has shape ``(n,)``. When computing linear interpolation,
125
+ the dtype is the same as the :attr:`values` array.
126
+ """
127
+
128
+ method = method or self.method
129
+ if method != "linear":
130
+ return super().__call__(xi, method)
131
+
132
+ out_of_bounds = self._prepare_xi_simple(xi)
133
+ xi_indices, norm_distances = rgi_cython.find_indices(self.grid, xi.T)
134
+
135
+ out = self._evaluate_linear(xi_indices, norm_distances)
136
+ return self._set_out_of_bounds(out, out_of_bounds)
137
+
138
+ def _set_out_of_bounds(
139
+ self,
140
+ out: npt.NDArray[np.float64],
141
+ out_of_bounds: npt.NDArray[np.bool_],
142
+ ) -> npt.NDArray[np.float64]:
143
+ """Set out-of-bounds values to the fill value.
144
+
145
+ Parameters
146
+ ----------
147
+ out : npt.NDArray[np.float64]
148
+ Values from interpolation. This is modified in-place.
149
+ out_of_bounds : npt.NDArray[np.bool_]
150
+ A 1-dimensional Boolean array indicating which points are out of bounds.
151
+
152
+ Returns
153
+ -------
154
+ out : npt.NDArray[np.float64]
155
+ A reference to the ``out`` array.
156
+ """
157
+ if self.fill_value is not None and np.any(out_of_bounds):
158
+ out[out_of_bounds] = self.fill_value
159
+
160
+ return out
161
+
162
+ def _evaluate_linear(
163
+ self,
164
+ indices: npt.NDArray[np.int64],
165
+ norm_distances: npt.NDArray[np.float64],
166
+ ) -> npt.NDArray[np.float64]:
167
+ """Evaluate the interpolator using linear interpolation.
168
+
169
+ This is a faster alternative to
170
+ :meth:`scipy.interpolate.RegularGridInterpolator._evaluate_linear`.
171
+
172
+ .. versionadded:: 0.24
173
+
174
+ .. versionchanged:: 0.40.0
175
+
176
+ Use Cython routines for evaluating the interpolation when the
177
+ dimension is 1, 2, 3, or 4.
178
+
179
+ Parameters
180
+ ----------
181
+ indices : npt.NDArray[np.int64]
182
+ Indices of the grid points to the left of the interpolation points.
183
+ Has shape ``(ndim, n_points)``.
184
+ norm_distances : npt.NDArray[np.float64]
185
+ Normalized distances between the interpolation points and the grid
186
+ points to the left. Has shape ``(ndim, n_points)``.
187
+
188
+ Returns
189
+ -------
190
+ npt.NDArray[np.float64]
191
+ Interpolated values with shape ``(n_points,)`` and the same dtype as
192
+ the :attr:`values` array.
193
+ """
194
+ # Let scipy "slow" implementation deal with high-dimensional grids
195
+ if indices.shape[0] > 4:
196
+ return super()._evaluate_linear(indices, norm_distances)
197
+
198
+ # Squeeze as much as possible
199
+ # Our cython implementation requires non-degenerate arrays
200
+ non_degen = tuple(s > 1 for s in self.values.shape)
201
+ values = self.values.squeeze()
202
+ indices = indices[non_degen, :]
203
+ norm_distances = norm_distances[non_degen, :]
204
+
205
+ ndim, n_points = indices.shape
206
+ out = np.empty(n_points, dtype=self.values.dtype)
207
+
208
+ if ndim == 4:
209
+ return rgi_cython.evaluate_linear_4d(values, indices, norm_distances, out)
210
+
211
+ if ndim == 3:
212
+ return rgi_cython.evaluate_linear_3d(values, indices, norm_distances, out)
213
+
214
+ if ndim == 2:
215
+ return rgi_cython.evaluate_linear_2d(values, indices, norm_distances, out)
216
+
217
+ if ndim == 1:
218
+ # np.interp could be better ... although that may also promote the dtype
219
+ # 1-d view is required for evaluate_linear_1d
220
+ return rgi_cython.evaluate_linear_1d(values, indices[0, :], norm_distances[0, :], out)
221
+
222
+ msg = f"Invalid number of dimensions: {ndim}"
223
+ raise ValueError(msg)
224
+
225
+
226
+ def _pick_method(scipy_version: str, method: str) -> str:
227
+ """Select an interpolation method.
228
+
229
+ For scipy versions 1.13.0 and later, fall back on legacy implementations
230
+ of tensor-product spline methods. The default implementations in 1.13.0
231
+ and later are incompatible with this class.
232
+
233
+ Parameters
234
+ ----------
235
+ scipy_version : str
236
+ scipy version (major.minor.patch)
237
+
238
+ method : str
239
+ Interpolation method. Passed into :class:`scipy.interpolate.RegularGridInterpolator`
240
+ as-is unless ``scipy_version`` is 1.13.0 or later and ``method`` is ``"slinear"``,
241
+ ``"cubic"``, or ``"quintic"``. In this case, ``"_legacy"`` is appended to ``method``.
242
+
243
+ Returns
244
+ -------
245
+ str
246
+ Interpolation method adjusted for compatibility with this class.
247
+ """
248
+ try:
249
+ version = scipy_version.split(".")
250
+ major = int(version[0])
251
+ minor = int(version[1])
252
+ except (IndexError, ValueError) as exc:
253
+ msg = f"Failed to parse major and minor version from {scipy_version}"
254
+ raise ValueError(msg) from exc
255
+
256
+ reimplemented_methods = ["slinear", "cubic", "quintic"]
257
+ if major > 1 or (major == 1 and minor >= 13) and method in reimplemented_methods:
258
+ return method + "_legacy"
259
+ return method
260
+
261
+
262
+ def _floatize_time(
263
+ time: npt.NDArray[np.datetime64], offset: np.datetime64
264
+ ) -> npt.NDArray[np.float64]:
265
+ """Convert an array of ``np.datetime64`` to an array of ``np.float64``.
266
+
267
+ In calls to :class:`scipy.interpolate.RegularGridInterpolator`, it's critical
268
+ that every coordinate be of same type. This creates complications: spatial
269
+ coordinates are float-like, whereas time coordinates are datetime-like. In
270
+ particular, it is not possible to cast an ``np.datetime64`` to a float
271
+ without losing information. In practice, this is not problematic because
272
+ ``np.float64`` has plenty of precision. Previously, this was more of an issue
273
+ because we used ``np.float32``.
274
+
275
+ This function uses a fixed time resolution (1 millisecond) to convert the time
276
+ coordinate to a float-like coordinate. The time resolution is taken to avoid
277
+ losing too much information for the time scales we encounter.
278
+
279
+ Care is taken to ensure "nat" values are converted to "nan".
280
+
281
+ Note that ``xarray`` also must confront this issue. They take a similar approach
282
+ in :func:`xarray.core.missing._floatize_x`. See
283
+ https://github.com/pydata/xarray/blob/d4db16699f30ad1dc3e6861601247abf4ac96567/xarray/core/missing.py#L572
284
+
285
+ .. versionchanged:: 0.40.0
286
+
287
+ No longer allow the option of converting to ``np.float32``. No longer
288
+ floor the time values to the preceding millisecond.
289
+
290
+ Parameters
291
+ ----------
292
+ time : npt.NDArray[np.datetime64]
293
+ Array of ``np.datetime64`` values.
294
+ offset : np.datetime64
295
+ The offset to subtract from ``time``.
296
+
297
+ Returns
298
+ -------
299
+ npt.NDArray[np.float64]
300
+ The number of milliseconds since ``offset``.
301
+ """
302
+ delta = time - offset
303
+ resolution = np.timedelta64(1, "ms")
304
+ return delta / resolution
305
+
306
+
307
+ def _localize(da: xr.DataArray, coords: dict[str, np.ndarray]) -> xr.DataArray:
308
+ """Clip ``da`` to the smallest bounding box that contains all of ``coords``.
309
+
310
+ Roughly follows approach taken by :func:`xarray.core.missing._localize`. See
311
+ https://github.com/pydata/xarray/blob/56f05c37924071eb4712479d47432aafd4dce38b/xarray/core/missing.py#L557
312
+
313
+ Parameters
314
+ ----------
315
+ da : xr.DataArray
316
+ DataArray to clip.
317
+ coords : dict[str, np.ndarray]
318
+ Coordinates to clip to.
319
+
320
+ Returns
321
+ -------
322
+ xr.DataArray
323
+ Clipped :class:`xarray.DataArray`. Has the same dimensions as the input ``da``.
324
+ In particular, each dimension of the returned DataArray is a slice of the
325
+ corresponding dimension of the input ``da``.
326
+ """
327
+ indexes: dict[str, Any] = {}
328
+ for dim, arr in coords.items():
329
+ dim_vals = da[dim].values
330
+
331
+ # Skip single level
332
+ if dim == "level" and dim_vals.size == 1 and dim_vals.item() == -1:
333
+ continue
334
+
335
+ # Create slice
336
+ minval = np.nanmin(arr)
337
+ maxval = np.nanmax(arr)
338
+ imin = np.searchsorted(dim_vals, minval, side="right") - 1
339
+ imin = max(0, imin)
340
+ imax = np.searchsorted(dim_vals, maxval, side="left") + 1
341
+ indexes[dim] = slice(imin, imax)
342
+
343
+ # Logging
344
+ n_in_bounds = np.sum((arr >= minval) & (arr <= maxval))
345
+ logger.debug(
346
+ "Interpolation in bounds along dimension %s: %d/%d",
347
+ dim,
348
+ n_in_bounds,
349
+ arr.size,
350
+ )
351
+
352
+ return da.isel(**indexes)
353
+
354
+
355
+ @overload
356
+ def interp(
357
+ longitude: npt.NDArray[np.float64],
358
+ latitude: npt.NDArray[np.float64],
359
+ level: npt.NDArray[np.float64],
360
+ time: npt.NDArray[np.datetime64],
361
+ da: xr.DataArray,
362
+ method: str,
363
+ bounds_error: bool,
364
+ fill_value: float | np.float64 | None,
365
+ localize: bool,
366
+ *,
367
+ indices: RGIArtifacts | None = ...,
368
+ return_indices: Literal[False] = ...,
369
+ ) -> npt.NDArray[np.float64]: ...
370
+
371
+
372
+ @overload
373
+ def interp(
374
+ longitude: npt.NDArray[np.float64],
375
+ latitude: npt.NDArray[np.float64],
376
+ level: npt.NDArray[np.float64],
377
+ time: npt.NDArray[np.datetime64],
378
+ da: xr.DataArray,
379
+ method: str,
380
+ bounds_error: bool,
381
+ fill_value: float | np.float64 | None,
382
+ localize: bool,
383
+ *,
384
+ indices: RGIArtifacts | None = ...,
385
+ return_indices: Literal[True],
386
+ ) -> tuple[npt.NDArray[np.float64], RGIArtifacts]: ...
387
+
388
+
389
+ @overload
390
+ def interp(
391
+ longitude: npt.NDArray[np.float64],
392
+ latitude: npt.NDArray[np.float64],
393
+ level: npt.NDArray[np.float64],
394
+ time: npt.NDArray[np.datetime64],
395
+ da: xr.DataArray,
396
+ method: str,
397
+ bounds_error: bool,
398
+ fill_value: float | np.float64 | None,
399
+ localize: bool,
400
+ *,
401
+ indices: RGIArtifacts | None = ...,
402
+ return_indices: bool = ...,
403
+ ) -> npt.NDArray[np.float64] | tuple[npt.NDArray[np.float64], RGIArtifacts]: ...
404
+
405
+
406
+ def interp(
407
+ longitude: npt.NDArray[np.float64],
408
+ latitude: npt.NDArray[np.float64],
409
+ level: npt.NDArray[np.float64],
410
+ time: npt.NDArray[np.datetime64],
411
+ da: xr.DataArray,
412
+ method: str,
413
+ bounds_error: bool,
414
+ fill_value: float | np.float64 | None,
415
+ localize: bool,
416
+ *,
417
+ indices: RGIArtifacts | None = None,
418
+ return_indices: bool = False,
419
+ ) -> npt.NDArray[np.float64] | tuple[npt.NDArray[np.float64], RGIArtifacts]:
420
+ """Interpolate over a grid with ``localize`` option.
421
+
422
+ .. versionchanged:: 0.25.6
423
+
424
+ Utilize scipy 1.9 upgrades to remove singleton dimensions.
425
+
426
+ .. versionchanged:: 0.26.0
427
+
428
+ Include ``indices`` and ``return_indices`` experimental parameters.
429
+ Currently, nan values in ``longitude``, ``latitude``, ``level``, or ``time``
430
+ are always propagated through to the output, regardless of ``bounds_error``.
431
+ In other words, a ValueError for an out of bounds coordinate is only raised
432
+ if a non-nan value is out of bounds.
433
+
434
+ .. versionchanged:: 0.40.0
435
+
436
+ When ``return_indices`` is True, an instance of :class:`RGIArtifacts`
437
+ is used to store the indices artifacts.
438
+
439
+ Parameters
440
+ ----------
441
+ longitude, latitude, level, time : np.ndarray
442
+ Coordinates of points to be interpolated. These parameters have the same
443
+ meaning as ``x`` in analogy with :func:`numpy.interp`. All four of these
444
+ arrays must be 1 dimensional of the same size.
445
+ da : xr.DataArray
446
+ Gridded data interpolated over. Must adhere to ``MetDataArray`` conventions.
447
+ In particular, the dimensions of ``da`` must be ``longitude``, ``latitude``,
448
+ ``level``, and ``time``. The three spatial dimensions must be monotonically
449
+ increasing with ``float64`` dtype. The ``time`` dimension must be
450
+ monotonically increasing with ``datetime64`` dtype.
451
+ Assumed to be cheap to load into memory (:attr:`xr.DataArray.values` is
452
+ used without hesitation).
453
+ method : str
454
+ Passed into :class:`scipy.interpolate.RegularGridInterpolator`.
455
+ bounds_error : bool
456
+ Passed into :class:`scipy.interpolate.RegularGridInterpolator`.
457
+ fill_value : float | np.float64 | None
458
+ Passed into :class:`scipy.interpolate.RegularGridInterpolator`.
459
+ localize : bool
460
+ If True, clip ``da`` to the smallest bounding box that contains all of
461
+ ``coords``.
462
+ indices : tuple | None, optional
463
+ Experimental. Provide intermediate artifacts computed by
464
+ :meth:``scipy.interpolate.RegularGridInterpolator._find_indices`
465
+ to avoid redundant computation. If known and provided, this can speed
466
+ up interpolation by avoiding an unnecessary call to ``_find_indices``.
467
+ By default, None. Must be used precisely.
468
+ return_indices : bool, optional
469
+ If True, return output of :meth:`scipy.interpolate.RegularGridInterpolator._find_indices`
470
+ in addition to interpolated values.
471
+
472
+ Returns
473
+ -------
474
+ npt.NDArray[np.float64] | tuple[npt.NDArray[np.float64], RGIArtifacts]
475
+ Interpolated values with same size as ``longitude``. If ``return_indices``
476
+ is True, return intermediate indices artifacts as well.
477
+
478
+ See Also
479
+ --------
480
+ - :meth:`MetDataArray.interpolate`
481
+ - :func:`scipy.interpolate.interpn`
482
+ - :class:`scipy.interpolate.RegularGridInterpolator`
483
+ """
484
+ if localize:
485
+ coords = {"longitude": longitude, "latitude": latitude, "level": level, "time": time}
486
+ da = _localize(da, coords)
487
+
488
+ indexes = da._indexes
489
+ x = indexes["longitude"].index.to_numpy() # type: ignore[attr-defined]
490
+ y = indexes["latitude"].index.to_numpy() # type: ignore[attr-defined]
491
+ z = indexes["level"].index.to_numpy() # type: ignore[attr-defined]
492
+ if any(v.dtype != np.float64 for v in (x, y, z)):
493
+ msg = "da must have float64 dtype for longitude, latitude, and level coordinates"
494
+ raise ValueError(msg)
495
+
496
+ # Convert t and time to float64
497
+ t = indexes["time"].index.to_numpy() # type: ignore[attr-defined]
498
+ offset = t[0]
499
+ t = _floatize_time(t, offset)
500
+
501
+ single_level = z.size == 1 and z.item() == -1.0
502
+ points: tuple[npt.NDArray[np.float64], ...]
503
+ if single_level:
504
+ values = da.values.squeeze(axis=2)
505
+ points = x, y, t
506
+ else:
507
+ values = da.values
508
+ points = x, y, z, t
509
+
510
+ interp_ = PycontrailsRegularGridInterpolator(
511
+ points=points,
512
+ values=values,
513
+ method=method,
514
+ bounds_error=bounds_error,
515
+ fill_value=fill_value,
516
+ )
517
+
518
+ if indices is None:
519
+ xi = _buildxi(longitude, latitude, level, time, offset, single_level)
520
+ if return_indices:
521
+ out, indices = _linear_interp_with_indices(interp_, xi, localize, None)
522
+ return out, indices
523
+ return interp_(xi)
524
+
525
+ out, indices = _linear_interp_with_indices(interp_, None, localize, indices)
526
+ if return_indices:
527
+ return out, indices
528
+ return out
529
+
530
+
531
+ def _buildxi(
532
+ longitude: npt.NDArray[np.float64],
533
+ latitude: npt.NDArray[np.float64],
534
+ level: npt.NDArray[np.float64],
535
+ time: npt.NDArray[np.datetime64],
536
+ offset: np.datetime64,
537
+ single_level: bool,
538
+ ) -> npt.NDArray[np.float64]:
539
+ """Build the input array for interpolation.
540
+
541
+ The implementation below achieves the same result as the following::
542
+
543
+ np.stack([longitude, latitude, level, time_float], axis=1])
544
+
545
+ This implementation is slightly faster than the above.
546
+ """
547
+
548
+ time_float = _floatize_time(time, offset)
549
+
550
+ ndim = 3 if single_level else 4
551
+ shape = longitude.size, ndim
552
+ xi = np.empty(shape, dtype=np.float64)
553
+
554
+ xi[:, 0] = longitude
555
+ xi[:, 1] = latitude
556
+ if not single_level:
557
+ xi[:, 2] = level
558
+ xi[:, -1] = time_float
559
+
560
+ return xi
561
+
562
+
563
+ def _linear_interp_with_indices(
564
+ interp: PycontrailsRegularGridInterpolator,
565
+ xi: npt.NDArray[np.float64] | None,
566
+ localize: bool,
567
+ indices: RGIArtifacts | None,
568
+ ) -> tuple[npt.NDArray[np.float64], RGIArtifacts]:
569
+ if interp.method != "linear":
570
+ msg = "Parameter 'indices' is only supported for 'method=linear'"
571
+ raise ValueError(msg)
572
+ if localize:
573
+ msg = "Parameter 'indices' is only supported for 'localize=False'"
574
+ raise ValueError(msg)
575
+
576
+ if indices is None:
577
+ assert xi is not None, "xi must be provided if indices is None"
578
+ out_of_bounds = interp._prepare_xi_simple(xi)
579
+ xi_indices, norm_distances = rgi_cython.find_indices(interp.grid, xi.T)
580
+ indices = RGIArtifacts(xi_indices, norm_distances, out_of_bounds)
581
+
582
+ out = interp._evaluate_linear(indices.xi_indices, indices.norm_distances)
583
+ out = interp._set_out_of_bounds(out, indices.out_of_bounds)
584
+ return out, indices
585
+
586
+
587
+ @dataclasses.dataclass
588
+ class RGIArtifacts:
589
+ """An interface to intermediate RGI interpolation artifacts."""
590
+
591
+ xi_indices: npt.NDArray[np.int64]
592
+ norm_distances: npt.NDArray[np.float64]
593
+ out_of_bounds: npt.NDArray[np.bool_]
594
+
595
+
596
+ # ------------------------------------------------------------------------------
597
+ # 1 dimensional interpolation
598
+ # ------------------------------------------------------------------------------
599
+
600
+
601
+ class EmissionsProfileInterpolator:
602
+ """Support for interpolating a profile on a linear or logarithmic scale.
603
+
604
+ This class simply wraps :func:`numpy.interp` with fixed values for the
605
+ ``xp`` and ``fp`` arguments. Unlike :class:`xarray.DataArray` interpolation,
606
+ the `numpy.interp` automatically clips values outside the range of the
607
+ ``xp`` array.
608
+
609
+ Parameters
610
+ ----------
611
+ xp : npt.NDarray[np.float64]
612
+ Array of x-values. These must be strictly increasing and free from
613
+ any nan values. Passed to :func:`numpy.interp`.
614
+ fp : npt.NDarray[np.float64]
615
+ Array of y-values. Passed to :func:`numpy.interp`.
616
+ drop_duplicates : bool, optional
617
+ Whether to drop duplicate values in ``xp``. Defaults to ``True``.
618
+
619
+ Examples
620
+ --------
621
+ >>> xp = np.array([3, 7, 10, 30], dtype=float)
622
+ >>> fp = np.array([0.1, 0.2, 0.3, 0.4], dtype=float)
623
+ >>> epi = EmissionsProfileInterpolator(xp, fp)
624
+ >>> # Interpolate a single value
625
+ >>> epi.interp(5)
626
+ np.float64(0.150000...)
627
+
628
+ >>> # Interpolate a single value on a logarithmic scale
629
+ >>> epi.log_interp(5)
630
+ np.float64(1.105171...)
631
+
632
+ >>> # Demonstrate speed up compared with xarray.DataArray interpolation
633
+ >>> import time, xarray as xr
634
+ >>> da = xr.DataArray(fp, dims=["x"], coords={"x": xp})
635
+
636
+ >>> inputs = [np.random.uniform(0, 31, size=200) for _ in range(1000)]
637
+ >>> t0 = time.perf_counter()
638
+ >>> xr_out = [da.interp(x=x.clip(3, 30)).values for x in inputs]
639
+ >>> t1 = time.perf_counter()
640
+ >>> np_out = [epi.interp(x) for x in inputs]
641
+ >>> t2 = time.perf_counter()
642
+ >>> assert np.allclose(xr_out, np_out)
643
+
644
+ >>> # We see a 100 fold speed up (more like 500x faster, but we don't
645
+ >>> # want the test to fail!)
646
+ >>> assert t2 - t1 < (t1 - t0) / 100
647
+ """
648
+
649
+ def __init__(
650
+ self,
651
+ xp: npt.NDArray[np.float64],
652
+ fp: npt.NDArray[np.float64],
653
+ drop_duplicates: bool = True,
654
+ ) -> None:
655
+ if drop_duplicates:
656
+ # Using np.diff to detect duplicates ... this assumes xp is sorted.
657
+ # If xp is not sorted, an ValueError will be raised in _validate
658
+ mask = np.abs(np.diff(xp, prepend=np.inf)) < 1e-15 # small tolerance
659
+ xp = xp[~mask]
660
+ fp = fp[~mask]
661
+
662
+ self.xp = np.asarray(xp)
663
+ self.fp = np.asarray(fp)
664
+ self._validate()
665
+
666
+ def __repr__(self) -> str:
667
+ return f"{self.__class__.__name__}(xp={self.xp}, fp={self.fp})"
668
+
669
+ def _validate(self) -> None:
670
+ if not len(self.xp):
671
+ msg = "xp must not be empty"
672
+ raise ValueError(msg)
673
+ if len(self.xp) != len(self.fp):
674
+ msg = "xp and fp must have the same length"
675
+ raise ValueError(msg)
676
+ if not np.all(np.diff(self.xp) > 0.0):
677
+ msg = "xp must be strictly increasing"
678
+ raise ValueError(msg)
679
+ if np.any(np.isnan(self.xp)):
680
+ msg = "xp must not contain nan values"
681
+ raise ValueError(msg)
682
+
683
+ def interp(self, x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
684
+ """Interpolate x against xp and fp.
685
+
686
+ Parameters
687
+ ----------
688
+ x : npt.NDArray[np.float64]
689
+ Array of x-values to interpolate.
690
+
691
+ Returns
692
+ -------
693
+ npt.NDArray[np.float64]
694
+ Array of interpolated y-values arising from the x-values. The ``dtype`` of
695
+ the output array is the same as the ``dtype`` of ``x``.
696
+ """
697
+ # Need to explicitly cast back to x.dtype
698
+ # https://github.com/numpy/numpy/issues/11214
699
+ x = np.asarray(x)
700
+ dtype = np.result_type(x, np.float32)
701
+ return np.interp(x, self.xp, self.fp).astype(dtype)
702
+
703
+ def log_interp(self, x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
704
+ """Interpolate x against xp and fp on a logarithmic scale.
705
+
706
+ This method composes the following three functions.
707
+ 1. :func:`numpy.log`
708
+ 2. :meth:`interp`
709
+ 3. :func:`numpy.exp`
710
+
711
+ Parameters
712
+ ----------
713
+ x : npt.NDArray[np.float64]
714
+ Array of x-values to interpolate.
715
+
716
+ Returns
717
+ -------
718
+ npt.NDArray[np.float64]
719
+ Array of interpolated y-values arising from the x-values.
720
+ """
721
+ return np.exp(self.interp(np.log(x)))