Merge pull request #39 from chrishavlin/handling_revsered_dims

chrishavlin · web-flow · commit 06f9722d0d17 · 2023-02-27T17:30:45.000-06:00
handling revesred dimension arrays
diff --git a/.github/workflows/build-test.yml b/.github/workflows/build-test.yml
@@ -11,9 +11,9 @@ jobs:
         python-version: ['3.8', '3.9', '3.10']
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
diff --git a/.github/workflows/check-build.yml b/.github/workflows/check-build.yml
@@ -6,9 +6,9 @@ jobs:
     name: twinecheck
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: 3.9
     - name: Install dependencies
diff --git a/.github/workflows/check-manifest.yml b/.github/workflows/check-manifest.yml
@@ -7,9 +7,9 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - name: Checkout Source
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
     - name: Setup Python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: 3.9
     - name: install check-manifest
diff --git a/.github/workflows/create-release.yml b/.github/workflows/create-release.yml
@@ -13,7 +13,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
       - name: Create Release
         id: create_release
         uses: actions/create-release@v1
diff --git a/.github/workflows/publish-to-pypi.yml b/.github/workflows/publish-to-pypi.yml
@@ -11,9 +11,9 @@ jobs:
     name: Build and Publish
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@master
+    - uses: actions/checkout@v3
     - name: Set up Python 3.9
-      uses: actions/setup-python@v1
+      uses: actions/setup-python@v4
       with:
         python-version: 3.9
     - name: Install dependencies
diff --git a/.github/workflows/style-checks.yml b/.github/workflows/style-checks.yml
@@ -6,11 +6,11 @@ jobs:
     name: flake8
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@master
+    - uses: actions/checkout@v3
     - name: Setup Python
-      uses: actions/setup-python@master
+      uses: actions/setup-python@v4
       with:
-        python-version: 3.8
+        python-version: 3.9
     - name: install
       id: install
       run : pip install -r tests/lint_requirements.txt
@@ -25,11 +25,11 @@ jobs:
     name: black
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@master
+    - uses: actions/checkout@v3
     - name: Setup Python
-      uses: actions/setup-python@master
+      uses: actions/setup-python@v4
       with:
-        python-version: 3.8
+        python-version: 3.9
     - name: install
       id: install
       run : pip install -r tests/lint_requirements.txt
@@ -44,11 +44,11 @@ jobs:
     name: isort
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@master
+    - uses: actions/checkout@v3
     - name: Setup Python
-      uses: actions/setup-python@master
+      uses: actions/setup-python@v4
       with:
-        python-version: 3.8
+        python-version: 3.9
     - name: install
       id: install
       run : pip install -r tests/lint_requirements.txt
diff --git a/HISTORY.md b/HISTORY.md
@@ -1,5 +1,14 @@
 # History
 
+## 0.1.2 (2022-02-27)
+
+Bug fix release.
+
+### Fixes:
+* fixes handling of xarray variables with dimensions in decreasing
+  order (e.g., latitude from 90 to -90) without
+having to re-index the whole variable (PR [39](https://github.com/data-exp-lab/yt_xarray/pull/39)).
+
 ## 0.1.1 (2022-02-07)
 
 This release builds out the loading functions introduced in v0.1.0 and includes
diff --git a/setup.cfg b/setup.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.1.1
+current_version = 0.1.2
 commit = True
 tag = True
 
diff --git a/setup.py b/setup.py
@@ -37,6 +37,6 @@
     test_suite="tests",
     tests_require=test_requirements,
     url="https://github.com/data-exp-lab/yt_xarray",
-    version="0.1.1",
+    version="0.1.2",
     zip_safe=False,
 )
diff --git a/tests/test_xr_to_yt.py b/tests/test_xr_to_yt.py
@@ -479,3 +479,64 @@ 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])
+
+
+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,
+        min_z=50,
+        max_z=650,
+        min_y=89,
+        max_y=-89,
+        n_x=50,
+        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, vals = _get_pixelized_slice(yt_ds)
+
+    pdy_lats = slc._generate_container_field("pdy")
+    assert np.all(pdy_lats > 0)
+    assert np.all(np.isfinite(vals))
diff --git a/yt_xarray/__init__.py b/yt_xarray/__init__.py
@@ -2,7 +2,7 @@
 
 __author__ = """Chris Havlin"""
 __email__ = "chris.havlin@gmail.com"
-__version__ = "0.1.1"
+__version__ = "0.1.2"
 
 
 # import the yt frontend and the xarray accessor so they are registered with
diff --git a/yt_xarray/accessor/_readers.py b/yt_xarray/accessor/_readers.py
@@ -18,22 +18,44 @@ 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
+        # 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):
@@ -63,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:
diff --git a/yt_xarray/accessor/_xr_to_yt.py b/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,31 +131,44 @@ 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 = []  # 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))
+
             # 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)
             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:
@@ -188,6 +202,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
@@ -413,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
diff --git a/yt_xarray/accessor/accessor.py b/yt_xarray/accessor/accessor.py
diff --git a/yt_xarray/utilities/_utilities.py b/yt_xarray/utilities/_utilities.py

Original file line number	Diff line number	Diff line change
`@@ -37,6 +37,6 @@`
`37`	`37`	`test_suite="tests",`
`38`	`38`	`tests_require=test_requirements,`
`39`	`39`	`url="https://github.com/data-exp-lab/yt_xarray",`
`40`		`- version="0.1.1",`
	`40`	`+ version="0.1.2",`
`41`	`41`	`zip_safe=False,`
`42`	`42`	`)`