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.