pycontrails 0.47.3__cp312-cp312-macosx_11_0_arm64.whl → 0.48.1__cp312-cp312-macosx_11_0_arm64.whl

pycontrails/datalib/ecmwf/hres.py

@@ -19,16 +19,9 @@ from overrides import overrides
 import pycontrails
 from pycontrails.core import cache, datalib
 from pycontrails.core.met import MetDataset, MetVariable
-from pycontrails.datalib.ecmwf.common import ECMWFAPI, rad_accumulated_to_average
-from pycontrails.datalib.ecmwf.variables import (
-    PRESSURE_LEVEL_VARIABLES,
-    SURFACE_VARIABLES,
-    TOAIncidentSolarRadiation,
-    TopNetSolarRadiation,
-    TopNetThermalRadiation,
-)
-from pycontrails.utils import dependencies, iteration
-from pycontrails.utils.temp import temp_file
+from pycontrails.datalib.ecmwf.common import ECMWFAPI
+from pycontrails.datalib.ecmwf.variables import PRESSURE_LEVEL_VARIABLES, SURFACE_VARIABLES
+from pycontrails.utils import dependencies, iteration, temp
 from pycontrails.utils.types import DatetimeLike
 
 if TYPE_CHECKING:
@@ -241,15 +234,6 @@ class HRES(ECMWFAPI):
     #: Forecast run time, either specified or assigned by the closest previous forecast run
     forecast_time: datetime
 
-    #: User provided ``ECMWFService`` url
-    url: str | None
-
-    #: User provided ``ECMWFService`` key
-    key: str | None
-
-    #: User provided ``ECMWFService`` email
-    email: str | None
-
     __marker = object()
 
     def __init__(
@@ -278,14 +262,11 @@ class HRES(ECMWFAPI):
                 pycontrails_optional_package="ecmwf",
             )
 
-        # constants
         # ERA5 now delays creating the server attribute until it is needed to download
         # from CDS. We could do the same here.
-        self.url = url
-        self.key = key
-        self.email = email
-        self.server = ECMWFService("mars", url=self.url, key=self.key, email=self.email)
+        self.server = ECMWFService("mars", url=url, key=key, email=email)
         self.paths = paths
+
         if cachestore is self.__marker:
             cachestore = cache.DiskCacheStore()
         self.cachestore = cachestore
@@ -300,10 +281,24 @@ class HRES(ECMWFAPI):
         self.variables = datalib.parse_variables(variables, self.supported_variables)
 
         self.grid = datalib.parse_grid(grid, [0.1, 0.25, 0.5, 1])  # lat/lon degree resolution
-        self.stream = stream  # "enfo" = ensemble forecast, "oper" = atmospheric model/HRES
-        self.field_type = (
-            field_type  # forecast (oper), perturbed or control forecast (enfo only), or analysis
-        )
+
+        # "enfo" = ensemble forecast
+        # "oper" = atmospheric model/HRES
+        if stream not in ("oper", "enfo"):
+            msg = "Parameter stream must be 'oper' or 'enfo'"
+            raise ValueError(msg)
+
+        self.stream = stream
+
+        # "fc" = forecast
+        # "pf" = perturbed forecast
+        # "cf" = control forecast
+        # "an" = analysis
+        if field_type not in ("fc", "pf", "cf", "an"):
+            msg = "Parameter field_type must be 'fc', 'pf', 'cf', or 'an'"
+            raise ValueError(msg)
+
+        self.field_type = field_type
 
         # set specific forecast time is requested
         if forecast_time is not None:
@@ -476,14 +471,12 @@ class HRES(ECMWFAPI):
         request = self.generate_mars_request(self.forecast_time, self.steps, request_type="list")
 
         # hold downloaded file in named temp file
-        with temp_file() as mars_temp_filename:
-            LOG.debug(f"Performing MARS request: {request}")
-            self.server.execute(request, mars_temp_filename)
+        with temp.temp_file() as mars_temp_filename:
+            LOG.debug("Performing MARS request: %s", request)
+            self.server.execute(request, target=mars_temp_filename)
 
             with open(mars_temp_filename, "r") as f:
-                txt = f.read()
-
-        return txt
+                return f.read()
 
     def generate_mars_request(
         self,
@@ -528,16 +521,16 @@ class HRES(ECMWFAPI):
             steps = self.steps
 
         # set date/time for file
-        _date = forecast_time.strftime("%Y%m%d")
-        _time = forecast_time.strftime("%H")
+        date = forecast_time.strftime("%Y%m%d")
+        time = forecast_time.strftime("%H")
 
         # make request of mars
         request: dict[str, Any] = {
             "class": "od",  # operational data
             "stream": self.stream,
             "expver": "1",  # production data only
-            "date": _date,
-            "time": _time,
+            "date": date,
+            "time": time,
             "type": self.field_type,
             "param": f"{'/'.join(self.variable_shortnames)}",
             "step": f"{'/'.join([str(s) for s in steps])}",
@@ -574,10 +567,8 @@ class HRES(ECMWFAPI):
         step = self.step_offset + self.timesteps.index(t)
 
         # single level or pressure level
-        if self.pressure_levels == [-1]:
-            suffix = f"hressl{self.grid}{self.stream}{self.field_type}"
-        else:
-            suffix = f"hrespl{self.grid}{self.stream}{self.field_type}"
+        levtype = "sl" if self.pressure_levels == [-1] else "pl"
+        suffix = f"hres{levtype}{self.grid}{self.stream}{self.field_type}"
 
         # return cache path
         return self.cachestore.path(f"{datestr}-{step}-{suffix}.nc")
@@ -635,47 +626,31 @@ class HRES(ECMWFAPI):
             # open cache files as xr.Dataset
             ds = self.open_dataset(disk_cachepaths, **xr_kwargs)
 
-            # TODO: corner case
-            # If any files are already cached, they will not have the version attached
-            if "pycontrails_version" not in ds.attrs:
-                ds.attrs["pycontrails_version"] = pycontrails.__version__
+            ds.attrs.setdefault("pycontrails_version", pycontrails.__version__)
 
         # run the same ECMWF-specific processing on the dataset
         mds = self._process_dataset(ds, **kwargs)
 
-        # convert accumulated radiation values to average instantaneous values
-        # set minimum for all values to 0
-        # !! Note that HRES accumulates from the *start of the forecast*,
-        # so we need to take the diff of each accumulated value
-        # the 0th value is set to the 1st value so each time step has a radiation value !!
-        dt_accumulation = 60 * 60
-
-        for key in [
-            TOAIncidentSolarRadiation.standard_name,
-            TopNetSolarRadiation.standard_name,
-            TopNetThermalRadiation.standard_name,
-        ]:
-            if key in mds.data:
-                if len(mds.data["time"]) < 2:
-                    raise RuntimeError(
-                        f"HRES datasets with data variable {key} must have at least two timesteps"
-                        f" to calculate the average instantaneous value of {key}"
-                    )
-
-                # take the difference between time slices
-                dkey_dt = mds.data[key].diff("time")
-
-                # set difference value back to the data model
-                mds.data[key] = dkey_dt
-
-                # set the 0th value of the data to the 1st difference value
-                # TODO: this assumption may not be universally applicable!
-                mds.data[key][dict(time=0)] = dkey_dt[dict(time=0)]
-
-                rad_accumulated_to_average(mds, key, dt_accumulation)
-
+        self.set_metadata(mds)
         return mds
 
+    @overrides
+    def set_metadata(self, ds: xr.Dataset | MetDataset) -> None:
+        if self.stream == "oper":
+            product = "forecast"
+        elif self.stream == "enfo":
+            product = "ensemble"
+        else:
+            msg = f"Unknown stream type {self.stream}"
+            raise ValueError(msg)
+
+        ds.attrs.update(
+            provider="ECMWF",
+            dataset="HRES",
+            product=product,
+            radiation_accumulated=True,
+        )
+
     def _open_and_cache(self, xr_kwargs: dict[str, Any]) -> xr.Dataset:
         """Open and cache :class:`xr.Dataset` from :attr:`self.paths`.
 
@@ -746,7 +721,7 @@ class HRES(ECMWFAPI):
         # Open ExitStack to control temp_file context manager
         with ExitStack() as stack:
             # hold downloaded file in named temp file
-            mars_temp_grib_filename = stack.enter_context(temp_file())
+            mars_temp_grib_filename = stack.enter_context(temp.temp_file())
 
             # retrieve data from MARS
             LOG.debug(f"Performing MARS request: {request}")
@@ -754,18 +729,8 @@ class HRES(ECMWFAPI):
 
             # translate into netcdf from grib
             LOG.debug("Translating file into netcdf")
-            mars_temp_nc_filename = stack.enter_context(temp_file())
             ds = stack.enter_context(xr.open_dataset(mars_temp_grib_filename, engine="cfgrib"))
 
-            ##### TODO: do we need to store intermediate netcdf file?
-            ds.to_netcdf(path=mars_temp_nc_filename, mode="w")
-
-            # open file, edit, and save for each hourly time step
-            ds = stack.enter_context(
-                xr.open_dataset(mars_temp_nc_filename, engine=datalib.NETCDF_ENGINE)
-            )
-            #####
-
             # run preprocessing before cache
             ds = self._preprocess_hres_dataset(ds)
 

pycontrails/datalib/ecmwf/ifs.py

@@ -27,7 +27,7 @@ class IFS(datalib.MetDataSource):
 
     .. warning::
 
-        This data source is fully implemented.
+        This data source is not fully implemented.
 
     Parameters
     ----------
@@ -186,9 +186,17 @@ class IFS(datalib.MetDataSource):
         # harmonize variable names
         ds = met.standardize_variables(ds, self.variables)
 
-        ds.attrs["met_source"] = type(self).__name__
+        self.set_metadata(ds)
         return met.MetDataset(ds, **kwargs)
 
+    @overrides
+    def set_metadata(self, ds: xr.Dataset | met.MetDataset) -> None:
+        ds.attrs.update(
+            provider="ECMWF",
+            dataset="IFS",
+            product="forecast",
+        )
+
     @overrides
     def download_dataset(self, times: list[datetime]) -> None:
         raise NotImplementedError("IFS download is not supported")

pycontrails/datalib/gfs/gfs.py

@@ -12,7 +12,7 @@ from __future__ import annotations
 import hashlib
 import logging
 import pathlib
-from contextlib import ExitStack
+import warnings
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Callable
 
@@ -30,15 +30,14 @@ from pycontrails.datalib.gfs.variables import (
     TOAUpwardShortwaveRadiation,
     Visibility,
 )
-from pycontrails.utils import dependencies
-from pycontrails.utils.temp import temp_file
+from pycontrails.utils import dependencies, temp
 from pycontrails.utils.types import DatetimeLike
 
 # optional imports
 if TYPE_CHECKING:
     import botocore
 
-LOG = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)
 
 #: Default GFS AWS bucket
 GFS_FORECAST_BUCKET = "noaa-gfs-bdp-pds"
@@ -72,8 +71,8 @@ class GFSForecast(datalib.MetDataSource):
         Specify latitude/longitude grid spacing in data.
         Defaults to 0.25.
     forecast_time : `DatetimeLike`, optional
-        Specify forecast run by runtime.
-        Defaults to None.
+        Specify forecast run by runtime. If None (default), the forecast time
+        is set to the 6 hour floor of the first timestep.
     cachestore : :class:`cache.CacheStore` | None, optional
         Cache data store for staging data files.
         Defaults to :class:`cache.DiskCacheStore`.
@@ -88,7 +87,7 @@ class GFSForecast(datalib.MetDataSource):
     >>> from pycontrails.datalib.gfs import GFSForecast
 
     >>> # Store data files to local disk (default behavior)
-    >>> times = ("2022-03-22 00:00:00", "2022-03-22 03:00:00", )
+    >>> times = ("2022-03-22 00:00:00", "2022-03-22 03:00:00")
     >>> gfs = GFSForecast(times, variables="air_temperature", pressure_levels=[300, 250])
     >>> gfs
     GFSForecast
@@ -97,7 +96,15 @@ class GFSForecast(datalib.MetDataSource):
         Pressure levels: [300, 250]
         Grid: 0.25
         Forecast time: 2022-03-22 00:00:00
-        Steps: [0, 1, 2, 3]
+
+    >>> gfs = GFSForecast(times, variables="air_temperature", pressure_levels=[300, 250], grid=0.5)
+    >>> gfs
+    GFSForecast
+        Timesteps: ['2022-03-22 00', '2022-03-22 03']
+        Variables: ['t']
+        Pressure levels: [300, 250]
+        Grid: 0.5
+        Forecast time: 2022-03-22 00:00:00
 
     Notes
     -----
@@ -164,12 +171,17 @@ class GFSForecast(datalib.MetDataSource):
         if time is None and paths is None:
             raise ValueError("Time input is required when paths is None")
 
-        self.timesteps = datalib.parse_timesteps(time, freq="1H")
+        # Forecast is available hourly for 0.25 degree grid,
+        # 3 hourly for 0.5 and 1 degree grid
+        # https://www.nco.ncep.noaa.gov/pmb/products/gfs/
+        freq = "1H" if grid == 0.25 else "3H"
+        self.timesteps = datalib.parse_timesteps(time, freq=freq)
+
         self.pressure_levels = datalib.parse_pressure_levels(
             pressure_levels, self.supported_pressure_levels
         )
         self.variables = datalib.parse_variables(variables, self.supported_variables)
-        self.grid = datalib.parse_grid(grid, [0.25, 0.5, 1])
+        self.grid = datalib.parse_grid(grid, (0.25, 0.5, 1))
 
         # note GFS allows unsigned requests (no credentials)
         # https://stackoverflow.com/questions/34865927/can-i-use-boto3-anonymously/34866092#34866092
@@ -180,24 +192,19 @@ class GFSForecast(datalib.MetDataSource):
         # set specific forecast time is requested
         if forecast_time is not None:
             forecast_time_pd = pd.to_datetime(forecast_time)
-            if forecast_time_pd.hour not in [0, 6, 12, 18]:
+            if forecast_time_pd.hour % 6:
                 raise ValueError("Forecast hour must be on one of 00, 06, 12, 18")
 
             self.forecast_time = datalib.round_hour(forecast_time_pd.to_pydatetime(), 6)
 
         # if no specific forecast is requested, set the forecast time using timesteps
-        elif self.timesteps:
+        else:
             # round first element to the nearest 6 hour time (00, 06, 12, 18 UTC) for forecast_time
             self.forecast_time = datalib.round_hour(self.timesteps[0], 6)
 
-        # when no forecast_time or time input, forecast_time is defined in _open_and_cache
-
     def __repr__(self) -> str:
         base = super().__repr__()
-        return (
-            f"{base}\n\tForecast time: {getattr(self, 'forecast_time', '')}\n\tSteps:"
-            f" {getattr(self, 'steps', '')}"
-        )
+        return f"{base}\n\tForecast time: {self.forecast_time}"
 
     @property
     def supported_pressure_levels(self) -> list[int]:
@@ -282,32 +289,6 @@ class GFSForecast(datalib.MetDataSource):
         )
         return hashlib.sha1(bytes(hashstr, "utf-8")).hexdigest()
 
-    @property
-    def step_offset(self) -> int:
-        """Difference between :attr:`forecast_time` and first timestep.
-
-        Returns
-        -------
-        int
-            Number of steps to offset in order to retrieve data starting from input time.
-            Returns 0 if :attr:`timesteps` is empty when loading from :attr:`paths`.
-        """
-        if self.timesteps:
-            return int((self.timesteps[0] - self.forecast_time).total_seconds() // 3600)
-
-        return 0
-
-    @property
-    def steps(self) -> list[int]:
-        """Forecast steps from :attr:`forecast_time` corresponding within input :attr:`time`.
-
-        Returns
-        -------
-        list[int]
-            List of forecast steps relative to :attr:`forecast_time`
-        """
-        return [self.step_offset + i for i in range(len(self.timesteps))]
-
     @property
     def _grid_string(self) -> str:
         """Return filename string for grid spacing."""
@@ -315,7 +296,7 @@ class GFSForecast(datalib.MetDataSource):
             return "0p25"
         if self.grid == 0.5:
             return "0p50"
-        if self.grid == 1:
+        if self.grid == 1.0:
             return "1p00"
         raise ValueError(f"Unsupported grid spacing {self.grid}. Must be one of 0.25, 0.5, or 1.0.")
 
@@ -336,7 +317,7 @@ class GFSForecast(datalib.MetDataSource):
         forecast_hour = str(self.forecast_time.hour).zfill(2)
         return f"gfs.{datestr}/{forecast_hour}/atmos"
 
-    def filename(self, step: int) -> str:
+    def filename(self, t: datetime) -> str:
         """Construct grib filename to retrieve from GFS bucket.
 
         String template:
@@ -349,8 +330,8 @@ class GFSForecast(datalib.MetDataSource):
 
         Parameters
         ----------
-        step : int
-            Integer step relative to forecast time
+        t : datetime
+            Timestep to download
 
         Returns
         -------
@@ -361,8 +342,10 @@ class GFSForecast(datalib.MetDataSource):
         ----------
         - https://www.nco.ncep.noaa.gov/pmb/products/gfs/
         """
+        step = pd.Timedelta(t - self.forecast_time) // pd.Timedelta(1, "h")
+        step_hour = str(step).zfill(3)
         forecast_hour = str(self.forecast_time.hour).zfill(2)
-        return f"gfs.t{forecast_hour}z.pgrb2.{self._grid_string}.f{str(step).zfill(3)}"
+        return f"gfs.t{forecast_hour}z.pgrb2.{self._grid_string}.f{step_hour}"
 
     @overrides
     def create_cachepath(self, t: datetime) -> str:
@@ -373,7 +356,7 @@ class GFSForecast(datalib.MetDataSource):
         datestr = self.forecast_time.strftime("%Y%m%d-%H")
 
         # get step relative to forecast forecast_time
-        step = self.step_offset + self.timesteps.index(t)
+        step = pd.Timedelta(t - self.forecast_time) // pd.Timedelta(1, "h")
 
         # single level or pressure level
         suffix = f"gfs{'sl' if self.pressure_levels == [-1] else 'pl'}{self.grid}"
@@ -384,7 +367,7 @@ class GFSForecast(datalib.MetDataSource):
     @overrides
     def download_dataset(self, times: list[datetime]) -> None:
         # get step relative to forecast forecast_time
-        LOG.debug(
+        logger.debug(
             f"Downloading GFS forecast for forecast time {self.forecast_time} and timesteps {times}"
         )
 
@@ -442,28 +425,15 @@ class GFSForecast(datalib.MetDataSource):
         ds.attrs.setdefault("pycontrails_version", pycontrails.__version__)
 
         # run the same GFS-specific processing on the dataset
-        mds = self._process_dataset(ds, **kwargs)
-
-        # set TOAUpwardShortwaveRadiation, TOAUpwardLongwaveRadiation step 0 == step 1
-        for key in [
-            TOAUpwardShortwaveRadiation.standard_name,
-            TOAUpwardLongwaveRadiation.standard_name,
-        ]:
-            # if step 0 (forecast time) exists in dimension
-            forecast_time = mds.data["forecast_time"].values
-            if key in mds.data and forecast_time in mds.data["time"]:
-                # make sure this isn't the only time in the dataset
-                if np.all(mds.data["time"].values == forecast_time):
-                    raise RuntimeError(
-                        f"GFS datasets with data variable {key} must have at least one timestep"
-                        f" after the forecast time to estimate the value of {key} at step 0"
-                    )
-
-                # set the 0th value of the data to the 1st value
-                # TODO: this assumption may not be universally applicable!
-                mds.data[key][dict(time=0)] = mds.data[key][dict(time=1)]
-
-        return mds
+        return self._process_dataset(ds, **kwargs)
+
+    @overrides
+    def set_metadata(self, ds: xr.Dataset | met.MetDataset) -> None:
+        ds.attrs.update(
+            provider="NCEP",
+            dataset="GFS",
+            product="forecast",
+        )
 
     def _download_file(self, t: datetime) -> None:
         """Download data file for forecast time and step.
@@ -487,17 +457,13 @@ class GFSForecast(datalib.MetDataSource):
             raise ValueError("Cachestore is required to download data")
 
         # construct filenames for each file
-        step = self.step_offset + self.timesteps.index(t)
-        filename = self.filename(step)
+        filename = self.filename(t)
         aws_key = f"{self.forecast_path}/{filename}"
 
-        # Open ExitStack to control temp_file context manager
-        with ExitStack() as stack:
-            # hold downloaded file in named temp file
-            temp_grib_filename = stack.enter_context(temp_file())
-
+        # Hold downloaded file in named temp file
+        with temp.temp_file() as temp_grib_filename:
             # retrieve data from AWS S3
-            LOG.debug(f"Downloading GFS file {filename} from AWS bucket to {temp_grib_filename}")
+            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
@@ -509,12 +475,8 @@ class GFSForecast(datalib.MetDataSource):
 
             ds = self._open_gfs_dataset(temp_grib_filename, t)
 
-            # write out data to temp, close grib file
-            temp_nc_filename = stack.enter_context(temp_file())
-            ds.to_netcdf(path=temp_nc_filename, mode="w")
-
-            # put each hourly file into cache
-            self.cachestore.put(temp_nc_filename, self.create_cachepath(t))
+            cache_path = self.create_cachepath(t)
+            ds.to_netcdf(cache_path)
 
     def _open_gfs_dataset(self, filepath: str | pathlib.Path, t: datetime) -> xr.Dataset:
         """Open GFS grib file for one forecast timestep.
@@ -532,22 +494,24 @@ class GFSForecast(datalib.MetDataSource):
             GFS dataset
         """
         # translate into netcdf from grib
-        LOG.debug(f"Translating {filepath} for timestep {str(t)} into netcdf")
+        logger.debug(f"Translating {filepath} for timestep {str(t)} into netcdf")
 
         # get step for timestep
-        step = self.step_offset + self.timesteps.index(t)
+        step = pd.Timedelta(t - self.forecast_time) // pd.Timedelta(1, "h")
 
         # open file for each variable short name individually
-        ds = xr.Dataset()
+        ds: xr.Dataset | None = None
         for variable in self.variables:
-            # radiation data is not available in the 0th step
-            if step == 0 and variable in [
+            # Radiation data is not available in the 0th step
+            is_radiation_step_zero = step == 0 and variable in (
                 TOAUpwardShortwaveRadiation,
                 TOAUpwardLongwaveRadiation,
-            ]:
-                LOG.debug(
-                    "Radiation data is not provided for the 0th step in GFS. Setting to np.nan"
-                    " using Visibility variable"
+            )
+
+            if is_radiation_step_zero:
+                warnings.warn(
+                    "Radiation data is not provided for the 0th step in GFS. "
+                    "Setting to np.nan using Visibility variable"
                 )
                 v = Visibility
             else:
@@ -559,23 +523,22 @@ class GFSForecast(datalib.MetDataSource):
                 engine="cfgrib",
             )
 
-            if not len(ds):
+            if ds is None:
                 ds = tmpds
             else:
                 ds[v.short_name] = tmpds[v.short_name]
 
             # set all radiation data to np.nan in the 0th step
-            if step == 0 and variable in [
-                TOAUpwardShortwaveRadiation,
-                TOAUpwardLongwaveRadiation,
-            ]:
+            if is_radiation_step_zero:
                 ds = ds.rename({Visibility.short_name: variable.short_name})
                 ds[variable.short_name] = np.nan
 
+        assert ds is not None, "No variables were loaded from grib file"
+
         # for pressure levels, need to rename "level" field and downselect
         if self.pressure_levels != [-1]:
             ds = ds.rename({"isobaricInhPa": "level"})
-            ds = ds.sel(dict(level=self.pressure_levels))
+            ds = ds.sel(level=self.pressure_levels)
 
         # for single level, and singular pressure levels, add the level dimension
         if len(self.pressure_levels) == 1:
@@ -623,15 +586,14 @@ class GFSForecast(datalib.MetDataSource):
             ds = ds.expand_dims({"level": self.pressure_levels})
 
         else:
-            ds = ds.sel(dict(level=self.pressure_levels))
+            ds = ds.sel(level=self.pressure_levels)
 
         # harmonize variable names
         ds = met.standardize_variables(ds, self.variables)
 
-        if "cachestore" not in kwargs:
-            kwargs["cachestore"] = self.cachestore
+        kwargs.setdefault("cachestore", self.cachestore)
 
-        ds.attrs["met_source"] = type(self).__name__
+        self.set_metadata(ds)
         return met.MetDataset(ds, **kwargs)