From 9f662ac6b34550f7d2817797a5cfcd82b77f31ab Mon Sep 17 00:00:00 2001
From: Zeb Engberg <zebengberg@gmail.com>
Date: Mon, 9 Dec 2024 10:04:19 -0700
Subject: [PATCH] update / ignore type hints for numpy 2.2 compatibility

---
 CHANGELOG.md                                  |   4 +
 pycontrails/core/aircraft_performance.py      | 116 ++--
 pycontrails/core/fleet.py                     |  16 +-
 pycontrails/core/flight.py                    | 140 ++--
 pycontrails/core/interpolation.py             | 110 +--
 pycontrails/core/met.py                       |  87 +--
 pycontrails/core/models.py                    |  20 +-
 pycontrails/core/polygon.py                   |  30 +-
 pycontrails/core/vector.py                    |  40 +-
 pycontrails/datalib/_met_utils/metsource.py   |  11 +-
 pycontrails/datalib/ecmwf/common.py           |   2 +-
 pycontrails/datalib/ecmwf/hres.py             |   4 +-
 pycontrails/datalib/ecmwf/ifs.py              |   2 +-
 pycontrails/datalib/gfs/gfs.py                |   2 +-
 pycontrails/datalib/goes.py                   |   8 +-
 pycontrails/ext/empirical_grid.py             |   2 +-
 pycontrails/models/cocip/cocip.py             |  20 +-
 pycontrails/models/cocip/cocip_uncertainty.py |   4 +-
 .../models/cocip/contrail_properties.py       | 647 +++++++++---------
 pycontrails/models/cocip/output_formats.py    | 106 +--
 pycontrails/models/cocip/radiative_forcing.py | 266 +++----
 pycontrails/models/cocip/radiative_heating.py | 270 ++++----
 .../models/cocip/unterstrasser_wake_vortex.py | 177 ++---
 pycontrails/models/cocip/wake_vortex.py       | 184 ++---
 pycontrails/models/cocip/wind_shear.py        |  16 +-
 pycontrails/models/cocipgrid/cocip_grid.py    |  24 +-
 pycontrails/models/emissions/black_carbon.py  | 216 +++---
 pycontrails/models/emissions/emissions.py     | 170 ++---
 pycontrails/models/emissions/ffm2.py          |  70 +-
 .../humidity_scaling/humidity_scaling.py      |  46 +-
 pycontrails/models/ps_model/ps_grid.py        |  25 +-
 pycontrails/models/ps_model/ps_model.py       |  26 +-
 .../models/ps_model/ps_operational_limits.py  |  38 +-
 pycontrails/physics/geo.py                    | 134 ++--
 pycontrails/physics/jet.py                    | 158 ++---
 pycontrails/physics/units.py                  |  28 +-
 pycontrails/utils/types.py                    |  11 +-
 37 files changed, 1633 insertions(+), 1597 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a6e033b03..3caec2604 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,10 @@
 
 - Fix the `ERA5` interface when making a pressure-level request with a single pressure level. This change accommodates CDS-Beta server behavior. Previously, a ValueError was raised in this case.
 
+### Internals
+
+- Update type hints for `numpy` 2.2 compatibility. Additional changes may be required after the next iteration of the `numpy` 2.2 series.
+
 ## v0.54.3
 
 ### Breaking changes
diff --git a/pycontrails/core/aircraft_performance.py b/pycontrails/core/aircraft_performance.py
index 074d6f67d..5d7d2c505 100644
--- a/pycontrails/core/aircraft_performance.py
+++ b/pycontrails/core/aircraft_performance.py
@@ -163,14 +163,14 @@ def simulate_fuel_and_performance(
         self,
         *,
         aircraft_type: str,
-        altitude_ft: npt.NDArray[np.float64],
+        altitude_ft: npt.NDArray[np.floating],
         time: npt.NDArray[np.datetime64],
-        true_airspeed: npt.NDArray[np.float64],
-        air_temperature: npt.NDArray[np.float64],
-        aircraft_mass: npt.NDArray[np.float64] | float | None,
-        thrust: npt.NDArray[np.float64] | float | None,
-        engine_efficiency: npt.NDArray[np.float64] | float | None,
-        fuel_flow: npt.NDArray[np.float64] | float | None,
+        true_airspeed: npt.NDArray[np.floating],
+        air_temperature: npt.NDArray[np.floating],
+        aircraft_mass: npt.NDArray[np.floating] | float | None,
+        thrust: npt.NDArray[np.floating] | float | None,
+        engine_efficiency: npt.NDArray[np.floating] | float | None,
+        fuel_flow: npt.NDArray[np.floating] | float | None,
         q_fuel: float,
         n_iter: int,
         amass_oew: float,
@@ -192,21 +192,21 @@ def simulate_fuel_and_performance(
         ----------
         aircraft_type: str
             Aircraft type designator used to query the underlying model database.
-        altitude_ft: npt.NDArray[np.float64]
+        altitude_ft: npt.NDArray[np.floating]
             Altitude at each waypoint, [:math:`ft`]
         time: npt.NDArray[np.datetime64]
             Waypoint time in ``np.datetime64`` format.
-        true_airspeed: npt.NDArray[np.float64]
+        true_airspeed: npt.NDArray[np.floating]
             True airspeed for each waypoint, [:math:`m s^{-1}`]
-        air_temperature : npt.NDArray[np.float64]
+        air_temperature : npt.NDArray[np.floating]
             Ambient temperature for each waypoint, [:math:`K`]
-        aircraft_mass : npt.NDArray[np.float64] | float | None
+        aircraft_mass : npt.NDArray[np.floating] | float | None
             Override the aircraft_mass at each waypoint, [:math:`kg`].
-        thrust : npt.NDArray[np.float64] | float | None
+        thrust : npt.NDArray[np.floating] | float | None
             Override the thrust setting at each waypoint, [:math: `N`].
-        engine_efficiency : npt.NDArray[np.float64] | float | None
+        engine_efficiency : npt.NDArray[np.floating] | float | None
             Override the engine efficiency at each waypoint.
-        fuel_flow : npt.NDArray[np.float64] | float | None
+        fuel_flow : npt.NDArray[np.floating] | float | None
             Override the fuel flow at each waypoint, [:math:`kg s^{-1}`].
         q_fuel : float
             Lower calorific value (LCV) of fuel, [:math:`J \ kg_{fuel}^{-1}`].
@@ -280,14 +280,14 @@ def _simulate_fuel_and_performance_known_aircraft_mass(
         self,
         *,
         aircraft_type: str,
-        altitude_ft: npt.NDArray[np.float64],
+        altitude_ft: npt.NDArray[np.floating],
         time: npt.NDArray[np.datetime64],
-        true_airspeed: npt.NDArray[np.float64],
-        air_temperature: npt.NDArray[np.float64],
-        aircraft_mass: npt.NDArray[np.float64] | float,
-        thrust: npt.NDArray[np.float64] | float | None,
-        engine_efficiency: npt.NDArray[np.float64] | float | None,
-        fuel_flow: npt.NDArray[np.float64] | float | None,
+        true_airspeed: npt.NDArray[np.floating],
+        air_temperature: npt.NDArray[np.floating],
+        aircraft_mass: npt.NDArray[np.floating] | float,
+        thrust: npt.NDArray[np.floating] | float | None,
+        engine_efficiency: npt.NDArray[np.floating] | float | None,
+        fuel_flow: npt.NDArray[np.floating] | float | None,
         q_fuel: float,
         **kwargs: Any,
     ) -> AircraftPerformanceData:
@@ -332,13 +332,13 @@ def _simulate_fuel_and_performance_unknown_aircraft_mass(
         self,
         *,
         aircraft_type: str,
-        altitude_ft: npt.NDArray[np.float64],
+        altitude_ft: npt.NDArray[np.floating],
         time: npt.NDArray[np.datetime64],
-        true_airspeed: npt.NDArray[np.float64],
-        air_temperature: npt.NDArray[np.float64],
-        thrust: npt.NDArray[np.float64] | float | None,
-        engine_efficiency: npt.NDArray[np.float64] | float | None,
-        fuel_flow: npt.NDArray[np.float64] | float | None,
+        true_airspeed: npt.NDArray[np.floating],
+        air_temperature: npt.NDArray[np.floating],
+        thrust: npt.NDArray[np.floating] | float | None,
+        engine_efficiency: npt.NDArray[np.floating] | float | None,
+        fuel_flow: npt.NDArray[np.floating] | float | None,
         q_fuel: float,
         n_iter: int,
         amass_oew: float,
@@ -351,7 +351,7 @@ def _simulate_fuel_and_performance_unknown_aircraft_mass(
         # Variable aircraft_mass will change dynamically after each iteration
         # Set the initial aircraft mass depending on a possible load factor
 
-        aircraft_mass: npt.NDArray[np.float64] | float
+        aircraft_mass: npt.NDArray[np.floating] | float
         if takeoff_mass is not None:
             aircraft_mass = takeoff_mass
         else:
@@ -411,14 +411,14 @@ def calculate_aircraft_performance(
         self,
         *,
         aircraft_type: str,
-        altitude_ft: npt.NDArray[np.float64],
-        air_temperature: npt.NDArray[np.float64],
+        altitude_ft: npt.NDArray[np.floating],
+        air_temperature: npt.NDArray[np.floating],
         time: npt.NDArray[np.datetime64] | None,
-        true_airspeed: npt.NDArray[np.float64] | float | None,
-        aircraft_mass: npt.NDArray[np.float64] | float,
-        engine_efficiency: npt.NDArray[np.float64] | float | None,
-        fuel_flow: npt.NDArray[np.float64] | float | None,
-        thrust: npt.NDArray[np.float64] | float | None,
+        true_airspeed: npt.NDArray[np.floating] | float | None,
+        aircraft_mass: npt.NDArray[np.floating] | float,
+        engine_efficiency: npt.NDArray[np.floating] | float | None,
+        fuel_flow: npt.NDArray[np.floating] | float | None,
+        thrust: npt.NDArray[np.floating] | float | None,
         q_fuel: float,
         **kwargs: Any,
     ) -> AircraftPerformanceData:
@@ -441,24 +441,24 @@ def calculate_aircraft_performance(
         ----------
         aircraft_type : str
             Used to query the underlying model database for aircraft engine parameters.
-        altitude_ft : npt.NDArray[np.float64]
+        altitude_ft : npt.NDArray[np.floating]
             Altitude at each waypoint, [:math:`ft`]
-        air_temperature : npt.NDArray[np.float64]
+        air_temperature : npt.NDArray[np.floating]
             Ambient temperature for each waypoint, [:math:`K`]
         time: npt.NDArray[np.datetime64] | None
             Waypoint time in ``np.datetime64`` format. If None, only drag force
             will is used in thrust calculations (ie, no vertical change and constant
             horizontal change). In addition, aircraft is assumed to be in cruise.
-        true_airspeed : npt.NDArray[np.float64] | float | None
+        true_airspeed : npt.NDArray[np.floating] | float | None
             True airspeed for each waypoint, [:math:`m s^{-1}`].
             If None, a nominal value is used.
-        aircraft_mass : npt.NDArray[np.float64] | float
+        aircraft_mass : npt.NDArray[np.floating] | float
             Aircraft mass for each waypoint, [:math:`kg`].
-        engine_efficiency : npt.NDArray[np.float64] | float | None
+        engine_efficiency : npt.NDArray[np.floating] | float | None
             Override the engine efficiency at each waypoint.
-        fuel_flow : npt.NDArray[np.float64] | float | None
+        fuel_flow : npt.NDArray[np.floating] | float | None
             Override the fuel flow at each waypoint, [:math:`kg s^{-1}`].
-        thrust : npt.NDArray[np.float64] | float | None
+        thrust : npt.NDArray[np.floating] | float | None
             Override the thrust setting at each waypoint, [:math: `N`].
         q_fuel : float
             Lower calorific value (LCV) of fuel, [:math:`J \ kg_{fuel}^{-1}`].
@@ -471,12 +471,12 @@ def calculate_aircraft_performance(
             Derived performance metrics at each waypoint.
         """
 
-    def ensure_true_airspeed_on_source(self) -> npt.NDArray[np.float64]:
+    def ensure_true_airspeed_on_source(self) -> npt.NDArray[np.floating]:
         """Add ``true_airspeed`` field to :attr:`source` data if not already present.
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             True airspeed, [:math:`m s^{-1}`]. If ``true_airspeed`` is already present
             on :attr:`source`, this is returned directly. Otherwise, it is calculated
             using :meth:`Flight.segment_true_airspeed`.
@@ -529,30 +529,30 @@ class AircraftPerformanceData:
 
     Parameters
     ----------
-    fuel_flow : npt.NDArray[np.float64]
+    fuel_flow : npt.NDArray[np.floating]
         Fuel mass flow rate for each waypoint, [:math:`kg s^{-1}`]
-    aircraft_mass : npt.NDArray[np.float64]
+    aircraft_mass : npt.NDArray[np.floating]
         Aircraft mass for each waypoint, [:math:`kg`]
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         True airspeed at each waypoint, [:math: `m s^{-1}`]
-    fuel_burn: npt.NDArray[np.float64]
+    fuel_burn: npt.NDArray[np.floating]
         Fuel consumption for each waypoint, [:math:`kg`]. Set to an array of
         all nan values if it cannot be computed (ie, working with gridpoints).
-    thrust: npt.NDArray[np.float64]
+    thrust: npt.NDArray[np.floating]
         Thrust force, [:math:`N`]
-    engine_efficiency: npt.NDArray[np.float64]
+    engine_efficiency: npt.NDArray[np.floating]
         Overall propulsion efficiency for each waypoint
-    rocd : npt.NDArray[np.float64]
+    rocd : npt.NDArray[np.floating]
         Rate of climb and descent, [:math:`ft min^{-1}`]
     """
 
-    fuel_flow: npt.NDArray[np.float64]
-    aircraft_mass: npt.NDArray[np.float64]
-    true_airspeed: npt.NDArray[np.float64]
-    fuel_burn: npt.NDArray[np.float64]
-    thrust: npt.NDArray[np.float64]
-    engine_efficiency: npt.NDArray[np.float64]
-    rocd: npt.NDArray[np.float64]
+    fuel_flow: npt.NDArray[np.floating]
+    aircraft_mass: npt.NDArray[np.floating]
+    true_airspeed: npt.NDArray[np.floating]
+    fuel_burn: npt.NDArray[np.floating]
+    thrust: npt.NDArray[np.floating]
+    engine_efficiency: npt.NDArray[np.floating]
+    rocd: npt.NDArray[np.floating]
 
 
 # --------------------------------
diff --git a/pycontrails/core/fleet.py b/pycontrails/core/fleet.py
index c345fc4dc..782d83ecf 100644
--- a/pycontrails/core/fleet.py
+++ b/pycontrails/core/fleet.py
@@ -265,12 +265,12 @@ def to_flight_list(self, copy: bool = True) -> list[Flight]:
 
     def segment_true_airspeed(
         self,
-        u_wind: npt.NDArray[np.float64] | float = 0.0,
-        v_wind: npt.NDArray[np.float64] | float = 0.0,
+        u_wind: npt.NDArray[np.floating] | float = 0.0,
+        v_wind: npt.NDArray[np.floating] | float = 0.0,
         smooth: bool = True,
         window_length: int = 7,
         polyorder: int = 1,
-    ) -> npt.NDArray[np.float64]:
+    ) -> npt.NDArray[np.floating]:
         """Calculate the true airspeed [:math:`m / s`] from the ground speed and horizontal winds.
 
         Because Flight.segment_true_airspeed uses a smoothing pattern, waypoints in :attr:`data`
@@ -302,7 +302,7 @@ def segment_true_airspeed(
             self[key] = v_wind
 
         # Calculate TAS on each flight individually
-        def calc_tas(fl: Flight) -> npt.NDArray[np.float64]:
+        def calc_tas(fl: Flight) -> npt.NDArray[np.floating]:
             u = fl.get("__u_wind", u_wind)
             v = fl.get("__v_wind", v_wind)
 
@@ -325,7 +325,7 @@ def calc_tas(fl: Flight) -> npt.NDArray[np.float64]:
         return np.concatenate(tas)
 
     @override
-    def segment_groundspeed(self, *args: Any, **kwargs: Any) -> npt.NDArray[np.float64]:
+    def segment_groundspeed(self, *args: Any, **kwargs: Any) -> npt.NDArray[np.floating]:
         fls = self.to_flight_list(copy=False)
         gs = [fl.segment_groundspeed(*args, **kwargs) for fl in fls]
         return np.concatenate(gs)
@@ -337,7 +337,7 @@ def resample_and_fill(self, *args: Any, **kwargs: Any) -> Fleet:
         return type(self).from_seq(flights, copy=False, broadcast_numeric=False, attrs=self.attrs)
 
     @override
-    def segment_length(self) -> npt.NDArray[np.float64]:
+    def segment_length(self) -> npt.NDArray[np.floating]:
         return np.where(self.final_waypoints, np.nan, super().segment_length())
 
     @property
@@ -346,11 +346,11 @@ def max_distance_gap(self) -> float:
         return np.nanmax(self.segment_length()).item()
 
     @override
-    def segment_azimuth(self) -> npt.NDArray[np.float64]:
+    def segment_azimuth(self) -> npt.NDArray[np.floating]:
         return np.where(self.final_waypoints, np.nan, super().segment_azimuth())
 
     @override
-    def segment_angle(self) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    def segment_angle(self) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
         sin_a, cos_a = super().segment_angle()
         sin_a[self.final_waypoints] = np.nan
         cos_a[self.final_waypoints] = np.nan
diff --git a/pycontrails/core/flight.py b/pycontrails/core/flight.py
index d2fec5ede..3a58003bd 100644
--- a/pycontrails/core/flight.py
+++ b/pycontrails/core/flight.py
@@ -414,7 +414,7 @@ def length(self) -> float:
     # Segment Properties
     # ------------
 
-    def segment_duration(self, dtype: npt.DTypeLike = np.float32) -> npt.NDArray[np.float64]:
+    def segment_duration(self, dtype: npt.DTypeLike = np.float32) -> npt.NDArray[np.floating]:
         r"""Compute time elapsed between waypoints in seconds.
 
         ``np.nan`` appended so the length of the output is the same as number of waypoints.
@@ -427,14 +427,14 @@ def segment_duration(self, dtype: npt.DTypeLike = np.float32) -> npt.NDArray[np.
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Time difference between waypoints, [:math:`s`].
             Returns an array with dtype specified by``dtype``
         """
 
         return segment_duration(self.data["time"], dtype=dtype)
 
-    def segment_haversine(self) -> npt.NDArray[np.float64]:
+    def segment_haversine(self) -> npt.NDArray[np.floating]:
         """Compute Haversine (great circle) distance between flight waypoints.
 
         Helper function used in :meth:`resample_and_fill`.
@@ -445,7 +445,7 @@ def segment_haversine(self) -> npt.NDArray[np.float64]:
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Array of great circle distances in [:math:`m`] between waypoints
 
         Examples
@@ -468,7 +468,7 @@ def segment_haversine(self) -> npt.NDArray[np.float64]:
         """
         return geo.segment_haversine(self["longitude"], self["latitude"])
 
-    def segment_length(self) -> npt.NDArray[np.float64]:
+    def segment_length(self) -> npt.NDArray[np.floating]:
         """Compute spherical distance between flight waypoints.
 
         Helper function used in :meth:`length` and :meth:`length_met`.
@@ -476,7 +476,7 @@ def segment_length(self) -> npt.NDArray[np.float64]:
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Array of distances in [:math:`m`] between waypoints
 
         Examples
@@ -498,7 +498,7 @@ def segment_length(self) -> npt.NDArray[np.float64]:
         """
         return geo.segment_length(self["longitude"], self["latitude"], self.altitude)
 
-    def segment_angle(self) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    def segment_angle(self) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
         """Calculate sine and cosine for the angle between each segment and the longitudinal axis.
 
         This is different from the usual navigational angle between two points known as *bearing*.
@@ -518,7 +518,7 @@ def segment_angle(self) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64
 
         Returns
         -------
-        npt.NDArray[np.float64], npt.NDArray[np.float64]
+        npt.NDArray[np.floating], npt.NDArray[np.floating]
             Returns ``sin(a), cos(a)``, where ``a`` is the angle between the segment and the
             longitudinal axis. The final values are of both arrays are ``np.nan``.
 
@@ -548,7 +548,7 @@ def segment_angle(self) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64
         """
         return geo.segment_angle(self["longitude"], self["latitude"])
 
-    def segment_azimuth(self) -> npt.NDArray[np.float64]:
+    def segment_azimuth(self) -> npt.NDArray[np.floating]:
         """Calculate (forward) azimuth at each waypoint.
 
         Method calls `pyproj.Geod.inv`, which is slow. See `geo.forward_azimuth`
@@ -560,7 +560,7 @@ def segment_azimuth(self) -> npt.NDArray[np.float64]:
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Array of azimuths.
 
         See Also
@@ -599,7 +599,7 @@ def segment_azimuth(self) -> npt.NDArray[np.float64]:
 
     def segment_groundspeed(
         self, smooth: bool = False, window_length: int = 7, polyorder: int = 1
-    ) -> npt.NDArray[np.float64]:
+    ) -> npt.NDArray[np.floating]:
         """Return groundspeed across segments.
 
         Calculate by dividing the horizontal segment length by the difference in waypoint times.
@@ -616,7 +616,7 @@ def segment_groundspeed(
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Groundspeed of the segment, [:math:`m s^{-1}`]
         """
         # get horizontal distance (altitude is ignored)
@@ -638,22 +638,22 @@ def segment_groundspeed(
 
     def segment_true_airspeed(
         self,
-        u_wind: npt.NDArray[np.float64] | float = 0.0,
-        v_wind: npt.NDArray[np.float64] | float = 0.0,
+        u_wind: npt.NDArray[np.floating] | float = 0.0,
+        v_wind: npt.NDArray[np.floating] | float = 0.0,
         smooth: bool = True,
         window_length: int = 7,
         polyorder: int = 1,
-    ) -> npt.NDArray[np.float64]:
+    ) -> npt.NDArray[np.floating]:
         r"""Calculate the true airspeed [:math:`m/s`] from the ground speed and horizontal winds.
 
         The calculated ground speed will first be smoothed with a Savitzky-Golay filter if enabled.
 
         Parameters
         ----------
-        u_wind : npt.NDArray[np.float64] | float
+        u_wind : npt.NDArray[np.floating] | float
             U wind speed, [:math:`m \ s^{-1}`].
             Defaults to 0 for all waypoints.
-        v_wind : npt.NDArray[np.float64] | float
+        v_wind : npt.NDArray[np.floating] | float
             V wind speed, [:math:`m \ s^{-1}`].
             Defaults to 0 for all waypoints.
         smooth : bool, optional
@@ -666,7 +666,7 @@ def segment_true_airspeed(
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             True wind speed of each segment, [:math:`m \ s^{-1}`]
         """
         groundspeed = self.segment_groundspeed(smooth, window_length, polyorder)
@@ -680,39 +680,39 @@ def segment_true_airspeed(
         return np.sqrt(tas_x * tas_x + tas_y * tas_y)
 
     def segment_mach_number(
-        self, true_airspeed: npt.NDArray[np.float64], air_temperature: npt.NDArray[np.float64]
-    ) -> npt.NDArray[np.float64]:
+        self, true_airspeed: npt.NDArray[np.floating], air_temperature: npt.NDArray[np.floating]
+    ) -> npt.NDArray[np.floating]:
         r"""Calculate the mach number of each segment.
 
         Parameters
         ----------
-        true_airspeed : npt.NDArray[np.float64]
+        true_airspeed : npt.NDArray[np.floating]
             True airspeed of the segment, [:math:`m \ s^{-1}`].
             See :meth:`segment_true_airspeed`.
-        air_temperature : npt.NDArray[np.float64]
+        air_temperature : npt.NDArray[np.floating]
             Average air temperature of each segment, [:math:`K`]
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Mach number of each segment
         """
         return units.tas_to_mach_number(true_airspeed, air_temperature)
 
     def segment_rocd(
         self,
-        air_temperature: None | npt.NDArray[np.float64] = None,
-    ) -> npt.NDArray[np.float64]:
+        air_temperature: None | npt.NDArray[np.floating] = None,
+    ) -> npt.NDArray[np.floating]:
         """Calculate the rate of climb and descent (ROCD).
 
         Parameters
         ----------
-        air_temperature: None | npt.NDArray[np.float64]
+        air_temperature: None | npt.NDArray[np.floating]
             Air temperature of each flight waypoint, [:math:`K`]
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Rate of climb and descent over segment, [:math:`ft min^{-1}`]
 
         See Also
@@ -725,7 +725,7 @@ def segment_phase(
         self,
         threshold_rocd: float = 250.0,
         min_cruise_altitude_ft: float = 20000.0,
-        air_temperature: None | npt.NDArray[np.float64] = None,
+        air_temperature: None | npt.NDArray[np.floating] = None,
     ) -> npt.NDArray[np.uint8]:
         """Identify the phase of flight (climb, cruise, descent) for each segment.
 
@@ -739,7 +739,7 @@ def segment_phase(
             This is specific for each aircraft type,
             and can be approximated as 50% of the altitude ceiling.
             Defaults to 20000 ft.
-        air_temperature: None | npt.NDArray[np.float64]
+        air_temperature: None | npt.NDArray[np.floating]
             Air temperature of each flight waypoint, [:math:`K`]
 
         Returns
@@ -1606,12 +1606,12 @@ def plot_profile(self, **kwargs: Any) -> matplotlib.axes.Axes:
         return ax
 
 
-def _return_linestring(data: dict[str, npt.NDArray[np.float64]]) -> list[list[float]]:
+def _return_linestring(data: dict[str, npt.NDArray[np.floating]]) -> list[list[float]]:
     """Return list of coordinates for geojson constructions.
 
     Parameters
     ----------
-    data : dict[str, npt.NDArray[np.float64]]
+    data : dict[str, npt.NDArray[np.floating]]
         :attr:`data` containing `longitude`, `latitude`, and `altitude` keys
 
     Returns
@@ -1659,15 +1659,15 @@ def _antimeridian_index(longitude: pd.Series) -> list[int]:
 
 
 def _sg_filter(
-    vals: npt.NDArray[np.float64], window_length: int = 7, polyorder: int = 1
-) -> npt.NDArray[np.float64]:
+    vals: npt.NDArray[np.floating], window_length: int = 7, polyorder: int = 1
+) -> npt.NDArray[np.floating]:
     """Apply Savitzky-Golay filter to smooth out noise in the time-series data.
 
     Used to smooth true airspeed, fuel flow, and altitude.
 
     Parameters
     ----------
-    vals : npt.NDArray[np.float64]
+    vals : npt.NDArray[np.floating]
         Input array
     window_length : int, optional
         Parameter for :func:`scipy.signal.savgol_filter`
@@ -1676,7 +1676,7 @@ def _sg_filter(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Smoothed values
 
     Raises
@@ -1702,11 +1702,11 @@ def _sg_filter(
 
 
 def _altitude_interpolation(
-    altitude: npt.NDArray[np.float64],
+    altitude: npt.NDArray[np.floating],
     nominal_rocd: float,
     freq: np.timedelta64,
     climb_or_descend_at_end: bool = False,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Interpolate nan values in ``altitude`` array.
 
     Suppose each group of consecutive nan values is enclosed by ``a0`` and ``a1`` with
@@ -1718,7 +1718,7 @@ def _altitude_interpolation(
 
     Parameters
     ----------
-    altitude : npt.NDArray[np.float64]
+    altitude : npt.NDArray[np.floating]
         Array of altitude values containing nan values. This function will raise
         an error if ``altitude`` does not contain nan values. Moreover, this function
         assumes the initial and final entries in ``altitude`` are not nan.
@@ -1732,7 +1732,7 @@ def _altitude_interpolation(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Altitude after nan values have been filled
     """
     # Determine nan state of altitude
@@ -1762,12 +1762,12 @@ def _altitude_interpolation(
 
 
 def _altitude_interpolation_climb_descend_end(
-    altitude: npt.NDArray[np.float64],
+    altitude: npt.NDArray[np.floating],
     na_group_size: npt.NDArray[np.intp],
     nominal_rocd: float,
     freq: np.timedelta64,
     isna: npt.NDArray[np.bool_],
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Interpolate altitude values by placing climbs/descents at end of nan sequences.
 
     The segment will remain at constant elevation until the end of the segment where
@@ -1776,7 +1776,7 @@ def _altitude_interpolation_climb_descend_end(
 
     Parameters
     ----------
-    altitude : npt.NDArray[np.float64]
+    altitude : npt.NDArray[np.floating]
         Array of altitude values containing nan values. This function will raise
         an error if ``altitude`` does not contain nan values. Moreover, this function
         assumes the initial and final entries in ``altitude`` are not nan.
@@ -1791,7 +1791,7 @@ def _altitude_interpolation_climb_descend_end(
         Array of boolean values indicating whether or not each entry in `altitude`
         is nan-valued.
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Altitude after nan values have been filled
     """
     cumalt_list = [np.flip(np.arange(1, size, dtype=float)) for size in na_group_size]
@@ -1823,14 +1823,14 @@ def _altitude_interpolation_climb_descend_end(
 
 
 def _altitude_interpolation_climb_descend_middle(
-    altitude: npt.NDArray[np.float64],
+    altitude: npt.NDArray[np.floating],
     start_na_idxs: npt.NDArray[np.intp],
     end_na_idxs: npt.NDArray[np.intp],
     na_group_size: npt.NDArray[np.intp],
     freq: np.timedelta64,
     nominal_rocd: float,
     isna: npt.NDArray[np.bool_],
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Interpolate nan altitude values based on step-climb logic.
 
     For short segments, the climb will be placed at the begining of the segment. For
@@ -1839,7 +1839,7 @@ def _altitude_interpolation_climb_descend_middle(
 
     Parameters
     ----------
-    altitude : npt.NDArray[np.float64]
+    altitude : npt.NDArray[np.floating]
         Array of altitude values containing nan values. This function will raise
         an error if ``altitude`` does not contain nan values. Moreover, this function
         assumes the initial and final entries in ``altitude`` are not nan.
@@ -1860,7 +1860,7 @@ def _altitude_interpolation_climb_descend_middle(
         Array of boolean values indicating whether or not each entry in `altitude`
         is nan-valued.
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Altitude after nan values have been filled
     """
     s = pd.Series(altitude)
@@ -1900,7 +1900,7 @@ def _altitude_interpolation_climb_descend_middle(
         if s[start_na_idx] <= s[end_na_idx]:
             cumalt_list.append(np.arange(1, size, dtype=float))
         else:
-            cumalt_list.append(np.flip(np.arange(1, size, dtype=float)))
+            cumalt_list.append(np.flip(np.arange(1, size, dtype=float)))  # type: ignore[arg-type]
 
     cumalt = np.concatenate(cumalt_list)
     cumalt = cumalt * nominal_rocd * (freq / np.timedelta64(1, "s"))
@@ -1925,13 +1925,13 @@ def _altitude_interpolation_climb_descend_middle(
 
 
 def _verify_altitude(
-    altitude: npt.NDArray[np.float64], nominal_rocd: float, freq: np.timedelta64
+    altitude: npt.NDArray[np.floating], nominal_rocd: float, freq: np.timedelta64
 ) -> None:
     """Confirm that the time derivative of `altitude` does not exceed twice `nominal_rocd`.
 
     Parameters
     ----------
-    altitude : npt.NDArray[np.float64]
+    altitude : npt.NDArray[np.floating]
         Array of filled altitude values containing nan values.
     nominal_rocd : float
         Nominal rate of climb/descent, in m/s
@@ -1954,11 +1954,11 @@ def _verify_altitude(
 
 def filter_altitude(
     time: npt.NDArray[np.datetime64],
-    altitude_ft: npt.NDArray[np.float64],
+    altitude_ft: npt.NDArray[np.floating],
     kernel_size: int = 17,
     cruise_threshold: float = 120,
-    air_temperature: None | npt.NDArray[np.float64] = None,
-) -> npt.NDArray[np.float64]:
+    air_temperature: None | npt.NDArray[np.floating] = None,
+) -> npt.NDArray[np.floating]:
     """
     Filter noisy altitude on a single flight.
 
@@ -1970,19 +1970,19 @@ def filter_altitude(
     ----------
     time : npt.NDArray[np.datetime64]
         Waypoint time in ``np.datetime64`` format.
-    altitude_ft : npt.NDArray[np.float64]
+    altitude_ft : npt.NDArray[np.floating]
         Altitude signal in feet
     kernel_size : int, optional
         Passed directly to :func:`scipy.signal.medfilt`, by default 11.
         Passed also to :func:`scipy.signal.medfilt`
     cruise_theshold : int, optional
         Minimal length of time, in seconds, for a flight to be in cruise to apply median filter
-    air_temperature: None | npt.NDArray[np.float64]
+    air_temperature: None | npt.NDArray[np.floating]
         Air temperature of each flight waypoint, [:math:`K`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Filtered altitude
 
     Notes
@@ -2072,7 +2072,7 @@ def filter_altitude(
 
 def segment_duration(
     time: npt.NDArray[np.datetime64], dtype: npt.DTypeLike = np.float32
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Calculate the time difference between waypoints.
 
     ``np.nan`` appended so the length of the output is the same as number of waypoints.
@@ -2087,7 +2087,7 @@ def segment_duration(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Time difference between waypoints, [:math:`s`].
         This returns an array with dtype specified by``dtype``.
     """
@@ -2098,8 +2098,8 @@ def segment_duration(
 
 
 def segment_phase(
-    rocd: npt.NDArray[np.float64],
-    altitude_ft: npt.NDArray[np.float64],
+    rocd: npt.NDArray[np.floating],
+    altitude_ft: npt.NDArray[np.floating],
     *,
     threshold_rocd: float = 250.0,
     min_cruise_altitude_ft: float = MIN_CRUISE_ALTITUDE,
@@ -2111,7 +2111,7 @@ def segment_phase(
     rocd: pt.NDArray[np.float64]
         Rate of climb and descent across segment, [:math:`ft min^{-1}`].
         See output from :func:`segment_rocd`.
-    altitude_ft: npt.NDArray[np.float64]
+    altitude_ft: npt.NDArray[np.floating]
         Altitude, [:math:`ft`]
     threshold_rocd: float, optional
         ROCD threshold to identify climb and descent, [:math:`ft min^{-1}`].
@@ -2161,26 +2161,26 @@ def segment_phase(
 
 
 def segment_rocd(
-    segment_duration: npt.NDArray[np.float64],
-    altitude_ft: npt.NDArray[np.float64],
-    air_temperature: None | npt.NDArray[np.float64] = None,
-) -> npt.NDArray[np.float64]:
+    segment_duration: npt.NDArray[np.floating],
+    altitude_ft: npt.NDArray[np.floating],
+    air_temperature: None | npt.NDArray[np.floating] = None,
+) -> npt.NDArray[np.floating]:
     """Calculate the rate of climb and descent (ROCD).
 
     Parameters
     ----------
-    segment_duration: npt.NDArray[np.float64]
+    segment_duration: npt.NDArray[np.floating]
         Time difference between waypoints, [:math:`s`].
         Expected to have numeric `dtype`, not `"timedelta64"`.
         See output from :func:`segment_duration`.
-    altitude_ft: npt.NDArray[np.float64]
+    altitude_ft: npt.NDArray[np.floating]
         Altitude of each waypoint, [:math:`ft`]
-    air_temperature: None | npt.NDArray[np.float64]
+    air_temperature: None | npt.NDArray[np.floating]
         Air temperature of each flight waypoint, [:math:`K`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Rate of climb and descent over segment, [:math:`ft min^{-1}`]
 
     Notes
@@ -2211,7 +2211,7 @@ def segment_rocd(
     T_correction[:-1] = (air_temperature[:-1] + air_temperature[1:]) / (T_isa[:-1] + T_isa[1:])
     T_correction[-1] = np.nan
 
-    return T_correction * out
+    return T_correction * out  # type: ignore[return-value]
 
 
 def _resample_to_freq(df: pd.DataFrame, freq: str) -> tuple[pd.DataFrame, pd.DatetimeIndex]:
diff --git a/pycontrails/core/interpolation.py b/pycontrails/core/interpolation.py
index 0708f2898..bb771e54b 100644
--- a/pycontrails/core/interpolation.py
+++ b/pycontrails/core/interpolation.py
@@ -43,9 +43,9 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
 
     Parameters
     ----------
-    points : tuple[npt.NDArray[np.float64], ...]
+    points : tuple[npt.NDArray[np.floating], ...]
         Coordinates of the grid points.
-    values : npt.NDArray[np.float64]
+    values : npt.NDArray[np.floating]
         Grid values. The shape of this array must be compatible with the
         coordinates. An error is raised if the dtype is not ``np.float32``
         or ``np.float64``.
@@ -59,8 +59,8 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
 
     def __init__(
         self,
-        points: tuple[npt.NDArray[np.float64], ...],
-        values: npt.NDArray[np.float64],
+        points: tuple[npt.NDArray[np.floating], ...],
+        values: npt.NDArray[np.floating],
         method: str,
         bounds_error: bool,
         fill_value: float | np.float64 | None,
@@ -78,12 +78,12 @@ def __init__(
         self.fill_value = fill_value
         self._spline = None
 
-    def _prepare_xi_simple(self, xi: npt.NDArray[np.float64]) -> npt.NDArray[np.bool_]:
+    def _prepare_xi_simple(self, xi: npt.NDArray[np.floating]) -> npt.NDArray[np.bool_]:
         """Run looser version of :meth:`_prepare_xi`.
 
         Parameters
         ----------
-        xi : npt.NDArray[np.float64]
+        xi : npt.NDArray[np.floating]
             Points at which to interpolate.
 
         Returns
@@ -106,13 +106,13 @@ def _prepare_xi_simple(self, xi: npt.NDArray[np.float64]) -> npt.NDArray[np.bool
         return self._find_out_of_bounds(xi.T)
 
     def __call__(
-        self, xi: npt.NDArray[np.float64], method: str | None = None
-    ) -> npt.NDArray[np.float64]:
+        self, xi: npt.NDArray[np.floating], method: str | None = None
+    ) -> npt.NDArray[np.floating]:
         """Evaluate the interpolator at the given points.
 
         Parameters
         ----------
-        xi : npt.NDArray[np.float64]
+        xi : npt.NDArray[np.floating]
             Points at which to interpolate. Must have shape ``(n, ndim)``, where
             ``ndim`` is the number of dimensions of the interpolator.
         method : str | None
@@ -120,7 +120,7 @@ def __call__(
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Interpolated values. Has shape ``(n,)``. When computing linear interpolation,
             the dtype is the same as the :attr:`values` array.
         """
@@ -137,21 +137,21 @@ def __call__(
 
     def _set_out_of_bounds(
         self,
-        out: npt.NDArray[np.float64],
+        out: npt.NDArray[np.floating],
         out_of_bounds: npt.NDArray[np.bool_],
-    ) -> npt.NDArray[np.float64]:
+    ) -> npt.NDArray[np.floating]:
         """Set out-of-bounds values to the fill value.
 
         Parameters
         ----------
-        out : npt.NDArray[np.float64]
+        out : npt.NDArray[np.floating]
             Values from interpolation. This is modified in-place.
         out_of_bounds : npt.NDArray[np.bool_]
             A 1-dimensional Boolean array indicating which points are out of bounds.
 
         Returns
         -------
-        out : npt.NDArray[np.float64]
+        out : npt.NDArray[np.floating]
             A reference to the ``out`` array.
         """
         if self.fill_value is not None and np.any(out_of_bounds):
@@ -162,8 +162,8 @@ def _set_out_of_bounds(
     def _evaluate_linear(
         self,
         indices: npt.NDArray[np.int64],
-        norm_distances: npt.NDArray[np.float64],
-    ) -> npt.NDArray[np.float64]:
+        norm_distances: npt.NDArray[np.floating],
+    ) -> npt.NDArray[np.floating]:
         """Evaluate the interpolator using linear interpolation.
 
         This is a faster alternative to
@@ -181,13 +181,13 @@ def _evaluate_linear(
         indices : npt.NDArray[np.int64]
             Indices of the grid points to the left of the interpolation points.
             Has shape ``(ndim, n_points)``.
-        norm_distances : npt.NDArray[np.float64]
+        norm_distances : npt.NDArray[np.floating]
             Normalized distances between the interpolation points and the grid
             points to the left. Has shape ``(ndim, n_points)``.
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Interpolated values with shape ``(n_points,)`` and the same dtype as
             the :attr:`values` array.
         """
@@ -261,7 +261,7 @@ def _pick_method(scipy_version: str, method: str) -> str:
 
 def _floatize_time(
     time: npt.NDArray[np.datetime64], offset: np.datetime64
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Convert an array of ``np.datetime64`` to an array of ``np.float64``.
 
     In calls to :class:`scipy.interpolate.RegularGridInterpolator`, it's critical
@@ -296,7 +296,7 @@ def _floatize_time(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         The number of milliseconds since ``offset``.
     """
     delta = time - offset
@@ -354,9 +354,9 @@ def _localize(da: xr.DataArray, coords: dict[str, np.ndarray]) -> xr.DataArray:
 
 @overload
 def interp(
-    longitude: npt.NDArray[np.float64],
-    latitude: npt.NDArray[np.float64],
-    level: npt.NDArray[np.float64],
+    longitude: npt.NDArray[np.floating],
+    latitude: npt.NDArray[np.floating],
+    level: npt.NDArray[np.floating],
     time: npt.NDArray[np.datetime64],
     da: xr.DataArray,
     method: str,
@@ -366,14 +366,14 @@ def interp(
     *,
     indices: RGIArtifacts | None = ...,
     return_indices: Literal[False] = ...,
-) -> npt.NDArray[np.float64]: ...
+) -> npt.NDArray[np.floating]: ...
 
 
 @overload
 def interp(
-    longitude: npt.NDArray[np.float64],
-    latitude: npt.NDArray[np.float64],
-    level: npt.NDArray[np.float64],
+    longitude: npt.NDArray[np.floating],
+    latitude: npt.NDArray[np.floating],
+    level: npt.NDArray[np.floating],
     time: npt.NDArray[np.datetime64],
     da: xr.DataArray,
     method: str,
@@ -383,14 +383,14 @@ def interp(
     *,
     indices: RGIArtifacts | None = ...,
     return_indices: Literal[True],
-) -> tuple[npt.NDArray[np.float64], RGIArtifacts]: ...
+) -> tuple[npt.NDArray[np.floating], RGIArtifacts]: ...
 
 
 @overload
 def interp(
-    longitude: npt.NDArray[np.float64],
-    latitude: npt.NDArray[np.float64],
-    level: npt.NDArray[np.float64],
+    longitude: npt.NDArray[np.floating],
+    latitude: npt.NDArray[np.floating],
+    level: npt.NDArray[np.floating],
     time: npt.NDArray[np.datetime64],
     da: xr.DataArray,
     method: str,
@@ -400,13 +400,13 @@ def interp(
     *,
     indices: RGIArtifacts | None = ...,
     return_indices: bool = ...,
-) -> npt.NDArray[np.float64] | tuple[npt.NDArray[np.float64], RGIArtifacts]: ...
+) -> npt.NDArray[np.floating] | tuple[npt.NDArray[np.floating], RGIArtifacts]: ...
 
 
 def interp(
-    longitude: npt.NDArray[np.float64],
-    latitude: npt.NDArray[np.float64],
-    level: npt.NDArray[np.float64],
+    longitude: npt.NDArray[np.floating],
+    latitude: npt.NDArray[np.floating],
+    level: npt.NDArray[np.floating],
     time: npt.NDArray[np.datetime64],
     da: xr.DataArray,
     method: str,
@@ -416,7 +416,7 @@ def interp(
     *,
     indices: RGIArtifacts | None = None,
     return_indices: bool = False,
-) -> npt.NDArray[np.float64] | tuple[npt.NDArray[np.float64], RGIArtifacts]:
+) -> npt.NDArray[np.floating] | tuple[npt.NDArray[np.floating], RGIArtifacts]:
     """Interpolate over a grid with ``localize`` option.
 
     .. versionchanged:: 0.25.6
@@ -471,7 +471,7 @@ def interp(
 
     Returns
     -------
-    npt.NDArray[np.float64] | tuple[npt.NDArray[np.float64], RGIArtifacts]
+    npt.NDArray[np.floating] | tuple[npt.NDArray[np.floating], RGIArtifacts]
         Interpolated values with same size as ``longitude``. If ``return_indices``
         is True, return intermediate indices artifacts as well.
 
@@ -499,7 +499,7 @@ def interp(
     t = _floatize_time(t, offset)
 
     single_level = z.size == 1 and z.item() == -1.0
-    points: tuple[npt.NDArray[np.float64], ...]
+    points: tuple[npt.NDArray[np.floating], ...]
     if single_level:
         values = da.values.squeeze(axis=2)
         points = x, y, t
@@ -529,13 +529,13 @@ def interp(
 
 
 def _buildxi(
-    longitude: npt.NDArray[np.float64],
-    latitude: npt.NDArray[np.float64],
-    level: npt.NDArray[np.float64],
+    longitude: npt.NDArray[np.floating],
+    latitude: npt.NDArray[np.floating],
+    level: npt.NDArray[np.floating],
     time: npt.NDArray[np.datetime64],
     offset: np.datetime64,
     single_level: bool,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Build the input array for interpolation.
 
     The implementation below achieves the same result as the following::
@@ -562,10 +562,10 @@ def _buildxi(
 
 def _linear_interp_with_indices(
     interp: PycontrailsRegularGridInterpolator,
-    xi: npt.NDArray[np.float64] | None,
+    xi: npt.NDArray[np.floating] | None,
     localize: bool,
     indices: RGIArtifacts | None,
-) -> tuple[npt.NDArray[np.float64], RGIArtifacts]:
+) -> tuple[npt.NDArray[np.floating], RGIArtifacts]:
     if interp.method != "linear":
         msg = "Parameter 'indices' is only supported for 'method=linear'"
         raise ValueError(msg)
@@ -589,7 +589,7 @@ class RGIArtifacts:
     """An interface to intermediate RGI interpolation artifacts."""
 
     xi_indices: npt.NDArray[np.int64]
-    norm_distances: npt.NDArray[np.float64]
+    norm_distances: npt.NDArray[np.floating]
     out_of_bounds: npt.NDArray[np.bool_]
 
 
@@ -608,10 +608,10 @@ class EmissionsProfileInterpolator:
 
     Parameters
     ----------
-    xp : npt.NDarray[np.float64]
+    xp : npt.NDArray[np.floating]
         Array of x-values. These must be strictly increasing and free from
         any nan values. Passed to :func:`numpy.interp`.
-    fp : npt.NDarray[np.float64]
+    fp : npt.NDArray[np.floating]
         Array of y-values. Passed to :func:`numpy.interp`.
     drop_duplicates : bool, optional
         Whether to drop duplicate values in ``xp``. Defaults to ``True``.
@@ -648,8 +648,8 @@ class EmissionsProfileInterpolator:
 
     def __init__(
         self,
-        xp: npt.NDArray[np.float64],
-        fp: npt.NDArray[np.float64],
+        xp: npt.NDArray[np.floating],
+        fp: npt.NDArray[np.floating],
         drop_duplicates: bool = True,
     ) -> None:
         if drop_duplicates:
@@ -680,17 +680,17 @@ def _validate(self) -> None:
             msg = "xp must not contain nan values"
             raise ValueError(msg)
 
-    def interp(self, x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+    def interp(self, x: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
         """Interpolate x against xp and fp.
 
         Parameters
         ----------
-        x : npt.NDArray[np.float64]
+        x : npt.NDArray[np.floating]
             Array of x-values to interpolate.
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Array of interpolated y-values arising from the x-values. The ``dtype`` of
             the output array is the same as the ``dtype`` of ``x``.
         """
@@ -700,7 +700,7 @@ def interp(self, x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
         dtype = np.result_type(x, np.float32)
         return np.interp(x, self.xp, self.fp).astype(dtype)
 
-    def log_interp(self, x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+    def log_interp(self, x: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
         """Interpolate x against xp and fp on a logarithmic scale.
 
         This method composes the following three functions.
@@ -710,12 +710,12 @@ def log_interp(self, x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
 
         Parameters
         ----------
-        x : npt.NDArray[np.float64]
+        x : npt.NDArray[np.floating]
             Array of x-values to interpolate.
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Array of interpolated y-values arising from the x-values.
         """
         return np.exp(self.interp(np.log(x)))
diff --git a/pycontrails/core/met.py b/pycontrails/core/met.py
index 49391713d..4b6ffe218 100644
--- a/pycontrails/core/met.py
+++ b/pycontrails/core/met.py
@@ -1495,9 +1495,9 @@ def in_memory(self) -> bool:
     @overload
     def interpolate(
         self,
-        longitude: float | npt.NDArray[np.float64],
-        latitude: float | npt.NDArray[np.float64],
-        level: float | npt.NDArray[np.float64],
+        longitude: float | npt.NDArray[np.floating],
+        latitude: float | npt.NDArray[np.floating],
+        level: float | npt.NDArray[np.floating],
         time: np.datetime64 | npt.NDArray[np.datetime64],
         *,
         method: str = ...,
@@ -1507,14 +1507,14 @@ def interpolate(
         lowmem: bool = ...,
         indices: interpolation.RGIArtifacts | None = ...,
         return_indices: Literal[False] = ...,
-    ) -> npt.NDArray[np.float64]: ...
+    ) -> npt.NDArray[np.floating]: ...
 
     @overload
     def interpolate(
         self,
-        longitude: float | npt.NDArray[np.float64],
-        latitude: float | npt.NDArray[np.float64],
-        level: float | npt.NDArray[np.float64],
+        longitude: float | npt.NDArray[np.floating],
+        latitude: float | npt.NDArray[np.floating],
+        level: float | npt.NDArray[np.floating],
         time: np.datetime64 | npt.NDArray[np.datetime64],
         *,
         method: str = ...,
@@ -1524,13 +1524,13 @@ def interpolate(
         lowmem: bool = ...,
         indices: interpolation.RGIArtifacts | None = ...,
         return_indices: Literal[True],
-    ) -> tuple[npt.NDArray[np.float64], interpolation.RGIArtifacts]: ...
+    ) -> tuple[npt.NDArray[np.floating], interpolation.RGIArtifacts]: ...
 
     def interpolate(
         self,
-        longitude: float | npt.NDArray[np.float64],
-        latitude: float | npt.NDArray[np.float64],
-        level: float | npt.NDArray[np.float64],
+        longitude: float | npt.NDArray[np.floating],
+        latitude: float | npt.NDArray[np.floating],
+        level: float | npt.NDArray[np.floating],
         time: np.datetime64 | npt.NDArray[np.datetime64],
         *,
         method: str = "linear",
@@ -1540,7 +1540,7 @@ def interpolate(
         lowmem: bool = False,
         indices: interpolation.RGIArtifacts | None = None,
         return_indices: bool = False,
-    ) -> npt.NDArray[np.float64] | tuple[npt.NDArray[np.float64], interpolation.RGIArtifacts]:
+    ) -> npt.NDArray[np.floating] | tuple[npt.NDArray[np.floating], interpolation.RGIArtifacts]:
         """Interpolate values over underlying DataArray.
 
         Zero dimensional coordinates are reshaped to 1D arrays.
@@ -1569,11 +1569,11 @@ def interpolate(
 
         Parameters
         ----------
-        longitude : float | npt.NDArray[np.float64]
+        longitude : float | npt.NDArray[np.floating]
             Longitude values to interpolate. Assumed to be 0 or 1 dimensional.
-        latitude : float | npt.NDArray[np.float64]
+        latitude : float | npt.NDArray[np.floating]
             Latitude values to interpolate. Assumed to be 0 or 1 dimensional.
-        level : float | npt.NDArray[np.float64]
+        level : float | npt.NDArray[np.floating]
             Level values to interpolate. Assumed to be 0 or 1 dimensional.
         time : np.datetime64 | npt.NDArray[np.datetime64]
             Time values to interpolate. Assumed to be 0 or 1 dimensional.
@@ -1696,18 +1696,17 @@ def interpolate(
 
     def _interp_lowmem(
         self,
-        longitude: float | npt.NDArray[np.float64],
-        latitude: float | npt.NDArray[np.float64],
-        level: float | npt.NDArray[np.float64],
+        longitude: float | npt.NDArray[np.floating],
+        latitude: float | npt.NDArray[np.floating],
+        level: float | npt.NDArray[np.floating],
         time: np.datetime64 | npt.NDArray[np.datetime64],
         *,
         method: str = "linear",
         bounds_error: bool = False,
         fill_value: float | np.float64 | None = np.nan,
-        minimize_memory: bool = False,
         indices: interpolation.RGIArtifacts | None = None,
         return_indices: bool = False,
-    ) -> npt.NDArray[np.float64] | tuple[npt.NDArray[np.float64], interpolation.RGIArtifacts]:
+    ) -> npt.NDArray[np.floating] | tuple[npt.NDArray[np.floating], interpolation.RGIArtifacts]:
         """Interpolate values against underlying DataArray.
 
         This method is used by :meth:`interpolate` when ``lowmem=True``.
@@ -1762,27 +1761,37 @@ def _interp_lowmem(
                 )
                 da.load()
 
-            tmp = interpolation.interp(
-                longitude=lon_sl,
-                latitude=lat_sl,
-                level=lev_sl,
-                time=t_sl,
-                da=da,
-                method=method,
-                bounds_error=bounds_error,
-                fill_value=fill_value,
-                localize=False,  # would be no-op; da is localized already
-                indices=indices_sl,
-                return_indices=return_indices,
-            )
-
             if return_indices:
-                out[mask], rgi_sl = tmp
+                out[mask], rgi_sl = interpolation.interp(
+                    longitude=lon_sl,
+                    latitude=lat_sl,
+                    level=lev_sl,
+                    time=t_sl,
+                    da=da,
+                    method=method,
+                    bounds_error=bounds_error,
+                    fill_value=fill_value,
+                    localize=False,  # would be no-op; da is localized already
+                    indices=indices_sl,
+                    return_indices=return_indices,
+                )
                 rgi_artifacts.xi_indices[:, mask] = rgi_sl.xi_indices
                 rgi_artifacts.norm_distances[:, mask] = rgi_sl.norm_distances
                 rgi_artifacts.out_of_bounds[mask] = rgi_sl.out_of_bounds
             else:
-                out[mask] = tmp
+                out[mask] = interpolation.interp(
+                    longitude=lon_sl,
+                    latitude=lat_sl,
+                    level=lev_sl,
+                    time=t_sl,
+                    da=da,
+                    method=method,
+                    bounds_error=bounds_error,
+                    fill_value=fill_value,
+                    localize=False,  # would be no-op; da is localized already
+                    indices=indices_sl,
+                    return_indices=return_indices,
+                )
 
         if return_indices:
             return out, rgi_artifacts
@@ -1878,7 +1887,7 @@ def proportion(self) -> float:
         if not self.binary:
             raise NotImplementedError("proportion method is only implemented for binary fields")
 
-        return self.data.sum().values.item() / self.data.count().values.item()
+        return self.data.sum().values.item() / self.data.count().values.item()  # type: ignore[operator]
 
     def find_edges(self) -> MetDataArray:
         """Find edges of regions.
@@ -2592,9 +2601,9 @@ def _extract_2d_arr_and_altitude(
     except KeyError:
         altitude = None
     else:
-        altitude = round(altitude)
+        altitude = round(altitude)  # type: ignore[call-overload]
 
-    return arr, altitude
+    return arr, altitude  # type: ignore[return-value]
 
 
 def downselect(data: XArrayType, bbox: tuple[float, ...]) -> XArrayType:
diff --git a/pycontrails/core/models.py b/pycontrails/core/models.py
index 2db061707..bb6f75c7c 100644
--- a/pycontrails/core/models.py
+++ b/pycontrails/core/models.py
@@ -825,7 +825,7 @@ def interpolate_met(
     *,
     q_method: str | None = None,
     **interp_kwargs: Any,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Interpolate ``vector`` against ``met`` gridded data.
 
     If ``vector_key`` (=``met_key`` by default) already exists,
@@ -854,7 +854,7 @@ def interpolate_met(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Interpolated values.
 
     Raises
@@ -933,15 +933,15 @@ def _extract_q(met: MetDataset, met_key: str, q_method: str) -> tuple[MetDataArr
 
 
 def _prepare_q(
-    mda: MetDataArray, level: npt.NDArray[np.float64], q_method: str, log_applied: bool
-) -> tuple[MetDataArray, npt.NDArray[np.float64]]:
+    mda: MetDataArray, level: npt.NDArray[np.floating], q_method: str, log_applied: bool
+) -> tuple[MetDataArray, npt.NDArray[np.floating]]:
     """Prepare specific humidity for interpolation with experimental ``q_method``.
 
     Parameters
     ----------
     mda : MetDataArray
         MetDataArray of specific humidity.
-    level : npt.NDArray[np.float64]
+    level : npt.NDArray[np.floating]
         Levels to interpolate to, [:math:`hPa`].
     q_method : str
         One of ``"log-q-log-p"`` or ``"cubic-spline"``.
@@ -952,7 +952,7 @@ def _prepare_q(
     -------
     mda : MetDataArray
         MetDataArray of specific humidity transformed for interpolation.
-    level : npt.NDArray[np.float64]
+    level : npt.NDArray[np.floating]
         Transformed levels for interpolation.
     """
     da = mda.data
@@ -975,8 +975,8 @@ def _prepare_q(
 
 
 def _prepare_q_log_q_log_p(
-    da: xr.DataArray, level: npt.NDArray[np.float64], log_applied: bool
-) -> tuple[MetDataArray, npt.NDArray[np.float64]]:
+    da: xr.DataArray, level: npt.NDArray[np.floating], log_applied: bool
+) -> tuple[MetDataArray, npt.NDArray[np.floating]]:
     da = da.assign_coords(level=np.log(da["level"]))
 
     if not log_applied:
@@ -994,8 +994,8 @@ def _prepare_q_log_q_log_p(
 
 
 def _prepare_q_cubic_spline(
-    da: xr.DataArray, level: npt.NDArray[np.float64]
-) -> tuple[MetDataArray, npt.NDArray[np.float64]]:
+    da: xr.DataArray, level: npt.NDArray[np.floating]
+) -> tuple[MetDataArray, npt.NDArray[np.floating]]:
     if da["level"][0] < 50.0 or da["level"][-1] > 1000.0:
         msg = "Cubic spline interpolation requires data to span 50-1000 hPa."
         raise ValueError(msg)
diff --git a/pycontrails/core/polygon.py b/pycontrails/core/polygon.py
index 1c721d9d8..53b47993e 100644
--- a/pycontrails/core/polygon.py
+++ b/pycontrails/core/polygon.py
@@ -42,7 +42,7 @@
 
 
 def buffer_and_clean(
-    contour: npt.NDArray[np.float64],
+    contour: npt.NDArray[np.floating],
     min_area: float,
     convex_hull: bool,
     epsilon: float,
@@ -54,7 +54,7 @@ def buffer_and_clean(
 
     Parameters
     ----------
-    contour : npt.NDArray[np.float64]
+    contour : npt.NDArray[np.floating]
         Contour to buffer and clean. A 2d array of shape (n, 2) where n is the number
         of vertices in the contour.
     min_area : float
@@ -157,13 +157,13 @@ def _round_polygon(polygon: shapely.Polygon, precision: int) -> shapely.Polygon:
 
 
 def _contours_to_polygons(
-    contours: Sequence[npt.NDArray[np.float64]],
+    contours: Sequence[npt.NDArray[np.floating]],
     hierarchy: npt.NDArray[np.int_],
     min_area: float,
     convex_hull: bool,
     epsilon: float,
-    longitude: npt.NDArray[np.float64] | None,
-    latitude: npt.NDArray[np.float64] | None,
+    longitude: npt.NDArray[np.floating] | None,
+    latitude: npt.NDArray[np.floating] | None,
     precision: int | None,
     buffer: float,
     i: int = 0,
@@ -172,7 +172,7 @@ def _contours_to_polygons(
 
     Parameters
     ----------
-    contours : Sequence[npt.NDArray[np.float64]]
+    contours : Sequence[npt.NDArray[np.floating]]
         The contours output from :func:`cv2.findContours`.
     hierarchy : npt.NDArray[np.int_]
         The hierarchy output from :func:`cv2.findContours`.
@@ -182,9 +182,9 @@ def _contours_to_polygons(
         Whether to take the convex hull of each polygon.
     epsilon : float
         Epsilon value to use when simplifying the polygons.
-    longitude : npt.NDArray[np.float64] | None
+    longitude : npt.NDArray[np.floating] | None
         Longitude values for the grid.
-    latitude : npt.NDArray[np.float64] | None
+    latitude : npt.NDArray[np.floating] | None
         Latitude values for the grid.
     precision : int | None
         Precision to use when rounding the coordinates.
@@ -254,7 +254,7 @@ def _contours_to_polygons(
 
 
 def determine_buffer(
-    longitude: npt.NDArray[np.float64], latitude: npt.NDArray[np.float64]
+    longitude: npt.NDArray[np.floating], latitude: npt.NDArray[np.floating]
 ) -> float:
     """Determine the proper buffer size to use when converting to polygons."""
 
@@ -279,22 +279,22 @@ def determine_buffer(
 
 
 def find_multipolygon(
-    arr: npt.NDArray[np.float64],
+    arr: npt.NDArray[np.floating],
     threshold: float,
     min_area: float,
     epsilon: float,
     lower_bound: bool = True,
     interiors: bool = True,
     convex_hull: bool = False,
-    longitude: npt.NDArray[np.float64] | None = None,
-    latitude: npt.NDArray[np.float64] | None = None,
+    longitude: npt.NDArray[np.floating] | None = None,
+    latitude: npt.NDArray[np.floating] | None = None,
     precision: int | None = None,
 ) -> shapely.MultiPolygon:
     """Compute a multipolygon from a 2d array.
 
     Parameters
     ----------
-    arr : npt.NDArray[np.float64]
+    arr : npt.NDArray[np.floating]
         Array to convert to a multipolygon. The array will be converted to a binary
         array by comparing each element to ``threshold``. This binary array is then
         passed into :func:`cv2.findContours` to find the contours.
@@ -312,11 +312,11 @@ def find_multipolygon(
         Whether to include interior polygons. By default, True.
     convex_hull : bool, optional
         Experimental. Whether to take the convex hull of each polygon. By default, False.
-    longitude : npt.NDArray[np.float64] | None, optional
+    longitude : npt.NDArray[np.floating] | None, optional
         If provided, the coordinates values corresponding to the longitude dimensions of ``arr``.
         The contour coordinates will be converted to longitude-latitude values by indexing
         into this array. Defaults to None.
-    latitude : npt.NDArray[np.float64] | None, optional
+    latitude : npt.NDArray[np.floating] | None, optional
         If provided, the coordinates values corresponding to the latitude dimensions of ``arr``.
     precision : int | None, optional
         If provided, the precision to use when rounding the coordinates. Defaults to None.
diff --git a/pycontrails/core/vector.py b/pycontrails/core/vector.py
index b6c59bc83..4d973e527 100644
--- a/pycontrails/core/vector.py
+++ b/pycontrails/core/vector.py
@@ -1392,7 +1392,7 @@ def _display_attrs(self) -> dict[str, str]:
         return attrs
 
     @property
-    def level(self) -> npt.NDArray[np.float64]:
+    def level(self) -> npt.NDArray[np.floating]:
         """Get pressure ``level`` values for points.
 
         Automatically calculates pressure level using :func:`units.m_to_pl` using ``altitude`` key.
@@ -1403,7 +1403,7 @@ def level(self) -> npt.NDArray[np.float64]:
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Point pressure level values, [:math:`hPa`]
         """
         try:
@@ -1412,7 +1412,7 @@ def level(self) -> npt.NDArray[np.float64]:
             return units.m_to_pl(self.altitude)
 
     @property
-    def altitude(self) -> npt.NDArray[np.float64]:
+    def altitude(self) -> npt.NDArray[np.floating]:
         """Get altitude.
 
         Automatically calculates altitude using :func:`units.pl_to_m` using ``level`` key.
@@ -1423,7 +1423,7 @@ def altitude(self) -> npt.NDArray[np.float64]:
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Altitude, [:math:`m`]
         """
         try:
@@ -1437,12 +1437,12 @@ def altitude(self) -> npt.NDArray[np.float64]:
             return units.ft_to_m(self["altitude_ft"])
 
     @property
-    def air_pressure(self) -> npt.NDArray[np.float64]:
+    def air_pressure(self) -> npt.NDArray[np.floating]:
         """Get ``air_pressure`` values for points.
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Point air pressure values, [:math:`Pa`]
         """
         try:
@@ -1451,12 +1451,12 @@ def air_pressure(self) -> npt.NDArray[np.float64]:
             return 100.0 * self.level
 
     @property
-    def altitude_ft(self) -> npt.NDArray[np.float64]:
+    def altitude_ft(self) -> npt.NDArray[np.floating]:
         """Get altitude in feet.
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Altitude, [:math:`ft`]
         """
         try:
@@ -1522,7 +1522,7 @@ def coords(self) -> dict[str, np.ndarray]:
     # Utilities
     # ------------
 
-    def transform_crs(self, crs: str) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    def transform_crs(self, crs: str) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
         """Transform trajectory data from one coordinate reference system (CRS) to another.
 
         Parameters
@@ -1535,7 +1535,7 @@ def transform_crs(self, crs: str) -> tuple[npt.NDArray[np.float64], npt.NDArray[
 
         Returns
         -------
-        tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]
+        tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]
             New x and y coordinates in the target CRS.
         """
         try:
@@ -1552,12 +1552,12 @@ def transform_crs(self, crs: str) -> tuple[npt.NDArray[np.float64], npt.NDArray[
         transformer = pyproj.Transformer.from_crs(crs_from, crs, always_xy=True)
         return transformer.transform(self["longitude"], self["latitude"])
 
-    def T_isa(self) -> npt.NDArray[np.float64]:
+    def T_isa(self) -> npt.NDArray[np.floating]:
         """Calculate the ICAO standard atmosphere temperature at each point.
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             ISA temperature, [:math:`K`]
 
         See Also
@@ -1610,24 +1610,24 @@ def intersect_met(
         self,
         mda: met_module.MetDataArray,
         *,
-        longitude: npt.NDArray[np.float64] | None = None,
-        latitude: npt.NDArray[np.float64] | None = None,
-        level: npt.NDArray[np.float64] | None = None,
+        longitude: npt.NDArray[np.floating] | None = None,
+        latitude: npt.NDArray[np.floating] | None = None,
+        level: npt.NDArray[np.floating] | None = None,
         time: npt.NDArray[np.datetime64] | None = None,
         use_indices: bool = False,
         **interp_kwargs: Any,
-    ) -> npt.NDArray[np.float64]:
+    ) -> npt.NDArray[np.floating]:
         """Intersect waypoints with MetDataArray.
 
         Parameters
         ----------
         mda : MetDataArray
             MetDataArray containing a meteorological variable at spatio-temporal coordinates.
-        longitude : npt.NDArray[np.float64], optional
+        longitude : npt.NDArray[np.floating], optional
             Override existing coordinates for met interpolation
-        latitude : npt.NDArray[np.float64], optional
+        latitude : npt.NDArray[np.floating], optional
             Override existing coordinates for met interpolation
-        level : npt.NDArray[np.float64], optional
+        level : npt.NDArray[np.floating], optional
             Override existing coordinates for met interpolation
         time : npt.NDArray[np.datetime64], optional
             Override existing coordinates for met interpolation
@@ -1646,7 +1646,7 @@ def intersect_met(
 
         Returns
         -------
-        npt.NDArray[np.float64]
+        npt.NDArray[np.floating]
             Interpolated values
 
         Examples
diff --git a/pycontrails/datalib/_met_utils/metsource.py b/pycontrails/datalib/_met_utils/metsource.py
index 2245795cd..8d03c540f 100644
--- a/pycontrails/datalib/_met_utils/metsource.py
+++ b/pycontrails/datalib/_met_utils/metsource.py
@@ -175,13 +175,16 @@ def parse_pressure_levels(
 
     out = arr.tolist()
     if supported is None:
-        return out
+        return out  # type: ignore[return-value]
 
-    if missing := set(out).difference(supported):
-        msg = f"Pressure levels {sorted(missing)} are not supported. Supported levels: {supported}"
+    if missing := set(out).difference(supported):  # type: ignore[arg-type]
+        msg = (
+            f"Pressure levels {sorted(missing)} are not supported. "  # type: ignore[type-var]
+            f"Supported levels: {supported}"
+        )
         raise ValueError(msg)
 
-    return out
+    return out  # type: ignore[return-value]
 
 
 def parse_variables(variables: VariableInput, supported: list[MetVariable]) -> list[MetVariable]:
diff --git a/pycontrails/datalib/ecmwf/common.py b/pycontrails/datalib/ecmwf/common.py
index 1f3cb89d4..706b7ed48 100644
--- a/pycontrails/datalib/ecmwf/common.py
+++ b/pycontrails/datalib/ecmwf/common.py
@@ -61,7 +61,7 @@ def _process_dataset(self, ds: xr.Dataset, **kwargs: Any) -> met.MetDataset:
 
         # downselect times
         if not self.timesteps:
-            self.timesteps = ds["time"].values.astype("datetime64[ns]").tolist()
+            self.timesteps = ds["time"].values.astype("datetime64[ns]").tolist()  # type: ignore[assignment]
         else:
             try:
                 ds = ds.sel(time=self.timesteps)
diff --git a/pycontrails/datalib/ecmwf/hres.py b/pycontrails/datalib/ecmwf/hres.py
index 311c3768e..e686c3c5c 100644
--- a/pycontrails/datalib/ecmwf/hres.py
+++ b/pycontrails/datalib/ecmwf/hres.py
@@ -691,7 +691,7 @@ def _open_and_cache(self, xr_kwargs: dict[str, Any]) -> xr.Dataset:
 
         # set forecast time if its not already defined
         if not getattr(self, "forecast_time", None):
-            self.forecast_time = ds["time"].values.astype("datetime64[s]").tolist()
+            self.forecast_time = ds["time"].values.astype("datetime64[s]").tolist()  # type: ignore[assignment]
 
         # check that forecast_time is correct if defined
         # note the "time" coordinate here is the HRES forecast_time
@@ -706,7 +706,7 @@ def _open_and_cache(self, xr_kwargs: dict[str, Any]) -> xr.Dataset:
         # set timesteps if not defined
         # note that "time" is now the actual timestep coordinates
         if not self.timesteps:
-            self.timesteps = ds["time"].values.astype("datetime64[s]").tolist()
+            self.timesteps = ds["time"].values.astype("datetime64[s]").tolist()  # type: ignore[assignment]
 
         self.cache_dataset(ds)
 
diff --git a/pycontrails/datalib/ecmwf/ifs.py b/pycontrails/datalib/ecmwf/ifs.py
index 6704fcd34..eb5db4df9 100644
--- a/pycontrails/datalib/ecmwf/ifs.py
+++ b/pycontrails/datalib/ecmwf/ifs.py
@@ -149,7 +149,7 @@ def open_metdataset(
         else:
             # set timesteps from dataset "time" coordinates
             # np.datetime64 doesn't covert to list[datetime] unless its unit is us
-            self.timesteps = ds["time"].values.astype("datetime64[us]").tolist()
+            self.timesteps = ds["time"].values.astype("datetime64[us]").tolist()  # type: ignore[assignment]
 
         # downselect hyam/hybm coefficients by the "lev" coordinate
         # (this is a 1-indexed verison of nhym)
diff --git a/pycontrails/datalib/gfs/gfs.py b/pycontrails/datalib/gfs/gfs.py
index 5d4d94e1c..60def95d0 100644
--- a/pycontrails/datalib/gfs/gfs.py
+++ b/pycontrails/datalib/gfs/gfs.py
@@ -597,7 +597,7 @@ def _process_dataset(self, ds: xr.Dataset, **kwargs: Any) -> met.MetDataset:
         else:
             # set timesteps from dataset "time" coordinates
             # np.datetime64 doesn't covert to list[datetime] unless its unit is us
-            self.timesteps = ds["time"].values.astype("datetime64[us]").tolist()
+            self.timesteps = ds["time"].values.astype("datetime64[us]").tolist()  # type: ignore[assignment]
 
         # if "level" is not in dims and
         # length of the requested pressure levels is 1
diff --git a/pycontrails/datalib/goes.py b/pycontrails/datalib/goes.py
index c597990d2..edd596da2 100644
--- a/pycontrails/datalib/goes.py
+++ b/pycontrails/datalib/goes.py
@@ -745,8 +745,8 @@ def to_ash(da: xr.DataArray, convention: str = "SEVIRI") -> npt.NDArray[np.float
 
 
 def _clip_and_scale(
-    arr: npt.NDArray[np.float64], low: float, high: float
-) -> npt.NDArray[np.float64]:
+    arr: npt.NDArray[np.floating], low: float, high: float
+) -> npt.NDArray[np.floating]:
     """Clip array and rescale to the interval [0, 1].
 
     Array is first clipped to the interval [low, high] and then linearly rescaled
@@ -757,7 +757,7 @@ def _clip_and_scale(
 
     Parameters
     ----------
-    arr : npt.NDArray[np.float64]
+    arr : npt.NDArray[np.floating]
         Array to clip and scale.
     low : float
         Lower clipping bound.
@@ -766,7 +766,7 @@ def _clip_and_scale(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Clipped and scaled array.
     """
     return (arr.clip(low, high) - low) / (high - low)
diff --git a/pycontrails/ext/empirical_grid.py b/pycontrails/ext/empirical_grid.py
index 13a7dcdaa..68265d70e 100644
--- a/pycontrails/ext/empirical_grid.py
+++ b/pycontrails/ext/empirical_grid.py
@@ -118,7 +118,7 @@ def _query_data(self) -> pd.DataFrame:
 
         return data[["altitude_ft", *columns]].drop(columns=["aircraft_type"])
 
-    def _sample(self, altitude_ft: npt.NDArray[np.float64]) -> None:
+    def _sample(self, altitude_ft: npt.NDArray[np.floating]) -> None:
         """Sample the data and update the source."""
 
         df = self._query_data()
diff --git a/pycontrails/models/cocip/cocip.py b/pycontrails/models/cocip/cocip.py
index 616cf7dee..92f429cde 100644
--- a/pycontrails/models/cocip/cocip.py
+++ b/pycontrails/models/cocip/cocip.py
@@ -2055,9 +2055,9 @@ def calc_radiative_properties(contrail: GeoVectorDataset, params: dict[str, Any]
 
 def calc_contrail_properties(
     contrail: GeoVectorDataset,
-    effective_vertical_resolution: float | npt.NDArray[np.float64],
-    wind_shear_enhancement_exponent: float | npt.NDArray[np.float64],
-    sedimentation_impact_factor: float | npt.NDArray[np.float64],
+    effective_vertical_resolution: float | npt.NDArray[np.floating],
+    wind_shear_enhancement_exponent: float | npt.NDArray[np.floating],
+    sedimentation_impact_factor: float | npt.NDArray[np.floating],
     radiative_heating_effects: bool,
 ) -> None:
     """Calculate geometric and ice-related properties of contrail.
@@ -2084,11 +2084,11 @@ def calc_contrail_properties(
     ----------
     contrail : GeoVectorDataset
         Grid points with many precomputed keys.
-    effective_vertical_resolution : float | npt.NDArray[np.float64]
+    effective_vertical_resolution : float | npt.NDArray[np.floating]
         Passed into :func:`wind_shear.wind_shear_enhancement_factor`.
-    wind_shear_enhancement_exponent : float | npt.NDArray[np.float64]
+    wind_shear_enhancement_exponent : float | npt.NDArray[np.floating]
         Passed into :func:`wind_shear.wind_shear_enhancement_factor`.
-    sedimentation_impact_factor: float | npt.NDArray[np.float64]
+    sedimentation_impact_factor: float | npt.NDArray[np.floating]
         Passed into `contrail_properties.vertical_diffusivity`.
     radiative_heating_effects: bool
         Include radiative heating effects on contrail cirrus properties.
@@ -2525,8 +2525,8 @@ def calc_timestep_contrail_evolution(
 def _rad_accumulation_to_average_instantaneous(
     rad: MetDataset,
     name: str,
-    arr: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    arr: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """Convert from radiation accumulation to average instantaneous values.
 
     .. versionadded:: 0.48.0
@@ -2537,12 +2537,12 @@ def _rad_accumulation_to_average_instantaneous(
         Radiation data
     name : str
         Variable name
-    arr : npt.NDArray[np.float64]
+    arr : npt.NDArray[np.floating]
         Array of values already interpolated from ``rad``
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Array of values converted from accumulation to average instantaneous values
 
     Raises
diff --git a/pycontrails/models/cocip/cocip_uncertainty.py b/pycontrails/models/cocip/cocip_uncertainty.py
index 5f230e57c..b4ddda651 100644
--- a/pycontrails/models/cocip/cocip_uncertainty.py
+++ b/pycontrails/models/cocip/cocip_uncertainty.py
@@ -212,7 +212,7 @@ def uncertainty_params(self) -> dict[str, rv_frozen]:
 
         return out
 
-    def rvs(self, size: None | int = None) -> dict[str, np.float64 | npt.NDArray[np.float64]]:
+    def rvs(self, size: None | int = None) -> dict[str, np.float64 | npt.NDArray[np.floating]]:
         """Call each distribution's `rvs` method to generate random parameters.
 
         Seed calls to `rvs` with class variable `rng`.
@@ -224,7 +224,7 @@ def rvs(self, size: None | int = None) -> dict[str, np.float64 | npt.NDArray[np.
 
         Returns
         -------
-        dict[str, float | npt.NDArray[np.float64]]
+        dict[str, float | npt.NDArray[np.floating]]
             Dictionary of random parameters. Dictionary keys consists of names of parameters in
             `CocipParams` to be overridden by random value.
 
diff --git a/pycontrails/models/cocip/contrail_properties.py b/pycontrails/models/cocip/contrail_properties.py
index 9da90969b..12810374a 100644
--- a/pycontrails/models/cocip/contrail_properties.py
+++ b/pycontrails/models/cocip/contrail_properties.py
@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, TypeVar
+from typing import Any, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -19,14 +19,14 @@
 
 
 def initial_iwc(
-    air_temperature: npt.NDArray[np.float64],
-    specific_humidity: npt.NDArray[np.float64],
-    air_pressure: npt.NDArray[np.float64],
-    fuel_dist: npt.NDArray[np.float64],
-    width: npt.NDArray[np.float64],
-    depth: npt.NDArray[np.float64],
+    air_temperature: npt.NDArray[np.floating],
+    specific_humidity: npt.NDArray[np.floating],
+    air_pressure: npt.NDArray[np.floating],
+    fuel_dist: npt.NDArray[np.floating],
+    width: npt.NDArray[np.floating],
+    depth: npt.NDArray[np.floating],
     ei_h2o: float,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     r"""
     Estimate the initial contrail ice water content (iwc) before the wake vortex phase.
 
@@ -35,24 +35,24 @@ def initial_iwc(
 
     Parameters
     ----------
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
-    specific_humidity : npt.NDArray[np.float64]
+    specific_humidity : npt.NDArray[np.floating]
         ambient specific humidity for each waypoint, [:math:`kg_{H_{2}O}/kg_{air}`]
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         initial pressure altitude at each waypoint, before the wake vortex phase, [:math:`Pa`]
-    fuel_dist : npt.NDArray[np.float64]
+    fuel_dist : npt.NDArray[np.floating]
         fuel consumption of the flight segment per distance traveled, [:math:`kg m^{-1}`]
-    width : npt.NDArray[np.float64]
+    width : npt.NDArray[np.floating]
         initial contrail width, [:math:`m`]
-    depth : npt.NDArray[np.float64]
+    depth : npt.NDArray[np.floating]
         initial contrail depth, [:math:`m`]
     ei_h2o : float
         water vapor emissions index of fuel, [:math:`kg_{H_{2}O} \ kg_{fuel}^{-1}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Initial contrail ice water content (iwc) at the original waypoint
         before the wake vortex phase, [:math:`kg_{H_{2}O}/kg_{air}`].
         Returns zero if iwc is is negative (dry air).
@@ -63,34 +63,34 @@ def initial_iwc(
 
 
 def q_exhaust(
-    air_temperature: npt.NDArray[np.float64],
-    air_pressure: npt.NDArray[np.float64],
-    fuel_dist: npt.NDArray[np.float64],
-    width: npt.NDArray[np.float64],
-    depth: npt.NDArray[np.float64],
+    air_temperature: npt.NDArray[np.floating],
+    air_pressure: npt.NDArray[np.floating],
+    fuel_dist: npt.NDArray[np.floating],
+    width: npt.NDArray[np.floating],
+    depth: npt.NDArray[np.floating],
     ei_h2o: float,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     r"""
     Calculate the specific humidity released by water vapor from aircraft emissions.
 
     Parameters
     ----------
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         initial pressure altitude at each waypoint, before the wake vortex phase, [:math:`Pa`]
-    fuel_dist : npt.NDArray[np.float64]
+    fuel_dist : npt.NDArray[np.floating]
         fuel consumption of the flight segment per distance travelled, [:math:`kg m^{-1}`]
-    width : npt.NDArray[np.float64]
+    width : npt.NDArray[np.floating]
         initial contrail width, [:math:`m`]
-    depth : npt.NDArray[np.float64]
+    depth : npt.NDArray[np.floating]
         initial contrail depth, [:math:`m`]
     ei_h2o : float
         water vapor emissions index of fuel, [:math:`kg_{H_{2}O} \ kg_{fuel}^{-1}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Humidity released by water vapour from aircraft emissions, [:math:`kg_{H_{2}O}/kg_{air}`]
     """
     rho_air = thermo.rho_d(air_temperature, air_pressure)
@@ -98,25 +98,25 @@ def q_exhaust(
 
 
 def iwc_adiabatic_heating(
-    air_temperature: npt.NDArray[np.float64],
-    air_pressure: npt.NDArray[np.float64],
-    air_pressure_1: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    air_temperature: npt.NDArray[np.floating],
+    air_pressure: npt.NDArray[np.floating],
+    air_pressure_1: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Calculate the change in ice water content due to adiabatic heating from the wake vortex phase.
 
     Parameters
     ----------
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         initial pressure altitude at each waypoint, before the wake vortex phase, [:math:`Pa`]
-    air_pressure_1 : npt.NDArray[np.float64]
+    air_pressure_1 : npt.NDArray[np.floating]
         pressure altitude at each waypoint, after the wake vortex phase, [:math:`Pa`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Change in ice water content due to adiabatic heating from the wake
         vortex phase, [:math:`kg_{H_{2}O}/kg_{air}`]
     """
@@ -130,26 +130,26 @@ def iwc_adiabatic_heating(
 
 
 def temperature_adiabatic_heating(
-    air_temperature: npt.NDArray[np.float64],
-    air_pressure: npt.NDArray[np.float64],
-    air_pressure_1: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    air_temperature: npt.NDArray[np.floating],
+    air_pressure: npt.NDArray[np.floating],
+    air_pressure_1: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """Calculate the ambient air temperature for each waypoint after the wake vortex phase.
 
     This calculation accounts for adiabatic heating.
 
     Parameters
     ----------
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         initial pressure altitude at each waypoint, before the wake vortex phase, [:math:`Pa`]
-    air_pressure_1 : npt.NDArray[np.float64]
+    air_pressure_1 : npt.NDArray[np.floating]
         pressure altitude at each waypoint, after the wake vortex phase, [:math:`Pa`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         ambient air temperature after the wake vortex phase, [:math:`K`]
 
     Notes
@@ -165,8 +165,8 @@ def temperature_adiabatic_heating(
 
 
 def iwc_post_wake_vortex(
-    iwc: npt.NDArray[np.float64], iwc_ad: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    iwc: npt.NDArray[np.floating], iwc_ad: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate the ice water content after the wake vortex phase (``iwc_1``).
 
@@ -178,16 +178,16 @@ def iwc_post_wake_vortex(
 
     Parameters
     ----------
-    iwc : npt.NDArray[np.float64]
+    iwc : npt.NDArray[np.floating]
         initial ice water content at each waypoint before the wake vortex
         phase, [:math:`kg_{H_{2}O}/kg_{air}`]
-    iwc_ad : npt.NDArray[np.float64]
+    iwc_ad : npt.NDArray[np.floating]
         change in iwc from adiabatic heating during the wake vortex
         phase, [:math:`kg_{H_{2}O}/kg_{air}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         ice water content after the wake vortex phase, ``iwc_1``, [:math:`kg_{H_{2}O}/kg_{air}`]
 
     Notes
@@ -202,13 +202,13 @@ def iwc_post_wake_vortex(
 
 
 def ice_particle_number(
-    nvpm_ei_n: npt.NDArray[np.float64],
-    fuel_dist: npt.NDArray[np.float64],
-    f_surv: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-    T_crit_sac: npt.NDArray[np.float64],
+    nvpm_ei_n: npt.NDArray[np.floating],
+    fuel_dist: npt.NDArray[np.floating],
+    f_surv: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+    T_crit_sac: npt.NDArray[np.floating],
     min_ice_particle_number_nvpm_ei_n: float,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Calculate the initial number of ice particles per distance after the wake vortex phase.
 
     The initial number of ice particle per distance is calculated from the black
@@ -218,15 +218,15 @@ def ice_particle_number(
 
     Parameters
     ----------
-    nvpm_ei_n : npt.NDArray[np.float64]
+    nvpm_ei_n : npt.NDArray[np.floating]
         black carbon number emissions index, [:math:`kg^{-1}`]
-    fuel_dist : npt.NDArray[np.float64]
+    fuel_dist : npt.NDArray[np.floating]
         fuel consumption of the flight segment per distance traveled, [:math:`kg m^{-1}`]
-    f_surv : npt.NDArray[np.float64]
+    f_surv : npt.NDArray[np.floating]
         Fraction of contrail ice particle number that survive the wake vortex phase.
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
-    T_crit_sac : npt.NDArray[np.float64]
+    T_crit_sac : npt.NDArray[np.floating]
         estimated Schmidt-Appleman temperature threshold for contrail formation, [:math:`K`]
     min_ice_particle_number_nvpm_ei_n : float
         lower bound for nvpm_ei_n to account for ambient aerosol particles for
@@ -234,7 +234,7 @@ def ice_particle_number(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         initial number of ice particles per distance after the wake vortex phase, [:math:`# m^{-1}`]
     """
     f_activation = ice_particle_activation_rate(air_temperature, T_crit_sac)
@@ -243,8 +243,8 @@ def ice_particle_number(
 
 
 def ice_particle_activation_rate(
-    air_temperature: npt.NDArray[np.float64], T_crit_sac: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    air_temperature: npt.NDArray[np.floating], T_crit_sac: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate the activation rate of black carbon particles to contrail ice crystals.
 
@@ -253,14 +253,14 @@ def ice_particle_activation_rate(
 
     Parameters
     ----------
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature at each waypoint before wake_wortex, [:math:`K`]
-    T_crit_sac : npt.NDArray[np.float64]
+    T_crit_sac : npt.NDArray[np.floating]
         estimated Schmidt-Appleman temperature threshold for contrail formation, [:math:`K`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Proportion of black carbon particles that activates to contrail ice parties.
 
     Notes
@@ -285,8 +285,8 @@ def ice_particle_activation_rate(
 
 
 def ice_particle_survival_fraction(
-    iwc: npt.NDArray[np.float64], iwc_1: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    iwc: npt.NDArray[np.floating], iwc_1: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Estimate the fraction of contrail ice particle number that survive the wake vortex phase.
 
@@ -295,15 +295,15 @@ def ice_particle_survival_fraction(
 
     Parameters
     ----------
-    iwc : npt.NDArray[np.float64]
+    iwc : npt.NDArray[np.floating]
         initial ice water content at each waypoint before the wake vortex
         phase, [:math:`kg_{H_{2}O}/kg_{air}`]
-    iwc_1 : npt.NDArray[np.float64]
+    iwc_1 : npt.NDArray[np.floating]
         ice water content after the wake vortex phase, [:math:`kg_{H_{2}O}/kg_{air}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Fraction of contrail ice particle number that survive the wake vortex phase.
     """
     f_surv = np.empty_like(iwc)
@@ -319,8 +319,8 @@ def ice_particle_survival_fraction(
 
 
 def initial_persistent(
-    iwc_1: npt.NDArray[np.float64], rhi_1: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    iwc_1: npt.NDArray[np.floating], rhi_1: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Determine if waypoints have persistent contrails.
 
@@ -335,14 +335,14 @@ def initial_persistent(
 
     Parameters
     ----------
-    iwc_1 : npt.NDArray[np.float64]
+    iwc_1 : npt.NDArray[np.floating]
         ice water content after the wake vortex phase, [:math:`kg_{H_{2}O}/kg_{air}`]
-    rhi_1 : npt.NDArray[np.float64]
+    rhi_1 : npt.NDArray[np.floating]
         relative humidity with respect to ice after the wake vortex phase
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Mask of waypoints with persistent contrails. Waypoints with persistent contrails
         will have value 1.
 
@@ -358,12 +358,12 @@ def initial_persistent(
 
 
 def contrail_persistent(
-    latitude: npt.NDArray[np.float64],
-    altitude: npt.NDArray[np.float64],
-    segment_length: npt.NDArray[np.float64],
+    latitude: npt.NDArray[np.floating],
+    altitude: npt.NDArray[np.floating],
+    segment_length: npt.NDArray[np.floating],
     age: npt.NDArray[np.timedelta64],
-    tau_contrail: npt.NDArray[np.float64],
-    n_ice_per_m3: npt.NDArray[np.float64],
+    tau_contrail: npt.NDArray[np.floating],
+    n_ice_per_m3: npt.NDArray[np.floating],
     params: dict[str, Any],
 ) -> npt.NDArray[np.bool_]:
     r"""
@@ -391,17 +391,17 @@ def contrail_persistent(
 
     Parameters
     ----------
-    latitude : npt.NDArray[np.float64]
+    latitude : npt.NDArray[np.floating]
         Contrail latitude, [:math:`\deg`]
-    altitude : npt.NDArray[np.float64]
+    altitude : npt.NDArray[np.floating]
         Contrail altitude, [:math:`m`]
-    segment_length : npt.NDArray[np.float64]
+    segment_length : npt.NDArray[np.floating]
         Contrail segment length, [:math:`m`]
     age : npt.NDArray[np.timedelta64]
         Contrail age
-    tau_contrail : npt.NDArray[np.float64]
+    tau_contrail : npt.NDArray[np.floating]
         Contrail optical depth
-    n_ice_per_m3 : npt.NDArray[np.float64]
+    n_ice_per_m3 : npt.NDArray[np.floating]
         Contrail ice particle number per volume of air, [:math:`# m^{-3}`]
     params : dict[str, Any]
         Dictionary of :class:`CocipParams` parameters determining the
@@ -420,7 +420,7 @@ def contrail_persistent(
     )
     status_4 = _within_range(altitude, max=params["max_altitude_m"], min=params["min_altitude_m"])
     status_5 = _within_range(segment_length, max=params["max_seg_length_m"])
-    status_6 = _within_range(latitude, max=89.0, min=-89.0)  # type: ignore[type-var]
+    status_6 = _within_range(latitude, max=89.0, min=-89.0)
 
     logger.debug(
         "Survival stats. age: %s, tau: %s, ice: %s, altitude: %s, segment: %s, latitude: %s",
@@ -434,13 +434,26 @@ def contrail_persistent(
     return status_1 & status_2 & status_3 & status_4 & status_5 & status_6
 
 
-T = TypeVar("T", np.float64, np.timedelta64)
+@overload
+def _within_range(
+    val: npt.NDArray[np.floating],
+    max: float | None = None,
+    min: float | None = None,
+) -> npt.NDArray[np.bool_]: ...
+
+
+@overload
+def _within_range(
+    val: npt.NDArray[np.timedelta64],
+    max: np.timedelta64 | None = None,
+    min: np.timedelta64 | None = None,
+) -> npt.NDArray[np.bool_]: ...
 
 
 def _within_range(
-    val: npt.NDArray[T],
-    max: npt.NDArray[T] | T | None = None,
-    min: npt.NDArray[T] | T | None = None,
+    val: np.ndarray,
+    max: float | np.timedelta64 | None = None,
+    min: float | np.timedelta64 | None = None,
 ) -> npt.NDArray[np.bool_]:
     """
     Check if the input values (val) are each within the specified range.
@@ -475,16 +488,16 @@ def _within_range(
 
 
 def contrail_edges(
-    lon: npt.NDArray[np.float64],
-    lat: npt.NDArray[np.float64],
-    sin_a: npt.NDArray[np.float64],
-    cos_a: npt.NDArray[np.float64],
-    width: npt.NDArray[np.float64],
+    lon: npt.NDArray[np.floating],
+    lat: npt.NDArray[np.floating],
+    sin_a: npt.NDArray[np.floating],
+    cos_a: npt.NDArray[np.floating],
+    width: npt.NDArray[np.floating],
 ) -> tuple[
-    npt.NDArray[np.float64],
-    npt.NDArray[np.float64],
-    npt.NDArray[np.float64],
-    npt.NDArray[np.float64],
+    npt.NDArray[np.floating],
+    npt.NDArray[np.floating],
+    npt.NDArray[np.floating],
+    npt.NDArray[np.floating],
 ]:
     """
     Calculate the longitude and latitude of the contrail edges to account for contrail spreading.
@@ -498,20 +511,20 @@ def contrail_edges(
 
     Parameters
     ----------
-    lon : npt.NDArray[np.float64]
+    lon : npt.NDArray[np.floating]
         longitude of contrail waypoint, degrees
-    lat : npt.NDArray[np.float64]
+    lat : npt.NDArray[np.floating]
         latitude of contrail waypoint, degrees
-    sin_a : npt.NDArray[np.float64]
+    sin_a : npt.NDArray[np.floating]
         sin(a), where a is the angle between the plume and the longitudinal axis
-    cos_a : npt.NDArray[np.float64]
+    cos_a : npt.NDArray[np.floating]
         cos(a), where a is the angle between the plume and the longitudinal axis
-    width : npt.NDArray[np.float64]
+    width : npt.NDArray[np.floating]
         contrail width at each waypoint, [:math:`m`]
 
     Returns
     -------
-    tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]]
+    tuple[npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.floating]]
         (lon_edge_l, lat_edge_l, lon_edge_r, lat_edge_r), longitudes and latitudes
         at the left and right edges of the contrail, degrees
     """  # noqa: E501
@@ -527,21 +540,21 @@ def contrail_edges(
 
 
 def contrail_vertices(
-    lon: npt.NDArray[np.float64],
-    lat: npt.NDArray[np.float64],
-    sin_a: npt.NDArray[np.float64],
-    cos_a: npt.NDArray[np.float64],
-    width: npt.NDArray[np.float64],
-    segment_length: npt.NDArray[np.float64],
+    lon: npt.NDArray[np.floating],
+    lat: npt.NDArray[np.floating],
+    sin_a: npt.NDArray[np.floating],
+    cos_a: npt.NDArray[np.floating],
+    width: npt.NDArray[np.floating],
+    segment_length: npt.NDArray[np.floating],
 ) -> tuple[
-    npt.NDArray[np.float64],
-    npt.NDArray[np.float64],
-    npt.NDArray[np.float64],
-    npt.NDArray[np.float64],
-    npt.NDArray[np.float64],
-    npt.NDArray[np.float64],
-    npt.NDArray[np.float64],
-    npt.NDArray[np.float64],
+    npt.NDArray[np.floating],
+    npt.NDArray[np.floating],
+    npt.NDArray[np.floating],
+    npt.NDArray[np.floating],
+    npt.NDArray[np.floating],
+    npt.NDArray[np.floating],
+    npt.NDArray[np.floating],
+    npt.NDArray[np.floating],
 ]:
     """
     Calculate the longitude and latitude of the contrail vertices.
@@ -559,22 +572,22 @@ def contrail_vertices(
 
     Parameters
     ----------
-    lon : npt.NDArray[np.float64]
+    lon : npt.NDArray[np.floating]
         longitude of contrail waypoint, degrees
-    lat : npt.NDArray[np.float64]
+    lat : npt.NDArray[np.floating]
         latitude of contrail waypoint, degrees
-    sin_a : npt.NDArray[np.float64]
+    sin_a : npt.NDArray[np.floating]
         sin(a), where a is the angle between the plume and the longitudinal axis
-    cos_a : npt.NDArray[np.float64]
+    cos_a : npt.NDArray[np.floating]
         cos(a), where a is the angle between the plume and the longitudinal axis
-    width : npt.NDArray[np.float64]
+    width : npt.NDArray[np.floating]
         contrail width at each waypoint, [:math:`m`]
-    segment_length : npt.NDArray[np.float64]
+    segment_length : npt.NDArray[np.floating]
         contrail length at each waypoint, [:math:`m`]
 
     Returns
     -------
-    tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]]
+    tuple[npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.floating]]
         (lon_1, lat_1, lon_2, lat_2, lon_3, lat_3, lon_4, lat_4) degrees
     """  # noqa: E501
     dlon_width = units.m_to_longitude_distance(width * sin_a * 0.5, lat)
@@ -598,10 +611,10 @@ def contrail_vertices(
 
 
 def plume_effective_cross_sectional_area(
-    width: npt.NDArray[np.float64],
-    depth: npt.NDArray[np.float64],
-    sigma_yz: npt.NDArray[np.float64] | float,
-) -> npt.NDArray[np.float64]:
+    width: npt.NDArray[np.floating],
+    depth: npt.NDArray[np.floating],
+    sigma_yz: npt.NDArray[np.floating] | float,
+) -> npt.NDArray[np.floating]:
     """
     Calculate the effective cross-sectional area of the contrail plume (``area_eff``).
 
@@ -610,16 +623,16 @@ def plume_effective_cross_sectional_area(
 
     Parameters
     ----------
-    width : npt.NDArray[np.float64]
+    width : npt.NDArray[np.floating]
         contrail width at each waypoint, [:math:`m`]
-    depth : npt.NDArray[np.float64]
+    depth : npt.NDArray[np.floating]
         contrail depth at each waypoint, [:math:`m`]
-    sigma_yz : npt.NDArray[np.float64] | float
+    sigma_yz : npt.NDArray[np.floating] | float
         temporal evolution of the contrail plume parameters
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         effective cross-sectional area of the contrail plume, [:math:`m^{2}`]
     """
     sigma_yy = 0.125 * (width**2)
@@ -628,8 +641,8 @@ def plume_effective_cross_sectional_area(
 
 
 def plume_effective_depth(
-    width: npt.NDArray[np.float64], area_eff: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    width: npt.NDArray[np.floating], area_eff: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate the effective depth of the contrail plume (``depth_eff``).
 
@@ -638,100 +651,100 @@ def plume_effective_depth(
 
     Parameters
     ----------
-    width : npt.NDArray[np.float64]
+    width : npt.NDArray[np.floating]
         contrail width at each waypoint, [:math:`m`]
-    area_eff : npt.NDArray[np.float64]
+    area_eff : npt.NDArray[np.floating]
         effective cross-sectional area of the contrail plume, [:math:`m^{2}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         effective depth of the contrail plume, [:math:`m`]
     """
     return area_eff / width
 
 
 def plume_mass_per_distance(
-    area_eff: npt.NDArray[np.float64], rho_air: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    area_eff: npt.NDArray[np.floating], rho_air: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate the contrail plume mass per unit length.
 
     Parameters
     ----------
-    area_eff : npt.NDArray[np.float64]
+    area_eff : npt.NDArray[np.floating]
         effective cross-sectional area of the contrail plume, [:math:`m^{2}`]
-    rho_air : npt.NDArray[np.float64]
+    rho_air : npt.NDArray[np.floating]
         density of air for each waypoint, [:math:`kg m^{-3}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         contrail plume mass per unit length, [:math:`kg m^{-1}`]
     """
     return area_eff * rho_air
 
 
 def ice_particle_number_per_volume_of_plume(
-    n_ice_per_m: npt.NDArray[np.float64], area_eff: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    n_ice_per_m: npt.NDArray[np.floating], area_eff: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate the number of contrail ice particles per volume of plume (``n_ice_per_vol``).
 
     Parameters
     ----------
-    n_ice_per_m : npt.NDArray[np.float64]
+    n_ice_per_m : npt.NDArray[np.floating]
         number of ice particles per distance at time t, [:math:`m^{-1}`]
-    area_eff : npt.NDArray[np.float64]
+    area_eff : npt.NDArray[np.floating]
         effective cross-sectional area of the contrail plume, [:math:`m^{2}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         number of ice particles per volume of contrail plume at time t, [:math:`# m^{-3}`]
     """
     return n_ice_per_m / area_eff
 
 
 def ice_particle_number_per_mass_of_air(
-    n_ice_per_vol: npt.NDArray[np.float64], rho_air: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    n_ice_per_vol: npt.NDArray[np.floating], rho_air: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate the number of contrail ice particles per mass of air.
 
     Parameters
     ----------
-    n_ice_per_vol : npt.NDArray[np.float64]
+    n_ice_per_vol : npt.NDArray[np.floating]
         number of ice particles per volume of contrail plume at time t, [:math:`# m^{-3}`]
-    rho_air : npt.NDArray[np.float64]
+    rho_air : npt.NDArray[np.floating]
         density of air for each waypoint, [:math:`kg m^{-3}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         number of ice particles per mass of air at time t, [:math:`# kg^{-1}`]
     """
     return n_ice_per_vol / rho_air
 
 
 def ice_particle_volume_mean_radius(
-    iwc: npt.NDArray[np.float64], n_ice_per_kg_air: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    iwc: npt.NDArray[np.floating], n_ice_per_kg_air: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate the ice particle volume mean radius.
 
     Parameters
     ----------
-    iwc : npt.NDArray[np.float64]
+    iwc : npt.NDArray[np.floating]
         contrail ice water content, i.e., contrail ice mass per
         kg of air, [:math:`kg_{H_{2}O}/kg_{air}`]
 
-    n_ice_per_kg_air : npt.NDArray[np.float64]
+    n_ice_per_kg_air : npt.NDArray[np.floating]
         number of ice particles per mass of air, [:math:`# kg^{-1}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         ice particle volume mean radius, [:math:`m`]
 
     Notes
@@ -752,10 +765,10 @@ def ice_particle_volume_mean_radius(
 
 
 def ice_particle_terminal_fall_speed(
-    air_pressure: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-    r_ice_vol: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    air_pressure: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+    r_ice_vol: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Calculate the terminal fall speed of contrail ice particles.
 
@@ -766,16 +779,16 @@ def ice_particle_terminal_fall_speed(
 
     Parameters
     ----------
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         Pressure altitude at each waypoint, [:math:`Pa`]
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         Ambient temperature for each waypoint, [:math:`K`]
-    r_ice_vol : npt.NDArray[np.float64]
+    r_ice_vol : npt.NDArray[np.floating]
         Ice particle volume mean radius, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Terminal fall speed of contrail ice particles, [:math:`m s^{-1}`]
 
     References
@@ -805,7 +818,7 @@ def ice_particle_terminal_fall_speed(
     return alpha * (30000.0 / air_pressure) ** 0.178 * (233.0 / air_temperature) ** 0.394
 
 
-def ice_particle_mass(r_ice_vol: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def ice_particle_mass(r_ice_vol: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     """
     Calculate the contrail ice particle mass.
 
@@ -813,34 +826,34 @@ def ice_particle_mass(r_ice_vol: npt.NDArray[np.float64]) -> npt.NDArray[np.floa
 
     Parameters
     ----------
-    r_ice_vol : npt.NDArray[np.float64]
+    r_ice_vol : npt.NDArray[np.floating]
         Ice particle volume mean radius, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Mean contrail ice particle mass, [:math:`kg`]
     """
     return ((4 / 3) * np.pi * r_ice_vol**3) * constants.rho_ice
 
 
 def horizontal_diffusivity(
-    ds_dz: npt.NDArray[np.float64], depth: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    ds_dz: npt.NDArray[np.floating], depth: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate contrail horizontal diffusivity.
 
     Parameters
     ----------
-    ds_dz : npt.NDArray[np.float64]
+    ds_dz : npt.NDArray[np.floating]
         Total wind shear (eastward and northward winds) with respect
         to altitude (``dz``), [:math:`m s^{-1} / Pa`]
-    depth : npt.NDArray[np.float64]
+    depth : npt.NDArray[np.floating]
         Contrail depth at each waypoint, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         horizontal diffusivity, [:math:`m^{2} s^{-1}`]
 
     References
@@ -856,39 +869,39 @@ def horizontal_diffusivity(
 
 
 def vertical_diffusivity(
-    air_pressure: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-    dT_dz: npt.NDArray[np.float64],
-    depth_eff: npt.NDArray[np.float64],
-    terminal_fall_speed: npt.NDArray[np.float64] | float,
-    sedimentation_impact_factor: npt.NDArray[np.float64] | float,
-    eff_heat_rate: npt.NDArray[np.float64] | None,
-) -> npt.NDArray[np.float64]:
+    air_pressure: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+    dT_dz: npt.NDArray[np.floating],
+    depth_eff: npt.NDArray[np.floating],
+    terminal_fall_speed: npt.NDArray[np.floating] | float,
+    sedimentation_impact_factor: npt.NDArray[np.floating] | float,
+    eff_heat_rate: npt.NDArray[np.floating] | None,
+) -> npt.NDArray[np.floating]:
     """
     Calculate contrail vertical diffusivity.
 
     Parameters
     ----------
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         Pressure altitude at each waypoint, [:math:`Pa`]
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         Ambient temperature for each waypoint, [:math:`K`]
-    dT_dz : npt.NDArray[np.float64]
+    dT_dz : npt.NDArray[np.floating]
         Temperature gradient with respect to altitude (dz), [:math:`K m^{-1}`]
-    depth_eff : npt.NDArray[np.float64]
+    depth_eff : npt.NDArray[np.floating]
         Effective depth of the contrail plume, [:math:`m`]
-    terminal_fall_speed : npt.NDArray[np.float64]
+    terminal_fall_speed : npt.NDArray[np.floating]
         Terminal fall speed of contrail ice particles, [:math:`m s^{-1}`]
     sedimentation_impact_factor : float
         Enhancement parameter denoted by `f_T` in eq. (35) Schumann (2012).
-    eff_heat_rate: npt.NDArray[np.float64] | None
+    eff_heat_rate: npt.NDArray[np.floating] | None
         Effective heating rate, i.e., rate of which the contrail plume
         is heated, [:math:`K s^{-1}`]. If None is passed, the radiative
         heating effects on contrail cirrus properties are not included.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         vertical diffusivity, [:math:`m^{2} s^{-1}`]
 
     References
@@ -913,7 +926,7 @@ def vertical_diffusivity(
     n_bv = thermo.brunt_vaisala_frequency(air_pressure, air_temperature, dT_dz)
     n_bv.clip(min=0.001, out=n_bv)
 
-    cvs: npt.NDArray[np.float64] | float
+    cvs: npt.NDArray[np.floating] | float
     if eff_heat_rate is not None:
         cvs = radiative_heating.convective_velocity_scale(depth_eff, eff_heat_rate, air_temperature)
         cvs.clip(min=0.01, out=cvs)
@@ -929,28 +942,28 @@ def vertical_diffusivity(
 
 
 def particle_losses_aggregation(
-    r_ice_vol: npt.NDArray[np.float64],
-    terminal_fall_speed: npt.NDArray[np.float64],
-    area_eff: npt.NDArray[np.float64],
+    r_ice_vol: npt.NDArray[np.floating],
+    terminal_fall_speed: npt.NDArray[np.floating],
+    area_eff: npt.NDArray[np.floating],
     agg_efficiency: float = 1.0,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """
     Calculate the rate of contrail ice particle losses due to sedimentation-induced aggregation.
 
     Parameters
     ----------
-    r_ice_vol : npt.NDArray[np.float64]
+    r_ice_vol : npt.NDArray[np.floating]
         Ice particle volume mean radius, [:math:`m`]
-    terminal_fall_speed : npt.NDArray[np.float64]
+    terminal_fall_speed : npt.NDArray[np.floating]
         Terminal fall speed of contrail ice particles, [:math:`m s^{-1}`]
-    area_eff : npt.NDArray[np.float64]
+    area_eff : npt.NDArray[np.floating]
         Effective cross-sectional area of the contrail plume, [:math:`m^{2}`]
     agg_efficiency : float, optional
         Aggregation efficiency
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Rate of contrail ice particle losses due to sedimentation-induced
         aggregation, [:math:`# s^{-1}`]
 
@@ -968,34 +981,34 @@ def particle_losses_aggregation(
 
 
 def particle_losses_turbulence(
-    width: npt.NDArray[np.float64],
-    depth: npt.NDArray[np.float64],
-    depth_eff: npt.NDArray[np.float64],
-    diffuse_h: npt.NDArray[np.float64],
-    diffuse_v: npt.NDArray[np.float64],
+    width: npt.NDArray[np.floating],
+    depth: npt.NDArray[np.floating],
+    depth_eff: npt.NDArray[np.floating],
+    diffuse_h: npt.NDArray[np.floating],
+    diffuse_v: npt.NDArray[np.floating],
     turb_efficiency: float = 0.1,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """
     Calculate the rate of contrail ice particle losses due to plume-internal turbulence.
 
     Parameters
     ----------
-    width : npt.NDArray[np.float64]
+    width : npt.NDArray[np.floating]
         Contrail width at each waypoint, [:math:`m`]
-    depth : npt.NDArray[np.float64]
+    depth : npt.NDArray[np.floating]
         Contrail depth at each waypoint, [:math:`m`]
-    depth_eff : npt.NDArray[np.float64]
+    depth_eff : npt.NDArray[np.floating]
         Effective depth of the contrail plume, [:math:`m`]
-    diffuse_h : npt.NDArray[np.float64]
+    diffuse_h : npt.NDArray[np.floating]
         Horizontal diffusivity, [:math:`m^{2} s^{-1}`]
-    diffuse_v : npt.NDArray[np.float64]
+    diffuse_v : npt.NDArray[np.floating]
         Vertical diffusivity, [:math:`m^{2} s^{-1}`]
     turb_efficiency : float, optional
         Turbulence sublimation efficiency
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Rate of contrail ice particle losses due to plume-internal turbulence, [:math:`# s^{-1}`]
 
     Notes
@@ -1018,25 +1031,25 @@ def particle_losses_turbulence(
 
 
 def contrail_optical_depth(
-    r_ice_vol: npt.NDArray[np.float64],
-    n_ice_per_m: npt.NDArray[np.float64],
-    width: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    r_ice_vol: npt.NDArray[np.floating],
+    n_ice_per_m: npt.NDArray[np.floating],
+    width: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Calculate the contrail optical depth for each waypoint.
 
     Parameters
     ----------
-    r_ice_vol : npt.NDArray[np.float64]
+    r_ice_vol : npt.NDArray[np.floating]
         ice particle volume mean radius, [:math:`m`]
-    n_ice_per_m : npt.NDArray[np.float64]
+    n_ice_per_m : npt.NDArray[np.floating]
         Number of contrail ice particles per distance, [:math:`m^{-1}`]
-    width : npt.NDArray[np.float64]
+    width : npt.NDArray[np.floating]
         Contrail width, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Contrail optical depth
     """
     q_ext = scattering_extinction_efficiency(r_ice_vol)
@@ -1048,18 +1061,20 @@ def contrail_optical_depth(
     return tau_contrail
 
 
-def scattering_extinction_efficiency(r_ice_vol: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def scattering_extinction_efficiency(
+    r_ice_vol: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Calculate the scattering extinction efficiency (``q_ext``) based on Mie-theory.
 
     Parameters
     ----------
-    r_ice_vol : npt.NDArray[np.float64]
+    r_ice_vol : npt.NDArray[np.floating]
         ice particle volume mean radius, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         scattering extinction efficiency
 
     References
@@ -1072,18 +1087,18 @@ def scattering_extinction_efficiency(r_ice_vol: npt.NDArray[np.float64]) -> npt.
     )
 
 
-def light_wave_phase_delay(r_ice_vol: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def light_wave_phase_delay(r_ice_vol: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     """
     Calculate the phase delay of the light wave passing through the contrail ice particle.
 
     Parameters
     ----------
-    r_ice_vol : npt.NDArray[np.float64]
+    r_ice_vol : npt.NDArray[np.floating]
         ice particle volume mean radius, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         phase delay of the light wave passing through the contrail ice particle
 
     References
@@ -1103,21 +1118,21 @@ def light_wave_phase_delay(r_ice_vol: npt.NDArray[np.float64]) -> npt.NDArray[np
 
 
 def segment_length_ratio(
-    seg_length_t1: npt.NDArray[np.float64],
-    seg_length_t2: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    seg_length_t1: npt.NDArray[np.floating],
+    seg_length_t2: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """Calculate the ratio of contrail segment length pre-advection to post-advection.
 
     Parameters
     ----------
-    seg_length_t1 : npt.NDArray[np.float64]
+    seg_length_t1 : npt.NDArray[np.floating]
         Segment length of contrail waypoint at the start of the time step, [:math:`m`]
-    seg_length_t2 : npt.NDArray[np.float64]
+    seg_length_t2 : npt.NDArray[np.floating]
         Segment length of contrail waypoint after time step and advection, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Ratio of segment length before advection to segment length after advection.
 
     Notes
@@ -1146,16 +1161,16 @@ def segment_length_ratio(
 
 
 def plume_temporal_evolution(
-    width_t1: npt.NDArray[np.float64],
-    depth_t1: npt.NDArray[np.float64],
-    sigma_yz_t1: npt.NDArray[np.float64],
-    dsn_dz_t1: npt.NDArray[np.float64],
-    diffuse_h_t1: npt.NDArray[np.float64],
-    diffuse_v_t1: npt.NDArray[np.float64],
-    seg_ratio: npt.NDArray[np.float64] | float,
+    width_t1: npt.NDArray[np.floating],
+    depth_t1: npt.NDArray[np.floating],
+    sigma_yz_t1: npt.NDArray[np.floating],
+    dsn_dz_t1: npt.NDArray[np.floating],
+    diffuse_h_t1: npt.NDArray[np.floating],
+    diffuse_v_t1: npt.NDArray[np.floating],
+    seg_ratio: npt.NDArray[np.floating] | float,
     dt: npt.NDArray[np.timedelta64] | np.timedelta64,
     max_depth: float | None,
-) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.floating]]:
     """
     Calculate the temporal evolution of the contrail plume parameters.
 
@@ -1163,13 +1178,13 @@ def plume_temporal_evolution(
 
     Parameters
     ----------
-    width_t1 : npt.NDArray[np.float64]
+    width_t1 : npt.NDArray[np.floating]
         contrail width at the start of the time step, [:math:`m`]
-    depth_t1 : npt.NDArray[np.float64]
+    depth_t1 : npt.NDArray[np.floating]
         contrail depth at the start of the time step, [:math:`m`]
-    sigma_yz_t1 : npt.NDArray[np.float64]
+    sigma_yz_t1 : npt.NDArray[np.floating]
         sigma_yz governs the contrail plume's temporal evolution at the start of the time step
-    dsn_dz_t1 : npt.NDArray[np.float64]
+    dsn_dz_t1 : npt.NDArray[np.floating]
         vertical gradient of the horizontal velocity (wind shear) normal to the contrail axis
         at the start of the time step, [:math:`m s^{-1} / Pa`]::
 
@@ -1178,11 +1193,11 @@ def plume_temporal_evolution(
                                         | (dsn_dz)                  |  <-- (dsn_dz)
                                         |                           |
                                                                     X
-    diffuse_h_t1 : npt.NDArray[np.float64]
+    diffuse_h_t1 : npt.NDArray[np.floating]
         horizontal diffusivity at the start of the time step, [:math:`m^{2} s^{-1}`]
-    diffuse_v_t1 : npt.NDArray[np.float64]
+    diffuse_v_t1 : npt.NDArray[np.floating]
         vertical diffusivity at the start of the time step, [:math:`m^{2} s^{-1}`]
-    seg_ratio : npt.NDArray[np.float64] | float
+    seg_ratio : npt.NDArray[np.floating] | float
         Segment length ratio before and after it is advected to the new location.
         See :func:`segment_length_ratio`.
     dt : npt.NDArray[np.timedelta64] | np.timedelta64
@@ -1193,11 +1208,11 @@ def plume_temporal_evolution(
 
     Returns
     -------
-    sigma_yy_t2 : npt.NDArray[np.float64]
+    sigma_yy_t2 : npt.NDArray[np.floating]
         The ``yy`` component of convariance matrix, [:math:`m^{2}`]
-    sigma_zz_t2 : npt.NDArray[np.float64]
+    sigma_zz_t2 : npt.NDArray[np.floating]
         The ``zz`` component of convariance matrix, [:math:`m^{2}`]
-    sigma_yz_t2 : npt.NDArray[np.float64]
+    sigma_yz_t2 : npt.NDArray[np.floating]
         The ``yz`` component of convariance matrix, [:math:`m^{2}`]
     """
     # Convert dt to seconds value and use dtype of other variables
@@ -1242,26 +1257,26 @@ def plume_temporal_evolution(
 
 
 def new_contrail_dimensions(
-    sigma_yy_t2: npt.NDArray[np.float64],
-    sigma_zz_t2: npt.NDArray[np.float64],
-) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    sigma_yy_t2: npt.NDArray[np.floating],
+    sigma_zz_t2: npt.NDArray[np.floating],
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
     """
     Calculate the new contrail width and depth.
 
     Parameters
     ----------
-    sigma_yy_t2 : npt.NDArray[np.float64]
+    sigma_yy_t2 : npt.NDArray[np.floating]
         element yy, covariance matrix of the Gaussian concentration
         field, Eq. (6) of Schumann (2012)
-    sigma_zz_t2 : npt.NDArray[np.float64]
+    sigma_zz_t2 : npt.NDArray[np.floating]
         element zz, covariance matrix of the Gaussian concentration
         field, Eq. (6) of Schumann (2012)
 
     Returns
     -------
-    width_t2 : npt.NDArray[np.float64]
+    width_t2 : npt.NDArray[np.floating]
         Contrail width at the end of the time step, [:math:`m`]
-    depth_t2 : npt.NDArray[np.float64]
+    depth_t2 : npt.NDArray[np.floating]
         Contrail depth at the end of the time step, [:math:`m`]
     """
     width_t2 = (8 * sigma_yy_t2) ** 0.5
@@ -1270,10 +1285,10 @@ def new_contrail_dimensions(
 
 
 def new_effective_area_from_sigma(
-    sigma_yy: npt.NDArray[np.float64],
-    sigma_zz: npt.NDArray[np.float64],
-    sigma_yz: npt.NDArray[np.float64] | float,
-) -> npt.NDArray[np.float64]:
+    sigma_yy: npt.NDArray[np.floating],
+    sigma_zz: npt.NDArray[np.floating],
+    sigma_yz: npt.NDArray[np.floating] | float,
+) -> npt.NDArray[np.floating]:
     """
     Calculate effective cross-sectional area of contrail plume (``area_eff``) from sigma parameters.
 
@@ -1282,19 +1297,19 @@ def new_effective_area_from_sigma(
 
     Parameters
     ----------
-    sigma_yy : npt.NDArray[np.float64]
+    sigma_yy : npt.NDArray[np.floating]
         element yy, covariance matrix of the Gaussian concentration
         field, Eq. (6) of Schumann (2012)
-    sigma_zz : npt.NDArray[np.float64]
+    sigma_zz : npt.NDArray[np.floating]
         element zz, covariance matrix of the Gaussian concentration
         field, Eq. (6) of Schumann (2012)
-    sigma_yz : npt.NDArray[np.float64] | float
+    sigma_yz : npt.NDArray[np.floating] | float
         element yz, covariance matrix of the Gaussian concentration
         field, Eq. (6) of Schumann (2012)
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Effective cross-sectional area of the contrail plume (area_eff)
     """
     det_sigma = sigma_yy * sigma_zz - sigma_yz**2
@@ -1302,44 +1317,44 @@ def new_effective_area_from_sigma(
 
 
 def new_ice_water_content(
-    iwc_t1: npt.NDArray[np.float64],
-    q_t1: npt.NDArray[np.float64],
-    q_t2: npt.NDArray[np.float64],
-    q_sat_t1: npt.NDArray[np.float64],
-    q_sat_t2: npt.NDArray[np.float64],
-    mass_plume_t1: npt.NDArray[np.float64],
-    mass_plume_t2: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    iwc_t1: npt.NDArray[np.floating],
+    q_t1: npt.NDArray[np.floating],
+    q_t2: npt.NDArray[np.floating],
+    q_sat_t1: npt.NDArray[np.floating],
+    q_sat_t2: npt.NDArray[np.floating],
+    mass_plume_t1: npt.NDArray[np.floating],
+    mass_plume_t2: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Calculate the new contrail ice water content after the time integration step (``iwc_t2``).
 
     Parameters
     ----------
-    iwc_t1 : npt.NDArray[np.float64]
+    iwc_t1 : npt.NDArray[np.floating]
         contrail ice water content, i.e., contrail ice mass per kg of air,
         at the start of the time step, [:math:`kg_{H_{2}O}/kg_{air}`]
-    q_t1 : npt.NDArray[np.float64]
+    q_t1 : npt.NDArray[np.floating]
         specific humidity for each waypoint at the start of the
         time step, [:math:`kg_{H_{2}O}/kg_{air}`]
-    q_t2 : npt.NDArray[np.float64]
+    q_t2 : npt.NDArray[np.floating]
         specific humidity for each waypoint at the end of the
         time step, [:math:`kg_{H_{2}O}/kg_{air}`]
-    q_sat_t1 : npt.NDArray[np.float64]
+    q_sat_t1 : npt.NDArray[np.floating]
         saturation humidity for each waypoint at the start of the
         time step, [:math:`kg_{H_{2}O}/kg_{air}`]
-    q_sat_t2 : npt.NDArray[np.float64]
+    q_sat_t2 : npt.NDArray[np.floating]
         saturation humidity for each waypoint at the end of the
         time step, [:math:`kg_{H_{2}O}/kg_{air}`]
-    mass_plume_t1 : npt.NDArray[np.float64]
+    mass_plume_t1 : npt.NDArray[np.floating]
         contrail plume mass per unit length at the start of the
         time step, [:math:`kg_{air} m^{-1}`]
-    mass_plume_t2 : npt.NDArray[np.float64]
+    mass_plume_t2 : npt.NDArray[np.floating]
         contrail plume mass per unit length at the end of the
         time step, [:math:`kg_{air} m^{-1}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Contrail ice water content at the end of the time step, [:math:`kg_{ice} kg_{air}^{-1}`]
 
     Notes
@@ -1363,31 +1378,31 @@ def new_ice_water_content(
 
 
 def new_ice_particle_number(
-    n_ice_per_m_t1: npt.NDArray[np.float64],
-    dn_dt_agg: npt.NDArray[np.float64],
-    dn_dt_turb: npt.NDArray[np.float64],
-    seg_ratio: npt.NDArray[np.float64] | float,
+    n_ice_per_m_t1: npt.NDArray[np.floating],
+    dn_dt_agg: npt.NDArray[np.floating],
+    dn_dt_turb: npt.NDArray[np.floating],
+    seg_ratio: npt.NDArray[np.floating] | float,
     dt: npt.NDArray[np.timedelta64] | np.timedelta64,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Calculate the number of ice particles per distance at the end of the time step.
 
     Parameters
     ----------
-    n_ice_per_m_t1 : npt.NDArray[np.float64]
+    n_ice_per_m_t1 : npt.NDArray[np.floating]
         number of contrail ice particles per distance at the start of
         the time step, [:math:`m^{-1}`]
-    dn_dt_agg : npt.NDArray[np.float64]
+    dn_dt_agg : npt.NDArray[np.floating]
         rate of ice particle losses due to sedimentation-induced aggregation, [:math:`# s^{-1}`]
-    dn_dt_turb : npt.NDArray[np.float64]
+    dn_dt_turb : npt.NDArray[np.floating]
         rate of contrail ice particle losses due to plume-internal turbulence, [:math:`# s^{-1}`]
-    seg_ratio : npt.NDArray[np.float64] | float
+    seg_ratio : npt.NDArray[np.floating] | float
         Segment length ratio before and after it is advected to the new location.
     dt : npt.NDArray[np.timedelta64] | np.timedelta64
         integrate contrails with time steps of dt, [:math:`s`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         number of ice particles per distance at the end of the time step, [:math:`m^{-1}`]
     """
     # Convert dt to seconds value and use dtype of other variables
@@ -1416,13 +1431,13 @@ def new_ice_particle_number(
 
 
 def energy_forcing(
-    rf_net_t1: npt.NDArray[np.float64],
-    rf_net_t2: npt.NDArray[np.float64],
-    width_t1: npt.NDArray[np.float64],
-    width_t2: npt.NDArray[np.float64],
-    seg_length_t2: npt.NDArray[np.float64] | float,
+    rf_net_t1: npt.NDArray[np.floating],
+    rf_net_t2: npt.NDArray[np.floating],
+    width_t1: npt.NDArray[np.floating],
+    width_t2: npt.NDArray[np.floating],
+    seg_length_t2: npt.NDArray[np.floating] | float,
     dt: npt.NDArray[np.timedelta64] | np.timedelta64,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Calculate the contrail energy forcing over time step.
 
     The contrail energy forcing is calculated as the local contrail net
@@ -1431,22 +1446,22 @@ def energy_forcing(
 
     Parameters
     ----------
-    rf_net_t1 : npt.NDArray[np.float64]
+    rf_net_t1 : npt.NDArray[np.floating]
         local contrail net radiative forcing at the start of the time step, [:math:`W m^{-2}`]
-    rf_net_t2 : npt.NDArray[np.float64]
+    rf_net_t2 : npt.NDArray[np.floating]
         local contrail net radiative forcing at the end of the time step, [:math:`W m^{-2}`]
-    width_t1 : npt.NDArray[np.float64]
+    width_t1 : npt.NDArray[np.floating]
         contrail width at the start of the time step, [:math:`m`]
-    width_t2 : npt.NDArray[np.float64]
+    width_t2 : npt.NDArray[np.floating]
         contrail width at the end of the time step, [:math:`m`]
-    seg_length_t2 : npt.NDArray[np.float64] | float
+    seg_length_t2 : npt.NDArray[np.floating] | float
         Segment length of contrail waypoint at the end of the time step, [:math:`m`]
     dt : npt.NDArray[np.timedelta64] | np.timedelta64
         integrate contrails with time steps of dt, [:math:`s`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Contrail energy forcing over time step dt, [:math:`J`].
     """
     rad_flux_per_m = mean_radiative_flux_per_m(rf_net_t1, rf_net_t2, width_t1, width_t2)
@@ -1455,27 +1470,27 @@ def energy_forcing(
 
 
 def mean_radiative_flux_per_m(
-    rf_net_t1: npt.NDArray[np.float64],
-    rf_net_t2: npt.NDArray[np.float64],
-    width_t1: npt.NDArray[np.float64],
-    width_t2: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    rf_net_t1: npt.NDArray[np.floating],
+    rf_net_t2: npt.NDArray[np.floating],
+    width_t1: npt.NDArray[np.floating],
+    width_t2: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """Calculate the mean radiative flux per length of contrail between two time steps.
 
     Parameters
     ----------
-    rf_net_t1 : npt.NDArray[np.float64]
+    rf_net_t1 : npt.NDArray[np.floating]
         local contrail net radiative forcing at the start of the time step, [:math:`W m^{-2}`]
-    rf_net_t2 : npt.NDArray[np.float64]
+    rf_net_t2 : npt.NDArray[np.floating]
         local contrail net radiative forcing at the end of the time step, [:math:`W m^{-2}`]
-    width_t1 : npt.NDArray[np.float64]
+    width_t1 : npt.NDArray[np.floating]
         contrail width at the start of the time step, [:math:`m`]
-    width_t2 : npt.NDArray[np.float64]
+    width_t2 : npt.NDArray[np.floating]
         contrail width at the end of the time step, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Mean radiative flux between time steps, [:math:`W m^{-1}`]
     """
     rad_flux_per_m_t1 = width_t1 * rf_net_t1
@@ -1484,13 +1499,13 @@ def mean_radiative_flux_per_m(
 
 
 def mean_energy_flux_per_m(
-    rad_flux_per_m: npt.NDArray[np.float64], dt: npt.NDArray[np.timedelta64] | np.timedelta64
-) -> npt.NDArray[np.float64]:
+    rad_flux_per_m: npt.NDArray[np.floating], dt: npt.NDArray[np.timedelta64] | np.timedelta64
+) -> npt.NDArray[np.floating]:
     """Calculate the mean energy flux per length of contrail on segment following waypoint.
 
     Parameters
     ----------
-    rad_flux_per_m : npt.NDArray[np.float64]
+    rad_flux_per_m : npt.NDArray[np.floating]
         Mean radiative flux between time steps for waypoint, [:math:`W m^{-1}`].
         See :func:`mean_radiative_flux_per_m`.
     dt : npt.NDArray[np.timedelta64]
@@ -1498,7 +1513,7 @@ def mean_energy_flux_per_m(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Mean energy flux per length of contrail after waypoint, [:math:`J m^{-1}`]
 
     Notes
diff --git a/pycontrails/models/cocip/output_formats.py b/pycontrails/models/cocip/output_formats.py
index 28d8596cb..cfdd7b471 100644
--- a/pycontrails/models/cocip/output_formats.py
+++ b/pycontrails/models/cocip/output_formats.py
@@ -1259,9 +1259,9 @@ def radiation_time_slice_statistics(
 
 
 def area_mean_ice_water_path(
-    iwc: npt.NDArray[np.float64],
-    plume_mass_per_m: npt.NDArray[np.float64],
-    segment_length: npt.NDArray[np.float64],
+    iwc: npt.NDArray[np.floating],
+    plume_mass_per_m: npt.NDArray[np.floating],
+    segment_length: npt.NDArray[np.floating],
     domain_area: float,
 ) -> float:
     """
@@ -1271,12 +1271,12 @@ def area_mean_ice_water_path(
 
     Parameters
     ----------
-    iwc : npt.NDArray[np.float64]
+    iwc : npt.NDArray[np.floating]
         Contrail ice water content, i.e., contrail ice mass per kg of
         air, [:math:`kg_{H_{2}O}/kg_{air}`]
-    plume_mass_per_m : npt.NDArray[np.float64]
+    plume_mass_per_m : npt.NDArray[np.floating]
         Contrail plume mass per unit length, [:math:`kg m^{-1}`]
-    segment_length : npt.NDArray[np.float64]
+    segment_length : npt.NDArray[np.floating]
         Contrail segment length for each waypoint, [:math:`m`]
     domain_area : float
         Domain surface area, [:math:`m^{2}`]
@@ -1290,20 +1290,20 @@ def area_mean_ice_water_path(
 
 
 def area_mean_ice_particle_radius(
-    r_ice_vol: npt.NDArray[np.float64],
-    n_ice_per_m: npt.NDArray[np.float64],
-    segment_length: npt.NDArray[np.float64],
+    r_ice_vol: npt.NDArray[np.floating],
+    n_ice_per_m: npt.NDArray[np.floating],
+    segment_length: npt.NDArray[np.floating],
 ) -> float:
     r"""
     Calculate the area-mean contrail ice particle radius.
 
     Parameters
     ----------
-    r_ice_vol : npt.NDArray[np.float64]
+    r_ice_vol : npt.NDArray[np.floating]
         Ice particle volume mean radius for each waypoint, [:math:`m`]
-    n_ice_per_m : npt.NDArray[np.float64]
+    n_ice_per_m : npt.NDArray[np.floating]
         Number of ice particles per distance for each waypoint, [:math:`m^{-1}`]
-    segment_length : npt.NDArray[np.float64]
+    segment_length : npt.NDArray[np.floating]
         Contrail segment length for each waypoint, [:math:`m`]
 
     Returns
@@ -1324,20 +1324,20 @@ def area_mean_ice_particle_radius(
 
 
 def volume_mean_ice_particle_radius(
-    r_ice_vol: npt.NDArray[np.float64],
-    n_ice_per_m: npt.NDArray[np.float64],
-    segment_length: npt.NDArray[np.float64],
+    r_ice_vol: npt.NDArray[np.floating],
+    n_ice_per_m: npt.NDArray[np.floating],
+    segment_length: npt.NDArray[np.floating],
 ) -> float:
     r"""
     Calculate the volume-mean contrail ice particle radius.
 
     Parameters
     ----------
-    r_ice_vol : npt.NDArray[np.float64]
+    r_ice_vol : npt.NDArray[np.floating]
         Ice particle volume mean radius for each waypoint, [:math:`m`]
-    n_ice_per_m : npt.NDArray[np.float64]
+    n_ice_per_m : npt.NDArray[np.floating]
         Number of ice particles per distance for each waypoint, [:math:`m^{-1}`]
-    segment_length : npt.NDArray[np.float64]
+    segment_length : npt.NDArray[np.floating]
         Contrail segment length for each waypoint, [:math:`m`]
 
     Returns
@@ -1356,20 +1356,20 @@ def volume_mean_ice_particle_radius(
 
 
 def mean_ice_particle_effective_radius(
-    r_ice_vol: npt.NDArray[np.float64],
-    n_ice_per_m: npt.NDArray[np.float64],
-    segment_length: npt.NDArray[np.float64],
+    r_ice_vol: npt.NDArray[np.floating],
+    n_ice_per_m: npt.NDArray[np.floating],
+    segment_length: npt.NDArray[np.floating],
 ) -> float:
     r"""
     Calculate the mean contrail ice particle effective radius.
 
     Parameters
     ----------
-    r_ice_vol : npt.NDArray[np.float64]
+    r_ice_vol : npt.NDArray[np.floating]
         Ice particle volume mean radius for each waypoint, [:math:`m`]
-    n_ice_per_m : npt.NDArray[np.float64]
+    n_ice_per_m : npt.NDArray[np.floating]
         Number of ice particles per distance for each waypoint, [:math:`m^{-1}`]
-    segment_length : npt.NDArray[np.float64]
+    segment_length : npt.NDArray[np.floating]
         Contrail segment length for each waypoint, [:math:`m`]
 
     Returns
@@ -1391,20 +1391,20 @@ def mean_ice_particle_effective_radius(
 
 
 def _total_ice_particle_cross_sectional_area(
-    r_ice_vol: npt.NDArray[np.float64],
-    n_ice_per_m: npt.NDArray[np.float64],
-    segment_length: npt.NDArray[np.float64],
+    r_ice_vol: npt.NDArray[np.floating],
+    n_ice_per_m: npt.NDArray[np.floating],
+    segment_length: npt.NDArray[np.floating],
 ) -> float:
     """
     Calculate total contrail ice particle cross-sectional area.
 
     Parameters
     ----------
-    r_ice_vol : npt.NDArray[np.float64]
+    r_ice_vol : npt.NDArray[np.floating]
         Ice particle volume mean radius for each waypoint, [:math:`m`]
-    n_ice_per_m : npt.NDArray[np.float64]
+    n_ice_per_m : npt.NDArray[np.floating]
         Number of ice particles per distance for each waypoint, [:math:`m^{-1}`]
-    segment_length : npt.NDArray[np.float64]
+    segment_length : npt.NDArray[np.floating]
         Contrail segment length for each waypoint, [:math:`m`]
 
     Returns
@@ -1417,20 +1417,20 @@ def _total_ice_particle_cross_sectional_area(
 
 
 def _total_ice_particle_volume(
-    r_ice_vol: npt.NDArray[np.float64],
-    n_ice_per_m: npt.NDArray[np.float64],
-    segment_length: npt.NDArray[np.float64],
+    r_ice_vol: npt.NDArray[np.floating],
+    n_ice_per_m: npt.NDArray[np.floating],
+    segment_length: npt.NDArray[np.floating],
 ) -> float:
     """
     Calculate total contrail ice particle volume.
 
     Parameters
     ----------
-    r_ice_vol : npt.NDArray[np.float64]
+    r_ice_vol : npt.NDArray[np.floating]
         Ice particle volume mean radius for each waypoint, [:math:`m`]
-    n_ice_per_m : npt.NDArray[np.float64]
+    n_ice_per_m : npt.NDArray[np.floating]
         Number of ice particles per distance for each waypoint, [:math:`m^{-1}`]
-    segment_length : npt.NDArray[np.float64]
+    segment_length : npt.NDArray[np.floating]
         Contrail segment length for each waypoint, [:math:`m`]
 
     Returns
@@ -1443,16 +1443,16 @@ def _total_ice_particle_volume(
 
 
 def _total_ice_particle_number(
-    n_ice_per_m: npt.NDArray[np.float64], segment_length: npt.NDArray[np.float64]
+    n_ice_per_m: npt.NDArray[np.floating], segment_length: npt.NDArray[np.floating]
 ) -> float:
     """
     Calculate total number of contrail ice particles.
 
     Parameters
     ----------
-    n_ice_per_m : npt.NDArray[np.float64]
+    n_ice_per_m : npt.NDArray[np.floating]
         Number of ice particles per distance for each waypoint, [:math:`m^{-1}`]
-    segment_length : npt.NDArray[np.float64]
+    segment_length : npt.NDArray[np.floating]
         Contrail segment length for each waypoint, [:math:`m`]
 
     Returns
@@ -1464,9 +1464,9 @@ def _total_ice_particle_number(
 
 
 def area_mean_contrail_property(
-    contrail_property: npt.NDArray[np.float64],
-    segment_length: npt.NDArray[np.float64],
-    width: npt.NDArray[np.float64],
+    contrail_property: npt.NDArray[np.floating],
+    segment_length: npt.NDArray[np.floating],
+    width: npt.NDArray[np.floating],
     domain_area: float,
 ) -> float:
     """
@@ -1477,11 +1477,11 @@ def area_mean_contrail_property(
 
     Parameters
     ----------
-    contrail_property : npt.NDArray[np.float64]
+    contrail_property : npt.NDArray[np.floating]
         Selected contrail property for each waypoint
-    segment_length : npt.NDArray[np.float64]
+    segment_length : npt.NDArray[np.floating]
         Contrail segment length for each waypoint, [:math:`m`]
-    width : npt.NDArray[np.float64]
+    width : npt.NDArray[np.floating]
         Contrail width for each waypoint, [:math:`m`]
     domain_area : float
         Domain surface area, [:math:`m^{2}`]
@@ -2009,26 +2009,26 @@ def natural_cirrus_properties_to_hi_res_grid(
 
 
 def _hi_res_grid_coordinates(
-    lon_coords: npt.NDArray[np.float64],
-    lat_coords: npt.NDArray[np.float64],
+    lon_coords: npt.NDArray[np.floating],
+    lat_coords: npt.NDArray[np.floating],
     *,
     spatial_grid_res: float = 0.05,
-) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
     r"""
     Calculate longitude and latitude coordinates for the high resolution grid.
 
     Parameters
     ----------
-    lon_coords : npt.NDArray[np.float64]
+    lon_coords : npt.NDArray[np.floating]
         Longitude coordinates provided by the original `MetDataset`.
-    lat_coords : npt.NDArray[np.float64]
+    lat_coords : npt.NDArray[np.floating]
         Latitude coordinates provided by the original `MetDataset`.
     spatial_grid_res : float
         Spatial grid resolution for the output, [:math:`\deg`]
 
     Returns
     -------
-    tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]
+    tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]
         Longitude and latitude coordinates for the high resolution grid.
     """
     d_lon = np.abs(np.diff(lon_coords)[0])
@@ -2058,8 +2058,8 @@ def _hi_res_grid_coordinates(
 
 
 def _repeat_rows_and_columns(
-    array_2d: npt.NDArray[np.float64], *, n_reps: int
-) -> npt.NDArray[np.float64]:
+    array_2d: npt.NDArray[np.floating], *, n_reps: int
+) -> npt.NDArray[np.floating]:
     """
     Repeat the elements in `array_2d` along each row and column.
 
diff --git a/pycontrails/models/cocip/radiative_forcing.py b/pycontrails/models/cocip/radiative_forcing.py
index 2f95a515c..34e0a6ed4 100644
--- a/pycontrails/models/cocip/radiative_forcing.py
+++ b/pycontrails/models/cocip/radiative_forcing.py
@@ -157,10 +157,10 @@ class RFConstants:
 
 
 def habit_weights(
-    r_vol_um: npt.NDArray[np.float64],
-    habit_distributions: npt.NDArray[np.float64],
-    radius_threshold_um: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    r_vol_um: npt.NDArray[np.floating],
+    habit_distributions: npt.NDArray[np.floating],
+    radius_threshold_um: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""Assign weights to different ice particle habits for each waypoint.
 
     For each waypoint, the distinct mix of ice particle habits are approximated
@@ -175,18 +175,18 @@ def habit_weights(
 
     Parameters
     ----------
-    r_vol_um : npt.NDArray[np.float64]
+    r_vol_um : npt.NDArray[np.floating]
         Contrail ice particle volume mean radius, [:math:`\mu m`]
-    habit_distributions : npt.NDArray[np.float64]
+    habit_distributions : npt.NDArray[np.floating]
         Habit weight distributions.
         See :attr:`CocipParams().habit_distributions`
-    radius_threshold_um : npt.NDArray[np.float64]
+    radius_threshold_um : npt.NDArray[np.floating]
         Radius thresholds for habit distributions.
         See :attr:`CocipParams.radius_threshold_um`
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Array with shape ``n_waypoints x 8 columns``, where each column is the weights to the ice
         particle habits, [:math:`[0 - 1]`], and the sum of each column should be equal to 1.
 
@@ -212,16 +212,16 @@ def habit_weights(
 
 
 def habit_weight_regime_idx(
-    r_vol_um: npt.NDArray[np.float64], radius_threshold_um: npt.NDArray[np.float64]
+    r_vol_um: npt.NDArray[np.floating], radius_threshold_um: npt.NDArray[np.floating]
 ) -> npt.NDArray[np.intp]:
     r"""
     Determine regime of ice particle habits based on contrail ice particle volume mean radius.
 
     Parameters
     ----------
-    r_vol_um : npt.NDArray[np.float64]
+    r_vol_um : npt.NDArray[np.floating]
         Contrail ice particle volume mean radius, [:math:`\mu m`]
-    radius_threshold_um : npt.NDArray[np.float64]
+    radius_threshold_um : npt.NDArray[np.floating]
         Radius thresholds for habit distributions.
         See :attr:`CocipParams.radius_threshold_um`
 
@@ -240,7 +240,7 @@ def habit_weight_regime_idx(
 
 
 def effective_radius_by_habit(
-    r_vol_um: npt.NDArray[np.float64], habit_idx: npt.NDArray[np.intp]
+    r_vol_um: npt.NDArray[np.floating], habit_idx: npt.NDArray[np.intp]
 ) -> np.ndarray:
     r"""Calculate the effective radius ``r_eff_um`` via the mean ice particle radius and habit type.
 
@@ -250,7 +250,7 @@ def effective_radius_by_habit(
 
     Parameters
     ----------
-    r_vol_um : npt.NDArray[np.float64]
+    r_vol_um : npt.NDArray[np.floating]
         Contrail ice particle volume mean radius, [:math:`\mu m`]
     habit_idx : npt.NDArray[np.intp]
         Habit type index for the contrail ice particle, corresponding to the
@@ -258,7 +258,7 @@ def effective_radius_by_habit(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Effective radius of ice particles for each combination of ``r_vol_um``
         and ``habit_idx``, [:math:`\mu m`]
 
@@ -290,35 +290,35 @@ def effective_radius_by_habit(
     return np.piecewise(r_vol_um, cond_list, func_list)
 
 
-def effective_radius_sphere(r_vol_um: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def effective_radius_sphere(r_vol_um: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     r"""
     Calculate the effective radius of contrail ice particles assuming a sphere particle habit.
 
     Parameters
     ----------
-    r_vol_um : npt.NDArray[np.float64]
+    r_vol_um : npt.NDArray[np.floating]
         Contrail ice particle volume mean radius, [:math:`\mu m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Effective radius, [:math:`\mu m`]
     """
     return np.minimum(r_vol_um, 25.0)
 
 
-def effective_radius_solid_column(r_vol_um: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def effective_radius_solid_column(r_vol_um: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     r"""
     Calculate the effective radius of contrail ice particles assuming a solid column particle habit.
 
     Parameters
     ----------
-    r_vol_um : npt.NDArray[np.float64]
+    r_vol_um : npt.NDArray[np.floating]
         Contrail ice particle volume mean radius, [:math:`\mu m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Effective radius, [:math:`\mu m`]
     """
     r_eff_um = (
@@ -329,17 +329,17 @@ def effective_radius_solid_column(r_vol_um: npt.NDArray[np.float64]) -> npt.NDAr
     return np.minimum(r_eff_um, 45.0)
 
 
-def effective_radius_hollow_column(r_vol_um: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def effective_radius_hollow_column(r_vol_um: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     r"""Calculate the effective radius of ice particles assuming a hollow column particle habit.
 
     Parameters
     ----------
-    r_vol_um : npt.NDArray[np.float64]
+    r_vol_um : npt.NDArray[np.floating]
         Contrail ice particle volume mean radius, [:math:`\mu m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Effective radius, [:math:`\mu m`]
     """
     r_eff_um = (
@@ -350,35 +350,37 @@ def effective_radius_hollow_column(r_vol_um: npt.NDArray[np.float64]) -> npt.NDA
     return np.minimum(r_eff_um, 45.0)
 
 
-def effective_radius_rough_aggregate(r_vol_um: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def effective_radius_rough_aggregate(
+    r_vol_um: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""Calculate the effective radius of ice particles assuming a rough aggregate particle habit.
 
     Parameters
     ----------
-    r_vol_um : npt.NDArray[np.float64]
+    r_vol_um : npt.NDArray[np.floating]
         Contrail ice particle volume mean radius, [:math:`\mu m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Effective radius, [:math:`\mu m`]
     """
     r_eff_um = 0.574 * r_vol_um
     return np.minimum(r_eff_um, 45.0)
 
 
-def effective_radius_rosette(r_vol_um: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def effective_radius_rosette(r_vol_um: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     r"""
     Calculate the effective radius of contrail ice particles assuming a rosette particle habit.
 
     Parameters
     ----------
-    r_vol_um : npt.NDArray[np.float64]
+    r_vol_um : npt.NDArray[np.floating]
         Contrail ice particle volume mean radius, [:math:`\mu m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Effective radius, [:math:`\mu m`]
     """
     r_eff_um = r_vol_um * (
@@ -387,18 +389,18 @@ def effective_radius_rosette(r_vol_um: npt.NDArray[np.float64]) -> npt.NDArray[n
     return np.minimum(r_eff_um, 45.0)
 
 
-def effective_radius_plate(r_vol_um: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def effective_radius_plate(r_vol_um: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     r"""
     Calculate the effective radius of contrail ice particles assuming a plate particle habit.
 
     Parameters
     ----------
-    r_vol_um : npt.NDArray[np.float64]
+    r_vol_um : npt.NDArray[np.floating]
         Contrail ice particle volume mean radius, [:math:`\mu m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Effective radius, [:math:`\mu m`]
     """
     r_eff_um = r_vol_um * (
@@ -407,36 +409,36 @@ def effective_radius_plate(r_vol_um: npt.NDArray[np.float64]) -> npt.NDArray[np.
     return np.minimum(r_eff_um, 45.0)
 
 
-def effective_radius_droxtal(r_vol_um: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def effective_radius_droxtal(r_vol_um: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     r"""
     Calculate the effective radius of contrail ice particles assuming a droxtal particle habit.
 
     Parameters
     ----------
-    r_vol_um : npt.NDArray[np.float64]
+    r_vol_um : npt.NDArray[np.floating]
         Contrail ice particle volume mean radius, [:math:`\mu m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Effective radius, [:math:`\mu m`]
     """
     r_eff_um = 0.94 * r_vol_um
     return np.minimum(r_eff_um, 45.0)
 
 
-def effective_radius_myhre(r_vol_um: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def effective_radius_myhre(r_vol_um: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     r"""
     Calculate the effective radius of contrail ice particles assuming a sphere particle habit.
 
     Parameters
     ----------
-    r_vol_um : npt.NDArray[np.float64]
+    r_vol_um : npt.NDArray[np.floating]
         Contrail ice particle volume mean radius, [:math:`\mu m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Effective radius, [:math:`\mu m`]
     """
     return np.minimum(r_vol_um, 45.0)
@@ -448,14 +450,14 @@ def effective_radius_myhre(r_vol_um: npt.NDArray[np.float64]) -> npt.NDArray[np.
 
 
 def longwave_radiative_forcing(
-    r_vol_um: npt.NDArray[np.float64],
-    olr: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-    tau_contrail: npt.NDArray[np.float64],
-    tau_cirrus: npt.NDArray[np.float64],
-    habit_weights_: npt.NDArray[np.float64],
-    r_eff_um: npt.NDArray[np.float64] | None = None,
-) -> npt.NDArray[np.float64]:
+    r_vol_um: npt.NDArray[np.floating],
+    olr: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+    tau_contrail: npt.NDArray[np.floating],
+    tau_cirrus: npt.NDArray[np.floating],
+    habit_weights_: npt.NDArray[np.floating],
+    r_eff_um: npt.NDArray[np.floating] | None = None,
+) -> npt.NDArray[np.floating]:
     r"""
     Calculate the local contrail longwave radiative forcing (:math:`RF_{LW}`).
 
@@ -463,28 +465,28 @@ def longwave_radiative_forcing(
 
     Parameters
     ----------
-    r_vol_um : npt.NDArray[np.float64]
+    r_vol_um : npt.NDArray[np.floating]
         Contrail ice particle volume mean radius, [:math:`\mu m`]
-    olr : npt.NDArray[np.float64]
+    olr : npt.NDArray[np.floating]
         Outgoing longwave radiation at each waypoint, [:math:`W m^{-2}`]
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         Ambient temperature at each waypoint, [:math:`K`]
-    tau_contrail : npt.NDArray[np.float64]
+    tau_contrail : npt.NDArray[np.floating]
         Contrail optical depth at each waypoint
-    tau_cirrus : npt.NDArray[np.float64]
+    tau_cirrus : npt.NDArray[np.floating]
         Optical depth of numerical weather prediction (NWP) cirrus above the
         contrail at each waypoint
-    habit_weights_ : npt.NDArray[np.float64]
+    habit_weights_ : npt.NDArray[np.floating]
         Weights to different ice particle habits for each waypoint,
         ``n_waypoints x 8`` (habit) columns, [:math:`[0 - 1]`]
-    r_eff_um : npt.NDArray[np.float64], optional
+    r_eff_um : npt.NDArray[np.floating], optional
         Provide effective radius corresponding to elements in ``r_vol_um``, [:math:`\mu m`].
         Defaults to None, which means the effective radius will be calculated using ``r_vol_um``
         and habit types in :func:`effective_radius_by_habit`.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Local contrail longwave radiative forcing (positive), [:math:`W m^{-2}`]
 
     Raises
@@ -549,15 +551,15 @@ def longwave_radiative_forcing(
 
 
 def shortwave_radiative_forcing(
-    r_vol_um: npt.NDArray[np.float64],
-    sdr: npt.NDArray[np.float64],
-    rsr: npt.NDArray[np.float64],
-    sd0: npt.NDArray[np.float64],
-    tau_contrail: npt.NDArray[np.float64],
-    tau_cirrus: npt.NDArray[np.float64],
-    habit_weights_: npt.NDArray[np.float64],
-    r_eff_um: npt.NDArray[np.float64] | None = None,
-) -> npt.NDArray[np.float64]:
+    r_vol_um: npt.NDArray[np.floating],
+    sdr: npt.NDArray[np.floating],
+    rsr: npt.NDArray[np.floating],
+    sd0: npt.NDArray[np.floating],
+    tau_contrail: npt.NDArray[np.floating],
+    tau_cirrus: npt.NDArray[np.floating],
+    habit_weights_: npt.NDArray[np.floating],
+    r_eff_um: npt.NDArray[np.floating] | None = None,
+) -> npt.NDArray[np.floating]:
     r"""
     Calculate the local contrail shortwave radiative forcing (:math:`RF_{SW}`).
 
@@ -565,30 +567,30 @@ def shortwave_radiative_forcing(
 
     Parameters
     ----------
-    r_vol_um : npt.NDArray[np.float64]
+    r_vol_um : npt.NDArray[np.floating]
         Contrail ice particle volume mean radius, [:math:`\mu m`]
-    sdr : npt.NDArray[np.float64]
+    sdr : npt.NDArray[np.floating]
         Solar direct radiation, [:math:`W m^{-2}`]
-    rsr : npt.NDArray[np.float64]
+    rsr : npt.NDArray[np.floating]
         Reflected solar radiation, [:math:`W m^{-2}`]
-    sd0 : npt.NDArray[np.float64]
+    sd0 : npt.NDArray[np.floating]
         Solar constant, [:math:`W m^{-2}`]
-    tau_contrail : npt.NDArray[np.float64]
+    tau_contrail : npt.NDArray[np.floating]
         Contrail optical depth for each waypoint
-    tau_cirrus : npt.NDArray[np.float64]
+    tau_cirrus : npt.NDArray[np.floating]
         Optical depth of numerical weather prediction (NWP) cirrus above the
         contrail for each waypoint.
-    habit_weights_ : npt.NDArray[np.float64]
+    habit_weights_ : npt.NDArray[np.floating]
         Weights to different ice particle habits for each waypoint,
         ``n_waypoints x 8`` (habit) columns, [:math:`[0 - 1]`]
-    r_eff_um : npt.NDArray[np.float64], optional
+    r_eff_um : npt.NDArray[np.floating], optional
         Provide effective radius corresponding to elements in ``r_vol_um``, [:math:`\mu m`].
         Defaults to None, which means the effective radius will be calculated using ``r_vol_um``
         and habit types in :func:`effective_radius_by_habit`.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Local contrail shortwave radiative forcing (negative), [:math:`W m^{-2}`]
 
     Raises
@@ -677,8 +679,8 @@ def shortwave_radiative_forcing(
 
 
 def net_radiative_forcing(
-    rf_lw: npt.NDArray[np.float64], rf_sw: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    rf_lw: npt.NDArray[np.floating], rf_sw: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate the local contrail net radiative forcing (rf_net).
 
@@ -686,22 +688,22 @@ def net_radiative_forcing(
 
     Parameters
     ----------
-    rf_lw : npt.NDArray[np.float64]
+    rf_lw : npt.NDArray[np.floating]
         local contrail longwave radiative forcing, [:math:`W m^{-2}`]
-    rf_sw : npt.NDArray[np.float64]
+    rf_sw : npt.NDArray[np.floating]
         local contrail shortwave radiative forcing, [:math:`W m^{-2}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         local contrail net radiative forcing, [:math:`W m^{-2}`]
     """
     return rf_lw + rf_sw
 
 
 def olr_reduction_natural_cirrus(
-    tau_cirrus: npt.NDArray[np.float64], delta_lc: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    tau_cirrus: npt.NDArray[np.floating], delta_lc: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate reduction in outgoing longwave radiation (OLR) due to the presence of natural cirrus.
 
@@ -710,16 +712,16 @@ def olr_reduction_natural_cirrus(
 
     Parameters
     ----------
-    tau_cirrus : npt.NDArray[np.float64]
+    tau_cirrus : npt.NDArray[np.floating]
         Optical depth of numerical weather prediction (NWP) cirrus above the
         contrail for each waypoint.
-    delta_lc : npt.NDArray[np.float64]
+    delta_lc : npt.NDArray[np.floating]
         Habit specific parameter to approximate the reduction of the outgoing
         longwave radiation at the contrail level due to natural cirrus above the contrail.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Reduction of outgoing longwave radiation
     """
     # e_lw calculations
@@ -727,30 +729,32 @@ def olr_reduction_natural_cirrus(
 
 
 def contrail_effective_emissivity(
-    r_eff_um: npt.NDArray[np.float64], delta_lr: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    r_eff_um: npt.NDArray[np.floating], delta_lr: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     r"""Calculate the effective emissivity of the contrail, ``f_lw``.
 
     Refer to Eq. (3) of Schumann et al. (2012).
 
     Parameters
     ----------
-    r_eff_um : npt.NDArray[np.float64]
+    r_eff_um : npt.NDArray[np.floating]
         Effective radius for each waypoint, n_waypoints x 8 (habit) columns, [:math:`\mu m`]
         See :func:`effective_radius_habit`.
-    delta_lr : npt.NDArray[np.float64]
+    delta_lr : npt.NDArray[np.floating]
         Habit specific parameter to approximate the effective emissivity of the contrail.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Effective emissivity of the contrail
     """
     # f_lw calculations
     return 1.0 - np.exp(-delta_lr * r_eff_um)
 
 
-def albedo(sdr: npt.NDArray[np.float64], rsr: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def albedo(
+    sdr: npt.NDArray[np.floating], rsr: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate albedo along contrail waypoint.
 
@@ -762,14 +766,14 @@ def albedo(sdr: npt.NDArray[np.float64], rsr: npt.NDArray[np.float64]) -> npt.ND
 
     Parameters
     ----------
-    sdr : npt.NDArray[np.float64]
+    sdr : npt.NDArray[np.floating]
         Solar direct radiation, [:math:`W m^{-2}`]
-    rsr : npt.NDArray[np.float64]
+    rsr : npt.NDArray[np.floating]
         Reflected solar radiation, [:math:`W m^{-2}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Albedo value, [:math:`[0 - 1]`]
     """
     day = sdr > 0.0
@@ -780,17 +784,17 @@ def albedo(sdr: npt.NDArray[np.float64], rsr: npt.NDArray[np.float64]) -> npt.ND
 
 
 def contrail_albedo(
-    tau_contrail: npt.NDArray[np.float64],
-    mue: npt.NDArray[np.float64],
-    r_eff_um: npt.NDArray[np.float64],
-    A_mu: npt.NDArray[np.float64],
-    B_mu: npt.NDArray[np.float64],
-    C_mu: npt.NDArray[np.float64],
-    delta_sr: npt.NDArray[np.float64],
-    F_r: npt.NDArray[np.float64],
-    gamma_lower: npt.NDArray[np.float64],
-    gamma_upper: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    tau_contrail: npt.NDArray[np.floating],
+    mue: npt.NDArray[np.floating],
+    r_eff_um: npt.NDArray[np.floating],
+    A_mu: npt.NDArray[np.floating],
+    B_mu: npt.NDArray[np.floating],
+    C_mu: npt.NDArray[np.floating],
+    delta_sr: npt.NDArray[np.floating],
+    F_r: npt.NDArray[np.floating],
+    gamma_lower: npt.NDArray[np.floating],
+    gamma_upper: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""
     Calculate the contrail albedo, ``alpha_c``.
 
@@ -798,31 +802,31 @@ def contrail_albedo(
 
     Parameters
     ----------
-    tau_contrail : npt.NDArray[np.float64]
+    tau_contrail : npt.NDArray[np.floating]
         Contrail optical depth for each waypoint
-    mue : npt.NDArray[np.float64]
+    mue : npt.NDArray[np.floating]
         Cosine of the solar zenith angle (theta), mue = cos(theta) = sdr/sd0
-    r_eff_um : npt.NDArray[np.float64]
+    r_eff_um : npt.NDArray[np.floating]
         Effective radius for each waypoint, n_waypoints x 8 (habit) columns, [:math:`\mu m`]
         See :func:`effective_radius_habit`.
-    A_mu : npt.NDArray[np.float64]
+    A_mu : npt.NDArray[np.floating]
         Habit-specific parameter to approximate the albedo of the contrail
-    B_mu : npt.NDArray[np.float64]
+    B_mu : npt.NDArray[np.floating]
         Habit-specific parameter to approximate the SZA-dependent contrail sideward scattering
-    C_mu : npt.NDArray[np.float64]
+    C_mu : npt.NDArray[np.floating]
         Habit-specific parameter to approximate the albedo of the contrail
-    delta_sr : npt.NDArray[np.float64]
+    delta_sr : npt.NDArray[np.floating]
         Habit-specific parameter to approximate the effective contrail optical depth
-    F_r : npt.NDArray[np.float64]
+    F_r : npt.NDArray[np.floating]
         Habit-specific parameter to approximate the effective contrail optical depth
-    gamma_lower : npt.NDArray[np.float64]
+    gamma_lower : npt.NDArray[np.floating]
         Habit-specific parameter to approximate the contrail reflectances
-    gamma_upper : npt.NDArray[np.float64]
+    gamma_upper : npt.NDArray[np.floating]
         Habit-specific parameter to approximate the contrail reflectances
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Contrail albedo for each waypoint and ice particle habit
     """
     tau_aps = tau_contrail * (1.0 - F_r * (1 - np.exp(-delta_sr * r_eff_um)))
@@ -835,11 +839,11 @@ def contrail_albedo(
 
 
 def effective_tau_cirrus(
-    tau_cirrus: npt.NDArray[np.float64],
-    mue: npt.NDArray[np.float64],
-    delta_sc: npt.NDArray[np.float64],
-    delta_sc_aps: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    tau_cirrus: npt.NDArray[np.floating],
+    mue: npt.NDArray[np.floating],
+    delta_sc: npt.NDArray[np.floating],
+    delta_sc_aps: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""
     Calculate the effective optical depth of natural cirrus above the contrail, ``e_sw``.
 
@@ -848,21 +852,21 @@ def effective_tau_cirrus(
 
     Parameters
     ----------
-    tau_cirrus : npt.NDArray[np.float64]
+    tau_cirrus : npt.NDArray[np.floating]
         Optical depth of numerical weather prediction (NWP) cirrus above the
         contrail for each waypoint.
-    mue : npt.NDArray[np.float64]
+    mue : npt.NDArray[np.floating]
         Cosine of the solar zenith angle (theta), mue = cos(theta) = sdr/sd0
-    delta_sc : npt.NDArray[np.float64]
+    delta_sc : npt.NDArray[np.floating]
         Habit-specific parameter to account for the optical depth of natural
         cirrus above the contrail
-    delta_sc_aps : npt.NDArray[np.float64]
+    delta_sc_aps : npt.NDArray[np.floating]
         Habit-specific parameter to account for the optical depth of natural
         cirrus above the contrail
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Effective optical depth of natural cirrus above the contrail,
         ``n_waypoints x 8`` (habit) columns.
 
@@ -883,8 +887,8 @@ def effective_tau_cirrus(
 
 def contrail_contrail_overlap_radiative_effects(
     contrails: GeoVectorDataset,
-    habit_distributions: npt.NDArray[np.float64],
-    radius_threshold_um: npt.NDArray[np.float64],
+    habit_distributions: npt.NDArray[np.floating],
+    radius_threshold_um: npt.NDArray[np.floating],
     *,
     min_altitude_m: float = 6000.0,
     max_altitude_m: float = 13000.0,
@@ -910,10 +914,10 @@ def contrail_contrail_overlap_radiative_effects(
         - rsr
         - olr
 
-    habit_distributions : npt.NDArray[np.float64]
+    habit_distributions : npt.NDArray[np.floating]
         Habit weight distributions.
         See :attr:`CocipParams.habit_distributions`
-    radius_threshold_um : npt.NDArray[np.float64]
+    radius_threshold_um : npt.NDArray[np.floating]
         Radius thresholds for habit distributions.
         See :attr:`CocipParams.radius_threshold_um`
     min_altitude_m : float
@@ -1154,8 +1158,8 @@ def _rsr_and_olr_with_contrail_overlap(
 
 def _local_sw_and_lw_rf_with_contrail_overlap(
     contrails_level: GeoVectorDataset,
-    habit_distributions: npt.NDArray[np.float64],
-    radius_threshold_um: npt.NDArray[np.float64],
+    habit_distributions: npt.NDArray[np.floating],
+    radius_threshold_um: npt.NDArray[np.floating],
 ) -> GeoVectorDataset:
     """
     Calculate local contrail SW and LW RF after accounting for contrail overlapping.
@@ -1164,10 +1168,10 @@ def _local_sw_and_lw_rf_with_contrail_overlap(
     ----------
     contrails_level : GeoVectorDataset
         Contrail waypoints at the current altitude layer.
-    habit_distributions : npt.NDArray[np.float64]
+    habit_distributions : npt.NDArray[np.floating]
         Habit weight distributions.
         See :attr:`CocipParams().habit_distributions`
-    radius_threshold_um : npt.NDArray[np.float64]
+    radius_threshold_um : npt.NDArray[np.floating]
         Radius thresholds for habit distributions.
         See :attr:`CocipParams.radius_threshold_um`
 
diff --git a/pycontrails/models/cocip/radiative_heating.py b/pycontrails/models/cocip/radiative_heating.py
index bca1b3365..a7cebf94c 100644
--- a/pycontrails/models/cocip/radiative_heating.py
+++ b/pycontrails/models/cocip/radiative_heating.py
@@ -98,25 +98,25 @@ class RadiativeHeatingConstants:
 
 
 def convective_velocity_scale(
-    depth_eff: npt.NDArray[np.float64],
-    eff_heat_rate: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    depth_eff: npt.NDArray[np.floating],
+    eff_heat_rate: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Calculate the convective velocity scale, i.e., vertical mixing rate.
 
     Parameters
     ----------
-    depth_eff : npt.NDArray[np.float64]
+    depth_eff : npt.NDArray[np.floating]
         Effective depth of the contrail plume, [:math:`m`]
-    eff_heat_rate: npt.NDArray[np.float64]
+    eff_heat_rate: npt.NDArray[np.floating]
         Effective heating rate, i.e., rate of which the contrail plume is heated, [:math:`K s^{-1}`]
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         Ambient temperature for each waypoint, [:math:`K`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Convective velocity scale, [:math:`m s^{-1}`]
     """
     return ((constants.g * depth_eff**2 * np.maximum(-eff_heat_rate, 0)) / air_temperature) ** (
@@ -125,11 +125,11 @@ def convective_velocity_scale(
 
 
 def effective_heating_rate(
-    d_heat_rate: npt.NDArray[np.float64],
-    cumul_rad_heat: npt.NDArray[np.float64],
-    dT_dz: npt.NDArray[np.float64],
-    depth: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    d_heat_rate: npt.NDArray[np.floating],
+    cumul_rad_heat: npt.NDArray[np.floating],
+    dT_dz: npt.NDArray[np.floating],
+    depth: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """Calculate effective heating rate.
 
     The effective heating rate accounts for the heat required to overcome the stable stratification.
@@ -137,20 +137,20 @@ def effective_heating_rate(
 
     Parameters
     ----------
-    d_heat_rate: npt.NDArray[np.float64]
+    d_heat_rate: npt.NDArray[np.floating]
         Differential heating rate, i.e., rate of which the contrail
         plume is heated, [:math:`K s^{-1}`]
-    cumul_rad_heat: npt.NDArray[np.float64]
+    cumul_rad_heat: npt.NDArray[np.floating]
         Cumulative solar and terrestrial radiative heating energy
         absorbed by the contrail, [:math:`K`]
-    dT_dz: npt.NDArray[np.float64]
+    dT_dz: npt.NDArray[np.floating]
         Temperature gradient with respect to altitude (dz), [:math:`K m^{-1}`]
-    depth : npt.NDArray[np.float64]
+    depth : npt.NDArray[np.floating]
         Contrail depth at each waypoint, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Effective heating rate, [:math:`K s^{-1}`]
     """
     filt = cumul_rad_heat > 0.0
@@ -164,18 +164,18 @@ def effective_heating_rate(
 
 
 def differential_heating_rate(
-    air_temperature: npt.NDArray[np.float64],
-    rhi: npt.NDArray[np.float64],
-    rho_air: npt.NDArray[np.float64],
-    r_ice_vol: npt.NDArray[np.float64],
-    depth_eff: npt.NDArray[np.float64],
-    tau_contrail: npt.NDArray[np.float64],
-    tau_cirrus: npt.NDArray[np.float64],
-    sd0: npt.NDArray[np.float64],
-    sdr: npt.NDArray[np.float64],
-    rsr: npt.NDArray[np.float64],
-    olr: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    air_temperature: npt.NDArray[np.floating],
+    rhi: npt.NDArray[np.floating],
+    rho_air: npt.NDArray[np.floating],
+    r_ice_vol: npt.NDArray[np.floating],
+    depth_eff: npt.NDArray[np.floating],
+    tau_contrail: npt.NDArray[np.floating],
+    tau_cirrus: npt.NDArray[np.floating],
+    sd0: npt.NDArray[np.floating],
+    sdr: npt.NDArray[np.floating],
+    rsr: npt.NDArray[np.floating],
+    olr: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Calculate the differential heating rate affecting the contrail plume.
 
@@ -186,32 +186,32 @@ def differential_heating_rate(
 
     Parameters
     ----------
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         Ambient temperature at each waypoint, [:math:`K`]
-    rhi : npt.NDArray[np.float64]
+    rhi : npt.NDArray[np.floating]
         Relative humidity with respect to ice at each waypoint
-    rho_air : npt.NDArray[np.float64]
+    rho_air : npt.NDArray[np.floating]
         Density of air for each waypoint, [:math:`kg m^{-3}`]
-    r_ice_vol : npt.NDArray[np.float64]
+    r_ice_vol : npt.NDArray[np.floating]
         Ice particle volume mean radius, [:math:`m`]
-    depth_eff : npt.NDArray[np.float64]
+    depth_eff : npt.NDArray[np.floating]
         Effective depth of the contrail plume, [:math:`m`]
-    tau_contrail : npt.NDArray[np.float64]
+    tau_contrail : npt.NDArray[np.floating]
         Contrail optical depth for each waypoint
-    tau_cirrus : npt.NDArray[np.float64]
+    tau_cirrus : npt.NDArray[np.floating]
         Optical depth of numerical weather prediction (NWP) cirrus above the contrail
-    sd0 : npt.NDArray[np.float64]
+    sd0 : npt.NDArray[np.floating]
         Solar constant, [:math:`W m^{-2}`]
-    sdr : npt.NDArray[np.float64]
+    sdr : npt.NDArray[np.floating]
         Solar direct radiation, [:math:`W m^{-2}`]
-    rsr : npt.NDArray[np.float64]
+    rsr : npt.NDArray[np.floating]
         Reflected solar radiation, [:math:`W m^{-2}`]
-    olr : npt.NDArray[np.float64]
+    olr : npt.NDArray[np.floating]
         Outgoing longwave radiation at each waypoint, [:math:`W m^{-2}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Differential heating rate, [:math:`K s^{-1}`]
     """
     r_ice_vol_um = r_ice_vol * 1e6
@@ -226,14 +226,14 @@ def differential_heating_rate(
 
 
 def differential_heating_rate_shortwave(
-    cp_contrail: npt.NDArray[np.float64],
-    r_ice_vol_um: npt.NDArray[np.float64],
-    tau_contrail: npt.NDArray[np.float64],
-    tau_cirrus: npt.NDArray[np.float64],
-    sd0: npt.NDArray[np.float64],
-    sdr: npt.NDArray[np.float64],
-    rsr: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    cp_contrail: npt.NDArray[np.floating],
+    r_ice_vol_um: npt.NDArray[np.floating],
+    tau_contrail: npt.NDArray[np.floating],
+    tau_cirrus: npt.NDArray[np.floating],
+    sd0: npt.NDArray[np.floating],
+    sdr: npt.NDArray[np.floating],
+    rsr: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""
     Calculate shortwave differential heating rate.
 
@@ -244,24 +244,24 @@ def differential_heating_rate_shortwave(
 
     Parameters
     ----------
-    cp_contrail : npt.NDArray[np.float64]
+    cp_contrail : npt.NDArray[np.floating]
         Contrail heat capacity per unit length and width, [:math:`J K^{-1} m^{-2}`]
-    r_ice_vol_um : npt.NDArray[np.float64]
+    r_ice_vol_um : npt.NDArray[np.floating]
         Ice particle volume mean radius, [:math:`\mu m`]
-    tau_contrail : npt.NDArray[np.float64]
+    tau_contrail : npt.NDArray[np.floating]
         Contrail optical depth for each waypoint
-    tau_cirrus : npt.NDArray[np.float64]
+    tau_cirrus : npt.NDArray[np.floating]
         Optical depth of numerical weather prediction (NWP) cirrus above the contrail
-    sd0 : npt.NDArray[np.float64]
+    sd0 : npt.NDArray[np.floating]
         Solar constant, [:math:`W m^{-2}`]
-    sdr : npt.NDArray[np.float64]
+    sdr : npt.NDArray[np.floating]
         Solar direct radiation, [:math:`W m^{-2}`]
-    rsr : npt.NDArray[np.float64]
+    rsr : npt.NDArray[np.floating]
         Reflected solar radiation, [:math:`W m^{-2}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Shortwave component of the differential heating rate, [:math:`K s^{-1}`]
     """
     # short circuit if no waypoints have sdr > 0
@@ -286,14 +286,14 @@ def differential_heating_rate_shortwave(
 
 
 def differential_heating_rate_longwave(
-    air_temperature: npt.NDArray[np.float64],
-    rhi: npt.NDArray[np.float64],
-    cp_contrail: npt.NDArray[np.float64],
-    r_ice_vol_um: npt.NDArray[np.float64],
-    tau_contrail: npt.NDArray[np.float64],
-    tau_cirrus: npt.NDArray[np.float64],
-    olr: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    air_temperature: npt.NDArray[np.floating],
+    rhi: npt.NDArray[np.floating],
+    cp_contrail: npt.NDArray[np.floating],
+    r_ice_vol_um: npt.NDArray[np.floating],
+    tau_contrail: npt.NDArray[np.floating],
+    tau_cirrus: npt.NDArray[np.floating],
+    olr: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""
     Calculate longwave differential heating rate.
 
@@ -305,24 +305,24 @@ def differential_heating_rate_longwave(
 
     Parameters
     ----------
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         Ambient temperature at each waypoint, [:math:`K`]
-    rhi : npt.NDArray[np.float64]
+    rhi : npt.NDArray[np.floating]
         Relative humidity with respect to ice at each waypoint
-    cp_contrail : npt.NDArray[np.float64]
+    cp_contrail : npt.NDArray[np.floating]
         Contrail heat capacity per unit length and width, [:math:`J K^{-1} m^{-2}`]
-    r_ice_vol_um : npt.NDArray[np.float64]
+    r_ice_vol_um : npt.NDArray[np.floating]
         Ice particle volume mean radius, [:math:`\mu m`]
-    tau_contrail : npt.NDArray[np.float64]
+    tau_contrail : npt.NDArray[np.floating]
         Contrail optical depth for each waypoint
-    tau_cirrus : npt.NDArray[np.float64]
+    tau_cirrus : npt.NDArray[np.floating]
         Optical depth of numerical weather prediction (NWP) cirrus above the contrail
-    olr : npt.NDArray[np.float64]
+    olr : npt.NDArray[np.floating]
         Outgoing longwave radiation at each waypoint, [:math:`W m^{-2}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Longwave component of the differential heating rate, [:math:`K s^{-1}`]
     """
     cool = RAD_HEAT.dsigma * air_temperature**RAD_HEAT.dak
@@ -341,18 +341,18 @@ def differential_heating_rate_longwave(
 
 
 def heating_rate(
-    air_temperature: npt.NDArray[np.float64],
-    rhi: npt.NDArray[np.float64],
-    rho_air: npt.NDArray[np.float64],
-    r_ice_vol: npt.NDArray[np.float64],
-    depth_eff: npt.NDArray[np.float64],
-    tau_contrail: npt.NDArray[np.float64],
-    tau_cirrus: npt.NDArray[np.float64],
-    sd0: npt.NDArray[np.float64],
-    sdr: npt.NDArray[np.float64],
-    rsr: npt.NDArray[np.float64],
-    olr: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    air_temperature: npt.NDArray[np.floating],
+    rhi: npt.NDArray[np.floating],
+    rho_air: npt.NDArray[np.floating],
+    r_ice_vol: npt.NDArray[np.floating],
+    depth_eff: npt.NDArray[np.floating],
+    tau_contrail: npt.NDArray[np.floating],
+    tau_cirrus: npt.NDArray[np.floating],
+    sd0: npt.NDArray[np.floating],
+    sdr: npt.NDArray[np.floating],
+    rsr: npt.NDArray[np.floating],
+    olr: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Calculate the heating rate affecting the contrail plume.
 
@@ -360,32 +360,32 @@ def heating_rate(
 
     Parameters
     ----------
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         Ambient temperature at each waypoint, [:math:`K`]
-    rhi : npt.NDArray[np.float64]
+    rhi : npt.NDArray[np.floating]
         Relative humidity with respect to ice at each waypoint
-    rho_air : npt.NDArray[np.float64]
+    rho_air : npt.NDArray[np.floating]
         Density of air for each waypoint, [:math:`kg m^{-3}`]
-    r_ice_vol : npt.NDArray[np.float64]
+    r_ice_vol : npt.NDArray[np.floating]
         Ice particle volume mean radius, [:math:`m`]
-    depth_eff : npt.NDArray[np.float64]
+    depth_eff : npt.NDArray[np.floating]
         Effective depth of the contrail plume, [:math:`m`]
-    tau_contrail : npt.NDArray[np.float64]
+    tau_contrail : npt.NDArray[np.floating]
         Contrail optical depth for each waypoint
-    tau_cirrus : npt.NDArray[np.float64]
+    tau_cirrus : npt.NDArray[np.floating]
         Optical depth of numerical weather prediction (NWP) cirrus above the contrail
-    sd0 : npt.NDArray[np.float64]
+    sd0 : npt.NDArray[np.floating]
         Solar constant, [:math:`W m^{-2}`]
-    sdr : npt.NDArray[np.float64]
+    sdr : npt.NDArray[np.floating]
         Solar direct radiation, [:math:`W m^{-2}`]
-    rsr : npt.NDArray[np.float64]
+    rsr : npt.NDArray[np.floating]
         Reflected solar radiation, [:math:`W m^{-2}`]
-    olr : npt.NDArray[np.float64]
+    olr : npt.NDArray[np.floating]
         Outgoing longwave radiation at each waypoint, [:math:`W m^{-2}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Heating rate, [:math:`K s^{-1}`]
     """
     r_ice_vol_um = r_ice_vol * 1e6
@@ -400,36 +400,36 @@ def heating_rate(
 
 
 def heating_rate_shortwave(
-    cp_contrail: npt.NDArray[np.float64],
-    r_ice_vol_um: npt.NDArray[np.float64],
-    tau_contrail: npt.NDArray[np.float64],
-    tau_cirrus: npt.NDArray[np.float64],
-    sd0: npt.NDArray[np.float64],
-    sdr: npt.NDArray[np.float64],
-    rsr: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    cp_contrail: npt.NDArray[np.floating],
+    r_ice_vol_um: npt.NDArray[np.floating],
+    tau_contrail: npt.NDArray[np.floating],
+    tau_cirrus: npt.NDArray[np.floating],
+    sd0: npt.NDArray[np.floating],
+    sdr: npt.NDArray[np.floating],
+    rsr: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""Calculate shortwave heating rate.
 
     Parameters
     ----------
-    cp_contrail : npt.NDArray[np.float64]
+    cp_contrail : npt.NDArray[np.floating]
         Contrail heat capacity per unit length and width, [:math:`J K^{-1} m^{-2}`]
-    r_ice_vol_um : npt.NDArray[np.float64]
+    r_ice_vol_um : npt.NDArray[np.floating]
         Ice particle volume mean radius, [:math:`\mu m`]
-    tau_contrail : npt.NDArray[np.float64]
+    tau_contrail : npt.NDArray[np.floating]
         Contrail optical depth for each waypoint
-    tau_cirrus : npt.NDArray[np.float64]
+    tau_cirrus : npt.NDArray[np.floating]
         Optical depth of numerical weather prediction (NWP) cirrus above the contrail
-    sd0 : npt.NDArray[np.float64]
+    sd0 : npt.NDArray[np.floating]
         Solar constant, [:math:`W m^{-2}`]
-    sdr : npt.NDArray[np.float64]
+    sdr : npt.NDArray[np.floating]
         Solar direct radiation, [:math:`W m^{-2}`]
-    rsr : npt.NDArray[np.float64]
+    rsr : npt.NDArray[np.floating]
         Reflected solar radiation, [:math:`W m^{-2}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Shortwave component of heating rate, [:math:`K s^{-1}`]
     """
     # short circuit if no waypoints have sdr > 0
@@ -452,36 +452,36 @@ def heating_rate_shortwave(
 
 
 def heating_rate_longwave(
-    air_temperature: npt.NDArray[np.float64],
-    rhi: npt.NDArray[np.float64],
-    cp_contrail: npt.NDArray[np.float64],
-    r_ice_vol_um: npt.NDArray[np.float64],
-    tau_contrail: npt.NDArray[np.float64],
-    tau_cirrus: npt.NDArray[np.float64],
-    olr: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    air_temperature: npt.NDArray[np.floating],
+    rhi: npt.NDArray[np.floating],
+    cp_contrail: npt.NDArray[np.floating],
+    r_ice_vol_um: npt.NDArray[np.floating],
+    tau_contrail: npt.NDArray[np.floating],
+    tau_cirrus: npt.NDArray[np.floating],
+    olr: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""Calculate longwave heating rate.
 
     Parameters
     ----------
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         Ambient temperature at each waypoint, [:math:`K`]
-    rhi : npt.NDArray[np.float64]
+    rhi : npt.NDArray[np.floating]
         Relative humidity with respect to ice at each waypoint
-    cp_contrail : npt.NDArray[np.float64]
+    cp_contrail : npt.NDArray[np.floating]
         Contrail heat capacity per unit length and width, [:math:`J K^{-1} m^{-2}`]
-    r_ice_vol_um : npt.NDArray[np.float64]
+    r_ice_vol_um : npt.NDArray[np.floating]
         Ice particle volume mean radius, [:math:`\mu m`]
-    tau_contrail : npt.NDArray[np.float64]
+    tau_contrail : npt.NDArray[np.floating]
         Contrail optical depth for each waypoint
-    tau_cirrus : npt.NDArray[np.float64]
+    tau_cirrus : npt.NDArray[np.floating]
         Optical depth of numerical weather prediction (NWP) cirrus above the contrail
-    olr : npt.NDArray[np.float64]
+    olr : npt.NDArray[np.floating]
         Outgoing longwave radiation at each waypoint, [:math:`W m^{-2}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Longwave component of heating rate, [:math:`K s^{-1}`]
     """
     fzlw = np.exp(-(rhi - 0.9) * RAD_HEAT.czlw)
@@ -500,21 +500,21 @@ def heating_rate_longwave(
 
 
 def contrail_heat_capacity(
-    rho_air: npt.NDArray[np.float64], depth_eff: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    rho_air: npt.NDArray[np.floating], depth_eff: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate contrail heat capacity per unit length and width.
 
     Parameters
     ----------
-    rho_air : npt.NDArray[np.float64]
+    rho_air : npt.NDArray[np.floating]
         density of air for each waypoint, [:math:`kg m^{-3}`]
-    depth_eff: npt.NDArray[np.float64]
+    depth_eff: npt.NDArray[np.floating]
         Effective depth of the contrail plume, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Contrail heat capacity per unit length and width, [:math:`J K^{-1} m^{-2}`]
     """
     return depth_eff * rho_air * constants.c_pd
diff --git a/pycontrails/models/cocip/unterstrasser_wake_vortex.py b/pycontrails/models/cocip/unterstrasser_wake_vortex.py
index c03f64e01..3b060abe3 100644
--- a/pycontrails/models/cocip/unterstrasser_wake_vortex.py
+++ b/pycontrails/models/cocip/unterstrasser_wake_vortex.py
@@ -26,15 +26,15 @@
 
 
 def ice_particle_number_survival_fraction(
-    air_temperature: npt.NDArray[np.float64],
-    rhi_0: npt.NDArray[np.float64],
-    ei_h2o: npt.NDArray[np.float64] | float,
-    wingspan: npt.NDArray[np.float64] | float,
-    true_airspeed: npt.NDArray[np.float64],
-    fuel_flow: npt.NDArray[np.float64],
-    aei_n: npt.NDArray[np.float64],
-    z_desc: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    air_temperature: npt.NDArray[np.floating],
+    rhi_0: npt.NDArray[np.floating],
+    ei_h2o: npt.NDArray[np.floating] | float,
+    wingspan: npt.NDArray[np.floating] | float,
+    true_airspeed: npt.NDArray[np.floating],
+    fuel_flow: npt.NDArray[np.floating],
+    aei_n: npt.NDArray[np.floating],
+    z_desc: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""
     Calculate fraction of ice particle number surviving the wake vortex phase and required inputs.
 
@@ -44,27 +44,27 @@ def ice_particle_number_survival_fraction(
 
     Parameters
     ----------
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
-    rhi_0: npt.NDArray[np.float64]
+    rhi_0: npt.NDArray[np.floating]
         Relative humidity with respect to ice at the flight waypoint
-    ei_h2o : npt.NDArray[np.float64] | float
+    ei_h2o : npt.NDArray[np.floating] | float
         Emission index of water vapor, [:math:`kg \ kg^{-1}`]
-    wingspan : npt.NDArray[np.float64] | float
+    wingspan : npt.NDArray[np.floating] | float
         aircraft wingspan, [:math:`m`]
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    fuel_flow : npt.NDArray[np.float64]
+    fuel_flow : npt.NDArray[np.floating]
         Fuel mass flow rate, [:math:`kg s^{-1}`]
-    aei_n : npt.NDArray[np.float64]
+    aei_n : npt.NDArray[np.floating]
         Apparent ice crystal number emissions index at contrail formation, [:math:`kg^{-1}`]
-    z_desc : npt.NDArray[np.float64]
+    z_desc : npt.NDArray[np.floating]
         Final vertical displacement of the wake vortex, ``dz_max`` in :mod:`wake_vortex.py`,
         [:math:`m`].
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Fraction of contrail ice particle number that survive the wake vortex phase.
 
     References
@@ -86,28 +86,28 @@ def ice_particle_number_survival_fraction(
 
 
 def z_total_length_scale(
-    aei_n: npt.NDArray[np.float64],
-    z_atm: npt.NDArray[np.float64],
-    z_emit: npt.NDArray[np.float64],
-    z_desc: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    aei_n: npt.NDArray[np.floating],
+    z_atm: npt.NDArray[np.floating],
+    z_emit: npt.NDArray[np.floating],
+    z_desc: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Calculate the total length-scale effect of the wake vortex downwash.
 
     Parameters
     ----------
-    aei_n : npt.NDArray[np.float64]
+    aei_n : npt.NDArray[np.floating]
         Apparent ice crystal number emissions index at contrail formation, [:math:`kg^{-1}`]
-    z_atm : npt.NDArray[np.float64]
+    z_atm : npt.NDArray[np.floating]
         Length-scale effect of ambient supersaturation on the ice crystal mass budget, [:math:`m`]
-    z_emit : npt.NDArray[np.float64]
+    z_emit : npt.NDArray[np.floating]
         Length-scale effect of water vapour emissions on the ice crystal mass budget, [:math:`m`]
-    z_desc : npt.NDArray[np.float64]
+    z_desc : npt.NDArray[np.floating]
         Final vertical displacement of the wake vortex, `dz_max` in `wake_vortex.py`, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Total length-scale effect of the wake vortex downwash, [:math:`m`]
     """
     alpha_base = (aei_n / 2.8e14) ** (-0.18)
@@ -121,25 +121,25 @@ def z_total_length_scale(
 
 
 def z_atm_length_scale(
-    air_temperature: npt.NDArray[np.float64],
-    rhi_0: npt.NDArray[np.float64],
+    air_temperature: npt.NDArray[np.floating],
+    rhi_0: npt.NDArray[np.floating],
     *,
     n_iter: int = 10,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Calculate the length-scale effect of ambient supersaturation on the ice crystal mass budget.
 
     Parameters
     ----------
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         Ambient temperature for each waypoint, [:math:`K`].
-    rhi_0 : npt.NDArray[np.float64]
+    rhi_0 : npt.NDArray[np.floating]
         Relative humidity with respect to ice at the flight waypoint.
     n_iter : int
         Number of iterations, set to 10 as default where ``z_atm`` is accurate to within +-1 m.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         The effect of the ambient supersaturation on the ice crystal mass budget,
         provided as a length scale equivalent, [:math:`m`].
 
@@ -174,28 +174,28 @@ def z_atm_length_scale(
 
 
 def emitted_water_vapour_concentration(
-    ei_h2o: npt.NDArray[np.float64] | float,
-    wingspan: npt.NDArray[np.float64] | float,
-    true_airspeed: npt.NDArray[np.float64],
-    fuel_flow: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    ei_h2o: npt.NDArray[np.floating] | float,
+    wingspan: npt.NDArray[np.floating] | float,
+    true_airspeed: npt.NDArray[np.floating],
+    fuel_flow: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""
     Calculate aircraft-emitted water vapour concentration in the plume.
 
     Parameters
     ----------
-    ei_h2o : npt.NDArray[np.float64] | float
+    ei_h2o : npt.NDArray[np.floating] | float
         Emission index of water vapor, [:math:`kg \ kg^{-1}`]
-    wingspan : npt.NDArray[np.float64] | float
+    wingspan : npt.NDArray[np.floating] | float
         aircraft wingspan, [:math:`m`]
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    fuel_flow : npt.NDArray[np.float64]
+    fuel_flow : npt.NDArray[np.floating]
         Fuel mass flow rate, [:math:`kg s^{-1}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Aircraft-emitted water vapour concentration in the plume, [:math:`kg m^{-3}`]
 
     Notes
@@ -208,22 +208,25 @@ def emitted_water_vapour_concentration(
 
 
 def z_emit_length_scale(
-    rho_emit: npt.NDArray[np.float64], air_temperature: npt.NDArray[np.float64], *, n_iter: int = 10
-) -> npt.NDArray[np.float64]:
+    rho_emit: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+    *,
+    n_iter: int = 10,
+) -> npt.NDArray[np.floating]:
     """Calculate the length-scale effect of water vapour emissions on the ice crystal mass budget.
 
     Parameters
     ----------
-    rho_emit : npt.NDArray[np.float64] | float
+    rho_emit : npt.NDArray[np.floating] | float
         Aircraft-emitted water vapour concentration in the plume, [:math:`kg m^{-3}`]
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
     n_iter : int
         Number of iterations, set to 10 as default where ``z_emit`` is accurate to within +-1 m.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         The effect of the aircraft water vapour emission on the ice crystal mass budget,
         provided as a length scale equivalent, [:math:`m`]
 
@@ -250,12 +253,12 @@ def z_emit_length_scale(
     return 0.5 * (z_1 + z_2)
 
 
-def plume_area(wingspan: npt.NDArray[np.float64] | float) -> npt.NDArray[np.float64] | float:
+def plume_area(wingspan: npt.NDArray[np.floating] | float) -> npt.NDArray[np.floating] | float:
     """Calculate area of the wake-vortex plume.
 
     Parameters
     ----------
-    wingspan : npt.NDArray[np.float64] | float
+    wingspan : npt.NDArray[np.floating] | float
         aircraft wingspan, [:math:`m`]
 
     Returns
@@ -272,33 +275,33 @@ def plume_area(wingspan: npt.NDArray[np.float64] | float) -> npt.NDArray[np.floa
 
 
 def z_desc_length_scale(
-    wingspan: npt.NDArray[np.float64] | float,
-    air_temperature: npt.NDArray[np.float64],
-    air_pressure: npt.NDArray[np.float64],
-    true_airspeed: npt.NDArray[np.float64],
-    aircraft_mass: npt.NDArray[np.float64],
-    dT_dz: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    wingspan: npt.NDArray[np.floating] | float,
+    air_temperature: npt.NDArray[np.floating],
+    air_pressure: npt.NDArray[np.floating],
+    true_airspeed: npt.NDArray[np.floating],
+    aircraft_mass: npt.NDArray[np.floating],
+    dT_dz: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """Calculate the final vertical displacement of the wake vortex.
 
     Parameters
     ----------
-    wingspan : npt.NDArray[np.float64] | float
+    wingspan : npt.NDArray[np.floating] | float
         aircraft wingspan, [:math:`m`]
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         pressure altitude at each waypoint, [:math:`Pa`]
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    aircraft_mass : npt.NDArray[np.float64] | float
+    aircraft_mass : npt.NDArray[np.floating] | float
         aircraft mass for each waypoint, [:math:`kg`]
-    dT_dz : npt.NDArray[np.float64]
+    dT_dz : npt.NDArray[np.floating]
         potential temperature gradient, [:math:`K m^{-1}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Final vertical displacement of the wake vortex, [:math:`m`]
 
     Notes
@@ -313,30 +316,30 @@ def z_desc_length_scale(
 
 
 def _initial_wake_vortex_circulation(
-    wingspan: npt.NDArray[np.float64] | float,
-    air_temperature: npt.NDArray[np.float64],
-    air_pressure: npt.NDArray[np.float64],
-    true_airspeed: npt.NDArray[np.float64],
-    aircraft_mass: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    wingspan: npt.NDArray[np.floating] | float,
+    air_temperature: npt.NDArray[np.floating],
+    air_pressure: npt.NDArray[np.floating],
+    true_airspeed: npt.NDArray[np.floating],
+    aircraft_mass: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """Calculate initial wake vortex circulation.
 
     Parameters
     ----------
-    wingspan : npt.NDArray[np.float64] | float
+    wingspan : npt.NDArray[np.floating] | float
         aircraft wingspan, [:math:`m`]
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         pressure altitude at each waypoint, [:math:`Pa`]
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    aircraft_mass : npt.NDArray[np.float64] | float
+    aircraft_mass : npt.NDArray[np.floating] | float
         aircraft mass for each waypoint, [:math:`kg`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Initial wake vortex circulation, [:math:`m^{2} s^{-1}`]
 
     Notes
@@ -350,19 +353,19 @@ def _initial_wake_vortex_circulation(
 
 
 def _survival_fraction_from_length_scale(
-    z_total: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    z_total: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Calculate fraction of ice particle number surviving the wake vortex phase.
 
     Parameters
     ----------
-    z_total : npt.NDArray[np.float64]
+    z_total : npt.NDArray[np.floating]
         Total length-scale effect of the wake vortex downwash, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Fraction of ice particle number surviving the wake vortex phase
     """
     f_surv = 0.45 + (1.19 / np.pi) * np.arctan(-1.35 + (z_total / 100.0))
@@ -371,23 +374,23 @@ def _survival_fraction_from_length_scale(
 
 
 def initial_contrail_depth(
-    z_desc: npt.NDArray[np.float64],
-    f_surv: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    z_desc: npt.NDArray[np.floating],
+    f_surv: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """Calculate initial contrail depth using :cite:`unterstrasserPropertiesYoungContrails2016`.
 
     Parameters
     ----------
-    z_desc : npt.NDArray[np.float64]
+    z_desc : npt.NDArray[np.floating]
         Final vertical displacement of the wake vortex, ``dz_max`` in :mod:`wake_vortex.py`,
         [:math:`m`].
-    f_surv : npt.NDArray[np.float64]
+    f_surv : npt.NDArray[np.floating]
         Fraction of contrail ice particle number that survive the wake vortex phase.
         See :func:`ice_particle_survival_fraction`.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Initial contrail depth, [:math:`m`]
 
     Notes
diff --git a/pycontrails/models/cocip/wake_vortex.py b/pycontrails/models/cocip/wake_vortex.py
index c0a26d611..86ad52531 100644
--- a/pycontrails/models/cocip/wake_vortex.py
+++ b/pycontrails/models/cocip/wake_vortex.py
@@ -31,43 +31,43 @@
 
 
 def max_downward_displacement(
-    wingspan: npt.NDArray[np.float64] | float,
-    true_airspeed: npt.NDArray[np.float64],
-    aircraft_mass: npt.NDArray[np.float64] | float,
-    air_temperature: npt.NDArray[np.float64],
-    dT_dz: npt.NDArray[np.float64],
-    ds_dz: npt.NDArray[np.float64],
-    air_pressure: npt.NDArray[np.float64],
+    wingspan: npt.NDArray[np.floating] | float,
+    true_airspeed: npt.NDArray[np.floating],
+    aircraft_mass: npt.NDArray[np.floating] | float,
+    air_temperature: npt.NDArray[np.floating],
+    dT_dz: npt.NDArray[np.floating],
+    ds_dz: npt.NDArray[np.floating],
+    air_pressure: npt.NDArray[np.floating],
     effective_vertical_resolution: float,
-    wind_shear_enhancement_exponent: float | npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    wind_shear_enhancement_exponent: float | npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Calculate the maximum contrail downward displacement after the wake vortex phase.
 
     Parameters
     ----------
-    wingspan : npt.NDArray[np.float64] | float
+    wingspan : npt.NDArray[np.floating] | float
         aircraft wingspan, [:math:`m`]
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    aircraft_mass : npt.NDArray[np.float64] | float
+    aircraft_mass : npt.NDArray[np.floating] | float
         aircraft mass for each waypoint, [:math:`kg`]
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
-    dT_dz : npt.NDArray[np.float64]
+    dT_dz : npt.NDArray[np.floating]
         potential temperature gradient, [:math:`K m^{-1}`]
-    ds_dz : npt.NDArray[np.float64]
+    ds_dz : npt.NDArray[np.floating]
         Difference in wind speed over dz in the atmosphere, [:math:`m s^{-1} / m`]
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         pressure altitude at each waypoint, [:math:`Pa`]
     effective_vertical_resolution: float
         Passed through to :func:`wind_shear.wind_shear_enhancement_factor`, [:math:`m`]
-    wind_shear_enhancement_exponent: npt.NDArray[np.float64] | float
+    wind_shear_enhancement_exponent: npt.NDArray[np.floating] | float
         Passed through to :func:`wind_shear.wind_shear_enhancement_factor`
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Max contrail downward displacement after the wake vortex phase, [:math:`m`]
 
     References
@@ -107,28 +107,28 @@ def max_downward_displacement(
 
 
 def effective_time_scale(
-    wingspan: npt.NDArray[np.float64] | float,
-    true_airspeed: npt.NDArray[np.float64],
-    aircraft_mass: npt.NDArray[np.float64] | float,
-    rho_air: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    wingspan: npt.NDArray[np.floating] | float,
+    true_airspeed: npt.NDArray[np.floating],
+    aircraft_mass: npt.NDArray[np.floating] | float,
+    rho_air: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""
     Calculate the effective time scale of the wake vortex.
 
     Parameters
     ----------
-    wingspan : npt.NDArray[np.float64]
+    wingspan : npt.NDArray[np.floating]
         aircraft wingspan, [:math:`m`]
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m \ s^{-1}`]
-    aircraft_mass : npt.NDArray[np.float64]
+    aircraft_mass : npt.NDArray[np.floating]
         aircraft mass for each waypoint, [:math:`kg`]
-    rho_air : npt.NDArray[np.float64]
+    rho_air : npt.NDArray[np.floating]
         density of air for each waypoint, [:math:`kg \ m^{-3}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Wake vortex effective time scale, [:math:`s`]
 
     Notes
@@ -144,31 +144,31 @@ def effective_time_scale(
 
 
 def downward_displacement_strongly_stratified(
-    wingspan: npt.NDArray[np.float64] | float,
-    true_airspeed: npt.NDArray[np.float64],
-    aircraft_mass: npt.NDArray[np.float64] | float,
-    rho_air: npt.NDArray[np.float64],
-    n_bv: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    wingspan: npt.NDArray[np.floating] | float,
+    true_airspeed: npt.NDArray[np.floating],
+    aircraft_mass: npt.NDArray[np.floating] | float,
+    rho_air: npt.NDArray[np.floating],
+    n_bv: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Calculate the maximum contrail downward displacement under strongly stratified conditions.
 
     Parameters
     ----------
-    wingspan : npt.NDArray[np.float64] | float
+    wingspan : npt.NDArray[np.floating] | float
         aircraft wingspan, [:math:`m`]
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    aircraft_mass : npt.NDArray[np.float64] | float
+    aircraft_mass : npt.NDArray[np.floating] | float
         aircraft mass for each waypoint, [:math:`kg`]
-    rho_air : npt.NDArray[np.float64]
+    rho_air : npt.NDArray[np.floating]
         density of air for each waypoint, [:math:`kg m^{-3}`]
-    n_bv : npt.NDArray[np.float64]
+    n_bv : npt.NDArray[np.floating]
         Brunt-Vaisaila frequency, [:math:`s^{-1}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Maximum contrail downward displacement, strongly stratified conditions, [:math:`m`]
 
     Notes
@@ -184,46 +184,46 @@ def downward_displacement_strongly_stratified(
 
 
 def downward_displacement_weakly_stratified(
-    wingspan: npt.NDArray[np.float64] | float,
-    true_airspeed: npt.NDArray[np.float64],
-    aircraft_mass: npt.NDArray[np.float64] | float,
-    rho_air: npt.NDArray[np.float64],
-    n_bv: npt.NDArray[np.float64],
-    dz_max_strong: npt.NDArray[np.float64],
-    ds_dz: npt.NDArray[np.float64],
-    t_0: npt.NDArray[np.float64],
+    wingspan: npt.NDArray[np.floating] | float,
+    true_airspeed: npt.NDArray[np.floating],
+    aircraft_mass: npt.NDArray[np.floating] | float,
+    rho_air: npt.NDArray[np.floating],
+    n_bv: npt.NDArray[np.floating],
+    dz_max_strong: npt.NDArray[np.floating],
+    ds_dz: npt.NDArray[np.floating],
+    t_0: npt.NDArray[np.floating],
     effective_vertical_resolution: float,
-    wind_shear_enhancement_exponent: npt.NDArray[np.float64] | float,
-) -> npt.NDArray[np.float64]:
+    wind_shear_enhancement_exponent: npt.NDArray[np.floating] | float,
+) -> npt.NDArray[np.floating]:
     """
     Calculate the maximum contrail downward displacement under weakly/stably stratified conditions.
 
     Parameters
     ----------
-    wingspan : npt.NDArray[np.float64] | float
+    wingspan : npt.NDArray[np.floating] | float
         aircraft wingspan, [:math:`m`]
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    aircraft_mass : npt.NDArray[np.float64] | float
+    aircraft_mass : npt.NDArray[np.floating] | float
         aircraft mass for each waypoint, [:math:`kg`]
-    rho_air : npt.NDArray[np.float64]
+    rho_air : npt.NDArray[np.floating]
         density of air for each waypoint, [:math:`kg m^{-3}`]
-    n_bv : npt.NDArray[np.float64]
+    n_bv : npt.NDArray[np.floating]
         Brunt-Vaisaila frequency, [:math:`s^{-1}`]
-    dz_max_strong : npt.NDArray[np.float64]
+    dz_max_strong : npt.NDArray[np.floating]
         Max contrail downward displacement under strongly stratified conditions, [:math:`m`]
-    ds_dz : npt.NDArray[np.float64]
+    ds_dz : npt.NDArray[np.floating]
         Difference in wind speed over dz in the atmosphere, [:math:`m s^{-1} / m`]
-    t_0 : npt.NDArray[np.float64]
+    t_0 : npt.NDArray[np.floating]
         Wake vortex effective time scale, [:math:`s`]
     effective_vertical_resolution: float
         Passed through to :func:`wind_shear.wind_shear_enhancement_factor`, [:math:`m`]
-    wind_shear_enhancement_exponent: npt.NDArray[np.float64] | float
+    wind_shear_enhancement_exponent: npt.NDArray[np.floating] | float
         Passed through to :func:`wind_shear.wind_shear_enhancement_factor`
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Maximum contrail downward displacement, weakly/stably stratified conditions, [:math:`m`]
 
     Notes
@@ -248,28 +248,28 @@ def downward_displacement_weakly_stratified(
 
 
 def wake_vortex_separation(
-    wingspan: npt.NDArray[np.float64] | float,
-) -> npt.NDArray[np.float64] | float:
+    wingspan: npt.NDArray[np.floating] | float,
+) -> npt.NDArray[np.floating] | float:
     """
     Calculate the wake vortex separation.
 
     Parameters
     ----------
-    wingspan : npt.NDArray[np.float64] | float
+    wingspan : npt.NDArray[np.floating] | float
         aircraft wingspan, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         wake vortex separation, [:math:`m`]
     """
     return (np.pi * wingspan) / 4.0
 
 
 def turbulent_kinetic_energy_dissipation_rate(
-    ds_dz: npt.NDArray[np.float64],
-    shear_enhancement_factor: npt.NDArray[np.float64] | float = 1.0,
-) -> npt.NDArray[np.float64]:
+    ds_dz: npt.NDArray[np.floating],
+    shear_enhancement_factor: npt.NDArray[np.floating] | float = 1.0,
+) -> npt.NDArray[np.floating]:
     """
     Calculate the turbulent kinetic energy dissipation rate (epsilon).
 
@@ -277,14 +277,14 @@ def turbulent_kinetic_energy_dissipation_rate(
 
     Parameters
     ----------
-    ds_dz : npt.NDArray[np.float64]
+    ds_dz : npt.NDArray[np.floating]
         Difference in wind speed over dz in the atmosphere, [:math:`m s^{-1} / m`]
-    shear_enhancement_factor : npt.NDArray[np.float64] | float
+    shear_enhancement_factor : npt.NDArray[np.floating] | float
         Multiplication factor to enhance the wind shear
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         turbulent kinetic energy dissipation rate, [:math:`m^{2} s^{-3}`]
 
     Notes
@@ -303,31 +303,31 @@ def turbulent_kinetic_energy_dissipation_rate(
 
 
 def normalized_dissipation_rate(
-    epsilon: npt.NDArray[np.float64],
-    wingspan: npt.NDArray[np.float64] | float,
-    true_airspeed: npt.NDArray[np.float64],
-    aircraft_mass: npt.NDArray[np.float64] | float,
-    rho_air: npt.NDArray[np.float64] | float,
-) -> npt.NDArray[np.float64]:
+    epsilon: npt.NDArray[np.floating],
+    wingspan: npt.NDArray[np.floating] | float,
+    true_airspeed: npt.NDArray[np.floating],
+    aircraft_mass: npt.NDArray[np.floating] | float,
+    rho_air: npt.NDArray[np.floating] | float,
+) -> npt.NDArray[np.floating]:
     """
     Calculate the normalized dissipation rate of the sinking wake vortex.
 
     Parameters
     ----------
-    epsilon: npt.NDArray[np.float64]
+    epsilon: npt.NDArray[np.floating]
         turbulent kinetic energy dissipation rate, [:math:`m^{2} s^{-3}`]
-    wingspan : npt.NDArray[np.float64] | float
+    wingspan : npt.NDArray[np.floating] | float
         aircraft wingspan, [:math:`m`]
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    aircraft_mass : npt.NDArray[np.float64]
+    aircraft_mass : npt.NDArray[np.floating]
         aircraft mass for each waypoint, [:math:`kg`]
-    rho_air : npt.NDArray[np.float64]
+    rho_air : npt.NDArray[np.floating]
         density of air for each waypoint, [:math:`kg m^{-3}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Normalized dissipation rate of the sinking wake vortex
 
     Notes
@@ -353,44 +353,44 @@ def normalized_dissipation_rate(
 
 
 def initial_contrail_width(
-    wingspan: npt.NDArray[np.float64] | float, dz_max: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    wingspan: npt.NDArray[np.floating] | float, dz_max: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate the initial contrail width.
 
     Parameters
     ----------
-    wingspan : npt.NDArray[np.float64] | float
+    wingspan : npt.NDArray[np.floating] | float
         aircraft wingspan, [:math:`m`]
-    dz_max : npt.NDArray[np.float64]
+    dz_max : npt.NDArray[np.floating]
         Max contrail downward displacement after the wake vortex phase, [:math:`m`]
         Only the size of this array is used; the values are ignored.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Initial contrail width, [:math:`m`]
     """
     return np.full_like(dz_max, np.pi / 4) * wingspan
 
 
 def initial_contrail_depth(
-    dz_max: npt.NDArray[np.float64], initial_wake_vortex_depth: float | npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    dz_max: npt.NDArray[np.floating], initial_wake_vortex_depth: float | npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate the initial contrail depth.
 
     Parameters
     ----------
-    dz_max : npt.NDArray[np.float64]
+    dz_max : npt.NDArray[np.floating]
         Max contrail downward displacement after the wake vortex phase, [:math:`m`]
-    initial_wake_vortex_depth : float | npt.NDArray[np.float64]
+    initial_wake_vortex_depth : float | npt.NDArray[np.floating]
         Initial wake vortex depth scaling factor.
         Denoted `C_D0` in eq (14) in :cite:`schumannContrailCirrusPrediction2012`.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Initial contrail depth, [:math:`m`]
     """
     return dz_max * initial_wake_vortex_depth
diff --git a/pycontrails/models/cocip/wind_shear.py b/pycontrails/models/cocip/wind_shear.py
index 281d9364a..8ad67cad7 100644
--- a/pycontrails/models/cocip/wind_shear.py
+++ b/pycontrails/models/cocip/wind_shear.py
@@ -9,10 +9,10 @@
 
 
 def wind_shear_enhancement_factor(
-    contrail_depth: npt.NDArray[np.float64],
-    effective_vertical_resolution: float | npt.NDArray[np.float64],
-    wind_shear_enhancement_exponent: float | npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    contrail_depth: npt.NDArray[np.floating],
+    effective_vertical_resolution: float | npt.NDArray[np.floating],
+    wind_shear_enhancement_exponent: float | npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""Calculate the multiplication factor to enhance the wind shear based on contrail depth.
 
     This factor accounts for any subgrid-scale that is not captured by the resolution
@@ -20,18 +20,18 @@ def wind_shear_enhancement_factor(
 
     Parameters
     ----------
-    contrail_depth : npt.NDArray[np.float64]
+    contrail_depth : npt.NDArray[np.floating]
         Contrail depth , [:math:`m`]. Expected to be positive and bounded away from 0.
-    effective_vertical_resolution : float | npt.NDArray[np.float64]
+    effective_vertical_resolution : float | npt.NDArray[np.floating]
         Vertical resolution of met data , [:math:`m`]
-    wind_shear_enhancement_exponent : float | npt.NDArray[np.float64]
+    wind_shear_enhancement_exponent : float | npt.NDArray[np.floating]
         Exponent used in calculation. Expected to be nonnegative.
         Discussed in paragraphs following eq. (39) in Schumann 2012 and referenced as `n`.
         When this parameter is 0, no enhancement occurs.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Wind shear enhancement factor
 
     Notes
diff --git a/pycontrails/models/cocipgrid/cocip_grid.py b/pycontrails/models/cocipgrid/cocip_grid.py
index 28656549e..a4df739df 100644
--- a/pycontrails/models/cocipgrid/cocip_grid.py
+++ b/pycontrails/models/cocipgrid/cocip_grid.py
@@ -806,10 +806,10 @@ def _get_source_param_override(self, key: str, vector: GeoVectorDataset) -> Any:
 
     @staticmethod
     def create_source(
-        level: npt.NDArray[np.float64] | list[float] | float,
+        level: npt.NDArray[np.floating] | list[float] | float,
         time: npt.NDArray[np.datetime64] | list[np.datetime64] | np.datetime64,
-        longitude: npt.NDArray[np.float64] | list[float] | None = None,
-        latitude: npt.NDArray[np.float64] | list[float] | None = None,
+        longitude: npt.NDArray[np.floating] | list[float] | None = None,
+        latitude: npt.NDArray[np.floating] | list[float] | None = None,
         lon_step: float = 1.0,
         lat_step: float = 1.0,
     ) -> MetDataset:
@@ -821,7 +821,7 @@ def create_source(
 
         Parameters
         ----------
-        level : level: npt.NDArray[np.float64] | list[float] | float
+        level : level: npt.NDArray[np.floating] | list[float] | float
             Pressure levels for gridded cocip.
             To avoid interpolating outside of the passed ``met`` and ``rad`` data, this
             parameter should avoid the extreme values of the ``met`` and `rad` levels.
@@ -829,7 +829,7 @@ def create_source(
             ``met.data['level'].values[1: -1]``.
         time: npt.NDArray[np.datetime64 | list[np.datetime64] | np.datetime64,
             One or more time values for gridded cocip.
-        longitude, latitude : npt.NDArray[np.float64] | list[float], optional
+        longitude, latitude : npt.NDArray[np.floating] | list[float], optional
             Longitude and latitude arrays, by default None. If not specified, values of
             ``lon_step`` and ``lat_step`` are used to define ``longitude`` and ``latitude``.
         lon_step, lat_step : float, optional
@@ -2173,7 +2173,7 @@ def _aggregate_ef_summary(vector_list: list[VectorDataset]) -> VectorDataset | N
 
 def result_to_metdataset(
     result: VectorDataset | None,
-    verbose_dict: dict[str, npt.NDArray[np.float64]],
+    verbose_dict: dict[str, npt.NDArray[np.floating]],
     source: MetDataset,
     nominal_segment_length: float,
     attrs: dict[str, str],
@@ -2185,7 +2185,7 @@ def result_to_metdataset(
     result : VectorDataset | None
         Aggregated data arising from contrail evolution. Expected to contain keys:
         ``index``, ``age``, ``ef``.
-    verbose_dict : dict[str, npt.NDArray[np.float64]]:
+    verbose_dict : dict[str, npt.NDArray[np.floating]]:
         Verbose outputs to attach to results.
     source : MetDataset
         :attr:`CocipGrid.`source` data on which to attach results.
@@ -2244,9 +2244,9 @@ def result_to_metdataset(
 
 def result_merge_source(
     result: VectorDataset | None,
-    verbose_dict: dict[str, npt.NDArray[np.float64]],
+    verbose_dict: dict[str, npt.NDArray[np.floating]],
     source: GeoVectorDataset,
-    nominal_segment_length: float | npt.NDArray[np.float64],
+    nominal_segment_length: float | npt.NDArray[np.floating],
     attrs: dict[str, str],
 ) -> GeoVectorDataset:
     """Merge ``results`` and ``verbose_dict`` onto ``source``."""
@@ -2282,7 +2282,7 @@ def _concat_verbose_dicts(
     verbose_dicts: list[dict[str, pd.Series]],
     source_size: int,
     verbose_outputs_formation: set[str],
-) -> dict[str, npt.NDArray[np.float64]]:
+) -> dict[str, npt.NDArray[np.floating]]:
     # Concatenate the values and return
     ret: dict[str, np.ndarray] = {}
     for key in verbose_outputs_formation:
@@ -2375,7 +2375,7 @@ def _warn_not_wrap(met: MetDataset) -> None:
         )
 
 
-def _get_uncertainty_params(contrail: VectorDataset) -> dict[str, npt.NDArray[np.float64]]:
+def _get_uncertainty_params(contrail: VectorDataset) -> dict[str, npt.NDArray[np.floating]]:
     """Return uncertainty parameters in ``contrail``.
 
     This function assumes the underlying humidity scaling model is
@@ -2398,7 +2398,7 @@ def _get_uncertainty_params(contrail: VectorDataset) -> dict[str, npt.NDArray[np
 
     Returns
     -------
-    dict[str, npt.NDArray[np.float64]]
+    dict[str, npt.NDArray[np.floating]]
         Dictionary of uncertainty parameters.
     """
     keys = (
diff --git a/pycontrails/models/emissions/black_carbon.py b/pycontrails/models/emissions/black_carbon.py
index dd5959fc2..8f1b94f68 100644
--- a/pycontrails/models/emissions/black_carbon.py
+++ b/pycontrails/models/emissions/black_carbon.py
@@ -14,29 +14,29 @@
 
 
 def mass_emissions_index_fox(
-    air_pressure: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-    true_airspeed: npt.NDArray[np.float64],
-    fuel_flow_per_engine: npt.NDArray[np.float64],
-    thrust_setting: npt.NDArray[np.float64],
+    air_pressure: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+    true_airspeed: npt.NDArray[np.floating],
+    fuel_flow_per_engine: npt.NDArray[np.floating],
+    thrust_setting: npt.NDArray[np.floating],
     pressure_ratio: float,
     *,
     comp_efficiency: float = 0.9,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     r"""
     Calculate the black carbon mass emissions index using the Formation and Oxidation Method (FOX).
 
     Parameters
     ----------
-    air_pressure: npt.NDArray[np.float64]
+    air_pressure: npt.NDArray[np.floating]
         Pressure altitude at each waypoint, [:math:`Pa`]
-    air_temperature: npt.NDArray[np.float64]
+    air_temperature: npt.NDArray[np.floating]
         Ambient temperature for each waypoint, [:math:`K`]
-    true_airspeed: npt.NDArray[np.float64]
+    true_airspeed: npt.NDArray[np.floating]
         True airspeed for each waypoint, [:math:`m s^{-1}`]
-    fuel_flow_per_engine: npt.NDArray[np.float64]
+    fuel_flow_per_engine: npt.NDArray[np.floating]
         Fuel mass flow rate per engine, [:math:`kg s^{-1}`]
-    thrust_setting: npt.NDArray[np.float64]
+    thrust_setting: npt.NDArray[np.floating]
         Engine thrust setting, which is the fuel mass flow rate divided by
         the maximum fuel mass flow rate
     pressure_ratio: float
@@ -46,7 +46,7 @@ def mass_emissions_index_fox(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Black carbon mass emissions index, [:math:`mg \ kg_{fuel}^{-1}`]
 
     References
@@ -102,10 +102,10 @@ def flame_temperature(t_3: ArrayScalarLike) -> ArrayScalarLike:
 
 
 def bc_mass_concentration_fox(
-    fuel_flow: npt.NDArray[np.float64],
-    t_fl: npt.NDArray[np.float64] | float,
-    afr: npt.NDArray[np.float64] | float,
-) -> npt.NDArray[np.float64]:
+    fuel_flow: npt.NDArray[np.floating],
+    t_fl: npt.NDArray[np.floating] | float,
+    afr: npt.NDArray[np.floating] | float,
+) -> npt.NDArray[np.floating]:
     """Calculate the black carbon mass concentration for ground conditions (``c_bc_ref``).
 
     This quantity is computed at the instrument sampling point without correcting
@@ -113,30 +113,30 @@ def bc_mass_concentration_fox(
 
     Parameters
     ----------
-    fuel_flow: npt.NDArray[np.float64]
+    fuel_flow: npt.NDArray[np.floating]
         Fuel mass flow rate, [:math:`kg s^{-1}`]
-    t_fl: npt.NDArray[np.float64] | float
+    t_fl: npt.NDArray[np.floating] | float
         Flame temperature at the combustion chamber, [:math:`K`]
-    afr: npt.NDArray[np.float64] | float
+    afr: npt.NDArray[np.floating] | float
         Air-to-fuel ratio
 
     Returns
     -------
-    npt.NDArray[np.float64]:
+    npt.NDArray[np.floating]:
         Black carbon mass concentration for ground conditions, [:math:`mg m^{-3}`]
     """
     return fuel_flow * (356 * np.exp(-6390 / t_fl) - 608 * afr * np.exp(-19778 / t_fl))
 
 
 def bc_mass_concentration_cruise_fox(
-    c_bc_ref: npt.NDArray[np.float64],
-    t_fl_cru: npt.NDArray[np.float64],
-    t_fl_ref: npt.NDArray[np.float64] | float,
-    p_3_cru: npt.NDArray[np.float64],
-    p_3_ref: npt.NDArray[np.float64] | float,
-    afr_cru: npt.NDArray[np.float64],
-    afr_ref: npt.NDArray[np.float64] | float,
-) -> npt.NDArray[np.float64]:
+    c_bc_ref: npt.NDArray[np.floating],
+    t_fl_cru: npt.NDArray[np.floating],
+    t_fl_ref: npt.NDArray[np.floating] | float,
+    p_3_cru: npt.NDArray[np.floating],
+    p_3_ref: npt.NDArray[np.floating] | float,
+    afr_cru: npt.NDArray[np.floating],
+    afr_ref: npt.NDArray[np.floating] | float,
+) -> npt.NDArray[np.floating]:
     """Calculate the black carbon mass concentration for cruise conditions (``c_bc_cru``).
 
     This quantity is computed at the instrument sampling point without correcting
@@ -144,24 +144,24 @@ def bc_mass_concentration_cruise_fox(
 
     Parameters
     ----------
-    c_bc_ref: npt.NDArray[np.float64]
+    c_bc_ref: npt.NDArray[np.floating]
         Black carbon mass concentration at reference conditions, [:math:`mg m^{-3}`]
-    t_fl_cru: npt.NDArray[np.float64]
+    t_fl_cru: npt.NDArray[np.floating]
         Flame temperature at cruise conditions, [:math:`K`]
-    t_fl_ref: npt.NDArray[np.float64] | float
+    t_fl_ref: npt.NDArray[np.floating] | float
         Flame temperature at reference conditions, [:math:`K`]
-    p_3_cru: npt.NDArray[np.float64]
+    p_3_cru: npt.NDArray[np.floating]
         Combustor inlet pressure at cruise conditions, [:math:`Pa`]
-    p_3_ref: npt.NDArray[np.float64] | float
+    p_3_ref: npt.NDArray[np.floating] | float
         Combustor inlet pressure at reference conditions, [:math:`Pa`]
-    afr_cru: npt.NDArray[np.float64]
+    afr_cru: npt.NDArray[np.floating]
         Air-to-fuel ratio at cruise conditions
-    afr_ref: npt.NDArray[np.float64] | float
+    afr_ref: npt.NDArray[np.floating] | float
         Air-to-fuel ratio at reference conditions
 
     Returns
     -------
-    npt.NDArray[np.float64]:
+    npt.NDArray[np.floating]:
         Black carbon mass concentration for cruise conditions, [:math:`mg m^{-3}`]
     """
     scaling_factor = dopelheuer_lecht_scaling_factor(
@@ -176,33 +176,33 @@ def bc_mass_concentration_cruise_fox(
 
 
 def dopelheuer_lecht_scaling_factor(
-    t_fl_cru: npt.NDArray[np.float64],
-    t_fl_ref: npt.NDArray[np.float64] | float,
-    p_3_cru: npt.NDArray[np.float64],
-    p_3_ref: npt.NDArray[np.float64] | float,
-    afr_cru: npt.NDArray[np.float64],
-    afr_ref: npt.NDArray[np.float64] | float,
-) -> npt.NDArray[np.float64]:
+    t_fl_cru: npt.NDArray[np.floating],
+    t_fl_ref: npt.NDArray[np.floating] | float,
+    p_3_cru: npt.NDArray[np.floating],
+    p_3_ref: npt.NDArray[np.floating] | float,
+    afr_cru: npt.NDArray[np.floating],
+    afr_ref: npt.NDArray[np.floating] | float,
+) -> npt.NDArray[np.floating]:
     """Estimate scaling factor to convert the reference BC mass concentration from ground to cruise.
 
     Parameters
     ----------
-    t_fl_cru: npt.NDArray[np.float64]
+    t_fl_cru: npt.NDArray[np.floating]
         Flame temperature at cruise conditions, [:math:`K`]
-    t_fl_ref: npt.NDArray[np.float64] | float
+    t_fl_ref: npt.NDArray[np.floating] | float
         Flame temperature at reference conditions, [:math:`K`]
-    p_3_cru: npt.NDArray[np.float64]
+    p_3_cru: npt.NDArray[np.floating]
         Combustor inlet pressure at cruise conditions, [:math:`Pa`]
-    p_3_ref: npt.NDArray[np.float64] | float
+    p_3_ref: npt.NDArray[np.floating] | float
         Combustor inlet pressure at reference conditions, [:math:`Pa`]
-    afr_cru: npt.NDArray[np.float64]
+    afr_cru: npt.NDArray[np.floating]
         Air-to-fuel ratio at cruise conditions
-    afr_ref: npt.NDArray[np.float64] | float
+    afr_ref: npt.NDArray[np.floating] | float
         Air-to-fuel ratio at reference conditions
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Dopelheuer & Lecht scaling factor
 
     References
@@ -219,17 +219,17 @@ def dopelheuer_lecht_scaling_factor(
 
 
 def mass_emissions_index_imfox(
-    fuel_flow_per_engine: npt.NDArray[np.float64],
-    thrust_setting: npt.NDArray[np.float64],
+    fuel_flow_per_engine: npt.NDArray[np.floating],
+    thrust_setting: npt.NDArray[np.floating],
     fuel_hydrogen: float,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     r"""Calculate the BC mass EI using the "Improved" Formation and Oxidation Method (ImFOX).
 
     Parameters
     ----------
-    fuel_flow_per_engine: npt.NDArray[np.float64]
+    fuel_flow_per_engine: npt.NDArray[np.floating]
         Fuel mass flow rate per engine, [:math:`kg s^{-1}`]
-    thrust_setting: npt.NDArray[np.float64]
+    thrust_setting: npt.NDArray[np.floating]
         Engine thrust setting, which is the fuel mass flow rate divided by the
         maximum fuel mass flow rate
     fuel_hydrogen: float
@@ -237,7 +237,7 @@ def mass_emissions_index_imfox(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Black carbon mass emissions index, [:math:`mg \ kg_{fuel}^{-1}`]
 
     References
@@ -253,20 +253,20 @@ def mass_emissions_index_imfox(
     return bc_mass_emissions_index(c_bc_cru, q_exhaust_cru)
 
 
-def air_to_fuel_ratio_imfox(thrust_setting: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def air_to_fuel_ratio_imfox(thrust_setting: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     """Calculate the air-to-fuel ratio at cruise conditions via Abrahamson's method.
 
     See Eq. (11) in :cite:`abrahamsonPredictiveModelDevelopment2016`.
 
     Parameters
     ----------
-    thrust_setting: npt.NDArray[np.float64]
+    thrust_setting: npt.NDArray[np.floating]
         Engine thrust setting, which is the fuel mass flow rate divided by
         the maximum fuel mass flow rate
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Air-to-fuel ratio at cruise conditions
 
     References
@@ -276,19 +276,19 @@ def air_to_fuel_ratio_imfox(thrust_setting: npt.NDArray[np.float64]) -> npt.NDAr
     return 55.4 - 30.8 * thrust_setting
 
 
-def turbine_inlet_temperature_imfox(afr: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def turbine_inlet_temperature_imfox(afr: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     """Calculate the turbine inlet temperature using Abrahamson's method.
 
     See Eq. (13) in :cite:`abrahamsonPredictiveModelDevelopment2016`.
 
     Parameters
     ----------
-    afr: npt.NDArray[np.float64]
+    afr: npt.NDArray[np.floating]
         air-to-fuel ratio at cruise conditions
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         turbine inlet temperature, [:math:`K`]
 
     References
@@ -299,11 +299,11 @@ def turbine_inlet_temperature_imfox(afr: npt.NDArray[np.float64]) -> npt.NDArray
 
 
 def bc_mass_concentration_imfox(
-    fuel_flow_per_engine: npt.NDArray[np.float64],
-    afr: npt.NDArray[np.float64],
-    t_4: npt.NDArray[np.float64],
+    fuel_flow_per_engine: npt.NDArray[np.floating],
+    afr: npt.NDArray[np.floating],
+    t_4: npt.NDArray[np.floating],
     fuel_hydrogen: float,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Calculate the BC mass concentration for ground and cruise conditions with ImFOX methodology.
 
     This quantity is computed at the instrument sampling point without
@@ -311,18 +311,18 @@ def bc_mass_concentration_imfox(
 
     Parameters
     ----------
-    fuel_flow_per_engine: npt.NDArray[np.float64]
+    fuel_flow_per_engine: npt.NDArray[np.floating]
         fuel mass flow rate per engine, [:math:`kg s^{-1}`]
-    afr: npt.NDArray[np.float64]
+    afr: npt.NDArray[np.floating]
         air-to-fuel ratio
-    t_4: npt.NDArray[np.float64]
+    t_4: npt.NDArray[np.floating]
         turbine inlet temperature, [:math:`K`]
     fuel_hydrogen: float
         percentage of hydrogen mass content in the fuel (13.8% for conventional Jet A-1 fuel)
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Black carbon mass concentration, [:math:`mg m^{-3}`]
     """
     exp_term = np.exp(13.6 - fuel_hydrogen)
@@ -336,18 +336,18 @@ def bc_mass_concentration_imfox(
 # ---------------------------------------------------------
 
 
-def exhaust_gas_volume_per_kg_fuel(afr: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def exhaust_gas_volume_per_kg_fuel(afr: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     """
     Calculate the volume of exhaust gas per mass of fuel burnt.
 
     Parameters
     ----------
-    afr: npt.NDArray[np.float64]
+    afr: npt.NDArray[np.floating]
         Air-to-fuel ratio
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Volume of exhaust gas per mass of fuel burnt, [:math:`m^{3}/kg_{fuel}`]
 
     References
@@ -358,21 +358,21 @@ def exhaust_gas_volume_per_kg_fuel(afr: npt.NDArray[np.float64]) -> npt.NDArray[
 
 
 def bc_mass_emissions_index(
-    c_bc: npt.NDArray[np.float64], q_exhaust: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    c_bc: npt.NDArray[np.floating], q_exhaust: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """
     Calculate the black carbon mass emissions index.
 
     Parameters
     ----------
-    c_bc: npt.NDArray[np.float64]
+    c_bc: npt.NDArray[np.floating]
         Black carbon mass concentration, [:math:`mg m^{-3}`]
-    q_exhaust: npt.NDArray[np.float64]
+    q_exhaust: npt.NDArray[np.floating]
         Volume of exhaust gas per mass of fuel burnt, [:math:`m^{3}/kg_{fuel}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Black carbon mass emissions index, [:math:`mg/kg_{fuel}`]
 
     References
@@ -388,17 +388,17 @@ def bc_mass_emissions_index(
 
 
 def geometric_mean_diameter_sac(
-    air_pressure: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-    true_airspeed: npt.NDArray[np.float64],
-    thrust_setting: npt.NDArray[np.float64],
+    air_pressure: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+    true_airspeed: npt.NDArray[np.floating],
+    thrust_setting: npt.NDArray[np.floating],
     pressure_ratio: float,
     q_fuel: float,
     *,
     comp_efficiency: float = 0.9,
     delta_loss: float = 5.75,
     cruise: bool = True,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     r"""Calculate the BC GMD for singular annular combustor (SAC) engines.
 
     The BC (black carbon) GMD (geometric mean diameter) is estimated using
@@ -407,13 +407,13 @@ def geometric_mean_diameter_sac(
 
     Parameters
     ----------
-    air_pressure: npt.NDArray[np.float64]
+    air_pressure: npt.NDArray[np.floating]
         Pressure altitude at each waypoint, [:math:`Pa`]
-    air_temperature: npt.NDArray[np.float64]
+    air_temperature: npt.NDArray[np.floating]
         Ambient temperature for each waypoint, [:math:`K`]
-    true_airspeed: npt.NDArray[np.float64]
+    true_airspeed: npt.NDArray[np.floating]
         True airspeed for each waypoint, [:math:`m s^{-1}`]
-    thrust_setting: npt.NDArray[np.float64]
+    thrust_setting: npt.NDArray[np.floating]
         Engine thrust setting, which is the fuel mass flow rate divided by the
         maximum fuel mass flow rate
     pressure_ratio: float
@@ -429,7 +429,7 @@ def geometric_mean_diameter_sac(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         black carbon geometric mean diameter, [:math:`nm`]
 
     References
@@ -450,15 +450,15 @@ def geometric_mean_diameter_sac(
 
 
 def number_emissions_index_fractal_aggregates(
-    nvpm_ei_m: npt.NDArray[np.float64],
-    gmd: npt.NDArray[np.float64],
+    nvpm_ei_m: npt.NDArray[np.floating],
+    gmd: npt.NDArray[np.floating],
     *,
-    gsd: float | npt.NDArray[np.float64] = 1.80,
+    gsd: float | npt.NDArray[np.floating] = 1.80,
     rho_bc: float = 1770,
     k_tem: float = 1.621e-5,
     d_tem: float = 0.39,
     d_fm: float = 2.76,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """
     Estimate the black carbon number emission index using the fractal aggregates (FA) model.
 
@@ -467,9 +467,9 @@ def number_emissions_index_fractal_aggregates(
 
     Parameters
     ----------
-    nvpm_ei_m: npt.NDArray[np.float64]
+    nvpm_ei_m: npt.NDArray[np.floating]
         Black carbon mass emissions index, [:math:`kg/kg_{fuel}`]
-    gmd: npt.NDArray[np.float64]
+    gmd: npt.NDArray[np.floating]
         Black carbon geometric mean diameter, [:math:`m`]
     gsd: float
         Black carbon geometric standard deviation (assumed to be 1.80)
@@ -484,7 +484,7 @@ def number_emissions_index_fractal_aggregates(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Black carbon number emissions index, [:math:`kg_{fuel}^{-1}`]
 
     References
@@ -506,9 +506,9 @@ def number_emissions_index_fractal_aggregates(
 
 
 def nvpm_number_ei_pct_reduction_due_to_saf(
-    hydrogen_content: float | npt.NDArray[np.float64],
-    thrust_setting: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    hydrogen_content: float | npt.NDArray[np.floating],
+    thrust_setting: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Adjust nvPM number emissions index to account for the effects of sustainable aviation fuels.
 
@@ -516,13 +516,13 @@ def nvpm_number_ei_pct_reduction_due_to_saf(
     ----------
     hydrogen_content: float
         The percentage of hydrogen mass content in the fuel.
-    thrust_setting: npt.NDArray[np.float64]
+    thrust_setting: npt.NDArray[np.floating]
         Engine thrust setting, where the equivalent fuel mass flow rate per engine at
         sea level, :math:`[0 - 1]`.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Percentage reduction in nvPM number emissions index
 
     References
@@ -537,9 +537,9 @@ def nvpm_number_ei_pct_reduction_due_to_saf(
 
 
 def nvpm_mass_ei_pct_reduction_due_to_saf(
-    hydrogen_content: float | npt.NDArray[np.float64],
-    thrust_setting: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    hydrogen_content: float | npt.NDArray[np.floating],
+    thrust_setting: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Adjust nvPM mass emissions index to account for the effects of sustainable aviation fuels.
 
@@ -550,13 +550,13 @@ def nvpm_mass_ei_pct_reduction_due_to_saf(
     ----------
     hydrogen_content: float
         The percentage of hydrogen mass content in the fuel.
-    thrust_setting: npt.NDArray[np.float64]
+    thrust_setting: npt.NDArray[np.floating]
         Engine thrust setting, where the equivalent fuel mass flow rate per engine at
         sea level, :math:`[0 - 1]`.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Percentage reduction in nvPM number emissions index
 
     References
@@ -571,12 +571,12 @@ def nvpm_mass_ei_pct_reduction_due_to_saf(
 
 
 def _template_saf_reduction(
-    hydrogen_content: float | npt.NDArray[np.float64],
-    thrust_setting: npt.NDArray[np.float64],
+    hydrogen_content: float | npt.NDArray[np.floating],
+    thrust_setting: npt.NDArray[np.floating],
     a0: float,
     a1: float,
     a2: float,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     # Thrust setting cannot be computed when engine data is not provided in
     # the ICAO EDB, so set default to 45% thrust.
     thrust_setting = np.nan_to_num(thrust_setting, nan=0.45)
diff --git a/pycontrails/models/emissions/emissions.py b/pycontrails/models/emissions/emissions.py
index 01349240c..e239de2bf 100644
--- a/pycontrails/models/emissions/emissions.py
+++ b/pycontrails/models/emissions/emissions.py
@@ -369,7 +369,7 @@ def _nvpm_emission_indices(self, engine_uid: str | None) -> None:
 
     def _nvpm_emission_indices_edb(
         self, edb_nvpm: EDBnvpm, fuel: Fuel
-    ) -> tuple[str, npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    ) -> tuple[str, npt.NDArray[np.floating], npt.NDArray[np.floating]]:
         """Calculate emission indices for nvPM mass and number.
 
         This method uses data from the ICAO EDB along with the T4/T2 methodology.
@@ -385,9 +385,9 @@ def _nvpm_emission_indices_edb(
         -------
         nvpm_data_source : str
             Source of nvpm data.
-        nvpm_ei_m : npt.NDArray[np.float64]
+        nvpm_ei_m : npt.NDArray[np.floating]
             Non-volatile particulate matter (nvPM) mass emissions index, [:math:`kg/kg_{fuel}`]
-        nvpm_ei_n : npt.NDArray[np.float64]
+        nvpm_ei_n : npt.NDArray[np.floating]
             Black carbon number emissions index, [:math:`kg_{fuel}^{-1}`]
 
         References
@@ -408,7 +408,7 @@ def _nvpm_emission_indices_edb(
 
     def _nvpm_emission_indices_sac(
         self, edb_gaseous: EDBGaseous, fuel: Fuel
-    ) -> tuple[str, npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    ) -> tuple[str, npt.NDArray[np.floating], npt.NDArray[np.floating]]:
         """Calculate EIs for nvPM mass and number assuming the profile of single annular combustors.
 
         nvPM EI_m is calculated using the FOX and ImFOX methods, while the nvPM EI_n
@@ -425,9 +425,9 @@ def _nvpm_emission_indices_sac(
         -------
         nvpm_data_source : str
             Source of nvpm data.
-        nvpm_ei_m : npt.NDArray[np.float64]
+        nvpm_ei_m : npt.NDArray[np.floating]
             Non-volatile particulate matter (nvPM) mass emissions index, [:math:`kg/kg_{fuel}`]
-        nvpm_ei_n : npt.NDArray[np.float64]
+        nvpm_ei_n : npt.NDArray[np.floating]
             Black carbon number emissions index, [:math:`kg_{fuel}^{-1}`]
 
         References
@@ -467,7 +467,7 @@ def _nvpm_emission_indices_sac(
 
     def _nvpm_emission_indices_constant(
         self,
-    ) -> tuple[str, npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    ) -> tuple[str, npt.NDArray[np.floating], npt.NDArray[np.floating]]:
         """
         Assume constant emission indices for nvPM mass and number.
 
@@ -481,9 +481,9 @@ def _nvpm_emission_indices_constant(
         -------
         nvpm_data_source : str
             Source of nvpm data.
-        nvpm_ei_m : npt.NDArray[np.float64]
+        nvpm_ei_m : npt.NDArray[np.floating]
             Non-volatile particulate matter (nvPM) mass emissions index, [:math:`kg/kg_{fuel}`]
-        nvpm_ei_n : npt.NDArray[np.float64]
+        nvpm_ei_n : npt.NDArray[np.floating]
             Black carbon number emissions index, [:math:`kg_{fuel}^{-1}`]
 
         References
@@ -605,12 +605,12 @@ def _check_edb_nvpm_availability(
 
 def nitrogen_oxide_emissions_index_ffm2(
     edb_gaseous: EDBGaseous,
-    fuel_flow_per_engine: npt.NDArray[np.float64],
-    true_airspeed: npt.NDArray[np.float64],
-    air_pressure: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-    specific_humidity: None | npt.NDArray[np.float64] = None,
-) -> npt.NDArray[np.float64]:
+    fuel_flow_per_engine: npt.NDArray[np.floating],
+    true_airspeed: npt.NDArray[np.floating],
+    air_pressure: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+    specific_humidity: None | npt.NDArray[np.floating] = None,
+) -> npt.NDArray[np.floating]:
     """
     Estimate the nitrogen oxide (NOx) emissions index (EI) using the Fuel Flow Method 2 (FFM2).
 
@@ -618,20 +618,20 @@ def nitrogen_oxide_emissions_index_ffm2(
     ----------
     edb_gaseous : EDBGaseous
         EDB gaseous data
-    fuel_flow_per_engine: npt.NDArray[np.float64]
+    fuel_flow_per_engine: npt.NDArray[np.floating]
         fuel mass flow rate per engine, [:math:`kg s^{-1}`]
-    true_airspeed: npt.NDArray[np.float64]
+    true_airspeed: npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         pressure altitude at each waypoint, [:math:`Pa`]
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
-    specific_humidity: npt.NDArray[np.float64]
+    specific_humidity: npt.NDArray[np.floating]
         specific humidity for each waypoint, [:math:`kg_{H_{2}O}/kg_{air}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Nitrogen oxide emissions index for each waypoint, [:math:`kg_{NO_{X}}/kg_{fuel}`]
     """
     res_nox = ffm2.estimate_nox(
@@ -647,11 +647,11 @@ def nitrogen_oxide_emissions_index_ffm2(
 
 def carbon_monoxide_emissions_index_ffm2(
     edb_gaseous: EDBGaseous,
-    fuel_flow_per_engine: npt.NDArray[np.float64],
-    true_airspeed: npt.NDArray[np.float64],
-    air_pressure: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    fuel_flow_per_engine: npt.NDArray[np.floating],
+    true_airspeed: npt.NDArray[np.floating],
+    air_pressure: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Estimate the carbon monoxide (CO) emissions index (EI) using the Fuel Flow Method 2 (FFM2).
 
@@ -659,18 +659,18 @@ def carbon_monoxide_emissions_index_ffm2(
     ----------
     edb_gaseous : EDBGaseous
         EDB gaseous data
-    fuel_flow_per_engine: npt.NDArray[np.float64]
+    fuel_flow_per_engine: npt.NDArray[np.floating]
         fuel mass flow rate per engine, [:math:`kg s^{-1}`]
-    true_airspeed: npt.NDArray[np.float64]
+    true_airspeed: npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         pressure altitude at each waypoint, [:math:`Pa`]
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Carbon monoxide emissions index for each waypoint, [:math:`kg_{CO}/kg_{fuel}`]
     """
     res_co = ffm2.estimate_ei(
@@ -685,11 +685,11 @@ def carbon_monoxide_emissions_index_ffm2(
 
 def hydrocarbon_emissions_index_ffm2(
     edb_gaseous: EDBGaseous,
-    fuel_flow_per_engine: npt.NDArray[np.float64],
-    true_airspeed: npt.NDArray[np.float64],
-    air_pressure: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    fuel_flow_per_engine: npt.NDArray[np.floating],
+    true_airspeed: npt.NDArray[np.floating],
+    air_pressure: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Estimate the hydrocarbon (HC) emissions index (EI) using the Fuel Flow Method 2 (FFM2).
 
@@ -697,18 +697,18 @@ def hydrocarbon_emissions_index_ffm2(
     ----------
     edb_gaseous : EDBGaseous
         EDB gaseous data
-    fuel_flow_per_engine: npt.NDArray[np.float64]
+    fuel_flow_per_engine: npt.NDArray[np.floating]
         fuel mass flow rate per engine, [:math:`kg s^{-1}`]
-    true_airspeed: npt.NDArray[np.float64]
+    true_airspeed: npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         pressure altitude at each waypoint, [:math:`Pa`]
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Hydrocarbon emissions index for each waypoint, [:math:`kg_{HC}/kg_{fuel}`]
     """
     res_hc = ffm2.estimate_ei(
@@ -723,12 +723,12 @@ def hydrocarbon_emissions_index_ffm2(
 
 def get_nvpm_emissions_index_edb(
     edb_nvpm: EDBnvpm,
-    true_airspeed: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-    air_pressure: npt.NDArray[np.float64],
-    thrust_setting: npt.NDArray[np.float64],
+    true_airspeed: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+    air_pressure: npt.NDArray[np.floating],
+    thrust_setting: npt.NDArray[np.floating],
     q_fuel: float,
-) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
     r"""Calculate nvPM mass emissions index (nvpm_ei_m) and number emissions index (nvpm_ei_n).
 
     Interpolate the non-volatile particulate matter (nvPM) mass and number emissions index from
@@ -741,24 +741,24 @@ def get_nvpm_emissions_index_edb(
     ----------
     edb_nvpm : EDBnvpm
         EDB nvPM data
-    true_airspeed: npt.NDArray[np.float64]
+    true_airspeed: npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    fuel_flow_per_engine: npt.NDArray[np.float64]
+    fuel_flow_per_engine: npt.NDArray[np.floating]
         fuel mass flow rate per engine, [:math:`kg s^{-1}`]
-    air_temperature: npt.NDArray[np.float64]
+    air_temperature: npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
-    air_pressure: npt.NDArray[np.float64]
+    air_pressure: npt.NDArray[np.floating]
         pressure altitude at each waypoint, [:math:`Pa`]
-    thrust_setting : npt.NDArray[np.float64]
+    thrust_setting : npt.NDArray[np.floating]
         thrust setting
     q_fuel : float
         Lower calorific value (LCV) of fuel, [:math:`J \ kg_{fuel}^{-1}`].
 
     Returns
     -------
-    nvpm_ei_m : npt.NDArray[np.float64]
+    nvpm_ei_m : npt.NDArray[np.floating]
         Non-volatile particulate matter (nvPM) mass emissions index, [:math:`kg/kg_{fuel}`]
-    nvpm_ei_n : npt.NDArray[np.float64]
+    nvpm_ei_n : npt.NDArray[np.floating]
         Black carbon number emissions index, [:math:`kg_{fuel}^{-1}`]
     """
     # Non-dimensionalized thrust setting
@@ -781,13 +781,13 @@ def get_nvpm_emissions_index_edb(
 
 def nvpm_mass_emissions_index_sac(
     edb_gaseous: EDBGaseous,
-    air_pressure: npt.NDArray[np.float64],
-    true_airspeed: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-    thrust_setting: npt.NDArray[np.float64],
-    fuel_flow_per_engine: npt.NDArray[np.float64],
+    air_pressure: npt.NDArray[np.floating],
+    true_airspeed: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+    thrust_setting: npt.NDArray[np.floating],
+    fuel_flow_per_engine: npt.NDArray[np.floating],
     hydrogen_content: float,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Estimate nvPM mass emission index for singular annular combustor (SAC) engines.
 
     Here, SAC should not be confused with the Schmidt-Appleman Criterion.
@@ -799,22 +799,22 @@ def nvpm_mass_emissions_index_sac(
     ----------
     edb_gaseous : EDBGaseous
         EDB gaseous data
-    air_pressure: npt.NDArray[np.float64]
+    air_pressure: npt.NDArray[np.floating]
         pressure altitude at each waypoint, [:math:`Pa`]
-    true_airspeed: npt.NDArray[np.float64]
+    true_airspeed: npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    air_temperature: npt.NDArray[np.float64]
+    air_temperature: npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
-    thrust_setting : npt.NDArray[np.float64]
+    thrust_setting : npt.NDArray[np.floating]
         thrust setting
-    fuel_flow_per_engine: npt.NDArray[np.float64]
+    fuel_flow_per_engine: npt.NDArray[np.floating]
         fuel mass flow rate per engine, [:math:`kg s^{-1}`]
     hydrogen_content : float
         Engine unique identification number from the ICAO EDB
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         nvPM mass emissions index, [:math:`kg/kg_{fuel}`]
     """
     nvpm_ei_m_fox = black_carbon.mass_emissions_index_fox(
@@ -834,12 +834,12 @@ def nvpm_mass_emissions_index_sac(
 
 def nvpm_geometric_mean_diameter_sac(
     edb_gaseous: EDBGaseous,
-    air_pressure: npt.NDArray[np.float64],
-    true_airspeed: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-    thrust_setting: npt.NDArray[np.float64],
+    air_pressure: npt.NDArray[np.floating],
+    true_airspeed: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+    thrust_setting: npt.NDArray[np.floating],
     q_fuel: float,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     r"""
     Estimate nvPM geometric mean diameter for singular annular combustor (SAC) engines.
 
@@ -847,20 +847,20 @@ def nvpm_geometric_mean_diameter_sac(
     ----------
     edb_gaseous : EDBGaseous
         EDB gaseous data
-    air_pressure: npt.NDArray[np.float64]
+    air_pressure: npt.NDArray[np.floating]
         pressure altitude at each waypoint, [:math:`Pa`]
-    true_airspeed: npt.NDArray[np.float64]
+    true_airspeed: npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    air_temperature: npt.NDArray[np.float64]
+    air_temperature: npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
-    thrust_setting : npt.NDArray[np.float64]
+    thrust_setting : npt.NDArray[np.floating]
         thrust setting
     q_fuel : float
         Lower calorific value (LCV) of fuel, [:math:`J \ kg_{fuel}^{-1}`].
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         nvPM geometric mean diameter, [:math:`m`]
     """
     nvpm_gmd = black_carbon.geometric_mean_diameter_sac(
@@ -877,11 +877,11 @@ def nvpm_geometric_mean_diameter_sac(
 
 def get_thrust_setting(
     edb_gaseous: EDBGaseous,
-    fuel_flow_per_engine: npt.NDArray[np.float64],
-    air_pressure: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-    true_airspeed: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    fuel_flow_per_engine: npt.NDArray[np.floating],
+    air_pressure: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+    true_airspeed: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """
     Approximate the engine thrust setting at cruise conditions.
 
@@ -893,18 +893,18 @@ def get_thrust_setting(
     ----------
     edb_gaseous : EDBGaseous
         EDB gaseous data
-    fuel_flow_per_engine: npt.NDArray[np.float64]
+    fuel_flow_per_engine: npt.NDArray[np.floating]
         Fuel mass flow rate per engine, [:math:`kg s^{-1}`]
-    air_pressure: npt.NDArray[np.float64]
+    air_pressure: npt.NDArray[np.floating]
         Pressure altitude at each waypoint, [:math:`Pa`]
-    air_temperature: npt.NDArray[np.float64]
+    air_temperature: npt.NDArray[np.floating]
         Ambient temperature for each waypoint, [:math:`K`]
-    true_airspeed: npt.NDArray[np.float64]
+    true_airspeed: npt.NDArray[np.floating]
         True airspeed for each waypoint, [:math:`m s^{-1}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Engine thrust setting. Returns ``np.nan`` if engine data is
         not available in the ICAO EDB dataset.
     """
diff --git a/pycontrails/models/emissions/ffm2.py b/pycontrails/models/emissions/ffm2.py
index 2d7b94c44..252af83b8 100644
--- a/pycontrails/models/emissions/ffm2.py
+++ b/pycontrails/models/emissions/ffm2.py
@@ -172,27 +172,27 @@ def co_hc_emissions_index_profile(
 
 def estimate_nox(
     log_ei_nox_profile: EmissionsProfileInterpolator,
-    fuel_flow_per_engine: npt.NDArray[np.float64],
-    true_airspeed: npt.NDArray[np.float64],
-    air_pressure: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-    specific_humidity: None | npt.NDArray[np.float64] = None,
-) -> npt.NDArray[np.float64]:
+    fuel_flow_per_engine: npt.NDArray[np.floating],
+    true_airspeed: npt.NDArray[np.floating],
+    air_pressure: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+    specific_humidity: None | npt.NDArray[np.floating] = None,
+) -> npt.NDArray[np.floating]:
     """Estimate the nitrogen oxide (NOx) emissions index (EI) at cruise conditions.
 
     Parameters
     ----------
     log_ei_nox_profile: EmissionsProfileInterpolator
         emissions profile containing the log of EI NOx versus log of fuel flow.
-    fuel_flow_per_engine: npt.NDArray[np.float64]
+    fuel_flow_per_engine: npt.NDArray[np.floating]
         fuel mass flow rate per engine, [:math:`kg s^{-1}`]
-    true_airspeed: npt.NDArray[np.float64]
+    true_airspeed: npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         pressure altitude at each waypoint, [:math:`Pa`]
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
-    specific_humidity: npt.NDArray[np.float64] | None
+    specific_humidity: npt.NDArray[np.floating] | None
         specific humidity for each waypoint, [:math:`kg_{H_{2}O}/kg_{air}`]
     """
 
@@ -215,24 +215,24 @@ def estimate_nox(
 
 def estimate_ei(
     log_ei_profile: EmissionsProfileInterpolator,
-    fuel_flow_per_engine: npt.NDArray[np.float64],
-    true_airspeed: npt.NDArray[np.float64],
-    air_pressure: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    fuel_flow_per_engine: npt.NDArray[np.floating],
+    true_airspeed: npt.NDArray[np.floating],
+    air_pressure: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     """Estimate carbon monoxide (CO) or hydrocarbon (HC) emissions index (EI).
 
     Parameters
     ----------
     log_ei_profile: EmissionsProfileInterpolator
         emissions profile containing the log of EI CO or EI HC versus log of fuel flow.
-    fuel_flow_per_engine: npt.NDArray[np.float64]
+    fuel_flow_per_engine: npt.NDArray[np.floating]
         fuel mass flow rate per engine, [:math:`kg s^{-1}`]
-    true_airspeed: npt.NDArray[np.float64]
+    true_airspeed: npt.NDArray[np.floating]
         true airspeed for each waypoint, [:math:`m s^{-1}`]
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         pressure altitude at each waypoint, [:math:`Pa`]
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         ambient temperature for each waypoint, [:math:`K`]
     """
     mach_num = units.tas_to_mach_number(true_airspeed, air_temperature)
@@ -252,29 +252,29 @@ def estimate_ei(
 
 
 def ei_at_cruise(
-    ei_sl: npt.NDArray[np.float64],
-    theta_amb: npt.NDArray[np.float64],
-    delta_amb: npt.NDArray[np.float64],
+    ei_sl: npt.NDArray[np.floating],
+    theta_amb: npt.NDArray[np.floating],
+    delta_amb: npt.NDArray[np.floating],
     ei_type: str,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Convert the estimated EI at sea level to cruise conditions.
 
     Refer to Eqs. (15) and (16) in DuBois & Paynter (2006).
 
     Parameters
     ----------
-    ei_sl : npt.NDArray[np.float64]
+    ei_sl : npt.NDArray[np.floating]
         Sea level EI values.
-    theta_amb : npt.NDArray[np.float64]
+    theta_amb : npt.NDArray[np.floating]
         Ratio of the ambient temperature to the temperature at mean sea-level.
-    delta_amb : npt.NDArray[np.float64]
+    delta_amb : npt.NDArray[np.floating]
         Ratio of the pressure altitude to the surface pressure.
     ei_type : str
         One of {"HC", "CO", "NOX"}
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Estimated cruise EI values.
 
     References
@@ -292,30 +292,30 @@ def ei_at_cruise(
 
 
 def _get_humidity_correction_factor(
-    specific_humidity: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    specific_humidity: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     return np.exp(-19 * (specific_humidity - 0.00634))
 
 
 def _estimate_specific_humidity(
-    air_temperature: npt.NDArray[np.float64], air_pressure: npt.NDArray[np.float64], rh: float
-) -> npt.NDArray[np.float64]:
+    air_temperature: npt.NDArray[np.floating], air_pressure: npt.NDArray[np.floating], rh: float
+) -> npt.NDArray[np.floating]:
     """Estimate the specific humidity by assuming a fixed relative humidity.
 
     Refer to Eqs. (43), (44) and (45) in DuBois & Paynter (2006).
 
     Parameters
     ----------
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         Air temperature, [:math:`K`]
-    air_pressure : npt.NDArray[np.float64]
+    air_pressure : npt.NDArray[np.floating]
         Air pressure, [:math:`Pa`]
     rh : float
         Relative humidity, [:math:`0 - 1`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Estimated specific humidity, [:math:`kg kg^{-1}`]
 
     References
diff --git a/pycontrails/models/humidity_scaling/humidity_scaling.py b/pycontrails/models/humidity_scaling/humidity_scaling.py
index abf2c7ccc..c9d7ac5bf 100644
--- a/pycontrails/models/humidity_scaling/humidity_scaling.py
+++ b/pycontrails/models/humidity_scaling/humidity_scaling.py
@@ -625,7 +625,7 @@ def histogram_matching(
     level_type: str,
     member: int | None,
     q_method: str | None,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Map ERA5-derived RHi to its corresponding IAGOS quantile via histogram matching.
 
     This matching is performed on a **single** ERA5 ensemble member.
@@ -649,7 +649,7 @@ def histogram_matching(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         The IAGOS quantiles corresponding to the ERA5-derived RHi values. Returned
         as a numpy array with the same shape and dtype as ``era5_rhi``.
     """
@@ -691,8 +691,8 @@ def histogram_matching(
 
 
 def histogram_matching_all_members(
-    era5_rhi_all_members: npt.NDArray[np.float64], member: int, q_method: str
-) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    era5_rhi_all_members: npt.NDArray[np.floating], member: int, q_method: str
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
     """Recalibrate ERA5-derived RHi values to IAGOS quantiles by histogram matching.
 
     This recalibration requires values for **all** ERA5 ensemble members. Currently, the
@@ -700,7 +700,7 @@ def histogram_matching_all_members(
 
     Parameters
     ----------
-    era5_rhi_all_members : npt.NDArray[np.float64]
+    era5_rhi_all_members : npt.NDArray[np.floating]
         ERA5-derived RHi values for all ensemble members. This array should have shape ``(n, 10)``.
     member : int
         The ERA5 ensemble member to use. Must be in the range ``[0, 10)``.
@@ -709,10 +709,10 @@ def histogram_matching_all_members(
 
     Returns
     -------
-    ensemble_mean_rhi : npt.NDArray[np.float64]
+    ensemble_mean_rhi : npt.NDArray[np.floating]
         The mean RHi values after histogram matching over all ensemble members.
         This is an array of shape ``(n,)``.
-    ensemble_member_rhi : npt.NDArray[np.float64]
+    ensemble_member_rhi : npt.NDArray[np.floating]
         The RHi values after histogram matching for the given ensemble member.
         This is an array of shape ``(n,)``.
     """
@@ -727,7 +727,7 @@ def histogram_matching_all_members(
 
     # Perform histogram matching on all other ensemble members
     # Add up the results into a single 'ensemble_mean_rhi' array
-    ensemble_mean_rhi: npt.NDArray[np.float64] = 0.0  # type: ignore[assignment]
+    ensemble_mean_rhi: npt.NDArray[np.floating] = 0.0  # type: ignore[assignment]
     for r in range(n_members):
         if r == member:
             ensemble_mean_rhi += ensemble_member_rhi
@@ -743,25 +743,25 @@ def histogram_matching_all_members(
 
 
 def eckel_scaling(
-    ensemble_mean_rhi: npt.NDArray[np.float64],
-    ensemble_member_rhi: npt.NDArray[np.float64],
+    ensemble_mean_rhi: npt.NDArray[np.floating],
+    ensemble_member_rhi: npt.NDArray[np.floating],
     q_method: str,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Apply Eckel scaling to the given RHi values.
 
     Parameters
     ----------
-    ensemble_mean_rhi : npt.NDArray[np.float64]
+    ensemble_mean_rhi : npt.NDArray[np.floating]
         The ensemble mean RHi values. This should be a 1D array with the same shape as
         ``ensemble_member_rhi``.
-    ensemble_member_rhi : npt.NDArray[np.float64]
+    ensemble_member_rhi : npt.NDArray[np.floating]
         The RHi values for a single ensemble member.
     q_method : {None, "cubic-spline", "log-q-log-p"}
         The interpolation method.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         The scaled RHi values. Values are manually clipped at 0 to ensure
         only non-negative values are returned.
 
@@ -984,11 +984,11 @@ def eval(
     @override
     def scale(  # type: ignore[override]
         self,
-        specific_humidity: npt.NDArray[np.float64],
-        air_temperature: npt.NDArray[np.float64],
-        air_pressure: npt.NDArray[np.float64],
+        specific_humidity: npt.NDArray[np.floating],
+        air_temperature: npt.NDArray[np.floating],
+        air_pressure: npt.NDArray[np.floating],
         **kwargs: Any,
-    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    ) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
         """Scale specific humidity values via histogram matching and Eckel scaling.
 
         Unlike the method on the base class, the method assumes each of the input
@@ -996,22 +996,22 @@ def scale(  # type: ignore[override]
 
         Parameters
         ----------
-        specific_humidity : npt.NDArray[np.float64]
+        specific_humidity : npt.NDArray[np.floating]
             A 2D array of specific humidity values for all ERA5 ensemble members.
             The shape of this array must be ``(n, 10)``, where ``n`` is the number
             of observations and ``10`` is the number of ERA5 ensemble members.
-        air_temperature : npt.NDArray[np.float64]
+        air_temperature : npt.NDArray[np.floating]
             A 1D array of air temperature values with shape ``(n,)``.
-        air_pressure : npt.NDArray[np.float64]
+        air_pressure : npt.NDArray[np.floating]
             A 1D array of air pressure values with shape ``(n,)``.
         kwargs: Any
             Unused, kept for compatibility with the base class.
 
         Returns
         -------
-        specific_humidity : npt.NDArray[np.float64]
+        specific_humidity : npt.NDArray[np.floating]
             The recalibrated specific humidity values. A 1D array with shape ``(n,)``.
-        rhi : npt.NDArray[np.float64]
+        rhi : npt.NDArray[np.floating]
             The recalibrated RHi values. A 1D array with shape ``(n,)``.
         """
 
diff --git a/pycontrails/models/ps_model/ps_grid.py b/pycontrails/models/ps_model/ps_grid.py
index d3221d65b..bb92a5a0e 100644
--- a/pycontrails/models/ps_model/ps_grid.py
+++ b/pycontrails/models/ps_model/ps_grid.py
@@ -178,9 +178,9 @@ def eval(
 @dataclasses.dataclass
 class _PerfVariables:
     atyp_param: PSAircraftEngineParams
-    air_pressure: npt.NDArray[np.float64] | float
-    air_temperature: npt.NDArray[np.float64] | float
-    mach_number: npt.NDArray[np.float64] | float
+    air_pressure: npt.NDArray[np.floating] | float
+    air_temperature: npt.NDArray[np.floating] | float
+    mach_number: npt.NDArray[np.floating] | float
     q_fuel: float
 
 
@@ -271,7 +271,7 @@ def _estimate_mass_extremes(
     atyp_param: PSAircraftEngineParams,
     perf: _PerfVariables,
     n_iter: int = 3,
-) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
     """Calculate the minimum and maximum mass for a given aircraft type."""
 
     oem = atyp_param.amass_oew  # operating empty mass
@@ -296,9 +296,9 @@ def _estimate_mass_extremes(
 
 
 def _parse_variables(
-    level: npt.NDArray[np.float64] | None,
-    air_temperature: xr.DataArray | npt.NDArray[np.float64] | None,
-) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    level: npt.NDArray[np.floating] | None,
+    air_temperature: xr.DataArray | npt.NDArray[np.floating] | None,
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
     """Parse the level and air temperature arguments."""
 
     if isinstance(air_temperature, xr.DataArray):
@@ -328,8 +328,8 @@ def _parse_variables(
 def ps_nominal_grid(
     aircraft_type: str,
     *,
-    level: npt.NDArray[np.float64] | None = None,
-    air_temperature: xr.DataArray | npt.NDArray[np.float64] | None = None,
+    level: npt.NDArray[np.floating] | None = None,
+    air_temperature: xr.DataArray | npt.NDArray[np.floating] | None = None,
     q_fuel: float = JetA.q_fuel,
     mach_number: float | None = None,
     maxiter: int = PSGridParams.maxiter,
@@ -345,10 +345,10 @@ def ps_nominal_grid(
     ----------
     aircraft_type : str
         The aircraft type.
-    level : npt.NDArray[np.float64] | None, optional
+    level : npt.NDArray[np.floating] | None, optional
         The pressure level, [:math:`hPa`]. If None, the ``air_temperature``
         argument must be a :class:`xarray.DataArray` with a ``level`` coordinate.
-    air_temperature : xr.DataArray | npt.NDArray[np.float64] | None, optional
+    air_temperature : xr.DataArray | npt.NDArray[np.floating] | None, optional
         The ambient air temperature, [:math:`K`]. If None (default), the ISA
         temperature is computed from the ``level`` argument. If a :class:`xarray.DataArray`,
         the ``level`` coordinate must be present and the ``level`` argument must be None
@@ -624,12 +624,11 @@ def ps_nominal_optimize_mach(
         altitude_m = units.pl_to_m(level)
         air_temperature = units.m_to_T_isa(altitude_m)
 
-    headwind: ArrayOrFloat
     if northward_wind is not None and eastward_wind is not None:
         if sin_a is None or cos_a is None:
             msg = "Segment angles must be provide if wind data is specified"
             raise ValueError(msg)
-        headwind = -(northward_wind * cos_a + eastward_wind * sin_a)
+        headwind = -(northward_wind * cos_a + eastward_wind * sin_a)  # type: ignore[misc]
     else:
         headwind = 0.0  # type: ignore
 
diff --git a/pycontrails/models/ps_model/ps_model.py b/pycontrails/models/ps_model/ps_model.py
index d79cbccd5..9757a1ee7 100644
--- a/pycontrails/models/ps_model/ps_model.py
+++ b/pycontrails/models/ps_model/ps_model.py
@@ -224,14 +224,14 @@ def calculate_aircraft_performance(
         self,
         *,
         aircraft_type: str,
-        altitude_ft: npt.NDArray[np.float64],
-        air_temperature: npt.NDArray[np.float64],
+        altitude_ft: npt.NDArray[np.floating],
+        air_temperature: npt.NDArray[np.floating],
         time: npt.NDArray[np.datetime64] | None,
-        true_airspeed: npt.NDArray[np.float64] | float | None,
-        aircraft_mass: npt.NDArray[np.float64] | float,
-        engine_efficiency: npt.NDArray[np.float64] | float | None,
-        fuel_flow: npt.NDArray[np.float64] | float | None,
-        thrust: npt.NDArray[np.float64] | float | None,
+        true_airspeed: npt.NDArray[np.floating] | float | None,
+        aircraft_mass: npt.NDArray[np.floating] | float,
+        engine_efficiency: npt.NDArray[np.floating] | float | None,
+        fuel_flow: npt.NDArray[np.floating] | float | None,
+        thrust: npt.NDArray[np.floating] | float | None,
         q_fuel: float,
         **kwargs: Any,
     ) -> AircraftPerformanceData:
@@ -266,8 +266,8 @@ def calculate_aircraft_performance(
         rn = reynolds_number(atyp_param.wing_surface_area, mach_num, air_temperature, air_pressure)
 
         # Allow array or None time
-        dv_dt: npt.NDArray[np.float64] | float
-        theta: npt.NDArray[np.float64] | float
+        dv_dt: npt.NDArray[np.floating] | float
+        theta: npt.NDArray[np.floating] | float
         if time is None:
             # Assume a nominal cruising state
             dt_sec = None
@@ -780,7 +780,7 @@ def overall_propulsion_efficiency(
     c_t_eta_b: ArrayOrFloat,
     atyp_param: PSAircraftEngineParams,
     eta_over_eta_b_min: float | None = None,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Calculate overall propulsion efficiency.
 
     Parameters
@@ -800,7 +800,7 @@ def overall_propulsion_efficiency(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Overall propulsion efficiency
     """
     eta_over_eta_b = propulsion_efficiency_over_max_propulsion_efficiency(mach_num, c_t, c_t_eta_b)
@@ -816,7 +816,7 @@ def propulsion_efficiency_over_max_propulsion_efficiency(
     mach_num: ArrayOrFloat,
     c_t: ArrayOrFloat,
     c_t_eta_b: ArrayOrFloat,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Calculate ratio of OPE to maximum OPE that can be attained for a given Mach number.
 
     Parameters
@@ -830,7 +830,7 @@ def propulsion_efficiency_over_max_propulsion_efficiency(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Ratio of OPE to maximum OPE, ``eta / eta_b``
 
     Notes
diff --git a/pycontrails/models/ps_model/ps_operational_limits.py b/pycontrails/models/ps_model/ps_operational_limits.py
index 45a1a8272..fc94abb35 100644
--- a/pycontrails/models/ps_model/ps_operational_limits.py
+++ b/pycontrails/models/ps_model/ps_operational_limits.py
@@ -155,34 +155,34 @@ def max_available_thrust_coefficient(
 
 
 def get_excess_thrust_available(
-    mach_number: float | npt.NDArray[np.float64],
-    air_temperature: float | npt.NDArray[np.float64],
-    air_pressure: float | npt.NDArray[np.float64],
-    aircraft_mass: float | npt.NDArray[np.float64],
-    theta: float | npt.NDArray[np.float64],
+    mach_number: ArrayOrFloat,
+    air_temperature: ArrayOrFloat,
+    air_pressure: ArrayOrFloat,
+    aircraft_mass: ArrayOrFloat,
+    theta: ArrayOrFloat,
     atyp_param: PSAircraftEngineParams,
-) -> float | npt.NDArray[np.float64]:
+) -> ArrayOrFloat:
     r"""
     Calculate the excess thrust coefficient available at specified operation condition.
 
     Parameters
     ----------
-    mach_number : float | npt.NDArray[np.float64]
+    mach_number : ArrayOrFloat
         Mach number at each waypoint
-    air_temperature : float | npt.NDArray[np.float64]
+    air_temperature : ArrayOrFloat
         Ambient temperature at each waypoint, [:math:`K`]
-    air_pressure : float | npt.NDArray[np.float64]
+    air_pressure : ArrayOrFloat
         Ambient pressure, [:math:`Pa`]
-    aircraft_mass : float | npt.NDArray[np.float64]
+    aircraft_mass : ArrayOrFloat
         Aircraft mass at each waypoint, [:math:`kg`]
-    theta : float | npt.NDArray[np.float64]
+    theta : ArrayOrFloat
         Climb (positive value) or descent (negative value) angle, [:math:`\deg`]
     atyp_param : PSAircraftEngineParams
         Extracted aircraft and engine parameters.
 
     Returns
     -------
-    float | npt.NDArray[np.float64]
+    ArrayOrFloat
         The difference between the maximum rated thrust coefficient and the thrust coefficient
         required to maintain the current mach_number.
     """
@@ -217,7 +217,7 @@ def get_excess_thrust_available(
     )
 
     tas = units.mach_number_to_tas(mach_number, air_temperature)
-    req_thrust_coeff = required_thrust_coefficient(c_lift, c_drag, tas)
+    req_thrust_coeff = required_thrust_coefficient(c_lift, c_drag, tas)  # type: ignore[type-var]
 
     c_t_eta_b = thrust_coefficient_at_max_efficiency(
         mach_number, atyp_param.m_des, atyp_param.c_t_des
@@ -226,7 +226,7 @@ def get_excess_thrust_available(
         air_temperature, mach_number, c_t_eta_b, atyp_param
     )
 
-    return max_thrust_coeff - req_thrust_coeff
+    return max_thrust_coeff - req_thrust_coeff  # type: ignore[return-value]
 
 
 def _normalised_max_throttle_parameter(
@@ -466,7 +466,9 @@ def maximum_mach_num(
 # ----------------
 
 
-def fuel_flow_idle(fuel_flow_idle_sls: float, altitude_ft: ArrayOrFloat) -> npt.NDArray[np.float64]:
+def fuel_flow_idle(
+    fuel_flow_idle_sls: float, altitude_ft: ArrayOrFloat
+) -> npt.NDArray[np.floating]:
     r"""Calculate minimum fuel mass flow rate at flight idle conditions.
 
     Parameters
@@ -478,7 +480,7 @@ def fuel_flow_idle(fuel_flow_idle_sls: float, altitude_ft: ArrayOrFloat) -> npt.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Fuel mass flow rate at flight idle conditions, [:math:`kg \ s^{-1}`]
     """
     x = altitude_ft / 10000.0
@@ -491,7 +493,7 @@ def max_fuel_flow(
     mach_number: ArrayOrFloat,
     fuel_flow_max_sls: float,
     flight_phase: npt.NDArray[np.uint8] | flight.FlightPhase,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     r"""Correct maximum fuel mass flow rate that can be supplied by the engine.
 
     Parameters
@@ -509,7 +511,7 @@ def max_fuel_flow(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Maximum allowable fuel mass flow rate, [:math:`kg \ s^{-1}`]
     """
     ff_max = jet.equivalent_fuel_flow_rate_at_cruise(
diff --git a/pycontrails/physics/geo.py b/pycontrails/physics/geo.py
index 9f1f281ce..961332c2e 100644
--- a/pycontrails/physics/geo.py
+++ b/pycontrails/physics/geo.py
@@ -62,20 +62,20 @@ def haversine(lons0: ArrayLike, lats0: ArrayLike, lons1: ArrayLike, lats1: Array
 
 
 def segment_haversine(
-    longitude: npt.NDArray[np.float64], latitude: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    longitude: npt.NDArray[np.floating], latitude: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     r"""Calculate haversine distance between consecutive points along path.
 
     Parameters
     ----------
-    longitude : npt.NDArray[np.float64]
+    longitude : npt.NDArray[np.floating]
         1D Longitude values with index corresponding to latitude inputs, [:math:`\deg`]
-    latitude : npt.NDArray[np.float64]
+    latitude : npt.NDArray[np.floating]
         1D Latitude values with index corresponding to longitude inputs, [:math:`\deg`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Haversine distance between (lat_i, lon_i) and (lat_i+1, lon_i+1), [:math:`m`]
         The final entry of the output is set to nan.
 
@@ -97,8 +97,8 @@ def segment_haversine(
 
 
 def azimuth_to_direction(
-    azimuth_: npt.NDArray[np.float64], latitude: npt.NDArray[np.float64]
-) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    azimuth_: npt.NDArray[np.floating], latitude: npt.NDArray[np.floating]
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
     r"""Calculate rectangular direction from spherical azimuth.
 
     This implementation uses the equation
@@ -109,14 +109,14 @@ def azimuth_to_direction(
 
     Parameters
     ----------
-    azimuth_ : npt.NDArray[np.float64]
+    azimuth_ : npt.NDArray[np.floating]
         Angle measured clockwise from true north, [:math:`\deg`]
-    latitude : npt.NDArray[np.float64]
+    latitude : npt.NDArray[np.floating]
         Latitude value of the point, [:math:`\deg`]
 
     Returns
     -------
-    tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]
+    tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]
         A tuple of sine and cosine values.
     """
     cos_lat = np.cos(units.degrees_to_radians(latitude))
@@ -138,22 +138,22 @@ def azimuth_to_direction(
 
 
 def azimuth(
-    lons0: npt.NDArray[np.float64],
-    lats0: npt.NDArray[np.float64],
-    lons1: npt.NDArray[np.float64],
-    lats1: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    lons0: npt.NDArray[np.floating],
+    lats0: npt.NDArray[np.floating],
+    lons1: npt.NDArray[np.floating],
+    lats1: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""Calculate angle relative to true north for set of coordinates.
 
     Parameters
     ----------
-    lons0 : npt.NDArray[np.float64]
+    lons0 : npt.NDArray[np.floating]
         Longitude values of initial endpoints, [:math:`\deg`].
-    lats0 : npt.NDArray[np.float64]
+    lats0 : npt.NDArray[np.floating]
         Latitude values of initial endpoints, [:math:`\deg`].
-    lons1 : npt.NDArray[np.float64]
+    lons1 : npt.NDArray[np.floating]
         Longitude values of terminal endpoints, [:math:`\deg`].
-    lats1 : npt.NDArray[np.float64]
+    lats1 : npt.NDArray[np.floating]
         Latitude values of terminal endpoints, [:math:`\deg`].
 
     References
@@ -162,7 +162,7 @@ def azimuth(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Azimuth relative to true north (:math:`0\deg`), [:math:`\deg`]
 
     See Also
@@ -186,22 +186,22 @@ def azimuth(
 
 
 def segment_azimuth(
-    longitude: npt.NDArray[np.float64], latitude: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    longitude: npt.NDArray[np.floating], latitude: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     r"""Calculate the angle between coordinate segments and true north.
 
     `np.nan` is added to the final value so the length of the output is the same as the inputs.
 
     Parameters
     ----------
-    longitude : npt.NDArray[np.float64]
+    longitude : npt.NDArray[np.floating]
         Longitude values, [:math:`\deg`]
-    latitude : npt.NDArray[np.float64]
+    latitude : npt.NDArray[np.floating]
         Latitude values, [:math:`\deg`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Azimuth relative to true north (:math:`0\deg`), [:math:`\deg`]
         Final entry of each array is set to `np.nan`.
 
@@ -227,22 +227,22 @@ def segment_azimuth(
 
 
 def longitudinal_angle(
-    lons0: npt.NDArray[np.float64],
-    lats0: npt.NDArray[np.float64],
-    lons1: npt.NDArray[np.float64],
-    lats1: npt.NDArray[np.float64],
-) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    lons0: npt.NDArray[np.floating],
+    lats0: npt.NDArray[np.floating],
+    lons1: npt.NDArray[np.floating],
+    lats1: npt.NDArray[np.floating],
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
     r"""Calculate angle with longitudinal axis for sequence of segments.
 
     Parameters
     ----------
-    lons0 : npt.NDArray[np.float64]
+    lons0 : npt.NDArray[np.floating]
         Longitude values of initial endpoints, [:math:`\deg`].
-    lats0 : npt.NDArray[np.float64]
+    lats0 : npt.NDArray[np.floating]
         Latitude values of initial endpoints, [:math:`\deg`].
-    lons1 : npt.NDArray[np.float64]
+    lons1 : npt.NDArray[np.floating]
         Longitude values of terminal endpoints, [:math:`\deg`].
-    lats1 : npt.NDArray[np.float64]
+    lats1 : npt.NDArray[np.floating]
         Latitude values of terminal endpoints, [:math:`\deg`].
 
     References
@@ -251,9 +251,9 @@ def longitudinal_angle(
 
     Returns
     -------
-    sin_a : npt.NDArray[np.float64]
+    sin_a : npt.NDArray[np.floating]
         Sine values.
-    cos_a : npt.NDArray[np.float64]
+    cos_a : npt.NDArray[np.floating]
         Cosine values.
     """
     lons0 = units.degrees_to_radians(lons0)
@@ -275,22 +275,22 @@ def longitudinal_angle(
 
 
 def segment_angle(
-    longitude: npt.NDArray[np.float64], latitude: npt.NDArray[np.float64]
-) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    longitude: npt.NDArray[np.floating], latitude: npt.NDArray[np.floating]
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
     r"""Calculate the angle between coordinate segments and the longitudinal axis.
 
     `np.nan` is added to the final value so the length of the output is the same as the inputs.
 
     Parameters
     ----------
-    longitude : npt.NDArray[np.float64]
+    longitude : npt.NDArray[np.floating]
         Longitude values, [:math:`\deg`]
-    latitude : npt.NDArray[np.float64]
+    latitude : npt.NDArray[np.floating]
         Latitude values, [:math:`\deg`]
 
     Returns
     -------
-    tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]
+    tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]
         sin(a), cos(a), where ``a`` is the angle between the segment and the longitudinal axis.
         Final entry of each array is set to `np.nan`.
 
@@ -332,10 +332,10 @@ def segment_angle(
 
 
 def segment_length(
-    longitude: npt.NDArray[np.float64],
-    latitude: npt.NDArray[np.float64],
-    altitude: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    longitude: npt.NDArray[np.floating],
+    latitude: npt.NDArray[np.floating],
+    altitude: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""Calculate the segment length between coordinates by assuming a great circle distance.
 
     Requires coordinates to be in EPSG:4326.
@@ -345,16 +345,16 @@ def segment_length(
 
     Parameters
     ----------
-    longitude : npt.NDArray[np.float64]
+    longitude : npt.NDArray[np.floating]
         Longitude values, [:math:`\deg`]
-    latitude : npt.NDArray[np.float64]
+    latitude : npt.NDArray[np.floating]
         Latitude values, [:math:`\deg`]
-    altitude : npt.NDArray[np.float64]
+    altitude : npt.NDArray[np.floating]
         Altitude values, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Array of distances in [:math:`m`] between coordinates.
         Final entry of each array is set to `np.nan`.
 
@@ -371,11 +371,11 @@ def segment_length(
 
 
 def forward_azimuth(
-    lons: npt.NDArray[np.float64],
-    lats: npt.NDArray[np.float64],
-    az: npt.NDArray[np.float64] | float,
-    dist: npt.NDArray[np.float64] | float,
-) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    lons: npt.NDArray[np.floating],
+    lats: npt.NDArray[np.floating],
+    az: npt.NDArray[np.floating] | float,
+    dist: npt.NDArray[np.floating] | float,
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
     r"""Calculate coordinates along forward azimuth.
 
     This function is identical to the `pyproj.Geod.fwd` method when working on
@@ -384,19 +384,19 @@ def forward_azimuth(
 
     Parameters
     ----------
-    lons : npt.NDArray[np.float64]
+    lons : npt.NDArray[np.floating]
         Array of longitude values.
-    lats : npt.NDArray[np.float64]
+    lats : npt.NDArray[np.floating]
         Array of latitude values.
-    az : npt.NDArray[np.float64] | float
+    az : npt.NDArray[np.floating] | float
         Azimuth, measured in [:math:`\deg`].
-    dist : npt.NDArray[np.float64] | float
+    dist : npt.NDArray[np.floating] | float
         Distance [:math:`m`] between initial longitude latitude values and
         point to be computed.
 
     Returns
     -------
-    tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]
+    tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]
         Tuple of longitude latitude arrays.
 
     See Also
@@ -1010,7 +1010,7 @@ def advect_horizontal(
 
 
 def spatial_bounding_box(
-    longitude: npt.NDArray[np.float64], latitude: npt.NDArray[np.float64], buffer: float = 1.0
+    longitude: npt.NDArray[np.floating], latitude: npt.NDArray[np.floating], buffer: float = 1.0
 ) -> tuple[float, float, float, float]:
     r"""
     Construct rectangular spatial bounding box from a set of waypoints.
@@ -1073,17 +1073,17 @@ def domain_surface_area(
 
 
 def grid_surface_area(
-    longitude: npt.NDArray[np.float64], latitude: npt.NDArray[np.float64]
+    longitude: npt.NDArray[np.floating], latitude: npt.NDArray[np.floating]
 ) -> xr.DataArray:
     r"""
     Calculate surface area that is covered by each pixel in a longitude-latitude grid.
 
     Parameters
     ----------
-    longitude: npt.NDArray[np.float64]
+    longitude: npt.NDArray[np.floating]
         Longitude coordinates in a longitude-latitude grid, [:math:`\deg`].
         Must be in ascending order.
-    latitude: npt.NDArray[np.float64]
+    latitude: npt.NDArray[np.floating]
         Latitude coordinates in a longitude-latitude grid, [:math:`\deg`].
         Must be in ascending order.
 
@@ -1119,19 +1119,19 @@ def grid_surface_area(
 
 
 def _area_between_latitude_and_north_pole(
-    latitude: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    latitude: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""
     Calculate surface area from the provided latitude to the North Pole.
 
     Parameters
     ----------
-    latitude: npt.NDArray[np.float64]
+    latitude: npt.NDArray[np.floating]
         1D Latitude values with index corresponding to latitude inputs, [:math:`\deg`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Surface area from latitude to North Pole, [:math:`m^{2}`]
     """
     lat_radians = units.degrees_to_radians(latitude)
diff --git a/pycontrails/physics/jet.py b/pycontrails/physics/jet.py
index 15ea6f202..25e0bd6b7 100644
--- a/pycontrails/physics/jet.py
+++ b/pycontrails/physics/jet.py
@@ -31,20 +31,20 @@
 
 
 def acceleration(
-    true_airspeed: npt.NDArray[np.float64], segment_duration: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    true_airspeed: npt.NDArray[np.floating], segment_duration: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     r"""Calculate the acceleration/deceleration at each waypoint.
 
     Parameters
     ----------
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         True airspeed, [:math:`m \ s^{-1}`]
-    segment_duration : npt.NDArray[np.float64]
+    segment_duration : npt.NDArray[np.floating]
         Time difference between waypoints, [:math:`s`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Acceleration/deceleration, [:math:`m \ s^{-2}`]
 
     See Also
@@ -59,20 +59,20 @@ def acceleration(
 
 
 def climb_descent_angle(
-    true_airspeed: npt.NDArray[np.float64], rocd: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    true_airspeed: npt.NDArray[np.floating], rocd: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     r"""Calculate angle between the horizontal plane and the actual flight path.
 
     Parameters
     ----------
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         True airspeed, [:math:`m \ s^{-1}`]
-    rocd : npt.NDArray[np.float64]
+    rocd : npt.NDArray[np.floating]
         Rate of climb/descent, [:math:`ft min^{-1}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Climb (positive value) or descent (negative value) angle, [:math:`\deg`]
 
     See Also
@@ -86,10 +86,10 @@ def climb_descent_angle(
 
 
 def clip_mach_number(
-    true_airspeed: npt.NDArray[np.float64],
-    air_temperature: npt.NDArray[np.float64],
+    true_airspeed: npt.NDArray[np.floating],
+    air_temperature: npt.NDArray[np.floating],
     max_mach_number: ArrayOrFloat,
-) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
     r"""Compute the Mach number from the true airspeed and ambient temperature.
 
     This method clips the computed Mach number to the value of ``max_mach_number``.
@@ -99,9 +99,9 @@ def clip_mach_number(
 
     Parameters
     ----------
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         Array of true airspeed, [:math:`m \ s^{-1}`]
-    air_temperature : npt.NDArray[np.float64]
+    air_temperature : npt.NDArray[np.floating]
         Array of ambient temperature, [:math: `K`]
     max_mach_number : ArrayOrFloat
         Maximum mach number associated to aircraft. If no clipping
@@ -109,10 +109,10 @@ def clip_mach_number(
 
     Returns
     -------
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         Array of true airspeed, [:math:`m \ s^{-1}`]. All values are clipped at
         ``max_mach_number``.
-    mach_num : npt.NDArray[np.float64]
+    mach_num : npt.NDArray[np.floating]
         Array of Mach numbers, [:math:`Ma`]. All values are clipped at
         ``max_mach_number``.
     """
@@ -137,13 +137,13 @@ def clip_mach_number(
 
 
 def overall_propulsion_efficiency(
-    true_airspeed: npt.NDArray[np.float64],
-    F_thrust: npt.NDArray[np.float64],
-    fuel_flow: npt.NDArray[np.float64],
+    true_airspeed: npt.NDArray[np.floating],
+    F_thrust: npt.NDArray[np.floating],
+    fuel_flow: npt.NDArray[np.floating],
     q_fuel: float,
     is_descent: npt.NDArray[np.bool_] | bool | None,
     threshold: float = 0.5,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     r"""Calculate the overall propulsion efficiency (OPE).
 
     Negative OPE values can occur during the descent phase and is clipped to a
@@ -152,22 +152,22 @@ def overall_propulsion_efficiency(
 
     Parameters
     ----------
-    true_airspeed: npt.NDArray[np.float64]
+    true_airspeed: npt.NDArray[np.floating]
         True airspeed for each waypoint, [:math:`m s^{-1}`].
-    F_thrust: npt.NDArray[np.float64]
+    F_thrust: npt.NDArray[np.floating]
         Thrust force provided by the engine, [:math:`N`].
-    fuel_flow: npt.NDArray[np.float64]
+    fuel_flow: npt.NDArray[np.floating]
         Fuel mass flow rate, [:math:`kg s^{-1}`].
     q_fuel : float
         Lower calorific value (LCV) of fuel, [:math:`J \ kg_{fuel}^{-1}`].
-    is_descent : npt.NDArray[np.float64] | None
+    is_descent : npt.NDArray[np.floating] | None
         Boolean array that indicates if a waypoint is in a descent phase.
     threshold : float
         Upper bound for realistic engine efficiency.
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Overall propulsion efficiency (OPE)
 
     References
@@ -196,49 +196,49 @@ def overall_propulsion_efficiency(
 
 
 def fuel_burn(
-    fuel_flow: npt.NDArray[np.float64], segment_duration: npt.NDArray[np.float64]
-) -> npt.NDArray[np.float64]:
+    fuel_flow: npt.NDArray[np.floating], segment_duration: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """Calculate the fuel consumption at each waypoint.
 
     Parameters
     ----------
-    fuel_flow: npt.NDArray[np.float64]
+    fuel_flow: npt.NDArray[np.floating]
         Fuel mass flow rate, [:math:`kg s^{-1}`]
-    segment_duration: npt.NDArray[np.float64]
+    segment_duration: npt.NDArray[np.floating]
         Time difference between waypoints, [:math:`s`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Fuel consumption at each waypoint, [:math:`kg`]
     """
     return fuel_flow * segment_duration
 
 
 def equivalent_fuel_flow_rate_at_sea_level(
-    fuel_flow_cruise: npt.NDArray[np.float64],
-    theta_amb: npt.NDArray[np.float64],
-    delta_amb: npt.NDArray[np.float64],
-    mach_num: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    fuel_flow_cruise: npt.NDArray[np.floating],
+    theta_amb: npt.NDArray[np.floating],
+    delta_amb: npt.NDArray[np.floating],
+    mach_num: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""Convert fuel mass flow rate at cruise conditions to equivalent flow rate at sea level.
 
     Refer to Eq. (40) in :cite:`duboisFuelFlowMethod22006`.
 
     Parameters
     ----------
-    fuel_flow_cruise : npt.NDArray[np.float64]
+    fuel_flow_cruise : npt.NDArray[np.floating]
         Fuel mass flow rate per engine, [:math:`kg s^{-1}`]
-    theta_amb : npt.NDArray[np.float64]
+    theta_amb : npt.NDArray[np.floating]
         Ratio of the ambient temperature to the temperature at mean sea-level.
-    delta_amb : npt.NDArray[np.float64]
+    delta_amb : npt.NDArray[np.floating]
         Ratio of the pressure altitude to the surface pressure.
-    mach_num : npt.NDArray[np.float64]
+    mach_num : npt.NDArray[np.floating]
         Mach number, [:math: `Ma`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Estimate of fuel flow per engine at sea level, [:math:`kg \ s^{-1}`].
 
     References
@@ -249,18 +249,18 @@ def equivalent_fuel_flow_rate_at_sea_level(
 
 
 def equivalent_fuel_flow_rate_at_cruise(
-    fuel_flow_sls: npt.NDArray[np.float64] | float,
+    fuel_flow_sls: npt.NDArray[np.floating] | float,
     theta_amb: ArrayOrFloat,
     delta_amb: ArrayOrFloat,
     mach_num: ArrayOrFloat,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     r"""Convert fuel mass flow rate at sea level to equivalent fuel flow rate at cruise conditions.
 
     Refer to Eq. (40) in :cite:`duboisFuelFlowMethod22006`.
 
     Parameters
     ----------
-    fuel_flow_sls : npt.NDArray[np.float64] | float
+    fuel_flow_sls : npt.NDArray[np.floating] | float
         Fuel mass flow rate, [:math:`kg s^{-1}`]
     theta_amb : ArrayOrFloat
         Ratio of the ambient temperature to the temperature at mean sea-level.
@@ -271,7 +271,7 @@ def equivalent_fuel_flow_rate_at_cruise(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Estimate of fuel mass flow rate at sea level, [:math:`kg \ s^{-1}`]
 
     References
@@ -285,28 +285,28 @@ def equivalent_fuel_flow_rate_at_cruise(
 
 
 def reserve_fuel_requirements(
-    rocd: npt.NDArray[np.float64],
-    altitude_ft: npt.NDArray[np.float64],
-    fuel_flow: npt.NDArray[np.float64],
-    fuel_burn: npt.NDArray[np.float64],
+    rocd: npt.NDArray[np.floating],
+    altitude_ft: npt.NDArray[np.floating],
+    fuel_flow: npt.NDArray[np.floating],
+    fuel_burn: npt.NDArray[np.floating],
 ) -> float:
     r"""
     Estimate reserve fuel requirements.
 
     Parameters
     ----------
-    rocd: npt.NDArray[np.float64]
+    rocd: npt.NDArray[np.floating]
         Rate of climb and descent, [:math:`ft \ min^{-1}`]
-    altitude_ft: npt.NDArray[np.float64]
+    altitude_ft: npt.NDArray[np.floating]
         Altitude, [:math:`ft`]
-    fuel_flow: npt.NDArray[np.float64]
+    fuel_flow: npt.NDArray[np.floating]
         Fuel mass flow rate, [:math:`kg \ s^{-1}`].
-    fuel_burn: npt.NDArray[np.float64]
+    fuel_burn: npt.NDArray[np.floating]
         Fuel consumption for each waypoint, [:math:`kg`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Reserve fuel requirements, [:math:`kg`]
 
     References
@@ -553,11 +553,11 @@ def update_aircraft_mass(
     operating_empty_weight: float,
     max_takeoff_weight: float,
     max_payload: float,
-    fuel_burn: npt.NDArray[np.float64],
+    fuel_burn: npt.NDArray[np.floating],
     total_reserve_fuel: float,
     load_factor: float,
     takeoff_mass: float | None,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Update aircraft mass based on the simulated total fuel consumption.
 
     Used internally for finding aircraft mass iteratively.
@@ -573,7 +573,7 @@ def update_aircraft_mass(
         Aircraft maximum take-off weight, [:math:`kg`].
     max_payload: float
         Aircraft maximum payload, [:math:`kg`]
-    fuel_burn: npt.NDArray[np.float64]
+    fuel_burn: npt.NDArray[np.floating]
         Fuel consumption for each waypoint, [:math:`kg`]
     total_reserve_fuel: float
         Total reserve fuel requirements, [:math:`kg`]
@@ -587,7 +587,7 @@ def update_aircraft_mass(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Updated aircraft mass, [:math:`kg`]
 
     See Also
@@ -761,30 +761,30 @@ def turbine_inlet_temperature(
 
 
 def thrust_force(
-    altitude: npt.NDArray[np.float64],
-    true_airspeed: npt.NDArray[np.float64],
-    segment_duration: npt.NDArray[np.float64],
-    aircraft_mass: npt.NDArray[np.float64],
-    F_drag: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    altitude: npt.NDArray[np.floating],
+    true_airspeed: npt.NDArray[np.floating],
+    segment_duration: npt.NDArray[np.floating],
+    aircraft_mass: npt.NDArray[np.floating],
+    F_drag: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
     r"""Calculate the thrust force at each waypoint.
 
     Parameters
     ----------
-    altitude : npt.NDArray[np.float64]
+    altitude : npt.NDArray[np.floating]
         Waypoint altitude, [:math:`m`]
-    true_airspeed : npt.NDArray[np.float64]
+    true_airspeed : npt.NDArray[np.floating]
         True airspeed, [:math:`m \ s^{-1}`]
-    segment_duration : npt.NDArray[np.float64]
+    segment_duration : npt.NDArray[np.floating]
         Time difference between waypoints, [:math:`s`]
-    aircraft_mass : npt.NDArray[np.float64]
+    aircraft_mass : npt.NDArray[np.floating]
         Aircraft mass, [:math:`kg`]
-    F_drag : npt.NDArray[np.float64]
+    F_drag : npt.NDArray[np.floating]
         Draft force, [:math:`N`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Thrust force, [:math:`N`]
 
     References
@@ -919,49 +919,49 @@ def air_to_fuel_ratio(
 # -------------------
 
 
-def temperature_ratio(T: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def temperature_ratio(T: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     """Calculate the ratio of ambient temperature relative to the temperature at mean sea level.
 
     Parameters
     ----------
-    T : npt.NDArray[np.float64]
+    T : npt.NDArray[np.floating]
         Air temperature, [:math:`K`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Ratio of the temperature to the temperature at mean sea-level (MSL).
     """
     return T / constants.T_msl
 
 
-def pressure_ratio(p: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def pressure_ratio(p: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     """Calculate the ratio of ambient pressure relative to the surface pressure.
 
     Parameters
     ----------
-    p : npt.NDArray[np.float64]
+    p : npt.NDArray[np.floating]
         Air pressure, [:math:`Pa`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Ratio of the pressure altitude to the surface pressure.
     """
     return p / constants.p_surface
 
 
-def density_ratio(rho: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def density_ratio(rho: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     r"""Calculate the ratio of air density relative to the air density at mean-sea-level.
 
     Parameters
     ----------
-    rho : npt.NDArray[np.float64]
+    rho : npt.NDArray[np.floating]
         Air density, [:math:`kg \ m^{3}`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Ratio of the density to the air density at mean sea-level (MSL).
     """
     return rho / constants.rho_msl
diff --git a/pycontrails/physics/units.py b/pycontrails/physics/units.py
index 869e88b80..583b6b270 100644
--- a/pycontrails/physics/units.py
+++ b/pycontrails/physics/units.py
@@ -106,13 +106,13 @@ def m_to_T_isa(h: ArrayScalarLike) -> ArrayScalarLike:
     return constants.T_msl + h_min * constants.T_lapse_rate  # type: ignore[return-value]
 
 
-def _low_altitude_m_to_pl(h: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def _low_altitude_m_to_pl(h: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     T_isa: np.ndarray = m_to_T_isa(h)
     power_term = -constants.g / (constants.T_lapse_rate * constants.R_d)
     return (constants.p_surface * (T_isa / constants.T_msl) ** power_term) / 100.0
 
 
-def _high_altitude_m_to_pl(h: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def _high_altitude_m_to_pl(h: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     T_tropopause_isa = m_to_T_isa(np.asarray(constants.h_tropopause))
     power_term = -constants.g / (constants.T_lapse_rate * constants.R_d)
     p_tropopause_isa = constants.p_surface * (T_tropopause_isa / constants.T_msl) ** power_term
@@ -121,17 +121,17 @@ def _high_altitude_m_to_pl(h: npt.NDArray[np.float64]) -> npt.NDArray[np.float64
 
 
 @support_arraylike
-def m_to_pl(h: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def m_to_pl(h: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     r"""Convert from altitude (m) to pressure level (hPa).
 
     Parameters
     ----------
-    h : npt.NDArray[np.float64]
+    h : npt.NDArray[np.floating]
         altitude, [:math:`m`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         pressure level, [:math:`hPa`], [:math:`mbar`]
 
     References
@@ -152,14 +152,14 @@ def m_to_pl(h: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
     return np.piecewise(h, condlist, funclist)
 
 
-def _low_altitude_pl_to_m(pl: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def _low_altitude_pl_to_m(pl: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     base = 100.0 * pl / constants.p_surface
     exponent = -constants.T_lapse_rate * constants.R_d / constants.g
     T_isa = constants.T_msl * base**exponent
     return (T_isa - constants.T_msl) / constants.T_lapse_rate
 
 
-def _high_altitude_pl_to_m(pl: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def _high_altitude_pl_to_m(pl: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     T_tropopause_isa = m_to_T_isa(np.asarray(constants.h_tropopause))
     power_term = -constants.g / (constants.T_lapse_rate * constants.R_d)
     p_tropopause_isa = constants.p_surface * (T_tropopause_isa / constants.T_msl) ** power_term
@@ -168,7 +168,7 @@ def _high_altitude_pl_to_m(pl: npt.NDArray[np.float64]) -> npt.NDArray[np.float6
 
 
 @support_arraylike
-def pl_to_m(pl: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+def pl_to_m(pl: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
     r"""Convert from pressure level (hPa) to altitude (m).
 
     Function is slightly different from the classical formula:
@@ -410,20 +410,20 @@ def tas_to_mach_number(true_airspeed: ArrayScalarLike, T: ArrayScalarLike) -> Ar
 
 
 def mach_number_to_tas(
-    mach_number: float | npt.NDArray[np.float64], T: float | npt.NDArray[np.float64]
-) -> float | npt.NDArray[np.float64]:
+    mach_number: float | npt.NDArray[np.floating], T: float | npt.NDArray[np.floating]
+) -> float | npt.NDArray[np.floating]:
     r"""Calculate true airspeed from the Mach number at a specified ambient temperature.
 
     Parameters
     ----------
-    mach_number : float | npt.NDArray[np.float64]
+    mach_number : float | npt.NDArray[np.floating]
         Mach number, [:math: `Ma`]
-    T : npt.NDArray[np.float64]
+    T : npt.NDArray[np.floating]
         Ambient temperature, [:math:`K`]
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         True airspeed, [:math:`m \ s^{-1}`]
 
     References
@@ -452,7 +452,7 @@ def lbs_to_kg(lbs: ArrayScalarLike) -> ArrayScalarLike:
 def dt_to_seconds(
     dt: npt.NDArray[np.timedelta64] | np.timedelta64,
     dtype: npt.DTypeLike = np.float64,
-) -> npt.NDArray[np.float64]:
+) -> npt.NDArray[np.floating]:
     """Convert a time delta to seconds as a float with specified ``dtype`` precision.
 
     Parameters
diff --git a/pycontrails/utils/types.py b/pycontrails/utils/types.py
index f049e3335..fde98b86d 100644
--- a/pycontrails/utils/types.py
+++ b/pycontrails/utils/types.py
@@ -17,16 +17,13 @@
 ArrayLike = TypeVar("ArrayLike", np.ndarray, xr.DataArray, xr.DataArray | np.ndarray)
 
 #: Array or Float (np.ndarray, float)
-ArrayOrFloat = TypeVar(
-    "ArrayOrFloat", npt.NDArray[np.float64], float, float | npt.NDArray[np.float64]
-)
+ArrayOrFloat = TypeVar("ArrayOrFloat", npt.NDArray[np.floating], float)
 
-#: Array like input (np.ndarray, xr.DataArray, np.float64, float)
+#: Array like input (np.ndarray, xr.DataArray, float)
 ArrayScalarLike = TypeVar(
     "ArrayScalarLike",
     np.ndarray,
     xr.DataArray,
-    np.float64,
     float,
     np.ndarray | float,
     xr.DataArray | np.ndarray,
@@ -44,7 +41,7 @@
 
 
 def support_arraylike(
-    func: Callable[[np.ndarray], np.ndarray],
+    func: Callable[[npt.NDArray[np.floating]], npt.NDArray[np.floating]],
 ) -> Callable[[ArrayScalarLike], ArrayScalarLike]:
     """Extend a numpy universal function operating on arrays of floats.
 
@@ -80,7 +77,7 @@ def wrapped(arr: ArrayScalarLike) -> ArrayScalarLike:
             return ret
 
         # Keep python native numeric types native
-        if isinstance(arr, float | int | np.float64):
+        if isinstance(arr, float | int):
             return ret.item()
 
         # Recreate pd.Series