Skip to content

Commit

Permalink
12k->15k
Browse files Browse the repository at this point in the history
  • Loading branch information
hmaarrfk committed May 5, 2024
1 parent 592e2bd commit 8c49aba
Show file tree
Hide file tree
Showing 3 changed files with 33 additions and 14 deletions.
6 changes: 5 additions & 1 deletion xarray/core/dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down
19 changes: 15 additions & 4 deletions xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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():
Expand Down
22 changes: 13 additions & 9 deletions xarray/core/indexes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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(
Expand Down

0 comments on commit 8c49aba

Please sign in to comment.