pycontrails 0.54.2__cp310-cp310-win_amd64.whl → 0.54.4__cp310-cp310-win_amd64.whl

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 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
         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 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
         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 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
 
         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 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
 
     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 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
     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 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
         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.
         """
@@ -254,14 +254,14 @@ def _pick_method(scipy_version: str, method: str) -> str:
         raise ValueError(msg) from exc
 
     reimplemented_methods = ["slinear", "cubic", "quintic"]
-    if major > 1 or (major == 1 and minor >= 13) and method in reimplemented_methods:
+    if major > 1 or ((major == 1 and minor >= 13) and method in reimplemented_methods):
         return method + "_legacy"
     return method
 
 
 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 @@ class EmissionsProfileInterpolator:
             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 @@ class EmissionsProfileInterpolator:
         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 @@ class EmissionsProfileInterpolator:
 
         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)))

pycontrails/core/met.py

@@ -1495,9 +1495,9 @@ class MetDataArray(MetBase):
     @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 @@ class MetDataArray(MetBase):
         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 @@ class MetDataArray(MetBase):
         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 @@ class MetDataArray(MetBase):
         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 @@ class MetDataArray(MetBase):
 
         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 @@ class MetDataArray(MetBase):
 
     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 @@ class MetDataArray(MetBase):
                 )
                 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 @@ class MetDataArray(MetBase):
         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:

pycontrails/core/models.py

@@ -135,6 +135,20 @@ class ModelParams:
         return {(name := field.name): getattr(self, name) for field in fields(self)}
 
 
+@dataclass
+class AdvectionBuffers(ModelParams):
+    """Override buffers in :class:`ModelParams` for advection models."""
+
+    #: Met longitude [WGS84] buffer for evolution by advection.
+    met_longitude_buffer: tuple[float, float] = (10.0, 10.0)
+
+    #: Met latitude buffer [WGS84] for evolution by advection.
+    met_latitude_buffer: tuple[float, float] = (10.0, 10.0)
+
+    #: Met level buffer [:math:`hPa`] for evolution by advection.
+    met_level_buffer: tuple[float, float] = (40.0, 40.0)
+
+
 # ------
 # Models
 # ------
@@ -146,7 +160,7 @@ class Model(ABC):
     Implementing classes must implement the :meth:`eval` method
     """
 
-    __slots__ = ("params", "met", "source")
+    __slots__ = ("met", "params", "source")
 
     #: Default model parameter dataclass
     default_params: type[ModelParams] = ModelParams
@@ -825,7 +839,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 +868,7 @@ def interpolate_met(
 
     Returns
     -------
-    npt.NDArray[np.float64]
+    npt.NDArray[np.floating]
         Interpolated values.
 
     Raises
@@ -933,15 +947,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 +966,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 +989,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 +1008,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)