refactor the index reversing

Author: chrishavlin

tests/test_xr_to_yt.py

@@ -481,19 +481,23 @@ def test_add_3rd_axis_name(yt_geom):
         _ = xr2yt._add_3rd_axis_name("bad_geometry", expected[:-1])
 
 
-@pytest.mark.parametrize(
-    "stretched,use_callable,chunksizes",
-    [
-        (True, False, None),
-        (False, False, None),
-        (False, True, None),
-        (False, True, 10),
-        (False, False, 10),
-    ],
-)
-def test_reversed_axis(stretched, use_callable, chunksizes):
-    # tests for when the incoming data is not positive-monotonic
+def _get_pixelized_slice(yt_ds):
+    slc = yt_ds.slice(
+        yt_ds.coordinates.axis_id["depth"],
+        yt_ds.domain_center[yt_ds.coordinates.axis_id["depth"]],
+        center=yt_ds.domain_center,
+    )
+    vals = yt_ds.coordinates.pixelize(
+        0,
+        slc,
+        ("stream", "test_field"),
+        yt_ds.arr([1, 359, -89, 89], "code_length"),
+        (400, 400),
+    )
+    return slc, vals
 
+
+def _get_ds_for_reverse_tests(stretched, use_callable, chunksizes):
     ds = construct_minimal_ds(
         min_x=1,
         max_x=359,
@@ -505,27 +509,34 @@ def test_reversed_axis(stretched, use_callable, chunksizes):
         n_y=100,
         n_z=30,
         z_stretched=stretched,
+        npseed=True,
     )
     yt_ds = ds.yt.load_grid(use_callable=use_callable, chunksizes=chunksizes)
+    return yt_ds
+
+
+@pytest.mark.parametrize(
+    "stretched,use_callable,chunksizes",
+    [
+        (True, False, None),
+        (False, False, None),
+        (False, True, None),
+        (False, True, 20),
+        (False, False, 20),
+    ],
+)
+def test_reversed_axis(stretched, use_callable, chunksizes):
+    # tests for when the incoming data is not positive-monotonic
+
+    yt_ds = _get_ds_for_reverse_tests(stretched, use_callable, chunksizes)
 
     if stretched:
         grid_obj = yt_ds.index.grids[0]
         ax_id = yt_ds.coordinates.axis_id["latitude"]
         assert np.all(grid_obj.cell_widths[ax_id] > 0)
 
-    slc = yt_ds.slice(
-        yt_ds.coordinates.axis_id["depth"],
-        yt_ds.domain_center[yt_ds.coordinates.axis_id["depth"]],
-        center=yt_ds.domain_center,
-    )
+    slc, vals = _get_pixelized_slice(yt_ds)
+
     pdy_lats = slc._generate_container_field("pdy")
     assert np.all(pdy_lats > 0)
-
-    vals = yt_ds.coordinates.pixelize(
-        0,
-        slc,
-        ("stream", "test_field"),
-        yt_ds.arr([1, 359, -89, 89], "code_length"),
-        (400, 400),
-    )
     assert np.all(np.isfinite(vals))

yt_xarray/accessor/_readers.py

@@ -18,42 +18,48 @@ def _reader(grid, field_name):
         # grid: a yt grid object
         _, fname = field_name
 
-        # this global start index accounts for indexing after any
+        # first get the internal yt index ranges
+        si = grid.get_global_startindex()
+        ei = si + grid.ActiveDimensions
+
+        # convert to the xarray indexing in 3 steps
+        # 1. account for dimension direction
+        # 2. account for any prior xarray subselections
+        # 3. account for interpolation requirements
+
+        # step 1 (if axis has been reversed, node ordering is also reversed)
+        for idim in range(sel_info.ndims):
+            if sel_info.reverse_axis[idim]:
+                # 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]
+
+        # step 2: this global start index accounts for indexing after any
         # subselections on the xarray DataArray are made. might
         # be a way to properly set this with yt grid objects.
         gsi = sel_info.starting_indices
         if sel_info.ndims == 2:
             gsi = np.append(gsi, 0)
+        si = si + gsi
+        ei = ei + gsi
 
-        si = grid.get_global_startindex() + gsi
-        ei = si + grid.ActiveDimensions
-        global_dims = sel_info.global_dims.copy()
+        # step 3: if we are interpolating to cell centers, grab some extra nodes
         if interp_required:
-            # if interpolating, si and ei must be node indices so
-            # we offset by an additional element
-            cell_to_node_offset = np.ones((3,), dtype=int)
-            ei = ei + cell_to_node_offset
+            for idim in range(sel_info.ndims):
+                if sel_info.reverse_axis[idim]:
+                    si[idim] = si[idim] - 1
+                else:
+                    ei[idim] = ei[idim] + 1
+
+        # now we can select the data (still accounting for any dimension reversal)
 
-        # build the index-selector for our yt grid object
+        # build the index-selector for each dimension
         c_list = sel_info.selected_coords  # the xarray coord names
         i_select_dict = {}
         for idim in range(sel_info.ndims):
-            if sel_info.reverse_axis[idim]:
-                # the xarray axis is in negative ordering. a yt index of
-                # 0 should point to the maximum index in xarray
-                si_idim = global_dims[idim] - si[idim]
-                ei_idim = global_dims[idim] - ei[idim]
-                # when reverse slicing, the final index will not be included...
-                # if the end index is within bounds, just bump it by one more,
-                # otherwise if end index is already 0, just pass in None to
-                # slice
-                if ei_idim > 0:
-                    ei_idim = ei_idim - 1
-                elif ei_idim == 0:
-                    ei_idim = None
-                i_select_dict[c_list[idim]] = slice(si_idim, ei_idim, -1)
-            else:
-                i_select_dict[c_list[idim]] = slice(si[idim], ei[idim])
+            i_select_dict[c_list[idim]] = slice(si[idim], ei[idim])
 
         # set any of the initial selections that will reduce the
         # dimensionality or size of the full DataArray
@@ -79,6 +85,11 @@ def _reader(grid, field_name):
             # interpolate from nodes to cell centers across all remaining dims
             datavals = _xr_to_yt._interpolate_to_cell_centers(datavals)
 
+        # final flips to account for all the index reversing
+        for idim in range(sel_info.ndims):
+            if sel_info.reverse_axis[idim]:
+                datavals = np.flip(datavals, axis=idim)
+
         # return the plain values
         vals = datavals.values.astype(np.float64)
         if sel_info.ndims == 2:

yt_xarray/accessor/_xr_to_yt.py

@@ -132,36 +132,43 @@ def _process_selection(self, xr_ds):
         cell_widths = []  # cell widths after selection
         grid_type = _GridType.UNIFORM  # start with uniform assumption
         reverse_axis = []  # axes must be positive-monitonic for yt
-        global_dims = []
+        global_dims = []  # the global shape
         for c in full_coords:
             coord_da = getattr(xr_ds, c)  # the full coordinate data array
+
+            # check if coordinate values are increasing
             rev_ax = coord_da[1] <= coord_da[0]
             reverse_axis.append(bool(rev_ax.values))
-            global_dims.append(coord_da.size)
+
             # store the global ranges
+            global_dims.append(coord_da.size)
             global_min = float(coord_da.min().values)
             global_max = float(coord_da.max().values)
             full_dimranges.append([global_min, global_max])
 
-            si = 0  # starting index
+            si = 0  # starting xarray-index for any pre-selections from user
             coord_select = {}  # the selection dictionary just for this coordinate
             if c in self.sel_dict:
                 coord_select[c] = self.sel_dict[c]
                 si = self._find_starting_index(c, coord_da, coord_select)
 
+            # apply any selections and extract coordinates
             sel_or_isel = getattr(coord_da, self.sel_dict_type)
             coord_vals = sel_or_isel(coord_select).values.astype(np.float64)
-            if reverse_axis[-1]:
-                coord_vals = coord_vals[::-1]
             is_time_dim = _check_for_time(c, coord_vals)
 
             if coord_vals.size > 1:
+
+                # not positive-monotonic? reverse it for cell width calculations
+                # changes to indexing are accounted for when extracting data.
+                if reverse_axis[-1]:
+                    coord_vals = coord_vals[::-1]
+
                 cell_widths.append(coord_vals[1:] - coord_vals[:-1])
                 dimranges.append([coord_vals.min(), coord_vals.max()])
                 n_edges.append(coord_vals.size)
                 n_cells.append(coord_vals.size - 1)
                 coord_list.append(c)
-                # coord_selected_arrays[c] = coord_vals
                 starting_indices.append(si)
 
                 if is_time_dim:
@@ -422,3 +429,23 @@ def _interpolate_to_cell_centers(data: xr.DataArray):
         dimvals = data.coords[dim].values
         interp_dict[dim] = (dimvals[1:] + dimvals[:-1]) / 2.0
     return data.interp(interp_dict)
+
+
+def _load_full_field_from_xr(
+    ds_xr, field: str, sel_info: Selection, interp_required: bool = False
+):
+    vals = sel_info.select_from_xr(ds_xr, field).load()
+
+    if interp_required:
+        vals = _interpolate_to_cell_centers(vals)
+    if any(sel_info.reverse_axis):
+        # if any dims are in decreaseing order, flip that axis
+        # after reading in the data
+        for idim, flip_it in enumerate(sel_info.reverse_axis):
+            if flip_it:
+                vals = np.flip(vals, axis=idim)
+
+    vals = vals.values.astype(np.float64)
+    if sel_info.ndims == 2:
+        vals = np.expand_dims(vals, axis=-1)
+    return vals

yt_xarray/accessor/accessor.py

@@ -8,6 +8,7 @@
 
 from yt_xarray.accessor import _xr_to_yt
 from yt_xarray.accessor._readers import _get_xarray_reader
+from yt_xarray.accessor._xr_to_yt import _load_full_field_from_xr
 from yt_xarray.utilities.logging import ytxr_log
 
 
@@ -243,15 +244,9 @@ def _load_single_grid(
         if use_callable:
             data[field] = (reader, units)
         else:
-            vals = sel_info.select_from_xr(ds_xr, field).load()
-            if interp_required:
-                vals = _xr_to_yt._interpolate_to_cell_centers(vals)
-            vals = vals.values.astype(np.float64)
-            for ax, rev_ax in enumerate(sel_info.reverse_axis):
-                if rev_ax:
-                    vals = np.flip(vals, axis=ax)
-            if sel_info.ndims == 2:
-                vals = np.expand_dims(vals, axis=-1)
+            vals = _load_full_field_from_xr(
+                ds_xr, field, sel_info, interp_required=interp_required
+            )
             data[field] = (vals, units)
 
     if sel_info.ndims == 2:
@@ -368,23 +363,19 @@ def _load_chunked_grid(
         c = cnames[idim]
         rev_ax = sel_info.reverse_axis[idim]
         if rev_ax is False:
-
             le_0 = ds_xr[fld].coords[c].isel({c: si_0}).values
-
             if interp_required is False:
-                # the left edges get bumped left since we are reading values
-                # again.
+                # move the edges so the node is now a cell center
                 le_0 = le_0 - dxyz[idim] / 2.0
 
-            # bbox value below already accounts for interp_required, no need to shift
+            # bbox value below already accounts for interp_required
             max_val = bbox[idim, 1]
             re_0 = np.concatenate([le_0[1:], [max_val]])
 
         else:
             re_0 = ds_xr[fld].coords[c].isel({c: si_0[::-1]}).values
             if interp_required is False:
-                # the left edges get bumped left since we are reading values
-                # again.
+                # move the edges so the node is now a cell center
                 re_0 = re_0 - dxyz[idim] / 2.0
             min_val = bbox[idim, 0]
             le_0 = np.concatenate([[min_val], re_0[:-1]])
@@ -424,14 +415,10 @@ def _load_chunked_grid(
     if use_callable is False:
         full_field_vals = {}
         for field in fields:
-            vals = sel_info.select_from_xr(ds_xr, field).load()
-            if interp_required:
-                vals = _xr_to_yt._interpolate_to_cell_centers(vals)
-            if any(sel_info.reverse_axis):
-                for idim, flip_it in enumerate(sel_info.reverse_axis):
-                    if flip_it:
-                        vals = np.flip(vals, axis=idim)
-            full_field_vals[field] = vals.values.astype(np.float64)
+            vals = _load_full_field_from_xr(
+                ds_xr, field, sel_info, interp_required=interp_required
+            )
+            full_field_vals[field] = vals
 
     for igrid in range(n_grids):
         gdict = {

yt_xarray/utilities/_utilities.py

@@ -29,6 +29,7 @@ def construct_minimal_ds(
     n_fields: int = 1,
     coord_order: Optional[Tuple[str, str, str]] = None,
     dtype: str = "float64",
+    npseed: bool = False,
 ) -> xr.Dataset:
 
     if coord_order is None:
@@ -78,6 +79,9 @@ def construct_minimal_ds(
         coord_order_rn += (cdict[cname],)
         var_shape += (n,)
 
+    if npseed:
+        np.random.seed(0)
+
     vals = np.random.random(var_shape).astype(dtype_to_use)
     if n_fields > 1:
         data_vars = {}