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 b464438 commit 2b5e936
Show file tree
Hide file tree
Showing 3 changed files with 54 additions and 12 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
41 changes: 34 additions & 7 deletions xarray/core/indexes.py
Original file line number Diff line number Diff line change
Expand Up @@ -1480,13 +1480,15 @@ def _id_index(self) -> dict[int, T_PandasOrXarrayIndex]:
@property
def _id_coord_names(self) -> dict[int, tuple[Hashable, ...]]:
if self.__id_coord_names is None:
id_coord_names: Mapping[int, list[Hashable]] = defaultdict(list)
for k, v in self._coord_name_id.items():
id_coord_names[v].append(k)
self.__id_coord_names = {k: tuple(v) for k, v in id_coord_names.items()}

self._create_id_coord_names()
return self.__id_coord_names

def _create_id_coord_names(self) -> None:
id_coord_names: Mapping[int, list[Hashable]] = defaultdict(list)
for k, v in self._coord_name_id.items():
id_coord_names[v].append(k)
self.__id_coord_names = {k: tuple(v) for k, v in id_coord_names.items()}

@property
def variables(self) -> Mapping[Hashable, Variable]:
return Frozen(self._variables)
Expand Down Expand Up @@ -1772,6 +1774,30 @@ def check_variables():
return not not_equal


def _apply_indexes_isel(
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 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


def _apply_indexes(
indexes: Indexes[Index],
args: Mapping[Any, Any],
Expand All @@ -1797,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(indexes, indexers, "isel")
return _apply_indexes_isel(indexes, coords, indexers)


def roll_indexes(
Expand Down

0 comments on commit 2b5e936

Please sign in to comment.