Allow non-numeric coordinates (e.g. time) as spatial axes (#65)

* fix: slicing should use indices, not coordinate values

* fix: use indices for mesh coordinates if coord dtype is not numeric

* feat: Add scales arg to specify scale of an index array made for non-numeric coords

* style: reformat with black

* fix: protect scales value when slicing is None

* fix: Modify slicing behavior

* test: Add two test cases for `test_source`

Author: annehaley

pvxarray/accessor.py

@@ -1,5 +1,6 @@
-from typing import Optional
+from typing import Dict, Optional
 
+import numpy as np
 import pyvista as pv
 import xarray as xr
 
@@ -38,9 +39,13 @@ def loc(self) -> _LocIndexer:
         """Attribute for location based indexing like pandas."""
         return _LocIndexer(self)
 
-    def _get_array(self, key):
+    def _get_array(self, key, scale=1):
         try:
-            return self._obj[key].values
+            values = self._obj[key].values
+            if "float" not in str(values.dtype) and "int" not in str(values.dtype):
+                # non-numeric coordinate, assign array of scaled indices
+                values = np.array(range(len(values))) * scale
+            return values
         except KeyError:
             raise KeyError(
                 f"Key {key} not present in DataArray. Choices are: {list(self._obj.coords.keys())}"
@@ -58,6 +63,7 @@ def mesh(
         order: Optional[str] = None,
         component: Optional[str] = None,
         mesh_type: Optional[str] = None,
+        scales: Optional[Dict] = None,
     ) -> pv.DataSet:
         ndim = 0
         if x is not None:
@@ -80,7 +86,7 @@ def mesh(
             meth = methods[mesh_type]
         except KeyError:
             raise KeyError
-        return meth(self, x=x, y=y, z=z, order=order, component=component)
+        return meth(self, x=x, y=y, z=z, order=order, component=component, scales=scales)
 
     def plot(
         self,

pvxarray/rectilinear.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Dict, Optional
 import warnings
 
 import pyvista as pv
@@ -13,6 +13,7 @@ def mesh(
     z: Optional[str] = None,
     order: Optional[str] = "C",
     component: Optional[str] = None,
+    scales: Optional[Dict] = None,
 ):
     if order is None:
         order = "C"
@@ -22,11 +23,11 @@ def mesh(
         raise ValueError("You must specify at least one dimension as X, Y, or Z.")
     # Construct the mesh
     if x is not None:
-        self._mesh.x = self._get_array(x)
+        self._mesh.x = self._get_array(x, scale=(scales and scales.get(x)) or 1)
     if y is not None:
-        self._mesh.y = self._get_array(y)
+        self._mesh.y = self._get_array(y, scale=(scales and scales.get(y)) or 1)
     if z is not None:
-        self._mesh.z = self._get_array(z)
+        self._mesh.z = self._get_array(z, scale=(scales and scales.get(z)) or 1)
     # Handle data values
     values = self.data
     values_dim = values.ndim

pvxarray/structured.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Dict, Optional
 import warnings
 
 import numpy as np
@@ -41,6 +41,7 @@ def _points(
     y: Optional[str] = None,
     z: Optional[str] = None,
     order: Optional[str] = "F",
+    scales: Optional[Dict] = None,
 ):
     """Generate structured points as new array."""
     if order is None:
@@ -52,11 +53,11 @@ def _points(
             raise ValueError("One dimensional structured grids should be rectilinear grids.")
         raise ValueError("You must specify at least two dimensions as X, Y, or Z.")
     if x is not None:
-        x = self._get_array(x)
+        x = self._get_array(x, scale=(scales and scales.get(x)) or 1)
     if y is not None:
-        y = self._get_array(y)
+        y = self._get_array(y, scale=(scales and scales.get(y)) or 1)
     if z is not None:
-        z = self._get_array(z)
+        z = self._get_array(z, scale=(scales and scales.get(z)) or 1)
     arrs = _coerce_shapes(x, y, z)
     x, y, z = arrs
     arr = [a for a in arrs if a is not None][0]
@@ -78,6 +79,7 @@ def mesh(
     z: Optional[str] = None,
     order: str = "F",
     component: Optional[str] = None,  # TODO
+    scales: Optional[Dict] = None,
 ):
     if order is None:
         order = "F"
@@ -88,7 +90,7 @@ def mesh(
             "StructuredGrid accessor duplicates data - VTK/PyVista data not shared with xarray."
         )
     )
-    points, shape = _points(self, x=x, y=y, z=z, order=order)
+    points, shape = _points(self, x=x, y=y, z=z, order=order, scales=scales)
     self._mesh.points = points
     self._mesh.dimensions = shape
     data = self.data

pvxarray/vtk_source.py

@@ -228,30 +228,21 @@ def _compute_sliced_data_array(self):
             self._sliced_data_array = None
             return None
 
+        indexing = {}
+        if self._slicing is not None:
+            indexing = {
+                k: slice(*v) for k, v in self._slicing.items() if k in [self.x, self.y, self.z]
+            }
+
         if self._time is not None:
-            da = self.data_array[{self._time: self.time_index}]
-        else:
-            da = self.data_array
-
-        if self._z and self._z_index is not None:
-            da = da[{self._z: self.z_index}]
-
-        if self._slicing:
-            indexing = {}
-            for axis in [
-                self.x,
-                self.y,
-                self.z,
-            ]:
-                if axis in self._slicing:
-                    s = self._slicing[axis]
-                    c = da.coords[axis]
-                    sliced_array = np.where(np.logical_and(c >= s[0], c <= s[1]))[0]
-                    sliced_array = sliced_array[:: int(s[2])]
-                    indexing[axis] = sliced_array
-            da = da.isel(**indexing)
-
-        elif self._resolution:
+            indexing.update(**{self._time: self.time_index})
+
+        if self.z and self.z_index is not None:
+            indexing.update(**{self.z: self.z_index})
+
+        da = self.data_array.isel(indexing)
+
+        if self._slicing is None and self._resolution is not None:
             rx, ry, rz = self.resolution_to_sampling_rate(da)
             if da.ndim <= 1:
                 da = da[::rx]
@@ -271,6 +262,7 @@ def _compute_mesh(self):
             order=self._order,
             component=self._component,
             mesh_type=self._mesh_type,
+            scales={k: v[2] for k, v in self._slicing.items()} if self._slicing else {},
         )
         return self._mesh
 

tests/test_source.py

@@ -28,3 +28,44 @@ def test_vtk_source():
     source.resolution = 0.5
     mesh = source.apply()
     assert mesh.n_points < 1325
+
+
+def test_vtk_source_time_as_spatial():
+    ds = xr.tutorial.load_dataset("air_temperature")
+
+    da = ds.air
+    source = PyVistaXarraySource(da, x="lon", y="lat", z="time")
+
+    mesh = source.apply()
+    assert mesh
+    assert mesh.n_points == 3869000
+    assert "air" in mesh.point_data
+
+    assert np.array_equal(mesh["air"], da.values.ravel())
+    assert np.array_equal(mesh.x, da.lon)
+    assert np.array_equal(mesh.y, da.lat)
+    # Z values are indexes instead of datetime objects
+    assert np.array_equal(mesh.z, list(range(da.time.size)))
+
+
+def test_vtk_source_slicing():
+    ds = xr.tutorial.load_dataset("eraint_uvz")
+
+    da = ds.z
+    source = PyVistaXarraySource(
+        da,
+        x="longitude",
+        y="latitude",
+        z="level",
+        time="month",
+    )
+    source.time_index = 1
+    source.slicing = {
+        "latitude": [0, 241, 2],
+        "longitude": [0, 480, 4],
+        "level": [0, 3, 1],
+        "month": [0, 2, 1],  # should be ignored in favor of t_index
+    }
+
+    sliced = source.sliced_data_array
+    assert sliced.shape == (3, 121, 120)