File: C:/Users/fred/anaconda3/Lib/site-packages/xarray/core/indexes.py
from __future__ import annotations
import collections.abc
import copy
from collections import defaultdict
from collections.abc import Hashable, Iterable, Iterator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Generic, 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,
emit_user_level_warning,
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" and not isinstance(index, CFTimeIndex):
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: dict[str, str] = {}
if hasattr(array, "dtype"):
if array.dtype.kind == "O":
kwargs["dtype"] = "object"
elif array.dtype == "float16":
emit_user_level_warning(
(
"`pandas.Index` does not support the `float16` dtype."
" Casting to `float64` for you, but in the future please"
" manually cast to either `float32` and `float64`."
),
category=DeprecationWarning,
)
kwargs["dtype"] = "float64"
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)
if flat_labels.dtype == "float16":
flat_labels = flat_labels.astype("float64")
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}"
)