from __future__ import annotations import collections.abc import copy from collections import defaultdict from typing import ( TYPE_CHECKING, Any, Dict, Generic, Hashable, Iterable, Iterator, Mapping, Sequence, TypeVar, cast, ) import numpy as np import pandas as pd from xarray.core import formatting, nputils, utils from xarray.core.indexing import ( IndexSelResult, PandasIndexingAdapter, PandasMultiIndexingAdapter, ) from xarray.core.utils import Frozen, get_valid_numpy_dtype, is_dict_like, is_scalar if TYPE_CHECKING: from xarray.core.types import ErrorOptions, T_Index from xarray.core.variable import Variable IndexVars = Dict[Any, "Variable"] class Index: """Base class inherited by all xarray-compatible indexes. Do not use this class directly for creating index objects. """ @classmethod def from_variables( cls, variables: Mapping[Any, Variable], *, options: Mapping[str, Any], ) -> Index: raise NotImplementedError() @classmethod def concat( cls: type[T_Index], indexes: Sequence[T_Index], dim: Hashable, positions: Iterable[Iterable[int]] | None = None, ) -> T_Index: raise NotImplementedError() @classmethod def stack(cls, variables: Mapping[Any, Variable], dim: Hashable) -> Index: raise NotImplementedError( f"{cls!r} cannot be used for creating an index of stacked coordinates" ) def unstack(self) -> tuple[dict[Hashable, Index], pd.MultiIndex]: raise NotImplementedError() def create_variables( self, variables: Mapping[Any, Variable] | None = None ) -> IndexVars: if variables is not None: # pass through return dict(**variables) else: return {} def to_pandas_index(self) -> pd.Index: """Cast this xarray index to a pandas.Index object or raise a TypeError if this is not supported. This method is used by all xarray operations that expect/require a pandas.Index object. """ raise TypeError(f"{self!r} cannot be cast to a pandas.Index object") def isel( self, indexers: Mapping[Any, int | slice | np.ndarray | Variable] ) -> Index | None: return None def sel(self, labels: dict[Any, Any]) -> IndexSelResult: raise NotImplementedError(f"{self!r} doesn't support label-based selection") def join(self: T_Index, other: T_Index, how: str = "inner") -> T_Index: raise NotImplementedError( f"{self!r} doesn't support alignment with inner/outer join method" ) def reindex_like(self: T_Index, other: T_Index) -> dict[Hashable, Any]: raise NotImplementedError(f"{self!r} doesn't support re-indexing labels") def equals(self, other): # pragma: no cover raise NotImplementedError() def roll(self, shifts: Mapping[Any, int]) -> Index | None: return None def rename( self, name_dict: Mapping[Any, Hashable], dims_dict: Mapping[Any, Hashable] ) -> Index: return self def __copy__(self) -> Index: return self._copy(deep=False) def __deepcopy__(self, memo: dict[int, Any] | None = None) -> Index: return self._copy(deep=True, memo=memo) def copy(self: T_Index, deep: bool = True) -> T_Index: return self._copy(deep=deep) def _copy( self: T_Index, deep: bool = True, memo: dict[int, Any] | None = None ) -> T_Index: cls = self.__class__ copied = cls.__new__(cls) if deep: for k, v in self.__dict__.items(): setattr(copied, k, copy.deepcopy(v, memo)) else: copied.__dict__.update(self.__dict__) return copied def __getitem__(self, indexer: Any): raise NotImplementedError() def _repr_inline_(self, max_width): return self.__class__.__name__ def _maybe_cast_to_cftimeindex(index: pd.Index) -> pd.Index: from xarray.coding.cftimeindex import CFTimeIndex if len(index) > 0 and index.dtype == "O": try: return CFTimeIndex(index) except (ImportError, TypeError): return index else: return index def safe_cast_to_index(array: Any) -> pd.Index: """Given an array, safely cast it to a pandas.Index. If it is already a pandas.Index, return it unchanged. Unlike pandas.Index, if the array has dtype=object or dtype=timedelta64, this function will not attempt to do automatic type conversion but will always return an index with dtype=object. """ from xarray.core.dataarray import DataArray from xarray.core.variable import Variable if isinstance(array, pd.Index): index = array elif isinstance(array, (DataArray, Variable)): # returns the original multi-index for pandas.MultiIndex level coordinates index = array._to_index() elif isinstance(array, Index): index = array.to_pandas_index() elif isinstance(array, PandasIndexingAdapter): index = array.array else: kwargs = {} if hasattr(array, "dtype") and array.dtype.kind == "O": kwargs["dtype"] = object index = pd.Index(np.asarray(array), **kwargs) return _maybe_cast_to_cftimeindex(index) def _sanitize_slice_element(x): from xarray.core.dataarray import DataArray from xarray.core.variable import Variable if not isinstance(x, tuple) and len(np.shape(x)) != 0: raise ValueError( f"cannot use non-scalar arrays in a slice for xarray indexing: {x}" ) if isinstance(x, (Variable, DataArray)): x = x.values if isinstance(x, np.ndarray): x = x[()] return x def _query_slice(index, label, coord_name="", method=None, tolerance=None): if method is not None or tolerance is not None: raise NotImplementedError( "cannot use ``method`` argument if any indexers are slice objects" ) indexer = index.slice_indexer( _sanitize_slice_element(label.start), _sanitize_slice_element(label.stop), _sanitize_slice_element(label.step), ) if not isinstance(indexer, slice): # unlike pandas, in xarray we never want to silently convert a # slice indexer into an array indexer raise KeyError( "cannot represent labeled-based slice indexer for coordinate " f"{coord_name!r} with a slice over integer positions; the index is " "unsorted or non-unique" ) return indexer def _asarray_tuplesafe(values): """ Convert values into a numpy array of at most 1-dimension, while preserving tuples. Adapted from pandas.core.common._asarray_tuplesafe """ if isinstance(values, tuple): result = utils.to_0d_object_array(values) else: result = np.asarray(values) if result.ndim == 2: result = np.empty(len(values), dtype=object) result[:] = values return result def _is_nested_tuple(possible_tuple): return isinstance(possible_tuple, tuple) and any( isinstance(value, (tuple, list, slice)) for value in possible_tuple ) def normalize_label(value, dtype=None) -> np.ndarray: if getattr(value, "ndim", 1) <= 1: value = _asarray_tuplesafe(value) if dtype is not None and dtype.kind == "f" and value.dtype.kind != "b": # pd.Index built from coordinate with float precision != 64 # see https://github.com/pydata/xarray/pull/3153 for details # bypass coercing dtype for boolean indexers (ignore index) # see https://github.com/pydata/xarray/issues/5727 value = np.asarray(value, dtype=dtype) return value def as_scalar(value: np.ndarray): # see https://github.com/pydata/xarray/pull/4292 for details return value[()] if value.dtype.kind in "mM" else value.item() def get_indexer_nd(index, labels, method=None, tolerance=None): """Wrapper around :meth:`pandas.Index.get_indexer` supporting n-dimensional labels """ flat_labels = np.ravel(labels) flat_indexer = index.get_indexer(flat_labels, method=method, tolerance=tolerance) indexer = flat_indexer.reshape(labels.shape) return indexer T_PandasIndex = TypeVar("T_PandasIndex", bound="PandasIndex") class PandasIndex(Index): """Wrap a pandas.Index as an xarray compatible index.""" index: pd.Index dim: Hashable coord_dtype: Any __slots__ = ("index", "dim", "coord_dtype") def __init__(self, array: Any, dim: Hashable, coord_dtype: Any = None): # make a shallow copy: cheap and because the index name may be updated # here or in other constructors (cannot use pd.Index.rename as this # constructor is also called from PandasMultiIndex) index = safe_cast_to_index(array).copy() if index.name is None: index.name = dim self.index = index self.dim = dim if coord_dtype is None: coord_dtype = get_valid_numpy_dtype(index) self.coord_dtype = coord_dtype def _replace(self, index, dim=None, coord_dtype=None): if dim is None: dim = self.dim if coord_dtype is None: coord_dtype = self.coord_dtype return type(self)(index, dim, coord_dtype) @classmethod def from_variables( cls, variables: Mapping[Any, Variable], *, options: Mapping[str, Any], ) -> PandasIndex: if len(variables) != 1: raise ValueError( f"PandasIndex only accepts one variable, found {len(variables)} variables" ) name, var = next(iter(variables.items())) if var.ndim != 1: raise ValueError( "PandasIndex only accepts a 1-dimensional variable, " f"variable {name!r} has {var.ndim} dimensions" ) dim = var.dims[0] # TODO: (benbovy - explicit indexes): add __index__ to ExplicitlyIndexesNDArrayMixin? # this could be eventually used by Variable.to_index() and would remove the need to perform # the checks below. # preserve wrapped pd.Index (if any) data = getattr(var._data, "array", var.data) # multi-index level variable: get level index if isinstance(var._data, PandasMultiIndexingAdapter): level = var._data.level if level is not None: data = var._data.array.get_level_values(level) obj = cls(data, dim, coord_dtype=var.dtype) assert not isinstance(obj.index, pd.MultiIndex) obj.index.name = name return obj @staticmethod def _concat_indexes(indexes, dim, positions=None) -> pd.Index: new_pd_index: pd.Index if not indexes: new_pd_index = pd.Index([]) else: if not all(idx.dim == dim for idx in indexes): dims = ",".join({f"{idx.dim!r}" for idx in indexes}) raise ValueError( f"Cannot concatenate along dimension {dim!r} indexes with " f"dimensions: {dims}" ) pd_indexes = [idx.index for idx in indexes] new_pd_index = pd_indexes[0].append(pd_indexes[1:]) if positions is not None: indices = nputils.inverse_permutation(np.concatenate(positions)) new_pd_index = new_pd_index.take(indices) return new_pd_index @classmethod def concat( cls, indexes: Sequence[PandasIndex], dim: Hashable, positions: Iterable[Iterable[int]] | None = None, ) -> PandasIndex: new_pd_index = cls._concat_indexes(indexes, dim, positions) if not indexes: coord_dtype = None else: coord_dtype = np.result_type(*[idx.coord_dtype for idx in indexes]) return cls(new_pd_index, dim=dim, coord_dtype=coord_dtype) def create_variables( self, variables: Mapping[Any, Variable] | None = None ) -> IndexVars: from xarray.core.variable import IndexVariable name = self.index.name attrs: Mapping[Hashable, Any] | None encoding: Mapping[Hashable, Any] | None if variables is not None and name in variables: var = variables[name] attrs = var.attrs encoding = var.encoding else: attrs = None encoding = None data = PandasIndexingAdapter(self.index, dtype=self.coord_dtype) var = IndexVariable(self.dim, data, attrs=attrs, encoding=encoding) return {name: var} def to_pandas_index(self) -> pd.Index: return self.index def isel( self, indexers: Mapping[Any, int | slice | np.ndarray | Variable] ) -> PandasIndex | None: from xarray.core.variable import Variable indxr = indexers[self.dim] if isinstance(indxr, Variable): if indxr.dims != (self.dim,): # can't preserve a index if result has new dimensions return None else: indxr = indxr.data if not isinstance(indxr, slice) and is_scalar(indxr): # scalar indexer: drop index return None return self._replace(self.index[indxr]) def sel( self, labels: dict[Any, Any], method=None, tolerance=None ) -> IndexSelResult: from xarray.core.dataarray import DataArray from xarray.core.variable import Variable if method is not None and not isinstance(method, str): raise TypeError("``method`` must be a string") assert len(labels) == 1 coord_name, label = next(iter(labels.items())) if isinstance(label, slice): indexer = _query_slice(self.index, label, coord_name, method, tolerance) elif is_dict_like(label): raise ValueError( "cannot use a dict-like object for selection on " "a dimension that does not have a MultiIndex" ) else: label_array = normalize_label(label, dtype=self.coord_dtype) if label_array.ndim == 0: label_value = as_scalar(label_array) if isinstance(self.index, pd.CategoricalIndex): if method is not None: raise ValueError( "'method' is not supported when indexing using a CategoricalIndex." ) if tolerance is not None: raise ValueError( "'tolerance' is not supported when indexing using a CategoricalIndex." ) indexer = self.index.get_loc(label_value) else: if method is not None: indexer = get_indexer_nd( self.index, label_array, method, tolerance ) if np.any(indexer < 0): raise KeyError( f"not all values found in index {coord_name!r}" ) else: try: indexer = self.index.get_loc(label_value) except KeyError as e: raise KeyError( f"not all values found in index {coord_name!r}. " "Try setting the `method` keyword argument (example: method='nearest')." ) from e elif label_array.dtype.kind == "b": indexer = label_array else: indexer = get_indexer_nd(self.index, label_array, method, tolerance) if np.any(indexer < 0): raise KeyError(f"not all values found in index {coord_name!r}") # attach dimension names and/or coordinates to positional indexer if isinstance(label, Variable): indexer = Variable(label.dims, indexer) elif isinstance(label, DataArray): indexer = DataArray(indexer, coords=label._coords, dims=label.dims) return IndexSelResult({self.dim: indexer}) def equals(self, other: Index): if not isinstance(other, PandasIndex): return False return self.index.equals(other.index) and self.dim == other.dim def join(self: PandasIndex, other: PandasIndex, how: str = "inner") -> PandasIndex: if how == "outer": index = self.index.union(other.index) else: # how = "inner" index = self.index.intersection(other.index) coord_dtype = np.result_type(self.coord_dtype, other.coord_dtype) return type(self)(index, self.dim, coord_dtype=coord_dtype) def reindex_like( self, other: PandasIndex, method=None, tolerance=None ) -> dict[Hashable, Any]: if not self.index.is_unique: raise ValueError( f"cannot reindex or align along dimension {self.dim!r} because the " "(pandas) index has duplicate values" ) return {self.dim: get_indexer_nd(self.index, other.index, method, tolerance)} def roll(self, shifts: Mapping[Any, int]) -> PandasIndex: shift = shifts[self.dim] % self.index.shape[0] if shift != 0: new_pd_idx = self.index[-shift:].append(self.index[:-shift]) else: new_pd_idx = self.index[:] return self._replace(new_pd_idx) def rename(self, name_dict, dims_dict): if self.index.name not in name_dict and self.dim not in dims_dict: return self new_name = name_dict.get(self.index.name, self.index.name) index = self.index.rename(new_name) new_dim = dims_dict.get(self.dim, self.dim) return self._replace(index, dim=new_dim) def _copy( self: T_PandasIndex, deep: bool = True, memo: dict[int, Any] | None = None ) -> T_PandasIndex: if deep: # pandas is not using the memo index = self.index.copy(deep=True) else: # index will be copied in constructor index = self.index return self._replace(index) def __getitem__(self, indexer: Any): return self._replace(self.index[indexer]) def __repr__(self): return f"PandasIndex({repr(self.index)})" def _check_dim_compat(variables: Mapping[Any, Variable], all_dims: str = "equal"): """Check that all multi-index variable candidates are 1-dimensional and either share the same (single) dimension or each have a different dimension. """ if any([var.ndim != 1 for var in variables.values()]): raise ValueError("PandasMultiIndex only accepts 1-dimensional variables") dims = {var.dims for var in variables.values()} if all_dims == "equal" and len(dims) > 1: raise ValueError( "unmatched dimensions for multi-index variables " + ", ".join([f"{k!r} {v.dims}" for k, v in variables.items()]) ) if all_dims == "different" and len(dims) < len(variables): raise ValueError( "conflicting dimensions for multi-index product variables " + ", ".join([f"{k!r} {v.dims}" for k, v in variables.items()]) ) def remove_unused_levels_categories(index: pd.Index) -> pd.Index: """ Remove unused levels from MultiIndex and unused categories from CategoricalIndex """ if isinstance(index, pd.MultiIndex): index = index.remove_unused_levels() # if it contains CategoricalIndex, we need to remove unused categories # manually. See https://github.com/pandas-dev/pandas/issues/30846 if any(isinstance(lev, pd.CategoricalIndex) for lev in index.levels): levels = [] for i, level in enumerate(index.levels): if isinstance(level, pd.CategoricalIndex): level = level[index.codes[i]].remove_unused_categories() else: level = level[index.codes[i]] levels.append(level) # TODO: calling from_array() reorders MultiIndex levels. It would # be best to avoid this, if possible, e.g., by using # MultiIndex.remove_unused_levels() (which does not reorder) on the # part of the MultiIndex that is not categorical, or by fixing this # upstream in pandas. index = pd.MultiIndex.from_arrays(levels, names=index.names) elif isinstance(index, pd.CategoricalIndex): index = index.remove_unused_categories() return index class PandasMultiIndex(PandasIndex): """Wrap a pandas.MultiIndex as an xarray compatible index.""" level_coords_dtype: dict[str, Any] __slots__ = ("index", "dim", "coord_dtype", "level_coords_dtype") def __init__(self, array: Any, dim: Hashable, level_coords_dtype: Any = None): super().__init__(array, dim) # default index level names names = [] for i, idx in enumerate(self.index.levels): name = idx.name or f"{dim}_level_{i}" if name == dim: raise ValueError( f"conflicting multi-index level name {name!r} with dimension {dim!r}" ) names.append(name) self.index.names = names if level_coords_dtype is None: level_coords_dtype = { idx.name: get_valid_numpy_dtype(idx) for idx in self.index.levels } self.level_coords_dtype = level_coords_dtype def _replace(self, index, dim=None, level_coords_dtype=None) -> PandasMultiIndex: if dim is None: dim = self.dim index.name = dim if level_coords_dtype is None: level_coords_dtype = self.level_coords_dtype return type(self)(index, dim, level_coords_dtype) @classmethod def from_variables( cls, variables: Mapping[Any, Variable], *, options: Mapping[str, Any], ) -> PandasMultiIndex: _check_dim_compat(variables) dim = next(iter(variables.values())).dims[0] index = pd.MultiIndex.from_arrays( [var.values for var in variables.values()], names=variables.keys() ) index.name = dim level_coords_dtype = {name: var.dtype for name, var in variables.items()} obj = cls(index, dim, level_coords_dtype=level_coords_dtype) return obj @classmethod def concat( # type: ignore[override] cls, indexes: Sequence[PandasMultiIndex], dim: Hashable, positions: Iterable[Iterable[int]] | None = None, ) -> PandasMultiIndex: new_pd_index = cls._concat_indexes(indexes, dim, positions) if not indexes: level_coords_dtype = None else: level_coords_dtype = {} for name in indexes[0].level_coords_dtype: level_coords_dtype[name] = np.result_type( *[idx.level_coords_dtype[name] for idx in indexes] ) return cls(new_pd_index, dim=dim, level_coords_dtype=level_coords_dtype) @classmethod def stack( cls, variables: Mapping[Any, Variable], dim: Hashable ) -> PandasMultiIndex: """Create a new Pandas MultiIndex from the product of 1-d variables (levels) along a new dimension. Level variables must have a dimension distinct from each other. Keeps levels the same (doesn't refactorize them) so that it gives back the original labels after a stack/unstack roundtrip. """ _check_dim_compat(variables, all_dims="different") level_indexes = [safe_cast_to_index(var) for var in variables.values()] for name, idx in zip(variables, level_indexes): if isinstance(idx, pd.MultiIndex): raise ValueError( f"cannot create a multi-index along stacked dimension {dim!r} " f"from variable {name!r} that wraps a multi-index" ) split_labels, levels = zip(*[lev.factorize() for lev in level_indexes]) labels_mesh = np.meshgrid(*split_labels, indexing="ij") labels = [x.ravel() for x in labels_mesh] index = pd.MultiIndex(levels, labels, sortorder=0, names=variables.keys()) level_coords_dtype = {k: var.dtype for k, var in variables.items()} return cls(index, dim, level_coords_dtype=level_coords_dtype) def unstack(self) -> tuple[dict[Hashable, Index], pd.MultiIndex]: clean_index = remove_unused_levels_categories(self.index) new_indexes: dict[Hashable, Index] = {} for name, lev in zip(clean_index.names, clean_index.levels): idx = PandasIndex( lev.copy(), name, coord_dtype=self.level_coords_dtype[name] ) new_indexes[name] = idx return new_indexes, clean_index @classmethod def from_variables_maybe_expand( cls, dim: Hashable, current_variables: Mapping[Any, Variable], variables: Mapping[Any, Variable], ) -> tuple[PandasMultiIndex, IndexVars]: """Create a new multi-index maybe by expanding an existing one with new variables as index levels. The index and its corresponding coordinates may be created along a new dimension. """ names: list[Hashable] = [] codes: list[list[int]] = [] levels: list[list[int]] = [] level_variables: dict[Any, Variable] = {} _check_dim_compat({**current_variables, **variables}) if len(current_variables) > 1: # expand from an existing multi-index data = cast( PandasMultiIndexingAdapter, next(iter(current_variables.values()))._data ) current_index = data.array names.extend(current_index.names) codes.extend(current_index.codes) levels.extend(current_index.levels) for name in current_index.names: level_variables[name] = current_variables[name] elif len(current_variables) == 1: # expand from one 1D variable (no multi-index): convert it to an index level var = next(iter(current_variables.values())) new_var_name = f"{dim}_level_0" names.append(new_var_name) cat = pd.Categorical(var.values, ordered=True) codes.append(cat.codes) levels.append(cat.categories) level_variables[new_var_name] = var for name, var in variables.items(): names.append(name) cat = pd.Categorical(var.values, ordered=True) codes.append(cat.codes) levels.append(cat.categories) level_variables[name] = var index = pd.MultiIndex(levels, codes, names=names) level_coords_dtype = {k: var.dtype for k, var in level_variables.items()} obj = cls(index, dim, level_coords_dtype=level_coords_dtype) index_vars = obj.create_variables(level_variables) return obj, index_vars def keep_levels( self, level_variables: Mapping[Any, Variable] ) -> PandasMultiIndex | PandasIndex: """Keep only the provided levels and return a new multi-index with its corresponding coordinates. """ index = self.index.droplevel( [k for k in self.index.names if k not in level_variables] ) if isinstance(index, pd.MultiIndex): level_coords_dtype = {k: self.level_coords_dtype[k] for k in index.names} return self._replace(index, level_coords_dtype=level_coords_dtype) else: # backward compatibility: rename the level coordinate to the dimension name return PandasIndex( index.rename(self.dim), self.dim, coord_dtype=self.level_coords_dtype[index.name], ) def reorder_levels( self, level_variables: Mapping[Any, Variable] ) -> PandasMultiIndex: """Re-arrange index levels using input order and return a new multi-index with its corresponding coordinates. """ index = self.index.reorder_levels(level_variables.keys()) level_coords_dtype = {k: self.level_coords_dtype[k] for k in index.names} return self._replace(index, level_coords_dtype=level_coords_dtype) def create_variables( self, variables: Mapping[Any, Variable] | None = None ) -> IndexVars: from xarray.core.variable import IndexVariable if variables is None: variables = {} index_vars: IndexVars = {} for name in (self.dim,) + self.index.names: if name == self.dim: level = None dtype = None else: level = name dtype = self.level_coords_dtype[name] var = variables.get(name, None) if var is not None: attrs = var.attrs encoding = var.encoding else: attrs = {} encoding = {} data = PandasMultiIndexingAdapter(self.index, dtype=dtype, level=level) index_vars[name] = IndexVariable( self.dim, data, attrs=attrs, encoding=encoding, fastpath=True, ) return index_vars def sel(self, labels, method=None, tolerance=None) -> IndexSelResult: from xarray.core.dataarray import DataArray from xarray.core.variable import Variable if method is not None or tolerance is not None: raise ValueError( "multi-index does not support ``method`` and ``tolerance``" ) new_index = None scalar_coord_values = {} # label(s) given for multi-index level(s) if all([lbl in self.index.names for lbl in labels]): label_values = {} for k, v in labels.items(): label_array = normalize_label(v, dtype=self.level_coords_dtype[k]) try: label_values[k] = as_scalar(label_array) except ValueError: # label should be an item not an array-like raise ValueError( "Vectorized selection is not " f"available along coordinate {k!r} (multi-index level)" ) has_slice = any([isinstance(v, slice) for v in label_values.values()]) if len(label_values) == self.index.nlevels and not has_slice: indexer = self.index.get_loc( tuple(label_values[k] for k in self.index.names) ) else: indexer, new_index = self.index.get_loc_level( tuple(label_values.values()), level=tuple(label_values.keys()) ) scalar_coord_values.update(label_values) # GH2619. Raise a KeyError if nothing is chosen if indexer.dtype.kind == "b" and indexer.sum() == 0: raise KeyError(f"{labels} not found") # assume one label value given for the multi-index "array" (dimension) else: if len(labels) > 1: coord_name = next(iter(set(labels) - set(self.index.names))) raise ValueError( f"cannot provide labels for both coordinate {coord_name!r} (multi-index array) " f"and one or more coordinates among {self.index.names!r} (multi-index levels)" ) coord_name, label = next(iter(labels.items())) if is_dict_like(label): invalid_levels = [ name for name in label if name not in self.index.names ] if invalid_levels: raise ValueError( f"invalid multi-index level names {invalid_levels}" ) return self.sel(label) elif isinstance(label, slice): indexer = _query_slice(self.index, label, coord_name) elif isinstance(label, tuple): if _is_nested_tuple(label): indexer = self.index.get_locs(label) elif len(label) == self.index.nlevels: indexer = self.index.get_loc(label) else: levels = [self.index.names[i] for i in range(len(label))] indexer, new_index = self.index.get_loc_level(label, level=levels) scalar_coord_values.update({k: v for k, v in zip(levels, label)}) else: label_array = normalize_label(label) if label_array.ndim == 0: label_value = as_scalar(label_array) indexer, new_index = self.index.get_loc_level(label_value, level=0) scalar_coord_values[self.index.names[0]] = label_value elif label_array.dtype.kind == "b": indexer = label_array else: if label_array.ndim > 1: raise ValueError( "Vectorized selection is not available along " f"coordinate {coord_name!r} with a multi-index" ) indexer = get_indexer_nd(self.index, label_array) if np.any(indexer < 0): raise KeyError(f"not all values found in index {coord_name!r}") # attach dimension names and/or coordinates to positional indexer if isinstance(label, Variable): indexer = Variable(label.dims, indexer) elif isinstance(label, DataArray): # do not include label-indexer DataArray coordinates that conflict # with the level names of this index coords = { k: v for k, v in label._coords.items() if k not in self.index.names } indexer = DataArray(indexer, coords=coords, dims=label.dims) if new_index is not None: if isinstance(new_index, pd.MultiIndex): level_coords_dtype = { k: self.level_coords_dtype[k] for k in new_index.names } new_index = self._replace( new_index, level_coords_dtype=level_coords_dtype ) dims_dict = {} drop_coords = [] else: new_index = PandasIndex( new_index, new_index.name, coord_dtype=self.level_coords_dtype[new_index.name], ) dims_dict = {self.dim: new_index.index.name} drop_coords = [self.dim] # variable(s) attrs and encoding metadata are propagated # when replacing the indexes in the resulting xarray object new_vars = new_index.create_variables() indexes = cast(Dict[Any, Index], {k: new_index for k in new_vars}) # add scalar variable for each dropped level variables = new_vars for name, val in scalar_coord_values.items(): variables[name] = Variable([], val) return IndexSelResult( {self.dim: indexer}, indexes=indexes, variables=variables, drop_indexes=list(scalar_coord_values), drop_coords=drop_coords, rename_dims=dims_dict, ) else: return IndexSelResult({self.dim: indexer}) def join(self, other, how: str = "inner"): if how == "outer": # bug in pandas? need to reset index.name other_index = other.index.copy() other_index.name = None index = self.index.union(other_index) index.name = self.dim else: # how = "inner" index = self.index.intersection(other.index) level_coords_dtype = { k: np.result_type(lvl_dtype, other.level_coords_dtype[k]) for k, lvl_dtype in self.level_coords_dtype.items() } return type(self)(index, self.dim, level_coords_dtype=level_coords_dtype) def rename(self, name_dict, dims_dict): if not set(self.index.names) & set(name_dict) and self.dim not in dims_dict: return self # pandas 1.3.0: could simply do `self.index.rename(names_dict)` new_names = [name_dict.get(k, k) for k in self.index.names] index = self.index.rename(new_names) new_dim = dims_dict.get(self.dim, self.dim) new_level_coords_dtype = { k: v for k, v in zip(new_names, self.level_coords_dtype.values()) } return self._replace( index, dim=new_dim, level_coords_dtype=new_level_coords_dtype ) def create_default_index_implicit( dim_variable: Variable, all_variables: Mapping | Iterable[Hashable] | None = None, ) -> tuple[PandasIndex, IndexVars]: """Create a default index from a dimension variable. Create a PandasMultiIndex if the given variable wraps a pandas.MultiIndex, otherwise create a PandasIndex (note that this will become obsolete once we depreciate implicitly passing a pandas.MultiIndex as a coordinate). """ if all_variables is None: all_variables = {} if not isinstance(all_variables, Mapping): all_variables = {k: None for k in all_variables} name = dim_variable.dims[0] array = getattr(dim_variable._data, "array", None) index: PandasIndex if isinstance(array, pd.MultiIndex): index = PandasMultiIndex(array, name) index_vars = index.create_variables() # check for conflict between level names and variable names duplicate_names = [k for k in index_vars if k in all_variables and k != name] if duplicate_names: # dirty workaround for an edge case where both the dimension # coordinate and the level coordinates are given for the same # multi-index object => do not raise an error # TODO: remove this check when removing the multi-index dimension coordinate if len(duplicate_names) < len(index.index.names): conflict = True else: duplicate_vars = [all_variables[k] for k in duplicate_names] conflict = any( v is None or not dim_variable.equals(v) for v in duplicate_vars ) if conflict: conflict_str = "\n".join(duplicate_names) raise ValueError( f"conflicting MultiIndex level / variable name(s):\n{conflict_str}" ) else: dim_var = {name: dim_variable} index = PandasIndex.from_variables(dim_var, options={}) index_vars = index.create_variables(dim_var) return index, index_vars # generic type that represents either a pandas or an xarray index T_PandasOrXarrayIndex = TypeVar("T_PandasOrXarrayIndex", Index, pd.Index) class Indexes(collections.abc.Mapping, Generic[T_PandasOrXarrayIndex]): """Immutable proxy for Dataset or DataArrary indexes. Keys are coordinate names and values may correspond to either pandas or xarray indexes. Also provides some utility methods. """ _indexes: dict[Any, T_PandasOrXarrayIndex] _variables: dict[Any, Variable] __slots__ = ( "_indexes", "_variables", "_dims", "__coord_name_id", "__id_index", "__id_coord_names", ) def __init__( self, indexes: dict[Any, T_PandasOrXarrayIndex], variables: dict[Any, Variable], ): """Constructor not for public consumption. Parameters ---------- indexes : dict Indexes held by this object. variables : dict Indexed coordinate variables in this object. """ self._indexes = indexes self._variables = variables self._dims: Mapping[Hashable, int] | None = None self.__coord_name_id: dict[Any, int] | None = None self.__id_index: dict[int, T_PandasOrXarrayIndex] | None = None self.__id_coord_names: dict[int, tuple[Hashable, ...]] | None = None @property def _coord_name_id(self) -> dict[Any, int]: if self.__coord_name_id is None: self.__coord_name_id = {k: id(idx) for k, idx in self._indexes.items()} return self.__coord_name_id @property def _id_index(self) -> dict[int, T_PandasOrXarrayIndex]: if self.__id_index is None: self.__id_index = {id(idx): idx for idx in self.get_unique()} return self.__id_index @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()} return self.__id_coord_names @property def variables(self) -> Mapping[Hashable, Variable]: return Frozen(self._variables) @property def dims(self) -> Mapping[Hashable, int]: from xarray.core.variable import calculate_dimensions if self._dims is None: self._dims = calculate_dimensions(self._variables) return Frozen(self._dims) def copy(self) -> Indexes: return type(self)(dict(self._indexes), dict(self._variables)) def get_unique(self) -> list[T_PandasOrXarrayIndex]: """Return a list of unique indexes, preserving order.""" unique_indexes: list[T_PandasOrXarrayIndex] = [] seen: set[int] = set() for index in self._indexes.values(): index_id = id(index) if index_id not in seen: unique_indexes.append(index) seen.add(index_id) return unique_indexes def is_multi(self, key: Hashable) -> bool: """Return True if ``key`` maps to a multi-coordinate index, False otherwise. """ return len(self._id_coord_names[self._coord_name_id[key]]) > 1 def get_all_coords( self, key: Hashable, errors: ErrorOptions = "raise" ) -> dict[Hashable, Variable]: """Return all coordinates having the same index. Parameters ---------- key : hashable Index key. errors : {"raise", "ignore"}, default: "raise" If "raise", raises a ValueError if `key` is not in indexes. If "ignore", an empty tuple is returned instead. Returns ------- coords : dict A dictionary of all coordinate variables having the same index. """ if errors not in ["raise", "ignore"]: raise ValueError('errors must be either "raise" or "ignore"') if key not in self._indexes: if errors == "raise": raise ValueError(f"no index found for {key!r} coordinate") else: return {} all_coord_names = self._id_coord_names[self._coord_name_id[key]] return {k: self._variables[k] for k in all_coord_names} def get_all_dims( self, key: Hashable, errors: ErrorOptions = "raise" ) -> Mapping[Hashable, int]: """Return all dimensions shared by an index. Parameters ---------- key : hashable Index key. errors : {"raise", "ignore"}, default: "raise" If "raise", raises a ValueError if `key` is not in indexes. If "ignore", an empty tuple is returned instead. Returns ------- dims : dict A dictionary of all dimensions shared by an index. """ from xarray.core.variable import calculate_dimensions return calculate_dimensions(self.get_all_coords(key, errors=errors)) def group_by_index( self, ) -> list[tuple[T_PandasOrXarrayIndex, dict[Hashable, Variable]]]: """Returns a list of unique indexes and their corresponding coordinates.""" index_coords = [] for i in self._id_index: index = self._id_index[i] coords = {k: self._variables[k] for k in self._id_coord_names[i]} index_coords.append((index, coords)) return index_coords def to_pandas_indexes(self) -> Indexes[pd.Index]: """Returns an immutable proxy for Dataset or DataArrary pandas indexes. Raises an error if this proxy contains indexes that cannot be coerced to pandas.Index objects. """ indexes: dict[Hashable, pd.Index] = {} for k, idx in self._indexes.items(): if isinstance(idx, pd.Index): indexes[k] = idx elif isinstance(idx, Index): indexes[k] = idx.to_pandas_index() return Indexes(indexes, self._variables) def copy_indexes( self, deep: bool = True, memo: dict[int, Any] | None = None ) -> tuple[dict[Hashable, T_PandasOrXarrayIndex], dict[Hashable, Variable]]: """Return a new dictionary with copies of indexes, preserving unique indexes. Parameters ---------- deep : bool, default: True Whether the indexes are deep or shallow copied onto the new object. memo : dict if object id to copied objects or None, optional To prevent infinite recursion deepcopy stores all copied elements in this dict. """ new_indexes = {} new_index_vars = {} for idx, coords in self.group_by_index(): if isinstance(idx, pd.Index): convert_new_idx = True dim = next(iter(coords.values())).dims[0] if isinstance(idx, pd.MultiIndex): idx = PandasMultiIndex(idx, dim) else: idx = PandasIndex(idx, dim) else: convert_new_idx = False new_idx = idx._copy(deep=deep, memo=memo) idx_vars = idx.create_variables(coords) if convert_new_idx: new_idx = cast(PandasIndex, new_idx).index new_indexes.update({k: new_idx for k in coords}) new_index_vars.update(idx_vars) return new_indexes, new_index_vars def __iter__(self) -> Iterator[T_PandasOrXarrayIndex]: return iter(self._indexes) def __len__(self) -> int: return len(self._indexes) def __contains__(self, key) -> bool: return key in self._indexes def __getitem__(self, key) -> T_PandasOrXarrayIndex: return self._indexes[key] def __repr__(self): return formatting.indexes_repr(self) def default_indexes( coords: Mapping[Any, Variable], dims: Iterable ) -> dict[Hashable, Index]: """Default indexes for a Dataset/DataArray. Parameters ---------- coords : Mapping[Any, xarray.Variable] Coordinate variables from which to draw default indexes. dims : iterable Iterable of dimension names. Returns ------- Mapping from indexing keys (levels/dimension names) to indexes used for indexing along that dimension. """ indexes: dict[Hashable, Index] = {} coord_names = set(coords) for name, var in coords.items(): if name in dims: index, index_vars = create_default_index_implicit(var, coords) if set(index_vars) <= coord_names: indexes.update({k: index for k in index_vars}) return indexes def indexes_equal( index: Index, other_index: Index, variable: Variable, other_variable: Variable, cache: dict[tuple[int, int], bool | None] | None = None, ) -> bool: """Check if two indexes are equal, possibly with cached results. If the two indexes are not of the same type or they do not implement equality, fallback to coordinate labels equality check. """ if cache is None: # dummy cache cache = {} key = (id(index), id(other_index)) equal: bool | None = None if key not in cache: if type(index) is type(other_index): try: equal = index.equals(other_index) except NotImplementedError: equal = None else: cache[key] = equal else: equal = None else: equal = cache[key] if equal is None: equal = variable.equals(other_variable) return cast(bool, equal) def indexes_all_equal( elements: Sequence[tuple[Index, dict[Hashable, Variable]]] ) -> bool: """Check if indexes are all equal. If they are not of the same type or they do not implement this check, check if their coordinate variables are all equal instead. """ def check_variables(): variables = [e[1] for e in elements] return any( not variables[0][k].equals(other_vars[k]) for other_vars in variables[1:] for k in variables[0] ) indexes = [e[0] for e in elements] same_objects = all(indexes[0] is other_idx for other_idx in indexes[1:]) if same_objects: return True same_type = all(type(indexes[0]) is type(other_idx) for other_idx in indexes[1:]) if same_type: try: not_equal = any( not indexes[0].equals(other_idx) for other_idx in indexes[1:] ) except NotImplementedError: not_equal = check_variables() else: not_equal = check_variables() return not not_equal def _apply_indexes( indexes: Indexes[Index], args: Mapping[Any, Any], func: str, ) -> tuple[dict[Hashable, Index], dict[Hashable, Variable]]: 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(): index_dims = {d for var in index_vars.values() for d in var.dims} index_args = {k: v for k, v in args.items() if k in index_dims} if index_args: new_index = getattr(index, func)(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) else: for k in index_vars: new_indexes.pop(k, None) return new_indexes, new_index_variables def isel_indexes( indexes: Indexes[Index], indexers: Mapping[Any, Any], ) -> tuple[dict[Hashable, Index], dict[Hashable, Variable]]: return _apply_indexes(indexes, indexers, "isel") def roll_indexes( indexes: Indexes[Index], shifts: Mapping[Any, int], ) -> tuple[dict[Hashable, Index], dict[Hashable, Variable]]: return _apply_indexes(indexes, shifts, "roll") def filter_indexes_from_coords( indexes: Mapping[Any, Index], filtered_coord_names: set, ) -> dict[Hashable, Index]: """Filter index items given a (sub)set of coordinate names. Drop all multi-coordinate related index items for any key missing in the set of coordinate names. """ filtered_indexes: dict[Any, Index] = dict(**indexes) index_coord_names: dict[Hashable, set[Hashable]] = defaultdict(set) for name, idx in indexes.items(): index_coord_names[id(idx)].add(name) for idx_coord_names in index_coord_names.values(): if not idx_coord_names <= filtered_coord_names: for k in idx_coord_names: del filtered_indexes[k] return filtered_indexes def assert_no_index_corrupted( indexes: Indexes[Index], coord_names: set[Hashable], action: str = "remove coordinate(s)", ) -> None: """Assert removing coordinates or indexes will not corrupt indexes.""" # An index may be corrupted when the set of its corresponding coordinate name(s) # partially overlaps the set of coordinate names to remove for index, index_coords in indexes.group_by_index(): common_names = set(index_coords) & coord_names if common_names and len(common_names) != len(index_coords): common_names_str = ", ".join(f"{k!r}" for k in common_names) index_names_str = ", ".join(f"{k!r}" for k in index_coords) raise ValueError( f"cannot {action} {common_names_str}, which would corrupt " f"the following index built from coordinates {index_names_str}:\n" f"{index}" )