pycontrails 0.47.3__cp310-cp310-win_amd64.whl → 0.48.1__cp310-cp310-win_amd64.whl

pycontrails/core/met.py

@@ -96,13 +96,14 @@ class MetBase(ABC, Generic[XArrayType]):
         for dim in self.dim_order:
             if dim not in self.data.dims:
                 if dim == "level":
-                    raise ValueError(
+                    msg = (
                         f"Meteorology data must contain dimension '{dim}'. "
                         "For single level data, set 'level' coordinate to constant -1 "
                         "using `ds = ds.expand_dims({'level': [-1]})`"
                     )
                 else:
-                    raise ValueError(f"Meteorology data must contain dimension '{dim}'.")
+                    msg = f"Meteorology data must contain dimension '{dim}'."
+                raise ValueError(msg)
 
     def _validate_longitude(self) -> None:
         """Check longitude bounds.
@@ -196,15 +197,15 @@ class MetBase(ABC, Generic[XArrayType]):
 
         dims_tuple = tuple(self.dim_order)
 
-        def _check_da(da: xr.DataArray, key: str | None = None) -> None:
+        def _check_da(da: xr.DataArray, key: Hashable | None = None) -> None:
             if da.dims != dims_tuple:
                 if key is not None:
                     msg = (
-                        "Data dimension not transposed on variable '{key}'. Initiate with"
-                        " `copy=True`."
+                        f"Data dimension not transposed on variable '{key}'. Initiate with"
+                        " 'copy=True'."
                     )
                 else:
-                    msg = "Data dimension not transposed. Initiate with `copy=True`."
+                    msg = "Data dimension not transposed. Initiate with 'copy=True'."
                 raise ValueError(msg)
 
         data = self.data
@@ -212,7 +213,7 @@ class MetBase(ABC, Generic[XArrayType]):
             _check_da(data)
             return
 
-        for key, da in self.data.data_vars.items():
+        for key, da in self.data.items():
             _check_da(da, key)
 
     def _validate_dims(self) -> None:
@@ -533,7 +534,7 @@ class MetBase(ABC, Generic[XArrayType]):
 
     @property
     def attrs(self) -> dict[Hashable, Any]:
-        """Pass through to `self.data.attrs`."""
+        """Pass through to :attr:`self.data.attrs`."""
         return self.data.attrs
 
     @abstractmethod
@@ -593,6 +594,16 @@ class MetDataset(MetBase):
         interval ``[-180, 180]``. Defaults to False.
     copy : bool, optional
         Copy data on construction. Defaults to True.
+    attrs : dict[str, Any], optional
+        Attributes to add to :attr:`data.attrs`. Defaults to None.
+        Generally, pycontrails :class:`Models` may use the following attributes:
+
+        - ``provider``: Name of the data provider (e.g. "ECMWF").
+        - ``dataset``: Name of the dataset (e.g. "ERA5").
+        - ``product``: Name of the product type (e.g. "reanalysis").
+
+    **attrs_kwargs : Any
+        Keyword arguments to add to :attr:`data.attrs`. Defaults to None.
 
     Examples
     --------
@@ -641,10 +652,13 @@ class MetDataset(MetBase):
         cachestore: CacheStore | None = None,
         wrap_longitude: bool = False,
         copy: bool = True,
-    ):
-        # init cache
+        attrs: dict[str, Any] | None = None,
+        **attrs_kwargs: Any,
+    ) -> None:
         self.cachestore = cachestore
 
+        data.attrs.update(attrs or {}, **attrs_kwargs)
+
         # if input is already a Dataset, copy into data
         if not isinstance(data, xr.Dataset):
             raise ValueError("Input `data` must be an xarray Dataset")
@@ -661,7 +675,7 @@ class MetDataset(MetBase):
             self._validate_dims()
 
     def __getitem__(self, key: Hashable) -> MetDataArray:
-        """Return DataArray of variable `key` cast to a `MetDataArray` object.
+        """Return DataArray of variable ``key`` cast to a :class:`MetDataArray` object.
 
         Parameters
         ----------
@@ -676,7 +690,7 @@ class MetDataset(MetBase):
         Raises
         ------
         KeyError
-            If `key` not found in :attr:`data`
+            If ``key`` not found in :attr:`data`
         """
         try:
             da = self.data[key]
@@ -789,7 +803,7 @@ class MetDataset(MetBase):
             yield str(key)
 
     def __contains__(self, key: Hashable) -> bool:
-        """Check if key `key` is in :attr:`data`.
+        """Check if key ``key`` is in :attr:`data`.
 
         Parameters
         ----------
@@ -799,7 +813,7 @@ class MetDataset(MetBase):
         Returns
         -------
         bool
-            True if `key` is in :attr:`data`, False otherwise
+            True if ``key`` is in :attr:`data`, False otherwise
         """
         return key in self.data
 
@@ -857,7 +871,7 @@ class MetDataset(MetBase):
             Raises when dataset does not contain variable in ``vars``
         """
         if isinstance(vars, (MetVariable, str)):
-            vars = [vars]
+            vars = (vars,)
 
         met_keys: list[str] = []
         for variable in vars:
@@ -1001,7 +1015,7 @@ class MetDataset(MetBase):
             time                [2022-03-01 00:00:00, 2022-03-01 01:00:00]
             longitude           [-180.0, 179.75]
             latitude            [-90.0, 90.0]
-            altitude            [10362.848672411146, 11783.938524404566]
+            altitude            [10362.8, 11783.9]
             crs                 EPSG:4326
 
         """
@@ -1022,6 +1036,98 @@ class MetDataset(MetBase):
             vector.attrs.update({str(k): v for k, v in self.attrs.items()})
         return vector
 
+    def _get_pycontrails_attr_template(
+        self,
+        name: str,
+        supported: tuple[str, ...],
+        examples: dict[str, str],
+    ) -> str:
+        """Look up an attribute with a custom error message."""
+        try:
+            out = self.attrs[name]
+        except KeyError as e:
+            msg = f"Specify '{name}' attribute on underlying dataset."
+
+            for i, (k, v) in enumerate(examples.items()):
+                if i == 0:
+                    msg = f"{msg} For example, set attrs['{name}'] = '{k}' for {v}."
+                else:
+                    msg = f"{msg} Set attrs['{name}'] = '{k}' for {v}."
+            raise KeyError(msg) from e
+
+        if out not in supported:
+            warnings.warn(
+                f"Unknown {name} '{out}'. Data may not be processed correctly. "
+                f"Known {name}s are {supported}. Contact the pycontrails "
+                "developers if you believe this is an error."
+            )
+
+        return out
+
+    @property
+    def provider_attr(self) -> str:
+        """Look up the 'provider' attribute with a custom error message.
+
+        Returns
+        -------
+        str
+            Provider of the data. If not one of 'ECMWF' or 'NCEP',
+            a warning is issued.
+        """
+        supported = ("ECMWF", "NCEP")
+        examples = {"ECMWF": "data provided by ECMWF", "NCEP": "GFS data"}
+        return self._get_pycontrails_attr_template("provider", supported, examples)
+
+    @property
+    def dataset_attr(self) -> str:
+        """Look up the 'dataset' attribute with a custom error message.
+
+        Returns
+        -------
+        str
+            Dataset of the data. If not one of 'ERA5', 'HRES', 'IFS',
+            or 'GFS', a warning is issued.
+        """
+        supported = ("ERA5", "HRES", "IFS", "GFS")
+        examples = {
+            "ERA5": "ECMWF ERA5 reanalysis data",
+            "HRES": "ECMWF HRES forecast data",
+            "GFS": "NCEP GFS forecast data",
+        }
+        return self._get_pycontrails_attr_template("dataset", supported, examples)
+
+    @property
+    def product_attr(self) -> str:
+        """Look up the 'product' attribute with a custom error message.
+
+        Returns
+        -------
+        str
+            Product of the data. If not one of 'forecast', 'ensemble', or 'reanalysis',
+            a warning is issued.
+
+        """
+        supported = ("reanalysis", "forecast", "ensemble")
+        examples = {
+            "reanalysis": "ECMWF ERA5 reanalysis data",
+            "ensemble": "ECMWF ERA5 ensemble member data",
+        }
+        return self._get_pycontrails_attr_template("product", supported, examples)
+
+    def standardize_variables(self, variables: Iterable[MetVariable]) -> None:
+        """Standardize variables **in-place**.
+
+        Parameters
+        ----------
+        variables : Iterable[MetVariable]
+            Data source variables
+
+        See Also
+        --------
+        :func:`standardize_variables`
+        """
+        standardize_variables(self, variables)
+
     @classmethod
     def from_coords(
         cls,
@@ -1177,9 +1283,10 @@ class MetDataArray(MetBase):
         in the case that `copy=True`. Validation only introduces a very small overhead.
         This parameter should only be set to `False` if working with data derived from an
         existing MetDataset or :class`MetDataArray`. By default `True`.
+    name : Hashable, optional
+        Name of the data variable. If not specified, the name will be set to "met".
     **kwargs
-        If `data` input is not a xr.DataArray, `data` will be passed
-        passed directly to xr.DataArray constructor with these keyword arguments.
+        To be removed in future versions. Passed directly to xr.DataArray constructor.
 
     Examples
     --------
@@ -1218,6 +1325,7 @@ class MetDataArray(MetBase):
         wrap_longitude: bool = False,
         copy: bool = True,
         validate: bool = True,
+        name: Hashable | None = None,
         **kwargs: Any,
     ) -> None:
         # init cache
@@ -1225,6 +1333,10 @@ class MetDataArray(MetBase):
 
         # try to create DataArray out of input data and **kwargs
         if not isinstance(data, xr.DataArray):
+            DeprecationWarning(
+                "Input 'data' must be an xarray DataArray. "
+                "Passing arbitrary kwargs will be removed in future versions."
+            )
             data = xr.DataArray(data, **kwargs)
 
         if copy:
@@ -1237,13 +1349,9 @@ class MetDataArray(MetBase):
             if validate:
                 self._validate_dims()
 
-        # if name is specified in kwargs, overwrite any other name in DataArray
-        if "name" in kwargs:
-            self.data.name = kwargs["name"]
-
-        # if at this point now "name" exists on DataArray, set default
-        if self.data.name is None:
-            self.data.name = "met"
+        # Priority: name > data.name > "met"
+        name = name or self.data.name or "met"
+        self.data.name = name
 
     @property
     def values(self) -> np.ndarray:
@@ -1263,7 +1371,7 @@ class MetDataArray(MetBase):
         """
         if not self.in_memory:
             self._check_memory("Extracting numpy array from")
-            self.data = self.data.load()
+            self.data.load()
 
         return self.data.values
 
@@ -1509,8 +1617,20 @@ class MetDataArray(MetBase):
         )
 
     def _check_memory(self, msg_start: str) -> None:
+        """Check the memory usage of the underlying data.
+
+        If the data is larger than 4 GB, a warning is issued. If the data is
+        larger than 32 GB, a RuntimeError is raised.
+        """
         n_bytes = self.data.nbytes
+        mb = round(n_bytes / int(1e6), 2)
+        logger.debug("Loading %s into memory consumes %s MB.", self.name, mb)
+
         n_gb = n_bytes // int(1e9)
+        if n_gb <= 4:
+            return
+
+        # Prevent something stupid
         msg = (
             f"{msg_start} MetDataArray {self.name} requires loading "
             f"at least {n_gb} GB of data into memory. Downselect data if possible. "
@@ -1518,13 +1638,9 @@ class MetDataArray(MetBase):
             "with the method 'downselect_met'."
         )
 
-        if n_gb > 32:  # Prevent something stupid
+        if n_gb > 32:
             raise RuntimeError(msg)
-        if n_gb > 4:
-            warnings.warn(msg)
-
-        mb = round(n_bytes / int(1e6), 2)
-        logger.debug("Loading %s into memory consumes %s MB.", self.name, mb)
+        warnings.warn(msg)
 
     def save(self, **kwargs: Any) -> list[str]:
         """Save intermediate to :attr:`cachestore` as netcdf.
@@ -2119,24 +2235,29 @@ def _is_zarr(ds: xr.Dataset | xr.DataArray) -> bool:
     return dask0.array.array.array.__class__.__name__ == "ZarrArrayWrapper"
 
 
-def shift_longitude(data: XArrayType) -> XArrayType:
-    """Shift longitude values from [0, 360) to [-180, 180) domain.
+def shift_longitude(data: XArrayType, bound: float = -180.0) -> XArrayType:
+    """Shift longitude values from any input domain to [bound, 360 + bound) domain.
 
     Sorts data by ascending longitude values.
 
+
     Parameters
     ----------
     data : XArrayType
         :class:`xr.Dataset` or :class:`xr.DataArray` with longitude dimension
+    bound : float, optional
+        Lower bound of the domain.
+        Output domain will be [bound, 360 + bound).
+        Defaults to -180, which results in longitude domain [-180, 180).
 
 
     Returns
     -------
     XArrayType
-        :class:`xr.Dataset` or :class:`xr.DataArray` with longitude values on [-180, 180).
+        :class:`xr.Dataset` or :class:`xr.DataArray` with longitude values on [a, 360 + a).
     """
     return data.assign_coords(
-        longitude=((data["longitude"].values + 180.0) % 360.0) - 180.0
+        longitude=((data["longitude"].values - bound) % 360.0) + bound
     ).sortby("longitude", ascending=True)
 
 
@@ -2389,9 +2510,12 @@ def originates_from_ecmwf(met: MetDataset | MetDataArray) -> bool:
     - :class:`HRES`
 
     """
-    return met.attrs.get("met_source") in ("ERA5", "HRES") or "ecmwf" in met.attrs.get(
-        "history", ""
-    )
+    if isinstance(met, MetDataset):
+        try:
+            return met.provider_attr == "ECMWF"
+        except KeyError:
+            pass
+    return "ecmwf" in met.attrs.get("history", "")
 
 
 def _load(hash: str, cachestore: CacheStore, chunks: dict[str, int]) -> xr.Dataset:

pycontrails/core/met_var.py

@@ -2,24 +2,24 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+import dataclasses
 
 
-@dataclass
+@dataclasses.dataclass(frozen=True)
 class MetVariable:
     """Met variable defined using CF, ECMWF, and WMO conventions.
 
     When there is a conflict between CF, ECMWF, and WMO conventions,
-    CF takes precendence, then WMO, then ECMWF.
+    CF takes precedence, then WMO, then ECMWF.
 
     References
     ----------
     - `CF Standard Names, version 77
       <https://cfconventions.org/Data/cf-standard-names/77/build/cf-standard-name-table.html>`_
     - `ECMWF Parameter Database <https://apps.ecmwf.int/codes/grib/param-db>`_
-    - `NCEP Grib v1 Code Table <https://www.nco.ncep.noaa.gov/pmb/docs/on388/table2.html>`
+    - `NCEP Grib v1 Code Table <https://www.nco.ncep.noaa.gov/pmb/docs/on388/table2.html>`_
     - `WMO Codes Registry, Grib Edition 2 <https://codes.wmo.int/_grib2>`_
-    - `NCEP Grib v2 Code Table <https://www.nco.ncep.noaa.gov/pmb/docs/grib2/grib2_doc/grib2_table4-2.shtml>`
+    - `NCEP Grib v2 Code Table <https://www.nco.ncep.noaa.gov/pmb/docs/grib2/grib2_doc/grib2_table4-2.shtml>`_
 
     Used for defining support parameters in a grib-like fashion.
     """  # noqa: E501
@@ -70,7 +70,7 @@ class MetVariable:
         ValueError
             If any of the inputs have an unknown :attr:`level_type`.
         """
-        level_types = ["surface", "isobaricInPa", "isobaricInhPa", "nominalTop"]
+        level_types = ("surface", "isobaricInPa", "isobaricInhPa", "nominalTop")
         if self.level_type is not None and self.level_type not in level_types:
             raise ValueError(f"`level_type` must be one of {level_types}")
 
@@ -102,8 +102,8 @@ class MetVariable:
         """
 
         # return only these keys if they are not None
-        keys = ["short_name", "standard_name", "long_name", "units"]
-        return {k: getattr(self, k) for k in keys if getattr(self, k, None) is not None}
+        keys = ("short_name", "standard_name", "long_name", "units")
+        return {k: v for k in keys if (v := getattr(self, k, None)) is not None}
 
 
 # ----
@@ -135,7 +135,7 @@ Altitude = MetVariable(
     amip="ta",
     description=(
         "Altitude is the (geometric) height above the geoid, which is the "
-        "reference  geopotential surface. The geoid is similar to mean sea level."
+        "reference geopotential surface. The geoid is similar to mean sea level."
     ),
 )
 

pycontrails/core/models.py

@@ -692,6 +692,33 @@ class Model(ABC):
         if self.params["interpolation_use_indices"] and isinstance(self.source, GeoVectorDataset):
             self.source._invalidate_indices()
 
+    def transfer_met_source_attrs(self, source: SourceType | None = None) -> None:
+        """Transfer met source metadata from :attr:`met` to ``source``."""
+
+        if self.met is None:
+            return
+
+        source = source or self.source
+        try:
+            source.attrs["met_source_provider"] = self.met.provider_attr
+        except KeyError:
+            pass
+
+        try:
+            source.attrs["met_source_dataset"] = self.met.dataset_attr
+        except KeyError:
+            pass
+
+        try:
+            source.attrs["met_source_product"] = self.met.product_attr
+        except KeyError:
+            pass
+
+        try:
+            source.attrs["met_source_forecast_time"] = self.met.attrs["forecast_time"]
+        except KeyError:
+            pass
+
 
 def _interp_grid_to_grid(
     met_key: str,