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,2616 @@
1
+ """Top level CoCiP classes and methods."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import logging
6
+ import warnings
7
+ from collections.abc import Sequence
8
+ from typing import Any, Literal, NoReturn, overload
9
+
10
+ import numpy as np
11
+ import numpy.typing as npt
12
+ import pandas as pd
13
+ import xarray as xr
14
+ from overrides import overrides
15
+
16
+ from pycontrails.core import met_var
17
+ from pycontrails.core.aircraft_performance import AircraftPerformance
18
+ from pycontrails.core.fleet import Fleet
19
+ from pycontrails.core.flight import Flight
20
+ from pycontrails.core.met import MetDataset
21
+ from pycontrails.core.models import Model, interpolate_met
22
+ from pycontrails.core.vector import GeoVectorDataset, VectorDataDict
23
+ from pycontrails.datalib import ecmwf, gfs
24
+ from pycontrails.models import sac, tau_cirrus
25
+ from pycontrails.models.cocip import (
26
+ contrail_properties,
27
+ radiative_forcing,
28
+ radiative_heating,
29
+ unterstrasser_wake_vortex,
30
+ wake_vortex,
31
+ wind_shear,
32
+ )
33
+ from pycontrails.models.cocip.cocip_params import CocipFlightParams
34
+ from pycontrails.models.emissions.emissions import Emissions
35
+ from pycontrails.physics import constants, geo, thermo, units
36
+
37
+ logger = logging.getLogger(__name__)
38
+
39
+
40
+ class Cocip(Model):
41
+ r"""Contrail Cirrus Prediction Model (CoCiP).
42
+
43
+ Published by Ulrich Schumann *et. al.*
44
+ (`DLR Institute of Atmospheric Physics <https://www.dlr.de/pa/en/>`_)
45
+ in :cite:`schumannContrailCirrusPrediction2012`, :cite:`schumannParametricRadiativeForcing2012`.
46
+
47
+ Parameters
48
+ ----------
49
+ met : MetDataset
50
+ Pressure level dataset containing :attr:`met_variables` variables.
51
+ See *Notes* for variable names by data source.
52
+ rad : MetDataset
53
+ Single level dataset containing top of atmosphere radiation fluxes.
54
+ See *Notes* for variable names by data source.
55
+ params : dict[str, Any], optional
56
+ Override Cocip model parameters with dictionary.
57
+ See :class:`CocipFlightParams` for model parameters.
58
+ **params_kwargs : Any
59
+ Override Cocip model parameters with keyword arguments.
60
+ See :class:`CocipFlightParams` for model parameters.
61
+
62
+ Notes
63
+ -----
64
+ **Inputs**
65
+
66
+ The required meteorology variables depend on the data source (e.g. ECMWF, GFS).
67
+
68
+ See :attr:`met_variables` and :attr:`rad_variables` for the list of required variables
69
+ to the ``met`` and ``rad`` parameters, respectively.
70
+ When an item in one of these arrays is a :class:`tuple`, variable keys depend on data source.
71
+
72
+ The current list of required variables (labelled by ``"standard_name"``):
73
+
74
+ .. list-table:: Variable keys for pressure level data
75
+ :header-rows: 1
76
+
77
+ * - Parameter
78
+ - ECMWF
79
+ - GFS
80
+ * - Air Temperature
81
+ - ``air_temperature``
82
+ - ``air_temperature``
83
+ * - Specific Humidity
84
+ - ``specific_humidity``
85
+ - ``specific_humidity``
86
+ * - Eastward wind
87
+ - ``eastward_wind``
88
+ - ``eastward_wind``
89
+ * - Northward wind
90
+ - ``northward_wind``
91
+ - ``northward_wind``
92
+ * - Vertical velocity
93
+ - ``lagrangian_tendency_of_air_pressure``
94
+ - ``lagrangian_tendency_of_air_pressure``
95
+ * - Ice water content
96
+ - ``specific_cloud_ice_water_content``
97
+ - ``ice_water_mixing_ratio``
98
+
99
+ .. list-table:: Variable keys for single-level radiation data
100
+ :header-rows: 1
101
+
102
+ * - Parameter
103
+ - ECMWF
104
+ - GFS
105
+ * - Top solar radiation
106
+ - ``top_net_solar_radiation``
107
+ - ``toa_upward_shortwave_flux``
108
+ * - Top thermal radiation
109
+ - ``top_net_thermal_radiation``
110
+ - ``toa_upward_longwave_flux``
111
+
112
+ **Modifications**
113
+
114
+ This implementation differs from original CoCiP (Fortran) implementation in a few places:
115
+
116
+ - This model uses aircraft performance and emissions models to calculate nvPM, fuel flow,
117
+ and overall propulsion efficiency, if not already provided.
118
+ - As described in :cite:`teohAviationContrailClimate2022`, this implementation sets
119
+ the initial ice particle activation rate to be a function of
120
+ the difference between the ambient temperature and the critical SAC threshold temperature.
121
+ See :func:`pycontrails.models.sac.T_critical_sac`.
122
+ - Isobaric heat capacity calculation.
123
+ The original model uses a constant value of 1004 :math:`J \ kg^{-1} \ K^{-1}`,
124
+ whereas this model calculates isobaric heat capacity as a function of specific humidity.
125
+ See :func:`pycontrails.physics.thermo.c_pm`.
126
+ - Solar direct radiation.
127
+ The original algorithm uses ECMWF radiation variable `tisr` (top incident solar radiation)
128
+ as solar direct radiation value.
129
+ This implementation calculates the theoretical solar direct radiation at
130
+ any arbitrary point in the atmosphere.
131
+ See :func:`pycontrails.physics.geo.solar_direct_radiation`.
132
+ - Segment angle.
133
+ The segment angle calculations for flights and contrail segments have been updated
134
+ to use more precise spherical geometry instead of a triangular approximation.
135
+ As the triangle approaches zero, the two calculations agree.
136
+ See :func:`pycontrails.physics.geo.segment_angle`.
137
+ - Integration.
138
+ This implementation consistently uses left-Riemann sums
139
+ in the time integration of contrail segments.
140
+ - Segment length ratio.
141
+ Instead of taking a geometric mean between contrail segments before/after advection,
142
+ a simple ratio is computed.
143
+ See :func:`contrail_properties.segment_length_ratio`.
144
+ - Segment energy flux.
145
+ This implementation does not average spatially contiguous contrail segments when calculating
146
+ the mean energy flux for the segment of interest.
147
+ See :func:`contrail_properties.mean_energy_flux_per_m`.
148
+
149
+ This implementation is regression tested against
150
+ results from :cite:`teohAviationContrailClimate2022`.
151
+
152
+ **Outputs**
153
+
154
+ NaN values may appear in model output. Specifically, ``np.nan`` values are used to indicate:
155
+
156
+ - Flight waypoint or contrail waypoint is not contained with the :attr:`met` domain.
157
+ - The variable was NOT computed during the model evaluation. For example, at flight waypoints
158
+ not producing any persistent contrails, "radiative" variables (``rsr``, ``olr``, ``rf_sw``,
159
+ ``rf_lw``, ``rf_net``) are not computed. Consequently, the corresponding values in the output
160
+ of :meth:`eval` are NaN. One exception to this rule is found on ``ef`` (energy forcing)
161
+ `contrail_age` predictions. For these two "cumulative" variables, waypoints not producing
162
+ any persistent contrails are assigned 0 values.
163
+
164
+ References
165
+ ----------
166
+ - :cite:`schumannDeterminationContrailsSatellite1990`
167
+ - :cite:`schumannContrailCirrusPrediction2010`
168
+ - :cite:`voigtInsituObservationsYoung2010`
169
+ - :cite:`schumannPotentialReduceClimate2011`
170
+ - :cite:`schumannContrailCirrusPrediction2012`
171
+ - :cite:`schumannParametricRadiativeForcing2012`
172
+ - :cite:`schumannContrailsVisibleAviation2012`
173
+ - :cite:`schumannEffectiveRadiusIce2011`
174
+ - :cite:`schumannDehydrationEffectsContrails2015`
175
+ - :cite:`teohMitigatingClimateForcing2020`
176
+ - :cite:`schumannAviationContrailCirrus2021`
177
+ - :cite:`schumannAirTrafficContrail2021`
178
+ - :cite:`teohAviationContrailClimate2022`
179
+
180
+ See Also
181
+ --------
182
+ :class:`CocipFlightParams`
183
+ :mod:`wake_vortex`
184
+ :mod:`contrail_properties`
185
+ :mod:`radiative_forcing`
186
+ :mod:`humidity_scaling`
187
+ :class:`Emissions`
188
+ :mod:`sac`
189
+ :mod:`tau_cirrus`
190
+ """
191
+
192
+ __slots__ = (
193
+ "rad",
194
+ "contrail",
195
+ "contrail_dataset",
196
+ "contrail_list",
197
+ "timesteps",
198
+ "_sac_flight",
199
+ "_downwash_flight",
200
+ "_downwash_contrail",
201
+ )
202
+
203
+ name = "cocip"
204
+ long_name = "Contrail Cirrus Prediction Model"
205
+ default_params = CocipFlightParams
206
+ met_variables = (
207
+ met_var.AirTemperature,
208
+ met_var.SpecificHumidity,
209
+ met_var.EastwardWind,
210
+ met_var.NorthwardWind,
211
+ met_var.VerticalVelocity,
212
+ (ecmwf.SpecificCloudIceWaterContent, gfs.CloudIceWaterMixingRatio),
213
+ )
214
+
215
+ #: Required single-level top of atmosphere radiation variables.
216
+ #: Variable keys depend on data source (e.g. ECMWF, GFS).
217
+ rad_variables = (
218
+ (ecmwf.TopNetSolarRadiation, gfs.TOAUpwardShortwaveRadiation),
219
+ (ecmwf.TopNetThermalRadiation, gfs.TOAUpwardLongwaveRadiation),
220
+ )
221
+
222
+ #: Minimal set of met variables needed to run the model after pre-processing.
223
+ #: The intention here is that ``ciwc`` is unnecessary after
224
+ #: ``tau_cirrus`` has already been calculated.
225
+ processed_met_variables = (
226
+ met_var.AirTemperature,
227
+ met_var.SpecificHumidity,
228
+ met_var.EastwardWind,
229
+ met_var.NorthwardWind,
230
+ met_var.VerticalVelocity,
231
+ tau_cirrus.TauCirrus,
232
+ )
233
+
234
+ #: Additional met variables used to support outputs
235
+ #:
236
+ #: .. versionchanged:: 0.48.0
237
+ #: Moved Geopotential from :attr:`met_variables` to :attr:`optional_met_variables`
238
+ optional_met_variables = (
239
+ (met_var.Geopotential, met_var.GeopotentialHeight),
240
+ (ecmwf.CloudAreaFractionInLayer, gfs.TotalCloudCoverIsobaric),
241
+ )
242
+
243
+ #: Met data is not optional
244
+ met: MetDataset
245
+ met_required = True
246
+
247
+ #: Radiation data formatted as a :class:`MetDataset` at a single pressure level [-1]
248
+ rad: MetDataset
249
+
250
+ #: Last Flight modeled in :meth:`eval`
251
+ source: Flight | Fleet
252
+
253
+ #: List of :class:`GeoVectorDataset` contrail objects - one for each timestep
254
+ contrail_list: list[GeoVectorDataset]
255
+
256
+ #: Contrail evolution output from model.
257
+ #:
258
+ #: Set to None when no contrails are formed.
259
+ #: Otherwise, this is a :class:`pandas.DataFrame` describing the evolution of the contrail.
260
+ #: Columns include:
261
+ #:
262
+ #: - ``waypoint``: The index of the waypoint in the original flight creating
263
+ #: the contrail. This can be used to join the contrail DataFrame to the :attr:`source`.
264
+ #: - ``formation_time``: Time of contrail formation. Agrees with the ``time`` column
265
+ #: in :attr:`source`.
266
+ #: - ``continuous``: Boolean indicating whether the contrail is continuous or not.
267
+ #: - ``persistent``: Boolean indicating whether the contrail is persistent or not.
268
+ #: A contrail segment is considered continuous if both the current and the next
269
+ #: contrail waypoint at the same time step persist.
270
+ #: - ``segment_length``: Length of the contrail segment, [:math:`m`].
271
+ #: - ``sin_a``, ``cos_a``: Sine and cosine of the segment angle.
272
+ #: - ``width``, ``depth``: Contrail width and depth, [:math:`m`].
273
+ #: - ``sigma_yz``: The ``yz`` component of the covariance matrix, [:math:`m^{2}`].
274
+ #: See :func:`contrail_properties.plume_temporal_evolution`.
275
+ #: - ``q_sat``: Saturation specific humidity over ice, [:math:`kg \ kg^{-1}`].
276
+ #: - ``n_ice_per_m``: Number of ice particles per distance, [:math:`m^{-1}`].
277
+ #: - ``iwc``: Ice water content, [:math:`kg_{ice} kg_{air}^{-1}`].
278
+ #: - ``tau_contrail``: Optical depth of the contrail. See
279
+ #: :func:`contrail_properties.contrail_optical_depth`.
280
+ #: - ``rf_sw``, ``rf_lw``, ``rf_net``: Shortwave, longwave, and net instantaneous
281
+ #: radiative forcing, [:math:`W \ m^{-2}`] at the contrail waypoint.
282
+ #: - ``ef``: Energy forcing, [:math:`J`] at the contrail waypoint. See
283
+ #: :func:`contrail_properties.energy_forcing`.
284
+ contrail: pd.DataFrame | None
285
+
286
+ #: :class:`xr.Dataset` representation of contrail evolution.
287
+ contrail_dataset: xr.Dataset | None
288
+
289
+ #: Array of :class:`numpy.datetime64` time steps for contrail evolution
290
+ timesteps: npt.NDArray[np.datetime64]
291
+
292
+ #: Parallel copy of flight waypoints after SAC filter applied
293
+ _sac_flight: Flight
294
+
295
+ #: Parallel copy of flight waypoints after wake vortex downwash applied
296
+ _downwash_flight: Flight
297
+
298
+ #: GeoVectorDataset representation of :attr:`downwash_flight`
299
+ _downwash_contrail: GeoVectorDataset
300
+
301
+ def __init__(
302
+ self,
303
+ met: MetDataset,
304
+ rad: MetDataset,
305
+ params: dict[str, Any] | None = None,
306
+ **params_kwargs: Any,
307
+ ) -> None:
308
+ # call Model init
309
+ super().__init__(met, params=params, **params_kwargs)
310
+
311
+ compute_tau_cirrus = self.params["compute_tau_cirrus_in_model_init"]
312
+ self.met, self.rad = process_met_datasets(met, rad, compute_tau_cirrus)
313
+
314
+ # initialize outputs to None
315
+ self.contrail = None
316
+ self.contrail_dataset = None
317
+
318
+ # ----------
319
+ # Public API
320
+ # ----------
321
+
322
+ @overload
323
+ def eval(self, source: Fleet, **params: Any) -> Fleet: ...
324
+
325
+ @overload
326
+ def eval(self, source: Flight, **params: Any) -> Flight: ...
327
+
328
+ @overload
329
+ def eval(self, source: Sequence[Flight], **params: Any) -> list[Flight]: ...
330
+
331
+ @overload
332
+ def eval(self, source: None = ..., **params: Any) -> NoReturn: ...
333
+
334
+ def eval(
335
+ self,
336
+ source: Flight | Sequence[Flight] | None = None,
337
+ **params: Any,
338
+ ) -> Flight | list[Flight]:
339
+ """Run CoCiP simulation on flight.
340
+
341
+ Simulates the formation and evolution of contrails from a Flight
342
+ using the contrail cirrus prediction model (CoCiP) from Schumann (2012)
343
+ :cite:`schumannContrailCirrusPrediction2012`.
344
+
345
+ .. versionchanged:: 0.25.11
346
+
347
+ Previously, any waypoint not surviving the wake vortex downwash phase of CoCiP
348
+ was assigned a nan-value in the ``ef`` array within the model
349
+ output. This is no longer the case. Instead, energy forcing is set to 0.0
350
+ for all waypoints which fail to produce persistent contrails. In particular,
351
+ nan values in the ``ef`` array are only used to indicate an
352
+ out-of-met-domain waypoint. The same convention is now used for output variables
353
+ ``contrail_age`` and ``cocip`` as well.
354
+
355
+ .. versionchanged::0.26.0
356
+
357
+ :attr:`met` and :attr:`rad` down-selection is now handled automatically.
358
+
359
+ Parameters
360
+ ----------
361
+ source : Flight | Sequence[Flight] | None
362
+ Input Flight(s) to model.
363
+ **params : Any
364
+ Overwrite model parameters before eval.
365
+
366
+ Returns
367
+ -------
368
+ Flight | list[Flight] | NoReturn
369
+ Flight(s) with updated Contrail data. The model parameter "verbose_outputs"
370
+ determines the variables on the return flight object.
371
+
372
+ References
373
+ ----------
374
+ - :cite:`schumannContrailCirrusPrediction2012`
375
+ """
376
+
377
+ self.update_params(params)
378
+ self.set_source(source)
379
+ self.source = self.require_source_type(Flight)
380
+ return_flight_list = isinstance(self.source, Fleet) and isinstance(source, Sequence)
381
+
382
+ self._set_timesteps()
383
+
384
+ # Downselect met for CoCiP initialization
385
+ # We only need to buffer in the negative vertical direction,
386
+ # which is the positive direction for level
387
+ logger.debug("Downselect met for Cocip initialization")
388
+ level_buffer = 0, self.params["met_level_buffer"][1]
389
+ met = self.source.downselect_met(self.met, level_buffer=level_buffer, copy=False)
390
+ met = add_tau_cirrus(met)
391
+
392
+ # Prepare flight for model
393
+ self._process_flight(met)
394
+
395
+ # Save humidity scaling type to output attrs
396
+ # NOTE: Do this after _process_flight because that method automatically
397
+ # broadcasts all numeric source params.
398
+ humidity_scaling = self.params["humidity_scaling"]
399
+ if humidity_scaling is not None:
400
+ for k, v in humidity_scaling.description.items():
401
+ self.source.attrs[f"humidity_scaling_{k}"] = v
402
+
403
+ if isinstance(self.source, Fleet):
404
+ label = f"fleet of {self.source.n_flights} flights"
405
+ else:
406
+ label = f"flight {self.source['flight_id'][0]}"
407
+
408
+ logger.debug("Finding initial linear contrails with SAC")
409
+ self._find_initial_contrail_regions()
410
+
411
+ if not self._sac_flight:
412
+ logger.debug("No linear contrails formed by %s", label)
413
+ return self._fill_empty_flight_results(return_flight_list)
414
+
415
+ self._simulate_wake_vortex_downwash(met)
416
+
417
+ self._find_initial_persistent_contrails(met)
418
+
419
+ if not self._downwash_flight:
420
+ logger.debug("No persistent contrails formed by flight %s", label)
421
+ return self._fill_empty_flight_results(return_flight_list)
422
+
423
+ self.contrail_list = []
424
+ self._simulate_contrail_evolution()
425
+
426
+ if not self.contrail_list:
427
+ logger.debug("No contrails formed by %s", label)
428
+ return self._fill_empty_flight_results(return_flight_list)
429
+
430
+ logger.debug("Complete contrail simulation for %s", label)
431
+
432
+ self._cleanup_indices()
433
+ self._bundle_results()
434
+
435
+ if return_flight_list:
436
+ return self.source.to_flight_list() # type: ignore[attr-defined]
437
+
438
+ return self.source
439
+
440
+ def _set_timesteps(self) -> None:
441
+ """Set the :attr:`timesteps` based on the ``source`` time range.
442
+
443
+ This method is called in :meth:`eval` before the flight is processed.
444
+ """
445
+ if isinstance(self.source, Fleet):
446
+ # time not sorted in Fleet instance
447
+ tmin = self.source["time"].min()
448
+ tmax = self.source["time"].max()
449
+ else:
450
+ tmin = self.source["time"][0]
451
+ tmax = self.source["time"][-1]
452
+
453
+ tmin = pd.to_datetime(tmin)
454
+ tmax = pd.to_datetime(tmax)
455
+ dt = pd.to_timedelta(self.params["dt_integration"])
456
+
457
+ t_start = tmin.ceil(dt)
458
+ t_end = tmax.floor(dt) + self.params["max_age"] + dt
459
+ self.timesteps = np.arange(t_start, t_end, dt)
460
+
461
+ # -------------
462
+ # Model Methods
463
+ # -------------
464
+
465
+ def _process_flight(self, met: MetDataset) -> None:
466
+ """Prepare :attr:`self.source` for use in model eval.
467
+
468
+ Missing flight values should be prefilled before calling this method.
469
+ This method modifies :attr:`self.source`.
470
+
471
+ .. versionchanged:: 0.35.2
472
+
473
+ No longer broadcast all numeric source params. Instead, numeric
474
+ source params can be accessed with :meth:`Flight.get_data_or_attr`.
475
+
476
+ Parameters
477
+ ----------
478
+ met : MetDataset
479
+ Meteorology data
480
+
481
+ Raises
482
+ ------
483
+ ValueError
484
+ If non-sequential waypoints are found in ``self.source["waypoint"]``
485
+ If there is no intersection between met domain and :attr:`source`.
486
+
487
+ See Also
488
+ --------
489
+ :method:`Flight.resample_and_fill`
490
+ """
491
+ logger.debug("Pre-processing flight parameters")
492
+
493
+ # STEP 1: Check for core columns
494
+ # Attach level and altitude to avoid some redundancy
495
+ self.source.setdefault("altitude", self.source.altitude)
496
+ self.source.setdefault("level", self.source.level)
497
+ self.source.setdefault("air_pressure", self.source.air_pressure)
498
+
499
+ core_columns = ("longitude", "latitude", "altitude", "time")
500
+ for col in core_columns:
501
+ if np.isnan(self.source[col]).any():
502
+ raise ValueError(
503
+ f"Parameter `flight` must not contain NaN values in {col} field."
504
+ "Call method `resample_and_fill` to clean up flight trajectory."
505
+ )
506
+
507
+ # STEP 2: Check for waypoints
508
+ # Note: Fleet instances always have a waypoint column
509
+ if "waypoint" not in self.source:
510
+ # Give flight a range index
511
+ self.source["waypoint"] = np.arange(self.source.size)
512
+ elif not isinstance(self.source, Fleet) and not np.all(
513
+ np.diff(self.source["waypoint"]) == 1
514
+ ):
515
+ # If self.source is a usual Flight (not the subclass Fleet)
516
+ # and has "waypoint" data, we want to give a warning if the waypoint
517
+ # data has an usual index. CoCiP uses the waypoint for its continuity
518
+ # convention, so the waypoint data is critical.
519
+ msg = (
520
+ "Found non-sequential waypoints in flight key 'waypoint'. "
521
+ "The CoCiP algorithm requires flight data key 'waypoint' "
522
+ "to contain sequential waypoints if defined."
523
+ )
524
+ raise ValueError(msg)
525
+
526
+ # STEP 3: Test met domain for some overlap
527
+ # We attach the intersection to the source.
528
+ # This is used in the function _fill_empty_flight_results
529
+ intersection = self.source.coords_intersect_met(met)
530
+ self.source["_met_intersection"] = intersection
531
+ logger.debug(
532
+ "Fraction of flight waypoints intersecting met domain: %s / %s",
533
+ intersection.sum(),
534
+ self.source.size,
535
+ )
536
+ if not intersection.any():
537
+ msg = (
538
+ "No intersection between flight waypoints and met domain. "
539
+ "Rerun Cocip with met overlapping flight."
540
+ )
541
+ raise ValueError(msg)
542
+
543
+ # STEP 4: Begin met interpolation
544
+ # Unfortunately we use both "u_wind" and "eastward_wind" to refer to the
545
+ # same variable, so the logic gets a bit more complex.
546
+ humidity_scaling = self.params["humidity_scaling"]
547
+ scale_humidity = humidity_scaling is not None and "specific_humidity" not in self.source
548
+ verbose_outputs = self.params["verbose_outputs"]
549
+
550
+ interp_kwargs = self.interp_kwargs
551
+ if self.params["preprocess_lowmem"]:
552
+ interp_kwargs["lowmem"] = True
553
+ interpolate_met(met, self.source, "air_temperature", **interp_kwargs)
554
+ interpolate_met(met, self.source, "specific_humidity", **interp_kwargs)
555
+ interpolate_met(met, self.source, "eastward_wind", "u_wind", **interp_kwargs)
556
+ interpolate_met(met, self.source, "northward_wind", "v_wind", **interp_kwargs)
557
+
558
+ if scale_humidity:
559
+ humidity_scaling.eval(self.source, copy_source=False)
560
+
561
+ # if humidity_scaling isn't defined, add rhi to source for verbose_outputs
562
+ elif verbose_outputs:
563
+ self.source["rhi"] = thermo.rhi(
564
+ self.source["specific_humidity"],
565
+ self.source["air_temperature"],
566
+ self.source.air_pressure,
567
+ )
568
+
569
+ # Cache extra met properties for post-analysis
570
+ if verbose_outputs:
571
+ interpolate_met(met, self.source, "tau_cirrus", **interp_kwargs)
572
+
573
+ # handle ECMWF/GFS ciwc variables
574
+ if (key := "specific_cloud_ice_water_content") in met: # noqa: SIM114
575
+ interpolate_met(met, self.source, key, **interp_kwargs)
576
+ elif (key := "ice_water_mixing_ratio") in met:
577
+ interpolate_met(met, self.source, key, **interp_kwargs)
578
+
579
+ self.source["rho_air"] = thermo.rho_d(
580
+ self.source["air_temperature"], self.source.air_pressure
581
+ )
582
+ self.source["sdr"] = geo.solar_direct_radiation(
583
+ self.source["longitude"], self.source["latitude"], self.source["time"]
584
+ )
585
+
586
+ # STEP 5: Calculate segment-specific properties if they are not already attached
587
+ if "true_airspeed" not in self.source:
588
+ self.source["true_airspeed"] = self.source.segment_true_airspeed(
589
+ u_wind=self.source["u_wind"],
590
+ v_wind=self.source["v_wind"],
591
+ smooth=self.params["smooth_true_airspeed"],
592
+ window_length=self.params["smooth_true_airspeed_window_length"],
593
+ )
594
+ if "segment_length" not in self.source:
595
+ self.source["segment_length"] = self.source.segment_length()
596
+
597
+ max_ = self.source.max_distance_gap
598
+ lim_ = self.params["max_seg_length_m"]
599
+ if max_ > 0.9 * lim_:
600
+ warnings.warn(
601
+ "Flight trajectory has segment lengths close to or exceeding the "
602
+ "'max_seg_length_m' parameter. Evolved contrail segments may reach "
603
+ "their end of life artificially early. Either resample the flight "
604
+ "with the 'resample_and_fill' method (recommended), or use a larger "
605
+ "'max_seg_length_m' parameter. Current values: "
606
+ f"max_seg_length_m={lim_}, max_seg_length_on_trajectory={max_}"
607
+ )
608
+
609
+ # if either angle is not provided, re-calculate both
610
+ # this is the one case where we will overwrite a flight data key
611
+ if "sin_a" not in self.source or "cos_a" not in self.source:
612
+ self.source["sin_a"], self.source["cos_a"] = self.source.segment_angle()
613
+
614
+ # STEP 6: Calculate emissions if requested, or if keys don't exist in Flight
615
+ if self.params["process_emissions"]:
616
+ self._process_emissions()
617
+
618
+ # STEP 7: Ensure that flight has the required variables defined as attrs or columns
619
+ self.source.ensure_vars(_emissions_variables())
620
+
621
+ def _process_emissions(self) -> None:
622
+ """Process flight emissions.
623
+
624
+ See :class:`Emissions`.
625
+
626
+ We should consider supporting OpenAP (https://github.com/TUDelft-CNS-ATM/openap)
627
+ and alternate performance models in the future.
628
+ """
629
+ logger.debug("Processing flight emissions")
630
+
631
+ # Call aircraft performance and Emissions models as needed
632
+ # NOTE: None of these sub-models actually do any met interpolation -- self.source already
633
+ # has all of the required met variables attached. Therefore, we don't need to worry about
634
+ # being consistent with passing in Cocip's interp_kwargs and humidity_scaling into
635
+ # the sub-models.
636
+ emissions = Emissions()
637
+ ap_model = self.params["aircraft_performance"]
638
+
639
+ # Run against a list of flights (Fleet)
640
+ if isinstance(self.source, Fleet):
641
+ # Rip the Fleet apart, run BADA on each, then reassemble
642
+ logger.debug("Separately running aircraft performance on each flight in fleet")
643
+ fls = self.source.to_flight_list(copy=False)
644
+ fls = [_eval_aircraft_performance(ap_model, fl) for fl in fls]
645
+
646
+ # In Fleet-mode, always call emissions
647
+ logger.debug("Separately running emissions on each flight in fleet")
648
+ fls = [_eval_emissions(emissions, fl) for fl in fls]
649
+
650
+ # Broadcast numeric AP and emissions variables back to Fleet.data
651
+ for fl in fls:
652
+ fl.broadcast_attrs(_emissions_variables(), raise_error=False)
653
+
654
+ # Convert back to fleet
655
+ attrs = self.source.attrs
656
+ attrs.pop("fl_attrs", None)
657
+ attrs.pop("data_keys", None)
658
+ self.source = Fleet.from_seq(fls, broadcast_numeric=False, copy=False, attrs=attrs)
659
+
660
+ # Single flight
661
+ else:
662
+ self.source = _eval_aircraft_performance(ap_model, self.source)
663
+ self.source = _eval_emissions(emissions, self.source)
664
+
665
+ # Scale nvPM with parameter (fleet / flight)
666
+ factor = self.params["nvpm_ei_n_enhancement_factor"]
667
+ try:
668
+ self.source["nvpm_ei_n"] *= factor
669
+ except KeyError:
670
+ self.source.attrs.update({"nvpm_ei_n": self.source.attrs["nvpm_ei_n"] * factor})
671
+
672
+ def _find_initial_contrail_regions(self) -> None:
673
+ """Apply Schmidt-Appleman criteria to determine regions of persistent contrail formation.
674
+
675
+ This method:
676
+ - Modifies :attr:`flight` in place by assigning additional columns arising from SAC
677
+ - Creates :attr:`_sac_flight` needed for :meth:`_simulate_wave_vortex_downwash`.
678
+ """
679
+ # calculate the SAC for each waypoint
680
+ sac_model = sac.SAC(
681
+ met=None, # self.source is already interpolated against met, so we don't need it
682
+ copy_source=False,
683
+ )
684
+ self.source = sac_model.eval(source=self.source)
685
+
686
+ # Estimate SAC threshold temperature along initial contrail
687
+ # This variable is not involved in calculating the SAC,
688
+ # but it is required by `contrail_properties.ice_particle_number`.
689
+ # The three variables used below are all computed in sac_model.eval
690
+ self.source["T_critical_sac"] = sac.T_critical_sac(
691
+ self.source["T_sat_liquid"],
692
+ self.source["rh"],
693
+ self.source["G"],
694
+ )
695
+
696
+ # create a new Flight only at points where "sac" == 1
697
+ filt = self.source["sac"] == 1.0
698
+ if self.params["filter_sac"]:
699
+ self._sac_flight = self.source.filter(filt)
700
+ logger.debug(
701
+ "Fraction of waypoints satisfying the SAC: %s / %s",
702
+ len(self._sac_flight),
703
+ len(self.source),
704
+ )
705
+ return
706
+
707
+ warnings.warn("Manually overriding SAC filter")
708
+ logger.info("Manually overriding SAC filter")
709
+ self._sac_flight = self.source.copy()
710
+ logger.debug(
711
+ "Fraction of waypoints satisfying the SAC: %s / %s",
712
+ np.sum(filt),
713
+ self._sac_flight.size,
714
+ )
715
+ logger.debug("None are filtered out!")
716
+
717
+ def _simulate_wake_vortex_downwash(self, met: MetDataset) -> None:
718
+ """Apply wake vortex model to calculate initial contrail geometry.
719
+
720
+ The calculation uses a parametric wake vortex transport and decay model to simulate the
721
+ wake vortex phase and obtain the initial contrail width, depth and downward displacement.
722
+
723
+ This method assigns additional columns "width" and "depth" to :attr:`_sac_flight`
724
+
725
+ Parameters
726
+ ----------
727
+ met : MetDataset
728
+ Meteorology data
729
+
730
+ References
731
+ ----------
732
+ - :cite:`holzapfelProbabilisticTwoPhaseWake2003`
733
+ """
734
+ air_temperature = self._sac_flight["air_temperature"]
735
+ u_wind = self._sac_flight["u_wind"]
736
+ v_wind = self._sac_flight["v_wind"]
737
+ air_pressure = self._sac_flight.air_pressure
738
+
739
+ # flight parameters
740
+ aircraft_mass = self._sac_flight.get_data_or_attr("aircraft_mass")
741
+ true_airspeed = self._sac_flight["true_airspeed"]
742
+
743
+ # In Fleet-mode, wingspan resides on `data`, and in Flight-mode,
744
+ # wingspan resides on `attrs`.
745
+ wingspan = self._sac_flight.get_data_or_attr("wingspan")
746
+
747
+ # get the pressure level `dz_m` lower than element pressure
748
+ dz_m = self.params["dz_m"]
749
+ air_pressure_lower = thermo.pressure_dz(air_temperature, air_pressure, dz_m)
750
+ level_lower = air_pressure_lower / 100.0
751
+
752
+ # get full met grid or flight data interpolated to the pressure level `p_dz`
753
+ interp_kwargs = self.interp_kwargs
754
+ if self.params["preprocess_lowmem"]:
755
+ interp_kwargs["lowmem"] = True
756
+ air_temperature_lower = interpolate_met(
757
+ met,
758
+ self._sac_flight,
759
+ "air_temperature",
760
+ "air_temperature_lower",
761
+ level=level_lower,
762
+ **interp_kwargs,
763
+ )
764
+ u_wind_lower = interpolate_met(
765
+ met,
766
+ self._sac_flight,
767
+ "eastward_wind",
768
+ "u_wind_lower",
769
+ level=level_lower,
770
+ **interp_kwargs,
771
+ )
772
+ v_wind_lower = interpolate_met(
773
+ met,
774
+ self._sac_flight,
775
+ "northward_wind",
776
+ "v_wind_lower",
777
+ level=level_lower,
778
+ **interp_kwargs,
779
+ )
780
+
781
+ # Temperature gradient
782
+ dT_dz = self._sac_flight["dT_dz"] = thermo.T_potential_gradient(
783
+ air_temperature,
784
+ air_pressure,
785
+ air_temperature_lower,
786
+ air_pressure_lower,
787
+ dz_m,
788
+ )
789
+
790
+ # wind shear
791
+ ds_dz = self._sac_flight["ds_dz"] = wind_shear.wind_shear(
792
+ u_wind, u_wind_lower, v_wind, v_wind_lower, dz_m
793
+ )
794
+
795
+ # Initial contrail width, depth and downward displacement
796
+ dz_max = self._sac_flight["dz_max"] = wake_vortex.max_downward_displacement(
797
+ wingspan,
798
+ true_airspeed,
799
+ aircraft_mass,
800
+ air_temperature,
801
+ dT_dz,
802
+ ds_dz,
803
+ air_pressure,
804
+ effective_vertical_resolution=self.params["effective_vertical_resolution"],
805
+ wind_shear_enhancement_exponent=self.params["wind_shear_enhancement_exponent"],
806
+ )
807
+
808
+ # derive downwash values and save to data model
809
+ self._sac_flight["width"] = wake_vortex.initial_contrail_width(wingspan, dz_max)
810
+ initial_wake_vortex_depth = self.params["initial_wake_vortex_depth"]
811
+ self._sac_flight["depth"] = wake_vortex.initial_contrail_depth(
812
+ dz_max, initial_wake_vortex_depth
813
+ )
814
+ # Initially, sigma_yz is set to 0
815
+ # See bottom left paragraph p. 552 Schumann 2012 beginning with:
816
+ # >>> "The contrail model starts from initial values ..."
817
+ self._sac_flight["sigma_yz"] = np.zeros_like(dz_max)
818
+
819
+ def _find_initial_persistent_contrails(self, met: MetDataset) -> None:
820
+ """Calculate the initial contrail properties after the wake vortex phase.
821
+
822
+ Determine points with initial persistent contrails.
823
+
824
+ This method first calculates ice water content (``iwc``) and number of ice particles per
825
+ distance (``n_ice_per_m_1``).
826
+
827
+ It then tests each contrail waypoint for initial persistence and calculates
828
+ parameters for the first iteration of the time integration.
829
+
830
+ Note that the subscript "_1" represents the conditions after the wake vortex phase.
831
+
832
+ Parameters
833
+ ----------
834
+ met : MetDataset
835
+ Meteorology data
836
+ """
837
+
838
+ # met parameters along Flight path
839
+ air_pressure = self._sac_flight.air_pressure
840
+ air_temperature = self._sac_flight["air_temperature"]
841
+ specific_humidity = self._sac_flight["specific_humidity"]
842
+ T_critical_sac = self._sac_flight["T_critical_sac"]
843
+
844
+ # Flight performance parameters
845
+ fuel_flow = self._sac_flight.get_data_or_attr("fuel_flow")
846
+ true_airspeed = self._sac_flight["true_airspeed"]
847
+ fuel_dist = fuel_flow / true_airspeed
848
+
849
+ nvpm_ei_n = self._sac_flight.get_data_or_attr("nvpm_ei_n")
850
+ ei_h2o = self._sac_flight.fuel.ei_h2o
851
+
852
+ # get initial contrail parameters from wake vortex simulation
853
+ width = self._sac_flight["width"]
854
+ depth = self._sac_flight["depth"]
855
+
856
+ # initial contrail altitude set to 0.5 * depth
857
+ contrail_1 = GeoVectorDataset(
858
+ longitude=self._sac_flight["longitude"],
859
+ latitude=self._sac_flight["latitude"],
860
+ altitude=self._sac_flight.altitude - 0.5 * depth,
861
+ time=self._sac_flight["time"],
862
+ copy=False,
863
+ )
864
+
865
+ # get met post wake vortex along initial contrail
866
+ interp_kwargs = self.interp_kwargs
867
+ if self.params["preprocess_lowmem"]:
868
+ interp_kwargs["lowmem"] = True
869
+ air_temperature_1 = interpolate_met(met, contrail_1, "air_temperature", **interp_kwargs)
870
+ interpolate_met(met, contrail_1, "specific_humidity", **interp_kwargs)
871
+
872
+ humidity_scaling = self.params["humidity_scaling"]
873
+ if humidity_scaling is not None:
874
+ humidity_scaling.eval(contrail_1, copy_source=False)
875
+ else:
876
+ contrail_1["air_pressure"] = contrail_1.air_pressure
877
+ contrail_1["rhi"] = thermo.rhi(
878
+ contrail_1["specific_humidity"], air_temperature_1, contrail_1["air_pressure"]
879
+ )
880
+
881
+ air_pressure_1 = contrail_1["air_pressure"]
882
+ specific_humidity_1 = contrail_1["specific_humidity"]
883
+ rhi_1 = contrail_1["rhi"]
884
+ level_1 = contrail_1.level
885
+ altitude_1 = contrail_1.altitude
886
+
887
+ # calculate thermo properties post wake vortex
888
+ q_sat_1 = thermo.q_sat_ice(air_temperature_1, air_pressure_1)
889
+ rho_air_1 = thermo.rho_d(air_temperature_1, air_pressure_1)
890
+
891
+ # Initialize initial contrail properties
892
+ iwc = contrail_properties.initial_iwc(
893
+ air_temperature, specific_humidity, air_pressure, fuel_dist, width, depth, ei_h2o
894
+ )
895
+ iwc_ad = contrail_properties.iwc_adiabatic_heating(
896
+ air_temperature, air_pressure, air_pressure_1
897
+ )
898
+ iwc_1 = contrail_properties.iwc_post_wake_vortex(iwc, iwc_ad)
899
+
900
+ if self.params["unterstrasser_ice_survival_fraction"]:
901
+ wingspan = self._sac_flight.get_data_or_attr("wingspan")
902
+ rhi_0 = thermo.rhi(specific_humidity, air_temperature, air_pressure)
903
+ f_surv = unterstrasser_wake_vortex.ice_particle_number_survival_fraction(
904
+ air_temperature,
905
+ rhi_0,
906
+ ei_h2o,
907
+ wingspan,
908
+ true_airspeed,
909
+ fuel_flow,
910
+ nvpm_ei_n,
911
+ 0.5 * depth, # Taking the mid-point of the contrail plume
912
+ )
913
+ else:
914
+ f_surv = contrail_properties.ice_particle_survival_fraction(iwc, iwc_1)
915
+
916
+ n_ice_per_m_1 = contrail_properties.ice_particle_number(
917
+ nvpm_ei_n=nvpm_ei_n,
918
+ fuel_dist=fuel_dist,
919
+ f_surv=f_surv,
920
+ air_temperature=air_temperature,
921
+ T_crit_sac=T_critical_sac,
922
+ min_ice_particle_number_nvpm_ei_n=self.params["min_ice_particle_number_nvpm_ei_n"],
923
+ )
924
+
925
+ # Check for persistent initial_contrails
926
+ persistent_1 = contrail_properties.initial_persistent(iwc_1, rhi_1)
927
+
928
+ self._sac_flight["altitude_1"] = altitude_1
929
+ self._sac_flight["level_1"] = level_1
930
+ self._sac_flight["air_temperature_1"] = air_temperature_1
931
+ self._sac_flight["specific_humidity_1"] = specific_humidity_1
932
+ self._sac_flight["q_sat_1"] = q_sat_1
933
+ self._sac_flight["air_pressure_1"] = air_pressure_1
934
+ self._sac_flight["rho_air_1"] = rho_air_1
935
+ self._sac_flight["rhi_1"] = rhi_1
936
+ self._sac_flight["iwc_1"] = iwc_1
937
+ self._sac_flight["n_ice_per_m_1"] = n_ice_per_m_1
938
+ self._sac_flight["persistent_1"] = persistent_1
939
+
940
+ # Create new Flight only at persistent points
941
+ if self.params["filter_initially_persistent"]:
942
+ self._downwash_flight = self._sac_flight.filter(persistent_1.astype(bool))
943
+ logger.debug(
944
+ "Fraction of waypoints with initially persistent contrails: %s / %s",
945
+ len(self._downwash_flight),
946
+ len(self._sac_flight),
947
+ )
948
+
949
+ else:
950
+ warnings.warn("Manually overriding initially persistent filter")
951
+ logger.info("Manually overriding initially persistent filter")
952
+ self._downwash_flight = self._sac_flight.copy()
953
+ logger.debug(
954
+ "Fraction of waypoints with initially persistent contrails: %s / %s",
955
+ persistent_1.sum(),
956
+ persistent_1.size,
957
+ )
958
+ logger.debug("None are filtered out!")
959
+
960
+ def _process_downwash_flight(self) -> tuple[MetDataset | None, MetDataset | None]:
961
+ """Create and calculate properties of contrails created by downwash vortex.
962
+
963
+ ``_downwash_contrail`` is a contrail representation of the waypoints of
964
+ ``_downwash_flight``, which has already been filtered for initial persistent waypoints.
965
+
966
+ Returns MetDatasets for subsequent use if ``preprocess_lowmem=False``.
967
+ """
968
+ self._downwash_contrail = self._create_downwash_contrail()
969
+ buffers = {
970
+ f"{coord}_buffer": self.params[f"met_{coord}_buffer"]
971
+ for coord in ("longitude", "latitude", "level")
972
+ }
973
+ logger.debug("Downselect met for start of Cocip evolution")
974
+ met = self._downwash_contrail.downselect_met(self.met, **buffers, copy=False)
975
+ met = add_tau_cirrus(met)
976
+ rad = self._downwash_contrail.downselect_met(self.rad, **buffers, copy=False)
977
+
978
+ calc_continuous(self._downwash_contrail)
979
+ calc_timestep_geometry(self._downwash_contrail)
980
+
981
+ interp_kwargs = self.interp_kwargs
982
+ if self.params["preprocess_lowmem"]:
983
+ interp_kwargs["lowmem"] = True
984
+ calc_timestep_meteorology(self._downwash_contrail, met, self.params, **interp_kwargs)
985
+ calc_shortwave_radiation(rad, self._downwash_contrail, **interp_kwargs)
986
+ calc_outgoing_longwave_radiation(rad, self._downwash_contrail, **interp_kwargs)
987
+ calc_contrail_properties(
988
+ self._downwash_contrail,
989
+ self.params["effective_vertical_resolution"],
990
+ self.params["wind_shear_enhancement_exponent"],
991
+ self.params["sedimentation_impact_factor"],
992
+ self.params["radiative_heating_effects"],
993
+ )
994
+
995
+ # Intersect with rad dataset
996
+ calc_radiative_properties(self._downwash_contrail, self.params)
997
+
998
+ if self.params["preprocess_lowmem"]:
999
+ return None, None
1000
+ return met, rad
1001
+
1002
+ def _simulate_contrail_evolution(self) -> None:
1003
+ """Simulate contrail evolution."""
1004
+
1005
+ met, rad = self._process_downwash_flight()
1006
+ interp_kwargs = self.interp_kwargs
1007
+
1008
+ contrail_contrail_overlapping = self.params["contrail_contrail_overlapping"]
1009
+ if contrail_contrail_overlapping and not isinstance(self.source, Fleet):
1010
+ warnings.warn("Contrail-Contrail Overlapping is only valid for Fleet mode.")
1011
+
1012
+ # Complete iteration at time_idx - 2
1013
+ for time_idx, time_end in enumerate(self.timesteps[:-1]):
1014
+ logger.debug("Start time integration step %s ending at time %s", time_idx, time_end)
1015
+
1016
+ # get the last evolution step of contrail waypoints, if it exists
1017
+ latest_contrail = self.contrail_list[-1] if self.contrail_list else GeoVectorDataset()
1018
+
1019
+ # load new contrail segments from downwash_flight
1020
+ contrail_2_segments = self._get_contrail_2_segments(time_idx)
1021
+ if contrail_2_segments:
1022
+ logger.debug(
1023
+ "Discover %s new contrail waypoints formed by downwash_flight waypoints.",
1024
+ contrail_2_segments.size,
1025
+ )
1026
+ logger.debug("Previously persistent contrail size: %s", latest_contrail.size)
1027
+
1028
+ # Append new waypoints to latest contrail, or set as first contrail waypoints
1029
+ latest_contrail = latest_contrail + contrail_2_segments
1030
+
1031
+ # CRITICAL: When running in "Fleet" mode, the latest_contrail is no longer
1032
+ # sorted by flight_id. Fixing that below.
1033
+ if isinstance(self.source, Fleet):
1034
+ latest_contrail = latest_contrail.sort(["flight_id", "time"])
1035
+
1036
+ # Check for an empty contrail
1037
+ if not latest_contrail:
1038
+ logger.debug("Empty latest_contrail at timestep %s", time_end)
1039
+ if np.all(time_end > self._downwash_contrail["time"]):
1040
+ logger.debug("No remaining downwash_contrail waypoints. Break.")
1041
+ break
1042
+ continue
1043
+
1044
+ # Update met, rad slices as needed
1045
+ met, rad = self._maybe_downselect_met_rad(met, rad, latest_contrail, time_end)
1046
+
1047
+ # Recalculate latest_contrail with new values
1048
+ # NOTE: We are doing a substantial amount of redundant computation here
1049
+ # At waypoints for which the continuity hasn't changed, there is nothing
1050
+ # new going on, and so we are overwriting variables in latest_contrail
1051
+ # with the same values
1052
+ # NOTE: Both latest_contrail and contrail_2_segments contain all
1053
+ # 54 variables. The only difference is the change in continuity.
1054
+ # The change in continuity impacts:
1055
+ # - sin_a, cos_a, segment_length (in calc_timestep_geometry)
1056
+ # - dsn_dz (in calc_timestep_meteorology, then enhanced in contrail_properties)
1057
+ # And there is huge room to optimize this
1058
+ calc_continuous(latest_contrail)
1059
+ calc_timestep_geometry(latest_contrail)
1060
+ calc_timestep_meteorology(latest_contrail, met, self.params, **interp_kwargs)
1061
+ calc_contrail_properties(
1062
+ latest_contrail,
1063
+ self.params["effective_vertical_resolution"],
1064
+ self.params["wind_shear_enhancement_exponent"],
1065
+ self.params["sedimentation_impact_factor"],
1066
+ self.params["radiative_heating_effects"],
1067
+ )
1068
+
1069
+ final_contrail = calc_timestep_contrail_evolution(
1070
+ met=met,
1071
+ rad=rad,
1072
+ contrail_1=latest_contrail,
1073
+ time_2=time_end,
1074
+ params=self.params,
1075
+ **interp_kwargs,
1076
+ )
1077
+
1078
+ if contrail_contrail_overlapping:
1079
+ final_contrail = _contrail_contrail_overlapping(final_contrail, self.params)
1080
+
1081
+ self.contrail_list.append(final_contrail)
1082
+
1083
+ def _maybe_downselect_met_rad(
1084
+ self,
1085
+ met: MetDataset | None,
1086
+ rad: MetDataset | None,
1087
+ latest_contrail: GeoVectorDataset,
1088
+ time_end: np.datetime64,
1089
+ ) -> tuple[MetDataset, MetDataset]:
1090
+ """Downselect ``self.met`` and ``self.rad`` if necessary to cover ``time_end``.
1091
+
1092
+ If current ``met`` and ``rad`` slices to not include ``time_end``, new slices are selected
1093
+ from ``self.met`` and ``self.rad``. Downselection in space will cover
1094
+ - locations of current contrails (``latest_contrail``),
1095
+ - locations of additional contrails that will be loaded from ``self._downwash_flight``
1096
+ before the new slices expire,
1097
+ plus a user-defined buffer.
1098
+ """
1099
+ if met is None or time_end > met.indexes["time"].to_numpy()[-1]:
1100
+ logger.debug("Downselect met at time_end %s within Cocip evolution", time_end)
1101
+ met = self._definitely_downselect_met_or_rad(self.met, latest_contrail, time_end)
1102
+ met = add_tau_cirrus(met)
1103
+
1104
+ if rad is None or time_end > rad.indexes["time"].to_numpy()[-1]:
1105
+ logger.debug("Downselect rad at time_end %s within Cocip evolution", time_end)
1106
+ rad = self._definitely_downselect_met_or_rad(self.rad, latest_contrail, time_end)
1107
+
1108
+ return met, rad
1109
+
1110
+ def _definitely_downselect_met_or_rad(
1111
+ self, met: MetDataset, latest_contrail: GeoVectorDataset, time_end: np.datetime64
1112
+ ) -> MetDataset:
1113
+ """Perform downselection when required by :meth:`_maybe_downselect_met_rad`.
1114
+
1115
+ Downselects ``met`` (which should be one of ``self.met`` or ``self.rad``)
1116
+ to cover ``time_end``. Downselection in space covers
1117
+ - locations of current contrails (``latest_contrail``),
1118
+ - locations of additional contrails that will be loaded from ``self._downwash_flight``
1119
+ before the new slices expire,
1120
+ plus a user-defined buffer, as described in :meth:`_maybe_downselect_met_rad`.
1121
+ """
1122
+ # compute lookahead for future contrails from downwash_flight
1123
+ met_time = met.indexes["time"].to_numpy()
1124
+ mask = met_time >= time_end
1125
+ lookahead = np.min(met_time[mask]) if np.any(mask) else time_end
1126
+
1127
+ # create vector for downselection based on current + future contrails
1128
+ future_contrails = self._downwash_flight.filter(
1129
+ (self._downwash_flight["time"] >= time_end)
1130
+ & (self._downwash_flight["time"] <= lookahead),
1131
+ copy=False,
1132
+ )
1133
+ vector = GeoVectorDataset(
1134
+ {
1135
+ key: np.concatenate((latest_contrail[key], future_contrails[key]))
1136
+ for key in ("longitude", "latitude", "level", "time")
1137
+ }
1138
+ )
1139
+
1140
+ # compute time buffer to ensure downselection extends to time_end
1141
+ buffers = {
1142
+ f"{coord}_buffer": self.params[f"met_{coord}_buffer"]
1143
+ for coord in ("longitude", "latitude", "level")
1144
+ }
1145
+ buffers["time_buffer"] = (
1146
+ np.timedelta64(0, "ns"),
1147
+ max(np.timedelta64(0, "ns"), time_end - vector["time"].max()),
1148
+ )
1149
+
1150
+ return vector.downselect_met(met, **buffers, copy=False)
1151
+
1152
+ def _create_downwash_contrail(self) -> GeoVectorDataset:
1153
+ """Get Contrail representation of downwash flight."""
1154
+
1155
+ downwash_contrail_data = {
1156
+ "waypoint": self._downwash_flight["waypoint"],
1157
+ "flight_id": self._downwash_flight["flight_id"],
1158
+ "time": self._downwash_flight["time"],
1159
+ "longitude": self._downwash_flight["longitude"],
1160
+ "latitude": self._downwash_flight["latitude"],
1161
+ # intentionally specify altitude and level to avoid pressure level calculations
1162
+ "altitude": self._downwash_flight["altitude_1"],
1163
+ "level": self._downwash_flight["level_1"],
1164
+ "air_pressure": self._downwash_flight["air_pressure_1"],
1165
+ "width": self._downwash_flight["width"],
1166
+ "depth": self._downwash_flight["depth"],
1167
+ "sigma_yz": self._downwash_flight["sigma_yz"],
1168
+ "air_temperature": self._downwash_flight["air_temperature_1"],
1169
+ "specific_humidity": self._downwash_flight["specific_humidity_1"],
1170
+ "q_sat": self._downwash_flight["q_sat_1"],
1171
+ "rho_air": self._downwash_flight["rho_air_1"],
1172
+ "rhi": self._downwash_flight["rhi_1"],
1173
+ "iwc": self._downwash_flight["iwc_1"],
1174
+ "n_ice_per_m": self._downwash_flight["n_ice_per_m_1"],
1175
+ "persistent": self._downwash_flight["persistent_1"],
1176
+ }
1177
+
1178
+ contrail = GeoVectorDataset(downwash_contrail_data, copy=True)
1179
+ contrail["formation_time"] = contrail["time"].copy()
1180
+ contrail["age"] = contrail["formation_time"] - contrail["time"]
1181
+
1182
+ # Heating rate, differential heating rate and
1183
+ # cumulative heat energy absorbed by the contrail
1184
+ if self.params["radiative_heating_effects"]:
1185
+ contrail["heat_rate"] = np.zeros_like(contrail["n_ice_per_m"])
1186
+ contrail["d_heat_rate"] = np.zeros_like(contrail["n_ice_per_m"])
1187
+
1188
+ # Increase in temperature averaged over the contrail plume
1189
+ contrail["cumul_heat"] = np.zeros_like(contrail["n_ice_per_m"])
1190
+
1191
+ # Temperature difference between the upper and lower half of the contrail plume
1192
+ contrail["cumul_differential_heat"] = np.zeros_like(contrail["n_ice_per_m"])
1193
+
1194
+ # Initially set energy forcing to 0 because the contrail just formed (age = 0)
1195
+ contrail["ef"] = np.zeros_like(contrail["n_ice_per_m"])
1196
+ if self.params["compute_atr20"]:
1197
+ contrail["global_yearly_mean_rf"] = np.zeros_like(contrail["n_ice_per_m"])
1198
+ contrail["atr20"] = np.zeros_like(contrail["n_ice_per_m"])
1199
+
1200
+ if not self.params["filter_sac"]:
1201
+ contrail["sac"] = self._downwash_flight["sac"]
1202
+ if not self.params["filter_initially_persistent"]:
1203
+ contrail["initially_persistent"] = self._downwash_flight["persistent_1"]
1204
+ if self.params["persistent_buffer"] is not None:
1205
+ contrail["end_of_life"] = np.full(contrail.size, np.datetime64("NaT", "ns"))
1206
+
1207
+ return contrail
1208
+
1209
+ def _get_contrail_2_segments(self, time_idx: int) -> GeoVectorDataset:
1210
+ """Get batch of newly formed contrail segments."""
1211
+ time = self._downwash_contrail["time"]
1212
+ t_cur = self.timesteps[time_idx]
1213
+ if time_idx == 0:
1214
+ filt = time < t_cur
1215
+ else:
1216
+ t_prev = self.timesteps[time_idx - 1]
1217
+ filt = (time < t_cur) & (time >= t_prev)
1218
+
1219
+ return self._downwash_contrail.filter(filt)
1220
+
1221
+ @overrides
1222
+ def _cleanup_indices(self) -> None:
1223
+ """Cleanup interpolation artifacts."""
1224
+
1225
+ if not self.params["interpolation_use_indices"]:
1226
+ return
1227
+
1228
+ if hasattr(self, "contrail_list"):
1229
+ for contrail in self.contrail_list:
1230
+ contrail._invalidate_indices()
1231
+
1232
+ self.source._invalidate_indices()
1233
+ self._sac_flight._invalidate_indices()
1234
+ if hasattr(self, "_downwash_flight"):
1235
+ self._downwash_flight._invalidate_indices()
1236
+ if hasattr(self, "_downwash_contrail"):
1237
+ self._downwash_contrail._invalidate_indices()
1238
+
1239
+ def _bundle_results(self) -> None:
1240
+ # ---
1241
+ # Create contrail dataframe (self.contrail)
1242
+ # ---
1243
+ self.contrail = GeoVectorDataset.sum(self.contrail_list).dataframe
1244
+ self.contrail["timestep"] = np.concatenate(
1245
+ [np.full(c.size, i) for i, c in enumerate(self.contrail_list)]
1246
+ )
1247
+
1248
+ # add age in hours to the contrail waypoint outputs
1249
+ age_hours = np.empty_like(self.contrail["ef"])
1250
+ np.divide(self.contrail["age"], np.timedelta64(1, "h"), out=age_hours)
1251
+ self.contrail["age_hours"] = age_hours
1252
+
1253
+ verbose_outputs = self.params["verbose_outputs"]
1254
+ if verbose_outputs:
1255
+ # Compute dt_integration -- logic is somewhat complicated, but
1256
+ # we're simply addressing that the first dt_integration
1257
+ # is different from the rest
1258
+
1259
+ # We call reset_index to introduces an `index` RangeIndex column,
1260
+ # Which we use in the `groupby` to identify the
1261
+ # index of the first evolution step at each waypoint.
1262
+ tmp = self.contrail.reset_index()
1263
+ cols = ["formation_time", "time", "index"]
1264
+ first_form_time = tmp.groupby("waypoint")[cols].first()
1265
+ first_dt = first_form_time["time"] - first_form_time["formation_time"]
1266
+ first_dt = first_dt.set_axis(first_form_time["index"])
1267
+
1268
+ self.contrail = tmp.set_index("index")
1269
+ self.contrail["dt_integration"] = first_dt
1270
+ self.contrail.fillna({"dt_integration": self.params["dt_integration"]}, inplace=True)
1271
+
1272
+ # ---
1273
+ # Create contrail xr.Dataset (self.contrail_dataset)
1274
+ # ---
1275
+ if isinstance(self.source, Fleet):
1276
+ keys = ["flight_id", "timestep", "waypoint"]
1277
+ else:
1278
+ keys = ["timestep", "waypoint"]
1279
+ self.contrail_dataset = xr.Dataset.from_dataframe(self.contrail.set_index(keys))
1280
+
1281
+ # ---
1282
+ # Create output Flight / Fleet (self.source)
1283
+ # ---
1284
+
1285
+ col_idx = ["flight_id", "waypoint"] if isinstance(self.source, Fleet) else ["waypoint"]
1286
+ del self.source["_met_intersection"]
1287
+
1288
+ # Attach intermediate calculations from `sac_flight` and `downwash_flight` to flight
1289
+ sac_cols = [
1290
+ "width",
1291
+ "depth",
1292
+ "rhi_1",
1293
+ "air_temperature_1",
1294
+ "specific_humidity_1",
1295
+ "altitude_1",
1296
+ "persistent_1",
1297
+ ]
1298
+
1299
+ # add additional columns
1300
+ if verbose_outputs:
1301
+ sac_cols += ["dT_dz", "ds_dz", "dz_max"]
1302
+
1303
+ downwash_cols = ["rho_air_1", "iwc_1", "n_ice_per_m_1"]
1304
+ df = pd.concat(
1305
+ [
1306
+ self.source.dataframe.set_index(col_idx),
1307
+ self._sac_flight.dataframe.set_index(col_idx)[sac_cols],
1308
+ self._downwash_flight.dataframe.set_index(col_idx)[downwash_cols],
1309
+ ],
1310
+ axis=1,
1311
+ )
1312
+
1313
+ # Aggregate contrail data back to flight
1314
+ grouped = self.contrail.groupby(col_idx)
1315
+
1316
+ # Perform all aggregations
1317
+ agg_dict = {"ef": ["sum"], "age": ["max"]}
1318
+ if self.params["compute_atr20"]:
1319
+ agg_dict["global_yearly_mean_rf"] = ["sum"]
1320
+ agg_dict["atr20"] = ["sum"]
1321
+
1322
+ rad_keys = ["sdr", "rsr", "olr", "rf_sw", "rf_lw", "rf_net"]
1323
+ for key in rad_keys:
1324
+ if verbose_outputs:
1325
+ agg_dict[key] = ["mean", "min", "max"]
1326
+ else:
1327
+ agg_dict[key] = ["mean"]
1328
+
1329
+ aggregated = grouped.agg(agg_dict)
1330
+ aggregated.columns = [f"{k1}_{k2}" for k1, k2 in aggregated.columns]
1331
+ aggregated = aggregated.rename(columns={"ef_sum": "ef", "age_max": "contrail_age"})
1332
+ if self.params["compute_atr20"]:
1333
+ aggregated = aggregated.rename(
1334
+ columns={"global_yearly_mean_rf_sum": "global_yearly_mean_rf", "atr20_sum": "atr20"}
1335
+ )
1336
+
1337
+ # Join the two
1338
+ df = df.join(aggregated)
1339
+
1340
+ # Fill missing values for ef and contrail_age per conventions
1341
+ # Mean, max, and min radiative values are *not* filled with 0
1342
+ df.fillna({"ef": 0.0, "contrail_age": np.timedelta64(0, "ns")}, inplace=True)
1343
+
1344
+ # cocip flag for each waypoint
1345
+ # -1 if negative EF, 0 if no EF, 1 if positive EF,
1346
+ # or NaN for outside of domain of flight waypoints that don't persist
1347
+ df["cocip"] = np.sign(df["ef"])
1348
+ logger.debug("Total number of waypoints with nonzero EF: %s", df["cocip"].ne(0.0).sum())
1349
+
1350
+ # reset the index
1351
+ df = df.reset_index()
1352
+
1353
+ # Reassign to source
1354
+ self.source.data = VectorDataDict({k: v.to_numpy() for k, v in df.items()})
1355
+
1356
+ def _fill_empty_flight_results(self, return_list_flight: bool) -> Flight | list[Flight]:
1357
+ """Fill empty results into flight / fleet and return.
1358
+
1359
+ This method attaches an all nan array to each of the variables:
1360
+ - sdr
1361
+ - rsr
1362
+ - olr
1363
+ - rf_sw
1364
+ - rf_lw
1365
+ - rf_net
1366
+
1367
+ This method also attaches zeros (for trajectory points contained within met grid)
1368
+ or nans (for trajectory points outside of the met grid) to the following variables.
1369
+ - ef
1370
+ - cocip
1371
+ - contrail_age
1372
+ - persistent_1
1373
+
1374
+ Parameters
1375
+ ----------
1376
+ return_list_flight : bool
1377
+ If True, a list of :class:`Flight` is returned. In this case, :attr:`source`
1378
+ is assumed to be a :class:`Fleet`.
1379
+
1380
+ Returns
1381
+ -------
1382
+ Flight | list[Flight]
1383
+ Flight or list of Flight objects with empty variables.
1384
+ """
1385
+ self._cleanup_indices()
1386
+
1387
+ intersection = self.source.data.pop("_met_intersection")
1388
+ zeros_and_nans = np.zeros(intersection.shape, dtype=np.float32)
1389
+ zeros_and_nans[~intersection] = np.nan
1390
+ self.source["ef"] = zeros_and_nans.copy()
1391
+ self.source["persistent_1"] = zeros_and_nans.copy()
1392
+ self.source["cocip"] = np.sign(zeros_and_nans)
1393
+ self.source["contrail_age"] = zeros_and_nans.astype("timedelta64[ns]")
1394
+
1395
+ if return_list_flight:
1396
+ return self.source.to_flight_list() # type: ignore[attr-defined]
1397
+
1398
+ return self.source
1399
+
1400
+
1401
+ # ----------------------------------------
1402
+ # Functions used in Cocip and CocipGrid
1403
+ # ----------------------------------------
1404
+
1405
+
1406
+ def process_met_datasets(
1407
+ met: MetDataset,
1408
+ rad: MetDataset,
1409
+ compute_tau_cirrus: bool | Literal["auto"] = "auto",
1410
+ ) -> tuple[MetDataset, MetDataset]:
1411
+ """Process and verify ERA5 data for :class:`Cocip` and :class:`CocipGrid`.
1412
+
1413
+ The implementation uses :class:`Cocip` for the source of truth in determining
1414
+ which met variables are required.
1415
+
1416
+ .. versionchanged:: 0.25.0
1417
+
1418
+ This method is called in the :class:`CocipGrid` constructor regardless
1419
+ of the ``process_met`` parameter. The same approach is also taken
1420
+ in :class:`Cocip` in version 0.27.0.
1421
+
1422
+ .. versionchanged:: 0.48.0
1423
+
1424
+ Remove the ``shift_radiation_time`` parameter. This parameter is now
1425
+ inferred from the metadata on the `rad` instance.
1426
+
1427
+ Parameters
1428
+ ----------
1429
+ met : MetDataset
1430
+ Met pressure-level data
1431
+ rad : MetDataset
1432
+ Rad single-level data
1433
+ compute_tau_cirrus : bool | Literal["auto"]
1434
+ Whether to add ``"tau_cirrus"`` variable to pressure-level met data. If set to
1435
+ ``"auto"``, ``"tau_cirrus"`` will be computed iff the met data is dask-backed.
1436
+
1437
+ Returns
1438
+ -------
1439
+ met : MetDataset
1440
+ Met data, possibly with "tau_cirrus" variable attached.
1441
+ rad : MetDataset
1442
+ Rad data with time shifted to account for accumulated values.
1443
+
1444
+ Raises
1445
+ ------
1446
+ If a previous version of pycontrails has already scaled the gridded humidity
1447
+ data.
1448
+ """
1449
+ # Check for remnants of previous scaling.
1450
+ if "_pycontrails_modified" in met["specific_humidity"].attrs:
1451
+ raise ValueError(
1452
+ "Specific humidity enhancement of the raw specific humidity values in "
1453
+ "the underlying met data is deprecated."
1454
+ )
1455
+
1456
+ if compute_tau_cirrus == "auto":
1457
+ # If met data is dask-backed, compute tau_cirrus
1458
+ compute_tau_cirrus = met.data["air_temperature"].chunks is not None
1459
+
1460
+ if "tau_cirrus" not in met.data:
1461
+ met.ensure_vars(Cocip.met_variables)
1462
+ if compute_tau_cirrus:
1463
+ met = add_tau_cirrus(met)
1464
+ else:
1465
+ met.ensure_vars(Cocip.processed_met_variables)
1466
+
1467
+ rad.ensure_vars(Cocip.rad_variables)
1468
+ rad = _process_rad(rad)
1469
+
1470
+ return met, rad
1471
+
1472
+
1473
+ def add_tau_cirrus(met: MetDataset) -> MetDataset:
1474
+ """Add "tau_cirrus" variable to weather data if it does not exist already.
1475
+
1476
+ Parameters
1477
+ ----------
1478
+ met : MetDataset
1479
+ Met pressure-level data.
1480
+
1481
+ Returns
1482
+ -------
1483
+ met : MetDataset
1484
+ Met data with "tau_cirrus" variable attached.
1485
+ """
1486
+ if "tau_cirrus" not in met.data:
1487
+ met.data["tau_cirrus"] = tau_cirrus.tau_cirrus(met)
1488
+ return met
1489
+
1490
+
1491
+ def _process_rad(rad: MetDataset) -> MetDataset:
1492
+ """Process radiation specific variables for model.
1493
+
1494
+ These variables are used to calculate the reflected solar radiation (RSR),
1495
+ and outgoing longwave radiation (OLR).
1496
+
1497
+ The time stamp is adjusted by ``shift_radiation_time`` to account for
1498
+ accumulated values being averaged over the time period.
1499
+
1500
+ Parameters
1501
+ ----------
1502
+ rad : MetDataset
1503
+ Rad single-level data
1504
+
1505
+ Returns
1506
+ -------
1507
+ MetDataset
1508
+ Rad data with time shifted.
1509
+
1510
+ Raises
1511
+ ------
1512
+ ValueError
1513
+ If a "radiation_accumulated" field is not found on ``rad.attrs``.
1514
+
1515
+ Notes
1516
+ -----
1517
+ - https://www.ecmwf.int/sites/default/files/elibrary/2015/18490-radiation-quantities-ecmwf-model-and-mars.pdf
1518
+ - https://confluence.ecmwf.int/pages/viewpage.action?pageId=155337784
1519
+ """
1520
+ # If the time coordinate has already been shifted, early return
1521
+ if "shift_radiation_time" in rad["time"].attrs:
1522
+ return rad
1523
+
1524
+ provider = rad.provider_attr
1525
+
1526
+ # Only shift ECMWF data -- exit for anything else
1527
+ # A warning is emitted upstream if the provider is not ECMWF or NCEP
1528
+ if provider != "ECMWF":
1529
+ return rad
1530
+
1531
+ dataset = rad.dataset_attr
1532
+ product = rad.product_attr
1533
+
1534
+ if dataset == "HRES":
1535
+ try:
1536
+ radiation_accumulated = rad.attrs["radiation_accumulated"]
1537
+ except KeyError as exc:
1538
+ msg = (
1539
+ "HRES data must have a boolean 'radiation_accumulated' attribute. "
1540
+ "This attribute is used to determine whether the radiation data "
1541
+ "has been accumulated over the time period. This is the case for "
1542
+ "HRES data taken from a common time of forecast with multiple "
1543
+ "forecast steps. If this is not the case, set the "
1544
+ "'radiation_accumulated' attribute to False."
1545
+ )
1546
+ raise ValueError(msg) from exc
1547
+ if radiation_accumulated:
1548
+ # Don't assume that radiation data is uniformly spaced in time
1549
+ # Instead, infer the appropriate time shift
1550
+ time_diff = rad.data["time"].diff("time", label="upper")
1551
+ time_shift = -time_diff / 2
1552
+
1553
+ # Keep the original attrs -- we need these later on
1554
+ old_attrs = {k: v.attrs for k, v in rad.data.items()}
1555
+
1556
+ # Also need to keep dataset-level attrs, which are lost
1557
+ # when dividing a Dataset by a DataArray
1558
+ old_rad_attrs = rad.data.attrs
1559
+
1560
+ # NOTE: Taking the diff will remove the first time step
1561
+ # This is typically what we want (forecast step 0 is all zeros)
1562
+ # But, if the data has been downselected for a particular Flight / Fleet,
1563
+ # we lose the first time step of the data.
1564
+ #
1565
+ # Other portions of the code convert HRES accumulated fluxes (J/m2)
1566
+ # to averaged fluxes (W/m2) assuming that accumulations are over
1567
+ # one hour. For those conversions to work correctly, must normalize
1568
+ # accumulations by number of hours between steps
1569
+ time_diff_h = time_diff / np.timedelta64(1, "h")
1570
+ rad.data = rad.data.diff("time", label="upper") / time_diff_h
1571
+ rad.data.attrs = old_rad_attrs
1572
+
1573
+ # Add back the original attrs
1574
+ for k, v in rad.data.items():
1575
+ v.attrs = old_attrs[k]
1576
+
1577
+ # Short-circuit to avoid idiot check
1578
+ rad.data = rad.data.assign_coords({"time": rad.data["time"] + time_shift})
1579
+ if np.unique(time_shift).size == 1:
1580
+ rad.data["time"].attrs["shift_radiation_time"] = str(time_shift.values[0])
1581
+ else:
1582
+ rad.data["time"].attrs["shift_radiation_time"] = "variable"
1583
+ return rad
1584
+
1585
+ else:
1586
+ shift_radiation_time = -np.timedelta64(30, "m")
1587
+
1588
+ elif dataset == "ERA5" and product == "ensemble":
1589
+ shift_radiation_time = -np.timedelta64(90, "m")
1590
+ else:
1591
+ shift_radiation_time = -np.timedelta64(30, "m")
1592
+
1593
+ # Do a final idiot check -- most likely, the time resolution of the data will
1594
+ # agree with the shift_radiation_time. If not, emit a warning. There could be
1595
+ # a false positive here if the data has been downsampled in time.
1596
+ logger.debug("Shifting rad time by %s", shift_radiation_time)
1597
+ rad_time_diff = np.diff(rad.data["time"])
1598
+ if not np.all(rad_time_diff / 2 == -shift_radiation_time):
1599
+ warnings.warn(
1600
+ f"Shifting radiation time dimension by unexpected interval {shift_radiation_time}. "
1601
+ f"The rad data has metadata indicating it is {product} ECMWF data. "
1602
+ f"This dataset should have time steps of {-2 * shift_radiation_time}."
1603
+ )
1604
+
1605
+ rad.data = rad.data.assign_coords({"time": rad.data["time"] + shift_radiation_time})
1606
+ rad.data["time"].attrs["shift_radiation_time"] = str(shift_radiation_time)
1607
+
1608
+ return rad
1609
+
1610
+
1611
+ def _eval_aircraft_performance(
1612
+ aircraft_performance: AircraftPerformance | None, flight: Flight
1613
+ ) -> Flight:
1614
+ """Evaluate the :class:`AircraftPerformance` model.
1615
+
1616
+ Parameters
1617
+ ----------
1618
+ aircraft_performance : AircraftPerformance | None
1619
+ Input aircraft performance model
1620
+ flight : Flight
1621
+ Flight to evaluate
1622
+
1623
+ Returns
1624
+ -------
1625
+ Flight
1626
+ Output from aircraft performance model
1627
+
1628
+ Raises
1629
+ ------
1630
+ ValueError
1631
+ If ``aircraft_performance`` is None
1632
+ """
1633
+
1634
+ ap_vars = {"wingspan", "engine_efficiency", "fuel_flow", "aircraft_mass"}
1635
+ missing = ap_vars.difference(flight).difference(flight.attrs)
1636
+ if not missing:
1637
+ return flight
1638
+
1639
+ if aircraft_performance is None:
1640
+ msg = (
1641
+ f"An AircraftPerformance model parameter is required if the flight does not contain "
1642
+ f"the following variables: {ap_vars}. This flight is missing: {missing}. "
1643
+ "Instantiate the Cocip model with an AircraftPerformance model. "
1644
+ "For example, 'Cocip(..., aircraft_performance=PSFlight(...))'."
1645
+ )
1646
+ raise ValueError(msg)
1647
+
1648
+ return aircraft_performance.eval(source=flight, copy_source=False)
1649
+
1650
+
1651
+ def _eval_emissions(emissions: Emissions, flight: Flight) -> Flight:
1652
+ """Evaluate the :class:`Emissions` model.
1653
+
1654
+ Parameters
1655
+ ----------
1656
+ emissions : Emissions
1657
+ Emissions model
1658
+ flight : Flight
1659
+ Flight to evaluate
1660
+
1661
+ Returns
1662
+ -------
1663
+ Flight
1664
+ Output from emissions model
1665
+ """
1666
+
1667
+ emissions_outputs = "nvpm_ei_n"
1668
+ if flight.ensure_vars(emissions_outputs, False):
1669
+ return flight
1670
+ return emissions.eval(source=flight, copy_source=False)
1671
+
1672
+
1673
+ def calc_continuous(contrail: GeoVectorDataset) -> None:
1674
+ """Calculate the continuous segments of this timestep.
1675
+
1676
+ Mutates parameter ``contrail`` in place by setting or updating the
1677
+ "continuous" variable.
1678
+
1679
+ Parameters
1680
+ ----------
1681
+ contrail : GeoVectorDataset
1682
+ GeoVectorDataset instance onto which "continuous" is set.
1683
+
1684
+ Raises
1685
+ ------
1686
+ ValueError
1687
+ If ``contrail`` is empty.
1688
+ """
1689
+ if not contrail:
1690
+ raise ValueError("Cannot calculate continuous on an empty contrail")
1691
+ same_flight = contrail["flight_id"][:-1] == contrail["flight_id"][1:]
1692
+ consecutive_waypoint = np.diff(contrail["waypoint"]) == 1
1693
+ continuous = np.empty(contrail.size, dtype=bool)
1694
+ continuous[:-1] = same_flight & consecutive_waypoint
1695
+ continuous[-1] = False # This fails if contrail is empty
1696
+ contrail.update(continuous=continuous) # overwrite continuous
1697
+
1698
+
1699
+ def calc_timestep_geometry(contrail: GeoVectorDataset) -> None:
1700
+ """Calculate contrail segment-specific properties.
1701
+
1702
+ Mutates parameter ``contrail`` in place by setting or updating the variables
1703
+ - "sin_a"
1704
+ - "cos_a"
1705
+ - "segment_length"
1706
+
1707
+ Any nan values in these variables are set to 0. This ensures any segment-based property
1708
+ derived from the segment geometry does not contribute to contrail energy forcing.
1709
+
1710
+ See Also
1711
+ --------
1712
+ - :func:`wind_shear.wind_shear_normal` to see how "sin_a" and "cos_a"
1713
+ are used to compute wind shear terms.
1714
+ - :func:`calc_timestep_contrail_evolution` to see how "segment_length" is used.
1715
+
1716
+ Parameters
1717
+ ----------
1718
+ contrail : GeoVectorDataset
1719
+ GeoVectorDataset instance onto which variables are set.
1720
+ """
1721
+ # get contrail waypoints
1722
+ longitude = contrail["longitude"]
1723
+ latitude = contrail["latitude"]
1724
+ altitude = contrail.altitude
1725
+ continuous = contrail["continuous"]
1726
+
1727
+ # calculate segment properties
1728
+ segment_length = geo.segment_length(longitude, latitude, altitude)
1729
+ sin_a, cos_a = geo.segment_angle(longitude, latitude)
1730
+
1731
+ # NOTE: Set all nan and discontinuous values to 0. With our current implementation, this
1732
+ # ensures degenerate or discontinuous segments do not contribute to contrail impact.
1733
+ # At the same time, this prevents nan values from propagating through evolving contrails.
1734
+ #
1735
+ # If we want to change this, take note of the following consequences.
1736
+ # nan values flow from one variable to another as follows:
1737
+ # segment_length
1738
+ # seg_ratio
1739
+ # sigma_yz
1740
+ # area_eff_2
1741
+ # iwc, n_ice_per_m, r_vol_ice
1742
+ # terminal_fall_speed
1743
+ # level [level advection]
1744
+ # u_wind, v_wind [interp]
1745
+ # longitude, latitude [advection]
1746
+ # Finally, at the next evolution step, the previous waypoint will accrue
1747
+ # a nan value after segment_length is recalculated
1748
+
1749
+ segment_length[~continuous] = 0.0
1750
+ sin_a[~continuous] = 0.0
1751
+ cos_a[~continuous] = 0.0
1752
+
1753
+ np.nan_to_num(segment_length, copy=False)
1754
+ np.nan_to_num(sin_a, copy=False)
1755
+ np.nan_to_num(cos_a, copy=False)
1756
+
1757
+ # override values on model
1758
+ contrail.update(segment_length=segment_length)
1759
+ contrail.update(sin_a=sin_a)
1760
+ contrail.update(cos_a=cos_a)
1761
+
1762
+
1763
+ def calc_timestep_meteorology(
1764
+ contrail: GeoVectorDataset,
1765
+ met: MetDataset,
1766
+ params: dict[str, Any],
1767
+ **interp_kwargs: Any,
1768
+ ) -> None:
1769
+ """Get and store meteorology parameters.
1770
+
1771
+ Mutates parameter ``contrail`` in place by setting or updating the variables
1772
+ - "sin_a"
1773
+ - "cos_a"
1774
+ - "segment_length"
1775
+
1776
+ Parameters
1777
+ ----------
1778
+ contrail : GeoVectorDataset
1779
+ GeoVectorDataset object onto which meterology variables are attached.
1780
+ met : MetDataset
1781
+ MetDataset with meteorology data variables.
1782
+ params : dict[str, Any]
1783
+ Cocip model ``params``.
1784
+ **interp_kwargs : Any
1785
+ Cocip model ``interp_kwargs``.
1786
+ """
1787
+ # get contrail geometry
1788
+ sin_a = contrail["sin_a"]
1789
+ cos_a = contrail["cos_a"]
1790
+
1791
+ # get standard met parameters for timestep
1792
+ air_pressure = contrail.air_pressure
1793
+ air_temperature = contrail["air_temperature"]
1794
+ u_wind = interpolate_met(met, contrail, "eastward_wind", "u_wind", **interp_kwargs)
1795
+ v_wind = interpolate_met(met, contrail, "northward_wind", "v_wind", **interp_kwargs)
1796
+ interpolate_met(
1797
+ met,
1798
+ contrail,
1799
+ "lagrangian_tendency_of_air_pressure",
1800
+ "vertical_velocity",
1801
+ **interp_kwargs,
1802
+ )
1803
+ interpolate_met(met, contrail, "tau_cirrus", **interp_kwargs)
1804
+
1805
+ # get the pressure level `dz_m` lower than element pressure
1806
+ air_pressure_lower = thermo.pressure_dz(air_temperature, air_pressure, params["dz_m"])
1807
+
1808
+ # get meteorology at contrail waypoints interpolated at the pressure level `air_pressure_lower`
1809
+ level_lower = air_pressure_lower / 100.0 # Pa -> hPa
1810
+
1811
+ # if met is already interpolated, this will automatically skip interpolation
1812
+ air_temperature_lower = interpolate_met(
1813
+ met,
1814
+ contrail,
1815
+ "air_temperature",
1816
+ "air_temperature_lower",
1817
+ level=level_lower,
1818
+ **interp_kwargs,
1819
+ )
1820
+ u_wind_lower = interpolate_met(
1821
+ met,
1822
+ contrail,
1823
+ "eastward_wind",
1824
+ "u_wind_lower",
1825
+ level=level_lower,
1826
+ **interp_kwargs,
1827
+ )
1828
+ v_wind_lower = interpolate_met(
1829
+ met,
1830
+ contrail,
1831
+ "northward_wind",
1832
+ "v_wind_lower",
1833
+ level=level_lower,
1834
+ **interp_kwargs,
1835
+ )
1836
+
1837
+ # Temperature gradient
1838
+ dT_dz = thermo.T_potential_gradient(
1839
+ air_temperature,
1840
+ air_pressure,
1841
+ air_temperature_lower,
1842
+ air_pressure_lower,
1843
+ params["dz_m"],
1844
+ )
1845
+
1846
+ # wind shear
1847
+ ds_dz = wind_shear.wind_shear(u_wind, u_wind_lower, v_wind, v_wind_lower, params["dz_m"])
1848
+
1849
+ # wind shear normal
1850
+ dsn_dz = wind_shear.wind_shear_normal(
1851
+ u_wind_top=u_wind,
1852
+ u_wind_btm=u_wind_lower,
1853
+ v_wind_top=v_wind,
1854
+ v_wind_btm=v_wind_lower,
1855
+ cos_a=cos_a,
1856
+ sin_a=sin_a,
1857
+ dz=params["dz_m"],
1858
+ )
1859
+
1860
+ # store values on contrail model
1861
+ contrail.update(dT_dz=dT_dz)
1862
+ contrail.update(ds_dz=ds_dz)
1863
+ contrail.update(dsn_dz=dsn_dz)
1864
+
1865
+
1866
+ def calc_shortwave_radiation(
1867
+ rad: MetDataset,
1868
+ vector: GeoVectorDataset,
1869
+ **interp_kwargs: Any,
1870
+ ) -> None:
1871
+ """Calculate shortwave radiation variables.
1872
+
1873
+ Calculates theoretical incident (``sdr``) and
1874
+ reflected shortwave radiation (``rsr``) from the radiation data provided.
1875
+
1876
+ Mutates input ``vector`` to include ``"sdr"`` and ``"rsr"``
1877
+ keys with calculated SDR, RSR values, respectively.
1878
+
1879
+ Parameters
1880
+ ----------
1881
+ rad : MetDataset
1882
+ Radiation data
1883
+ vector : GeoVectorDataset
1884
+ Flight or GeoVectorDataset instance
1885
+ **interp_kwargs : Any
1886
+ Interpolation keyword arguments
1887
+
1888
+ Raises
1889
+ ------
1890
+ ValueError
1891
+ If ``rad`` does not contain ``"toa_upward_shortwave_flux"`` or
1892
+ ``"top_net_solar_radiation"`` variable.
1893
+
1894
+ Notes
1895
+ -----
1896
+ In accordance with the original CoCiP Fortran algorithm,
1897
+ the SDR is set to 0 when the cosine of the solar zenith angle is below 0.01.
1898
+
1899
+ See Also
1900
+ --------
1901
+ :func:`geo.solar_direct_radiation`
1902
+ """
1903
+ if "sdr" in vector and "rsr" in vector:
1904
+ return
1905
+
1906
+ try:
1907
+ sdr = vector["sdr"]
1908
+ except KeyError:
1909
+ # calculate instantaneous theoretical solar direct radiation based on geo position and time
1910
+ longitude = vector["longitude"]
1911
+ latitude = vector["latitude"]
1912
+ time = vector["time"]
1913
+ sdr = geo.solar_direct_radiation(longitude, latitude, time, threshold_cos_sza=0.01)
1914
+ vector["sdr"] = sdr
1915
+
1916
+ # GFS contains RSR (toa_upward_shortwave_flux) variable directly
1917
+ gfs_key = "toa_upward_shortwave_flux"
1918
+ if gfs_key in rad:
1919
+ interpolate_met(rad, vector, gfs_key, "rsr", **interp_kwargs)
1920
+ return
1921
+
1922
+ ecmwf_key = "top_net_solar_radiation"
1923
+ if ecmwf_key not in rad:
1924
+ msg = f"'rad' data must contain either '{gfs_key}' or '{ecmwf_key}' (ECMWF) variable."
1925
+ raise ValueError(msg)
1926
+
1927
+ # ECMWF contains "top_net_solar_radiation" which is SDR - RSR
1928
+ tnsr = interpolate_met(rad, vector, ecmwf_key, **interp_kwargs)
1929
+ tnsr = _rad_accumulation_to_average_instantaneous(rad, ecmwf_key, tnsr)
1930
+ vector.update({ecmwf_key: tnsr})
1931
+
1932
+ vector["rsr"] = np.maximum(sdr - tnsr, 0.0)
1933
+
1934
+
1935
+ def calc_outgoing_longwave_radiation(
1936
+ rad: MetDataset,
1937
+ vector: GeoVectorDataset,
1938
+ **interp_kwargs: Any,
1939
+ ) -> None:
1940
+ """Calculate outgoing longwave radiation (``olr``) from the radiation data provided.
1941
+
1942
+ Mutates input ``vector`` to include ``"olr"`` key with calculated OLR values.
1943
+
1944
+ Parameters
1945
+ ----------
1946
+ rad : MetDataset
1947
+ Radiation data
1948
+ vector : GeoVectorDataset
1949
+ Flight or GeoVectorDataset instance
1950
+ **interp_kwargs : Any
1951
+ Interpolation keyword arguments
1952
+
1953
+ Raises
1954
+ ------
1955
+ ValueError
1956
+ If ``rad`` does not contain a ``"toa_upward_longwave_flux"``
1957
+ or ``"top_net_thermal_radiation"`` variable.
1958
+ """
1959
+
1960
+ if "olr" in vector:
1961
+ return None
1962
+
1963
+ # GFS contains OLR (toa_upward_longwave_flux) variable directly
1964
+ gfs_key = "toa_upward_longwave_flux"
1965
+ if gfs_key in rad:
1966
+ interpolate_met(rad, vector, gfs_key, "olr", **interp_kwargs)
1967
+ return
1968
+
1969
+ # ECMWF contains "top_net_thermal_radiation" which is -1 * OLR
1970
+ ecmwf_key = "top_net_thermal_radiation"
1971
+ if ecmwf_key not in rad:
1972
+ msg = f"'rad' data must contain either '{gfs_key}' or '{ecmwf_key}' (ECMWF) variable."
1973
+ raise ValueError(msg)
1974
+
1975
+ tntr = interpolate_met(rad, vector, ecmwf_key, **interp_kwargs)
1976
+ tntr = _rad_accumulation_to_average_instantaneous(rad, ecmwf_key, tntr)
1977
+ vector.update({ecmwf_key: tntr})
1978
+
1979
+ vector["olr"] = np.maximum(-tntr, 0.0)
1980
+
1981
+
1982
+ def calc_radiative_properties(contrail: GeoVectorDataset, params: dict[str, Any]) -> None:
1983
+ """Calculate radiative properties for contrail.
1984
+
1985
+ This function is used by both :class:`Cocip` and :class`CocipGrid`.
1986
+
1987
+ Mutates original `contrail` parameter with additional keys:
1988
+ - "rf_sw"
1989
+ - "rf_lw"
1990
+ - "rf_net"
1991
+
1992
+ Parameters
1993
+ ----------
1994
+ contrail : GeoVectorDataset
1995
+ Grid points already interpolated against met and rad data. In particular,
1996
+ the variables
1997
+ - "air_temperature"
1998
+ - "r_ice_vol"
1999
+ - "tau_contrail"
2000
+ - "tau_cirrus"
2001
+ - "olr"
2002
+ - "rsr"
2003
+ - "sdr"
2004
+
2005
+ are required on the parameter `contrail`.
2006
+ params : dict[str, Any]
2007
+ Model parameters
2008
+ """
2009
+ time = contrail["time"]
2010
+ air_temperature = contrail["air_temperature"]
2011
+
2012
+ if params["radiative_heating_effects"]:
2013
+ air_temperature += contrail["cumul_heat"]
2014
+
2015
+ r_ice_vol = contrail["r_ice_vol"]
2016
+ tau_contrail = contrail["tau_contrail"]
2017
+ tau_cirrus_ = contrail["tau_cirrus"]
2018
+
2019
+ # calculate solar constant
2020
+ theta_rad = geo.orbital_position(time)
2021
+ sd0 = geo.solar_constant(theta_rad)
2022
+
2023
+ # radiation dataset with contrail waypoints at timestep
2024
+ sdr = contrail["sdr"]
2025
+ rsr = contrail["rsr"]
2026
+ olr = contrail["olr"]
2027
+
2028
+ # radiation calculations
2029
+ r_vol_um = r_ice_vol * 1e6
2030
+ habit_weights = radiative_forcing.habit_weights(
2031
+ r_vol_um, params["habit_distributions"], params["radius_threshold_um"]
2032
+ )
2033
+ rf_lw = radiative_forcing.longwave_radiative_forcing(
2034
+ r_vol_um, olr, air_temperature, tau_contrail, tau_cirrus_, habit_weights
2035
+ )
2036
+ rf_sw = radiative_forcing.shortwave_radiative_forcing(
2037
+ r_vol_um, sdr, rsr, sd0, tau_contrail, tau_cirrus_, habit_weights
2038
+ )
2039
+
2040
+ # scale RF by enhancement factors
2041
+ rf_lw_scaled = rf_lw * params["rf_lw_enhancement_factor"]
2042
+ rf_sw_scaled = rf_sw * params["rf_sw_enhancement_factor"]
2043
+ rf_net = radiative_forcing.net_radiative_forcing(rf_lw_scaled, rf_sw_scaled)
2044
+
2045
+ # store values on contrail
2046
+ contrail["rf_sw"] = rf_sw
2047
+ contrail["rf_lw"] = rf_lw
2048
+ contrail["rf_net"] = rf_net
2049
+
2050
+
2051
+ def calc_contrail_properties(
2052
+ contrail: GeoVectorDataset,
2053
+ effective_vertical_resolution: float | npt.NDArray[np.float64],
2054
+ wind_shear_enhancement_exponent: float | npt.NDArray[np.float64],
2055
+ sedimentation_impact_factor: float | npt.NDArray[np.float64],
2056
+ radiative_heating_effects: bool,
2057
+ ) -> None:
2058
+ """Calculate geometric and ice-related properties of contrail.
2059
+
2060
+ This function is used by both :class:`Cocip` and :class`CocipGrid`.
2061
+
2062
+ This function modifies parameter `contrail` in place:
2063
+ - Mutates contrail data variables:
2064
+ - "ds_dz"
2065
+ - "dsn_dz"
2066
+ - Attaches additional variables:
2067
+ - "area_eff"
2068
+ - "plume_mass_per_m"
2069
+ - "r_ice_vol"
2070
+ - "terminal_fall_speed"
2071
+ - "diffuse_h"
2072
+ - "diffuse_v"
2073
+ - "n_ice_per_vol"
2074
+ - "tau_contrail"
2075
+ - "dn_dt_agg"
2076
+ - "dn_dt_turb"
2077
+
2078
+ Parameters
2079
+ ----------
2080
+ contrail : GeoVectorDataset
2081
+ Grid points with many precomputed keys.
2082
+ effective_vertical_resolution : float | npt.NDArray[np.float64]
2083
+ Passed into :func:`wind_shear.wind_shear_enhancement_factor`.
2084
+ wind_shear_enhancement_exponent : float | npt.NDArray[np.float64]
2085
+ Passed into :func:`wind_shear.wind_shear_enhancement_factor`.
2086
+ sedimentation_impact_factor: float | npt.NDArray[np.float64]
2087
+ Passed into `contrail_properties.vertical_diffusivity`.
2088
+ radiative_heating_effects: bool
2089
+ Include radiative heating effects on contrail cirrus properties.
2090
+ """
2091
+ time = contrail["time"]
2092
+ iwc = contrail["iwc"]
2093
+ depth = contrail["depth"]
2094
+ width = contrail["width"]
2095
+ n_ice_per_m = contrail["n_ice_per_m"]
2096
+ t_cirrus_ = contrail["tau_cirrus"]
2097
+
2098
+ # get required meteorology
2099
+ air_temperature = contrail["air_temperature"]
2100
+ air_pressure = contrail.air_pressure
2101
+ rhi = contrail["rhi"]
2102
+ rho_air = contrail["rho_air"]
2103
+ dT_dz = contrail["dT_dz"]
2104
+ ds_dz = contrail["ds_dz"]
2105
+ dsn_dz = contrail["dsn_dz"]
2106
+ sigma_yz = contrail["sigma_yz"]
2107
+
2108
+ if radiative_heating_effects:
2109
+ air_temperature += contrail["cumul_heat"]
2110
+
2111
+ # get required radiation
2112
+ theta_rad = geo.orbital_position(time)
2113
+ sd0 = geo.solar_constant(theta_rad)
2114
+ sdr = contrail["sdr"]
2115
+ rsr = contrail["rsr"]
2116
+ olr = contrail["olr"]
2117
+
2118
+ # shear enhancements
2119
+ shear_enhancement = wind_shear.wind_shear_enhancement_factor(
2120
+ contrail_depth=depth,
2121
+ effective_vertical_resolution=effective_vertical_resolution,
2122
+ wind_shear_enhancement_exponent=wind_shear_enhancement_exponent,
2123
+ )
2124
+ ds_dz = ds_dz * shear_enhancement
2125
+ dsn_dz = dsn_dz * shear_enhancement
2126
+
2127
+ # effective area
2128
+ area_eff = contrail_properties.plume_effective_cross_sectional_area(width, depth, sigma_yz)
2129
+ depth_eff = contrail_properties.plume_effective_depth(width, area_eff)
2130
+
2131
+ # ice particles
2132
+ n_ice_per_vol = contrail_properties.ice_particle_number_per_volume_of_plume(
2133
+ n_ice_per_m, area_eff
2134
+ )
2135
+ n_ice_per_kg_air = contrail_properties.ice_particle_number_per_mass_of_air(
2136
+ n_ice_per_vol, rho_air
2137
+ )
2138
+ plume_mass_per_m = contrail_properties.plume_mass_per_distance(area_eff, rho_air)
2139
+ r_ice_vol = contrail_properties.ice_particle_volume_mean_radius(iwc, n_ice_per_kg_air)
2140
+ tau_contrail = contrail_properties.contrail_optical_depth(r_ice_vol, n_ice_per_m, width)
2141
+
2142
+ terminal_fall_speed = contrail_properties.ice_particle_terminal_fall_speed(
2143
+ air_pressure, air_temperature, r_ice_vol
2144
+ )
2145
+ diffuse_h = contrail_properties.horizontal_diffusivity(ds_dz, depth)
2146
+
2147
+ if radiative_heating_effects:
2148
+ heat_rate = radiative_heating.heating_rate(
2149
+ air_temperature=air_temperature,
2150
+ rhi=rhi,
2151
+ rho_air=rho_air,
2152
+ r_ice_vol=r_ice_vol,
2153
+ depth_eff=depth_eff,
2154
+ tau_contrail=tau_contrail,
2155
+ tau_cirrus=t_cirrus_,
2156
+ sd0=sd0,
2157
+ sdr=sdr,
2158
+ rsr=rsr,
2159
+ olr=olr,
2160
+ )
2161
+
2162
+ cumul_differential_heat = contrail["cumul_differential_heat"]
2163
+ d_heat_rate = radiative_heating.differential_heating_rate(
2164
+ air_temperature=air_temperature,
2165
+ rhi=rhi,
2166
+ rho_air=rho_air,
2167
+ r_ice_vol=r_ice_vol,
2168
+ depth_eff=depth_eff,
2169
+ tau_contrail=tau_contrail,
2170
+ tau_cirrus=t_cirrus_,
2171
+ sd0=sd0,
2172
+ sdr=sdr,
2173
+ rsr=rsr,
2174
+ olr=olr,
2175
+ )
2176
+
2177
+ eff_heat_rate = radiative_heating.effective_heating_rate(
2178
+ d_heat_rate, cumul_differential_heat, dT_dz, depth
2179
+ )
2180
+ else:
2181
+ eff_heat_rate = None
2182
+
2183
+ diffuse_v = contrail_properties.vertical_diffusivity(
2184
+ air_pressure=air_pressure,
2185
+ air_temperature=air_temperature,
2186
+ dT_dz=dT_dz,
2187
+ depth_eff=depth_eff,
2188
+ terminal_fall_speed=terminal_fall_speed,
2189
+ sedimentation_impact_factor=sedimentation_impact_factor,
2190
+ eff_heat_rate=eff_heat_rate,
2191
+ )
2192
+
2193
+ dn_dt_agg = contrail_properties.particle_losses_aggregation(
2194
+ r_ice_vol, terminal_fall_speed, area_eff
2195
+ )
2196
+ dn_dt_turb = contrail_properties.particle_losses_turbulence(
2197
+ width, depth, depth_eff, diffuse_h, diffuse_v
2198
+ )
2199
+
2200
+ # Set properties to model
2201
+ contrail.update(ds_dz=ds_dz)
2202
+ contrail.update(dsn_dz=dsn_dz)
2203
+ contrail.update(area_eff=area_eff)
2204
+ contrail.update(plume_mass_per_m=plume_mass_per_m)
2205
+ contrail.update(r_ice_vol=r_ice_vol)
2206
+ contrail.update(terminal_fall_speed=terminal_fall_speed)
2207
+ contrail.update(diffuse_h=diffuse_h)
2208
+ contrail.update(diffuse_v=diffuse_v)
2209
+ contrail.update(n_ice_per_vol=n_ice_per_vol)
2210
+ contrail.update(tau_contrail=tau_contrail)
2211
+ contrail.update(dn_dt_agg=dn_dt_agg)
2212
+ contrail.update(dn_dt_turb=dn_dt_turb)
2213
+ if radiative_heating_effects:
2214
+ contrail.update(heat_rate=heat_rate, d_heat_rate=d_heat_rate)
2215
+
2216
+
2217
+ def calc_timestep_contrail_evolution(
2218
+ met: MetDataset,
2219
+ rad: MetDataset,
2220
+ contrail_1: GeoVectorDataset,
2221
+ time_2: np.datetime64,
2222
+ params: dict[str, Any],
2223
+ **interp_kwargs: Any,
2224
+ ) -> GeoVectorDataset:
2225
+ """Calculate the contrail evolution across timestep (t1 -> t2).
2226
+
2227
+ Note the variable suffix "_1" is used to reference the current time
2228
+ and the suffix "_2" is used to refer to the time at the next timestep.
2229
+
2230
+ Parameters
2231
+ ----------
2232
+ met : MetDataset
2233
+ Meteorology data
2234
+ rad : MetDataset
2235
+ Radiation data
2236
+ contrail_1 : GeoVectorDataset
2237
+ Contrail waypoints at current timestep (1)
2238
+ time_2 : np.datetime64
2239
+ Time at the end of the evolution step (2)
2240
+ params : dict[str, Any]
2241
+ Model parameters
2242
+ **interp_kwargs : Any
2243
+ Interpolation keyword arguments
2244
+
2245
+ Returns
2246
+ -------
2247
+ GeoVectorDataset
2248
+ The contrail evolved to ``time_2``.
2249
+ """
2250
+
2251
+ # get lat/lon for current timestep (t1)
2252
+ longitude_1 = contrail_1["longitude"]
2253
+ latitude_1 = contrail_1["latitude"]
2254
+ level_1 = contrail_1.level
2255
+ time_1 = contrail_1["time"]
2256
+
2257
+ # get contrail_1 geometry
2258
+ segment_length_1 = contrail_1["segment_length"]
2259
+ width_1 = contrail_1["width"]
2260
+ depth_1 = contrail_1["depth"]
2261
+
2262
+ # get required met values for evolution calculations
2263
+ q_sat_1 = contrail_1["q_sat"]
2264
+ rho_air_1 = contrail_1["rho_air"]
2265
+ u_wind_1 = contrail_1["u_wind"]
2266
+ v_wind_1 = contrail_1["v_wind"]
2267
+
2268
+ specific_humidity_1 = contrail_1["specific_humidity"]
2269
+ vertical_velocity_1 = contrail_1["vertical_velocity"]
2270
+ iwc_1 = contrail_1["iwc"]
2271
+
2272
+ # get required contrail_1 properties
2273
+ sigma_yz_1 = contrail_1["sigma_yz"]
2274
+ dsn_dz_1 = contrail_1["dsn_dz"]
2275
+ terminal_fall_speed_1 = contrail_1["terminal_fall_speed"]
2276
+ diffuse_h_1 = contrail_1["diffuse_h"]
2277
+ diffuse_v_1 = contrail_1["diffuse_v"]
2278
+ plume_mass_per_m_1 = contrail_1["plume_mass_per_m"]
2279
+ dn_dt_agg_1 = contrail_1["dn_dt_agg"]
2280
+ dn_dt_turb_1 = contrail_1["dn_dt_turb"]
2281
+ n_ice_per_m_1 = contrail_1["n_ice_per_m"]
2282
+
2283
+ # get contrail_1 radiative properties
2284
+ rf_net_1 = contrail_1["rf_net"]
2285
+
2286
+ # initialize new timestep with evolved coordinates
2287
+ # assume waypoints are the same to start
2288
+ waypoint_2 = contrail_1["waypoint"]
2289
+ formation_time_2 = contrail_1["formation_time"]
2290
+ time_2_array = np.full_like(time_1, time_2)
2291
+ dt = time_2_array - time_1
2292
+
2293
+ # get new contrail location & segment properties after t_step
2294
+ longitude_2 = geo.advect_longitude(longitude_1, latitude_1, u_wind_1, dt)
2295
+ latitude_2 = geo.advect_latitude(latitude_1, v_wind_1, dt)
2296
+ level_2 = geo.advect_level(level_1, vertical_velocity_1, rho_air_1, terminal_fall_speed_1, dt)
2297
+ altitude_2 = units.pl_to_m(level_2)
2298
+
2299
+ contrail_2 = GeoVectorDataset(
2300
+ {
2301
+ "waypoint": waypoint_2,
2302
+ "flight_id": contrail_1["flight_id"],
2303
+ "formation_time": formation_time_2,
2304
+ "time": time_2_array,
2305
+ "age": time_2_array - formation_time_2,
2306
+ "longitude": longitude_2,
2307
+ "latitude": latitude_2,
2308
+ "altitude": altitude_2,
2309
+ "level": level_2,
2310
+ },
2311
+ copy=False,
2312
+ )
2313
+ intersection = contrail_2.coords_intersect_met(met)
2314
+ if not np.any(intersection):
2315
+ warnings.warn(
2316
+ f"At time {time_2}, the contrail has no intersection with the met data. "
2317
+ "This is likely due to the contrail being advected outside the met domain."
2318
+ )
2319
+
2320
+ # Update cumulative radiative heating energy absorbed by the contrail
2321
+ # This will always be zero if radiative_heating_effects is not activated in cocip_params
2322
+ if params["radiative_heating_effects"]:
2323
+ dt_sec = dt / np.timedelta64(1, "s")
2324
+ heat_rate_1 = contrail_1["heat_rate"]
2325
+ cumul_heat = contrail_1["cumul_heat"]
2326
+ cumul_heat += heat_rate_1 * dt_sec
2327
+ cumul_heat.clip(max=1.5, out=cumul_heat) # Constrain additional heat to 1.5 K as precaution
2328
+ contrail_2["cumul_heat"] = cumul_heat
2329
+
2330
+ d_heat_rate_1 = contrail_1["d_heat_rate"]
2331
+ cumul_differential_heat = contrail_1["cumul_differential_heat"]
2332
+ cumul_differential_heat += -d_heat_rate_1 * dt_sec
2333
+ contrail_2["cumul_differential_heat"] = cumul_differential_heat
2334
+
2335
+ # Attach a few more artifacts for disabled filtering
2336
+ if not params["filter_sac"]:
2337
+ contrail_2["sac"] = contrail_1["sac"]
2338
+ if not params["filter_initially_persistent"]:
2339
+ contrail_2["initially_persistent"] = contrail_1["initially_persistent"]
2340
+ if params["persistent_buffer"] is not None:
2341
+ contrail_2["end_of_life"] = contrail_1["end_of_life"]
2342
+
2343
+ # calculate initial contrail properties for the next timestep
2344
+ calc_continuous(contrail_2)
2345
+ calc_timestep_geometry(contrail_2)
2346
+
2347
+ # get next segment lengths
2348
+ segment_length_2 = contrail_2["segment_length"]
2349
+
2350
+ # new contrail dimensions
2351
+ seg_ratio_2 = contrail_properties.segment_length_ratio(segment_length_1, segment_length_2)
2352
+ sigma_yy_2, sigma_zz_2, sigma_yz_2 = contrail_properties.plume_temporal_evolution(
2353
+ width_1,
2354
+ depth_1,
2355
+ sigma_yz_1,
2356
+ dsn_dz_1,
2357
+ diffuse_h_1,
2358
+ diffuse_v_1,
2359
+ seg_ratio_2,
2360
+ dt,
2361
+ max_depth=params["max_depth"],
2362
+ )
2363
+ width_2, depth_2 = contrail_properties.new_contrail_dimensions(sigma_yy_2, sigma_zz_2)
2364
+
2365
+ # store data back in model to support next time step calculations
2366
+ contrail_2["width"] = width_2
2367
+ contrail_2["depth"] = depth_2
2368
+ contrail_2["sigma_yz"] = sigma_yz_2
2369
+
2370
+ # new contrail meteorology parameters
2371
+ air_temperature_2 = interpolate_met(met, contrail_2, "air_temperature", **interp_kwargs)
2372
+
2373
+ if params["radiative_heating_effects"]:
2374
+ air_temperature_2 += contrail_2["cumul_heat"]
2375
+
2376
+ interpolate_met(met, contrail_2, "specific_humidity", **interp_kwargs)
2377
+
2378
+ humidity_scaling = params["humidity_scaling"]
2379
+ if humidity_scaling is not None:
2380
+ humidity_scaling.eval(contrail_2, copy_source=False)
2381
+ else:
2382
+ contrail_2["air_pressure"] = contrail_2.air_pressure
2383
+ contrail_2["rhi"] = thermo.rhi(
2384
+ contrail_2["specific_humidity"], air_temperature_2, contrail_2["air_pressure"]
2385
+ )
2386
+
2387
+ air_pressure_2 = contrail_2["air_pressure"]
2388
+ specific_humidity_2 = contrail_2["specific_humidity"]
2389
+ rho_air_2 = thermo.rho_d(air_temperature_2, air_pressure_2)
2390
+ q_sat_2 = thermo.q_sat_ice(air_temperature_2, air_pressure_2)
2391
+
2392
+ # store data back in model to support next ``calc_timestep_meteorology``
2393
+ contrail_2["rho_air"] = rho_air_2
2394
+ contrail_2["q_sat"] = q_sat_2
2395
+
2396
+ # New contrail ice particle mass and number
2397
+ area_eff_2 = contrail_properties.new_effective_area_from_sigma(
2398
+ sigma_yy_2, sigma_zz_2, sigma_yz_2
2399
+ )
2400
+ plume_mass_per_m_2 = contrail_properties.plume_mass_per_distance(area_eff_2, rho_air_2)
2401
+ iwc_2 = contrail_properties.new_ice_water_content(
2402
+ iwc_1,
2403
+ specific_humidity_1,
2404
+ specific_humidity_2,
2405
+ q_sat_1,
2406
+ q_sat_2,
2407
+ plume_mass_per_m_1,
2408
+ plume_mass_per_m_2,
2409
+ )
2410
+ n_ice_per_m_2 = contrail_properties.new_ice_particle_number(
2411
+ n_ice_per_m_1, dn_dt_agg_1, dn_dt_turb_1, seg_ratio_2, dt
2412
+ )
2413
+
2414
+ contrail_2["n_ice_per_m"] = n_ice_per_m_2
2415
+ contrail_2["iwc"] = iwc_2
2416
+
2417
+ # calculate next timestep meteorology, contrail, and radiative properties
2418
+ calc_timestep_meteorology(contrail_2, met, params, **interp_kwargs)
2419
+
2420
+ # Intersect with rad dataset
2421
+ calc_shortwave_radiation(rad, contrail_2, **interp_kwargs)
2422
+ calc_outgoing_longwave_radiation(rad, contrail_2, **interp_kwargs)
2423
+ calc_contrail_properties(
2424
+ contrail_2,
2425
+ params["effective_vertical_resolution"],
2426
+ params["wind_shear_enhancement_exponent"],
2427
+ params["sedimentation_impact_factor"],
2428
+ params["radiative_heating_effects"],
2429
+ )
2430
+ calc_radiative_properties(contrail_2, params)
2431
+
2432
+ # get properties to measure persistence
2433
+ latitude_2 = contrail_2["latitude"]
2434
+ altitude_2 = contrail_2.altitude
2435
+ age_2 = contrail_2["age"]
2436
+ tau_contrail_2 = contrail_2["tau_contrail"]
2437
+ n_ice_per_vol_2 = contrail_2["n_ice_per_vol"]
2438
+
2439
+ # calculate next persistent
2440
+ persistent_2 = contrail_2["persistent"] = contrail_properties.contrail_persistent(
2441
+ latitude=latitude_2,
2442
+ altitude=altitude_2,
2443
+ segment_length=segment_length_2,
2444
+ age=age_2,
2445
+ tau_contrail=tau_contrail_2,
2446
+ n_ice_per_m3=n_ice_per_vol_2,
2447
+ params=params,
2448
+ )
2449
+
2450
+ # Get energy forcing by looking forward to the next time step radiative forcing
2451
+ rf_net_2 = contrail_2["rf_net"]
2452
+ energy_forcing_2 = contrail_properties.energy_forcing(
2453
+ rf_net_t1=rf_net_1,
2454
+ rf_net_t2=rf_net_2,
2455
+ width_t1=width_1,
2456
+ width_t2=width_2,
2457
+ seg_length_t2=segment_length_2,
2458
+ dt=dt,
2459
+ )
2460
+
2461
+ # NOTE: Because of our geometry-continuity conventions, any waypoint without
2462
+ # continuity automatically has 0 EF. This is because calc_timestep_geometry
2463
+ # sets the segment_length entries to 0, thereby eliminating EF.
2464
+ # If we change conventions in calc_timestep_geometry, we may want to
2465
+ # explicitly set the discontinuous entries here to 0. We do this for age
2466
+ # here as well. It's somewhat of a hack, but it ensures nonzero contrail_age
2467
+ # is consistent with nonzero EF.
2468
+ contrail_2["ef"] = energy_forcing_2
2469
+ contrail_2["age"][energy_forcing_2 == 0.0] = np.timedelta64(0, "ns")
2470
+
2471
+ if params["compute_atr20"]:
2472
+ contrail_2["global_yearly_mean_rf"] = (
2473
+ contrail_2["ef"] / constants.surface_area_earth / constants.seconds_per_year
2474
+ )
2475
+ contrail_2["atr20"] = (
2476
+ params["global_rf_to_atr20_factor"] * contrail_2["global_yearly_mean_rf"]
2477
+ )
2478
+
2479
+ # filter contrail_2 by persistent waypoints, if any continuous segments are left
2480
+ logger.debug(
2481
+ "Fraction of waypoints surviving: %s / %s",
2482
+ persistent_2.sum(),
2483
+ persistent_2.size,
2484
+ )
2485
+
2486
+ if (buff := params["persistent_buffer"]) is not None:
2487
+ # Here mask gets waypoints that are just now losing persistence
2488
+ mask = (~persistent_2) & np.isnat(contrail_2["end_of_life"])
2489
+ contrail_2["end_of_life"][mask] = time_2
2490
+
2491
+ # Keep waypoints that are still persistent, which is determined by filt2
2492
+ # And waypoints within the persistent buffer, which is determined by filt1
2493
+ # So we only drop waypoints that are outside of the persistent buffer
2494
+ filt1 = contrail_2["time"] - contrail_2["end_of_life"] < buff
2495
+ filt2 = np.isnat(contrail_2["end_of_life"])
2496
+ filt = filt1 | filt2
2497
+ logger.debug(
2498
+ "Fraction of waypoints surviving with buffer %s: %s / %s",
2499
+ buff,
2500
+ filt.sum(),
2501
+ filt.size,
2502
+ )
2503
+ return contrail_2.filter(filt)
2504
+
2505
+ # filter persistent contrails
2506
+ final_contrail = contrail_2.filter(persistent_2)
2507
+
2508
+ # recalculate continuous contrail segments
2509
+ # and set EF, age for any newly discontinuous segments to 0
2510
+ if final_contrail:
2511
+ calc_continuous(final_contrail)
2512
+ continuous = final_contrail["continuous"]
2513
+ final_contrail["ef"][~continuous] = 0.0
2514
+ final_contrail["age"][~continuous] = np.timedelta64(0, "ns")
2515
+ if params["compute_atr20"]:
2516
+ final_contrail["global_yearly_mean_rf"][~continuous] = 0.0
2517
+ final_contrail["atr20"][~continuous] = 0.0
2518
+ return final_contrail
2519
+
2520
+
2521
+ def _rad_accumulation_to_average_instantaneous(
2522
+ rad: MetDataset,
2523
+ name: str,
2524
+ arr: npt.NDArray[np.float64],
2525
+ ) -> npt.NDArray[np.float64]:
2526
+ """Convert from radiation accumulation to average instantaneous values.
2527
+
2528
+ .. versionadded:: 0.48.0
2529
+
2530
+ Parameters
2531
+ ----------
2532
+ rad : MetDataset
2533
+ Radiation data
2534
+ name : str
2535
+ Variable name
2536
+ arr : npt.NDArray[np.float64]
2537
+ Array of values already interpolated from ``rad``
2538
+
2539
+ Returns
2540
+ -------
2541
+ npt.NDArray[np.float64]
2542
+ Array of values converted from accumulation to average instantaneous values
2543
+
2544
+ Raises
2545
+ ------
2546
+ KeyError
2547
+ If units are not provided on ``rad``.
2548
+ ValueError
2549
+ If unknown units are provided on ``rad``.
2550
+ """
2551
+ mda = rad[name]
2552
+ try:
2553
+ unit = mda.attrs["units"]
2554
+ except KeyError as e:
2555
+ msg = (
2556
+ f"Radiation data contains '{name}' variable "
2557
+ "but units are not specified. Provide units in the "
2558
+ f"rad['{name}'].attrs passed into Cocip."
2559
+ )
2560
+ raise KeyError(msg) from e
2561
+
2562
+ # The unit is already instantaneous
2563
+ if unit == "W m**-2":
2564
+ return arr
2565
+
2566
+ if unit != "J m**-2":
2567
+ msg = f"Unexpected units '{unit}' for '{name}' variable. Expected 'J m**-2' or 'W m**-2'."
2568
+ raise ValueError(msg)
2569
+
2570
+ # Convert from J m**-2 to W m**-2
2571
+ if rad.dataset_attr == "ERA5" and rad.product_attr == "ensemble":
2572
+ n_seconds = 3.0 * 3600.0 # 3 hour interval
2573
+ else:
2574
+ n_seconds = 3600.0 # 1 hour interval
2575
+
2576
+ return arr / n_seconds
2577
+
2578
+
2579
+ def _emissions_variables() -> tuple[str, ...]:
2580
+ """Return variables required for emissions calculation."""
2581
+ return (
2582
+ "engine_efficiency",
2583
+ "fuel_flow",
2584
+ "aircraft_mass",
2585
+ "nvpm_ei_n",
2586
+ "wingspan",
2587
+ )
2588
+
2589
+
2590
+ def _contrail_contrail_overlapping(
2591
+ contrail: GeoVectorDataset, params: dict[str, Any]
2592
+ ) -> GeoVectorDataset:
2593
+ """Mutate ``contrail`` to account for contrail-contrail overlapping effects."""
2594
+
2595
+ if not contrail:
2596
+ return contrail
2597
+
2598
+ contrail = radiative_forcing.contrail_contrail_overlap_radiative_effects(
2599
+ contrail,
2600
+ habit_distributions=params["habit_distributions"],
2601
+ radius_threshold_um=params["radius_threshold_um"],
2602
+ min_altitude_m=params["min_altitude_m"],
2603
+ max_altitude_m=params["max_altitude_m"],
2604
+ dz_overlap_m=params["dz_overlap_m"],
2605
+ )
2606
+
2607
+ contrail.update(
2608
+ rsr=contrail.data.pop("rsr_overlap"),
2609
+ olr=contrail.data.pop("olr_overlap"),
2610
+ tau_cirrus=contrail.data.pop("tau_cirrus_overlap"),
2611
+ rf_sw=contrail.data.pop("rf_sw_overlap"),
2612
+ rf_lw=contrail.data.pop("rf_lw_overlap"),
2613
+ rf_net=contrail.data.pop("rf_net_overlap"),
2614
+ )
2615
+
2616
+ return contrail