pycontrails 0.54.1__cp310-cp310-win_amd64.whl → 0.54.3__cp310-cp310-win_amd64.whl

pycontrails/datalib/gfs/gfs.py

@@ -9,18 +9,24 @@ References
 
 from __future__ import annotations
 
+import contextlib
 import hashlib
 import logging
 import pathlib
+import sys
 import warnings
 from collections.abc import Callable
 from datetime import datetime
 from typing import TYPE_CHECKING, Any
 
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
 import numpy as np
 import pandas as pd
 import xarray as xr
-from overrides import overrides
 
 import pycontrails
 from pycontrails.core import cache, met
@@ -82,6 +88,9 @@ class GFSForecast(metsource.MetDataSource):
     show_progress : bool, optional
         Show progress when downloading files from GFS AWS Bucket.
         Defaults to False
+    cache_download: bool, optional
+        If True, cache downloaded grib files rather than storing them in a temporary file.
+        By default, False.
 
     Examples
     --------
@@ -116,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")
+    __slots__ = ("client", "grid", "cachestore", "show_progress", "forecast_time", "cache_download")
 
     #: S3 client for accessing GFS bucket
     client: botocore.client.S3
@@ -142,7 +151,8 @@ class GFSForecast(metsource.MetDataSource):
         forecast_time: DatetimeLike | None = None,
         cachestore: cache.CacheStore | None = __marker,  # type: ignore[assignment]
         show_progress: bool = False,
-    ):
+        cache_download: bool = False,
+    ) -> None:
         try:
             import boto3
         except ModuleNotFoundError as e:
@@ -169,6 +179,7 @@ class GFSForecast(metsource.MetDataSource):
             cachestore = cache.DiskCacheStore()
         self.cachestore = cachestore
         self.show_progress = show_progress
+        self.cache_download = cache_download
 
         if time is None and paths is None:
             raise ValueError("Time input is required when paths is None")
@@ -349,7 +360,7 @@ class GFSForecast(metsource.MetDataSource):
         forecast_hour = str(self.forecast_time.hour).zfill(2)
         return f"gfs.t{forecast_hour}z.pgrb2.{self._grid_string}.f{step_hour}"
 
-    @overrides
+    @override
     def create_cachepath(self, t: datetime) -> str:
         if self.cachestore is None:
             raise ValueError("self.cachestore attribute must be defined to create cache path")
@@ -366,7 +377,7 @@ class GFSForecast(metsource.MetDataSource):
         # return cache path
         return self.cachestore.path(f"{datestr}-{step}-{suffix}.nc")
 
-    @overrides
+    @override
     def download_dataset(self, times: list[datetime]) -> None:
         # get step relative to forecast forecast_time
         logger.debug(
@@ -377,7 +388,7 @@ class GFSForecast(metsource.MetDataSource):
         for t in times:
             self._download_file(t)
 
-    @overrides
+    @override
     def cache_dataset(self, dataset: xr.Dataset) -> None:
         # if self.cachestore is None:
         #     LOG.debug("Cache is turned off, skipping")
@@ -385,7 +396,7 @@ class GFSForecast(metsource.MetDataSource):
 
         raise NotImplementedError("GFS caching only implemented with download")
 
-    @overrides
+    @override
     def open_metdataset(
         self,
         dataset: xr.Dataset | None = None,
@@ -429,7 +440,7 @@ class GFSForecast(metsource.MetDataSource):
         # run the same GFS-specific processing on the dataset
         return self._process_dataset(ds, **kwargs)
 
-    @overrides
+    @override
     def set_metadata(self, ds: xr.Dataset | met.MetDataset) -> None:
         ds.attrs.update(
             provider="NCEP",
@@ -462,24 +473,30 @@ class GFSForecast(metsource.MetDataSource):
         filename = self.filename(t)
         aws_key = f"{self.forecast_path}/{filename}"
 
+        stack = contextlib.ExitStack()
+        if self.cache_download:
+            target = self.cachestore.path(aws_key.replace("/", "-"))
+        else:
+            target = stack.enter_context(temp.temp_file())
+
         # Hold downloaded file in named temp file
-        with temp.temp_file() as temp_grib_filename:
+        with stack:
             # retrieve data from AWS S3
-            logger.debug(f"Downloading GFS file {filename} from AWS bucket to {temp_grib_filename}")
-            if self.show_progress:
-                _download_with_progress(
-                    self.client, GFS_FORECAST_BUCKET, aws_key, temp_grib_filename, filename
-                )
-            else:
-                self.client.download_file(
-                    Bucket=GFS_FORECAST_BUCKET, Key=aws_key, Filename=temp_grib_filename
-                )
+            logger.debug(f"Downloading GFS file {filename} from AWS bucket to {target}")
+            if not self.cache_download or not self.cachestore.exists(target):
+                self._make_download(aws_key, target, filename)
 
-            ds = self._open_gfs_dataset(temp_grib_filename, t)
+            ds = self._open_gfs_dataset(target, t)
 
             cache_path = self.create_cachepath(t)
             ds.to_netcdf(cache_path)
 
+    def _make_download(self, aws_key: str, target: str, filename: str) -> None:
+        if self.show_progress:
+            _download_with_progress(self.client, GFS_FORECAST_BUCKET, aws_key, target, filename)
+        else:
+            self.client.download_file(Bucket=GFS_FORECAST_BUCKET, Key=aws_key, Filename=target)
+
     def _open_gfs_dataset(self, filepath: str | pathlib.Path, t: datetime) -> xr.Dataset:
         """Open GFS grib file for one forecast timestep.
 
@@ -502,7 +519,7 @@ class GFSForecast(metsource.MetDataSource):
         step = pd.Timedelta(t - self.forecast_time) // pd.Timedelta(1, "h")
 
         # open file for each variable short name individually
-        ds: xr.Dataset | None = None
+        da_dict = {}
         for variable in self.variables:
             # Radiation data is not available in the 0th step
             is_radiation_step_zero = step == 0 and variable in (
@@ -519,23 +536,24 @@ class GFSForecast(metsource.MetDataSource):
             else:
                 v = variable
 
-            tmpds = xr.open_dataset(
-                filepath,
-                filter_by_keys={"typeOfLevel": v.level_type, "shortName": v.short_name},
-                engine="cfgrib",
-            )
-
-            if ds is None:
-                ds = tmpds
-            else:
-                ds[v.short_name] = tmpds[v.short_name]
+            try:
+                da = xr.open_dataarray(
+                    filepath,
+                    filter_by_keys={"typeOfLevel": v.level_type, "shortName": v.short_name},
+                    engine="cfgrib",
+                )
+            except ValueError as exc:
+                # To debug this situation, you can use:
+                # import cfgrib
+                # cfgrib.open_datasets(filepath)
+                msg = f"Variable {v.short_name} not found in {filepath}"
+                raise ValueError(msg) from exc
 
-            # set all radiation data to np.nan in the 0th step
             if is_radiation_step_zero:
-                ds = ds.rename({Visibility.short_name: variable.short_name})
-                ds[variable.short_name] = np.nan
+                da = xr.full_like(da, np.nan)  # set all radiation data to np.nan in the 0th step
+            da_dict[variable.short_name] = da
 
-        assert ds is not None, "No variables were loaded from grib file"
+        ds = xr.Dataset(da_dict)
 
         # for pressure levels, need to rename "level" field and downselect
         if self.pressure_levels != [-1]:

pycontrails/datalib/gfs/variables.py

@@ -43,7 +43,9 @@ CloudIceWaterMixingRatio = MetVariable(
 
 
 TOAUpwardShortwaveRadiation = MetVariable(
-    short_name="uswrf",
+    # Note the variable in the GFS Grib file is "uswrf" for the "nominalTop" level
+    # eccodes > 2.38 rewrites to `suswrf` on loading
+    short_name="suswrf",
     standard_name="toa_upward_shortwave_flux",
     long_name="Top of atmosphere upward shortwave radiation",
     units="W m**-2",
@@ -56,7 +58,9 @@ TOAUpwardShortwaveRadiation = MetVariable(
 )
 
 TOAUpwardLongwaveRadiation = MetVariable(
-    short_name="ulwrf",
+    # Note the variable in the GFS Grib file is "ulwrf" for the "nominalTop" level
+    # eccodes > 2.38 rewrites to `sulwrf` on loading
+    short_name="sulwrf",
     standard_name="toa_upward_longwave_flux",
     long_name="Top of atmosphere upward longwave radiation",
     units="W m**-2",

pycontrails/datalib/landsat.py

@@ -152,7 +152,7 @@ class Landsat:
         are used. Bands must share a common resolution. The resolutions of each band are:
 
         - B1-B7, B9: 30 m
-        - B9: 15 m
+        - B8: 15 m
         - B10, B11: 30 m (upsampled from true resolution of 100 m)
 
     cachestore : cache.CacheStore, optional
@@ -291,9 +291,7 @@ def _check_band_resolution(bands: set[str]) -> None:
     there are two valid cases: only band 8, or any bands except band 8.
     """
     groups = [
-        {
-            "B8",
-        },  # 15 m
+        {"B8"},  # 15 m
         {f"B{i}" for i in range(1, 12) if i != 8},  # 30 m
     ]
     if not any(bands.issubset(group) for group in groups):
@@ -313,10 +311,9 @@ def _read(path: str, meta: str, band: str, processing: str) -> xr.DataArray:
             pycontrails_optional_package="sat",
         )
 
-    src = rasterio.open(path)
-    img = src.read(1)
-    crs = pyproj.CRS.from_epsg(src.crs.to_epsg())
-    src.close()
+    with rasterio.open(path) as src:
+        img = src.read(1)
+        crs = pyproj.CRS.from_epsg(src.crs.to_epsg())
 
     if processing == "reflectance":
         mult, add = _read_band_reflectance_rescaling(meta, band)

pycontrails/datalib/sentinel.py

@@ -313,9 +313,8 @@ def _check_band_resolution(bands: set[str]) -> None:
 def _read(path: str, granule_meta: str, safe_meta: str, band: str, processing: str) -> xr.DataArray:
     """Read imagery data from Sentinel-2 files."""
     Image.MAX_IMAGE_PIXELS = None  # avoid decompression bomb warning
-    src = Image.open(path)
-    img = np.asarray(src)
-    src.close()
+    with Image.open(path) as src:
+        img = np.asarray(src)
 
     if processing == "reflectance":
         gain, offset = _read_band_reflectance_rescaling(safe_meta, band)
@@ -357,10 +356,9 @@ def _band_id(band: str) -> int:
     """Get band ID used in some metadata files."""
     if band in (f"B{i:2d}" for i in range(1, 9)):
         return int(band[1:]) - 1
-    elif band == "B8A":
+    if band == "B8A":
         return 8
-    else:
-        return int(band[1:])
+    return int(band[1:])
 
 
 def _read_band_reflectance_rescaling(meta: str, band: str) -> tuple[float, float]:
@@ -389,12 +387,10 @@ def _read_band_reflectance_rescaling(meta: str, band: str) -> tuple[float, float
     for elem in elems:
         if int(elem.attrib["band_id"]) == band_id and elem.text is not None:
             offset = float(elem.text)
-            break
-    else:
-        msg = f"Could not find reflectance offset for band {band} (band ID {band_id})"
-        raise ValueError(msg)
+            return gain, offset
 
-    return gain, offset
+    msg = f"Could not find reflectance offset for band {band} (band ID {band_id})"
+    raise ValueError(msg)
 
 
 def _read_image_coordinates(meta: str, band: str) -> tuple[np.ndarray, np.ndarray]:

pycontrails/ext/bada.py

@@ -22,8 +22,9 @@ try:
 
 except ImportError as e:
     raise ImportError(
-        "Failed to import the 'pycontrails-bada' package. Install with 'pip install"
-        ' "pycontrails-bada @ git+ssh://git@github.com/contrailcirrus/pycontrails-bada.git"\'.'
+        "Failed to import the 'pycontrails-bada' package. Install with 'pip install "
+        "--index-url https://us-central1-python.pkg.dev/contrails-301217/pycontrails/simple "
+        "pycontrails-bada'."
     ) from e
 else:
     __all__ = [

pycontrails/ext/synthetic_flight.py

@@ -20,8 +20,9 @@ try:
     from pycontrails.ext.bada import bada_model
 except ImportError as e:
     raise ImportError(
-        'SyntheticFlight requires BADA extension. Install with `pip install "pycontrails-bada @'
-        ' git+ssh://git@github.com/contrailcirrus/pycontrails-bada.git"`'
+        "Failed to import the 'pycontrails-bada' package. Install with 'pip install "
+        "--index-url https://us-central1-python.pkg.dev/contrails-301217/pycontrails/simple "
+        "pycontrails-bada'."
     ) from e
 
 logger = logging.getLogger(__name__)

pycontrails/models/accf.py

@@ -88,13 +88,16 @@ class ACCFParams(ModelParams):
     h2o_scaling: float = 1.0
     o3_scaling: float = 1.0
 
-    forecast_step: float = 6.0
+    forecast_step: float | None = None
 
     sep_ri_rw: bool = False
 
     climate_indicator: str = "ATR"
 
-    horizontal_resolution: float = 0.5
+    #: The horizontal resolution of the meteorological data in degrees.
+    #: If None, it will be inferred from the ``met`` dataset for :class:`MetDataset`
+    #: source, otherwise it will be set to 0.5.
+    horizontal_resolution: float | None = None
 
     emission_scenario: str = "pulse"
 
@@ -116,6 +119,8 @@ class ACCFParams(ModelParams):
 
     PMO: bool = False
 
+    unit_K_per_kg_fuel: bool = False
+
 
 class ACCF(Model):
     """Compute Algorithmic Climate Change Functions (ACCF).
@@ -146,16 +151,13 @@ class ACCF(Model):
         SpecificHumidity,
         ecmwf.PotentialVorticity,
         Geopotential,
-        RelativeHumidity,
+        (RelativeHumidity, ecmwf.RelativeHumidity),
         NorthwardWind,
         EastwardWind,
-        ecmwf.PotentialVorticity,
     )
     sur_variables = (ecmwf.SurfaceSolarDownwardRadiation, ecmwf.TopNetThermalRadiation)
     default_params = ACCFParams
 
-    short_vars = frozenset(v.short_name for v in (*met_variables, *sur_variables))
-
     # This variable won't get used since we are not writing the output
     # anywhere, but the library will complain if it's not defined
     path_lib = "./"
@@ -168,7 +170,13 @@ class ACCF(Model):
         **params_kwargs: Any,
     ) -> None:
         # Normalize ECMWF variables
-        met = standardize_variables(met, self.met_variables)
+        variables = (v[0] if isinstance(v, tuple) else v for v in self.met_variables)
+        met = standardize_variables(met, variables)
+
+        # If relative humidity is in percentage, convert to a proportion
+        if met["relative_humidity"].attrs.get("units") == "%":
+            met.data["relative_humidity"] /= 100.0
+            met.data["relative_humidity"].attrs["units"] = "1"
 
         # Ignore humidity scaling warning
         with warnings.catch_warnings():
@@ -231,18 +239,21 @@ class ACCF(Model):
             if hasattr(self, "surface"):
                 self.surface = self.source.downselect_met(self.surface)
 
-        if isinstance(self.source, MetDataset):
-            # Overwrite horizontal resolution to match met
-            longitude = self.source.data["longitude"].values
-            if longitude.size > 1:
-                hres = abs(longitude[1] - longitude[0])
-                self.params["horizontal_resolution"] = float(hres)
-
-            else:
+        if self.params["horizontal_resolution"] is None:
+            if isinstance(self.source, MetDataset):
+                # Overwrite horizontal resolution to match met
+                longitude = self.source.data["longitude"].values
                 latitude = self.source.data["latitude"].values
-                if latitude.size > 1:
+                if longitude.size > 1:
+                    hres = abs(longitude[1] - longitude[0])
+                    self.params["horizontal_resolution"] = float(hres)
+                elif latitude.size > 1:
                     hres = abs(latitude[1] - latitude[0])
                     self.params["horizontal_resolution"] = float(hres)
+                else:
+                    self.params["horizontal_resolution"] = 0.5
+            else:
+                self.params["horizontal_resolution"] = 0.5
 
         p_settings = _get_accf_config(self.params)
 
@@ -267,10 +278,14 @@ class ACCF(Model):
         aCCFs, _ = clim_imp.get_xarray()
 
         # assign ACCF outputs to source
+        skip = {
+            v[0].short_name if isinstance(v, tuple) else v.short_name
+            for v in (*self.met_variables, *self.sur_variables)
+        }
         maCCFs = MetDataset(aCCFs)
         for key, arr in maCCFs.data.items():
             # skip met variables
-            if key in self.short_vars:
+            if key in skip:
                 continue
 
             assert isinstance(key, str)
@@ -292,7 +307,12 @@ class ACCF(Model):
         # It also needs variables to have the ECMWF short name
         if isinstance(self.met, MetDataset):
             ds_met = self.met.data.transpose("time", "level", "latitude", "longitude")
-            name_dict = {v.standard_name: v.short_name for v in self.met_variables}
+            name_dict = {
+                v[0].standard_name if isinstance(v, tuple) else v.standard_name: v[0].short_name
+                if isinstance(v, tuple)
+                else v.short_name
+                for v in self.met_variables
+            }
             ds_met = ds_met.rename(name_dict)
         else:
             ds_met = None
@@ -340,7 +360,7 @@ def _get_accf_config(params: dict[str, Any]) -> dict[str, Any]:
         "horizontal_resolution": params["horizontal_resolution"],
         "forecast_step": params["forecast_step"],
         "NOx_aCCF": True,
-        "NOx&inverse_EIs": params["nox_ei"],
+        "NOx_EI&F_km": params["nox_ei"],
         "output_format": "netCDF",
         "mean": False,
         "std": False,
@@ -361,6 +381,7 @@ def _get_accf_config(params: dict[str, Any]) -> dict[str, Any]:
             "H2O": params["h2o_scaling"],
             "O3": params["o3_scaling"],
         },
+        "unit_K/kg(fuel)": params["unit_K_per_kg_fuel"],
         "PCFA": params["pfca"],
         "PCFA-ISSR": {
             "rhi_threshold": params["issr_rhi_threshold"],

pycontrails/models/apcemm/apcemm.py

@@ -28,6 +28,7 @@ from pycontrails.core.met_var import (
     SpecificHumidity,
     VerticalVelocity,
 )
+from pycontrails.core.vector import GeoVectorDataset
 from pycontrails.models.apcemm import utils
 from pycontrails.models.apcemm.inputs import APCEMMInput
 from pycontrails.models.dry_advection import DryAdvection
@@ -314,7 +315,7 @@ class APCEMM(models.Model):
     source: Flight
 
     #: Output from trajectory calculation
-    trajectories: Flight | None
+    trajectories: GeoVectorDataset | None
 
     #: Time series output from the APCEMM early plume model
     vortex: pd.DataFrame | None

pycontrails/models/cocip/cocip.py

@@ -3,15 +3,20 @@
 from __future__ import annotations
 
 import logging
+import sys
 import warnings
 from collections.abc import Sequence
 from typing import Any, Literal, NoReturn, overload
 
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
 import xarray as xr
-from overrides import overrides
 
 from pycontrails.core import met_var
 from pycontrails.core.aircraft_performance import AircraftPerformance
@@ -1218,7 +1223,7 @@ class Cocip(Model):
 
         return self._downwash_contrail.filter(filt)
 
-    @overrides
+    @override
     def _cleanup_indices(self) -> None:
         """Cleanup interpolation artifacts."""
 
@@ -2291,8 +2296,7 @@ def calc_timestep_contrail_evolution(
     dt = time_2_array - time_1
 
     # get new contrail location & segment properties after t_step
-    longitude_2 = geo.advect_longitude(longitude_1, latitude_1, u_wind_1, dt)
-    latitude_2 = geo.advect_latitude(latitude_1, v_wind_1, dt)
+    longitude_2, latitude_2 = geo.advect_horizontal(longitude_1, latitude_1, u_wind_1, v_wind_1, dt)
     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)
 

pycontrails/models/cocipgrid/cocip_grid.py

@@ -816,6 +816,9 @@ class CocipGrid(models.Model):
         """
         Shortcut to create a :class:`MetDataset` source from coordinate arrays.
 
+        .. versionchanged:: 0.54.3
+            By default, the returned latitude values now extend to the poles.
+
         Parameters
         ----------
         level : level: npt.NDArray[np.float64] | list[float] | float
@@ -829,8 +832,6 @@ class CocipGrid(models.Model):
         longitude, latitude : npt.NDArray[np.float64] | 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``.
-            To avoid model degradation at the poles, latitude values are expected to be
-            between -80 and 80 degrees.
         lon_step, lat_step : float, optional
             Longitude and latitude resolution, by default 1.0.
             Only used if parameter ``longitude`` (respective ``latitude``) not specified.
@@ -847,15 +848,11 @@ class CocipGrid(models.Model):
         if longitude is None:
             longitude = np.arange(-180, 180, lon_step, dtype=float)
         if latitude is None:
-            latitude = np.arange(-80, 80.000001, lat_step, dtype=float)
-
-        out = MetDataset.from_coords(longitude=longitude, latitude=latitude, level=level, time=time)
-
-        if np.any(out.data.latitude > 80.0001) or np.any(out.data.latitude < -80.0001):
-            msg = "Model only supports latitude between -80 and 80."
-            raise ValueError(msg)
+            latitude = np.arange(-90, 90.000001, lat_step, dtype=float)
 
-        return out
+        return MetDataset.from_coords(
+            longitude=longitude, latitude=latitude, level=level, time=time
+        )
 
 
 ################################
@@ -2054,10 +2051,13 @@ def advect(
     time_t2 = time + dt
     age_t2 = age + dt
 
-    longitude_t2 = geo.advect_longitude(
-        longitude=longitude, latitude=latitude, u_wind=u_wind, dt=dt
+    longitude_t2, latitude_t2 = geo.advect_horizontal(
+        longitude=longitude,
+        latitude=latitude,
+        u_wind=u_wind,
+        v_wind=v_wind,
+        dt=dt,
     )
-    latitude_t2 = geo.advect_latitude(latitude=latitude, v_wind=v_wind, dt=dt)
     level_t2 = geo.advect_level(level, vertical_velocity, rho_air, terminal_fall_speed, dt)
     altitude_t2 = units.pl_to_m(level_t2)
 
@@ -2089,15 +2089,20 @@ def advect(
     u_wind_tail = contrail["eastward_wind_tail"]
     v_wind_tail = contrail["northward_wind_tail"]
 
-    longitude_head_t2 = geo.advect_longitude(
-        longitude=longitude_head, latitude=latitude_head, u_wind=u_wind_head, dt=dt_head
+    longitude_head_t2, latitude_head_t2 = geo.advect_horizontal(
+        longitude=longitude_head,
+        latitude=latitude_head,
+        u_wind=u_wind_head,
+        v_wind=v_wind_head,
+        dt=dt_head,
     )
-    latitude_head_t2 = geo.advect_latitude(latitude=latitude_head, v_wind=v_wind_head, dt=dt_head)
-
-    longitude_tail_t2 = geo.advect_longitude(
-        longitude=longitude_tail, latitude=latitude_tail, u_wind=u_wind_tail, dt=dt_tail
+    longitude_tail_t2, latitude_tail_t2 = geo.advect_horizontal(
+        longitude=longitude_tail,
+        latitude=latitude_tail,
+        u_wind=u_wind_tail,
+        v_wind=v_wind_tail,
+        dt=dt_tail,
     )
-    latitude_tail_t2 = geo.advect_latitude(latitude=latitude_tail, v_wind=v_wind_tail, dt=dt_tail)
 
     segment_length_t2 = geo.haversine(
         lons0=longitude_head_t2,