pycontrails 0.49.3__cp310-cp310-macosx_10_9_x86_64.whl → 0.49.5__cp310-cp310-macosx_10_9_x86_64.whl

pycontrails/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.49.3'
-__version_tuple__ = version_tuple = (0, 49, 3)
+__version__ = version = '0.49.5'
+__version_tuple__ = version_tuple = (0, 49, 5)

pycontrails/core/datalib.py

@@ -353,7 +353,7 @@ class MetDataSource(abc.ABC):
         self,
         time: TimeInput | None,
         variables: VariableInput,
-        pressure_levels: PressureLevelInput = [-1],
+        pressure_levels: PressureLevelInput = -1,
         paths: str | list[str] | pathlib.Path | list[pathlib.Path] | None = None,
         grid: float | None = None,
         **kwargs: Any,

pycontrails/core/flight.py

@@ -936,8 +936,7 @@ class Flight(GeoVectorDataset):
 
         # Remove original index if requested
         if not keep_original_index:
-            filt = df.index.isin(t)
-            df = df.loc[filt]
+            df = df.loc[t]
 
         # finally reset index
         df = df.reset_index()
@@ -1200,11 +1199,13 @@ class Flight(GeoVectorDataset):
         dict[str, Any]
             Python representation of geojson FeatureCollection
         """
-        points = _return_linestring({
-            "longitude": self["longitude"],
-            "latitude": self["latitude"],
-            "altitude": self.altitude,
-        })
+        points = _return_linestring(
+            {
+                "longitude": self["longitude"],
+                "latitude": self["latitude"],
+                "altitude": self.altitude,
+            }
+        )
         geometry = {"type": "LineString", "coordinates": points}
         properties = {
             "start_time": self.time_start.isoformat(),
@@ -2041,7 +2042,7 @@ def _resample_to_freq(df: pd.DataFrame, freq: str) -> tuple[pd.DataFrame, pd.Dat
     """Resample a DataFrame to a given frequency.
 
     This function is used to resample a DataFrame to a given frequency. The new
-    index will include all the original index values and the new esampled-to-freq
+    index will include all the original index values and the new resampled-to-freq
     index values. The "longitude" and "latitude" columns will be linearly interpolated
     to the new index values.
 
@@ -2064,7 +2065,7 @@ def _resample_to_freq(df: pd.DataFrame, freq: str) -> tuple[pd.DataFrame, pd.Dat
     # and the resampled-to-freq index values.
     t0 = df.index[0]
     t1 = df.index[-1]
-    t = pd.date_range(t0, t1, freq=freq).floor(freq)
+    t = pd.date_range(t0, t1, freq=freq, name="time").floor(freq)
     if t[0] < t0:
         t = t[1:]
 
@@ -2072,8 +2073,7 @@ def _resample_to_freq(df: pd.DataFrame, freq: str) -> tuple[pd.DataFrame, pd.Dat
     concat_arr = np.unique(concat_arr)
     concat_index = pd.DatetimeIndex(concat_arr, name="time", copy=False)
 
-    out = pd.DataFrame(index=concat_index, columns=df.columns, dtype=float)
-    out.loc[df.index] = df
+    out = df.reindex(concat_index)
 
     # Linearly interpolate small horizontal gap
     coords = ["longitude", "latitude"]

pycontrails/core/interpolation.py

@@ -66,7 +66,8 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
         fill_value: float | np.float64 | None,
     ):
         if values.dtype not in (np.float32, np.float64):
-            raise ValueError("values must be a float array")
+            msg = f"values must be a float array, not {values.dtype}"
+            raise ValueError(msg)
 
         self.grid = points
         self.values = values
@@ -94,7 +95,8 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
                 g0 = self.grid[i][0]
                 g1 = self.grid[i][-1]
                 if not (np.all(p >= g0) and np.all(p <= g1)):
-                    raise ValueError(f"One of the requested xi is out of bounds in dimension {i}")
+                    msg = f"One of the requested xi is out of bounds in dimension {i}"
+                    raise ValueError(msg)
 
             return np.zeros(xi.shape[0], dtype=bool)
 
@@ -213,7 +215,8 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
             # np.interp could be better ... although that may also promote the dtype
             return rgi_cython.evaluate_linear_1d(values, indices, norm_distances, out)
 
-        raise ValueError(f"Invalid number of dimensions: {ndim}")
+        msg = f"Invalid number of dimensions: {ndim}"
+        raise ValueError(msg)
 
 
 def _floatize_time(
@@ -442,18 +445,16 @@ def interp(
         coords = {"longitude": longitude, "latitude": latitude, "level": level, "time": time}
         da = _localize(da, coords)
 
-    # Using da.coords.variables is slightly more performant than da["longitude"].values
-    variables = da.coords.variables
-    x = variables["longitude"].values
-    y = variables["latitude"].values
-    z = variables["level"].values
+    indexes = da._indexes
+    x = indexes["longitude"].index.to_numpy()  # type: ignore[attr-defined]
+    y = indexes["latitude"].index.to_numpy()  # type: ignore[attr-defined]
+    z = indexes["level"].index.to_numpy()  # type: ignore[attr-defined]
     if any(v.dtype != np.float64 for v in (x, y, z)):
-        raise ValueError(
-            "da must have float64 dtype for longitude, latitude, and level coordinates"
-        )
+        msg = "da must have float64 dtype for longitude, latitude, and level coordinates"
+        raise ValueError(msg)
 
     # Convert t and time to float64
-    t = variables["time"].values
+    t = indexes["time"].index.to_numpy()  # type: ignore[attr-defined]
     offset = t[0]
     t = _floatize_time(t, offset)
 
@@ -526,9 +527,11 @@ def _linear_interp_with_indices(
     indices: RGIArtifacts | None,
 ) -> tuple[npt.NDArray[np.float64], RGIArtifacts]:
     if interp.method != "linear":
-        raise ValueError("Parameter 'indices' is only supported for 'method=linear'")
+        msg = "Parameter 'indices' is only supported for 'method=linear'"
+        raise ValueError(msg)
     if localize:
-        raise ValueError("Parameter 'indices' is only supported for 'localize=False'")
+        msg = "Parameter 'indices' is only supported for 'localize=False'"
+        raise ValueError(msg)
 
     if indices is None:
         assert xi is not None, "xi must be provided if indices is None"
@@ -604,7 +607,10 @@ class EmissionsProfileInterpolator:
     """
 
     def __init__(
-        self, xp: npt.NDArray[np.float64], fp: npt.NDArray[np.float64], drop_duplicates: bool = True
+        self,
+        xp: npt.NDArray[np.float64],
+        fp: npt.NDArray[np.float64],
+        drop_duplicates: bool = True,
     ) -> None:
         if drop_duplicates:
             # Using np.diff to detect duplicates ... this assumes xp is sorted.
@@ -622,13 +628,17 @@ class EmissionsProfileInterpolator:
 
     def _validate(self) -> None:
         if not len(self.xp):
-            raise ValueError("xp must not be empty")
+            msg = "xp must not be empty"
+            raise ValueError(msg)
         if len(self.xp) != len(self.fp):
-            raise ValueError("xp and fp must have the same length")
+            msg = "xp and fp must have the same length"
+            raise ValueError(msg)
         if not np.all(np.diff(self.xp) > 0.0):
-            raise ValueError("xp must be strictly increasing")
+            msg = "xp must be strictly increasing"
+            raise ValueError(msg)
         if np.any(np.isnan(self.xp)):
-            raise ValueError("xp must not contain nan values")
+            msg = "xp must not contain nan values"
+            raise ValueError(msg)
 
     def interp(self, x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
         """Interpolate x against xp and fp.

pycontrails/core/met.py

@@ -23,6 +23,7 @@ from typing import (
 
 import numpy as np
 import numpy.typing as npt
+import pandas as pd
 import xarray as xr
 from overrides import overrides
 
@@ -110,7 +111,7 @@ class MetBase(ABC, Generic[XArrayType]):
         ValueError
             If longitude values are not contained in the interval [-180, 180].
         """
-        longitude = self.variables["longitude"].values
+        longitude = self.indexes["longitude"].to_numpy()
         if longitude.dtype != COORD_DTYPE:
             raise ValueError(
                 "Longitude values must be of type float64. "
@@ -154,7 +155,7 @@ class MetBase(ABC, Generic[XArrayType]):
         ValueError
             If latitude values are not contained in the interval [-90, 90].
         """
-        latitude = self.variables["latitude"].values
+        latitude = self.indexes["latitude"].to_numpy()
         if latitude.dtype != COORD_DTYPE:
             raise ValueError(
                 "Latitude values must be of type float64. "
@@ -181,10 +182,11 @@ class MetBase(ABC, Generic[XArrayType]):
         ValueError
             If one of the coordinates is not sorted.
         """
-        if not np.all(np.diff(self.variables["time"]) > np.timedelta64(0, "ns")):
+        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`.")
         for coord in self.dim_order[:3]:  # exclude time, the 4th dimension
-            if not np.all(np.diff(self.variables[coord]) > 0.0):
+            if not np.all(np.diff(indexes[coord]) > 0.0):
                 raise ValueError(f"Coordinate '{coord}' not sorted. Initiate with 'copy=True'.")
 
     def _validate_transpose(self) -> None:
@@ -251,8 +253,9 @@ class MetBase(ABC, Generic[XArrayType]):
         self._validate_dim_contains_coords()
 
         # Ensure spatial coordinates all have dtype COORD_DTYPE
+        indexes = self.indexes
         for coord in ("longitude", "latitude", "level"):
-            arr = self.variables[coord].values
+            arr = indexes[coord].to_numpy()
             if arr.dtype != COORD_DTYPE:
                 self.data[coord] = arr.astype(COORD_DTYPE)
 
@@ -265,7 +268,7 @@ class MetBase(ABC, Generic[XArrayType]):
         if not self.is_wrapped:
             # Ensure longitude is contained in interval [-180, 180)
             # If longitude has value at 180, we might not want to shift it?
-            lon = self.variables["longitude"].values
+            lon = self.indexes["longitude"].to_numpy()
 
             # This longitude shifting can give rise to precision errors with float32
             # Only shift if necessary
@@ -288,48 +291,12 @@ class MetBase(ABC, Generic[XArrayType]):
         # single level data
         if self.is_single_level:
             # add level attributes to reflect surface level
-            self.data["level"].attrs.update(units="", long_name="Single Level")
+            level_attrs = self.data["level"].attrs
+            if not level_attrs:
+                level_attrs.update(units="", long_name="Single Level")
             return
 
-        # pressure level data
-        level = self.variables["level"].values
-
-        # add pressure level attributes
-        self.data["level"].attrs.update(units="hPa", long_name="Pressure", positive="down")
-
-        # add altitude and air_pressure
-
-        # XXX: use the dtype of the data to determine the precision of these coordinates
-        # There are two competing conventions here:
-        # - coordinate data should be float64
-        # - gridded data is typically float32
-        # - air_pressure and altitude often play both roles
-        # It is more important for air_pressure and altitude to be grid-aligned than to be
-        # coordinate-aligned, so we use the dtype of the data to determine the precision of
-        # these coordinates
-        if isinstance(self.data, xr.Dataset):
-            dtype = np.result_type(*self.data.data_vars.values(), np.float32)
-        else:
-            dtype = self.data.dtype
-
-        level = level.astype(dtype)
-        air_pressure = level * 100.0
-        altitude = units.pl_to_m(level)
-        self.data = self.data.assign_coords({"air_pressure": ("level", air_pressure)})
-        self.data = self.data.assign_coords({"altitude": ("level", altitude)})
-
-        # add air_pressure units and long name attributes
-        self.data.coords["air_pressure"].attrs.update(
-            standard_name=AirPressure.standard_name,
-            long_name=AirPressure.long_name,
-            units=AirPressure.units,
-        )
-        # add altitude units and long name attributes
-        self.data.coords["altitude"].attrs.update(
-            standard_name=Altitude.standard_name,
-            long_name=Altitude.long_name,
-            units=Altitude.units,
-        )
+        self.data = _add_vertical_coords(self.data)
 
     @property
     def hash(self) -> str:
@@ -382,29 +349,36 @@ class MetBase(ABC, Generic[XArrayType]):
         dict[str, np.ndarray]
             Dictionary of coordinates
         """
+        variables = self.indexes
         return {
-            "longitude": self.variables["longitude"].values,
-            "latitude": self.variables["latitude"].values,
-            "level": self.variables["level"].values,
-            "time": self.variables["time"].values,
+            "longitude": variables["longitude"].to_numpy(),
+            "latitude": variables["latitude"].to_numpy(),
+            "level": variables["level"].to_numpy(),
+            "time": variables["time"].to_numpy(),
         }
 
     @property
-    def variables(self) -> dict[Hashable, xr.Variable]:
-        """Low level access to underlying :attr:`data` variables.
+    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
 
-        This method is typically is faster for accessing coordinate variables.
+    @property
+    def indexes(self) -> dict[Hashable, pd.Index]:
+        """Low level access to underlying :attr:`data` indexes.
+
+        This method is typically is faster for accessing coordinate indexes.
 
         .. versionadded:: 0.25.2
 
         Returns
         -------
-        dict[Hashable, xr.Variable]
-            Dictionary of variables. The type is actually..
-
-                xarray.core.utils.Frozen[Any, xr.Variable]
-
-            In practice, this behaves like a dictionary.
+        dict[Hashable, pd.Index]
+            Dictionary of indexes.
 
         Examples
         --------
@@ -414,16 +388,14 @@ class MetBase(ABC, Generic[XArrayType]):
         >>> levels = [200, 300]
         >>> era5 = ERA5(times, variables, levels)
         >>> mds = era5.open_metdataset()
-        >>> mds.variables["level"].values  # faster access than mds.data["level"]
+        >>> mds.indexes["level"].to_numpy()
         array([200., 300.])
 
         >>> mda = mds["air_temperature"]
-        >>> mda.variables["level"].values  # faster access than mda.data["level"]
+        >>> mda.indexes["level"].to_numpy()
         array([200., 300.])
         """
-        if isinstance(self.data, xr.Dataset):
-            return self.data.variables  # type: ignore[return-value]
-        return self.data.coords.variables
+        return {k: v.index for k, v in self.data._indexes.items()}  # type: ignore[attr-defined]
 
     @property
     def is_wrapped(self) -> bool:
@@ -449,7 +421,7 @@ class MetBase(ABC, Generic[XArrayType]):
         --------
         :func:`pycontrails.physics.geo.advect_longitude`
         """
-        longitude = self.variables["longitude"].values
+        longitude = self.indexes["longitude"].to_numpy()
         return _is_wrapped(longitude)
 
     @property
@@ -464,7 +436,7 @@ class MetBase(ABC, Generic[XArrayType]):
         bool
             If instance contains single level data.
         """
-        level = self.variables["level"].values
+        level = self.indexes["level"].to_numpy()
         return len(level) == 1 and level[0] == -1
 
     @abstractmethod
@@ -569,6 +541,82 @@ class MetBase(ABC, Generic[XArrayType]):
         """
         return _is_zarr(self.data)
 
+    def downselect_met(
+        self,
+        met: MetDataType,
+        *,
+        longitude_buffer: tuple[float, float] = (0.0, 0.0),
+        latitude_buffer: tuple[float, float] = (0.0, 0.0),
+        level_buffer: tuple[float, float] = (0.0, 0.0),
+        time_buffer: tuple[np.timedelta64, np.timedelta64] = (
+            np.timedelta64(0, "h"),
+            np.timedelta64(0, "h"),
+        ),
+        copy: bool = True,
+    ) -> MetDataType:
+        """Downselect ``met`` to encompass a spatiotemporal region of the data.
+
+        .. warning::
+
+            This method is analogous to :meth:`GeoVectorDataset.downselect_met`.
+            It does not change the instance data, but instead operates on the
+            ``met`` input. This method is different from :meth:`downselect` which
+            operates on the instance data.
+
+        Parameters
+        ----------
+        met : MetDataset | MetDataArray
+            MetDataset or MetDataArray to downselect.
+        longitude_buffer : tuple[float, float], optional
+            Extend longitude domain past by ``longitude_buffer[0]`` on the low side
+            and ``longitude_buffer[1]`` on the high side.
+            Units must be the same as class coordinates.
+            Defaults to ``(0, 0)`` degrees.
+        latitude_buffer : tuple[float, float], optional
+            Extend latitude domain past by ``latitude_buffer[0]`` on the low side
+            and ``latitude_buffer[1]`` on the high side.
+            Units must be the same as class coordinates.
+            Defaults to ``(0, 0)`` degrees.
+        level_buffer : tuple[float, float], optional
+            Extend level domain past by ``level_buffer[0]`` on the low side
+            and ``level_buffer[1]`` on the high side.
+            Units must be the same as class coordinates.
+            Defaults to ``(0, 0)`` [:math:`hPa`].
+        time_buffer : tuple[np.timedelta64, np.timedelta64], optional
+            Extend time domain past by ``time_buffer[0]`` on the low side
+            and ``time_buffer[1]`` on the high side.
+            Units must be the same as class coordinates.
+            Defaults to ``(np.timedelta64(0, "h"), np.timedelta64(0, "h"))``.
+        copy : bool
+            If returned object is a copy or view of the original. True by default.
+
+        Returns
+        -------
+        MetDataset | MetDataArray
+            Copy of downselected MetDataset or MetDataArray.
+        """
+        indexes = self.indexes
+        lon = indexes["longitude"].to_numpy()
+        lat = indexes["latitude"].to_numpy()
+        level = indexes["level"].to_numpy()
+        time = indexes["time"].to_numpy()
+
+        vector = vector_module.GeoVectorDataset(
+            longitude=[lon.min(), lon.max()],
+            latitude=[lat.min(), lat.max()],
+            level=[level.min(), level.max()],
+            time=[time.min(), time.max()],
+        )
+
+        return vector.downselect_met(
+            met,
+            longitude_buffer=longitude_buffer,
+            latitude_buffer=latitude_buffer,
+            level_buffer=level_buffer,
+            time_buffer=time_buffer,
+            copy=copy,
+        )
+
 
 class MetDataset(MetBase):
     """Meteorological dataset with multiple variables.
@@ -668,6 +716,8 @@ class MetDataset(MetBase):
                 raise ValueError("Set 'copy=True' when using 'wrap_longitude=True'.")
             self.data = data
             self._validate_dims()
+            if not self.is_single_level:
+                self.data = _add_vertical_coords(self.data)
 
     def __getitem__(self, key: Hashable) -> MetDataArray:
         """Return DataArray of variable ``key`` cast to a :class:`MetDataArray` object.
@@ -1014,22 +1064,22 @@ class MetDataset(MetBase):
             crs                 EPSG:4326
 
         """
-        coords_keys = [str(key) for key in self.data.dims]  # str not in Hashable
-        coords_vals = [self.variables[key].values for key in coords_keys]
+        coords_keys = self.data.dims
+        variables = self.indexes
+        coords_vals = [variables[key].values for key in coords_keys]
         coords_meshes = np.meshgrid(*coords_vals, indexing="ij")
-        raveled_coords = [mesh.ravel() for mesh in coords_meshes]
+        raveled_coords = (mesh.ravel() for mesh in coords_meshes)
         data = dict(zip(coords_keys, raveled_coords))
 
-        vector = vector_module.GeoVectorDataset(data, copy=False)
-        for key in self:
+        out = vector_module.GeoVectorDataset(data, copy=False)
+        for key, da in self.data.items():
             # The call to .values here will load the data if it is lazy
-            vector[key] = self[key].data.values.ravel()
+            out[key] = da.values.ravel()  # type: ignore[index]
 
         if transfer_attrs:
-            # vector.attrs expects keys to be strings .... we'll get an error
-            # if we cannot cast here
-            vector.attrs.update({str(k): v for k, v in self.attrs.items()})
-        return vector
+            out.attrs.update(self.attrs)  # type: ignore[arg-type]
+
+        return out
 
     def _get_pycontrails_attr_template(
         self,
@@ -1157,15 +1207,15 @@ class MetDataset(MetBase):
         >>> met = MetDataset.from_coords(longitude, latitude, level, time)
         >>> met
         MetDataset with data:
-        <xarray.Dataset>
+        <xarray.Dataset> Size: 360B
         Dimensions:       (longitude: 20, latitude: 20, level: 2, time: 1)
         Coordinates:
-          * longitude     (longitude) float64 0.0 0.5 1.0 1.5 2.0 ... 8.0 8.5 9.0 9.5
-          * latitude      (latitude) float64 0.0 0.5 1.0 1.5 2.0 ... 7.5 8.0 8.5 9.0 9.5
-          * level         (level) float64 250.0 300.0
-          * time          (time) datetime64[ns] 2019-01-01
-            air_pressure  (level) float32 2.5e+04 3e+04
-            altitude      (level) float32 1.036e+04 9.164e+03
+          * longitude     (longitude) float64 160B 0.0 0.5 1.0 1.5 ... 8.0 8.5 9.0 9.5
+          * latitude      (latitude) float64 160B 0.0 0.5 1.0 1.5 ... 8.0 8.5 9.0 9.5
+          * level         (level) float64 16B 250.0 300.0
+          * time          (time) datetime64[ns] 8B 2019-01-01
+            air_pressure  (level) float32 8B 2.5e+04 3e+04
+            altitude      (level) float32 8B 1.036e+04 9.164e+03
         Data variables:
             *empty*
 
@@ -1180,18 +1230,18 @@ class MetDataset(MetBase):
         >>> met["humidity"] = xr.DataArray(np.full(met.shape, 0.5), coords=met.coords)
         >>> met
         MetDataset with data:
-        <xarray.Dataset>
+        <xarray.Dataset> Size: 13kB
         Dimensions:       (longitude: 20, latitude: 20, level: 2, time: 1)
         Coordinates:
-          * longitude     (longitude) float64 0.0 0.5 1.0 1.5 2.0 ... 8.0 8.5 9.0 9.5
-          * latitude      (latitude) float64 0.0 0.5 1.0 1.5 2.0 ... 7.5 8.0 8.5 9.0 9.5
-          * level         (level) float64 250.0 300.0
-          * time          (time) datetime64[ns] 2019-01-01
-            air_pressure  (level) float32 2.5e+04 3e+04
-            altitude      (level) float32 1.036e+04 9.164e+03
+          * longitude     (longitude) float64 160B 0.0 0.5 1.0 1.5 ... 8.0 8.5 9.0 9.5
+          * latitude      (latitude) float64 160B 0.0 0.5 1.0 1.5 ... 8.0 8.5 9.0 9.5
+          * level         (level) float64 16B 250.0 300.0
+          * time          (time) datetime64[ns] 8B 2019-01-01
+            air_pressure  (level) float32 8B 2.5e+04 3e+04
+            altitude      (level) float32 8B 1.036e+04 9.164e+03
         Data variables:
-            temperature   (longitude, latitude, level, time) float64 234.5 ... 234.5
-            humidity      (longitude, latitude, level, time) float64 0.5 0.5 ... 0.5 0.5
+            temperature   (longitude, latitude, level, time) float64 6kB 234.5 ... 234.5
+            humidity      (longitude, latitude, level, time) float64 6kB 0.5 0.5 ... 0.5
 
         >>> # Convert to a GeoVectorDataset
         >>> vector = met.to_vector()
@@ -1888,8 +1938,9 @@ class MetDataArray(MetBase):
         from pycontrails.core import polygon
 
         # Convert to nested lists of coordinates for GeoJSON representation
-        longitude: npt.NDArray[np.float64] = self.variables["longitude"].values
-        latitude: npt.NDArray[np.float64] = self.variables["latitude"].values
+        indexes = self.indexes
+        longitude = indexes["longitude"].to_numpy()
+        latitude = indexes["latitude"].to_numpy()
 
         mp = polygon.find_multipolygon(
             arr,
@@ -2097,9 +2148,9 @@ class MetDataArray(MetBase):
         volume = self.data.sel(time=time).values
 
         # convert from array index back to coordinates
-        longitude = self.variables["longitude"].values
-        latitude = self.variables["latitude"].values
-        altitude = units.pl_to_m(self.variables["level"].values)
+        longitude = self.indexes["longitude"].values
+        latitude = self.indexes["latitude"].values
+        altitude = units.pl_to_m(self.indexes["level"].values)
 
         # Pad volume on all axes to close the volumes
         if closed:
@@ -2290,11 +2341,7 @@ def _wrap_longitude(data: XArrayType) -> XArrayType:
     ValueError
         If longitude values are already wrapped.
     """
-    if isinstance(data, xr.Dataset):
-        lon = data.variables["longitude"].values
-    else:
-        lon = data.coords.variables["longitude"].values
-
+    lon = data._indexes["longitude"].index.to_numpy()  # type: ignore[attr-defined]
     if _is_wrapped(lon):
         raise ValueError("Longitude values are already wrapped")
 
@@ -2359,7 +2406,7 @@ def _extract_2d_arr_and_altitude(
     """
     # Determine level if not specified
     if level is None:
-        level_coord = mda.variables["level"].values
+        level_coord = mda.indexes["level"].values
         if len(level_coord) == 1:
             level = level_coord[0]
         else:
@@ -2370,7 +2417,7 @@ def _extract_2d_arr_and_altitude(
 
     # Determine time if not specified
     if time is None:
-        time_coord = mda.variables["time"].values
+        time_coord = mda.indexes["time"].values
         if len(time_coord) == 1:
             time = time_coord[0]
         else:
@@ -2533,3 +2580,44 @@ def _load(hash: str, cachestore: CacheStore, chunks: dict[str, int]) -> xr.Datas
     """
     disk_path = cachestore.get(f"{hash}*.nc")
     return xr.open_mfdataset(disk_path, chunks=chunks)
+
+
+def _add_vertical_coords(data: XArrayType) -> XArrayType:
+    """Add "air_pressure" and "altitude" coordinates to data."""
+
+    data["level"].attrs.update(units="hPa", long_name="Pressure", positive="down")
+
+    coords = data.coords
+    if "air_pressure" in coords and "altitude" in coords:
+        return data
+
+    # XXX: use the dtype of the data to determine the precision of these coordinates
+    # There are two competing conventions here:
+    # - coordinate data should be float64
+    # - gridded data is typically float32
+    # - air_pressure and altitude often play both roles
+    # It is more important for air_pressure and altitude to be grid-aligned than to be
+    # coordinate-aligned, so we use the dtype of the data to determine the precision of
+    # these coordinates
+    if isinstance(data, xr.Dataset):
+        dtype = np.result_type(*data.data_vars.values(), np.float32)
+    else:
+        dtype = data.dtype
+    level = data["level"].values.astype(dtype, copy=False)
+
+    if "air_pressure" not in coords:
+        data = data.assign_coords(air_pressure=("level", level * 100.0))
+        data.coords["air_pressure"].attrs.update(
+            standard_name=AirPressure.standard_name,
+            long_name=AirPressure.long_name,
+            units=AirPressure.units,
+        )
+    if "altitude" not in coords:
+        data = data.assign_coords(altitude=("level", units.pl_to_m(level)))
+        data.coords["altitude"].attrs.update(
+            standard_name=Altitude.standard_name,
+            long_name=Altitude.long_name,
+            units=Altitude.units,
+        )
+
+    return data