File: C:/Users/fred/anaconda3/Lib/site-packages/xarray/core/alignment.py
from __future__ import annotations
import functools
import operator
from collections import defaultdict
from collections.abc import Hashable, Iterable, Mapping
from contextlib import suppress
from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, cast
import numpy as np
import pandas as pd
from xarray.core import dtypes
from xarray.core.common import DataWithCoords
from xarray.core.indexes import (
Index,
Indexes,
PandasIndex,
PandasMultiIndex,
indexes_all_equal,
safe_cast_to_index,
)
from xarray.core.utils import is_dict_like, is_full_slice
from xarray.core.variable import Variable, as_compatible_data, calculate_dimensions
if TYPE_CHECKING:
from xarray.core.dataarray import DataArray
from xarray.core.dataset import Dataset
from xarray.core.types import JoinOptions, T_DataArray, T_Dataset, T_DataWithCoords
DataAlignable = TypeVar("DataAlignable", bound=DataWithCoords)
def reindex_variables(
variables: Mapping[Any, Variable],
dim_pos_indexers: Mapping[Any, Any],
copy: bool = True,
fill_value: Any = dtypes.NA,
sparse: bool = False,
) -> dict[Hashable, Variable]:
"""Conform a dictionary of variables onto a new set of variables reindexed
with dimension positional indexers and possibly filled with missing values.
Not public API.
"""
new_variables = {}
dim_sizes = calculate_dimensions(variables)
masked_dims = set()
unchanged_dims = set()
for dim, indxr in dim_pos_indexers.items():
# Negative values in dim_pos_indexers mean values missing in the new index
# See ``Index.reindex_like``.
if (indxr < 0).any():
masked_dims.add(dim)
elif np.array_equal(indxr, np.arange(dim_sizes.get(dim, 0))):
unchanged_dims.add(dim)
for name, var in variables.items():
if isinstance(fill_value, dict):
fill_value_ = fill_value.get(name, dtypes.NA)
else:
fill_value_ = fill_value
if sparse:
var = var._as_sparse(fill_value=fill_value_)
indxr = tuple(
slice(None) if d in unchanged_dims else dim_pos_indexers.get(d, slice(None))
for d in var.dims
)
needs_masking = any(d in masked_dims for d in var.dims)
if needs_masking:
new_var = var._getitem_with_mask(indxr, fill_value=fill_value_)
elif all(is_full_slice(k) for k in indxr):
# no reindexing necessary
# here we need to manually deal with copying data, since
# we neither created a new ndarray nor used fancy indexing
new_var = var.copy(deep=copy)
else:
new_var = var[indxr]
new_variables[name] = new_var
return new_variables
CoordNamesAndDims = tuple[tuple[Hashable, tuple[Hashable, ...]], ...]
MatchingIndexKey = tuple[CoordNamesAndDims, type[Index]]
NormalizedIndexes = dict[MatchingIndexKey, Index]
NormalizedIndexVars = dict[MatchingIndexKey, dict[Hashable, Variable]]
class Aligner(Generic[DataAlignable]):
"""Implements all the complex logic for the re-indexing and alignment of Xarray
objects.
For internal use only, not public API.
Usage:
aligner = Aligner(*objects, **kwargs)
aligner.align()
aligned_objects = aligner.results
"""
objects: tuple[DataAlignable, ...]
results: tuple[DataAlignable, ...]
objects_matching_indexes: tuple[dict[MatchingIndexKey, Index], ...]
join: str
exclude_dims: frozenset[Hashable]
exclude_vars: frozenset[Hashable]
copy: bool
fill_value: Any
sparse: bool
indexes: dict[MatchingIndexKey, Index]
index_vars: dict[MatchingIndexKey, dict[Hashable, Variable]]
all_indexes: dict[MatchingIndexKey, list[Index]]
all_index_vars: dict[MatchingIndexKey, list[dict[Hashable, Variable]]]
aligned_indexes: dict[MatchingIndexKey, Index]
aligned_index_vars: dict[MatchingIndexKey, dict[Hashable, Variable]]
reindex: dict[MatchingIndexKey, bool]
reindex_kwargs: dict[str, Any]
unindexed_dim_sizes: dict[Hashable, set]
new_indexes: Indexes[Index]
def __init__(
self,
objects: Iterable[DataAlignable],
join: str = "inner",
indexes: Mapping[Any, Any] | None = None,
exclude_dims: Iterable = frozenset(),
exclude_vars: Iterable[Hashable] = frozenset(),
method: str | None = None,
tolerance: int | float | Iterable[int | float] | None = None,
copy: bool = True,
fill_value: Any = dtypes.NA,
sparse: bool = False,
):
self.objects = tuple(objects)
self.objects_matching_indexes = ()
if join not in ["inner", "outer", "override", "exact", "left", "right"]:
raise ValueError(f"invalid value for join: {join}")
self.join = join
self.copy = copy
self.fill_value = fill_value
self.sparse = sparse
if method is None and tolerance is None:
self.reindex_kwargs = {}
else:
self.reindex_kwargs = {"method": method, "tolerance": tolerance}
if isinstance(exclude_dims, str):
exclude_dims = [exclude_dims]
self.exclude_dims = frozenset(exclude_dims)
self.exclude_vars = frozenset(exclude_vars)
if indexes is None:
indexes = {}
self.indexes, self.index_vars = self._normalize_indexes(indexes)
self.all_indexes = {}
self.all_index_vars = {}
self.unindexed_dim_sizes = {}
self.aligned_indexes = {}
self.aligned_index_vars = {}
self.reindex = {}
self.results = tuple()
def _normalize_indexes(
self,
indexes: Mapping[Any, Any],
) -> tuple[NormalizedIndexes, NormalizedIndexVars]:
"""Normalize the indexes/indexers used for re-indexing or alignment.
Return dictionaries of xarray Index objects and coordinate variables
such that we can group matching indexes based on the dictionary keys.
"""
if isinstance(indexes, Indexes):
xr_variables = dict(indexes.variables)
else:
xr_variables = {}
xr_indexes: dict[Hashable, Index] = {}
for k, idx in indexes.items():
if not isinstance(idx, Index):
if getattr(idx, "dims", (k,)) != (k,):
raise ValueError(
f"Indexer has dimensions {idx.dims} that are different "
f"from that to be indexed along '{k}'"
)
data = as_compatible_data(idx)
pd_idx = safe_cast_to_index(data)
pd_idx.name = k
if isinstance(pd_idx, pd.MultiIndex):
idx = PandasMultiIndex(pd_idx, k)
else:
idx = PandasIndex(pd_idx, k, coord_dtype=data.dtype)
xr_variables.update(idx.create_variables())
xr_indexes[k] = idx
normalized_indexes = {}
normalized_index_vars = {}
for idx, index_vars in Indexes(xr_indexes, xr_variables).group_by_index():
coord_names_and_dims = []
all_dims: set[Hashable] = set()
for name, var in index_vars.items():
dims = var.dims
coord_names_and_dims.append((name, dims))
all_dims.update(dims)
exclude_dims = all_dims & self.exclude_dims
if exclude_dims == all_dims:
continue
elif exclude_dims:
excl_dims_str = ", ".join(str(d) for d in exclude_dims)
incl_dims_str = ", ".join(str(d) for d in all_dims - exclude_dims)
raise ValueError(
f"cannot exclude dimension(s) {excl_dims_str} from alignment because "
"these are used by an index together with non-excluded dimensions "
f"{incl_dims_str}"
)
key = (tuple(coord_names_and_dims), type(idx))
normalized_indexes[key] = idx
normalized_index_vars[key] = index_vars
return normalized_indexes, normalized_index_vars
def find_matching_indexes(self) -> None:
all_indexes: dict[MatchingIndexKey, list[Index]]
all_index_vars: dict[MatchingIndexKey, list[dict[Hashable, Variable]]]
all_indexes_dim_sizes: dict[MatchingIndexKey, dict[Hashable, set]]
objects_matching_indexes: list[dict[MatchingIndexKey, Index]]
all_indexes = defaultdict(list)
all_index_vars = defaultdict(list)
all_indexes_dim_sizes = defaultdict(lambda: defaultdict(set))
objects_matching_indexes = []
for obj in self.objects:
obj_indexes, obj_index_vars = self._normalize_indexes(obj.xindexes)
objects_matching_indexes.append(obj_indexes)
for key, idx in obj_indexes.items():
all_indexes[key].append(idx)
for key, index_vars in obj_index_vars.items():
all_index_vars[key].append(index_vars)
for dim, size in calculate_dimensions(index_vars).items():
all_indexes_dim_sizes[key][dim].add(size)
self.objects_matching_indexes = tuple(objects_matching_indexes)
self.all_indexes = all_indexes
self.all_index_vars = all_index_vars
if self.join == "override":
for dim_sizes in all_indexes_dim_sizes.values():
for dim, sizes in dim_sizes.items():
if len(sizes) > 1:
raise ValueError(
"cannot align objects with join='override' with matching indexes "
f"along dimension {dim!r} that don't have the same size"
)
def find_matching_unindexed_dims(self) -> None:
unindexed_dim_sizes = defaultdict(set)
for obj in self.objects:
for dim in obj.dims:
if dim not in self.exclude_dims and dim not in obj.xindexes.dims:
unindexed_dim_sizes[dim].add(obj.sizes[dim])
self.unindexed_dim_sizes = unindexed_dim_sizes
def assert_no_index_conflict(self) -> None:
"""Check for uniqueness of both coordinate and dimension names across all sets
of matching indexes.
We need to make sure that all indexes used for re-indexing or alignment
are fully compatible and do not conflict each other.
Note: perhaps we could choose less restrictive constraints and instead
check for conflicts among the dimension (position) indexers returned by
`Index.reindex_like()` for each matching pair of object index / aligned
index?
(ref: https://github.com/pydata/xarray/issues/1603#issuecomment-442965602)
"""
matching_keys = set(self.all_indexes) | set(self.indexes)
coord_count: dict[Hashable, int] = defaultdict(int)
dim_count: dict[Hashable, int] = defaultdict(int)
for coord_names_dims, _ in matching_keys:
dims_set: set[Hashable] = set()
for name, dims in coord_names_dims:
coord_count[name] += 1
dims_set.update(dims)
for dim in dims_set:
dim_count[dim] += 1
for count, msg in [(coord_count, "coordinates"), (dim_count, "dimensions")]:
dup = {k: v for k, v in count.items() if v > 1}
if dup:
items_msg = ", ".join(
f"{k!r} ({v} conflicting indexes)" for k, v in dup.items()
)
raise ValueError(
"cannot re-index or align objects with conflicting indexes found for "
f"the following {msg}: {items_msg}\n"
"Conflicting indexes may occur when\n"
"- they relate to different sets of coordinate and/or dimension names\n"
"- they don't have the same type\n"
"- they may be used to reindex data along common dimensions"
)
def _need_reindex(self, dims, cmp_indexes) -> bool:
"""Whether or not we need to reindex variables for a set of
matching indexes.
We don't reindex when all matching indexes are equal for two reasons:
- It's faster for the usual case (already aligned objects).
- It ensures it's possible to do operations that don't require alignment
on indexes with duplicate values (which cannot be reindexed with
pandas). This is useful, e.g., for overwriting such duplicate indexes.
"""
if not indexes_all_equal(cmp_indexes):
# always reindex when matching indexes are not equal
return True
unindexed_dims_sizes = {}
for dim in dims:
if dim in self.unindexed_dim_sizes:
sizes = self.unindexed_dim_sizes[dim]
if len(sizes) > 1:
# reindex if different sizes are found for unindexed dims
return True
else:
unindexed_dims_sizes[dim] = next(iter(sizes))
if unindexed_dims_sizes:
indexed_dims_sizes = {}
for cmp in cmp_indexes:
index_vars = cmp[1]
for var in index_vars.values():
indexed_dims_sizes.update(var.sizes)
for dim, size in unindexed_dims_sizes.items():
if indexed_dims_sizes.get(dim, -1) != size:
# reindex if unindexed dimension size doesn't match
return True
return False
def _get_index_joiner(self, index_cls) -> Callable:
if self.join in ["outer", "inner"]:
return functools.partial(
functools.reduce,
functools.partial(index_cls.join, how=self.join),
)
elif self.join == "left":
return operator.itemgetter(0)
elif self.join == "right":
return operator.itemgetter(-1)
elif self.join == "override":
# We rewrite all indexes and then use join='left'
return operator.itemgetter(0)
else:
# join='exact' return dummy lambda (error is raised)
return lambda _: None
def align_indexes(self) -> None:
"""Compute all aligned indexes and their corresponding coordinate variables."""
aligned_indexes = {}
aligned_index_vars = {}
reindex = {}
new_indexes = {}
new_index_vars = {}
for key, matching_indexes in self.all_indexes.items():
matching_index_vars = self.all_index_vars[key]
dims = {d for coord in matching_index_vars[0].values() for d in coord.dims}
index_cls = key[1]
if self.join == "override":
joined_index = matching_indexes[0]
joined_index_vars = matching_index_vars[0]
need_reindex = False
elif key in self.indexes:
joined_index = self.indexes[key]
joined_index_vars = self.index_vars[key]
cmp_indexes = list(
zip(
[joined_index] + matching_indexes,
[joined_index_vars] + matching_index_vars,
)
)
need_reindex = self._need_reindex(dims, cmp_indexes)
else:
if len(matching_indexes) > 1:
need_reindex = self._need_reindex(
dims,
list(zip(matching_indexes, matching_index_vars)),
)
else:
need_reindex = False
if need_reindex:
if self.join == "exact":
raise ValueError(
"cannot align objects with join='exact' where "
"index/labels/sizes are not equal along "
"these coordinates (dimensions): "
+ ", ".join(f"{name!r} {dims!r}" for name, dims in key[0])
)
joiner = self._get_index_joiner(index_cls)
joined_index = joiner(matching_indexes)
if self.join == "left":
joined_index_vars = matching_index_vars[0]
elif self.join == "right":
joined_index_vars = matching_index_vars[-1]
else:
joined_index_vars = joined_index.create_variables()
else:
joined_index = matching_indexes[0]
joined_index_vars = matching_index_vars[0]
reindex[key] = need_reindex
aligned_indexes[key] = joined_index
aligned_index_vars[key] = joined_index_vars
for name, var in joined_index_vars.items():
new_indexes[name] = joined_index
new_index_vars[name] = var
# Explicitly provided indexes that are not found in objects to align
# may relate to unindexed dimensions so we add them too
for key, idx in self.indexes.items():
if key not in aligned_indexes:
index_vars = self.index_vars[key]
reindex[key] = False
aligned_indexes[key] = idx
aligned_index_vars[key] = index_vars
for name, var in index_vars.items():
new_indexes[name] = idx
new_index_vars[name] = var
self.aligned_indexes = aligned_indexes
self.aligned_index_vars = aligned_index_vars
self.reindex = reindex
self.new_indexes = Indexes(new_indexes, new_index_vars)
def assert_unindexed_dim_sizes_equal(self) -> None:
for dim, sizes in self.unindexed_dim_sizes.items():
index_size = self.new_indexes.dims.get(dim)
if index_size is not None:
sizes.add(index_size)
add_err_msg = (
f" (note: an index is found along that dimension "
f"with size={index_size!r})"
)
else:
add_err_msg = ""
if len(sizes) > 1:
raise ValueError(
f"cannot reindex or align along dimension {dim!r} "
f"because of conflicting dimension sizes: {sizes!r}" + add_err_msg
)
def override_indexes(self) -> None:
objects = list(self.objects)
for i, obj in enumerate(objects[1:]):
new_indexes = {}
new_variables = {}
matching_indexes = self.objects_matching_indexes[i + 1]
for key, aligned_idx in self.aligned_indexes.items():
obj_idx = matching_indexes.get(key)
if obj_idx is not None:
for name, var in self.aligned_index_vars[key].items():
new_indexes[name] = aligned_idx
new_variables[name] = var.copy(deep=self.copy)
objects[i + 1] = obj._overwrite_indexes(new_indexes, new_variables)
self.results = tuple(objects)
def _get_dim_pos_indexers(
self,
matching_indexes: dict[MatchingIndexKey, Index],
) -> dict[Hashable, Any]:
dim_pos_indexers = {}
for key, aligned_idx in self.aligned_indexes.items():
obj_idx = matching_indexes.get(key)
if obj_idx is not None:
if self.reindex[key]:
indexers = obj_idx.reindex_like(aligned_idx, **self.reindex_kwargs)
dim_pos_indexers.update(indexers)
return dim_pos_indexers
def _get_indexes_and_vars(
self,
obj: DataAlignable,
matching_indexes: dict[MatchingIndexKey, Index],
) -> tuple[dict[Hashable, Index], dict[Hashable, Variable]]:
new_indexes = {}
new_variables = {}
for key, aligned_idx in self.aligned_indexes.items():
index_vars = self.aligned_index_vars[key]
obj_idx = matching_indexes.get(key)
if obj_idx is None:
# add the index if it relates to unindexed dimensions in obj
index_vars_dims = {d for var in index_vars.values() for d in var.dims}
if index_vars_dims <= set(obj.dims):
obj_idx = aligned_idx
if obj_idx is not None:
for name, var in index_vars.items():
new_indexes[name] = aligned_idx
new_variables[name] = var.copy(deep=self.copy)
return new_indexes, new_variables
def _reindex_one(
self,
obj: DataAlignable,
matching_indexes: dict[MatchingIndexKey, Index],
) -> DataAlignable:
new_indexes, new_variables = self._get_indexes_and_vars(obj, matching_indexes)
dim_pos_indexers = self._get_dim_pos_indexers(matching_indexes)
new_obj = obj._reindex_callback(
self,
dim_pos_indexers,
new_variables,
new_indexes,
self.fill_value,
self.exclude_dims,
self.exclude_vars,
)
new_obj.encoding = obj.encoding
return new_obj
def reindex_all(self) -> None:
self.results = tuple(
self._reindex_one(obj, matching_indexes)
for obj, matching_indexes in zip(
self.objects, self.objects_matching_indexes
)
)
def align(self) -> None:
if not self.indexes and len(self.objects) == 1:
# fast path for the trivial case
(obj,) = self.objects
self.results = (obj.copy(deep=self.copy),)
return
self.find_matching_indexes()
self.find_matching_unindexed_dims()
self.assert_no_index_conflict()
self.align_indexes()
self.assert_unindexed_dim_sizes_equal()
if self.join == "override":
self.override_indexes()
elif self.join == "exact" and not self.copy:
self.results = self.objects
else:
self.reindex_all()
def align(
*objects: DataAlignable,
join: JoinOptions = "inner",
copy: bool = True,
indexes=None,
exclude=frozenset(),
fill_value=dtypes.NA,
) -> tuple[DataAlignable, ...]:
"""
Given any number of Dataset and/or DataArray objects, returns new
objects with aligned indexes and dimension sizes.
Array from the aligned objects are suitable as input to mathematical
operators, because along each dimension they have the same index and size.
Missing values (if ``join != 'inner'``) are filled with ``fill_value``.
The default fill value is NaN.
Parameters
----------
*objects : Dataset or DataArray
Objects to align.
join : {"outer", "inner", "left", "right", "exact", "override"}, optional
Method for joining the indexes of the passed objects along each
dimension:
- "outer": use the union of object indexes
- "inner": use the intersection of object indexes
- "left": use indexes from the first object with each dimension
- "right": use indexes from the last object with each dimension
- "exact": instead of aligning, raise `ValueError` when indexes to be
aligned are not equal
- "override": if indexes are of same size, rewrite indexes to be
those of the first object with that dimension. Indexes for the same
dimension must have the same size in all objects.
copy : bool, default: True
If ``copy=True``, data in the return values is always copied. If
``copy=False`` and reindexing is unnecessary, or can be performed with
only slice operations, then the output may share memory with the input.
In either case, new xarray objects are always returned.
indexes : dict-like, optional
Any indexes explicitly provided with the `indexes` argument should be
used in preference to the aligned indexes.
exclude : sequence of str, optional
Dimensions that must be excluded from alignment
fill_value : scalar or dict-like, optional
Value to use for newly missing values. If a dict-like, maps
variable names to fill values. Use a data array's name to
refer to its values.
Returns
-------
aligned : tuple of DataArray or Dataset
Tuple of objects with the same type as `*objects` with aligned
coordinates.
Raises
------
ValueError
If any dimensions without labels on the arguments have different sizes,
or a different size than the size of the aligned dimension labels.
Examples
--------
>>> x = xr.DataArray(
... [[25, 35], [10, 24]],
... dims=("lat", "lon"),
... coords={"lat": [35.0, 40.0], "lon": [100.0, 120.0]},
... )
>>> y = xr.DataArray(
... [[20, 5], [7, 13]],
... dims=("lat", "lon"),
... coords={"lat": [35.0, 42.0], "lon": [100.0, 120.0]},
... )
>>> x
<xarray.DataArray (lat: 2, lon: 2)>
array([[25, 35],
[10, 24]])
Coordinates:
* lat (lat) float64 35.0 40.0
* lon (lon) float64 100.0 120.0
>>> y
<xarray.DataArray (lat: 2, lon: 2)>
array([[20, 5],
[ 7, 13]])
Coordinates:
* lat (lat) float64 35.0 42.0
* lon (lon) float64 100.0 120.0
>>> a, b = xr.align(x, y)
>>> a
<xarray.DataArray (lat: 1, lon: 2)>
array([[25, 35]])
Coordinates:
* lat (lat) float64 35.0
* lon (lon) float64 100.0 120.0
>>> b
<xarray.DataArray (lat: 1, lon: 2)>
array([[20, 5]])
Coordinates:
* lat (lat) float64 35.0
* lon (lon) float64 100.0 120.0
>>> a, b = xr.align(x, y, join="outer")
>>> a
<xarray.DataArray (lat: 3, lon: 2)>
array([[25., 35.],
[10., 24.],
[nan, nan]])
Coordinates:
* lat (lat) float64 35.0 40.0 42.0
* lon (lon) float64 100.0 120.0
>>> b
<xarray.DataArray (lat: 3, lon: 2)>
array([[20., 5.],
[nan, nan],
[ 7., 13.]])
Coordinates:
* lat (lat) float64 35.0 40.0 42.0
* lon (lon) float64 100.0 120.0
>>> a, b = xr.align(x, y, join="outer", fill_value=-999)
>>> a
<xarray.DataArray (lat: 3, lon: 2)>
array([[ 25, 35],
[ 10, 24],
[-999, -999]])
Coordinates:
* lat (lat) float64 35.0 40.0 42.0
* lon (lon) float64 100.0 120.0
>>> b
<xarray.DataArray (lat: 3, lon: 2)>
array([[ 20, 5],
[-999, -999],
[ 7, 13]])
Coordinates:
* lat (lat) float64 35.0 40.0 42.0
* lon (lon) float64 100.0 120.0
>>> a, b = xr.align(x, y, join="left")
>>> a
<xarray.DataArray (lat: 2, lon: 2)>
array([[25, 35],
[10, 24]])
Coordinates:
* lat (lat) float64 35.0 40.0
* lon (lon) float64 100.0 120.0
>>> b
<xarray.DataArray (lat: 2, lon: 2)>
array([[20., 5.],
[nan, nan]])
Coordinates:
* lat (lat) float64 35.0 40.0
* lon (lon) float64 100.0 120.0
>>> a, b = xr.align(x, y, join="right")
>>> a
<xarray.DataArray (lat: 2, lon: 2)>
array([[25., 35.],
[nan, nan]])
Coordinates:
* lat (lat) float64 35.0 42.0
* lon (lon) float64 100.0 120.0
>>> b
<xarray.DataArray (lat: 2, lon: 2)>
array([[20, 5],
[ 7, 13]])
Coordinates:
* lat (lat) float64 35.0 42.0
* lon (lon) float64 100.0 120.0
>>> a, b = xr.align(x, y, join="exact")
Traceback (most recent call last):
...
ValueError: cannot align objects with join='exact' ...
>>> a, b = xr.align(x, y, join="override")
>>> a
<xarray.DataArray (lat: 2, lon: 2)>
array([[25, 35],
[10, 24]])
Coordinates:
* lat (lat) float64 35.0 40.0
* lon (lon) float64 100.0 120.0
>>> b
<xarray.DataArray (lat: 2, lon: 2)>
array([[20, 5],
[ 7, 13]])
Coordinates:
* lat (lat) float64 35.0 40.0
* lon (lon) float64 100.0 120.0
"""
aligner = Aligner(
objects,
join=join,
copy=copy,
indexes=indexes,
exclude_dims=exclude,
fill_value=fill_value,
)
aligner.align()
return aligner.results
def deep_align(
objects: Iterable[Any],
join: JoinOptions = "inner",
copy=True,
indexes=None,
exclude=frozenset(),
raise_on_invalid=True,
fill_value=dtypes.NA,
):
"""Align objects for merging, recursing into dictionary values.
This function is not public API.
"""
from xarray.core.dataarray import DataArray
from xarray.core.dataset import Dataset
if indexes is None:
indexes = {}
def is_alignable(obj):
return isinstance(obj, (DataArray, Dataset))
positions = []
keys = []
out = []
targets = []
no_key = object()
not_replaced = object()
for position, variables in enumerate(objects):
if is_alignable(variables):
positions.append(position)
keys.append(no_key)
targets.append(variables)
out.append(not_replaced)
elif is_dict_like(variables):
current_out = {}
for k, v in variables.items():
if is_alignable(v) and k not in indexes:
# Skip variables in indexes for alignment, because these
# should to be overwritten instead:
# https://github.com/pydata/xarray/issues/725
# https://github.com/pydata/xarray/issues/3377
# TODO(shoyer): doing this here feels super-hacky -- can we
# move it explicitly into merge instead?
positions.append(position)
keys.append(k)
targets.append(v)
current_out[k] = not_replaced
else:
current_out[k] = v
out.append(current_out)
elif raise_on_invalid:
raise ValueError(
"object to align is neither an xarray.Dataset, "
"an xarray.DataArray nor a dictionary: {!r}".format(variables)
)
else:
out.append(variables)
aligned = align(
*targets,
join=join,
copy=copy,
indexes=indexes,
exclude=exclude,
fill_value=fill_value,
)
for position, key, aligned_obj in zip(positions, keys, aligned):
if key is no_key:
out[position] = aligned_obj
else:
out[position][key] = aligned_obj # type: ignore[index] # maybe someone can fix this?
return out
def reindex(
obj: DataAlignable,
indexers: Mapping[Any, Any],
method: str | None = None,
tolerance: int | float | Iterable[int | float] | None = None,
copy: bool = True,
fill_value: Any = dtypes.NA,
sparse: bool = False,
exclude_vars: Iterable[Hashable] = frozenset(),
) -> DataAlignable:
"""Re-index either a Dataset or a DataArray.
Not public API.
"""
# TODO: (benbovy - explicit indexes): uncomment?
# --> from reindex docstrings: "any mis-matched dimension is simply ignored"
# bad_keys = [k for k in indexers if k not in obj._indexes and k not in obj.dims]
# if bad_keys:
# raise ValueError(
# f"indexer keys {bad_keys} do not correspond to any indexed coordinate "
# "or unindexed dimension in the object to reindex"
# )
aligner = Aligner(
(obj,),
indexes=indexers,
method=method,
tolerance=tolerance,
copy=copy,
fill_value=fill_value,
sparse=sparse,
exclude_vars=exclude_vars,
)
aligner.align()
return aligner.results[0]
def reindex_like(
obj: DataAlignable,
other: Dataset | DataArray,
method: str | None = None,
tolerance: int | float | Iterable[int | float] | None = None,
copy: bool = True,
fill_value: Any = dtypes.NA,
) -> DataAlignable:
"""Re-index either a Dataset or a DataArray like another Dataset/DataArray.
Not public API.
"""
if not other._indexes:
# This check is not performed in Aligner.
for dim in other.dims:
if dim in obj.dims:
other_size = other.sizes[dim]
obj_size = obj.sizes[dim]
if other_size != obj_size:
raise ValueError(
"different size for unlabeled "
f"dimension on argument {dim!r}: {other_size!r} vs {obj_size!r}"
)
return reindex(
obj,
indexers=other.xindexes,
method=method,
tolerance=tolerance,
copy=copy,
fill_value=fill_value,
)
def _get_broadcast_dims_map_common_coords(args, exclude):
common_coords = {}
dims_map = {}
for arg in args:
for dim in arg.dims:
if dim not in common_coords and dim not in exclude:
dims_map[dim] = arg.sizes[dim]
if dim in arg._indexes:
common_coords.update(arg.xindexes.get_all_coords(dim))
return dims_map, common_coords
def _broadcast_helper(
arg: T_DataWithCoords, exclude, dims_map, common_coords
) -> T_DataWithCoords:
from xarray.core.dataarray import DataArray
from xarray.core.dataset import Dataset
def _set_dims(var):
# Add excluded dims to a copy of dims_map
var_dims_map = dims_map.copy()
for dim in exclude:
with suppress(ValueError):
# ignore dim not in var.dims
var_dims_map[dim] = var.shape[var.dims.index(dim)]
return var.set_dims(var_dims_map)
def _broadcast_array(array: T_DataArray) -> T_DataArray:
data = _set_dims(array.variable)
coords = dict(array.coords)
coords.update(common_coords)
return array.__class__(
data, coords, data.dims, name=array.name, attrs=array.attrs
)
def _broadcast_dataset(ds: T_Dataset) -> T_Dataset:
data_vars = {k: _set_dims(ds.variables[k]) for k in ds.data_vars}
coords = dict(ds.coords)
coords.update(common_coords)
return ds.__class__(data_vars, coords, ds.attrs)
# remove casts once https://github.com/python/mypy/issues/12800 is resolved
if isinstance(arg, DataArray):
return cast("T_DataWithCoords", _broadcast_array(arg))
elif isinstance(arg, Dataset):
return cast("T_DataWithCoords", _broadcast_dataset(arg))
else:
raise ValueError("all input must be Dataset or DataArray objects")
# TODO: this typing is too restrictive since it cannot deal with mixed
# DataArray and Dataset types...? Is this a problem?
def broadcast(*args: T_DataWithCoords, exclude=None) -> tuple[T_DataWithCoords, ...]:
"""Explicitly broadcast any number of DataArray or Dataset objects against
one another.
xarray objects automatically broadcast against each other in arithmetic
operations, so this function should not be necessary for normal use.
If no change is needed, the input data is returned to the output without
being copied.
Parameters
----------
*args : DataArray or Dataset
Arrays to broadcast against each other.
exclude : sequence of str, optional
Dimensions that must not be broadcasted
Returns
-------
broadcast : tuple of DataArray or tuple of Dataset
The same data as the input arrays, but with additional dimensions
inserted so that all data arrays have the same dimensions and shape.
Examples
--------
Broadcast two data arrays against one another to fill out their dimensions:
>>> a = xr.DataArray([1, 2, 3], dims="x")
>>> b = xr.DataArray([5, 6], dims="y")
>>> a
<xarray.DataArray (x: 3)>
array([1, 2, 3])
Dimensions without coordinates: x
>>> b
<xarray.DataArray (y: 2)>
array([5, 6])
Dimensions without coordinates: y
>>> a2, b2 = xr.broadcast(a, b)
>>> a2
<xarray.DataArray (x: 3, y: 2)>
array([[1, 1],
[2, 2],
[3, 3]])
Dimensions without coordinates: x, y
>>> b2
<xarray.DataArray (x: 3, y: 2)>
array([[5, 6],
[5, 6],
[5, 6]])
Dimensions without coordinates: x, y
Fill out the dimensions of all data variables in a dataset:
>>> ds = xr.Dataset({"a": a, "b": b})
>>> (ds2,) = xr.broadcast(ds) # use tuple unpacking to extract one dataset
>>> ds2
<xarray.Dataset>
Dimensions: (x: 3, y: 2)
Dimensions without coordinates: x, y
Data variables:
a (x, y) int64 1 1 2 2 3 3
b (x, y) int64 5 6 5 6 5 6
"""
if exclude is None:
exclude = set()
args = align(*args, join="outer", copy=False, exclude=exclude)
dims_map, common_coords = _get_broadcast_dims_map_common_coords(args, exclude)
result = [_broadcast_helper(arg, exclude, dims_map, common_coords) for arg in args]
return tuple(result)