handling revesred dimension arrays

Author: chrishavlin

tests/test_xr_to_yt.py

@@ -479,3 +479,53 @@ def test_add_3rd_axis_name(yt_geom):
 
     with pytest.raises(ValueError, match="Unsupported geometry type"):
         _ = 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
+
+    ds = construct_minimal_ds(
+        min_x=1,
+        max_x=359,
+        min_z=50,
+        max_z=650,
+        min_y=89,
+        max_y=-89,
+        n_x=50,
+        n_y=100,
+        n_z=30,
+        z_stretched=stretched,
+    )
+    yt_ds = ds.yt.load_grid(use_callable=use_callable, chunksizes=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,
+    )
+    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

@@ -27,6 +27,7 @@ def _reader(grid, field_name):
 
         si = grid.get_global_startindex() + gsi
         ei = si + grid.ActiveDimensions
+        global_dims = sel_info.global_dims.copy()
         if interp_required:
             # if interpolating, si and ei must be node indices so
             # we offset by an additional element
@@ -37,7 +38,22 @@ def _reader(grid, field_name):
         c_list = sel_info.selected_coords  # the xarray coord names
         i_select_dict = {}
         for idim in range(sel_info.ndims):
-            i_select_dict[c_list[idim]] = slice(si[idim], ei[idim])
+            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])
 
         # set any of the initial selections that will reduce the
         # dimensionality or size of the full DataArray

yt_xarray/accessor/_xr_to_yt.py

@@ -54,6 +54,7 @@ def __init__(
         self.ndims: int = None
         self.grid_type = None  # one of _GridType members
         self.cell_widths: list = None
+        self.global_dims: list = None
         self._process_selection(xr_ds)
 
         self.yt_coord_names = _convert_to_yt_internal_coords(self.selected_coords)
@@ -130,9 +131,13 @@ def _process_selection(self, xr_ds):
         starting_indices = []  # global starting index
         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 = []
         for c in full_coords:
             coord_da = getattr(xr_ds, c)  # the full coordinate data array
-
+            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_min = float(coord_da.min().values)
             global_max = float(coord_da.max().values)
@@ -146,6 +151,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)
+            if reverse_axis[-1]:
+                coord_vals = coord_vals[::-1]
             is_time_dim = _check_for_time(c, coord_vals)
 
             if coord_vals.size > 1:
@@ -188,6 +195,8 @@ 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.global_dims = np.array(global_dims)
         # self.coord_selected_arrays = coord_selected_arrays
 
         # set the yt grid dictionary

yt_xarray/accessor/accessor.py

@@ -247,6 +247,9 @@ def _load_single_grid(
             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)
             data[field] = (vals, units)
@@ -363,15 +366,28 @@ def _load_chunked_grid(
             )
 
         c = cnames[idim]
-        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.
-            le_0 = le_0 - dxyz[idim] / 2.0
+        rev_ax = sel_info.reverse_axis[idim]
+        if rev_ax is False:
 
-        # bbox value below already accounts for interp_required, no need to shift
-        max_val = bbox[idim, 1]
-        re_0 = np.concatenate([le_0[1:], [max_val]])
+            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.
+                le_0 = le_0 - dxyz[idim] / 2.0
+
+            # bbox value below already accounts for interp_required, no need to shift
+            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.
+                re_0 = re_0 - dxyz[idim] / 2.0
+            min_val = bbox[idim, 0]
+            le_0 = np.concatenate([[min_val], re_0[:-1]])
 
         # sizes also already account for interp_required
         subgrid_size = ei_0 - si_0
@@ -411,6 +427,10 @@ def _load_chunked_grid(
             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)
 
     for igrid in range(n_grids):
@@ -425,10 +445,8 @@ def _load_chunked_grid(
             if use_callable:
                 gdict[field] = (reader, units)
             else:
-                # NO these values need to be chunked too.
                 si = subgrid_start[igrid]
                 ei = subgrid_end[igrid]
-                # this needs to be fixed for two2 fields...
                 gridvals = full_field_vals[field][
                     si[0] : ei[0], si[1] : ei[1], si[2] : ei[2]
                 ]