diff --git a/test/core/test_indexing.py b/test/core/test_indexing.py index f70725267..8c0faffd1 100644 --- a/test/core/test_indexing.py +++ b/test/core/test_indexing.py @@ -39,17 +39,17 @@ def test_sel_indexes_grid(): def test_sel_uses_grid_dim_labels(): """ensure obj.sel({grid_dim: ...}) actually utilizes coordinate labels on that grid dim, for UxDataArrays and UxDatasets. Regression test for #1641. - TODO: fix #1714 then uncomment the UxDataset tests below + Also contains regression test for #1714. """ # test corresponding to the workflow described in #1641, but for UxDataset uxds = ux.tutorial.open_dataset("outCSne30-vortex") - # (uncomment the next few lines after fixing #1714) - # uxds1 = uxds.assign_coords(n_face=np.arange(uxds.n_face.size)) - # uxds2 = uxds1.isel(n_face=range(0, 100, 5)) - # uxds3 = uxds2 + 7 - # # "check what the results look like on what were originally faces 20, 30, and 40" - # uxds4 = uxds3.sel(n_face=[20,30,40]) - # assert uxds4.sizes['n_face'] == uxds4.uxgrid.n_face == 3 + # (the next few lines, through the `assert`, also serve as a regression test for #1714) + uxds1 = uxds.assign_coords(n_face=np.arange(uxds.n_face.size)) + uxds2 = uxds1.isel(n_face=range(0, 100, 5)) + uxds3 = uxds2 + 7 + # "check what the results look like on what were originally faces 20, 30, and 40" + uxds4 = uxds3.sel(n_face=[20,30,40]) + assert uxds4.sizes['n_face'] == uxds4.uxgrid.n_face == 3 # test corresponding to the workflow described in #1641, for UxDataArray uxarr = uxds['psi'] @@ -70,15 +70,14 @@ def test_sel_uses_grid_dim_labels(): def test_can_index_grid_dim_not_in_data(): """ensure isel() and sel() can both index a grid dim even if that dim is not present in the data itself; - for UxDataArrays and UxDatasets. TODO: fix #1713 then uncomment the UxDataset tests below. + for UxDataArrays and UxDatasets. The UxDataset checks serve as a regression test for #1713. """ ds = ux.tutorial.open_dataset("outCSne30-vortex") - # (uncomment the next few lines after fixing #1713) - # assert "n_face" in ds.dims - # result = ds.isel(n_edge=7) - # assert result.sizes["n_face"] == result.uxgrid.n_face == 2 - # result = ds.sel(n_edge=7) - # assert result.sizes["n_face"] == result.uxgrid.n_face == 2 + assert "n_face" in ds.dims + result = ds.isel(n_edge=7) + assert result.sizes["n_face"] == result.uxgrid.n_face == 2 + result = ds.sel(n_edge=7) + assert result.sizes["n_face"] == result.uxgrid.n_face == 2 arr = ds["psi"] result = arr.isel(n_edge=7) @@ -110,7 +109,7 @@ def test_sel_can_use_slice(): """ensure sel() can use slice() objects as indexers, and provides expected results, with expected sizes, for UxDataArrays and UxDatasets. Regression test inspired by reviewer comment in #1641, also related to #1639. - TODO: fix #1714 then uncomment the relevant UxDataset tests below + Also contains regression test for #1714. """ grid = ux.Grid.from_healpix(zoom=0) # 12 faces arr = ux.UxDataArray( @@ -129,10 +128,10 @@ def test_sel_can_use_slice(): uxds = ux.UxDataset({'data': arr.to_xarray()}, uxgrid=grid) result = uxds.sel(n_face=slice(0, 2)) assert result.n_face.size == result.uxgrid.n_face == 2 - # (uncomment the next few lines after fixing #1714) - # labeled_ds = uxds.assign_coords(n_face=np.arange(grid.n_face)) - # result = labeled_ds.sel(n_face=slice(0, 2)) - # assert result.n_face.size == result.uxgrid.n_face == 3 + # (the remaining lines also serve as a regression test for #1714) + labeled_ds = uxds.assign_coords(n_face=np.arange(grid.n_face)) + result = labeled_ds.sel(n_face=slice(0, 2)) + assert result.n_face.size == result.uxgrid.n_face == 3 def test_isel_can_use_bool(): """ensure isel() supports indexing by a boolean indexer array. @@ -151,54 +150,322 @@ def test_isel_can_use_bool(): result = arr.isel(n_face=[False, False, False, False]) assert result.sizes['n_face'] == result.uxgrid.n_face == 0 + +def test_isel_can_use_bool_with_coords(): + """ensure isel() supports indexing by a boolean indexer array with coords. + Regression test for bug (1) discovered during review of PR #1759. + """ + ds = ux.tutorial.open_dataset("quad-hexagon") + arr = ds['t2m'] + + # simplest case: boolean xr.DataArray mask with coords. + tokeep0 = xr.DataArray([True, False, True, False], coords={'tokeep': [2,4,6,8]}) + # (also want to test UxDataArray boolean mask (ensure no infinite recursion)) + tokeep1 = arr * xr.DataArray([True, False, True, True], dims={'n_face'}) > 0 + assert isinstance(tokeep1, ux.UxDataArray) + # (also want to test UxDataArray boolean mask with coords + tokeep2 = tokeep1.assign_coords(tokeep=('n_face', [1,3,5,7])) + assert isinstance(tokeep2, ux.UxDataArray) + assert np.all(tokeep2.coords['tokeep'] == [1,3,5,7]) + + # test on UxDataArray + result0 = arr.isel(n_face=tokeep0) + assert result0.sizes['n_face'] == result0.uxgrid.n_face == 2 + assert np.all(result0.coords['tokeep'] == [2,6]) + result1 = arr.isel(n_face=tokeep1) + assert result1.sizes['n_face'] == result1.uxgrid.n_face == 3 + result2 = arr.isel(n_face=tokeep2) + assert result2.sizes['n_face'] == result2.uxgrid.n_face == 3 + assert np.all(result2.coords['tokeep'] == [1,5,7]) + # (reviewer found bug on arr.where(), so doing a spot check for that here too) + result_where2 = arr.where(tokeep2, drop=True) + assert result2.equals(result_where2) + + # repeat tests for UxDataset + result0 = ds.isel(n_face=tokeep0) + assert result0.sizes['n_face'] == result0.uxgrid.n_face == 2 + assert np.all(result0.coords['tokeep'] == [2,6]) + result1 = ds.isel(n_face=tokeep1) + assert result1.sizes['n_face'] == result1.uxgrid.n_face == 3 + result2 = ds.isel(n_face=tokeep2) + assert result2.sizes['n_face'] == result2.uxgrid.n_face == 3 + assert np.all(result2.coords['tokeep'] == [1,5,7]) + result_where2 = ds.where(tokeep2, drop=True) + assert result2.equals(result_where2) + + # also test second reviewer's example (2) from PR #1759: + uxds = ux.tutorial.open_dataset("quad-hexagon") + uxds = uxds.assign_coords(node_id=("n_node", np.arange(uxds.uxgrid.n_node))) + mask = xr.DataArray([True, False, True, False], dims="n_face", + coords={"lab": ("n_face", [10, 20, 30, 40])}) + uxds.isel(n_face=mask) # (just ensuring it doesn't crash) + + # also test what happens if uxarray object has grid dim coords: + ds1 = ux.tutorial.open_dataset("quad-hexagon").assign_coords(n_face=[0, 10, 20, 30]) + arr1 = ds1['t2m'] + wherebig = arr1 > 297.6 # hard-coding just to make the example easily. + # (if quad-hexagon file changes, rework this example) + assert np.all(wherebig.values == [False, True, True, False]) + resultA = arr1.isel(n_face=wherebig) # shouldn't crash! + assert np.all(resultA.coords['n_face'] == [10, 20]) + resultB = ds1.isel(n_face=wherebig) # shouldn't crash! + assert np.all(resultB.coords['n_face'] == [10, 20]) + + def test_indexing_by_dataarray(): """ensure isel() and sel() with indexer=xr.DataArray(...) both work as expected. The dims/coords of the Grid object should never incorporate indexer's dims/coords. - The dims/coords of the data object (UxDataArray or UxDataset) should not incorporate + + The dims of the data object (UxDataArray or UxDataset) should not incorporate the indexer's dims when indexing along a grid dim (e.g. 'n_face') (this is already true), - but should probably incorporate its coords (this isn't true yet; see issue #1712). + (e.g. don't rename 'n_face' to match the indexer's dim name!) see issue #1712 for details. + + Though, the coords of the data object *should* incorporate the indexer's coords when possible: + - Always incorporate the indexer's scalar coords. + - For 1D coords, it is more complicated: + only incorporate 1D coords if grid dim is 'n_face' + (because 'n_edge' and 'n_node' indexing don't necessarily lead to same + size indexer as result) + and when data is located along 'n_face' + (because otherwise the inder's dim doesn't align with the result's grid dim). + I.e., only incorporate 1D coords when indexing face-centered data along 'n_face'. + See #1712 for more details. Regression test for bug in branch (fixed before merging to main) for PR 1729. + + TODO: update accordingly after fixing #1758. """ + # --- n_face indexing of n_face data --- # # ensure grid's dims/coords do not incorporate indexer's dims/coords: indexer0 = xr.DataArray(0, coords={"newcoord": 7}) indexer1 = xr.DataArray([1,2], dims="newdim", coords={"newdim": [7,8]}) + indexer2 = indexer1.assign_coords({"other1dcoord": ("newdim", [9,10]), "scalarcoord": 100}) ds = ux.tutorial.open_dataset("quad-hexagon") + assert "n_face" in ds.dims # behavior for #1712 depends on data location. result0_isel = ds.isel(n_face=indexer0) assert "newcoord" not in result0_isel.uxgrid._ds.coords - # assert "newcoord" in result0_isel.coords # uncomment after fixing #1712 + assert "newcoord" in result0_isel.coords # regression test for #1712 result0_sel = ds.sel(n_face=indexer0) assert "newcoord" not in result0_sel.uxgrid._ds.coords - # assert "newcoord" in result0_sel.coords # uncomment after fixing #1712 + assert "newcoord" in result0_sel.coords # regression test for #1712 result1_isel = ds.isel(n_face=indexer1) assert "newdim" not in result1_isel.uxgrid._ds.dims assert "newdim" not in result1_isel.uxgrid._ds.coords assert "newdim" not in result1_isel.dims and "n_face" in result1_isel.dims # didn't rename 'n_face'. - # assert "newdim" in result1_isel.coords # uncomment after fixing #1712 + assert "newdim" in result1_isel.coords # regression test for #1712 result1_sel = ds.sel(n_face=indexer1) assert "newdim" not in result1_sel.uxgrid._ds.dims assert "newdim" not in result1_sel.uxgrid._ds.coords assert "newdim" not in result1_sel.dims and "n_face" in result1_sel.dims - # assert "newdim" in result1_sel.coords # uncomment after fixing #1712 + assert "newdim" in result1_sel.coords # regression test for #1712 + result2_isel = ds.isel(n_face=indexer2) + assert np.all(result2_isel.coords["other1dcoord"] == [9,10]) + assert result2_isel.coords["scalarcoord"] == 100 + result2_sel = ds.sel(n_face=indexer2) + assert np.all(result2_sel.coords["other1dcoord"] == [9,10]) + assert result2_sel.coords["scalarcoord"] == 100 # repeat tests but with UxDataArray: arr = ds['t2m'] result0_isel = arr.isel(n_face=indexer0) assert "newcoord" not in result0_isel.uxgrid._ds.coords - # assert "newcoord" in result0_isel.coords + assert "newcoord" in result0_isel.coords result0_sel = arr.sel(n_face=indexer0) assert "newcoord" not in result0_sel.uxgrid._ds.coords - # assert "newcoord" in result0_sel.coords + assert "newcoord" in result0_sel.coords result1_isel = arr.isel(n_face=indexer1) assert "newdim" not in result1_isel.uxgrid._ds.dims assert "newdim" not in result1_isel.uxgrid._ds.coords assert "newdim" not in result1_isel.dims and "n_face" in result1_isel.dims - # assert "newdim" in result1_isel.coords + assert "newdim" in result1_isel.coords result1_sel = arr.sel(n_face=indexer1) assert "newdim" not in result1_sel.uxgrid._ds.dims assert "newdim" not in result1_sel.uxgrid._ds.coords assert "newdim" not in result1_sel.dims and "n_face" in result1_sel.dims - # assert "newdim" in result1_sel.coords + assert "newdim" in result1_sel.coords + result2_isel = arr.isel(n_face=indexer2) + assert np.all(result2_isel.coords["other1dcoord"] == [9,10]) + assert result2_isel.coords["scalarcoord"] == 100 + result2_sel = arr.sel(n_face=indexer2) + assert np.all(result2_sel.coords["other1dcoord"] == [9,10]) + assert result2_sel.coords["scalarcoord"] == 100 + + # --- non-n_face indexing and/or non-n_face data --- # + # using loops to avoid writing extremely long test; + # loops are slightly harder to debug but worthwhile to include in at least one test, + # to cover more combinations of cases (e.g., discovered #1758 while making this test). + ds_face = ds + ds_node = ux.tutorial.open_dataset("quad-hexagon-random-node") + #ds_edge = ux.tutorial.open_dataset("quad-hexagon-random-edge") # uncomment after fixing #1758 + assert "n_face" in ds_face.dims + assert "n_node" in ds_node.dims + #assert "n_edge" in ds_edge.dims # uncomment after fixing #1758 + counter = 0 # (count up during loop to make sure nothing is skipped unexpectedly) + for method in "isel", "sel": + for dataset in [ds_face, ds_node]: # include ds_edge after fixing #1758 + for grid_dim in ("n_face", "n_edge", "n_node"): + for to_array in [False, True]: + if grid_dim == "n_face" and "n_face" in dataset.dims: + continue # already tested this above! + counter += 1 + obj = dataset[list(dataset.data_vars)[0]] if to_array else dataset + result0 = getattr(obj, method)({grid_dim: indexer0}) + assert "newcoord" not in result0.uxgrid._ds.coords + assert "newcoord" in result0.coords # 0D coord should always show up! + result1 = getattr(obj, method)({grid_dim: indexer1}) + assert "newdim" not in result1.uxgrid._ds.dims + assert "newdim" not in result1.uxgrid._ds.coords + assert "newdim" not in result1.dims + assert "newdim" not in result1.coords + result2 = getattr(obj, method)({grid_dim: indexer2}) + assert "other1dcoord" not in result2.uxgrid._ds.coords + assert "other1dcoord" not in result2.coords + assert result2.coords["scalarcoord"] == 100 + n_data_grid_dim_combos = 2 * 3 - 1 # after fixing #1758, update to: 3 * 3 - 1. + # the -1 accounts for skipping when both are "n_face" above. + assert counter == 2 * n_data_grid_dim_combos * 2 + + +def test_indexing_of_uxarray_obj_with_grid_dim_coords(): + """ensure obj.isel() and obj.sel() work properly when obj has coords along grid dim. + Regression test for second reviewer's example (4) from PR 1759. + """ + # example (4) from PR 1759 + uxds = ux.tutorial.open_dataset("quad-hexagon").assign_coords(n_face=[0, 10, 20, 30]) + pick_by_label = xr.DataArray([0, 20], dims="station", coords={"station": ["A", "B"]}) + pick_by_index = xr.DataArray([0, 2], dims="station", coords={"station": ["A", "B"]}) + + by_index = uxds.isel(n_face=pick_by_index) + assert set(by_index.coords) == {'n_face', 'station'} + assert np.all(by_index.coords['n_face'] == [0, 20]) + assert np.all(by_index.coords['station'] == ['A', 'B']) + by_label = uxds.sel(n_face=pick_by_label) + assert by_label.identical(by_index) + + # repeat above, for UxDataArray: + uxarr = uxds['t2m'] + by_index = uxarr.isel(n_face=pick_by_index) + assert set(by_index.coords) == {'n_face', 'station'} + assert np.all(by_index.coords['n_face'] == [0, 20]) + assert np.all(by_index.coords['station'] == ['A', 'B']) + by_label = uxarr.sel(n_face=pick_by_label) + assert by_label.identical(by_index) + + # sanity check / spot tests: indexer with grid dim and coords: + # (A) should be allowed if coords don't conflict with obj + # (B) should crash if coords conflict with obj + indexerA_sel = xr.DataArray([0, 20], dims="n_face", coords={"n_face": [0, 20]}) + indexerB_sel = xr.DataArray([0, 20], dims="n_face", coords={"n_face": [7, 99]}) + indexerA_isel = xr.DataArray([0, 2], dims="n_face", coords={"n_face": [0, 20]}) + indexerB_isel = xr.DataArray([0, 2], dims="n_face", coords={"n_face": [0, 2]}) + resultA_sel = uxds.sel(n_face=indexerA_sel) + assert np.all(resultA_sel.coords['n_face'] == [0, 20]) + with pytest.raises(IndexError, match="dimension coordinate 'n_face' conflicts"): + _resultB_sel = uxds.sel(n_face=indexerB_sel) + resultA_isel = uxds.isel(n_face=indexerA_isel) + assert np.all(resultA_isel.coords['n_face'] == [0, 20]) + with pytest.raises(IndexError, match="dimension coordinate 'n_face' conflicts"): + _resultA_sel = uxds.isel(n_face=indexerB_isel) + + # repeat sanity checks, for UxDataArray: + resultA_sel = uxarr.sel(n_face=indexerA_sel) + assert np.all(resultA_sel.coords['n_face'] == [0, 20]) + with pytest.raises(IndexError, match="dimension coordinate 'n_face' conflicts"): + _resultB_sel = uxarr.sel(n_face=indexerB_sel) + resultA_isel = uxarr.isel(n_face=indexerA_isel) + assert np.all(resultA_isel.coords['n_face'] == [0, 20]) + with pytest.raises(IndexError, match="dimension coordinate 'n_face' conflicts"): + _resultA_sel = uxarr.isel(n_face=indexerB_isel) + + +def test_dataset_isel_keeps_bonus_coords(): + """ensure UxDataset.isel() keeps "bonus" coords, + i.e. coords in the dataset which do not actually appear in any data var. + Regression test for bug (3) discovered during review of PR #1759. + """ + ds = ux.tutorial.open_dataset('quad-hexagon') + ds = ds.assign_coords({'bonus_coord': xr.DataArray(['a', 'b'], dims=['bonus_dim'])}) + assert 'bonus_coord' in ds.coords + assert all('bonus_coord' not in arr for arr in ds.data_vars.values()) + result = ds.isel(n_face=0) + assert 'bonus_coord' in result.coords and 'bonus_dim' in result.dims + + # also test using second reviewer's example (1) from PR #1759: + 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]) + assert 'node_id' in sub.coords + + +def test_indexing_by_size_0_array(): + """ensure indexing by size 0 indexer works as expected (size 0 result). + Regression test for second reviewer's example (3) from PR 1759. + """ + def _as_array(obj): # return UxDataArray from UxDataArray, or UxDataset's data_vars[0]. + return obj[list(obj.data_vars)[0]] if isinstance(obj, ux.UxDataset) else obj + + # example (3) from PR 1759 + uxds = ux.tutorial.open_dataset("quad-hexagon") + assert set(uxds.dims) == {'n_face'} + uxds = uxds.assign_coords(node_id=("n_node", np.arange(uxds.uxgrid.n_node))) + assert set(uxds.dims) == {'n_face', 'n_node'} + hits = np.where(uxds.uxgrid.edge_lon > 1e9)[0] # empty + assert hits.size == 0 + result = uxds.isel(n_edge=xr.DataArray(hits, dims="selected")) + assert _as_array(result).size == 0 + assert 'node_id' in result.coords + assert set(result.dims) == {'n_face', 'n_node'} + assert result.sizes['n_face'] == result.sizes['n_node'] == 0 + assert result.uxgrid.n_face == result.uxgrid.n_node == result.uxgrid.n_edge == 0 + + # followup to example (3): fancy empty array (with dim that has coords (which are empty)) + uxds = ux.tutorial.open_dataset("quad-hexagon") + fancy_empty = xr.DataArray(np.array([]), dims="selected", + coords={"selected": np.array([], dtype=int)}) + for grid_dim in ("n_face", "n_edge", "n_node"): + resultA = uxds.isel({grid_dim: fancy_empty}) + assert resultA.sizes == {'n_face': 0} # and, must not pick up "selected" dim from indexer. + resultB = uxds.sel({grid_dim: fancy_empty}) + assert resultB.sizes == {'n_face': 0} + # repeat for UxDataArray: + resultC = uxds['t2m'].isel({grid_dim: fancy_empty}) + assert resultC.sizes == {'n_face': 0} + resultD = uxds['t2m'].sel({grid_dim: fancy_empty}) + assert resultD.sizes == {'n_face': 0} + + # simpler tests, but applied across isel, sel, UxDataArray, UxDataset, n_edge, n_node, and n_face: + ds_face = ux.tutorial.open_dataset("quad-hexagon-random-face") + ds_node = ux.tutorial.open_dataset("quad-hexagon-random-node") + #ds_edge = ux.tutorial.open_dataset("quad-hexagon-random-edge") # uncomment after fixing #1758 + assert "n_face" in ds_face.dims + assert "n_node" in ds_node.dims + #assert "n_edge" in ds_edge.dims # uncomment after fixing #1758 + counter = 0 # (count up during loop to make sure nothing is skipped unexpectedly) + for method in "isel", "sel": + for dataset in [ds_face, ds_node]: # include ds_edge after fixing #1758 + for grid_dim in ("n_face", "n_edge", "n_node"): + for to_array in [False, True]: + counter += 1 + obj = _as_array(dataset) if to_array else dataset + # index by empty list: + result = getattr(obj, method)({grid_dim: []}) + assert _as_array(result).size == 0 + assert set(result.dims) == set(obj.dims) + # index by empty numpy array: + result = getattr(obj, method)({grid_dim: np.array([])}) + assert _as_array(result).size == 0 + assert set(result.dims) == set(obj.dims) + # index by empty xr.DataArray: + result = getattr(obj, method)({grid_dim: xr.DataArray([])}) + assert _as_array(result).size == 0 + assert set(result.dims) == set(obj.dims) + # index by empty xr.DataArray with a scalar coord: + indexer = xr.DataArray([]).assign_coords({"scalar7": 7}) + result = getattr(obj, method)({grid_dim: indexer}) + assert _as_array(result).size == 0 + assert set(result.dims) == set(obj.dims) + assert result.coords["scalar7"] == 7 + def test_indexing_does_not_edit_indexers_dict(): """ensure isel() and sel() do not edit the provided indexers dict. @@ -253,7 +520,7 @@ def test_sel_crash_if_provided_selection_options_with_coordless_dims(): whenever any of the indexed dims have no associated coordinates. (Tests below also demonstrate that this behavior is consistent with xarray.) Regression test inspired by reviewer comment in #1641. - TODO: fix #1714 then uncomment the relevant UxDataset tests below + Also includes a regression test for #1714. """ kw_options = ({"method": "nearest"}, {"method": "nearest", "tolerance": 0.1}) @@ -264,13 +531,13 @@ def test_sel_crash_if_provided_selection_options_with_coordless_dims(): assert set(ds0.dims) == {'n_face'} ds0_labeled = ds0.assign_coords({'n_face': [0,10,20,30]}) - # (uncomment the next few lines after fixing #1714) - # ds0.sel(n_face=[0,1]) # (sanity check: no crash when no options provided) - # for kw in kw_options: - # with pytest.raises(ValueError, match=r"cannot supply selection options.+for dimension 'n_face'"): - # ds0.sel(n_face=[0,1], **kw) # provides method, tolerance, or both. - # # separately: checking to ensure that passing these options is fine in "labeled" case. - # ds0_labeled.sel(n_face=[0,10], **kw) + # (the next few lines also serve as a regression test for #1714) + ds0.sel(n_face=[0,1]) # (sanity check: no crash when no options provided) + for kw in kw_options: + with pytest.raises(ValueError, match=r"cannot supply selection options.+for dimension 'n_face'"): + ds0.sel(n_face=[0,1], **kw) # provides method, tolerance, or both. + # separately: checking to ensure that passing these options is fine in "labeled" case. + ds0_labeled.sel(n_face=[0,10], **kw) # ensure same behavior for xarray objects: ds0.to_xarray().sel(n_face=[0,1]) @@ -280,9 +547,9 @@ def test_sel_crash_if_provided_selection_options_with_coordless_dims(): ds0_labeled.to_xarray().sel(n_face=[0,10], **kw) # ensure supplying just tolerance raises a different error, if indexing is otherwise valid: - # (uncomment the next few lines after fixing #1714) - # with pytest.raises(ValueError, match=r"tolerance argument only valid if doing.+"): - # ds0_labeled.sel(n_face=[0,10], tolerance=0.1) + # (the following pytest.raises statement also serves as a regression test for #1714) + with pytest.raises(ValueError, match=r"tolerance argument only valid if doing.+"): + ds0_labeled.sel(n_face=[0,10], tolerance=0.1) with pytest.raises(ValueError, match=r"tolerance argument only valid if doing.+"): ds0_labeled.to_xarray().sel(n_face=[0,10], tolerance=0.1) @@ -388,3 +655,83 @@ def test_sel_crash_if_provided_selection_options_with_coordless_dims(): arr1.to_xarray().sel(time=4, **kw) with pytest.raises(ValueError, match=r"cannot supply selection options"): arr1.to_xarray().sel(time=4, n_face=[3], **kw) + +def test_isel_crash_if_coordinates_conflict(): + """Ensure isel crashes if there is a coordinates conflict, + such as indexing an array with time dim by an array with a scalar time coord. + Regression test for bug (2) discovered during review of PR #1759. + (Also checks indexing an array with scalar time coord by an array with a time dim.) + """ + ds = ux.tutorial.open_dataset("outCSne30-timeseries") + arr = ds['psi'] + indexer0 = xr.DataArray(0, coords={'time': arr['time'][0].item()}) + indexer1 = arr.isel(time=0).argmax('n_face') + indexer2 = xr.DataArray([0,1], dims='time', coords={'time': arr['time'][:2].values}) + assert isinstance(indexer1, ux.UxDataArray) + assert 'time' in indexer1.coords and 'time' not in indexer1.dims + MATCH_ERRMSG = "dimension coordinate 'time' conflicts between indexed and indexing objects" + MATCH_ERRMSG_2 = "The indexer's dimension \('time'\) already exists as a scalar" + with pytest.raises(IndexError, match=MATCH_ERRMSG): + arr.to_xarray().isel(n_face=indexer0) # sanity check that xarray also crashes here. + with pytest.raises(IndexError, match=MATCH_ERRMSG): + arr.isel(n_face=indexer0) + with pytest.raises(IndexError, match=MATCH_ERRMSG): + arr.to_xarray().isel(n_face=indexer1) # sanity check that xarray also crashes here. + with pytest.raises(IndexError, match=MATCH_ERRMSG): + arr.isel(n_face=indexer1) + arr_t0 = arr.isel(time=0) + assert 'time' in arr_t0.coords and 'time' not in arr_t0.dims + with pytest.raises(ux.errors.DimensionError, match=MATCH_ERRMSG_2): + arr_t0.isel(n_face=indexer2) + + # repeat tests for UxDataArray: + with pytest.raises(IndexError, match=MATCH_ERRMSG): + ds.to_xarray().isel(n_face=indexer0) + with pytest.raises(IndexError, match=MATCH_ERRMSG): + ds.isel(n_face=indexer0) + with pytest.raises(IndexError, match=MATCH_ERRMSG): + ds.to_xarray().isel(n_face=indexer1) + with pytest.raises(IndexError, match=MATCH_ERRMSG): + ds.isel(n_face=indexer1) + ds_t0 = ds.isel(time=0) + assert 'time' in ds_t0.coords and 'time' not in ds_t0.dims + with pytest.raises(ux.errors.DimensionError, match=MATCH_ERRMSG_2): + ds_t0.isel(n_face=indexer2) + +def test_isel_when_indexer_dim_in_uxarray_obj(): + """Ensure isel() crashes with NotImplementedError when the indexer's dim has + the same name as a dim or non-scalar coordinate in the uxarray object being indexed. + Regression test for follow-up to bug (2) discovered during review of PR #1759. + """ + ds = ux.tutorial.open_dataset('quad-hexagon').expand_dims(time=[100,200]) + arr = ds['t2m'] + imax = arr.argmax('n_face') + assert set(imax.dims) == {'time'} and set(imax.coords) == {'time'} + assert set(ds.dims) == {'time', 'n_face'} and set(ds.coords) == {'time'} + assert imax.coords['time'].equals(ds.coords['time']) + # pure xarray indexing here returns a result with only 'time' dimension; + # that's the main reason uxarray should raise NotImplementedError in this case. + xr_result = ds.to_xarray().isel(n_face=imax.to_xarray()) + assert set(xr_result.dims) == {'time'} + ERRMSG = ( + r"Indexing a {typestr} .+ using an xarray DataArray whose dimension \('time'\) " + r"is already present .+ is not yet supported" + ) + with pytest.raises(NotImplementedError, match=ERRMSG.format(typestr="UxDataset")): + ds.isel(n_face=imax) + ds1 = ds.swap_dims({'time': 'otherdim'}) + assert 'time' in ds1.coords and 'time' not in ds1.dims + with pytest.raises(NotImplementedError, match=ERRMSG.format(typestr="UxDataset")): + ds1.isel(n_face=imax) + + # repeat tests for UxDataArray: + assert set(arr.dims) == {'time', 'n_face'} and set(arr.coords) == {'time'} + assert imax.coords['time'].equals(arr.coords['time']) + xr_result = arr.to_xarray().isel(n_face=imax.to_xarray()) + assert set(xr_result.dims) == {'time'} + with pytest.raises(NotImplementedError, match=ERRMSG.format(typestr="UxDataArray")): + arr.isel(n_face=imax) + arr1 = arr.swap_dims({'time': 'otherdim'}) + assert 'time' in arr1.coords and 'time' not in arr1.dims + with pytest.raises(NotImplementedError, match=ERRMSG.format(typestr="UxDataArray")): + arr1.isel(n_face=imax) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 80a023c25..96adfe4d9 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -11,7 +11,6 @@ from xarray.core.options import OPTIONS from xarray.core.utils import UncachedAccessor -import uxarray from uxarray.constants import GRID_DIMS from uxarray.core.aggregation import _uxda_grid_aggregate from uxarray.core.gradient import ( @@ -45,7 +44,12 @@ from uxarray.plot.accessor import UxDataArrayPlotAccessor from uxarray.remap.accessor import RemapAccessor from uxarray.subset import DataArraySubsetAccessor -from uxarray.utils.coords import _preserve_valid_coords +from uxarray.utils.coords import ( + _assert_grid_dim_coord_consistent_if_in_both, + _assign_grid_dim_indexer_coords_if_appropriate, + _crash_if_1d_xarray_indexer_dim_in_uxarray_obj, + _preserve_valid_coords, +) from uxarray.utils.imports import _raise_hint_if_optional_deps_missing if TYPE_CHECKING: @@ -586,8 +590,10 @@ def to_dataset( ------- uxds: UxDataSet """ + from uxarray.core.dataset import UxDataset + xrds = super().to_dataset(dim=dim, name=name, promote_attrs=promote_attrs) - uxds = uxarray.core.dataset.UxDataset(xrds, uxgrid=self._uxgrid) + uxds = UxDataset(xrds, uxgrid=self._uxgrid) return uxds @@ -1996,7 +2002,10 @@ def isel( using n_edge=7 selects just the two faces touching edge 7. For data on 'n_face', the result would have 'n_face' with just those two faces. For data on 'n_edge', the result would have 'n_edge' with all edges located on either of those two faces. + Grid dimension indexers cannot have more than 1 dimension (such as a 2D DataArray). + Grid dimensions are never renamed (even if indexed by 1D DataArray with different dim name). + Grid dimension indexer cannot have a non-grid dimension which exists in the original UxDataArray. Parameters ---------- @@ -2032,6 +2041,10 @@ def isel( ------- UxDataArray A new UxDataArray indexed according to `indexers` and updated grid if applicable. + If indexer DataArrays have coordinates that do not conflict with + this object, then these coordinates will be attached, + except that 1D coordinates of indexers applied along a grid dimension will + only be included if it is 'n_face' and the data also has 'n_face' dimension. Raises ------ @@ -2061,22 +2074,27 @@ def isel( indexers = indexers.copy() # don't modify the original dict grid_indexer = indexers.pop(grid_dim) + _crash_if_1d_xarray_indexer_dim_in_uxarray_obj(self, grid_dim, grid_indexer) + sliced_grid = self.uxgrid.isel( **{grid_dim: grid_indexer}, inverse_indices=inverse_indices ) - da = self._slice_from_grid(sliced_grid) + result = self._slice_from_grid(sliced_grid) + + result = _assign_grid_dim_indexer_coords_if_appropriate( + result, grid_dim, grid_indexer + ) # if there are any remaining indexers, apply them if indexers: - xarr = super(UxDataArray, da).isel( + result = super(UxDataArray, result).isel( indexers=indexers, drop=drop, missing_dims=missing_dims ) # re‐wrap so the grid sticks around - return type(self)(xarr, uxgrid=sliced_grid) + result = type(self)(result, uxgrid=sliced_grid) - # no other dims, return the grid‐sliced da - return da + return result else: # len(grid_dims)>1; _validate_indexers should have crashed. raise AssertionError("internal implementation error if reached this line") @@ -2096,7 +2114,10 @@ def sel( using n_edge=7 selects just the two faces touching edge 7. For data on 'n_face', the result would have 'n_face' with just those two faces. For data on 'n_edge', the result would have 'n_edge' with all edges located on either of those two faces. + Grid dimension indexers cannot have more than 1 dimension (such as a 2D DataArray). + Grid dimensions are never renamed (even if indexed by 1D DataArray with different dim name). + Grid dimension indexer cannot have a non-grid dimension which exists in the original UxDataArray. By default, grid dims do not have coordinates assigned. But, if they have been assigned, `.sel()` respects them in the intuitive way. For example, @@ -2157,7 +2178,8 @@ def sel( and the uxgrid indexed appropriately as well, if indexing any grid dim. If indexer DataArrays have coordinates that do not conflict with this object, then these coordinates will be attached, - except for indexers along a grid dimension (see issue #1712). + except that 1D coordinates of indexers applied along a grid dimension will + only be included if it is 'n_face' and the data also has 'n_face' dimension. In general, the result's data will be a view of the data in this array, unless indexing along a grid dimension or otherwise triggering vectorized indexing by using an array indexer, @@ -2203,6 +2225,10 @@ def sel( # offload the grid-indexing work to isel(): result = self.isel({grid_dim: grid_indices}, drop=drop) + # special case: if grid_dim in indexer and result.coords, ensure consistency. + # (all other coords' consistency checks already occurred in isel().) + _assert_grid_dim_coord_consistent_if_in_both(result, grid_dim, grid_indexer) + # index by other dims if any remain: ds = result.to_xarray().sel( indexers=indexers, # (grid_dim indexer was popped) @@ -2311,7 +2337,7 @@ def _slice_from_grid(self, sliced_grid): "Data variable must be either node, edge, or face centered." ) - return UxDataArray(da_sliced, uxgrid=sliced_grid) + return type(self)(da_sliced, uxgrid=sliced_grid) def get_dual(self): """Compute the dual mesh for a data array, returns a new data array @@ -2349,9 +2375,7 @@ def get_dual(self): dims = [dim_map.get(dim, dim) for dim in self.dims] # Construct the new data array - uxda = uxarray.UxDataArray( - uxgrid=dual, data=self.data, dims=dims, name=self.name - ) + uxda = type(self)(uxgrid=dual, data=self.data, dims=dims, name=self.name) return uxda diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index e81431707..9de3be7bd 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -11,7 +11,7 @@ from xarray.core.options import OPTIONS from xarray.core.utils import UncachedAccessor -import uxarray +from uxarray.constants import GRID_DIMS from uxarray.core.dataarray import UxDataArray from uxarray.core.utils import ( _map_dims_to_ugrid, @@ -28,6 +28,11 @@ from uxarray.io._healpix import get_zoom_from_cells from uxarray.plot.accessor import UxDatasetPlotAccessor from uxarray.remap.accessor import RemapAccessor +from uxarray.utils.coords import ( + _assert_grid_dim_coord_consistent_if_in_both, + _assign_grid_dim_indexer_coords_if_appropriate, + _crash_if_1d_xarray_indexer_dim_in_uxarray_obj, +) class UxDataset(xr.Dataset): @@ -387,33 +392,39 @@ def from_healpix( return cls.from_xarray(ds, uxgrid, {face_dim: "n_face"}) - def _slice_dataset_from_grid(self, sliced_grid, grid_dim: str, grid_indexer): + def _slice_from_grid(self, sliced_grid): + """returns UxDataset based on slicing self according to sliced_grid. + sliced_grid should be a ``Grid`` which came directly from self.uxgrid.isel(...) + (or from slicing something equal to self.uxgrid), else behavior is undefined. + """ data_vars = {} for name, da in self.data_vars.items(): - if grid_dim in da.dims: - if hasattr(da, "_slice_from_grid"): - data_vars[name] = da._slice_from_grid(sliced_grid) - else: - data_vars[name] = da.isel({grid_dim: grid_indexer}) + if hasattr(da, "_slice_from_grid") and any( + dim in da.dims for dim in GRID_DIMS + ): + data_vars[name] = da._slice_from_grid(sliced_grid) else: data_vars[name] = da - coords = {} - for cname, cda in self.coords.items(): - if grid_dim in cda.dims: - # Prefer authoritative coords from the sliced grid if available - replacement = getattr(sliced_grid, cname, None) - coords[cname] = ( - replacement - if replacement is not None - else cda.isel({grid_dim: grid_indexer}) - ) - else: - coords[cname] = cda - - ds = xr.Dataset(data_vars=data_vars, coords=coords, attrs=self.attrs) + # Also account for any coords which aren't attached to any data_var: + bonus_coords = {} + for coord in self.coords: + for data_var in data_vars.values(): + if coord in data_var.coords: + break + else: # didn't break + da = self.coords[coord] + if hasattr(da, "_slice_from_grid") and any( + dim in da.dims for dim in GRID_DIMS + ): + bonus_coords[coord] = da._slice_from_grid(sliced_grid) + else: + bonus_coords[coord] = da - return ds + ds_sliced = xr.Dataset( + data_vars=data_vars, coords=bonus_coords, attrs=self.attrs + ) + return type(self)(ds_sliced, uxgrid=sliced_grid) def isel( self, @@ -435,7 +446,10 @@ def isel( using n_edge=7 selects just the two faces touching edge 7. For data on 'n_face', the result would have 'n_face' with just those two faces. For data on 'n_edge', the result would have 'n_edge' with all edges located on either of those two faces. + Grid dimension indexers cannot have more than 1 dimension (such as a 2D DataArray). + Grid dimensions are never renamed (even if indexed by 1D DataArray with different dim name). + Grid dimension indexer cannot have a non-grid dimension which exists in the original UxDataset. Parameters ---------- @@ -471,6 +485,10 @@ def isel( ------- UxDataset A new UxDataset indexed according to `indexers` and updated grid if applicable. + If indexer DataArrays have coordinates that do not conflict with + this object, then these coordinates will be attached, + except that 1D coordinates of indexers applied along a grid dimension will + only be included if it is 'n_face' and the data also has 'n_face' dimension. Raises ------ @@ -500,23 +518,27 @@ def isel( indexers = indexers.copy() # don't modify the original dict grid_indexer = indexers.pop(grid_dim) - # slice the grid + _crash_if_1d_xarray_indexer_dim_in_uxarray_obj(self, grid_dim, grid_indexer) + sliced_grid = self.uxgrid.isel( **{grid_dim: grid_indexer}, inverse_indices=inverse_indices ) - ds = self._slice_dataset_from_grid( - sliced_grid=sliced_grid, - grid_dim=grid_dim, - grid_indexer=grid_indexer, + result = self._slice_from_grid(sliced_grid) + + result = _assign_grid_dim_indexer_coords_if_appropriate( + result, grid_dim, grid_indexer ) + # if there are any remaining indexers, apply them if indexers: - ds = xr.Dataset.isel( - ds, indexers=indexers, drop=drop, missing_dims=missing_dims + result = super(UxDataset, result).isel( + indexers=indexers, drop=drop, missing_dims=missing_dims ) + # re‐wrap so the grid sticks around + result = type(self)(result, uxgrid=sliced_grid) - return type(self)(ds, uxgrid=sliced_grid) + return result else: # len(grid_dims)>1; _validate_indexers should have crashed. raise AssertionError("internal implementation error if reached this line") @@ -537,7 +559,10 @@ def sel( using n_edge=7 selects just the two faces touching edge 7. For data on 'n_face', the result would have 'n_face' with just those two faces. For data on 'n_edge', the result would have 'n_edge' with all edges located on either of those two faces. + Grid dimension indexers cannot have more than 1 dimension (such as a 2D DataArray). + Grid dimensions are never renamed (even if indexed by 1D DataArray with different dim name). + Grid dimension indexer cannot have a non-grid dimension which exists in the original UxDataset. By default, grid dims do not have coordinates assigned. But, if they have been assigned, `.sel()` respects them in the intuitive way. For example, @@ -566,7 +591,8 @@ def sel( multi-index, the indexer may also be a dict-like object with keys matching index level names. If DataArrays are passed as indexers, xarray-style indexing will be - carried out. See :ref:`indexing` for the details. + carried out (see :ref:`indexing` for the details), + with one exception: grid dimensions will never be renamed. One of indexers or indexers_kwargs must be provided. method : {None, "nearest", "pad", "ffill", "backfill", "bfill"}, optional Method to use for inexact matches: @@ -599,7 +625,8 @@ def sel( and the uxgrid indexed appropriately as well, if indexing any grid dim. If indexer DataArrays have coordinates that do not conflict with this object, then these coordinates will be attached, - except for indexers along a grid dimension (see issue #1712). + except that 1D coordinates of indexers applied along a grid dimension will + only be included if it is 'n_face' and the data also has 'n_face' dimension. In general, each array's data will be a view of the array's data in this dataset, unless indexing along a grid dimension or otherwise triggering vectorized indexing by using an array indexer, @@ -646,6 +673,10 @@ def sel( # offload the grid-indexing work to isel(): result = self.isel({grid_dim: grid_indices}, drop=drop) + # special case: if grid_dim in indexer and result.coords, ensure consistency. + # (all other coords' consistency checks already occurred in isel().) + _assert_grid_dim_coord_consistent_if_in_both(result, grid_dim, grid_indexer) + # index by other dims if any remain: ds = result.to_xarray().sel( indexers=indexers, # (grid_dim indexer was popped) @@ -921,7 +952,7 @@ def get_dual(self): ) # Initialize new dataset - dataset = uxarray.UxDataset(uxgrid=dual) + dataset = type(self)(uxgrid=dual) # Dictionary to swap dimensions dim_map = {"n_face": "n_node", "n_node": "n_face"} @@ -932,9 +963,7 @@ def get_dual(self): dims = [dim_map.get(dim, dim) for dim in self[var].dims] # Construct the new data array - uxda = uxarray.UxDataArray( - uxgrid=dual, data=self[var].data, dims=dims, name=var - ) + uxda = UxDataArray(uxgrid=dual, data=self[var].data, dims=dims, name=var) # Add data array to dataset dataset[var] = uxda diff --git a/uxarray/core/utils.py b/uxarray/core/utils.py index c82d69928..84c0a2dc6 100644 --- a/uxarray/core/utils.py +++ b/uxarray/core/utils.py @@ -1,3 +1,5 @@ +import warnings + import numpy as np import xarray as xr from xarray.core.utils import either_dict_or_kwargs @@ -186,10 +188,7 @@ def _resolve_coordinate_labels_to_indices( ): """returns indices which would be selected by coord_array.sel({dim: labels_to_sel}, ...) coord_array.isel({dim: result}) should be equivalent to coord_array.sel({dim: labels_to_sel}, ...). - - (Implementation here drops extra coordinates from any indexers, - but if it is being applied to grid dims for sel() then it will produce behavior - which is consistent with isel(), unless issue #1712 gets fixed.) + If labels_to_sel is an xr.DataArray, its coords/dims will also be attached to the result. dim: str dimension name to select along @@ -203,10 +202,26 @@ def _resolve_coordinate_labels_to_indices( # (Maybe a more efficient implementation exists, but this is simple and gives correct results.) indices = xr.DataArray(np.arange(coord_array.sizes[dim]), dims=dim) _indices_coord_name = f"__{dim}_indices__" # just needs to be any unused name. + if _indices_coord_name in coord_array.coords: + warnings.warn( + f"Coordinate {_indices_coord_name!r} already exists in coord_array.coords " + "and will be overwritten, which may cause errors or subtly incorrect results..." + ) if hasattr(coord_array, "to_xarray"): # convert to xarray to avoid recursive sel() coord_array = coord_array.to_xarray() coord_with_indices = coord_array.assign_coords({_indices_coord_name: indices}) selected = coord_with_indices.sel( {dim: labels_to_sel}, method=method, tolerance=tolerance ) - return selected[_indices_coord_name].values # (return as np.ndarray, not DataArray) + result = selected[_indices_coord_name] + if isinstance(labels_to_sel, xr.DataArray): + # handle coords appropriately + result = result.drop_vars((dim, _indices_coord_name)) + # (drop grid dim coords because the caller is expected to handle those directly; + # here the goal is just to properly propagate any other coords from labels_to_sel.) + result = result.rename(None) # no reason to keep the _indices_coord_name + # (and keeping it for longer could maybe cause confusing error later?) + else: + # drop all coords/name info which was added internally during this method. + result = result.values + return result diff --git a/uxarray/utils/coords.py b/uxarray/utils/coords.py index caaccf525..15c0bcf2e 100644 --- a/uxarray/utils/coords.py +++ b/uxarray/utils/coords.py @@ -6,8 +6,11 @@ import numpy as np import xarray as xr +import xarray.core.coordinates import xarray.core.utils as xr_core_utils +from uxarray.errors import DimensionError + def _preserve_valid_coords( obj: xr.DataArray | xr.Dataset, @@ -84,3 +87,126 @@ def _indices1d_from_indexing(xarray_obj, dim, indexer): if dim in xarray_obj.coords: xarray_obj = xarray_obj.drop_vars(dim) return xarray_obj[dim].isel({dim: indexer}).values + + +def _crash_if_1d_xarray_indexer_dim_in_uxarray_obj(uxarray_obj, grid_dim, indexer): + """if xarray indexer's dim is not grid_dim and is in uxarray object, raise NotImplementedError + rather than silently treating indexer as indexer.values, or silently giving a confusing result. + (Should only be applied where indexer is indexing uxarray_obj along grid_dim.) + + In this case, uxarray_obj.to_xarray().isel({grid_dim: indexer}) produces "pointwise" indexing, + but that is difficult to support reliably along a grid dimension. + + Example: uxarr.isel(n_face=uxarr.argmax('n_face')) for uxarr a UxDataArray with 'time' dimension, + should maybe produce "the maximum values of uxarr across all faces, at every face where there is + a maximum at any given time, for all times" or something like that? + (It very confusing, and unclear if any user is doing something like this.) + + In xarray, uxarr.to_xarray().isel(n_face=uxarr.argmax('n_face')) produces an array with only the + 'time' dimension, telling the maximum value of uxarray across all faces, at each time. + + (Only handles 1D xr.DataArray indexers; 0D, 2D+, and non-DataArray indexers are handled elsewhere.) + """ + if isinstance(indexer, xr.DataArray): + if indexer.ndim == 1: + the_dim = indexer.dims[0] + nonscalar_coords = [ + c for c in uxarray_obj.coords if len(uxarray_obj.coords[c].dims) > 0 + ] + if the_dim != grid_dim and ( + the_dim in uxarray_obj.dims or the_dim in nonscalar_coords + ): + raise NotImplementedError( + f"Indexing a {type(uxarray_obj).__name__} along a grid dimension ({grid_dim!r}), using " + f"an xarray DataArray whose dimension ({the_dim!r}) is already present in the original " + f"{type(uxarray_obj).__name__}'s dims or non-scalar coords, is not yet supported. " + f"Consider using obj.to_xarray() for basic xarray indexing, " + f"using indexer.data to convert to indexer to a non-DataArray object, " + f"or using indexer.rename({{{the_dim!r}: 'any_unused_dim_name'}}) to avoid matching dims." + ) + # all other cases handled elsewhere. + + +def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexer): + """returns uxarray_obj but with coords assigned from indexer if appropriate. + "appropriate" `indexer` is a 0D or 1D xr.DataArray and has any relevant coordinates to assign. + 2D+ xr.DataArray indexers are not supported for grid dimensions and cause DimensionError here. + All other indexers do not provide coordinate info, so uxarray_obj gets returned unchanged. + + For 0D indexer, just assign indexer.coords if nonempty (else, return uxarray_obj unchanged). + For 1D indexer, depends on grid_dim and uxarray_obj. + if grid_dim=="n_face" and "n_face" in uxarray_obj: + swap indexer's 1 dim to be "n_face" instead of its original name, + then assign indexer.coords. + in all other cases: + drop indexer's 1 dim, then assign indexer.coords if nonempty. + (For "n_edge" and "n_node" indexing, the result's shape won't necessarily + match the indexer's shape, so coords along the grid dim can't be assigned. + Meanwhile, if grid_dim not in uxarray_obj, it is impossible to assign coords + along that dim, so once again, coords along the grid dim can't be assigned.) + Note: if indexer is booleans, instead use indexer.isel(indexer_dim=indexer).coords, + because the result will only keep values wherever indexer value is True. + (Also in this case, if indexer.to_xarray() exists, call it, to avoid recursion.) + """ + if isinstance(indexer, xr.DataArray): + indexing_1d_bool_along_n_face = ( + indexer.ndim == 1 + and indexer.dtype == bool + and grid_dim == "n_face" + and "n_face" in uxarray_obj.dims + ) + if not indexing_1d_bool_along_n_face: + # make sure indexer.coords are consistent with uxarray_obj's coords. + # e.g., if uxarray_obj has time dim but indexer has time scalar coord, crash! + xr.core.coordinates.assert_coordinate_consistent( + uxarray_obj, indexer.coords.variables + ) + # else: handle that check below (using indexer.isel(...) instead.) + if indexer.ndim == 0: + coords = indexer.coords + elif indexer.ndim == 1: + the_dim = indexer.dims[0] + if grid_dim == "n_face" and "n_face" in uxarray_obj.dims: + if indexer.dtype == bool: + if hasattr(indexer, "to_xarray"): + indexer = indexer.to_xarray() + indexer = indexer.isel({the_dim: indexer}) + xr.core.coordinates.assert_coordinate_consistent( + uxarray_obj, indexer.coords.variables + ) + if ( + the_dim in uxarray_obj.coords + and len(uxarray_obj.coords[the_dim].dims) == 0 + ): + raise DimensionError( + f"The indexer's dimension ({the_dim!r}) already exists as a scalar " + f"coordinate in the {type(uxarray_obj).__name__} object being indexed." + ) # + coords = indexer.swap_dims({the_dim: "n_face"}).coords + else: + # remove any 1D coords (but keep scalar coords). + # Drop by name to properly handle any subtleties of a 0D indexer + # (which would crash .isel({the_dim: 0}, drop=True, + # and would subtly add a new dim to the result if the_dim isn't dropped.) + coords = indexer.drop_vars( + [c for c in indexer.coords if the_dim in indexer[c].dims] + ).coords + else: + raise DimensionError( + f"2D+ indexers are not supported for grid dimensions. Got xr.DataArray " + f"indexer with ndim={indexer.ndim}, along grid_dim={grid_dim!r}." + ) + if coords: + return uxarray_obj.assign_coords(coords) + return uxarray_obj + + +def _assert_grid_dim_coord_consistent_if_in_both(uxarray_obj, grid_dim, indexer): + """assert grid_dim's coordinate is consistent if in both uxarray_obj and indexer. + Otherwise, does nothing. + """ + if isinstance(indexer, xr.DataArray): + if grid_dim in uxarray_obj.coords and grid_dim in indexer.coords: + xr.core.coordinates.assert_coordinate_consistent( + uxarray_obj, {grid_dim: indexer.coords.variables[grid_dim]} + )