pycontrails 0.49.3__cp312-cp312-win_amd64.whl → 0.49.5__cp312-cp312-win_amd64.whl

pycontrails/core/models.py

@@ -447,7 +447,8 @@ class Model(ABC):
 
         # Raise error if source is not a MetDataset or GeoVectorDataset
         if not isinstance(source, (MetDataset, GeoVectorDataset)):
-            raise TypeError(f"Unable to handle input eval data {source}")
+            msg = f"Unknown source type: {type(source)}"
+            raise TypeError(msg)
 
         if self.params["copy_source"]:
             source = source.copy()
@@ -516,11 +517,14 @@ class Model(ABC):
         TypeError
             Raised if :attr:`met` is not a :class:`MetDataset`.
         """
-        if not hasattr(self, "source"):
-            raise ValueError("Attribute `source` must be defined before calling `downselect_met`")
+        try:
+            source = self.source
+        except AttributeError as exc:
+            msg = "Attribute 'source' must be defined before calling 'downselect_met'."
+            raise AttributeError(msg) from exc
 
         # TODO: This could be generalized for a MetDataset source
-        if not isinstance(self.source, GeoVectorDataset):
+        if not isinstance(source, GeoVectorDataset):
             msg = "Attribute 'source' must be a GeoVectorDataset"
             raise TypeError(msg)
 
@@ -543,7 +547,7 @@ class Model(ABC):
         }
         kwargs = {k: v for k, v in buffers.items() if v is not None}
 
-        self.met = self.source.downselect_met(self.met, **kwargs, copy=False)
+        self.met = source.downselect_met(self.met, **kwargs, copy=False)
 
     def set_source_met(
         self,
@@ -614,7 +618,8 @@ class Model(ABC):
                 continue
 
             if not isinstance(self.source, MetDataset):
-                raise TypeError(f"Unknown source type: {type(self.source)}")
+                msg = f"Unknown source type: {type(self.source)}"
+                raise TypeError(msg)
 
             da = self.met.data[met_key].reset_coords(drop=True)
             try:
@@ -770,7 +775,8 @@ def _interp_grid_to_grid(
         interped = da.interp(coords, **interp_kwargs).load().astype(da.dtype, copy=False)
         return interped.assign_coords(level=level0)
 
-    raise NotImplementedError(f"Unsupported q_method: {q_method}")
+    msg = f"Unsupported q_method: {q_method}"
+    raise NotImplementedError(msg)
 
 
 def _raise_missing_met_var(var: MetVariable | Sequence[MetVariable]) -> NoReturn:
@@ -786,15 +792,17 @@ def _raise_missing_met_var(var: MetVariable | Sequence[MetVariable]) -> NoReturn
     KeyError
     """
     if isinstance(var, MetVariable):
-        raise KeyError(
+        msg = (
             f"Variable `{var.standard_name}` not found. Either pass parameter `met`"
             f"in model constructor, or define `{var.standard_name}` data on input data."
         )
+        raise KeyError(msg)
     missing_keys = [v.standard_name for v in var]
-    raise KeyError(
+    msg = (
         f"One of `{missing_keys}` is required. Either pass parameter `met`"
         f"in model constructor, or define one of `{missing_keys}` data on input data."
     )
+    raise KeyError(msg)
 
 
 def interpolate_met(
@@ -847,7 +855,8 @@ def interpolate_met(
         return out
 
     if met is None:
-        raise ValueError(f"No variable key '{vector_key}' in 'vector' and 'met' is None")
+        msg = f"No variable key '{vector_key}' in 'vector' and 'met' is None"
+        raise KeyError(msg)
 
     if met_key in ("q", "specific_humidity") and q_method is not None:
         mda, log_applied = _extract_q(met, met_key, q_method)
@@ -859,7 +868,8 @@ def interpolate_met(
         try:
             mda = met[met_key]
         except KeyError as exc:
-            raise KeyError(f"No variable key '{met_key}' in 'met'.") from exc
+            msg = f"No variable key '{met_key}' in 'met'."
+            raise KeyError(msg) from exc
 
         out = vector.intersect_met(mda, **interp_kwargs)
 
@@ -890,7 +900,8 @@ def _extract_q(met: MetDataset, met_key: str, q_method: str) -> tuple[MetDataArr
         try:
             return met[met_key], False
         except KeyError as exc:
-            raise KeyError(f"No variable key '{met_key}' in 'met'.") from exc
+            msg = f"No variable key '{met_key}' in 'met'."
+            raise KeyError(msg) from exc
 
     try:
         return met["log_specific_humidity"], True
@@ -904,7 +915,8 @@ def _extract_q(met: MetDataset, met_key: str, q_method: str) -> tuple[MetDataArr
     try:
         return met[met_key], False
     except KeyError as exc:
-        raise KeyError(f"No variable key '{met_key}' in 'met'.") from exc
+        msg = f"No variable key '{met_key}' in 'met'."
+        raise KeyError(msg) from exc
 
 
 def _prepare_q(
@@ -972,7 +984,8 @@ def _prepare_q_cubic_spline(
     da: xr.DataArray, level: npt.NDArray[np.float64]
 ) -> tuple[MetDataArray, npt.NDArray[np.float64]]:
     if da["level"][0] < 50.0 or da["level"][-1] > 1000.0:
-        raise ValueError("Cubic spline interpolation requires data to span 50-1000 hPa.")
+        msg = "Cubic spline interpolation requires data to span 50-1000 hPa."
+        raise ValueError(msg)
     ppoly = _load_spline()
 
     da = da.assign_coords(level=ppoly(da["level"]))
@@ -1038,7 +1051,8 @@ def raise_invalid_q_method_error(q_method: str) -> NoReturn:
         ``q_method`` is not one of ``None``, ``"log-q-log-p"``, or ``"cubic-spline"``.
     """
     available = None, "log-q-log-p", "cubic-spline"
-    raise ValueError(f"Invalid 'q_method' value '{q_method}'. Must be one of {available}.")
+    msg = f"Invalid 'q_method' value '{q_method}'. Must be one of {available}."
+    raise ValueError(msg)
 
 
 @functools.cache

pycontrails/core/vector.py

@@ -956,10 +956,8 @@ class VectorDataset:
                 continue
 
             min_dtype = np.min_scalar_type(scalar)
-            if np.can_cast(min_dtype, np.float32):
-                self.data.update({key: np.full(self.size, scalar, dtype=np.float32)})
-            else:
-                self.data.update({key: np.full(self.size, scalar)})
+            dtype = np.float32 if np.can_cast(min_dtype, np.float32) else None
+            self.data.update({key: np.full(self.size, scalar, dtype=dtype)})
 
     def broadcast_numeric_attrs(
         self, ignore_keys: str | Iterable[str] | None = None, overwrite: bool = False
@@ -1607,22 +1605,22 @@ class GeoVectorDataset(VectorDataset):
         npt.NDArray[np.bool_]
             True if point is inside the bounding box defined by ``met``.
         """
-        variables = met.variables
+        indexes = met.indexes
 
         lat_intersect = coordinates.intersect_domain(
-            variables["latitude"].values,
+            indexes["latitude"].to_numpy(),
             self["latitude"],
         )
         lon_intersect = coordinates.intersect_domain(
-            variables["longitude"].values,
+            indexes["longitude"].to_numpy(),
             self["longitude"],
         )
         level_intersect = coordinates.intersect_domain(
-            variables["level"].values,
+            indexes["level"].to_numpy(),
             self.level,
         )
         time_intersect = coordinates.intersect_domain(
-            variables["time"].values,
+            indexes["time"].to_numpy(),
             self["time"],
         )
 
@@ -1901,18 +1899,19 @@ class GeoVectorDataset(VectorDataset):
         MetDataset | MetDataArray
             Copy of downselected MetDataset or MetDataArray.
         """
+        variables = met.indexes
         lon_slice = coordinates.slice_domain(
-            met.variables["longitude"].values,
+            variables["longitude"].to_numpy(),
             self["longitude"],
             buffer=longitude_buffer,
         )
         lat_slice = coordinates.slice_domain(
-            met.variables["latitude"].values,
+            variables["latitude"].to_numpy(),
             self["latitude"],
             buffer=latitude_buffer,
         )
         time_slice = coordinates.slice_domain(
-            met.variables["time"].values,
+            variables["time"].to_numpy(),
             self["time"],
             buffer=time_buffer,
         )
@@ -1922,7 +1921,7 @@ class GeoVectorDataset(VectorDataset):
             level_slice = slice(None)
         else:
             level_slice = coordinates.slice_domain(
-                met.variables["level"].values,
+                variables["level"].to_numpy(),
                 self.level,
                 buffer=level_buffer,
             )
@@ -2043,8 +2042,8 @@ def vector_to_lon_lat_grid(
     >>> da = ds["foo"]
     >>> da.coords
     Coordinates:
-      * longitude  (longitude) float64 -10.0 -9.5 -9.0 -8.5 -8.0 ... 8.0 8.5 9.0 9.5
-      * latitude   (latitude) float64 -10.0 -9.5 -9.0 -8.5 -8.0 ... 8.0 8.5 9.0 9.5
+      * longitude  (longitude) float64 320B -10.0 -9.5 -9.0 -8.5 ... 8.0 8.5 9.0 9.5
+      * latitude   (latitude) float64 320B -10.0 -9.5 -9.0 -8.5 ... 8.0 8.5 9.0 9.5
 
     >>> da.values.round(2)
     array([[2.23, 0.67, 1.29, ..., 4.66, 3.91, 1.93],

pycontrails/datalib/gfs/gfs.py

@@ -134,7 +134,7 @@ class GFSForecast(datalib.MetDataSource):
         self,
         time: datalib.TimeInput | None,
         variables: datalib.VariableInput,
-        pressure_levels: datalib.PressureLevelInput = [-1],
+        pressure_levels: datalib.PressureLevelInput = -1,
         paths: str | list[str] | pathlib.Path | list[pathlib.Path] | None = None,
         grid: float = 0.25,
         forecast_time: DatetimeLike | None = None,

pycontrails/datalib/spire/spire.py

@@ -198,17 +198,19 @@ def identify_flights(messages: pd.DataFrame) -> pd.Series:
         index=messages.index,
     )
 
-    for idx, gp in messages[[
-        "icao_address",
-        "tail_number",
-        "aircraft_type_icao",
-        "callsign",
-        "timestamp",
-        "longitude",
-        "latitude",
-        "altitude_baro",
-        "on_ground",
-    ]].groupby(["icao_address", "tail_number", "aircraft_type_icao", "callsign"], sort=False):
+    for idx, gp in messages[
+        [
+            "icao_address",
+            "tail_number",
+            "aircraft_type_icao",
+            "callsign",
+            "timestamp",
+            "longitude",
+            "latitude",
+            "altitude_baro",
+            "on_ground",
+        ]
+    ].groupby(["icao_address", "tail_number", "aircraft_type_icao", "callsign"], sort=False):
         # minimum # of messages > TRAJECTORY_MINIMUM_MESSAGES
         if len(gp) < TRAJECTORY_MINIMUM_MESSAGES:
             logger.debug(f"Message {idx} group too small to create flight ids")
@@ -542,14 +544,16 @@ def validate_flights(messages: pd.DataFrame) -> pd.Series:
         index=messages.index,
     )
 
-    for _, gp in messages[[
-        "flight_id",
-        "aircraft_type_icao",
-        "timestamp",
-        "altitude_baro",
-        "on_ground",
-        "speed",
-    ]].groupby("flight_id", sort=False):
+    for _, gp in messages[
+        [
+            "flight_id",
+            "aircraft_type_icao",
+            "timestamp",
+            "altitude_baro",
+            "on_ground",
+            "speed",
+        ]
+    ].groupby("flight_id", sort=False):
         # save flight ids
         valid.loc[gp.index] = is_valid_trajectory(gp)
 

pycontrails/ext/synthetic_flight.py

@@ -162,7 +162,7 @@ class SyntheticFlight:
         msg += f"\nmax_queue_size: {self.max_queue_size}  min_n_waypoints: {self.min_n_waypoints}"
         return msg
 
-    def __call__(self, timestep: np.timedelta64 = np.timedelta64(1, "m")) -> Flight:
+    def __call__(self, timestep: np.timedelta64 | None = None) -> Flight:
         """Create random flight within `bounds` at a constant altitude.
 
         BADA4 data is used to determine flight speed at a randomly chosen altitude within
@@ -183,6 +183,8 @@ class SyntheticFlight:
         Flight
             Random `Flight` instance constrained by bounds.
         """
+        if timestep is None:
+            timestep = np.timedelta64(1, "m")
         # Building flights with `u_wind` and `v_wind` involved in the true airspeed calculation is
         # slow. BUT, we can do it in a vectorized way. So we maintain a short queue that gets
         # repeatedly replenished.

pycontrails/models/accf.py

@@ -312,10 +312,12 @@ class ACCF(Model):
         )
 
         if p_settings["lat_bound"] and p_settings["lon_bound"]:
-            ws.reduce_domain({
-                "latitude": p_settings["lat_bound"],
-                "longitude": p_settings["lon_bound"],
-            })
+            ws.reduce_domain(
+                {
+                    "latitude": p_settings["lat_bound"],
+                    "longitude": p_settings["lon_bound"],
+                }
+            )
 
         self.ds = ws.get_xarray()
         self.variable_names = ws.variable_names

pycontrails/models/cocip/cocip.py

@@ -302,7 +302,7 @@ class Cocip(Model):
         self,
         met: MetDataset,
         rad: MetDataset,
-        params: dict[str, Any] = {},
+        params: dict[str, Any] | None = None,
         **params_kwargs: Any,
     ) -> None:
         # call Model init
@@ -379,15 +379,14 @@ class Cocip(Model):
         self.source = self.require_source_type(Flight)
         return_flight_list = isinstance(self.source, Fleet) and isinstance(source, Sequence)
 
-        self._calc_timesteps()
+        self._set_timesteps()
 
         # Downselect met for CoCiP initialization
         # We only need to buffer in the negative vertical direction,
         # which is the positive direction for level
         logger.debug("Downselect met for Cocip initialization")
-        buffer = 0, self.params["met_level_buffer"][1]
-        met = self.source.downselect_met(self.met, copy=False, level_buffer=buffer)
-        # Add tau_cirrus if it doesn't exist already.
+        level_buffer = 0, self.params["met_level_buffer"][1]
+        met = self.source.downselect_met(self.met, level_buffer=level_buffer, copy=False)
         met = add_tau_cirrus(met)
 
         # Prepare flight for model
@@ -438,8 +437,11 @@ class Cocip(Model):
 
         return self.source
 
-    def _calc_timesteps(self) -> None:
-        """Calculate :attr:`timesteps`."""
+    def _set_timesteps(self) -> None:
+        """Set the :attr:`timesteps` based on the ``source`` time range.
+
+        This method is called in :meth:`eval` before the flight is processed.
+        """
         if isinstance(self.source, Fleet):
             # time not sorted in Fleet instance
             tmin = self.source["time"].min()
@@ -1011,11 +1013,11 @@ class Cocip(Model):
                 np.timedelta64(0, "ns"),
                 time_end - latest_contrail["time"].max(),
             )
-            if time_end > met.variables["time"].values[-1]:
+            if time_end > met.indexes["time"].to_numpy()[-1]:
                 logger.debug("Downselect met at time_end %s within Cocip evolution", time_end)
                 met = latest_contrail.downselect_met(self.met, **buffers, copy=False)
                 met = add_tau_cirrus(met)
-            if time_end > rad.variables["time"].values[-1]:
+            if time_end > rad.indexes["time"].to_numpy()[-1]:
                 logger.debug("Downselect rad at time_end %s within Cocip evolution", time_end)
                 rad = latest_contrail.downselect_met(self.rad, **buffers, copy=False)
 
@@ -1446,20 +1448,46 @@ def _process_rad(rad: MetDataset) -> MetDataset:
             )
             raise ValueError(msg) from exc
         if radiation_accumulated:
+
+            # Don't assume that radiation data is uniformly spaced in time
+            # Instead, infer the appropriate time shift
+            time_diff = rad.data["time"].diff("time", label="upper")
+            time_shift = -time_diff / 2
+
             # Keep the original attrs -- we need these later on
             old_attrs = {k: v.attrs for k, v in rad.data.items()}
 
+            # Also need to keep dataset-level attrs, which are lost
+            # when dividing a Dataset by a DataArray
+            old_rad_attrs = rad.data.attrs
+
             # NOTE: Taking the diff will remove the first time step
             # This is typically what we want (forecast step 0 is all zeros)
             # But, if the data has been downselected for a particular Flight / Fleet,
             # we lose the first time step of the data.
-            rad.data = rad.data.diff("time", label="upper")
+            #
+            # Other portions of the code convert HRES accumulated fluxes (J/m2)
+            # to averaged fluxes (W/m2) assuming that accumulations are over
+            # one hour. For those conversions to work correctly, must normalize
+            # accumulations by number of hours between steps
+            time_diff_h = time_diff / np.timedelta64(1, "h")
+            rad.data = rad.data.diff("time", label="upper") / time_diff_h
+            rad.data.attrs = old_rad_attrs
 
             # Add back the original attrs
             for k, v in rad.data.items():
                 v.attrs = old_attrs[k]
 
-        shift_radiation_time = -np.timedelta64(30, "m")
+            # Short-circuit to avoid idiot check
+            rad.data = rad.data.assign_coords({"time": rad.data["time"] + time_shift})
+            if np.unique(time_shift).size == 1:
+                rad.data["time"].attrs["shift_radiation_time"] = str(time_shift.values[0])
+            else:
+                rad.data["time"].attrs["shift_radiation_time"] = "variable"
+            return rad
+
+        else:
+            shift_radiation_time = -np.timedelta64(30, "m")
 
     elif dataset == "ERA5" and product == "ensemble":
         shift_radiation_time = -np.timedelta64(90, "m")
@@ -1507,17 +1535,19 @@ def _eval_aircraft_performance(
         If ``aircraft_performance`` is None
     """
 
-    aircraft_performance_outputs = "wingspan", "engine_efficiency", "fuel_flow", "aircraft_mass"
-    if flight.ensure_vars(aircraft_performance_outputs, False):
+    ap_vars = {"wingspan", "engine_efficiency", "fuel_flow", "aircraft_mass"}
+    missing = ap_vars.difference(flight).difference(flight.attrs)
+    if not missing:
         return flight
 
     if aircraft_performance is None:
-        raise ValueError(
-            "An AircraftPerformance model parameter is required if the flight "
-            f"does not contain the following variables: {aircraft_performance_outputs}. "
-            "For example, instantiate the Cocip model with "
-            "'Cocip(..., aircraft_performance=PSFlight(...))'."
+        msg = (
+            f"An AircraftPerformance model parameter is required if the flight does not contain "
+            f"the following variables: {ap_vars}. This flight is missing: {missing}. "
+            "Instantiate the Cocip model with an AircraftPerformance model. "
+            "For example, 'Cocip(..., aircraft_performance=PSFlight(...))'."
         )
+        raise ValueError(msg)
 
     return aircraft_performance.eval(source=flight, copy_source=False)
 

pycontrails/models/cocip/cocip_params.py

@@ -35,16 +35,18 @@ def _habit_distributions() -> npt.NDArray[np.float32]:
 
 
 def _habits() -> npt.NDArray[np.str_]:
-    return np.array([
-        "Sphere",
-        "Solid column",
-        "Hollow column",
-        "Rough aggregate",
-        "Rosette-6",
-        "Plate",
-        "Droxtal",
-        "Myhre",
-    ])
+    return np.array(
+        [
+            "Sphere",
+            "Solid column",
+            "Hollow column",
+            "Rough aggregate",
+            "Rosette-6",
+            "Plate",
+            "Droxtal",
+            "Myhre",
+        ]
+    )
 
 
 @dataclasses.dataclass
@@ -184,6 +186,7 @@ class CocipParams(ModelParams):
     #: Experimental. Radiative effects due to contrail-contrail overlapping
     #: Account for change in local contrail shortwave and longwave radiative forcing
     #: due to contrail-contrail overlapping.
+    #:
     #: .. versionadded:: 0.45
     contrail_contrail_overlapping: bool = False
 

pycontrails/models/cocip/output_formats.py

@@ -293,16 +293,18 @@ def longitude_latitude_grid(
     """
     # Ensure the required columns are included in `flight_waypoints`, `contrails` and `met`
     flight_waypoints.ensure_vars(("segment_length", "ef"))
-    contrails.ensure_vars((
-        "formation_time",
-        "segment_length",
-        "width",
-        "tau_contrail",
-        "rf_sw",
-        "rf_lw",
-        "rf_net",
-        "ef",
-    ))
+    contrails.ensure_vars(
+        (
+            "formation_time",
+            "segment_length",
+            "width",
+            "tau_contrail",
+            "rf_sw",
+            "rf_lw",
+            "rf_net",
+            "ef",
+        )
+    )
     met.ensure_vars(
         ("air_temperature", "specific_humidity", "specific_cloud_ice_water_content", "geopotential")
     )
@@ -857,35 +859,39 @@ def time_slice_statistics(
     - ``mean_albedo_at_contrail_wypts``, [dimensionless]
     """
     # Ensure the required columns are included in `flight_waypoints`, `contrails`, `met` and `rad`
-    flight_waypoints.ensure_vars((
-        "flight_id",
-        "segment_length",
-        "true_airspeed",
-        "fuel_flow",
-        "engine_efficiency",
-        "nvpm_ei_n",
-        "sac",
-        "persistent_1",
-    ))
-    contrails.ensure_vars((
-        "flight_id",
-        "segment_length",
-        "air_temperature",
-        "iwc",
-        "r_ice_vol",
-        "n_ice_per_m",
-        "tau_contrail",
-        "tau_cirrus",
-        "width",
-        "area_eff",
-        "sdr",
-        "rsr",
-        "olr",
-        "rf_sw",
-        "rf_lw",
-        "rf_net",
-        "ef",
-    ))
+    flight_waypoints.ensure_vars(
+        (
+            "flight_id",
+            "segment_length",
+            "true_airspeed",
+            "fuel_flow",
+            "engine_efficiency",
+            "nvpm_ei_n",
+            "sac",
+            "persistent_1",
+        )
+    )
+    contrails.ensure_vars(
+        (
+            "flight_id",
+            "segment_length",
+            "air_temperature",
+            "iwc",
+            "r_ice_vol",
+            "n_ice_per_m",
+            "tau_contrail",
+            "tau_cirrus",
+            "width",
+            "area_eff",
+            "sdr",
+            "rsr",
+            "olr",
+            "rf_sw",
+            "rf_lw",
+            "rf_net",
+            "ef",
+        )
+    )
 
     # Ensure that the waypoints are within `t_start` and `t_end`
     is_in_time = flight_waypoints.dataframe["time"].between(t_start, t_end, inclusive="right")
@@ -920,12 +926,14 @@ def time_slice_statistics(
 
     # Meteorology domain statistics
     if met is not None:
-        met.ensure_vars((
-            "air_temperature",
-            "specific_humidity",
-            "specific_cloud_ice_water_content",
-            "geopotential",
-        ))
+        met.ensure_vars(
+            (
+                "air_temperature",
+                "specific_humidity",
+                "specific_cloud_ice_water_content",
+                "geopotential",
+            )
+        )
         met = met.downselect(spatial_bbox)
         met_stats = meteorological_time_slice_statistics(t_end, contrails, met, humidity_scaling)
 
@@ -1929,13 +1937,15 @@ def natural_cirrus_properties_to_hi_res_grid(
         relatively narrow contrails.
     """
     # Ensure the required columns are included in `met`
-    met.ensure_vars((
-        "air_temperature",
-        "specific_humidity",
-        "specific_cloud_ice_water_content",
-        "geopotential",
-        "fraction_of_cloud_cover",
-    ))
+    met.ensure_vars(
+        (
+            "air_temperature",
+            "specific_humidity",
+            "specific_cloud_ice_water_content",
+            "geopotential",
+            "fraction_of_cloud_cover",
+        )
+    )
 
     # Ensure `met` only contains one time step, constraint can be relaxed in the future.
     if len(met["time"].data) > 1: