pycontrails 0.58.0__cp314-cp314-macosx_11_0_arm64.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 (122) hide show
  1. pycontrails/__init__.py +70 -0
  2. pycontrails/_version.py +34 -0
  3. pycontrails/core/__init__.py +30 -0
  4. pycontrails/core/aircraft_performance.py +679 -0
  5. pycontrails/core/airports.py +228 -0
  6. pycontrails/core/cache.py +889 -0
  7. pycontrails/core/coordinates.py +174 -0
  8. pycontrails/core/fleet.py +483 -0
  9. pycontrails/core/flight.py +2185 -0
  10. pycontrails/core/flightplan.py +228 -0
  11. pycontrails/core/fuel.py +140 -0
  12. pycontrails/core/interpolation.py +702 -0
  13. pycontrails/core/met.py +2931 -0
  14. pycontrails/core/met_var.py +387 -0
  15. pycontrails/core/models.py +1321 -0
  16. pycontrails/core/polygon.py +549 -0
  17. pycontrails/core/rgi_cython.cpython-314-darwin.so +0 -0
  18. pycontrails/core/vector.py +2249 -0
  19. pycontrails/datalib/__init__.py +12 -0
  20. pycontrails/datalib/_met_utils/metsource.py +746 -0
  21. pycontrails/datalib/ecmwf/__init__.py +73 -0
  22. pycontrails/datalib/ecmwf/arco_era5.py +345 -0
  23. pycontrails/datalib/ecmwf/common.py +114 -0
  24. pycontrails/datalib/ecmwf/era5.py +554 -0
  25. pycontrails/datalib/ecmwf/era5_model_level.py +490 -0
  26. pycontrails/datalib/ecmwf/hres.py +804 -0
  27. pycontrails/datalib/ecmwf/hres_model_level.py +466 -0
  28. pycontrails/datalib/ecmwf/ifs.py +287 -0
  29. pycontrails/datalib/ecmwf/model_levels.py +435 -0
  30. pycontrails/datalib/ecmwf/static/model_level_dataframe_v20240418.csv +139 -0
  31. pycontrails/datalib/ecmwf/variables.py +268 -0
  32. pycontrails/datalib/geo_utils.py +261 -0
  33. pycontrails/datalib/gfs/__init__.py +28 -0
  34. pycontrails/datalib/gfs/gfs.py +656 -0
  35. pycontrails/datalib/gfs/variables.py +104 -0
  36. pycontrails/datalib/goes.py +757 -0
  37. pycontrails/datalib/himawari/__init__.py +27 -0
  38. pycontrails/datalib/himawari/header_struct.py +266 -0
  39. pycontrails/datalib/himawari/himawari.py +667 -0
  40. pycontrails/datalib/landsat.py +589 -0
  41. pycontrails/datalib/leo_utils/__init__.py +5 -0
  42. pycontrails/datalib/leo_utils/correction.py +266 -0
  43. pycontrails/datalib/leo_utils/landsat_metadata.py +300 -0
  44. pycontrails/datalib/leo_utils/search.py +250 -0
  45. pycontrails/datalib/leo_utils/sentinel_metadata.py +748 -0
  46. pycontrails/datalib/leo_utils/static/bq_roi_query.sql +6 -0
  47. pycontrails/datalib/leo_utils/vis.py +59 -0
  48. pycontrails/datalib/sentinel.py +650 -0
  49. pycontrails/datalib/spire/__init__.py +5 -0
  50. pycontrails/datalib/spire/exceptions.py +62 -0
  51. pycontrails/datalib/spire/spire.py +604 -0
  52. pycontrails/ext/bada.py +42 -0
  53. pycontrails/ext/cirium.py +14 -0
  54. pycontrails/ext/empirical_grid.py +140 -0
  55. pycontrails/ext/synthetic_flight.py +431 -0
  56. pycontrails/models/__init__.py +1 -0
  57. pycontrails/models/accf.py +425 -0
  58. pycontrails/models/apcemm/__init__.py +8 -0
  59. pycontrails/models/apcemm/apcemm.py +983 -0
  60. pycontrails/models/apcemm/inputs.py +226 -0
  61. pycontrails/models/apcemm/static/apcemm_yaml_template.yaml +183 -0
  62. pycontrails/models/apcemm/utils.py +437 -0
  63. pycontrails/models/cocip/__init__.py +29 -0
  64. pycontrails/models/cocip/cocip.py +2742 -0
  65. pycontrails/models/cocip/cocip_params.py +305 -0
  66. pycontrails/models/cocip/cocip_uncertainty.py +291 -0
  67. pycontrails/models/cocip/contrail_properties.py +1530 -0
  68. pycontrails/models/cocip/output_formats.py +2270 -0
  69. pycontrails/models/cocip/radiative_forcing.py +1260 -0
  70. pycontrails/models/cocip/radiative_heating.py +520 -0
  71. pycontrails/models/cocip/unterstrasser_wake_vortex.py +508 -0
  72. pycontrails/models/cocip/wake_vortex.py +396 -0
  73. pycontrails/models/cocip/wind_shear.py +120 -0
  74. pycontrails/models/cocipgrid/__init__.py +9 -0
  75. pycontrails/models/cocipgrid/cocip_grid.py +2552 -0
  76. pycontrails/models/cocipgrid/cocip_grid_params.py +138 -0
  77. pycontrails/models/dry_advection.py +602 -0
  78. pycontrails/models/emissions/__init__.py +21 -0
  79. pycontrails/models/emissions/black_carbon.py +599 -0
  80. pycontrails/models/emissions/emissions.py +1353 -0
  81. pycontrails/models/emissions/ffm2.py +336 -0
  82. pycontrails/models/emissions/static/default-engine-uids.csv +239 -0
  83. pycontrails/models/emissions/static/edb-gaseous-v29b-engines.csv +596 -0
  84. pycontrails/models/emissions/static/edb-nvpm-v29b-engines.csv +215 -0
  85. pycontrails/models/extended_k15.py +1327 -0
  86. pycontrails/models/humidity_scaling/__init__.py +37 -0
  87. pycontrails/models/humidity_scaling/humidity_scaling.py +1075 -0
  88. pycontrails/models/humidity_scaling/quantiles/era5-model-level-quantiles.pq +0 -0
  89. pycontrails/models/humidity_scaling/quantiles/era5-pressure-level-quantiles.pq +0 -0
  90. pycontrails/models/issr.py +210 -0
  91. pycontrails/models/pcc.py +326 -0
  92. pycontrails/models/pcr.py +154 -0
  93. pycontrails/models/ps_model/__init__.py +18 -0
  94. pycontrails/models/ps_model/ps_aircraft_params.py +381 -0
  95. pycontrails/models/ps_model/ps_grid.py +701 -0
  96. pycontrails/models/ps_model/ps_model.py +1000 -0
  97. pycontrails/models/ps_model/ps_operational_limits.py +525 -0
  98. pycontrails/models/ps_model/static/ps-aircraft-params-20250328.csv +69 -0
  99. pycontrails/models/ps_model/static/ps-synonym-list-20250328.csv +104 -0
  100. pycontrails/models/sac.py +442 -0
  101. pycontrails/models/tau_cirrus.py +183 -0
  102. pycontrails/physics/__init__.py +1 -0
  103. pycontrails/physics/constants.py +117 -0
  104. pycontrails/physics/geo.py +1138 -0
  105. pycontrails/physics/jet.py +968 -0
  106. pycontrails/physics/static/iata-cargo-load-factors-20250221.csv +74 -0
  107. pycontrails/physics/static/iata-passenger-load-factors-20250221.csv +74 -0
  108. pycontrails/physics/thermo.py +551 -0
  109. pycontrails/physics/units.py +472 -0
  110. pycontrails/py.typed +0 -0
  111. pycontrails/utils/__init__.py +1 -0
  112. pycontrails/utils/dependencies.py +66 -0
  113. pycontrails/utils/iteration.py +13 -0
  114. pycontrails/utils/json.py +187 -0
  115. pycontrails/utils/temp.py +50 -0
  116. pycontrails/utils/types.py +163 -0
  117. pycontrails-0.58.0.dist-info/METADATA +180 -0
  118. pycontrails-0.58.0.dist-info/RECORD +122 -0
  119. pycontrails-0.58.0.dist-info/WHEEL +6 -0
  120. pycontrails-0.58.0.dist-info/licenses/LICENSE +178 -0
  121. pycontrails-0.58.0.dist-info/licenses/NOTICE +43 -0
  122. pycontrails-0.58.0.dist-info/top_level.txt +3 -0
@@ -0,0 +1,435 @@
1
+ """Utilities for working with ECMWF model-level data."""
2
+
3
+ import datetime
4
+ import pathlib
5
+ import warnings
6
+
7
+ import dask.array
8
+ import numpy as np
9
+ import numpy.typing as npt
10
+ import pandas as pd
11
+ import xarray as xr
12
+
13
+ from pycontrails.physics import units
14
+
15
+ _path_to_static = pathlib.Path(__file__).parent / "static"
16
+ MODEL_LEVELS_PATH = _path_to_static / "model_level_dataframe_v20240418.csv"
17
+
18
+
19
+ def model_level_reference_pressure(
20
+ alt_ft_min: float | None = None,
21
+ alt_ft_max: float | None = None,
22
+ ) -> list[int]:
23
+ """Return the pressure levels at each model level assuming a constant surface pressure.
24
+
25
+ This function assumes
26
+ `137 model levels <https://confluence.ecmwf.int/display/UDOC/L137+model+level+definitions>`_
27
+ and the constant ICAO ISA surface pressure of 1013.25 hPa.
28
+
29
+ The returned pressure levels are rounded to the nearest hPa.
30
+
31
+ Parameters
32
+ ----------
33
+ alt_ft_min : float | None
34
+ Minimum altitude, [:math:`ft`]. If None, there is no minimum altitude
35
+ used in filtering the ``MODEL_LEVELS_PATH`` table.
36
+ alt_ft_max : float | None
37
+ Maximum altitude, [:math:`ft`]. If None, there is no maximum altitude
38
+ used in filtering the ``MODEL_LEVELS_PATH`` table.
39
+
40
+ Returns
41
+ -------
42
+ list[int]
43
+ List of pressure levels, [:math:`hPa`] between the minimum and maximum altitudes.
44
+
45
+ See Also
46
+ --------
47
+ model_level_pressure
48
+ """
49
+ usecols = ["n", "Geometric Altitude [m]", "pf [hPa]"]
50
+ df = pd.read_csv(MODEL_LEVELS_PATH, usecols=usecols, index_col="n")
51
+
52
+ filt = df.index >= 1 # exclude degenerate model level 0
53
+ if alt_ft_min is not None:
54
+ alt_m_min = units.ft_to_m(alt_ft_min)
55
+ filt &= df["Geometric Altitude [m]"] >= alt_m_min
56
+ if alt_ft_max is not None:
57
+ alt_m_max = units.ft_to_m(alt_ft_max)
58
+ filt &= df["Geometric Altitude [m]"] <= alt_m_max
59
+
60
+ return df.loc[filt, "pf [hPa]"].round().astype(int).tolist()
61
+
62
+
63
+ def _cache_model_level_dataframe() -> None:
64
+ """Regenerate static model level data file.
65
+
66
+ Read the ERA5 L137 model level definitions published by ECMWF
67
+ and cache it in a static file for use by this module.
68
+ This should only be used by model developers, and only if ECMWF model
69
+ level definitions change. ``MODEL_LEVEL_PATH`` must be manually
70
+ updated to use newly-cached files.
71
+
72
+ Requires the `lxml <https://lxml.de/>`_ package to be installed.
73
+ """
74
+
75
+ url = "https://confluence.ecmwf.int/display/UDOC/L137+model+level+definitions"
76
+ df = pd.read_html(url, na_values="-", index_col="n")[0]
77
+
78
+ today = datetime.datetime.now()
79
+ new_file_path = _path_to_static / f"model_level_dataframe_v{today.strftime('%Y%m%d')}.csv"
80
+ if new_file_path.is_file():
81
+ msg = f"Static file already exists at {new_file_path}"
82
+ raise ValueError(msg)
83
+
84
+ df.to_csv(new_file_path)
85
+
86
+
87
+ def model_level_pressure(sp: xr.DataArray, model_levels: npt.ArrayLike) -> xr.DataArray:
88
+ r"""Return the pressure levels at each model level given the surface pressure.
89
+
90
+ This function assumes
91
+ `137 model levels <https://confluence.ecmwf.int/display/UDOC/L137+model+level+definitions>`_.
92
+ Unlike :func:`model_level_reference_pressure`, this function
93
+ does not assume constant pressure. Instead, it uses the
94
+ `half-level pressure formula <https://confluence.ecmwf.int/x/JJh0CQ#heading-Pressureonmodellevels>`_
95
+ :math:`p = a + b \cdot \text{sp}` where :math:`a` and :math:`b` are constants
96
+ for each model level.
97
+
98
+ Parameters
99
+ ----------
100
+ sp : xr.DataArray
101
+ Surface pressure, [:math:`\text{Pa}`]. A warning is issued if the minimum
102
+ value of ``sp`` is less than 30320.0 Pa. Such low values are unrealistic.
103
+ model_levels : npt.ArrayLike
104
+ Target model levels. Expected to be a one-dimensional array of integers between 1 and 137.
105
+
106
+ Returns
107
+ -------
108
+ xr.DataArray
109
+ Pressure levels at each model level, [:math:`hPa`]. The shape of the output is
110
+ the product of the shape of the input and the length of ``model_levels``. In
111
+ other words, the output will have dimensions of the input plus a new dimension
112
+ for ``model_levels``.
113
+
114
+ If ``sp`` is not dask-backed, the output will be computed eagerly. In particular,
115
+ if ``sp`` has a large size and ``model_levels`` is a large range, this function
116
+ may consume a large amount of memory.
117
+
118
+ The ``dtype`` of the output is the same as the ``dtype`` of the ``sp`` parameter.
119
+
120
+ Examples
121
+ --------
122
+ >>> import numpy as np
123
+ >>> import xarray as xr
124
+
125
+ >>> sp_arr = np.linspace(101325.0, 90000.0, 16).reshape(4, 4)
126
+ >>> longitude = np.linspace(-180, 180, 4)
127
+ >>> latitude = np.linspace(-90, 90, 4)
128
+ >>> sp = xr.DataArray(sp_arr, coords={"longitude": longitude, "latitude": latitude})
129
+
130
+ >>> model_levels = [80, 100]
131
+ >>> model_level_pressure(sp, model_levels)
132
+ <xarray.DataArray (model_level: 2, longitude: 4, latitude: 4)> Size: 256B
133
+ array([[[259.75493944, 259.27107504, 258.78721064, 258.30334624],
134
+ [257.81948184, 257.33561744, 256.85175304, 256.36788864],
135
+ [255.88402424, 255.40015984, 254.91629544, 254.43243104],
136
+ [253.94856664, 253.46470224, 252.98083784, 252.49697344]],
137
+ [[589.67975444, 586.47283154, 583.26590864, 580.05898574],
138
+ [576.85206284, 573.64513994, 570.43821704, 567.23129414],
139
+ [564.02437124, 560.81744834, 557.61052544, 554.40360254],
140
+ [551.19667964, 547.98975674, 544.78283384, 541.57591094]]])
141
+ Coordinates:
142
+ * longitude (longitude) float64 32B -180.0 -60.0 60.0 180.0
143
+ * latitude (latitude) float64 32B -90.0 -30.0 30.0 90.0
144
+ * model_level (model_level) int64 16B 80 100
145
+
146
+ See Also
147
+ --------
148
+ model_level_reference_pressure
149
+ """
150
+ # When sp is too low, the pressure up the vertical column will not monotonically decreasing.
151
+ # The first example of this occurs when sp is close to 30320.0 Pa between model
152
+ # levels 114 and 115. Issue a warning here to alert the user.
153
+ if (sp < 30320.0).any():
154
+ msg = (
155
+ "The 'sp' parameter appears to be low. The calculated pressure levels will "
156
+ "not be monotonically decreasing. The 'sp' parameter has units of Pa. "
157
+ "Most surface pressure data should be in the range of 50000.0 to 105000.0 Pa."
158
+ )
159
+ warnings.warn(msg)
160
+
161
+ model_levels = np.asarray(model_levels, dtype=int)
162
+ if not np.all((model_levels >= 1) & (model_levels <= 137)):
163
+ msg = "model_levels must be integers between 1 and 137"
164
+ raise ValueError(msg)
165
+
166
+ usecols = ["n", "a [Pa]", "b"]
167
+ df = (
168
+ pd.read_csv(MODEL_LEVELS_PATH, usecols=usecols)
169
+ .rename(columns={"n": "model_level", "a [Pa]": "a"})
170
+ .set_index("model_level")
171
+ )
172
+
173
+ a = df["a"].to_xarray()
174
+ b = df["b"].to_xarray()
175
+
176
+ if "model_level" in sp.dims:
177
+ sp_model_levels = sp["model_level"]
178
+ if len(sp_model_levels) != 1:
179
+ msg = "Found multiple model levels in sp, expected at most one"
180
+ raise ValueError(msg)
181
+ if sp_model_levels.item() != 1:
182
+ msg = f"sp must be at model level 1, found model level {sp_model_levels.item()}"
183
+ raise ValueError(msg)
184
+ # Remove the model level dimension to allow automatic broadcasting below
185
+ sp = sp.squeeze("model_level")
186
+
187
+ dtype = sp.dtype
188
+ a = a.astype(dtype, copy=False)
189
+ b = b.astype(dtype, copy=False)
190
+
191
+ indexer = {"model_level": model_levels}
192
+ p_half_below = a.sel(indexer) + b.sel(indexer) * sp
193
+
194
+ indexer = {"model_level": model_levels - 1}
195
+ p_half_above = (a.sel(indexer) + b.sel(indexer) * sp).assign_coords(model_level=model_levels)
196
+
197
+ p_full = (p_half_above + p_half_below) / 2.0
198
+ return p_full / 100.0 # Pa -> hPa
199
+
200
+
201
+ def searchsorted2d(
202
+ a: npt.NDArray[np.floating],
203
+ v: npt.NDArray[np.floating],
204
+ ) -> npt.NDArray[np.int64]:
205
+ """Return the indices where elements in ``v`` would be inserted in ``a`` along its second axis.
206
+
207
+ Implementation based on a `StackOverflow answer <https://stackoverflow.com/a/40588862>`_.
208
+
209
+ Parameters
210
+ ----------
211
+ a : npt.NDArray[np.floating]
212
+ 2D array of shape ``(m, n)`` that is sorted along its second axis. This is not checked.
213
+ v : npt.NDArray[np.floating]
214
+ 1D array of values of shape ``(k,)`` to insert into the second axis of ``a``.
215
+ The current implementation could be extended to handle 2D arrays as well.
216
+
217
+ Returns
218
+ -------
219
+ npt.NDArray[np.int64]
220
+ 2D array of indices where elements in ``v`` would be inserted in ``a`` along its
221
+ second axis to keep the second axis of ``a`` sorted. The shape of the output is ``(m, k)``.
222
+
223
+ Examples
224
+ --------
225
+ >>> a = np.array([
226
+ ... [ 1., 8., 11., 12.],
227
+ ... [ 5., 8., 9., 14.],
228
+ ... [ 4., 5., 6., 17.],
229
+ ... ])
230
+ >>> v = np.array([3., 7., 10., 13., 15.])
231
+ >>> searchsorted2d(a, v)
232
+ array([[1, 1, 2, 4, 4],
233
+ [0, 1, 3, 3, 4],
234
+ [0, 3, 3, 3, 3]])
235
+ """
236
+ if a.ndim != 2:
237
+ msg = "The parameter 'a' must be a 2D array"
238
+ raise ValueError(msg)
239
+ if v.ndim != 1:
240
+ msg = "The parameter 'v' must be a 1D array"
241
+ raise ValueError(msg)
242
+
243
+ m, n = a.shape
244
+
245
+ offset_scalar = max(np.ptp(a).item(), np.ptp(v).item()) + 1.0
246
+
247
+ # IMPORTANT: Keep the dtype as float64 to avoid round-off error
248
+ # when computing a_scaled and v_scaled
249
+ # If we used float32 here, the searchsorted output below can be off by 1
250
+ # or 2 if offset_scalar is large and m is large
251
+ steps = np.arange(m, dtype=np.float64).reshape(-1, 1)
252
+ offset = steps * offset_scalar
253
+ a_scaled = a + offset # float32 + float64 = float64
254
+ v_scaled = v + offset # float32 + float64 = float64
255
+
256
+ idx_scaled = np.searchsorted(a_scaled.reshape(-1), v_scaled.reshape(-1)).reshape(v_scaled.shape)
257
+ return idx_scaled - n * steps.astype(np.int64)
258
+
259
+
260
+ def _interp_artifacts(
261
+ xp: npt.NDArray[np.floating], x: npt.NDArray[np.floating]
262
+ ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], npt.NDArray[np.bool_]]:
263
+ """Compute the indices and distances for linear interpolation."""
264
+ idx = searchsorted2d(xp, x)
265
+ out_of_bounds = (idx == 0) | (idx == xp.shape[1])
266
+ idx.clip(1, xp.shape[1] - 1, out=idx)
267
+
268
+ x0 = np.take_along_axis(xp, idx - 1, axis=1)
269
+ x1 = np.take_along_axis(xp, idx, axis=1)
270
+ dist = (x.reshape(1, -1) - x0) / (x1 - x0)
271
+
272
+ return idx, dist, out_of_bounds
273
+
274
+
275
+ def _interp_on_chunk(ds_chunk: xr.Dataset, target_pl: npt.NDArray[np.floating]) -> xr.Dataset:
276
+ """Interpolate the data on a chunk to the target pressure levels.
277
+
278
+ Parameters
279
+ ----------
280
+ ds_chunk : xr.Dataset
281
+ Chunk of the dataset. The last dimension must be "model_level".
282
+ The dataset from which ``ds_chunk`` is taken must not split the
283
+ "model_level" dimension across chunks.
284
+ target_pl : npt.NDArray[np.floating]
285
+ Target pressure levels, [:math:`hPa`].
286
+
287
+ Returns
288
+ -------
289
+ xr.Dataset
290
+ Interpolated data on the target pressure levels. This has the same
291
+ dimensions as ``ds_chunk`` except that the "model_level" dimension
292
+ is replaced with "level". The shape of the "level" dimension is
293
+ the length of ``target_pl``.
294
+ """
295
+ if any(da_chunk.dims[-1] != "model_level" for da_chunk in ds_chunk.values()):
296
+ msg = "The last dimension of the dataset must be 'model_level'"
297
+ raise ValueError(msg)
298
+
299
+ pl_chunk = ds_chunk["pressure_level"]
300
+
301
+ # Put the model_level column in the second dimension
302
+ # And stack the horizontal dimensions into the first dimension
303
+ xp = pl_chunk.values.reshape(-1, len(pl_chunk["model_level"]))
304
+
305
+ # AFAICT, metview performs linear interpolation in xp and target_pl by default
306
+ # However, the conversion_from_ml_to_pl.py script in https://confluence.ecmwf.int/x/JJh0CQ
307
+ # suggests interpolating in the log space. If using consecutive model levels,
308
+ # the difference between the two methods is negligible. We use the log space
309
+ # method here for consistency with the ECMWF script. This only changes
310
+ # the `dist` calculation below.
311
+ idx, dist, out_of_bounds = _interp_artifacts(np.log(xp), np.log(target_pl))
312
+
313
+ shape4d = pl_chunk.shape[:-1] + target_pl.shape
314
+ idx = idx.reshape(shape4d)
315
+ dist = dist.reshape(shape4d)
316
+ out_of_bounds = out_of_bounds.reshape(shape4d)
317
+
318
+ interped_dict = {}
319
+
320
+ for name, da in ds_chunk.items():
321
+ if name == "pressure_level":
322
+ continue
323
+
324
+ fp = da.values
325
+ f0 = np.take_along_axis(fp, idx - 1, axis=-1)
326
+ f1 = np.take_along_axis(fp, idx, axis=-1)
327
+ interped = f0 + dist * (f1 - f0)
328
+ interped[out_of_bounds] = np.nan # we could extrapolate here like RGI(..., fill_value=None)
329
+
330
+ coords = {k: da.coords[k] for k in da.dims[:-1]}
331
+ coords["level"] = target_pl
332
+
333
+ interped_dict[name] = xr.DataArray(
334
+ interped,
335
+ dims=tuple(coords),
336
+ coords=coords,
337
+ attrs=da.attrs,
338
+ )
339
+
340
+ return xr.Dataset(interped_dict)
341
+
342
+
343
+ def _build_template(ds: xr.Dataset, target_pl: npt.NDArray[np.floating]) -> xr.Dataset:
344
+ """Build the template dataset for the interpolated data."""
345
+ coords = {k: ds.coords[k] for k in ds.dims if k != "model_level"} | {"level": target_pl}
346
+
347
+ dims = tuple(coords)
348
+ shape = tuple(len(v) for v in coords.values())
349
+
350
+ vars = {
351
+ k: (dims, dask.array.empty(shape=shape, dtype=da.dtype))
352
+ for k, da in ds.items()
353
+ if k != "pressure_level"
354
+ }
355
+
356
+ chunks = {k: v for k, v in ds.chunks.items() if k != "model_level"}
357
+ chunks["level"] = (len(target_pl),)
358
+
359
+ return xr.Dataset(data_vars=vars, coords=coords, attrs=ds.attrs).chunk(chunks)
360
+
361
+
362
+ def ml_to_pl(
363
+ ds: xr.Dataset,
364
+ target_pl: npt.ArrayLike,
365
+ *,
366
+ lnsp: xr.DataArray | None = None,
367
+ sp: xr.DataArray | None = None,
368
+ ) -> xr.Dataset:
369
+ r"""Interpolate L137 model-level meteorology data to pressure levels.
370
+
371
+ The implementation is here is consistent with ECMWF's
372
+ `suggested implementation <https://confluence.ecmwf.int/x/JJh0CQ#heading-Step2Interpolatevariablesonmodellevelstocustompressurelevels>`_.
373
+
374
+ Parameters
375
+ ----------
376
+ ds : xr.Dataset
377
+ Dataset with model-level meteorology data. Must include a "model_level" dimension
378
+ which is not split across chunks. The non-"model_level" dimensions must be
379
+ aligned with the "lnsp" parameter. Can include any number of variables.
380
+ Any `non-dimension coordinates <https://docs.xarray.dev/en/latest/user-guide/terminology.html#term-Non-dimension-coordinate>`_
381
+ will be dropped.
382
+ target_pl : npt.ArrayLike
383
+ Target pressure levels, [:math:`hPa`].
384
+ lnsp : xr.DataArray
385
+ Natural logarithm of surface pressure, [:math:`\ln(\text{Pa})`]. If provided,
386
+ ``sp`` is ignored. At least one of ``lnsp`` or ``sp`` must be provided.
387
+ The chunking over dimensions in common with ``ds`` must be the same as ``ds``.
388
+ sp : xr.DataArray
389
+ Surface pressure, [:math:`\text{Pa}`]. At least one of ``lnsp`` or ``sp`` must be provided.
390
+ The chunking over dimensions in common with ``ds`` must be the same as ``ds``.
391
+
392
+ Returns
393
+ -------
394
+ xr.Dataset
395
+ Interpolated data on the target pressure levels. This has the same
396
+ dimensions as ``ds`` except that the "model_level" dimension
397
+ is replaced with "level". The shape of the "level" dimension is
398
+ the length of ``target_pl``. If ``ds`` is dask-backed, the output
399
+ will be as well. Call ``.compute()`` to compute the result eagerly.
400
+ """
401
+ if sp is None:
402
+ if lnsp is None:
403
+ msg = "At least one of 'lnsp' or 'sp' must be provided"
404
+ raise ValueError(msg)
405
+ sp = dask.array.exp(lnsp)
406
+
407
+ model_levels = ds["model_level"]
408
+ pl = model_level_pressure(sp, model_levels)
409
+
410
+ if "pressure_level" in ds:
411
+ msg = "The dataset must not contain a 'pressure_level' variable"
412
+ raise ValueError(msg)
413
+ ds = ds.assign(pressure_level=pl)
414
+
415
+ ds = ds.reset_coords(drop=True) # drop "expver"
416
+
417
+ # If there are any variables which do not have the "model_level" dimension,
418
+ # issue a warning and drop them
419
+ for name, da in ds.items():
420
+ if "model_level" not in da.dims:
421
+ msg = f"Variable '{name}' does not have a 'model_level' dimension"
422
+ warnings.warn(msg)
423
+ ds = ds.drop_vars([name])
424
+
425
+ # IMPORTANT: model_level must be the last dimension for _interp_on_chunk
426
+ ds = ds.transpose(..., "model_level")
427
+
428
+ # Raise if chunks over model level
429
+ if ds.chunks and len(ds.chunks["model_level"]) > 1:
430
+ msg = "The 'model_level' dimension must not be split across chunks"
431
+ raise ValueError(msg)
432
+
433
+ target_pl = np.asarray(target_pl, dtype=sp.dtype)
434
+ template = _build_template(ds, target_pl)
435
+ return xr.map_blocks(_interp_on_chunk, ds, (target_pl,), template=template)
@@ -0,0 +1,139 @@
1
+ n,a [Pa],b,ph [hPa],pf [hPa],Geopotential Altitude [m],Geometric Altitude [m],Temperature [K],Density [kg/m^3]
2
+ 0,0.0,0.0,0.0,,,,,
3
+ 1,2.000365,0.0,0.02,0.01,79301.79,80301.65,198.05,1.8e-05
4
+ 2,3.102241,0.0,0.031,0.0255,73721.58,74584.91,209.21,4.2e-05
5
+ 3,4.666084,0.0,0.0467,0.0388,71115.75,71918.79,214.42,6.3e-05
6
+ 4,6.827977,0.0,0.0683,0.0575,68618.43,69365.77,221.32,9e-05
7
+ 5,9.746966,0.0,0.0975,0.0829,66210.99,66906.53,228.06,0.000127
8
+ 6,13.605424,0.0,0.1361,0.1168,63890.03,64537.43,234.56,0.000173
9
+ 7,18.608931,0.0,0.1861,0.1611,61651.77,62254.39,240.83,0.000233
10
+ 8,24.985718,0.0,0.2499,0.218,59492.5,60053.46,246.87,0.000308
11
+ 9,32.98571,0.0,0.3299,0.2899,57408.61,57930.78,252.71,0.0004
12
+ 10,42.879242,0.0,0.4288,0.3793,55396.62,55882.68,258.34,0.000512
13
+ 11,54.955463,0.0,0.5496,0.4892,53453.2,53905.62,263.78,0.000646
14
+ 12,69.520576,0.0,0.6952,0.6224,51575.15,51996.21,269.04,0.000806
15
+ 13,86.895882,0.0,0.869,0.7821,49767.41,50159.36,270.65,0.001007
16
+ 14,107.415741,0.0,1.0742,0.9716,48048.7,48413.94,270.65,0.001251
17
+ 15,131.425507,0.0,1.3143,1.1942,46416.22,46756.98,269.02,0.001546
18
+ 16,159.279404,0.0,1.5928,1.4535,44881.17,45199.69,264.72,0.001913
19
+ 17,191.338562,0.0,1.9134,1.7531,43440.23,43738.55,260.68,0.002343
20
+ 18,227.968948,0.0,2.2797,2.0965,42085.0,42364.93,256.89,0.002843
21
+ 19,269.539581,0.0,2.6954,2.4875,40808.05,41071.2,253.31,0.003421
22
+ 20,316.420746,0.0,3.1642,2.9298,39602.76,39850.56,249.94,0.004084
23
+ 21,368.982361,0.0,3.6898,3.427,38463.25,38696.94,246.75,0.004838
24
+ 22,427.592499,0.0,4.2759,3.9829,37384.22,37604.95,243.73,0.005693
25
+ 23,492.616028,0.0,4.9262,4.601,36360.94,36569.72,240.86,0.006655
26
+ 24,564.413452,0.0,5.6441,5.2851,35389.15,35586.89,238.14,0.007731
27
+ 25,643.339905,0.0,6.4334,6.0388,34465.0,34652.52,235.55,0.008931
28
+ 26,729.744141,0.0,7.2974,6.8654,33585.02,33763.05,233.09,0.010261
29
+ 27,823.967834,0.0,8.2397,7.7686,32746.04,32915.27,230.74,0.011729
30
+ 28,926.34491,0.0,9.2634,8.7516,31945.53,32106.57,228.6,0.013337
31
+ 29,1037.201172,0.0,10.372,9.8177,31177.59,31330.96,227.83,0.015012
32
+ 30,1156.853638,0.0,11.5685,10.9703,30438.54,30584.71,227.09,0.016829
33
+ 31,1285.610352,0.0,12.8561,12.2123,29726.69,29866.09,226.38,0.018793
34
+ 32,1423.770142,0.0,14.2377,13.5469,29040.48,29173.5,225.69,0.02091
35
+ 33,1571.622925,0.0,15.7162,14.977,28378.46,28505.47,225.03,0.023186
36
+ 34,1729.448975,0.0,17.2945,16.5054,27739.29,27860.64,224.39,0.025624
37
+ 35,1897.519287,0.0,18.9752,18.1348,27121.74,27237.73,223.77,0.028232
38
+ 36,2076.095947,0.0,20.761,19.8681,26524.63,26635.56,223.17,0.031013
39
+ 37,2265.431641,0.0,22.6543,21.7076,25946.9,26053.04,222.6,0.033972
40
+ 38,2465.770508,0.0,24.6577,23.656,25387.55,25489.15,222.04,0.037115
41
+ 39,2677.348145,0.0,26.7735,25.7156,24845.63,24942.93,221.5,0.040445
42
+ 40,2900.391357,0.0,29.0039,27.8887,24320.28,24413.5,220.97,0.043967
43
+ 41,3135.119385,0.0,31.3512,30.1776,23810.67,23900.02,220.46,0.047685
44
+ 42,3381.743652,0.0,33.8174,32.5843,23316.04,23401.71,219.97,0.051604
45
+ 43,3640.468262,0.0,36.4047,35.1111,22835.68,22917.85,219.49,0.055727
46
+ 44,3911.490479,0.0,39.1149,37.7598,22368.91,22447.75,219.02,0.060059
47
+ 45,4194.930664,0.0,41.9493,40.5321,21915.16,21990.82,218.57,0.064602
48
+ 46,4490.817383,0.0,44.9082,43.4287,21473.98,21546.62,218.12,0.069359
49
+ 47,4799.149414,0.0,47.9915,46.4498,21045.0,21114.77,217.7,0.07433
50
+ 48,5119.89502,0.0,51.199,49.5952,20627.87,20694.9,217.28,0.079516
51
+ 49,5452.990723,0.0,54.5299,52.8644,20222.24,20286.66,216.87,0.084916
52
+ 50,5798.344727,0.0,57.9834,56.2567,19827.95,19889.88,216.65,0.090458
53
+ 51,6156.074219,0.0,61.5607,59.7721,19443.55,19503.09,216.65,0.09611
54
+ 52,6526.946777,0.0,65.2695,63.4151,19068.35,19125.61,216.65,0.101968
55
+ 53,6911.870605,0.0,69.1187,67.1941,18701.27,18756.34,216.65,0.108045
56
+ 54,7311.869141,0.0,73.1187,71.1187,18341.27,18394.25,216.65,0.114355
57
+ 55,7727.412109,7e-06,77.281,75.1999,17987.41,18038.35,216.65,0.120917
58
+ 56,8159.354004,2.4e-05,81.6182,79.4496,17638.78,17687.77,216.65,0.127751
59
+ 57,8608.525391,5.9e-05,86.145,83.8816,17294.53,17341.62,216.65,0.134877
60
+ 58,9076.400391,0.000112,90.8774,88.5112,16953.83,16999.08,216.65,0.142321
61
+ 59,9562.682617,0.000199,95.828,93.3527,16616.09,16659.55,216.65,0.150106
62
+ 60,10065.978516,0.00034,101.0047,98.4164,16281.1,16322.83,216.65,0.158248
63
+ 61,10584.631836,0.000562,106.4153,103.71,15948.85,15988.88,216.65,0.16676
64
+ 62,11116.662109,0.00089,112.0681,109.2417,15619.3,15657.7,216.65,0.175655
65
+ 63,11660.067383,0.001353,117.9714,115.0198,15292.44,15329.24,216.65,0.184946
66
+ 64,12211.547852,0.001992,124.1337,121.0526,14968.24,15003.5,216.65,0.194646
67
+ 65,12766.873047,0.002857,130.5637,127.3487,14646.68,14680.44,216.65,0.20477
68
+ 66,13324.668945,0.003971,137.2703,133.917,14327.75,14360.05,216.65,0.215331
69
+ 67,13881.331055,0.005378,144.2624,140.7663,14011.41,14042.3,216.65,0.226345
70
+ 68,14432.139648,0.007133,151.5493,147.9058,13697.65,13727.18,216.65,0.237825
71
+ 69,14975.615234,0.009261,159.1403,155.3448,13386.45,13414.65,216.65,0.249786
72
+ 70,15508.256836,0.011806,167.045,163.0927,13077.79,13104.7,216.65,0.262244
73
+ 71,16026.115234,0.014816,175.2731,171.1591,12771.64,12797.3,216.65,0.275215
74
+ 72,16527.322266,0.018318,183.8344,179.5537,12467.99,12492.44,216.65,0.288713
75
+ 73,17008.789063,0.022355,192.7389,188.2867,12166.81,12190.1,216.65,0.302755
76
+ 74,17467.613281,0.026964,201.9969,197.3679,11868.08,11890.24,216.65,0.317357
77
+ 75,17901.621094,0.032176,211.6186,206.8078,11571.79,11592.86,216.65,0.332536
78
+ 76,18308.433594,0.038026,221.6146,216.6166,11277.92,11297.93,216.65,0.348308
79
+ 77,18685.71875,0.044548,231.9954,226.805,10986.7,11005.69,216.74,0.364545
80
+ 78,19031.289063,0.051773,242.7719,237.3837,10696.22,10714.22,218.62,0.378253
81
+ 79,19343.511719,0.059728,253.9549,248.3634,10405.61,10422.64,220.51,0.392358
82
+ 80,19620.042969,0.068448,265.5556,259.7553,10114.89,10130.98,222.4,0.406868
83
+ 81,19859.390625,0.077958,277.5852,271.5704,9824.08,9839.26,224.29,0.42179
84
+ 82,20059.931641,0.088286,290.0548,283.82,9533.2,9547.49,226.18,0.43713
85
+ 83,20219.664063,0.099462,302.9762,296.5155,9242.26,9255.7,228.08,0.452897
86
+ 84,20337.863281,0.111505,316.3607,309.6684,8951.3,8963.9,229.97,0.469097
87
+ 85,20412.308594,0.124448,330.2202,323.2904,8660.32,8672.11,231.86,0.485737
88
+ 86,20442.078125,0.138313,344.5663,337.3932,8369.35,8380.36,233.75,0.502825
89
+ 87,20425.71875,0.153125,359.4111,351.9887,8078.41,8088.67,235.64,0.520367
90
+ 88,20361.816406,0.16891,374.7666,367.0889,7787.51,7797.04,237.53,0.53837
91
+ 89,20249.511719,0.185689,390.645,382.7058,7496.68,7505.51,239.42,0.556842
92
+ 90,20087.085938,0.203491,407.0583,398.8516,7205.93,7214.09,241.31,0.57579
93
+ 91,19874.025391,0.222333,424.019,415.5387,6915.29,6922.8,243.2,0.595219
94
+ 92,19608.572266,0.242244,441.5395,432.7792,6624.76,6631.66,245.09,0.615138
95
+ 93,19290.226563,0.263242,459.6321,450.5858,6334.38,6340.68,246.98,0.635553
96
+ 94,18917.460938,0.285354,478.3096,468.9708,6044.15,6049.89,248.86,0.656471
97
+ 95,18489.707031,0.308598,497.5845,487.947,5754.1,5759.3,250.75,0.677899
98
+ 96,18006.925781,0.332939,517.4198,507.5021,5464.6,5469.3,252.63,0.699815
99
+ 97,17471.839844,0.358254,537.7195,527.5696,5176.77,5180.98,254.5,0.722139
100
+ 98,16888.6875,0.384363,558.343,548.0312,4892.26,4896.02,256.35,0.744735
101
+ 99,16262.046875,0.411125,579.1926,568.7678,4612.58,4615.92,258.17,0.767472
102
+ 100,15596.695313,0.438391,600.1668,589.6797,4338.77,4341.73,259.95,0.790242
103
+ 101,14898.453125,0.466003,621.1624,610.6646,4071.8,4074.41,261.68,0.812937
104
+ 102,14173.324219,0.4938,642.0764,631.6194,3812.53,3814.82,263.37,0.835453
105
+ 103,13427.769531,0.521619,662.8084,652.4424,3561.7,3563.69,265.0,0.857686
106
+ 104,12668.257813,0.549301,683.262,673.0352,3319.94,3321.67,266.57,0.879541
107
+ 105,11901.339844,0.576692,703.3467,693.3043,3087.75,3089.25,268.08,0.900929
108
+ 106,11133.304688,0.603648,722.9795,713.1631,2865.54,2866.83,269.52,0.921768
109
+ 107,10370.175781,0.630036,742.0855,732.5325,2653.58,2654.69,270.9,0.941988
110
+ 108,9617.515625,0.655736,760.5996,751.3426,2452.04,2452.99,272.21,0.961527
111
+ 109,8880.453125,0.680643,778.4661,769.5329,2260.99,2261.8,273.45,0.980334
112
+ 110,8163.375,0.704669,795.6396,787.0528,2080.41,2081.09,274.63,0.998368
113
+ 111,7470.34375,0.727739,812.0847,803.8622,1910.19,1910.76,275.73,1.015598
114
+ 112,6804.421875,0.749797,827.7756,819.9302,1750.14,1750.63,276.77,1.032005
115
+ 113,6168.53125,0.770798,842.6959,835.2358,1600.04,1600.44,277.75,1.047576
116
+ 114,5564.382813,0.790717,856.8376,849.7668,1459.58,1459.91,278.66,1.06231
117
+ 115,4993.796875,0.809536,870.2004,863.519,1328.43,1328.7,279.52,1.076209
118
+ 116,4457.375,0.827256,882.791,876.4957,1206.21,1206.44,280.31,1.089286
119
+ 117,3955.960938,0.843881,894.6222,888.7066,1092.54,1092.73,281.05,1.101558
120
+ 118,3489.234375,0.859432,905.7116,900.1669,987.0,987.15,281.73,1.113047
121
+ 119,3057.265625,0.873929,916.0815,910.8965,889.17,889.29,282.37,1.123777
122
+ 120,2659.140625,0.887408,925.7571,920.9193,798.62,798.72,282.96,1.133779
123
+ 121,2294.242188,0.8999,934.7666,930.2618,714.94,715.02,283.5,1.143084
124
+ 122,1961.5,0.911448,943.1399,938.9532,637.7,637.76,284.0,1.151724
125
+ 123,1659.476563,0.922096,950.9082,947.024,566.49,566.54,284.47,1.159733
126
+ 124,1387.546875,0.931881,958.1037,954.5059,500.91,500.95,284.89,1.167147
127
+ 125,1143.25,0.94086,964.7584,961.4311,440.58,440.61,285.29,1.173999
128
+ 126,926.507813,0.949064,970.9046,967.8315,385.14,385.16,285.65,1.180323
129
+ 127,734.992188,0.95655,976.5737,973.7392,334.22,334.24,285.98,1.186154
130
+ 128,568.0625,0.963352,981.7968,979.1852,287.51,287.52,286.28,1.191523
131
+ 129,424.414063,0.969513,986.6036,984.2002,244.68,244.69,286.56,1.196462
132
+ 130,302.476563,0.975078,991.023,988.8133,205.44,205.44,286.81,1.201001
133
+ 131,202.484375,0.980072,995.0824,993.0527,169.5,169.51,287.05,1.205168
134
+ 132,122.101563,0.984542,998.8081,996.9452,136.62,136.62,287.26,1.208992
135
+ 133,62.78125,0.9885,1002.225,1000.5165,106.54,106.54,287.46,1.212498
136
+ 134,22.835938,0.991984,1005.3562,1003.7906,79.04,79.04,287.64,1.21571
137
+ 135,3.757813,0.995003,1008.2239,1006.79,53.92,53.92,287.8,1.21865
138
+ 136,0.0,0.99763,1010.8487,1009.5363,30.96,30.96,287.95,1.221341
139
+ 137,0.0,1.0,1013.25,1012.0494,10.0,10.0,288.09,1.223803