Merge pull request #87 from chrishavlin/fix_indexing_with_time_isel

fix handling of 2d+time datasets with reversed axes

Author: chrishavlin

pyproject.toml

@@ -29,7 +29,7 @@ dependencies=['yt>=4.2.0', 'h5py>=3.4.0', 'pooch>=1.5.1', 'xarray']
 
 [project.optional-dependencies]
 full = ["netCDF4", "scipy", "dask[complete]"]
-test = ["pytest", "pytest-cov"]
+test = ["pytest", "pytest-cov", "cartopy"]
 docs = ["Sphinx==7.2.6", "jinja2==3.1.2", "nbsphinx==0.9.3"]
 
 [tool.black]

yt_xarray/accessor/_readers.py

@@ -32,8 +32,8 @@ def _reader(grid, field_name):
                 # note that the si, ei are exchanged!
                 si0 = si.copy()
                 ei0 = ei.copy()
-                si[idim] = sel_info.global_dims[idim] - ei0[idim]
-                ei[idim] = sel_info.global_dims[idim] - si0[idim]
+                si[idim] = sel_info.global_dims_no_time[idim] - ei0[idim]
+                ei[idim] = sel_info.global_dims_no_time[idim] - si0[idim]
 
         # step 2: this global start index accounts for indexing after any
         # subselections on the xarray DataArray are made. might

yt_xarray/accessor/_xr_to_yt.py

@@ -54,6 +54,7 @@ def __init__(
         self.grid_type = None  # one of _GridType members
         self.cell_widths: list = None
         self.global_dims: list = None
+        self.time_index_number: int = None
         self._process_selection(xr_ds)
 
         self.yt_coord_names = _convert_to_yt_internal_coords(self.selected_coords)
@@ -130,7 +131,8 @@ def _process_selection(self, xr_ds):
         reverse_axis = []  # axes must be positive-monitonic for yt
         reverse_axis_names = []
         global_dims = []  # the global shape
-        for c in full_coords:
+        time_index_number = None
+        for icoord, c in enumerate(full_coords):
             coord_da = getattr(xr_ds, c)  # the full coordinate data array
 
             # check if coordinate values are increasing
@@ -156,6 +158,8 @@ def _process_selection(self, xr_ds):
             sel_or_isel = getattr(coord_da, self.sel_dict_type)
             coord_vals = sel_or_isel(coord_select).values.astype(np.float64)
             is_time_dim = _check_for_time(c, coord_vals)
+            if is_time_dim:
+                time_index_number = icoord
 
             if coord_vals.size > 1:
                 # not positive-monotonic? reverse it for cell width calculations
@@ -193,6 +197,8 @@ def _process_selection(self, xr_ds):
         self.ndims = len(n_edges)
         self.selected_shape = tuple(n_edges)
         self.select_shape_cells = tuple(n_cells)
+        if time_index_number is not None:
+            _ = full_dimranges.pop(time_index_number)
         self.full_bbox = np.array(full_dimranges).astype(np.float64)
         self.selected_bbox = np.array(dimranges).astype(np.float64)
         self.full_coords = tuple(full_coords)
@@ -201,9 +207,15 @@ def _process_selection(self, xr_ds):
         self.selected_time = time
         self.grid_type = grid_type
         self.cell_widths = cell_widths
-        self.reverse_axis = reverse_axis
         self.reverse_axis_names = reverse_axis_names
         self.global_dims = np.array(global_dims)
+        if time_index_number is not None:
+            _ = global_dims.pop(time_index_number)
+            _ = reverse_axis.pop(time_index_number)
+        self.reverse_axis = reverse_axis
+        self.time_index_number = time_index_number
+        self.global_dims_no_time = np.array(global_dims)
+
         # self.coord_selected_arrays = coord_selected_arrays
 
         # set the yt grid dictionary

yt_xarray/tests/test_xr_to_yt.py

@@ -533,3 +533,25 @@ def test_reversed_axis(stretched, use_callable, chunksizes):
     pdy_lats = slc._generate_container_field("pdy")
     assert np.all(pdy_lats > 0)
     assert np.all(np.isfinite(vals))
+
+
+def test_reader_with_2d_space_time_and_reverse_axis():
+
+    # test for https://github.com/data-exp-lab/yt_xarray/issues/86
+
+    # a base xarray ds to be used in various places.
+    ds = construct_ds_with_extra_dim(
+        3,
+        ncoords=5,
+        nd_space=2,
+        reverse_indices=[
+            1,
+        ],
+    )
+
+    field = ("stream", "test_case_3")
+    ds_yt = ds.yt.load_grid(
+        "test_case_3", sel_dict={"time": 0}, geometry="geographic", use_callable=True
+    )
+    slc = yt.SlicePlot(ds_yt, "altitude", field)
+    assert np.all(np.isfinite(slc.frb[field]))

yt_xarray/utilities/_utilities.py

@@ -1,5 +1,5 @@
 import os.path
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple
 
 import numpy as np
 import xarray as xr
@@ -111,7 +111,7 @@ def _test_time_coord(nt=5):
 
 def _get_test_coord(
     cname, n, minv: Optional[float] = None, maxv: Optional[float] = None
-):
+) -> np.ndarray:
     if cname in known_coord_aliases:
         cname = known_coord_aliases[cname]
 
@@ -140,7 +140,13 @@ def _get_test_coord(
     return np.linspace(minv, maxv, n)
 
 
-def construct_ds_with_extra_dim(icoord: int, dim_name: str = "time"):
+def construct_ds_with_extra_dim(
+    icoord: int,
+    dim_name: str = "time",
+    ncoords: Optional[int] = None,
+    nd_space: int = 3,
+    reverse_indices: Optional[List[int]] = None,
+):
     coord_configs = {
         0: (dim_name, "x", "y", "z"),
         1: (dim_name, "z", "y", "x"),
@@ -150,11 +156,22 @@ def construct_ds_with_extra_dim(icoord: int, dim_name: str = "time"):
     }
 
     data_vars = {}
-    coords = {c: _get_test_coord(c, icoord + 4) for c in coord_configs[icoord]}
+    if ncoords is None:
+        ncoords = icoord + 4
+
+    full_dims = coord_configs[icoord]
+    dim_order = [full_dims[idim] for idim in range(nd_space + 1)]
+    coords = {c: _get_test_coord(c, ncoords) for c in dim_order}
+
+    if reverse_indices is not None:
+        for indx in reverse_indices:
+            dim = dim_order[indx]
+            coords[dim] = coords[dim][::-1]
+
     var_shape = tuple([len(c) for c in coords.values()])
     vals = np.random.random(var_shape)
     fname = f"test_case_{icoord}"
-    da = xr.DataArray(vals, coords=coords, dims=coord_configs[icoord])
+    da = xr.DataArray(vals, coords=coords, dims=dim_order)
     data_vars[fname] = da
 
     return xr.Dataset(data_vars=data_vars)