Fix isel() bugs: missing indexers' coords; UxDataset.isel() silently skips slicing / fails if grid dim coords exist

[Sevans711]

Closes #1712;
Closes #1713 / Closes #1714

Overview

Mapping code changes to original issues:

• The fix for isel() with xr.DataArray indexer along grid dimension fails to include indexer's coordinates in result #1712 was to define _assign_grid_dim_indexer_coords_if_appropriate in coords.py, and utilize it during UxDataArray.isel() and UxDataset.isel(). Also updated isel() and sel() docstrings accordingly.
• The basic fix for UxDataset.isel() silently fails to slice when providing grid dim not in the dataset #1713 was incredibly simple, just needed to check if any grid dim is in da.dims when looping through data arrays in the dataset, in UxDataset._slice_dataset_from_grid. While touching this method though, I decided to clean it up further than this (see relevant "expansion of scope" details below).
• The fix for In UxDataset.isel() if "n_face" coords assigned: ValueError: dimension 'n_face' already exists as a scalar variable #1714 was simply to delete the code commented as "Prefer authoritative coords from the sliced grid if available", in UxDataset._slice_dataset_from_grid. I fully deleted those lines instead of "trying to fix them" because I don't think they made sense to have in the first place. See point (2) below for more details.

In all cases, also updated test_indexing.py accordingly. For 1713 and 1714, it was sufficient to just uncomment some existing lines which now serve as regression tests for these issues. For 1712, it also included building new tests inside test_indexing_by_dataarray(), which led to discovering #1758.

Expansion of scope: updated UxDataset._slice_dataset_from_grid in a few different ways to increase consistency with UxDataArray indexing methods:

1. Renamed to UxDataset._slice_from_grid (to match naming conventions of UxDataArray._slice_from_grid)
2. Deleted the "Prefer authoritative coords from the sliced grid if available" code. I don't think it would ever be desirable to overwrite dataset coords for a grid dimension based on the underlying grid's coords, during isel(). The PR which originally added those lines also contains no description of their purpose (see Add UxDataset.isel()  #1352), and they were not being covered by any example within the test suite (removing them does not cause any existing tests to fail).
3. Deleted the logical branch for handling data arrays in self.data_vars when grid_dim in da.dims and not hasattr(da, "_slice_from_grid"). This branch would never be reached. (Right now, all data_vars are UxDataArrays. Eventually, maybe only data_vars with a grid dimension will be UxDataArrays. Either way, there should never be a data_var which both has a grid dim and is an xr.DataArray.)
4. No longer accepts grid_dim or grid_indexer args; these were only being used in the now-removed parts of the code from points (2) and (3) above.
5. Now returns a UxDataset instead of an xr.Dataset. (UxDataArray._slice_from_grid returns a UxDataArray, not an xr.DataArray.)

Tiny expansions of scope:

6. Updated UxDataArray._slice_from_grid() to return an object of type(self), instead of UxDataArray. This has no extra cost, and makes uxarray slightly more object-oriented friendly. Before this change, on subclasses like class MyCoolArray(UxDataArray), isel() would return a UxDataArray instance if indexing by exactly 1 grid dim (e.g. isel(n_face=7)), but an instance of the subclass (MyCoolArray) if indexing by 0 grid dims (e.g. isel(time=0)) or if additionally indexing by something other than a grid dim (e.g. isel(time=0, n_face=7)). After this change it will always return an instance of the subclass (MyCoolArray) regardless of which dimensions get indexed.
7. Minor syntax change in the UxDataArray.isel() code (i.e. calling the result result instead of da or xarr, and updating the value of result inside the if indexers:... block instead of returning from directly inside that block) to emphasize that the implementation is nearly identical to UxDataset.isel(). Maybe some day the full implementation, or at least a majority of it, could be moved to a shared parent class (something like UxDataContainer?), as a way to help avoid any more unexpected behavior differences between UxDataArray.isel() and UxDataset.isel().

Tiny expansions of scope unrelated to original issue:

8. Small code cleanup with no functionality changes: noticed and removed import uxarray from top of dataset.py and dataarray.py files. Pretty sure it doesn't belong there; importing a full package from within itself is not standard practice. Removing this made it necessary to make a few other tiny updates:

• (8a) updated UxDataset.get_dual() to use type(self) instead of uxarray.UxDataset, and to use UxDataArray instead of uxarray.UxDataArray.
• (8b) Similarly, UxDataArray.get_dual() now uses type(self) instead of uxarray.UxDataArray.
• (8c) Also updated UxDataArray.to_dataset() to explicitly include from uxarray.core.dataset import UxDataset instead of using uxarray.core.dataset.UxDataset.

PR Checklist

General

• An issue is created and linked
• Added appropriate labels (if your uxarray repo permissions allow it)
• Filled out Overview and Expected Usage (if applicable) sections

Testing & Benchmarking

• There is adequate test coverage of changes from this PR (add new tests if needed)
• [N/A] If this PR could affect performance, ran ASV benchmarks and confirmed they show expected behavior (add a new benchmark if necessary)

Documentation and Examples

• Docstrings updated with any function changes, and included in all new functions
• User (public) functions added to docs/api.rst; internal (private) function names start with an underscore (_)
• [N/A] If touched any notebook files, cleared the output of all cells before committing
• [N/A] If added new notebook files, put into appropriate directories and referenced in appropriate files

AI Disclosure

AI Usage: GitHub Copilot's inline code suggestions, plus a small discussion with Claude for debugging/understanding syntax of super().

• I have tested and take responsibility for all AI-generated content in my PR.

[dylannelson]

@Sevans711
I noticed 3 things that changed before/after that may be of concern:

Notebooks for current main (before) and this branch (after):
testing1759-BEFORE.html
testing1759-AFTER.html

1) A boolean mask with coords crashes .isel and .where(drop=True)

Trying this crashes when before it worked and had a clean output:

gds = ux.open_dataset(path + r"\geos-cs\c12\test-c12.native.nc4", path + r"\geos-cs\c12\test-c12.native.nc4")
phis = gds["PHIS"].isel(time=0)
phis.where(phis > phis.mean(), drop=True)

AlignmentError: cannot reindex or align along dimension 'n_face' because of conflicting dimension sizes: {864, 217}

vs before it returned a data array of size 217:

2) Indexer coords replace the data's own coords

Trying to do

ds = ux.open_dataset(path + r"\ugrid\outCSne30\outCSne30.ug", path + r"\ugrid\outCSne30\outCSne30_sel_timeseries.nc")
psi = ds["psi"]
series = psi.isel(n_face=psi.isel(time=0).argmax("n_face"))
series.time

Results in an empty data array, whereas before it kept it's time

Before:

After:

3) .isel dropping coords

Some coordinates dropped, like Xdim and Ydim in this case:

gds = ux.open_dataset(path + r"\geos-cs\c12\test-c12.native.nc4", path + r"\geos-cs\c12\test-c12.native.nc4")
print(sorted(gds.coords))
print(sorted(gds.isel(n_face=[0, 1]).coords))

Before

After

[Sevans711]

@dylannelson Thank you for reviewing and for pointing out these bugs! I believe I was able to fix them all. Notes below:

1. Corrected logic for handling boolean indexer arrays inside of _assign_grid_dim_indexer_coords_if_appropriate; these need to be handled specially because the output's size does not match the indexer array's size. Added corresponding regression test test_isel_can_use_bool_with_coords.
2. To me this looks like a bug on main, too, because it is inconsistent with xarray behavior. E.g., psi.to_xarray().isel(n_face=psi.isel(time=0).argmax("n_face")) crashes with IndexError: dimension coordinate 'time' conflicts between indexed and indexing objects:.... Fixed here by adding a check for consistency between coordinates, using the relevant method directly from xarray (xr.core.coordinates.assert_coordinate_consistent(uxarray_obj, indexer.coords)) (the check was added into _assign_grid_dim_indexer_coords_if_appropriate). So, now, this case will (correctly) cause a crash, instead of incorrectly choosing to keep coordinates from the indexer or from the original array. Added corresponding regression test test_isel_crash_if_coordinates_conflict.
3. Ah, this was happening due to my rewrite of UxDataset._slice_from_grid, which failed to account for coordinates which might not be attached to any of the dataset's data_vars. Updated that function to fix this, and added regression test test_dataset_isel_keeps_extra_coords.

[dylannelson]

1 and 3 looks fixed and 2 may still be a bit of a concern. The error seems like it can still come up even when the time dimensions match.

I'll leave it up to you though if you want to leave this be because I'm not certain if this is some unrealistic user behavior, or if this could maybe be a separate issue. Or maybe the error could be updated to be more broad, like suggesting users to remove time if it's not needed (?)

[Sevans711]

1 and 3 looks fixed and 2 may still be a bit of a concern. The error seems like it can still come up even when the time dimensions match.

I'll leave it up to you though if you want to leave this be because I'm not certain if this is some unrealistic user behavior, or if this could maybe be a separate issue. Or maybe the error could be updated to be more broad, like suggesting users to remove time if it's not needed (?)

This doesn't seem unrealistic to me; uxarray should either provide correct results or a much less confusing error message. Looking into it further, I decided to go with a much less confusing error message. Leaving it as a NotImplementedError for now, because maybe it could be implemented in the future. The _crash_if_1d_xarray_indexer_dim_in_uxarray_obj function now handles this case, and its docstring explains the reasoning a bit more. Added test_isel_when_indexer_dim_in_uxarray_obj as a regression test for this case.

Also fixed: now raise DimensionError when uxarray object has a scalar coordinate which indexer has as a 1D coordinate, instead of silently overwriting the uxarray object's coordinate. (This matches xarray's behavior.)

[cmdupuis3]

Four crashes:

1

import uxarray as ux
import numpy as np

uxds = ux.tutorial.open_dataset("quad-hexagon")
uxds = uxds.assign_coords(node_id=("n_node", np.arange(uxds.uxgrid.n_node)))
sub = uxds.isel(n_face=[0, 1])
print(list(sub.coords))   # []   (main: ['node_id'])
sub.node_id               # AttributeError

2

import xarray as xr

mask = xr.DataArray([True, False, True, False], dims="n_face",
                    coords={"lab": ("n_face", [10, 20, 30, 40])})
uxds.isel(n_face=mask)    # ValueError: conflicting sizes

3

hits = np.where(uxds.uxgrid.edge_lon > 1e9)[0]   # empty
uxds.isel(n_edge=xr.DataArray(hits, dims="selected"))   # IndexError

4

This one seems to be wrong on main as well.

uxds = ux.tutorial.open_dataset("quad-hexagon").assign_coords(n_face=[0, 10, 20, 30])
picks = xr.DataArray([0, 20], dims="station", coords={"station": ["A", "B"]})

by_index = uxds.isel(n_face=xr.DataArray([0, 2], dims="station",
                                         coords={"station": ["A", "B"]}))
print("isel keeps indexer coords:", list(by_index.coords))   # ['n_face', 'station']

by_label = uxds.sel(n_face=picks)
print("sel  keeps indexer coords:", list(by_label.coords))   # ['n_face']  <-- 'station' gone

by_label.sel(station="A")                                    # ValueError: no 'station'

[Sevans711]

@cmdupuis3 Thank you for reviewing and finding these bugs!

It looks like your examples (1) and (2) were actually already solved by the commits I pushed just a few hours before your comment. Though, it doesn't hurt to add these to the test suite, so I added them into test_indexing_by_dataarray() and test_isel_can_use_bool_with_coords, respectively.

Example (3) was definitely still a bug; I failed to account for size 0 indexers in the _assign_grid_dim_indexer_coords_if_appropriate() function. This should now be fixed, and also I added your example along with more extensive regression tests for this in test_indexing_by_size_0_array().

Example (4) was also still a bug; looking into uxarray.core.utils._resolve_coordinate_labels_to_indices I see that there was even a note suggesting it should be updated when fixing #1712, but I forgot to do so (even though I am the one who left the note… 😅). I've now updated that method appropriately (also added a warning just in case the assumption that f"__{dim}_indices__" isn't in coords ever fails, to avoid silently giving wrong results in that case). I also added regression tests in test_indexing_of_uxarray_obj_with_grid_dim_coords(), which led me to discover and fix another subtle issue (if uxarray_obj and indexer both contain grid_dim coords but they are incompatible, should crash, but was not previously crashing).

I believe all the points are addressed now, please let me know if you find anything else!

[cmdupuis3]

@Sevans711 1 and 4 seem to be resolved. I'm still getting crashes on variants of 2 and 3 though.

import uxarray as ux
ds_lab = ux.tutorial.open_dataset("quad-hexagon").assign_coords(n_face=[0, 10, 20, 30])
ds_lab.t2m.isel(n_face=ds_lab.t2m > 297.6)
# IndexError: dimension coordinate 'n_face' conflicts between indexed and indexing objects

import numpy as np
import xarray as xr
import uxarray as ux

ds = ux.tutorial.open_dataset("quad-hexagon")
empty = xr.DataArray(np.array([], dtype=int), dims="selected",
                     coords={"selected": np.array([], dtype=int)})

print(ds.isel(n_edge=empty).sizes)   # no crash, but gains a stray 'selected' dim of size 0
ds.t2m.isel(n_node=empty)            # CoordinateValidationError

[Sevans711]

1 and 4 seem to be resolved. I'm still getting crashes on variants of 2 and 3 though.

I believe these are fixed now! The issue with 2 is because boolean indexing changes the shape of the indexer's coords to be compared with the result, which I wasn't accounting for. Now there is a if not indexing_1d_bool_along_n_face check to guard against calling xr.core.coordinates.assert_coordinate_consistent too early in this special case (and it is called later, once the indexer has been isel'd too, to have the same shape as the uxarray_obj). Regression test added inside test_isel_can_use_bool_with_coords().

The issue with 3 was some subtle bugs occurring with 0D indexers. I checked the implementation you suggested separately and it worked well; adopted it here, now using coords = indexer.drop_vars([c for c in indexer.coords if the_dim in indexer[c].dims]).coords instead of special-casing 0D indexers explicitly. Regression test added inside test_indexing_by_size_0_array(). Existing tests seem to pass after these changes, suggesting that swapping to this implementation doesn't re-introduce any old bugs, at least.

Please continue to let me know if you find anything else! I appreciate being able to fix all of these in one place instead of needing to find and fix them over many PRs.

[cmdupuis3]

Looks good now, no more crashes that I see here.