Coverage for C:\Users\Bob\Documents\python_projects\ecape\src\ecape\calc.py: 99%
88 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-06-04 23:10 -0600
« prev ^ index » next coverage.py v7.2.7, created at 2023-06-04 23:10 -0600
1# SPDX-FileCopyrightText: 2023-present Robert Capella <bob.capella@gmail.com>
2# SPDX-License-Identifier: MIT
4"""Calculate the entraining CAPE (ECAPE) of a parcel"""
5from typing import Callable, Tuple
7import metpy.calc as mpcalc
8import numpy as np
9import pint
10from metpy.constants import dry_air_spec_heat_press, earth_gravity
11from metpy.units import check_units, units
13PintList = np.ndarray[pint.Quantity]
16@check_units("[pressure]", "[temperature]", "[temperature]")
17def _get_parcel_profile(
18 pressure: PintList, temperature: PintList, dew_point_temperature: PintList, parcel_func: Callable = None
19) -> PintList:
20 """
21 Retrieve a parcel's temperature profile.
23 Args:
24 pressure:
25 Total atmospheric pressure
26 temperature:
27 Air temperature
28 dew_point_temperature:
29 Dew point temperature
30 parcel_func:
31 parcel profile retrieval callable via MetPy
33 Returns:
34 parcel_profile
36 """
38 # if surface-based, skip this process, None is default for lfc, el MetPy calcs
39 if parcel_func:
40 # calculate the parcel's starting temperature, then parcel temperature profile
41 parcel_p, parcel_t, parcel_td, *parcel_i = parcel_func(pressure, temperature, dew_point_temperature)
42 parcel_profile = mpcalc.parcel_profile(pressure, parcel_t, parcel_td)
43 else:
44 parcel_profile = None
46 return parcel_profile
49@check_units("[pressure]", "[length]", "[temperature]", "[temperature]")
50def calc_lfc_height(
51 pressure: PintList, height: PintList, temperature: PintList, dew_point_temperature: PintList, parcel_func: Callable
52) -> Tuple[int, pint.Quantity]:
53 """
54 Retrieve a parcel's level of free convection (lfc).
56 Args:
57 pressure:
58 Total atmospheric pressure
59 height:
60 Atmospheric heights at the levels given by 'pressure'.
61 temperature:
62 Air temperature
63 dew_point_temperature:
64 Dew point temperature
65 parcel_func:
66 parcel profile retrieval callable via MetPy
68 Returns:
69 lfc:
70 index of the last instance of negative buoyancy below the lfc
71 lfc_z:
72 height of the last instance of negative buoyancy below the lfc
74 """
76 # calculate the parcel's temperature profile
77 parcel_profile = _get_parcel_profile(pressure, temperature, dew_point_temperature, parcel_func)
79 # calculate the lfc, select the appropriate index & associated height
80 lfc_p, lfc_t = mpcalc.lfc(pressure, temperature, dew_point_temperature, parcel_temperature_profile=parcel_profile)
81 lfc_idx = (pressure - lfc_p > 0).nonzero()[0][-1]
82 lfc_z = height[lfc_idx]
84 return lfc_idx, lfc_z
87@check_units("[pressure]", "[length]", "[temperature]", "[temperature]")
88def calc_el_height(
89 pressure: PintList, height: PintList, temperature: PintList, dew_point_temperature: PintList, parcel_func: Callable
90) -> Tuple[int, pint.Quantity]:
91 """
92 Retrieve a parcel's equilibrium level (el).
94 Args:
95 pressure:
96 Total atmospheric pressure
97 height:
98 Atmospheric heights at the levels given by 'pressure'.
99 temperature:
100 Air temperature
101 dew_point_temperature:
102 Dew point temperature
103 parcel_func:
104 parcel profile retrieval callable via MetPy
106 Returns:
107 el_idx:
108 index of the last instance of positive buoyancy below the el
109 el_z:
110 height of the last instance of positive buoyancy below the el
112 """
114 # calculate the parcel's temperature profile
115 parcel_profile = _get_parcel_profile(pressure, temperature, dew_point_temperature, parcel_func)
117 # calculate the el, select the appropriate index & associated height
118 el_p, el_t = mpcalc.el(pressure, temperature, dew_point_temperature, parcel_temperature_profile=parcel_profile)
119 el_idx = (pressure - el_p > 0).nonzero()[0][-1]
120 el_z = height[el_idx]
122 return el_idx, el_z
125@check_units("[pressure]", "[speed]", "[speed]", "[length]")
126def calc_sr_wind(pressure: PintList, u_wind: PintList, v_wind: PintList, height: PintList) -> pint.Quantity:
127 """
128 Calculate the mean storm relative (as compared to Bunkers right motion) wind magnitude in the 0-1 km AGL layer
130 Args:
131 pressure:
132 Total atmospheric pressure
133 u_wind:
134 X component of the wind
135 v_wind
136 Y component of the wind
137 height:
138 Atmospheric heights at the levels given by 'pressure'.
140 Returns:
141 sr_wind:
142 0-1 km AGL average storm relative wind magnitude
144 """
146 bunkers_right, _, _ = mpcalc.bunkers_storm_motion(pressure, u_wind, v_wind, height) # right, left, mean
148 u_sr = u_wind - bunkers_right[0] # u-component
149 v_sr = v_wind - bunkers_right[1] # v-component
151 u_sr_1km = u_sr[np.nonzero(height <= 1000 * units("m"))]
152 v_sr_1km = v_sr[np.nonzero(height <= 1000 * units("m"))]
154 sr_wind = np.mean(mpcalc.wind_speed(u_sr_1km, v_sr_1km))
156 return sr_wind
159@check_units("[pressure]", "[length]", "[temperature]", "[mass]/[mass]")
160def calc_mse(
161 pressure: PintList, height: PintList, temperature: PintList, specific_humidity: PintList
162) -> Tuple[PintList, PintList]:
163 """
164 Calculate the moist static energy terms of interest.
166 Args:
167 pressure:
168 Total atmospheric pressure
169 height:
170 Atmospheric heights at the levels given by 'pressure'.
171 temperature:
172 Air temperature
173 specific_humidity:
174 Specific humidity
176 Returns:
177 moist_static_energy_bar:
178 Mean moist static energy from the surface to a layer
179 moist_static_energy_star:
180 Saturated moist static energy
181 """
183 # calculate MSE_bar
184 moist_static_energy = mpcalc.moist_static_energy(height, temperature, specific_humidity)
185 moist_static_energy_bar = np.cumsum(moist_static_energy) / np.arange(1, len(moist_static_energy) + 1)
186 moist_static_energy_bar = moist_static_energy_bar.to("J/kg")
188 # calculate MSE*
189 saturation_mixing_ratio = mpcalc.saturation_mixing_ratio(pressure, temperature)
190 moist_static_energy_star = mpcalc.moist_static_energy(height, temperature, saturation_mixing_ratio)
191 moist_static_energy_star = moist_static_energy_star.to("J/kg")
193 return moist_static_energy_bar, moist_static_energy_star
196@check_units("[energy]/[mass]", "[energy]/[mass]", "[temperature]")
197def calc_integral_arg(moist_static_energy_bar, moist_static_energy_star, temperature) -> PintList:
198 """
199 Calculate the contents of the integral defined in the NCAPE equation (54).
201 Args:
202 moist_static_energy_bar:
203 Mean moist static energy from the surface to a layer
204 moist_static_energy_star:
205 Saturated moist static energy
206 temperature:
207 Air temperature
209 Returns:
210 integral_arg:
211 Contents of integral defined in NCAPE eqn. 54
213 """
215 # NCAPE eqn 54 integrand, see compute_NCAPE.m L32
216 integral_arg = -(earth_gravity / (dry_air_spec_heat_press * temperature)) * (
217 moist_static_energy_bar - moist_static_energy_star
218 )
220 return integral_arg
223@check_units("[length]/[time]**2", "[length]", "[dimensionless]", "[dimensionless]")
224def calc_ncape(integral_arg: PintList, height: PintList, lfc_idx: int, el_idx: int) -> pint.Quantity:
225 """
226 Calculate the buoyancy dilution potential (NCAPE)
228 Args:
229 integral_arg:
230 Contents of integral defined in NCAPE eqn. 54
231 height:
232 Atmospheric heights at the levels given by 'pressure'.
233 lfc_idx:
234 Index of the last instance of negative buoyancy below the lfc
235 el_idx:
236 Index of the last instance of positive buoyancy below the el
238 Returns:
239 ncape:
240 Buoyancy dilution potential of the free troposphere (eqn. 54)
241 """
243 # see compute_NCAPE.m L41
244 ncape = np.sum(
245 (0.5 * integral_arg[lfc_idx:el_idx] + 0.5 * integral_arg[lfc_idx + 1 : el_idx + 1])
246 * (height[lfc_idx + 1 : el_idx + 1] - height[lfc_idx:el_idx])
247 )
249 return ncape
252@check_units("[speed]", "[dimensionless]", "[length]**2/[time]**2", "[energy]/[mass]")
253def calc_ecape_a(sr_wind: PintList, psi: pint.Quantity, ncape: pint.Quantity, cape: pint.Quantity) -> pint.Quantity:
254 """
255 Calculate the entraining cape of a parcel
257 Args:
258 sr_wind:
259 0-1 km AGL average storm relative wind magnitude
260 psi:
261 Parameter defined in eqn. 52, constant for a given equilibrium level
262 ncape:
263 Buoyancy dilution potential of the free troposphere (eqn. 54)
264 cape:
265 Convective available potential energy (CAPE, user-defined type)
266 Returns:
267 ecape:
268 Entraining CAPE (eqn. 55)
269 """
271 # broken into terms for readability
272 term_a = sr_wind**2 / 2.0
273 term_b = (-1 - psi - (2 * psi / sr_wind**2) * ncape) / (4 * psi / sr_wind**2)
274 term_c = (
275 np.sqrt((1 + psi + (2 * psi / sr_wind**2) * ncape) ** 2 + 8 * (psi / sr_wind**2) * (cape - (psi * ncape)))
276 ) / (4 * psi / sr_wind**2)
278 ecape_a = term_a + term_b + term_c
280 # set to 0 if negative
281 return ecape_a.to("J/kg") if ecape_a >= 0 else 0
284@check_units("[length]")
285def calc_psi(el_z: pint.Quantity) -> pint.Quantity:
286 """
287 Calculate the constant psi as denoted in eqn. 52
289 Args:
290 el_z:
291 height of the last instance of positive buoyancy below the el
293 Returns:
294 psi:
295 Parameter defined in eqn. 52, constant for a given equilibrium level, see COMPUTE_ECAPE.m L88 (pitchfork)
296 """
298 # additional constants as denoted in section 4 step 1.
299 sigma = 1.6 * units("dimensionless")
300 alpha = 0.8 * units("dimensionless")
301 l_mix = 120.0 * units("m")
302 pr = (1.0 / 3.0) * units("dimensionless") # prandtl number
303 ksq = 0.18 * units("dimensionless") # von karman constant
305 psi = (ksq * alpha**2 * np.pi**2 * l_mix) / (4.0 * pr * sigma**2 * el_z)
307 return psi
310@check_units("[length]", "[pressure]", "[temperature]", "[mass]/[mass]", "[speed]", "[speed]")
311def calc_ecape(
312 height: PintList,
313 pressure: PintList,
314 temperature: PintList,
315 specific_humidity: PintList,
316 u_wind: PintList,
317 v_wind: PintList,
318 cape_type: str = "most_unstable",
319) -> pint.Quantity:
320 """
321 Calculate the entraining CAPE (ECAPE) of a parcel
323 Parameters:
324 ------------
325 height: np.ndarray[pint.Quantity]
326 Atmospheric heights at the levels given by 'pressure'.
327 pressure: np.ndarray[pint.Quantity]
328 Total atmospheric pressure
329 temperature: np.ndarray[pint.Quantity]
330 Air temperature
331 specific humidity: np.ndarray[pint.Quantity]
332 Specific humidity
333 u_wind: np.ndarray[pint.Quantity]
334 X component of the wind
335 v_wind np.ndarray[pint.Quantity]
336 Y component of the wind
337 cape_type: np.ndarray[pint.Quantity]
338 Variation of CAPE desired. 'most_unstable' (default), 'surface_based', or 'mixed_layer'
340 Returns:
341 ----------
342 ecape : 'pint.Quantity'
343 Entraining CAPE
344 """
346 cape_func = {
347 "most_unstable": mpcalc.most_unstable_cape_cin,
348 "surface_based": mpcalc.surface_based_cape_cin,
349 "mixed_layer": mpcalc.mixed_layer_cape_cin,
350 }
352 parcel_func = {
353 "most_unstable": mpcalc.most_unstable_parcel,
354 "surface_based": None,
355 "mixed_layer": mpcalc.mixed_parcel,
356 }
358 # calculate cape
359 dew_point_temperature = mpcalc.dewpoint_from_specific_humidity(pressure, temperature, specific_humidity)
360 cape, _ = cape_func[cape_type](pressure, temperature, dew_point_temperature)
362 # calculate the level of free convection (lfc) and equilibrium level (el) indexes
363 lfc_idx, _ = calc_lfc_height(pressure, height, temperature, dew_point_temperature, parcel_func[cape_type])
364 el_idx, el_z = calc_el_height(pressure, height, temperature, dew_point_temperature, parcel_func[cape_type])
366 # calculate the buoyancy dilution potential (ncape)
367 moist_static_energy_bar, moist_static_energy_star = calc_mse(pressure, height, temperature, specific_humidity)
368 integral_arg = calc_integral_arg(moist_static_energy_bar, moist_static_energy_star, temperature)
369 ncape = calc_ncape(integral_arg, height, lfc_idx, el_idx)
371 # calculate the storm relative (sr) wind
372 sr_wind = calc_sr_wind(pressure, u_wind, v_wind, height)
374 # calculate the entraining cape (ecape)
375 psi = calc_psi(el_z)
376 ecape_a = calc_ecape_a(sr_wind, psi, ncape, cape)
378 return ecape_a
380if __name__ == "__main__":
381 pass