12k->15k

pydata · May 5, 2024 · 8c49aba · 8c49aba
1 parent 592e2bd
commit 8c49aba
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 14 deletions.
diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py
@@ -1503,7 +1503,11 @@ def isel(
         # lists, or zero or one-dimensional np.ndarray's
 
         variable = self._variable.isel(indexers, missing_dims=missing_dims)
-        indexes, index_variables = isel_indexes(self.xindexes, indexers)
+        indexes, index_variables = isel_indexes(
+            self._indexes, 
+            self._coords,
+            indexers,
+        )
 
         coords = {}
         for coord_name, coord_value in self._coords.items():

diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
@@ -2979,8 +2979,11 @@ def isel(
         dims: dict[Hashable, int] = {}
         coord_names = self._coord_names.copy()
 
-        indexes, index_variables = isel_indexes(self.xindexes, indexers)
-
+        indexes, index_variables = isel_indexes(
+            self._indexes,
+            self.coords,
+            indexers,
+        )
         for name, var in self._variables.items():
             # preserve variable order
             if name in index_variables:
@@ -3015,7 +3018,11 @@ def _isel_fancy(
         valid_indexers = dict(self._validate_indexers(indexers, missing_dims))
 
         variables: dict[Hashable, Variable] = {}
-        indexes, index_variables = isel_indexes(self.xindexes, valid_indexers)
+        indexes, index_variables = isel_indexes(
+            self._indexes,
+            self.coords,
+            valid_indexers,
+        )
 
         for name, var in self.variables.items():
             if name in index_variables:
@@ -7828,7 +7835,11 @@ def diff(
         else:
             raise ValueError("The 'label' argument has to be either 'upper' or 'lower'")
 
-        indexes, index_vars = isel_indexes(self.xindexes, slice_new)
+        indexes, index_vars = isel_indexes(
+            self._indexes,
+            self.coords,
+            slice_new
+        )
         variables = {}
 
         for name, var in self.variables.items():

diff --git a/xarray/core/indexes.py b/xarray/core/indexes.py
@@ -1775,23 +1775,26 @@ def check_variables():
 
 
 def _apply_indexes_isel(
-    indexes: Indexes[Index],
+    indexes: dict[Hashable, Index],
+    coords,
     args: Mapping[Any, Any],
 ):
     new_indexes: dict[Hashable, Index] = {k: v for k, v in indexes.items()}
     new_index_variables: dict[Hashable, Variable] = {}
-    #     for index, index_vars in indexes.group_by_index():
-    for index in indexes._id_index.values():
-        # index = indexes._id_index[i]
-        index_dims = (index.index.name,)
-        index_args = {k: v for k, v in args.items() if k in index_dims}
+    for name, index in indexes.items():
+        index_args = {k: v for k, v in args.items() if k == name}
+        index_dims = (name,)
+        index_vars = {name: coords[name]}
         if index_args:
             new_index = index.isel(index_args)
             if new_index is not None:
                 new_indexes.update({k: new_index for k in index_vars})
                 new_index_vars = new_index.create_variables(index_vars)
                 new_index_variables.update(new_index_vars)
-
+                new_index_variables.update(new_index_vars)
+            else:
+                for k in index_vars:
+                    new_indexes.pop(k, None)
     return new_indexes, new_index_variables
 
 
@@ -1820,10 +1823,11 @@ def _apply_indexes(
 
 
 def isel_indexes(
-    indexes: Indexes[Index],
+    indexes,
+    coords,
     indexers: Mapping[Any, Any],
 ) -> tuple[dict[Hashable, Index], dict[Hashable, Variable]]:
-    return _apply_indexes_isel(indexes, indexers)
+    return _apply_indexes_isel(indexes, coords, indexers)
 
 
 def roll_indexes(