pycontrails 0.54.3__cp313-cp313-win_amd64.whl → 0.54.5__cp313-cp313-win_amd64.whl

pycontrails/datalib/ecmwf/__init__.py

@@ -40,21 +40,21 @@ from pycontrails.datalib.ecmwf.variables import (
 )
 
 __all__ = [
-    "ERA5ARCO",
-    "CDSCredentialsNotFound",
+    "ECMWF_VARIABLES",
     "ERA5",
-    "ERA5ModelLevel",
+    "ERA5ARCO",
     "HRES",
-    "HRESModelLevel",
     "IFS",
-    "model_level_reference_pressure",
-    "model_level_pressure",
-    "ml_to_pl",
-    "open_arco_era5_model_level_data",
-    "open_arco_era5_single_level",
+    "MODEL_LEVELS_PATH",
+    "MODEL_LEVEL_VARIABLES",
+    "PRESSURE_LEVEL_VARIABLES",
+    "SURFACE_VARIABLES",
+    "CDSCredentialsNotFound",
     "CloudAreaFraction",
     "CloudAreaFractionInLayer",
     "Divergence",
+    "ERA5ModelLevel",
+    "HRESModelLevel",
     "OzoneMassMixingRatio",
     "PotentialVorticity",
     "RelativeHumidity",
@@ -65,9 +65,9 @@ __all__ = [
     "TOAIncidentSolarRadiation",
     "TopNetSolarRadiation",
     "TopNetThermalRadiation",
-    "ECMWF_VARIABLES",
-    "MODEL_LEVELS_PATH",
-    "MODEL_LEVEL_VARIABLES",
-    "PRESSURE_LEVEL_VARIABLES",
-    "SURFACE_VARIABLES",
+    "ml_to_pl",
+    "model_level_pressure",
+    "model_level_reference_pressure",
+    "open_arco_era5_model_level_data",
+    "open_arco_era5_single_level",
 ]

pycontrails/datalib/ecmwf/common.py

@@ -61,7 +61,7 @@ class ECMWFAPI(metsource.MetDataSource):
 
         # 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)

pycontrails/datalib/ecmwf/era5.py

@@ -137,10 +137,10 @@ class ERA5(ECMWFAPI):
     """
 
     __slots__ = (
-        "product_type",
         "cds",
-        "url",
         "key",
+        "product_type",
+        "url",
     )
 
     #: Product type, one of "reanalysis", "ensemble_mean", "ensemble_members", "ensemble_spread"
@@ -319,9 +319,9 @@ class ERA5(ECMWFAPI):
         str
             ERA5 dataset name in CDS
         """
-        if self.pressure_levels != [-1]:
-            return "reanalysis-era5-pressure-levels"
-        return "reanalysis-era5-single-levels"
+        if self.is_single_level:
+            return "reanalysis-era5-single-levels"
+        return "reanalysis-era5-pressure-levels"
 
     def create_cachepath(self, t: datetime | pd.Timestamp) -> str:
         """Return cachepath to local ERA5 data file based on datetime.
@@ -539,9 +539,9 @@ class ERA5(ECMWFAPI):
             LOG.debug("Input dataset processed with pycontrails > 0.29")
             return ds
 
-        # For "reanalysis-era5-single-levels" or if self.pressure_levels length == 1,
+        # For "reanalysis-era5-single-levels"
         # then the netcdf file does not contain the dimension "level"
-        if len(self.pressure_levels) == 1:
+        if self.is_single_level:
             ds = ds.expand_dims(level=self.pressure_levels)
 
         # New CDS-Beta gives "valid_time" instead of "time"

pycontrails/datalib/ecmwf/hres.py

@@ -228,7 +228,7 @@ class HRES(ECMWFAPI):
     ... )
     """
 
-    __slots__ = ("server", "stream", "field_type", "forecast_time", "url", "key", "email")
+    __slots__ = ("email", "field_type", "forecast_time", "key", "server", "stream", "url")
 
     #: stream type, "oper" = atmospheric model/HRES, "enfo" = ensemble forecast.
     stream: str
@@ -691,7 +691,7 @@ class HRES(ECMWFAPI):
 
         # 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 @@ class HRES(ECMWFAPI):
         # 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)
 

pycontrails/datalib/ecmwf/ifs.py

@@ -149,7 +149,7 @@ class IFS(metsource.MetDataSource):
         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)

pycontrails/datalib/gfs/__init__.py

@@ -16,13 +16,13 @@ from pycontrails.datalib.gfs.variables import (
 
 __all__ = [
     "GFS_FORECAST_BUCKET",
-    "GFSForecast",
-    "CloudIceWaterMixingRatio",
-    "TotalCloudCoverIsobaric",
-    "Visibility",
-    "TOAUpwardShortwaveRadiation",
-    "TOAUpwardLongwaveRadiation",
     "GFS_VARIABLES",
     "PRESSURE_LEVEL_VARIABLES",
     "SURFACE_VARIABLES",
+    "CloudIceWaterMixingRatio",
+    "GFSForecast",
+    "TOAUpwardLongwaveRadiation",
+    "TOAUpwardShortwaveRadiation",
+    "TotalCloudCoverIsobaric",
+    "Visibility",
 ]

pycontrails/datalib/gfs/gfs.py

@@ -125,7 +125,7 @@ class GFSForecast(metsource.MetDataSource):
     - `GFS Documentation <https://www.emc.ncep.noaa.gov/emc/pages/numerical_forecast_systems/gfs/documentation.php>`_
     """
 
-    __slots__ = ("client", "grid", "cachestore", "show_progress", "forecast_time", "cache_download")
+    __slots__ = ("cache_download", "cachestore", "client", "forecast_time", "grid", "show_progress")
 
     #: S3 client for accessing GFS bucket
     client: botocore.client.S3
@@ -597,7 +597,7 @@ class GFSForecast(metsource.MetDataSource):
         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

pycontrails/datalib/goes.py

@@ -384,7 +384,7 @@ class GOES:
            [277.24512, 277.45377, 278.18408, ..., 274.6369 , 274.01093,
             274.06308],
            [276.8278 , 277.14078, 277.7146 , ..., 274.6369 , 273.9066 ,
-            274.16742]], dtype=float32)
+            274.16742]], shape=(500, 500), dtype=float32)
 
     """
 
@@ -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)

pycontrails/ext/empirical_grid.py

@@ -118,7 +118,7 @@ class EmpiricalGrid(AircraftPerformanceGrid):
 
         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()

pycontrails/models/apcemm/apcemm.py

@@ -277,13 +277,13 @@ class APCEMM(models.Model):
     """
 
     __slots__ = (
-        "apcemm_path",
+        "_trajectory_downsampling",
         "apcemm_input_params",
+        "apcemm_path",
         "cachestore",
+        "contrail",
         "trajectories",
         "vortex",
-        "contrail",
-        "_trajectory_downsampling",
     )
 
     name = "apcemm"
@@ -474,7 +474,7 @@ class APCEMM(models.Model):
             for coord in ("longitude", "latitude", "level")
         }
         buffers["time_buffer"] = (0, self.params["max_age"] + self.params["dt_lagrangian"])
-        met = self.source.downselect_met(self.met, **buffers, copy=False)
+        met = self.source.downselect_met(self.met, **buffers)
         model = DryAdvection(
             met=met,
             dt_integration=self.params["dt_lagrangian"],
@@ -816,7 +816,7 @@ class APCEMM(models.Model):
         # Ensure required met data is present.
         # No buffers needed for interpolation!
         vars = ap_model.met_variables + ap_model.optional_met_variables + emissions.met_variables
-        met = self.source.downselect_met(self.met, copy=False)
+        met = self.source.downselect_met(self.met)
         met.ensure_vars(vars)
         met.standardize_variables(vars)
         for var in vars:

pycontrails/models/apcemm/utils.py

@@ -214,7 +214,7 @@ def generate_apcemm_input_met(
     )
 
     # Downselect met before interpolation
-    met = vector.downselect_met(met, copy=False)
+    met = vector.downselect_met(met)
 
     # Interpolate meteorology data onto vector
     scale_humidity = humidity_scaling is not None and "specific_humidity" not in vector

pycontrails/models/cocip/__init__.py

@@ -15,9 +15,10 @@ from pycontrails.models.cocip.output_formats import (
 
 __all__ = [
     "Cocip",
+    "CocipFlightParams",
     "CocipParams",
     "CocipUncertaintyParams",
-    "CocipFlightParams",
+    "compare_cocip_with_goes",
     "contrail_flight_summary_statistics",
     "contrails_to_hi_res_grid",
     "flight_waypoint_summary_statistics",
@@ -25,5 +26,4 @@ __all__ = [
     "longitude_latitude_grid",
     "natural_cirrus_properties_to_hi_res_grid",
     "time_slice_statistics",
-    "compare_cocip_with_goes",
 ]

pycontrails/models/cocip/cocip.py

@@ -195,14 +195,14 @@ class Cocip(Model):
     """
 
     __slots__ = (
-        "rad",
+        "_downwash_contrail",
+        "_downwash_flight",
+        "_sac_flight",
         "contrail",
         "contrail_dataset",
         "contrail_list",
+        "rad",
         "timesteps",
-        "_sac_flight",
-        "_downwash_flight",
-        "_downwash_contrail",
     )
 
     name = "cocip"
@@ -391,7 +391,7 @@ class Cocip(Model):
         # which is the positive direction for level
         logger.debug("Downselect met for Cocip initialization")
         level_buffer = 0, self.params["met_level_buffer"][1]
-        met = self.source.downselect_met(self.met, level_buffer=level_buffer, copy=False)
+        met = self.source.downselect_met(self.met, level_buffer=level_buffer)
         met = add_tau_cirrus(met)
 
         # Prepare flight for model
@@ -660,7 +660,7 @@ class Cocip(Model):
             attrs = self.source.attrs
             attrs.pop("fl_attrs", None)
             attrs.pop("data_keys", None)
-            self.source = Fleet.from_seq(fls, broadcast_numeric=False, copy=False, attrs=attrs)
+            self.source = Fleet.from_seq(fls, broadcast_numeric=False, attrs=attrs)
 
         # Single flight
         else:
@@ -976,9 +976,9 @@ class Cocip(Model):
             for coord in ("longitude", "latitude", "level")
         }
         logger.debug("Downselect met for start of Cocip evolution")
-        met = self._downwash_contrail.downselect_met(self.met, **buffers, copy=False)
+        met = self._downwash_contrail.downselect_met(self.met, **buffers)
         met = add_tau_cirrus(met)
-        rad = self._downwash_contrail.downselect_met(self.rad, **buffers, copy=False)
+        rad = self._downwash_contrail.downselect_met(self.rad, **buffers)
 
         calc_continuous(self._downwash_contrail)
         calc_timestep_geometry(self._downwash_contrail)
@@ -1135,11 +1135,11 @@ class Cocip(Model):
             & (self._downwash_flight["time"] <= lookahead),
             copy=False,
         )
-        vector = GeoVectorDataset(
+        vector = GeoVectorDataset._from_fastpath(
             {
                 key: np.concatenate((latest_contrail[key], future_contrails[key]))
                 for key in ("longitude", "latitude", "level", "time")
-            }
+            },
         )
 
         # compute time buffer to ensure downselection extends to time_end
@@ -1152,7 +1152,7 @@ class Cocip(Model):
             max(np.timedelta64(0, "ns"), time_end - vector["time"].max()),
         )
 
-        return vector.downselect_met(met, **buffers, copy=False)
+        return vector.downselect_met(met, **buffers)
 
     def _create_downwash_contrail(self) -> GeoVectorDataset:
         """Get Contrail representation of downwash flight."""
@@ -1180,7 +1180,7 @@ class Cocip(Model):
             "persistent": self._downwash_flight["persistent_1"],
         }
 
-        contrail = GeoVectorDataset(downwash_contrail_data, copy=True)
+        contrail = GeoVectorDataset._from_fastpath(downwash_contrail_data).copy()
         contrail["formation_time"] = contrail["time"].copy()
         contrail["age"] = contrail["formation_time"] - contrail["time"]
 
@@ -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.
@@ -2300,7 +2300,7 @@ def calc_timestep_contrail_evolution(
     level_2 = geo.advect_level(level_1, vertical_velocity_1, rho_air_1, terminal_fall_speed_1, dt)
     altitude_2 = units.pl_to_m(level_2)
 
-    contrail_2 = GeoVectorDataset(
+    contrail_2 = GeoVectorDataset._from_fastpath(
         {
             "waypoint": waypoint_2,
             "flight_id": contrail_1["flight_id"],
@@ -2312,7 +2312,6 @@ def calc_timestep_contrail_evolution(
             "altitude": altitude_2,
             "level": level_2,
         },
-        copy=False,
     )
     intersection = contrail_2.coords_intersect_met(met)
     if not np.any(intersection):
@@ -2525,8 +2524,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 +2536,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

pycontrails/models/cocip/cocip_params.py

@@ -11,7 +11,7 @@ import numpy as np
 import numpy.typing as npt
 
 from pycontrails.core.aircraft_performance import AircraftPerformance
-from pycontrails.core.models import ModelParams
+from pycontrails.core.models import AdvectionBuffers
 from pycontrails.models.emissions.emissions import EmissionsParams
 from pycontrails.models.humidity_scaling import HumidityScaling
 
@@ -50,7 +50,7 @@ def _habits() -> npt.NDArray[np.str_]:
 
 
 @dataclasses.dataclass
-class CocipParams(ModelParams):
+class CocipParams(AdvectionBuffers):
     """Model parameters required by the CoCiP models."""
 
     # -------------------------
@@ -116,15 +116,6 @@ class CocipParams(ModelParams):
     #: evaluation after the met data is downselected.
     compute_tau_cirrus_in_model_init: bool | str = "auto"
 
-    #: Met longitude [WGS84] buffer for Cocip evolution.
-    met_longitude_buffer: tuple[float, float] = (10.0, 10.0)
-
-    #: Met latitude buffer [WGS84] for Cocip evolution.
-    met_latitude_buffer: tuple[float, float] = (10.0, 10.0)
-
-    #: Met level buffer [:math:`hPa`] for Cocip initialization and evolution.
-    met_level_buffer: tuple[float, float] = (40.0, 40.0)
-
     # ---------
     # Filtering
     # ---------

pycontrails/models/cocip/cocip_uncertainty.py

@@ -2,8 +2,8 @@
 
 from __future__ import annotations
 
+import dataclasses
 import logging
-from dataclasses import asdict, dataclass
 from typing import Any, ClassVar
 
 import numpy as np
@@ -70,7 +70,7 @@ class habit_dirichlet(rv_frozen):
         return habit_weights
 
 
-@dataclass
+@dataclasses.dataclass
 class CocipUncertaintyParams(CocipParams):
     """Model parameters for CoCiP epistemic uncertainty.
 
@@ -117,26 +117,30 @@ class CocipUncertaintyParams(CocipParams):
     #: Schumann takes ``wind_shear_enhancement_exponent`` = 0.5 and discusses the case of 0 and 2/3
     #: as possibilities.
     #: With a value of 0, wind shear is not enhanced.
-    wind_shear_enhancement_exponent_uncertainty: rv_frozen | None = stats.triang(
-        loc=0.0, c=CocipParams.wind_shear_enhancement_exponent, scale=1.0
+    wind_shear_enhancement_exponent_uncertainty: rv_frozen | None = dataclasses.field(
+        default_factory=lambda: stats.triang(
+            loc=0.0, c=CocipParams.wind_shear_enhancement_exponent, scale=1.0
+        )
     )
 
     #: Schumann takes ``initial_wake_vortex_depth`` = 0.5 and discusses some
     #: uncertainty in this value. This parameter should be non-negative.
-    initial_wake_vortex_depth_uncertainty: rv_frozen | None = stats.triang(
-        loc=0.3, c=CocipParams.initial_wake_vortex_depth, scale=0.4
+    initial_wake_vortex_depth_uncertainty: rv_frozen | None = dataclasses.field(
+        default_factory=lambda: stats.triang(
+            loc=0.3, c=CocipParams.initial_wake_vortex_depth, scale=0.4
+        )
     )
 
     #: Schumann takes a default value of 0.1 and describes it as an "important adjustable parameter"
     #: Currently, `CocipParams` uses a default value of 0.5
-    sedimentation_impact_factor_uncertainty: rv_frozen | None = stats.norm(
-        loc=CocipParams.sedimentation_impact_factor, scale=0.1
+    sedimentation_impact_factor_uncertainty: rv_frozen | None = dataclasses.field(
+        default_factory=lambda: stats.norm(loc=CocipParams.sedimentation_impact_factor, scale=0.1)
     )
 
     #: Teoh 2022 (to appear) takes values between 70% decrease and 100% increase.
     #: This coincides with the log normal distribution defined below.
-    nvpm_ei_n_enhancement_factor_uncertainty: rv_frozen | None = stats.lognorm(
-        s=0.15, scale=1 / stats.lognorm(s=0.15).mean()
+    nvpm_ei_n_enhancement_factor_uncertainty: rv_frozen | None = dataclasses.field(
+        default_factory=lambda: stats.lognorm(s=0.15, scale=1 / stats.lognorm(s=0.15).mean())
     )
 
     #: Scale shortwave radiative forcing.
@@ -145,8 +149,8 @@ class CocipUncertaintyParams(CocipParams):
     #: by `libRadTran <http://www.libradtran.org/doku.php>`_.
     #: We use the average RMS error across all habit types (pg 1397) as the standard deviation
     #: of a normally distributed scaling factor for SW forcing
-    rf_sw_enhancement_factor_uncertainty: rv_frozen | None = stats.norm(
-        loc=CocipParams.rf_sw_enhancement_factor, scale=0.106
+    rf_sw_enhancement_factor_uncertainty: rv_frozen | None = dataclasses.field(
+        default_factory=lambda: stats.norm(loc=CocipParams.rf_sw_enhancement_factor, scale=0.106)
     )
 
     #: Scale longwave radiative forcing.
@@ -154,8 +158,8 @@ class CocipUncertaintyParams(CocipParams):
     #: fit to the data generated by `libRadTran <http://www.libradtran.org/doku.php>`_.
     #: We use the average RMS error across all habit types (pg 1397) as the standard deviation
     #: of a normally distributed scaling factor for LW forcing.
-    rf_lw_enhancement_factor_uncertainty: rv_frozen | None = stats.norm(
-        loc=CocipParams.rf_lw_enhancement_factor, scale=0.071
+    rf_lw_enhancement_factor_uncertainty: rv_frozen | None = dataclasses.field(
+        default_factory=lambda: stats.norm(loc=CocipParams.rf_lw_enhancement_factor, scale=0.071)
     )
 
     #: Scale the habit distributions by a dirichlet distribution
@@ -163,7 +167,9 @@ class CocipUncertaintyParams(CocipParams):
     #: where :math:`\text{G}_{i}` is the approximate habit weight distributions
     #: defined in :attr:`CocipParams().habit_distributions`.
     #: Higher values of :math:`\text{C}` correspond to higher confidence in initial estimates.
-    habit_distributions_uncertainty: rv_frozen | None = habit_dirichlet(C=96.0)
+    habit_distributions_uncertainty: rv_frozen | None = dataclasses.field(
+        default_factory=habit_dirichlet
+    )
 
     def __post_init__(self) -> None:
         """Override values of model parameters according to ranges."""
@@ -192,7 +198,7 @@ class CocipUncertaintyParams(CocipParams):
 
         out = {}
 
-        param_dict = asdict(self)
+        param_dict = dataclasses.asdict(self)
         for uncertainty_param, dist in param_dict.items():
             if uncertainty_param.endswith("_uncertainty") and dist is not None:
                 param = uncertainty_param.split("_uncertainty")[0]
@@ -212,7 +218,7 @@ class CocipUncertaintyParams(CocipParams):
 
         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 +230,7 @@ class CocipUncertaintyParams(CocipParams):
 
         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.