from __future__ import annotations import functools import operator from collections import defaultdict from contextlib import suppress from typing import ( TYPE_CHECKING, Any, Callable, Dict, Generic, Hashable, Iterable, Mapping, Tuple, Type, 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() 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 array([[25, 35], [10, 24]]) Coordinates: * lat (lat) float64 35.0 40.0 * lon (lon) float64 100.0 120.0 >>> y 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 array([[25, 35]]) Coordinates: * lat (lat) float64 35.0 * lon (lon) float64 100.0 120.0 >>> b 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 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 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 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 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 array([[25, 35], [10, 24]]) Coordinates: * lat (lat) float64 35.0 40.0 * lon (lon) float64 100.0 120.0 >>> b 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 array([[25., 35.], [nan, nan]]) Coordinates: * lat (lat) float64 35.0 42.0 * lon (lon) float64 100.0 120.0 >>> b 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 array([[25, 35], [10, 24]]) Coordinates: * lat (lat) float64 35.0 40.0 * lon (lon) float64 100.0 120.0 >>> b 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 array([1, 2, 3]) Dimensions without coordinates: x >>> b array([5, 6]) Dimensions without coordinates: y >>> a2, b2 = xr.broadcast(a, b) >>> a2 array([[1, 1], [2, 2], [3, 3]]) Dimensions without coordinates: x, y >>> b2 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 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)