pycontrails 0.42.0__cp311-cp311-macosx_11_0_arm64.whl → 0.42.2__cp311-cp311-macosx_11_0_arm64.whl

pycontrails/models/cocip/wake_vortex.py

@@ -3,47 +3,48 @@
 from __future__ import annotations
 
 import numpy as np
+import numpy.typing as npt
 
 from pycontrails.models.cocip import wind_shear
 from pycontrails.physics import constants, thermo
 
 
 def max_downward_displacement(
-    wingspan: np.ndarray | float,
-    true_airspeed: np.ndarray,
-    aircraft_mass: np.ndarray | float,
-    air_temperature: np.ndarray,
-    dT_dz: np.ndarray,
-    ds_dz: np.ndarray,
-    air_pressure: np.ndarray,
+    wingspan: npt.NDArray[np.float_] | float,
+    true_airspeed: npt.NDArray[np.float_],
+    aircraft_mass: npt.NDArray[np.float_] | float,
+    air_temperature: npt.NDArray[np.float_],
+    dT_dz: npt.NDArray[np.float_],
+    ds_dz: npt.NDArray[np.float_],
+    air_pressure: npt.NDArray[np.float_],
     effective_vertical_resolution: float,
-    wind_shear_enhancement_exponent: float | np.ndarray,
-) -> np.ndarray:
+    wind_shear_enhancement_exponent: float | npt.NDArray[np.float_],
+) -> npt.NDArray[np.float_]:
     """
     Calculate the maximum contrail downward displacement after the wake vortex phase.
 
     Parameters
     ----------
-    wingspan : np.ndarray | float
+    wingspan : npt.NDArray[np.float_] | float
         aircraft wingspan, [:math:`m`]
-    true_airspeed : np.ndarray
+    true_airspeed : npt.NDArray[np.float_]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    aircraft_mass : np.ndarray | float
+    aircraft_mass : npt.NDArray[np.float_] | float
         aircraft mass for each waypoint, [:math:`kg`]
-    air_temperature : np.ndarray
+    air_temperature : npt.NDArray[np.float_]
         ambient temperature for each waypoint, [:math:`K`]
-    dT_dz : np.ndarray
+    dT_dz : npt.NDArray[np.float_]
         potential temperature gradient, [:math:`K m^{-1}`]
-    ds_dz : np.ndarray
+    ds_dz : npt.NDArray[np.float_]
         Difference in wind speed over dz in the atmosphere, [:math:`m s^{-1} / m`]
-    air_pressure : np.ndarray
+    air_pressure : npt.NDArray[np.float_]
         pressure altitude at each waypoint, [:math:`Pa`]
     effective_vertical_resolution, wind_shear_enhancement_exponent: float
         Passed through to :func:`wind_shear.wind_shear_enhancement_factor`
 
     Returns
     -------
-    np.ndarray
+    npt.NDArray[np.float_]
         Max contrail downward displacement after the wake vortex phase, [:math:`m`]
 
     References
@@ -81,28 +82,28 @@ def max_downward_displacement(
 
 
 def get_effective_time_scale(
-    wingspan: np.ndarray | float,
-    true_airspeed: np.ndarray,
-    aircraft_mass: np.ndarray | float,
-    rho_air: np.ndarray,
-) -> np.ndarray:
+    wingspan: npt.NDArray[np.float_] | float,
+    true_airspeed: npt.NDArray[np.float_],
+    aircraft_mass: npt.NDArray[np.float_] | float,
+    rho_air: npt.NDArray[np.float_],
+) -> npt.NDArray[np.float_]:
     r"""
     Calculate the effective time scale of the wake vortex.
 
     Parameters
     ----------
-    wingspan : np.ndarray
+    wingspan : npt.NDArray[np.float_]
         aircraft wingspan, [:math:`m`]
-    true_airspeed : np.ndarray
+    true_airspeed : npt.NDArray[np.float_]
         true airspeed for each waypoint, [:math:`m \ s^{-1}`]
-    aircraft_mass : np.ndarray
+    aircraft_mass : npt.NDArray[np.float_]
         aircraft mass for each waypoint, [:math:`kg`]
-    rho_air : np.ndarray
+    rho_air : npt.NDArray[np.float_]
         density of air for each waypoint, [:math:`kg \ m^{-3}`]
 
     Returns
     -------
-    np.ndarray
+    npt.NDArray[np.float_]
         Wake vortex effective time scale
 
     Notes
@@ -118,31 +119,31 @@ def get_effective_time_scale(
 
 
 def get_downward_displacement_strongly_stratified(
-    wingspan: np.ndarray | float,
-    true_airspeed: np.ndarray,
-    aircraft_mass: np.ndarray | float,
-    rho_air: np.ndarray,
-    n_bv: np.ndarray,
-) -> np.ndarray:
+    wingspan: npt.NDArray[np.float_] | float,
+    true_airspeed: npt.NDArray[np.float_],
+    aircraft_mass: npt.NDArray[np.float_] | float,
+    rho_air: npt.NDArray[np.float_],
+    n_bv: npt.NDArray[np.float_],
+) -> npt.NDArray[np.float_]:
     """
     Calculate the maximum contrail downward displacement under strongly stratified conditions.
 
     Parameters
     ----------
-    wingspan : np.ndarray | float
+    wingspan : npt.NDArray[np.float_] | float
         aircraft wingspan, [:math:`m`]
-    true_airspeed : np.ndarray
+    true_airspeed : npt.NDArray[np.float_]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    aircraft_mass : np.ndarray | float
+    aircraft_mass : npt.NDArray[np.float_] | float
         aircraft mass for each waypoint, [:math:`kg`]
-    rho_air : np.ndarray
+    rho_air : npt.NDArray[np.float_]
         density of air for each waypoint, [:math:`kg m^{-3}`]
-    n_bv : np.ndarray
+    n_bv : npt.NDArray[np.float_]
         Brunt-Vaisaila frequency, [:math:`s^{-1}`]
 
     Returns
     -------
-    np.ndarray
+    npt.NDArray[np.float_]
         Maximum contrail downward displacement, strongly stratified conditions.
 
     Notes
@@ -158,46 +159,46 @@ def get_downward_displacement_strongly_stratified(
 
 
 def get_downward_displacement_weakly_stratified(
-    wingspan: np.ndarray | float,
-    true_airspeed: np.ndarray,
-    aircraft_mass: np.ndarray | float,
-    rho_air: np.ndarray,
-    n_bv: np.ndarray,
-    dz_max_strong: np.ndarray,
-    ds_dz: np.ndarray,
-    t_0: np.ndarray,
+    wingspan: npt.NDArray[np.float_] | float,
+    true_airspeed: npt.NDArray[np.float_],
+    aircraft_mass: npt.NDArray[np.float_] | float,
+    rho_air: npt.NDArray[np.float_],
+    n_bv: npt.NDArray[np.float_],
+    dz_max_strong: npt.NDArray[np.float_],
+    ds_dz: npt.NDArray[np.float_],
+    t_0: npt.NDArray[np.float_],
     effective_vertical_resolution: float,
-    wind_shear_enhancement_exponent: np.ndarray | float,
-) -> np.ndarray:
+    wind_shear_enhancement_exponent: npt.NDArray[np.float_] | float,
+) -> npt.NDArray[np.float_]:
     """
     Calculate the maximum contrail downward displacement under weakly/stably stratified conditions.
 
     Parameters
     ----------
-    wingspan : np.ndarray | float
+    wingspan : npt.NDArray[np.float_] | float
         aircraft wingspan, [:math:`m`]
-    true_airspeed : np.ndarray
+    true_airspeed : npt.NDArray[np.float_]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    aircraft_mass : np.ndarray | float
+    aircraft_mass : npt.NDArray[np.float_] | float
         aircraft mass for each waypoint, [:math:`kg`]
-    rho_air : np.ndarray
+    rho_air : npt.NDArray[np.float_]
         density of air for each waypoint, [:math:`kg m^{-3}`]
-    n_bv : np.ndarray
+    n_bv : npt.NDArray[np.float_]
         Brunt-Vaisaila frequency, [:math:`s^{-1}`]
-    dz_max_strong : np.ndarray
+    dz_max_strong : npt.NDArray[np.float_]
         Max contrail downward displacement under strongly stratified conditions, [:math:`m`]
-    ds_dz : np.ndarray
+    ds_dz : npt.NDArray[np.float_]
         Difference in wind speed over dz in the atmosphere, [:math:`m s^{-1} / m`]
-    t_0 : np.ndarray
+    t_0 : npt.NDArray[np.float_]
         Wake vortex effective time scale
     effective_vertical_resolution: float
         Passed through to :func:`wind_shear.wind_shear_enhancement_factor`
-    wind_shear_enhancement_exponent: np.ndarray | float
+    wind_shear_enhancement_exponent: npt.NDArray[np.float_] | float
         Passed through to :func:`wind_shear.wind_shear_enhancement_factor`
 
     Returns
     -------
-    np.ndarray
+    npt.NDArray[np.float_]
         Maximum contrail downward displacement, weakly/stably stratified conditions.
 
     Notes
@@ -209,7 +210,7 @@ def get_downward_displacement_weakly_stratified(
     - :cite:`schumannContrailCirrusPrediction2012`
     """
     b_0 = get_wake_vortex_separation(wingspan)
-    dz_max = np.maximum(dz_max_strong, 10)
+    dz_max = np.maximum(dz_max_strong, 10.0)
     shear_enhancement_factor = wind_shear.wind_shear_enhancement_factor(
         dz_max, effective_vertical_resolution, wind_shear_enhancement_exponent
     )
@@ -221,27 +222,29 @@ def get_downward_displacement_weakly_stratified(
     return b_0 * (7.68 * (1 - 4.07 * epsn_st + 5.67 * epsn_st**2) * (0.79 - n_bv * t_0) + 1.88)
 
 
-def get_wake_vortex_separation(wingspan: np.ndarray | float) -> np.ndarray | float:
+def get_wake_vortex_separation(
+    wingspan: npt.NDArray[np.float_] | float,
+) -> npt.NDArray[np.float_] | float:
     """
     Calculate the wake vortex separation.
 
     Parameters
     ----------
-    wingspan : np.ndarray | float
+    wingspan : npt.NDArray[np.float_] | float
         aircraft wingspan, [:math:`m`]
 
     Returns
     -------
-    np.ndarray
+    npt.NDArray[np.float_]
         wake vortex separation, [:math:`m`]
     """
-    return (np.pi * wingspan) / 4
+    return (np.pi * wingspan) / 4.0
 
 
 def turbulent_kinetic_energy_dissipation_rate(
-    ds_dz: np.ndarray,
-    shear_enhancement_factor: np.ndarray | float = 1.0,
-) -> np.ndarray:
+    ds_dz: npt.NDArray[np.float_],
+    shear_enhancement_factor: npt.NDArray[np.float_] | float = 1.0,
+) -> npt.NDArray[np.float_]:
     """
     Calculate the turbulent kinetic energy dissipation rate (epsilon).
 
@@ -249,14 +252,14 @@ def turbulent_kinetic_energy_dissipation_rate(
 
     Parameters
     ----------
-    ds_dz : np.ndarray
+    ds_dz : npt.NDArray[np.float_]
         Difference in wind speed over dz in the atmosphere, [:math:`m s^{-1} / m`]
-    shear_enhancement_factor : np.ndarray | float
+    shear_enhancement_factor : npt.NDArray[np.float_] | float
         Multiplication factor to enhance the wind shear
 
     Returns
     -------
-    np.ndarray
+    npt.NDArray[np.float_]
         turbulent kinetic energy dissipation rate
 
     Notes
@@ -271,31 +274,31 @@ def turbulent_kinetic_energy_dissipation_rate(
 
 
 def normalized_dissipation_rate(
-    epsilon: np.ndarray,
-    wingspan: np.ndarray | float,
-    true_airspeed: np.ndarray,
-    aircraft_mass: np.ndarray | float,
-    rho_air: np.ndarray | float,
-) -> np.ndarray:
+    epsilon: npt.NDArray[np.float_],
+    wingspan: npt.NDArray[np.float_] | float,
+    true_airspeed: npt.NDArray[np.float_],
+    aircraft_mass: npt.NDArray[np.float_] | float,
+    rho_air: npt.NDArray[np.float_] | float,
+) -> npt.NDArray[np.float_]:
     """
     Calculate the normalized dissipation rate of the sinking wake vortex.
 
     Parameters
     ----------
-    epsilon: np.ndarray
+    epsilon: npt.NDArray[np.float_]
         turbulent kinetic energy dissipation rate
-    wingspan : np.ndarray | float
+    wingspan : npt.NDArray[np.float_] | float
         aircraft wingspan, [:math:`m`]
-    true_airspeed : np.ndarray
+    true_airspeed : npt.NDArray[np.float_]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    aircraft_mass : np.ndarray
+    aircraft_mass : npt.NDArray[np.float_]
         aircraft mass for each waypoint, [:math:`kg`]
-    rho_air : np.ndarray
+    rho_air : npt.NDArray[np.float_]
         density of air for each waypoint, [:math:`kg m^{-3}`]
 
     Returns
     -------
-    np.ndarray
+    npt.NDArray[np.float_]
         Normalized dissipation rate of the sinking wake vortex
 
     Notes
@@ -320,43 +323,45 @@ def normalized_dissipation_rate(
     return np.minimum(epsn_st, 0.36)
 
 
-def initial_contrail_width(wingspan: np.ndarray | float, dz_max: np.ndarray) -> np.ndarray:
+def initial_contrail_width(
+    wingspan: npt.NDArray[np.float_] | float, dz_max: npt.NDArray[np.float_]
+) -> npt.NDArray[np.float_]:
     """
     Calculate the initial contrail width.
 
     Parameters
     ----------
-    wingspan : np.ndarray | float
+    wingspan : npt.NDArray[np.float_] | float
         aircraft wingspan, [:math:`m`]
-    dz_max : np.ndarray
+    dz_max : npt.NDArray[np.float_]
         Max contrail downward displacement after the wake vortex phase, [:math:`m`]
         Only the size of this array is used; the values are ignored.
 
     Returns
     -------
-    np.ndarray
+    npt.NDArray[np.float_]
         Initial contrail width, [:math:`m`]
     """
     return np.full_like(dz_max, np.pi / 4) * wingspan
 
 
 def initial_contrail_depth(
-    dz_max: np.ndarray, initial_wake_vortex_depth: float | np.ndarray
-) -> np.ndarray:
+    dz_max: npt.NDArray[np.float_], initial_wake_vortex_depth: float | npt.NDArray[np.float_]
+) -> npt.NDArray[np.float_]:
     """
     Calculate the initial contrail depth.
 
     Parameters
     ----------
-    dz_max : np.ndarray
+    dz_max : npt.NDArray[np.float_]
         Max contrail downward displacement after the wake vortex phase, [:math:`m`]
-    initial_wake_vortex_depth : float | np.ndarray
+    initial_wake_vortex_depth : float | npt.NDArray[np.float_]
         Initial wake vortex depth scaling factor.
         Denoted `C_D0` in eq (14) in :cite:`schumannContrailCirrusPrediction2012`.
 
     Returns
     -------
-    np.ndarray
+    npt.NDArray[np.float_]
         Initial contrail depth [m]
     """
     return dz_max * initial_wake_vortex_depth