pycontrails 0.54.0__cp311-cp311-macosx_11_0_arm64.whl → 0.54.2__cp311-cp311-macosx_11_0_arm64.whl

pycontrails/core/met.py

@@ -6,6 +6,7 @@ import hashlib
 import json
 import logging
 import pathlib
+import sys
 import typing
 import warnings
 from abc import ABC, abstractmethod
@@ -29,11 +30,20 @@ from typing import (
     overload,
 )
 
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+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 interpolation
 from pycontrails.core import vector as vector_module
@@ -70,7 +80,7 @@ class MetBase(ABC, Generic[XArrayType]):
     cachestore: CacheStore | None
 
     #: Default dimension order for DataArray or Dataset (x, y, z, t)
-    dim_order: tuple[Hashable, Hashable, Hashable, Hashable] = (
+    dim_order = (
         "longitude",
         "latitude",
         "level",
@@ -97,17 +107,18 @@ class MetBase(ABC, Generic[XArrayType]):
         ValueError
             If data does not contain all four coordinates (longitude, latitude, level, time).
         """
-        for dim in self.dim_order:
-            if dim not in self.data.dims:
-                if dim == "level":
-                    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:
-                    msg = f"Meteorology data must contain dimension '{dim}'."
-                raise ValueError(msg)
+        missing = set(self.dim_order).difference(self.data.dims)
+        if not missing:
+            return
+
+        dim = sorted(missing)
+        msg = f"Meteorology data must contain dimension(s): {dim}."
+        if "level" in dim:
+            msg += (
+                " For single level data, set 'level' coordinate to constant -1 "
+                "using `ds = ds.expand_dims({'level': [-1]})`"
+            )
+        raise ValueError(msg)
 
     def _validate_longitude(self) -> None:
         """Check longitude bounds.
@@ -123,8 +134,8 @@ class MetBase(ABC, Generic[XArrayType]):
         if longitude.dtype != COORD_DTYPE:
             raise ValueError(
                 "Longitude values must be of type float64. "
-                "Initiate with 'copy=True' to convert to float64. "
-                "Initiate with 'validate=False' to skip validation."
+                "Instantiate with 'copy=True' to convert to float64. "
+                "Instantiate with 'validate=False' to skip validation."
             )
 
         if self.is_wrapped:
@@ -167,8 +178,8 @@ class MetBase(ABC, Generic[XArrayType]):
         if latitude.dtype != COORD_DTYPE:
             raise ValueError(
                 "Latitude values must be of type float64. "
-                "Initiate with 'copy=True' to convert to float64. "
-                "Initiate with 'validate=False' to skip validation."
+                "Instantiate with 'copy=True' to convert to float64. "
+                "Instantiate with 'validate=False' to skip validation."
             )
 
         if latitude[0] < -90.0:
@@ -192,10 +203,10 @@ class MetBase(ABC, Generic[XArrayType]):
         """
         indexes = self.indexes
         if not np.all(np.diff(indexes["time"]) > np.timedelta64(0, "ns")):
-            raise ValueError("Coordinate `time` not sorted. Initiate with `copy=True`.")
+            raise ValueError("Coordinate 'time' not sorted. Instantiate with 'copy=True'.")
         for coord in self.dim_order[:3]:  # exclude time, the 4th dimension
             if not np.all(np.diff(indexes[coord]) > 0.0):
-                raise ValueError(f"Coordinate '{coord}' not sorted. Initiate with 'copy=True'.")
+                raise ValueError(f"Coordinate '{coord}' not sorted. Instantiate with 'copy=True'.")
 
     def _validate_transpose(self) -> None:
         """Check that data is transposed according to :attr:`dim_order`."""
@@ -204,11 +215,11 @@ class MetBase(ABC, Generic[XArrayType]):
             if da.dims != self.dim_order:
                 if key is not None:
                     msg = (
-                        f"Data dimension not transposed on variable '{key}'. Initiate with"
+                        f"Data dimension not transposed on variable '{key}'. Instantiate with"
                         " 'copy=True'."
                     )
                 else:
-                    msg = "Data dimension not transposed. Initiate with 'copy=True'."
+                    msg = "Data dimension not transposed. Instantiate with 'copy=True'."
                 raise ValueError(msg)
 
         data = self.data
@@ -228,6 +239,12 @@ class MetBase(ABC, Generic[XArrayType]):
         self._validate_longitude()
         self._validate_latitude()
         self._validate_transpose()
+        if self.data["level"].dtype != COORD_DTYPE:
+            raise ValueError(
+                "Level values must be of type float64. "
+                "Instantiate with 'copy=True' to convert to float64. "
+                "Instantiate with 'validate=False' to skip validation."
+            )
 
     def _preprocess_dims(self, wrap_longitude: bool) -> None:
         """Confirm DataArray or Dataset include required dimension in a consistent format.
@@ -363,16 +380,6 @@ class MetBase(ABC, Generic[XArrayType]):
             "time": variables["time"].to_numpy(),
         }
 
-    @property
-    def variables(self) -> dict[Hashable, pd.Index]:
-        """See :attr:`indexes`."""
-        warnings.warn(
-            "The 'variables' property is deprecated and will be removed in a future release. "
-            "Use 'indexes' instead.",
-            DeprecationWarning,
-        )
-        return self.indexes
-
     @property
     def indexes(self) -> dict[Hashable, pd.Index]:
         """Low level access to underlying :attr:`data` indexes.
@@ -745,8 +752,8 @@ class MetDataset(MetBase):
         except KeyError as e:
             raise KeyError(
                 f"Variable {key} not found. Available variables: {', '.join(self.data.data_vars)}. "
-                "To get items (e.g. `time` or `level`) from underlying `xr.Dataset` object, "
-                "use the `data` attribute."
+                "To get items (e.g. 'time' or 'level') from underlying xr.Dataset object, "
+                "use the 'data' attribute."
             ) from e
         return MetDataArray(da, copy=False, validate=False)
 
@@ -866,13 +873,13 @@ class MetDataset(MetBase):
         return key in self.data
 
     @property
-    @overrides
+    @override
     def shape(self) -> tuple[int, int, int, int]:
         sizes = self.data.sizes
         return sizes["longitude"], sizes["latitude"], sizes["level"], sizes["time"]
 
     @property
-    @overrides
+    @override
     def size(self) -> int:
         return np.prod(self.shape).item()
 
@@ -981,7 +988,7 @@ class MetDataset(MetBase):
         hash: str,
         cachestore: CacheStore | None = None,
         chunks: dict[str, int] | None = None,
-    ) -> MetDataset:
+    ) -> Self:
         """Load saved intermediate from :attr:`cachestore`.
 
         Parameters
@@ -996,7 +1003,7 @@ class MetDataset(MetBase):
 
         Returns
         -------
-        MetDataset
+        Self
             New MetDataArray with loaded data.
         """
         cachestore = cachestore or DiskCacheStore()
@@ -1018,14 +1025,14 @@ class MetDataset(MetBase):
             cachestore=self.cachestore,
         )
 
-    @overrides
+    @override
     def broadcast_coords(self, name: str) -> xr.DataArray:
         da = xr.ones_like(self.data[next(iter(self.data.keys()))]) * self.data[name]
         da.name = name
 
         return da
 
-    @overrides
+    @override
     def downselect(self, bbox: tuple[float, ...]) -> MetDataset:
         data = downselect(self.data, bbox)
         return MetDataset(data, cachestore=self.cachestore, copy=False)
@@ -1057,14 +1064,13 @@ class MetDataset(MetBase):
         >>> era5 = ERA5(time=times, variables=variables, pressure_levels=levels)
         >>> met = era5.open_metdataset()
         >>> met.to_vector(transfer_attrs=False)
-        GeoVectorDataset [6 keys x 4152960 length, 1 attributes]
+        GeoVectorDataset [6 keys x 4152960 length, 0 attributes]
             Keys: longitude, latitude, level, time, air_temperature, ..., specific_humidity
             Attributes:
             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.8, 11783.9]
-            crs                 EPSG:4326
 
         """
         coords_keys = self.data.dims
@@ -1183,7 +1189,7 @@ class MetDataset(MetBase):
         latitude: npt.ArrayLike | float,
         level: npt.ArrayLike | float,
         time: npt.ArrayLike | np.datetime64,
-    ) -> MetDataset:
+    ) -> Self:
         r"""Create a :class:`MetDataset` containing a coordinate skeleton from coordinate arrays.
 
         Parameters
@@ -1197,7 +1203,7 @@ class MetDataset(MetBase):
 
         Returns
         -------
-        MetDataset
+        Self
             MetDataset with no variables.
 
         Examples
@@ -1283,7 +1289,7 @@ class MetDataset(MetBase):
         return cls(xr.Dataset({}, coords=coords))
 
     @classmethod
-    def from_zarr(cls, store: Any, **kwargs: Any) -> MetDataset:
+    def from_zarr(cls, store: Any, **kwargs: Any) -> Self:
         """Create a :class:`MetDataset` from a path to a Zarr store.
 
         Parameters
@@ -1295,7 +1301,7 @@ class MetDataset(MetBase):
 
         Returns
         -------
-        MetDataset
+        Self
             MetDataset with data from Zarr store.
         """
         kwargs.setdefault("storage_options", {"read_only": True})
@@ -1374,20 +1380,9 @@ class MetDataArray(MetBase):
         copy: bool = True,
         validate: bool = True,
         name: Hashable | None = None,
-        **kwargs: Any,
     ) -> None:
-        # init cache
         self.cachestore = cachestore
 
-        # try to create DataArray out of input data and **kwargs
-        if not isinstance(data, xr.DataArray):
-            warnings.warn(
-                "Input 'data' must be an xarray DataArray. "
-                "Passing arbitrary kwargs will be removed in future versions.",
-                DeprecationWarning,
-            )
-            data = xr.DataArray(data, **kwargs)
-
         if copy:
             self.data = data.copy()
             self._preprocess_dims(wrap_longitude)
@@ -1415,8 +1410,9 @@ class MetDataArray(MetBase):
 
         See Also
         --------
-        - :meth:`xr.Dataset.load`
-        - :meth:`xr.DataArray.load`
+        :meth:`xarray.Dataset.load`
+        :meth:`xarray.DataArray.load`
+
         """
         if not self.in_memory:
             self._check_memory("Extracting numpy array from")
@@ -1447,12 +1443,12 @@ class MetDataArray(MetBase):
         return np.array_equal(self.data, self.data.astype(bool))
 
     @property
-    @overrides
+    @override
     def size(self) -> int:
         return self.data.size
 
     @property
-    @overrides
+    @override
     def shape(self) -> tuple[int, int, int, int]:
         # https://github.com/python/mypy/issues/1178
         return typing.cast(tuple[int, int, int, int], self.data.shape)
@@ -1842,7 +1838,7 @@ class MetDataArray(MetBase):
         hash: str,
         cachestore: CacheStore | None = None,
         chunks: dict[str, int] | None = None,
-    ) -> MetDataArray:
+    ) -> Self:
         """Load saved intermediate from :attr:`cachestore`.
 
         Parameters
@@ -2031,7 +2027,7 @@ class MetDataArray(MetBase):
         See Also
         --------
         :meth:`to_polyhedra`
-        :func:`pycontrails.core.polygons.find_multipolygons`
+        :func:`polygons.find_multipolygons`
 
         Examples
         --------
@@ -2235,7 +2231,7 @@ class MetDataArray(MetBase):
 
         Returns
         -------
-        dict | :class:`o3d.geometry.TriangleMesh`
+        dict | open3d.geometry.TriangleMesh
             Python representation of geojson object or `Open3D Triangle Mesh
             <http://www.open3d.org/docs/release/tutorial/geometry/mesh.html>`_ depending on the
             `return_type` parameter.
@@ -2249,8 +2245,9 @@ class MetDataArray(MetBase):
 
         See Also
         --------
-        :meth:`to_polygons`
-        `skimage.measure.marching_cubes <https://scikit-image.org/docs/dev/api/skimage.measure.html#skimage.measure.marching_cubes>`_
+        :meth:`to_polygon_feature`
+        :func:`skimage.measure.marching_cubes`
+        :class:`open3d.geometry.TriangleMesh`
 
         Notes
         -----
@@ -2395,14 +2392,14 @@ class MetDataArray(MetBase):
             )
         return mesh
 
-    @overrides
+    @override
     def broadcast_coords(self, name: str) -> xr.DataArray:
         da = xr.ones_like(self.data) * self.data[name]
         da.name = name
 
         return da
 
-    @overrides
+    @override
     def downselect(self, bbox: tuple[float, ...]) -> MetDataArray:
         data = downselect(self.data, bbox)
         return MetDataArray(data, cachestore=self.cachestore)