import copy import datetime import functools import sys import warnings from collections import defaultdict from html import escape from numbers import Number from operator import methodcaller from pathlib import Path from typing import ( TYPE_CHECKING, Any, Callable, DefaultDict, Dict, Hashable, Iterable, Iterator, List, Mapping, MutableMapping, Optional, Sequence, Set, Tuple, TypeVar, Union, cast, overload, ) import numpy as np import pandas as pd import xarray as xr from ..coding.cftimeindex import _parse_array_of_cftime_strings from ..plot.dataset_plot import _Dataset_PlotMethods from . import ( alignment, dtypes, duck_array_ops, formatting, formatting_html, groupby, ops, resample, rolling, utils, weighted, ) from .alignment import _broadcast_helper, _get_broadcast_dims_map_common_coords, align from .common import ( DataWithCoords, ImplementsDatasetReduce, _contains_datetime_like_objects, ) from .coordinates import ( DatasetCoordinates, LevelCoordinatesSource, assert_coordinate_consistent, remap_label_indexers, ) from .duck_array_ops import datetime_to_numeric from .indexes import ( Indexes, default_indexes, isel_variable_and_index, propagate_indexes, remove_unused_levels_categories, roll_index, ) from .indexing import is_fancy_indexer from .merge import ( dataset_merge_method, dataset_update_method, merge_coordinates_without_align, merge_data_and_coords, ) from .missing import get_clean_interp_index from .options import OPTIONS, _get_keep_attrs from .pycompat import is_duck_dask_array from .utils import ( Default, Frozen, SortedKeysDict, _check_inplace, _default, decode_numpy_dict_values, drop_dims_from_indexers, either_dict_or_kwargs, hashable, infix_dims, is_dict_like, is_scalar, maybe_wrap_array, ) from .variable import ( IndexVariable, Variable, as_variable, assert_unique_multiindex_level_names, broadcast_variables, ) if TYPE_CHECKING: from ..backends import AbstractDataStore, ZarrStore from .dataarray import DataArray from .merge import CoercibleMapping T_DSorDA = TypeVar("T_DSorDA", DataArray, "Dataset") try: from dask.delayed import Delayed except ImportError: Delayed = None # list of attributes of pd.DatetimeIndex that are ndarrays of time info _DATETIMEINDEX_COMPONENTS = [ "year", "month", "day", "hour", "minute", "second", "microsecond", "nanosecond", "date", "time", "dayofyear", "weekofyear", "dayofweek", "quarter", ] def _get_virtual_variable( variables, key: Hashable, level_vars: Mapping = None, dim_sizes: Mapping = None ) -> Tuple[Hashable, Hashable, Variable]: """Get a virtual variable (e.g., 'time.year' or a MultiIndex level) from a dict of xarray.Variable objects (if possible) """ if level_vars is None: level_vars = {} if dim_sizes is None: dim_sizes = {} if key in dim_sizes: data = pd.Index(range(dim_sizes[key]), name=key) variable = IndexVariable((key,), data) return key, key, variable if not isinstance(key, str): raise KeyError(key) split_key = key.split(".", 1) var_name: Optional[str] if len(split_key) == 2: ref_name, var_name = split_key elif len(split_key) == 1: ref_name, var_name = key, None else: raise KeyError(key) if ref_name in level_vars: dim_var = variables[level_vars[ref_name]] ref_var = dim_var.to_index_variable().get_level_variable(ref_name) else: ref_var = variables[ref_name] if var_name is None: virtual_var = ref_var var_name = key else: if _contains_datetime_like_objects(ref_var): ref_var = xr.DataArray(ref_var) data = getattr(ref_var.dt, var_name).data else: data = getattr(ref_var, var_name).data virtual_var = Variable(ref_var.dims, data) return ref_name, var_name, virtual_var def calculate_dimensions(variables: Mapping[Hashable, Variable]) -> Dict[Hashable, int]: """Calculate the dimensions corresponding to a set of variables. Returns dictionary mapping from dimension names to sizes. Raises ValueError if any of the dimension sizes conflict. """ dims: Dict[Hashable, int] = {} last_used = {} scalar_vars = {k for k, v in variables.items() if not v.dims} for k, var in variables.items(): for dim, size in zip(var.dims, var.shape): if dim in scalar_vars: raise ValueError( "dimension %r already exists as a scalar variable" % dim ) if dim not in dims: dims[dim] = size last_used[dim] = k elif dims[dim] != size: raise ValueError( "conflicting sizes for dimension %r: " "length %s on %r and length %s on %r" % (dim, size, k, dims[dim], last_used[dim]) ) return dims def merge_indexes( indexes: Mapping[Hashable, Union[Hashable, Sequence[Hashable]]], variables: Mapping[Hashable, Variable], coord_names: Set[Hashable], append: bool = False, ) -> Tuple[Dict[Hashable, Variable], Set[Hashable]]: """Merge variables into multi-indexes. Not public API. Used in Dataset and DataArray set_index methods. """ vars_to_replace: Dict[Hashable, Variable] = {} vars_to_remove: List[Hashable] = [] dims_to_replace: Dict[Hashable, Hashable] = {} error_msg = "{} is not the name of an existing variable." for dim, var_names in indexes.items(): if isinstance(var_names, str) or not isinstance(var_names, Sequence): var_names = [var_names] names: List[Hashable] = [] codes: List[List[int]] = [] levels: List[List[int]] = [] current_index_variable = variables.get(dim) for n in var_names: try: var = variables[n] except KeyError: raise ValueError(error_msg.format(n)) if ( current_index_variable is not None and var.dims != current_index_variable.dims ): raise ValueError( "dimension mismatch between %r %s and %r %s" % (dim, current_index_variable.dims, n, var.dims) ) if current_index_variable is not None and append: current_index = current_index_variable.to_index() if isinstance(current_index, pd.MultiIndex): names.extend(current_index.names) codes.extend(current_index.codes) levels.extend(current_index.levels) else: names.append("%s_level_0" % dim) cat = pd.Categorical(current_index.values, ordered=True) codes.append(cat.codes) levels.append(cat.categories) if not len(names) and len(var_names) == 1: idx = pd.Index(variables[var_names[0]].values) else: # MultiIndex for n in var_names: try: var = variables[n] except KeyError: raise ValueError(error_msg.format(n)) names.append(n) cat = pd.Categorical(var.values, ordered=True) codes.append(cat.codes) levels.append(cat.categories) idx = pd.MultiIndex(levels, codes, names=names) for n in names: dims_to_replace[n] = dim vars_to_replace[dim] = IndexVariable(dim, idx) vars_to_remove.extend(var_names) new_variables = {k: v for k, v in variables.items() if k not in vars_to_remove} new_variables.update(vars_to_replace) # update dimensions if necessary, GH: 3512 for k, v in new_variables.items(): if any(d in dims_to_replace for d in v.dims): new_dims = [dims_to_replace.get(d, d) for d in v.dims] new_variables[k] = v._replace(dims=new_dims) new_coord_names = coord_names | set(vars_to_replace) new_coord_names -= set(vars_to_remove) return new_variables, new_coord_names def split_indexes( dims_or_levels: Union[Hashable, Sequence[Hashable]], variables: Mapping[Hashable, Variable], coord_names: Set[Hashable], level_coords: Mapping[Hashable, Hashable], drop: bool = False, ) -> Tuple[Dict[Hashable, Variable], Set[Hashable]]: """Extract (multi-)indexes (levels) as variables. Not public API. Used in Dataset and DataArray reset_index methods. """ if isinstance(dims_or_levels, str) or not isinstance(dims_or_levels, Sequence): dims_or_levels = [dims_or_levels] dim_levels: DefaultDict[Any, List[Hashable]] = defaultdict(list) dims = [] for k in dims_or_levels: if k in level_coords: dim_levels[level_coords[k]].append(k) else: dims.append(k) vars_to_replace = {} vars_to_create: Dict[Hashable, Variable] = {} vars_to_remove = [] for d in dims: index = variables[d].to_index() if isinstance(index, pd.MultiIndex): dim_levels[d] = index.names else: vars_to_remove.append(d) if not drop: vars_to_create[str(d) + "_"] = Variable(d, index, variables[d].attrs) for d, levs in dim_levels.items(): index = variables[d].to_index() if len(levs) == index.nlevels: vars_to_remove.append(d) else: vars_to_replace[d] = IndexVariable(d, index.droplevel(levs)) if not drop: for lev in levs: idx = index.get_level_values(lev) vars_to_create[idx.name] = Variable(d, idx, variables[d].attrs) new_variables = dict(variables) for v in set(vars_to_remove): del new_variables[v] new_variables.update(vars_to_replace) new_variables.update(vars_to_create) new_coord_names = (coord_names | set(vars_to_create)) - set(vars_to_remove) return new_variables, new_coord_names def _assert_empty(args: tuple, msg: str = "%s") -> None: if args: raise ValueError(msg % args) def _maybe_chunk( name, var, chunks=None, token=None, lock=None, name_prefix="xarray-", overwrite_encoded_chunks=False, ): from dask.base import tokenize if chunks is not None: chunks = {dim: chunks[dim] for dim in var.dims if dim in chunks} if var.ndim: # when rechunking by different amounts, make sure dask names change # by provinding chunks as an input to tokenize. # subtle bugs result otherwise. see GH3350 token2 = tokenize(name, token if token else var._data, chunks) name2 = f"{name_prefix}{name}-{token2}" var = var.chunk(chunks, name=name2, lock=lock) if overwrite_encoded_chunks and var.chunks is not None: var.encoding["chunks"] = tuple(x[0] for x in var.chunks) return var else: return var def as_dataset(obj: Any) -> "Dataset": """Cast the given object to a Dataset. Handles Datasets, DataArrays and dictionaries of variables. A new Dataset object is only created if the provided object is not already one. """ if hasattr(obj, "to_dataset"): obj = obj.to_dataset() if not isinstance(obj, Dataset): obj = Dataset(obj) return obj class DataVariables(Mapping[Hashable, "DataArray"]): __slots__ = ("_dataset",) def __init__(self, dataset: "Dataset"): self._dataset = dataset def __iter__(self) -> Iterator[Hashable]: return ( key for key in self._dataset._variables if key not in self._dataset._coord_names ) def __len__(self) -> int: return len(self._dataset._variables) - len(self._dataset._coord_names) def __contains__(self, key: Hashable) -> bool: return key in self._dataset._variables and key not in self._dataset._coord_names def __getitem__(self, key: Hashable) -> "DataArray": if key not in self._dataset._coord_names: return cast("DataArray", self._dataset[key]) raise KeyError(key) def __repr__(self) -> str: return formatting.data_vars_repr(self) @property def variables(self) -> Mapping[Hashable, Variable]: all_variables = self._dataset.variables return Frozen({k: all_variables[k] for k in self}) def _ipython_key_completions_(self): """Provide method for the key-autocompletions in IPython. """ return [ key for key in self._dataset._ipython_key_completions_() if key not in self._dataset._coord_names ] class _LocIndexer: __slots__ = ("dataset",) def __init__(self, dataset: "Dataset"): self.dataset = dataset def __getitem__(self, key: Mapping[Hashable, Any]) -> "Dataset": if not utils.is_dict_like(key): raise TypeError("can only lookup dictionaries from Dataset.loc") return self.dataset.sel(key) class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords): """A multi-dimensional, in memory, array database. A dataset resembles an in-memory representation of a NetCDF file, and consists of variables, coordinates and attributes which together form a self describing dataset. Dataset implements the mapping interface with keys given by variable names and values given by DataArray objects for each variable name. One dimensional variables with name equal to their dimension are index coordinates used for label based indexing. To load data from a file or file-like object, use the `open_dataset` function. Parameters ---------- data_vars : dict-like, optional A mapping from variable names to :py:class:`~xarray.DataArray` objects, :py:class:`~xarray.Variable` objects or to tuples of the form ``(dims, data[, attrs])`` which can be used as arguments to create a new ``Variable``. Each dimension must have the same length in all variables in which it appears. The following notations are accepted: - mapping {var name: DataArray} - mapping {var name: Variable} - mapping {var name: (dimension name, array-like)} - mapping {var name: (tuple of dimension names, array-like)} - mapping {dimension name: array-like} (it will be automatically moved to coords, see below) Each dimension must have the same length in all variables in which it appears. coords : dict-like, optional Another mapping in similar form as the `data_vars` argument, except the each item is saved on the dataset as a "coordinate". These variables have an associated meaning: they describe constant/fixed/independent quantities, unlike the varying/measured/dependent quantities that belong in `variables`. Coordinates values may be given by 1-dimensional arrays or scalars, in which case `dims` do not need to be supplied: 1D arrays will be assumed to give index values along the dimension with the same name. The following notations are accepted: - mapping {coord name: DataArray} - mapping {coord name: Variable} - mapping {coord name: (dimension name, array-like)} - mapping {coord name: (tuple of dimension names, array-like)} - mapping {dimension name: array-like} (the dimension name is implicitly set to be the same as the coord name) The last notation implies that the coord name is the same as the dimension name. attrs : dict-like, optional Global attributes to save on this dataset. Examples -------- Create data: >>> np.random.seed(0) >>> temperature = 15 + 8 * np.random.randn(2, 2, 3) >>> precipitation = 10 * np.random.rand(2, 2, 3) >>> lon = [[-99.83, -99.32], [-99.79, -99.23]] >>> lat = [[42.25, 42.21], [42.63, 42.59]] >>> time = pd.date_range("2014-09-06", periods=3) >>> reference_time = pd.Timestamp("2014-09-05") Initialize a dataset with multiple dimensions: >>> ds = xr.Dataset( ... data_vars=dict( ... temperature=(["x", "y", "time"], temperature), ... precipitation=(["x", "y", "time"], precipitation), ... ), ... coords=dict( ... lon=(["x", "y"], lon), ... lat=(["x", "y"], lat), ... time=time, ... reference_time=reference_time, ... ), ... attrs=dict(description="Weather related data."), ... ) >>> ds Dimensions: (time: 3, x: 2, y: 2) Coordinates: lon (x, y) float64 -99.83 -99.32 -99.79 -99.23 lat (x, y) float64 42.25 42.21 42.63 42.59 * time (time) datetime64[ns] 2014-09-06 2014-09-07 2014-09-08 reference_time datetime64[ns] 2014-09-05 Dimensions without coordinates: x, y Data variables: temperature (x, y, time) float64 29.11 18.2 22.83 ... 18.28 16.15 26.63 precipitation (x, y, time) float64 5.68 9.256 0.7104 ... 7.992 4.615 7.805 Attributes: description: Weather related data. Find out where the coldest temperature was and what values the other variables had: >>> ds.isel(ds.temperature.argmin(...)) Dimensions: () Coordinates: lon float64 -99.32 lat float64 42.21 time datetime64[ns] 2014-09-08 reference_time datetime64[ns] 2014-09-05 Data variables: temperature float64 7.182 precipitation float64 8.326 Attributes: description: Weather related data. """ _attrs: Optional[Dict[Hashable, Any]] _cache: Dict[str, Any] _coord_names: Set[Hashable] _dims: Dict[Hashable, int] _encoding: Optional[Dict[Hashable, Any]] _indexes: Optional[Dict[Hashable, pd.Index]] _variables: Dict[Hashable, Variable] __slots__ = ( "_attrs", "_cache", "_coord_names", "_dims", "_encoding", "_file_obj", "_indexes", "_variables", "__weakref__", ) _groupby_cls = groupby.DatasetGroupBy _rolling_cls = rolling.DatasetRolling _coarsen_cls = rolling.DatasetCoarsen _resample_cls = resample.DatasetResample _weighted_cls = weighted.DatasetWeighted def __init__( self, # could make a VariableArgs to use more generally, and refine these # categories data_vars: Mapping[Hashable, Any] = None, coords: Mapping[Hashable, Any] = None, attrs: Mapping[Hashable, Any] = None, ): # TODO(shoyer): expose indexes as a public argument in __init__ if data_vars is None: data_vars = {} if coords is None: coords = {} both_data_and_coords = set(data_vars) & set(coords) if both_data_and_coords: raise ValueError( "variables %r are found in both data_vars and coords" % both_data_and_coords ) if isinstance(coords, Dataset): coords = coords.variables variables, coord_names, dims, indexes, _ = merge_data_and_coords( data_vars, coords, compat="broadcast_equals" ) self._attrs = dict(attrs) if attrs is not None else None self._file_obj = None self._encoding = None self._variables = variables self._coord_names = coord_names self._dims = dims self._indexes = indexes @classmethod def load_store(cls, store, decoder=None) -> "Dataset": """Create a new dataset from the contents of a backends.*DataStore object """ variables, attributes = store.load() if decoder: variables, attributes = decoder(variables, attributes) obj = cls(variables, attrs=attributes) obj._file_obj = store return obj @property def variables(self) -> Mapping[Hashable, Variable]: """Low level interface to Dataset contents as dict of Variable objects. This ordered dictionary is frozen to prevent mutation that could violate Dataset invariants. It contains all variable objects constituting the Dataset, including both data variables and coordinates. """ return Frozen(self._variables) @property def attrs(self) -> Dict[Hashable, Any]: """Dictionary of global attributes on this dataset""" if self._attrs is None: self._attrs = {} return self._attrs @attrs.setter def attrs(self, value: Mapping[Hashable, Any]) -> None: self._attrs = dict(value) @property def encoding(self) -> Dict: """Dictionary of global encoding attributes on this dataset""" if self._encoding is None: self._encoding = {} return self._encoding @encoding.setter def encoding(self, value: Mapping) -> None: self._encoding = dict(value) @property def dims(self) -> Mapping[Hashable, int]: """Mapping from dimension names to lengths. Cannot be modified directly, but is updated when adding new variables. Note that type of this object differs from `DataArray.dims`. See `Dataset.sizes` and `DataArray.sizes` for consistently named properties. """ return Frozen(SortedKeysDict(self._dims)) @property def sizes(self) -> Mapping[Hashable, int]: """Mapping from dimension names to lengths. Cannot be modified directly, but is updated when adding new variables. This is an alias for `Dataset.dims` provided for the benefit of consistency with `DataArray.sizes`. See also -------- DataArray.sizes """ return self.dims def load(self, **kwargs) -> "Dataset": """Manually trigger loading and/or computation of this dataset's data from disk or a remote source into memory and return this dataset. Unlike compute, the original dataset is modified and returned. Normally, it should not be necessary to call this method in user code, because all xarray functions should either work on deferred data or load data automatically. However, this method can be necessary when working with many file objects on disk. Parameters ---------- **kwargs : dict Additional keyword arguments passed on to ``dask.array.compute``. See Also -------- dask.array.compute """ # access .data to coerce everything to numpy or dask arrays lazy_data = { k: v._data for k, v in self.variables.items() if is_duck_dask_array(v._data) } if lazy_data: import dask.array as da # evaluate all the dask arrays simultaneously evaluated_data = da.compute(*lazy_data.values(), **kwargs) for k, data in zip(lazy_data, evaluated_data): self.variables[k].data = data # load everything else sequentially for k, v in self.variables.items(): if k not in lazy_data: v.load() return self def __dask_tokenize__(self): from dask.base import normalize_token return normalize_token( (type(self), self._variables, self._coord_names, self._attrs) ) def __dask_graph__(self): graphs = {k: v.__dask_graph__() for k, v in self.variables.items()} graphs = {k: v for k, v in graphs.items() if v is not None} if not graphs: return None else: try: from dask.highlevelgraph import HighLevelGraph return HighLevelGraph.merge(*graphs.values()) except ImportError: from dask import sharedict return sharedict.merge(*graphs.values()) def __dask_keys__(self): import dask return [ v.__dask_keys__() for v in self.variables.values() if dask.is_dask_collection(v) ] def __dask_layers__(self): import dask return sum( [ v.__dask_layers__() for v in self.variables.values() if dask.is_dask_collection(v) ], (), ) @property def __dask_optimize__(self): import dask.array as da return da.Array.__dask_optimize__ @property def __dask_scheduler__(self): import dask.array as da return da.Array.__dask_scheduler__ def __dask_postcompute__(self): import dask info = [ (True, k, v.__dask_postcompute__()) if dask.is_dask_collection(v) else (False, k, v) for k, v in self._variables.items() ] args = ( info, self._coord_names, self._dims, self._attrs, self._indexes, self._encoding, self._file_obj, ) return self._dask_postcompute, args def __dask_postpersist__(self): import dask info = [ (True, k, v.__dask_postpersist__()) if dask.is_dask_collection(v) else (False, k, v) for k, v in self._variables.items() ] args = ( info, self._coord_names, self._dims, self._attrs, self._indexes, self._encoding, self._file_obj, ) return self._dask_postpersist, args @staticmethod def _dask_postcompute(results, info, *args): variables = {} results2 = list(results[::-1]) for is_dask, k, v in info: if is_dask: func, args2 = v r = results2.pop() result = func(r, *args2) else: result = v variables[k] = result final = Dataset._construct_direct(variables, *args) return final @staticmethod def _dask_postpersist(dsk, info, *args): variables = {} # postpersist is called in both dask.optimize and dask.persist # When persisting, we want to filter out unrelated keys for # each Variable's task graph. is_persist = len(dsk) == len(info) for is_dask, k, v in info: if is_dask: func, args2 = v if is_persist: name = args2[1][0] dsk2 = {k: v for k, v in dsk.items() if k[0] == name} else: dsk2 = dsk result = func(dsk2, *args2) else: result = v variables[k] = result return Dataset._construct_direct(variables, *args) def compute(self, **kwargs) -> "Dataset": """Manually trigger loading and/or computation of this dataset's data from disk or a remote source into memory and return a new dataset. Unlike load, the original dataset is left unaltered. Normally, it should not be necessary to call this method in user code, because all xarray functions should either work on deferred data or load data automatically. However, this method can be necessary when working with many file objects on disk. Parameters ---------- **kwargs : dict Additional keyword arguments passed on to ``dask.array.compute``. See Also -------- dask.array.compute """ new = self.copy(deep=False) return new.load(**kwargs) def _persist_inplace(self, **kwargs) -> "Dataset": """Persist all Dask arrays in memory""" # access .data to coerce everything to numpy or dask arrays lazy_data = { k: v._data for k, v in self.variables.items() if is_duck_dask_array(v._data) } if lazy_data: import dask # evaluate all the dask arrays simultaneously evaluated_data = dask.persist(*lazy_data.values(), **kwargs) for k, data in zip(lazy_data, evaluated_data): self.variables[k].data = data return self def persist(self, **kwargs) -> "Dataset": """Trigger computation, keeping data as dask arrays This operation can be used to trigger computation on underlying dask arrays, similar to ``.compute()`` or ``.load()``. However this operation keeps the data as dask arrays. This is particularly useful when using the dask.distributed scheduler and you want to load a large amount of data into distributed memory. Parameters ---------- **kwargs : dict Additional keyword arguments passed on to ``dask.persist``. See Also -------- dask.persist """ new = self.copy(deep=False) return new._persist_inplace(**kwargs) @classmethod def _construct_direct( cls, variables, coord_names, dims=None, attrs=None, indexes=None, encoding=None, file_obj=None, ): """Shortcut around __init__ for internal use when we want to skip costly validation """ if dims is None: dims = calculate_dimensions(variables) obj = object.__new__(cls) obj._variables = variables obj._coord_names = coord_names obj._dims = dims obj._indexes = indexes obj._attrs = attrs obj._file_obj = file_obj obj._encoding = encoding return obj def _replace( self, variables: Dict[Hashable, Variable] = None, coord_names: Set[Hashable] = None, dims: Dict[Any, int] = None, attrs: Union[Dict[Hashable, Any], None, Default] = _default, indexes: Union[Dict[Any, pd.Index], None, Default] = _default, encoding: Union[dict, None, Default] = _default, inplace: bool = False, ) -> "Dataset": """Fastpath constructor for internal use. Returns an object with optionally with replaced attributes. Explicitly passed arguments are *not* copied when placed on the new dataset. It is up to the caller to ensure that they have the right type and are not used elsewhere. """ if inplace: if variables is not None: self._variables = variables if coord_names is not None: self._coord_names = coord_names if dims is not None: self._dims = dims if attrs is not _default: self._attrs = attrs if indexes is not _default: self._indexes = indexes if encoding is not _default: self._encoding = encoding obj = self else: if variables is None: variables = self._variables.copy() if coord_names is None: coord_names = self._coord_names.copy() if dims is None: dims = self._dims.copy() if attrs is _default: attrs = copy.copy(self._attrs) if indexes is _default: indexes = copy.copy(self._indexes) if encoding is _default: encoding = copy.copy(self._encoding) obj = self._construct_direct( variables, coord_names, dims, attrs, indexes, encoding ) return obj def _replace_with_new_dims( self, variables: Dict[Hashable, Variable], coord_names: set = None, attrs: Union[Dict[Hashable, Any], None, Default] = _default, indexes: Union[Dict[Hashable, pd.Index], None, Default] = _default, inplace: bool = False, ) -> "Dataset": """Replace variables with recalculated dimensions.""" dims = calculate_dimensions(variables) return self._replace( variables, coord_names, dims, attrs, indexes, inplace=inplace ) def _replace_vars_and_dims( self, variables: Dict[Hashable, Variable], coord_names: set = None, dims: Dict[Hashable, int] = None, attrs: Union[Dict[Hashable, Any], None, Default] = _default, inplace: bool = False, ) -> "Dataset": """Deprecated version of _replace_with_new_dims(). Unlike _replace_with_new_dims(), this method always recalculates indexes from variables. """ if dims is None: dims = calculate_dimensions(variables) return self._replace( variables, coord_names, dims, attrs, indexes=None, inplace=inplace ) def _overwrite_indexes(self, indexes: Mapping[Any, pd.Index]) -> "Dataset": if not indexes: return self variables = self._variables.copy() new_indexes = dict(self.indexes) for name, idx in indexes.items(): variables[name] = IndexVariable(name, idx) new_indexes[name] = idx obj = self._replace(variables, indexes=new_indexes) # switch from dimension to level names, if necessary dim_names: Dict[Hashable, str] = {} for dim, idx in indexes.items(): if not isinstance(idx, pd.MultiIndex) and idx.name != dim: dim_names[dim] = idx.name if dim_names: obj = obj.rename(dim_names) return obj def copy(self, deep: bool = False, data: Mapping = None) -> "Dataset": """Returns a copy of this dataset. If `deep=True`, a deep copy is made of each of the component variables. Otherwise, a shallow copy of each of the component variable is made, so that the underlying memory region of the new dataset is the same as in the original dataset. Use `data` to create a new object with the same structure as original but entirely new data. Parameters ---------- deep : bool, optional Whether each component variable is loaded into memory and copied onto the new object. Default is False. data : dict-like, optional Data to use in the new object. Each item in `data` must have same shape as corresponding data variable in original. When `data` is used, `deep` is ignored for the data variables and only used for coords. Returns ------- object : Dataset New object with dimensions, attributes, coordinates, name, encoding, and optionally data copied from original. Examples -------- Shallow copy versus deep copy >>> da = xr.DataArray(np.random.randn(2, 3)) >>> ds = xr.Dataset( ... {"foo": da, "bar": ("x", [-1, 2])}, ... coords={"x": ["one", "two"]}, ... ) >>> ds.copy() Dimensions: (dim_0: 2, dim_1: 3, x: 2) Coordinates: * x (x) >> ds_0 = ds.copy(deep=False) >>> ds_0["foo"][0, 0] = 7 >>> ds_0 Dimensions: (dim_0: 2, dim_1: 3, x: 2) Coordinates: * x (x) >> ds Dimensions: (dim_0: 2, dim_1: 3, x: 2) Coordinates: * x (x) >> ds.copy(data={"foo": np.arange(6).reshape(2, 3), "bar": ["a", "b"]}) Dimensions: (dim_0: 2, dim_1: 3, x: 2) Coordinates: * x (x) >> ds Dimensions: (dim_0: 2, dim_1: 3, x: 2) Coordinates: * x (x) Dict[str, Hashable]: """Return a mapping of all MultiIndex levels and their corresponding coordinate name. """ level_coords: Dict[str, Hashable] = {} for name, index in self.indexes.items(): if isinstance(index, pd.MultiIndex): level_names = index.names (dim,) = self.variables[name].dims level_coords.update({lname: dim for lname in level_names}) return level_coords def _copy_listed(self, names: Iterable[Hashable]) -> "Dataset": """Create a new Dataset with the listed variables from this dataset and the all relevant coordinates. Skips all validation. """ variables: Dict[Hashable, Variable] = {} coord_names = set() indexes: Dict[Hashable, pd.Index] = {} for name in names: try: variables[name] = self._variables[name] except KeyError: ref_name, var_name, var = _get_virtual_variable( self._variables, name, self._level_coords, self.dims ) variables[var_name] = var if ref_name in self._coord_names or ref_name in self.dims: coord_names.add(var_name) if (var_name,) == var.dims: indexes[var_name] = var.to_index() needed_dims: Set[Hashable] = set() for v in variables.values(): needed_dims.update(v.dims) dims = {k: self.dims[k] for k in needed_dims} # preserves ordering of coordinates for k in self._variables: if k not in self._coord_names: continue if set(self.variables[k].dims) <= needed_dims: variables[k] = self._variables[k] coord_names.add(k) if k in self.indexes: indexes[k] = self.indexes[k] return self._replace(variables, coord_names, dims, indexes=indexes) def _construct_dataarray(self, name: Hashable) -> "DataArray": """Construct a DataArray by indexing this dataset""" from .dataarray import DataArray try: variable = self._variables[name] except KeyError: _, name, variable = _get_virtual_variable( self._variables, name, self._level_coords, self.dims ) needed_dims = set(variable.dims) coords: Dict[Hashable, Variable] = {} for k in self.coords: if set(self.variables[k].dims) <= needed_dims: coords[k] = self.variables[k] if self._indexes is None: indexes = None else: indexes = {k: v for k, v in self._indexes.items() if k in coords} return DataArray(variable, coords, name=name, indexes=indexes, fastpath=True) def __copy__(self) -> "Dataset": return self.copy(deep=False) def __deepcopy__(self, memo=None) -> "Dataset": # memo does nothing but is required for compatibility with # copy.deepcopy return self.copy(deep=True) @property def _attr_sources(self) -> List[Mapping[Hashable, Any]]: """List of places to look-up items for attribute-style access""" return self._item_sources + [self.attrs] @property def _item_sources(self) -> List[Mapping[Hashable, Any]]: """List of places to look-up items for key-completion""" return [ self.data_vars, self.coords, {d: self[d] for d in self.dims}, LevelCoordinatesSource(self), ] def __contains__(self, key: object) -> bool: """The 'in' operator will return true or false depending on whether 'key' is an array in the dataset or not. """ return key in self._variables def __len__(self) -> int: return len(self.data_vars) def __bool__(self) -> bool: return bool(self.data_vars) def __iter__(self) -> Iterator[Hashable]: return iter(self.data_vars) def __array__(self, dtype=None): raise TypeError( "cannot directly convert an xarray.Dataset into a " "numpy array. Instead, create an xarray.DataArray " "first, either with indexing on the Dataset or by " "invoking the `to_array()` method." ) @property def nbytes(self) -> int: return sum(v.nbytes for v in self.variables.values()) @property def loc(self) -> _LocIndexer: """Attribute for location based indexing. Only supports __getitem__, and only when the key is a dict of the form {dim: labels}. """ return _LocIndexer(self) # FIXME https://github.com/python/mypy/issues/7328 @overload def __getitem__(self, key: Mapping) -> "Dataset": # type: ignore ... @overload def __getitem__(self, key: Hashable) -> "DataArray": # type: ignore ... @overload def __getitem__(self, key: Any) -> "Dataset": ... def __getitem__(self, key): """Access variables or coordinates this dataset as a :py:class:`~xarray.DataArray`. Indexing with a list of names will return a new ``Dataset`` object. """ if utils.is_dict_like(key): return self.isel(**cast(Mapping, key)) if hashable(key): return self._construct_dataarray(key) else: return self._copy_listed(np.asarray(key)) def __setitem__(self, key: Hashable, value) -> None: """Add an array to this dataset. If value is a `DataArray`, call its `select_vars()` method, rename it to `key` and merge the contents of the resulting dataset into this dataset. If value is an `Variable` object (or tuple of form ``(dims, data[, attrs])``), add it to this dataset as a new variable. """ if utils.is_dict_like(key): raise NotImplementedError( "cannot yet use a dictionary as a key to set Dataset values" ) self.update({key: value}) def __delitem__(self, key: Hashable) -> None: """Remove a variable from this dataset.""" del self._variables[key] self._coord_names.discard(key) if key in self.indexes: assert self._indexes is not None del self._indexes[key] self._dims = calculate_dimensions(self._variables) # mutable objects should not be hashable # https://github.com/python/mypy/issues/4266 __hash__ = None # type: ignore def _all_compat(self, other: "Dataset", compat_str: str) -> bool: """Helper function for equals and identical""" # some stores (e.g., scipy) do not seem to preserve order, so don't # require matching order for equality def compat(x: Variable, y: Variable) -> bool: return getattr(x, compat_str)(y) return self._coord_names == other._coord_names and utils.dict_equiv( self._variables, other._variables, compat=compat ) def broadcast_equals(self, other: "Dataset") -> bool: """Two Datasets are broadcast equal if they are equal after broadcasting all variables against each other. For example, variables that are scalar in one dataset but non-scalar in the other dataset can still be broadcast equal if the the non-scalar variable is a constant. See Also -------- Dataset.equals Dataset.identical """ try: return self._all_compat(other, "broadcast_equals") except (TypeError, AttributeError): return False def equals(self, other: "Dataset") -> bool: """Two Datasets are equal if they have matching variables and coordinates, all of which are equal. Datasets can still be equal (like pandas objects) if they have NaN values in the same locations. This method is necessary because `v1 == v2` for ``Dataset`` does element-wise comparisons (like numpy.ndarrays). See Also -------- Dataset.broadcast_equals Dataset.identical """ try: return self._all_compat(other, "equals") except (TypeError, AttributeError): return False def identical(self, other: "Dataset") -> bool: """Like equals, but also checks all dataset attributes and the attributes on all variables and coordinates. See Also -------- Dataset.broadcast_equals Dataset.equals """ try: return utils.dict_equiv(self.attrs, other.attrs) and self._all_compat( other, "identical" ) except (TypeError, AttributeError): return False @property def indexes(self) -> Indexes: """Mapping of pandas.Index objects used for label based indexing""" if self._indexes is None: self._indexes = default_indexes(self._variables, self._dims) return Indexes(self._indexes) @property def coords(self) -> DatasetCoordinates: """Dictionary of xarray.DataArray objects corresponding to coordinate variables """ return DatasetCoordinates(self) @property def data_vars(self) -> DataVariables: """Dictionary of DataArray objects corresponding to data variables""" return DataVariables(self) def set_coords( self, names: "Union[Hashable, Iterable[Hashable]]", inplace: bool = None ) -> "Dataset": """Given names of one or more variables, set them as coordinates Parameters ---------- names : hashable or iterable of hashable Name(s) of variables in this dataset to convert into coordinates. Returns ------- Dataset See also -------- Dataset.swap_dims """ # TODO: allow inserting new coordinates with this method, like # DataFrame.set_index? # nb. check in self._variables, not self.data_vars to insure that the # operation is idempotent _check_inplace(inplace) if isinstance(names, str) or not isinstance(names, Iterable): names = [names] else: names = list(names) self._assert_all_in_dataset(names) obj = self.copy() obj._coord_names.update(names) return obj def reset_coords( self, names: "Union[Hashable, Iterable[Hashable], None]" = None, drop: bool = False, inplace: bool = None, ) -> "Dataset": """Given names of coordinates, reset them to become variables Parameters ---------- names : hashable or iterable of hashable, optional Name(s) of non-index coordinates in this dataset to reset into variables. By default, all non-index coordinates are reset. drop : bool, optional If True, remove coordinates instead of converting them into variables. Returns ------- Dataset """ _check_inplace(inplace) if names is None: names = self._coord_names - set(self.dims) else: if isinstance(names, str) or not isinstance(names, Iterable): names = [names] else: names = list(names) self._assert_all_in_dataset(names) bad_coords = set(names) & set(self.dims) if bad_coords: raise ValueError( "cannot remove index coordinates with reset_coords: %s" % bad_coords ) obj = self.copy() obj._coord_names.difference_update(names) if drop: for name in names: del obj._variables[name] return obj def dump_to_store(self, store: "AbstractDataStore", **kwargs) -> None: """Store dataset contents to a backends.*DataStore object.""" from ..backends.api import dump_to_store # TODO: rename and/or cleanup this method to make it more consistent # with to_netcdf() dump_to_store(self, store, **kwargs) def to_netcdf( self, path=None, mode: str = "w", format: str = None, group: str = None, engine: str = None, encoding: Mapping = None, unlimited_dims: Iterable[Hashable] = None, compute: bool = True, invalid_netcdf: bool = False, ) -> Union[bytes, "Delayed", None]: """Write dataset contents to a netCDF file. Parameters ---------- path : str, Path or file-like, optional Path to which to save this dataset. File-like objects are only supported by the scipy engine. If no path is provided, this function returns the resulting netCDF file as bytes; in this case, we need to use scipy, which does not support netCDF version 4 (the default format becomes NETCDF3_64BIT). mode : {"w", "a"}, default: "w" Write ('w') or append ('a') mode. If mode='w', any existing file at this location will be overwritten. If mode='a', existing variables will be overwritten. format : {"NETCDF4", "NETCDF4_CLASSIC", "NETCDF3_64BIT", \ "NETCDF3_CLASSIC"}, optional File format for the resulting netCDF file: * NETCDF4: Data is stored in an HDF5 file, using netCDF4 API features. * NETCDF4_CLASSIC: Data is stored in an HDF5 file, using only netCDF 3 compatible API features. * NETCDF3_64BIT: 64-bit offset version of the netCDF 3 file format, which fully supports 2+ GB files, but is only compatible with clients linked against netCDF version 3.6.0 or later. * NETCDF3_CLASSIC: The classic netCDF 3 file format. It does not handle 2+ GB files very well. All formats are supported by the netCDF4-python library. scipy.io.netcdf only supports the last two formats. The default format is NETCDF4 if you are saving a file to disk and have the netCDF4-python library available. Otherwise, xarray falls back to using scipy to write netCDF files and defaults to the NETCDF3_64BIT format (scipy does not support netCDF4). group : str, optional Path to the netCDF4 group in the given file to open (only works for format='NETCDF4'). The group(s) will be created if necessary. engine : {"netcdf4", "scipy", "h5netcdf"}, optional Engine to use when writing netCDF files. If not provided, the default engine is chosen based on available dependencies, with a preference for 'netcdf4' if writing to a file on disk. encoding : dict, optional Nested dictionary with variable names as keys and dictionaries of variable specific encodings as values, e.g., ``{"my_variable": {"dtype": "int16", "scale_factor": 0.1, "zlib": True}, ...}`` The `h5netcdf` engine supports both the NetCDF4-style compression encoding parameters ``{"zlib": True, "complevel": 9}`` and the h5py ones ``{"compression": "gzip", "compression_opts": 9}``. This allows using any compression plugin installed in the HDF5 library, e.g. LZF. unlimited_dims : iterable of hashable, optional Dimension(s) that should be serialized as unlimited dimensions. By default, no dimensions are treated as unlimited dimensions. Note that unlimited_dims may also be set via ``dataset.encoding["unlimited_dims"]``. compute: bool, default: True If true compute immediately, otherwise return a ``dask.delayed.Delayed`` object that can be computed later. invalid_netcdf: bool, default: False Only valid along with ``engine="h5netcdf"``. If True, allow writing hdf5 files which are invalid netcdf as described in https://github.com/shoyer/h5netcdf. """ if encoding is None: encoding = {} from ..backends.api import to_netcdf return to_netcdf( self, path, mode, format=format, group=group, engine=engine, encoding=encoding, unlimited_dims=unlimited_dims, compute=compute, invalid_netcdf=invalid_netcdf, ) def to_zarr( self, store: Union[MutableMapping, str, Path] = None, chunk_store: Union[MutableMapping, str, Path] = None, mode: str = None, synchronizer=None, group: str = None, encoding: Mapping = None, compute: bool = True, consolidated: bool = False, append_dim: Hashable = None, region: Mapping[str, slice] = None, ) -> "ZarrStore": """Write dataset contents to a zarr group. .. note:: Experimental The Zarr backend is new and experimental. Please report any unexpected behavior via github issues. Parameters ---------- store : MutableMapping, str or Path, optional Store or path to directory in file system. chunk_store : MutableMapping, str or Path, optional Store or path to directory in file system only for Zarr array chunks. Requires zarr-python v2.4.0 or later. mode : {"w", "w-", "a", None}, optional Persistence mode: "w" means create (overwrite if exists); "w-" means create (fail if exists); "a" means override existing variables (create if does not exist). If ``append_dim`` is set, ``mode`` can be omitted as it is internally set to ``"a"``. Otherwise, ``mode`` will default to `w-` if not set. synchronizer : object, optional Zarr array synchronizer. group : str, optional Group path. (a.k.a. `path` in zarr terminology.) encoding : dict, optional Nested dictionary with variable names as keys and dictionaries of variable specific encodings as values, e.g., ``{"my_variable": {"dtype": "int16", "scale_factor": 0.1,}, ...}`` compute: bool, optional If True write array data immediately, otherwise return a ``dask.delayed.Delayed`` object that can be computed to write array data later. Metadata is always updated eagerly. consolidated: bool, optional If True, apply zarr's `consolidate_metadata` function to the store after writing metadata. append_dim: hashable, optional If set, the dimension along which the data will be appended. All other dimensions on overriden variables must remain the same size. region: dict, optional Optional mapping from dimension names to integer slices along dataset dimensions to indicate the region of existing zarr array(s) in which to write this dataset's data. For example, ``{'x': slice(0, 1000), 'y': slice(10000, 11000)}`` would indicate that values should be written to the region ``0:1000`` along ``x`` and ``10000:11000`` along ``y``. Two restrictions apply to the use of ``region``: - If ``region`` is set, _all_ variables in a dataset must have at least one dimension in common with the region. Other variables should be written in a separate call to ``to_zarr()``. - Dimensions cannot be included in both ``region`` and ``append_dim`` at the same time. To create empty arrays to fill in with ``region``, use a separate call to ``to_zarr()`` with ``compute=False``. See "Appending to existing Zarr stores" in the reference documentation for full details. References ---------- https://zarr.readthedocs.io/ Notes ----- Zarr chunking behavior: If chunks are found in the encoding argument or attribute corresponding to any DataArray, those chunks are used. If a DataArray is a dask array, it is written with those chunks. If not other chunks are found, Zarr uses its own heuristics to choose automatic chunk sizes. """ from ..backends.api import to_zarr if encoding is None: encoding = {} return to_zarr( self, store=store, chunk_store=chunk_store, mode=mode, synchronizer=synchronizer, group=group, encoding=encoding, compute=compute, consolidated=consolidated, append_dim=append_dim, region=region, ) def __repr__(self) -> str: return formatting.dataset_repr(self) def _repr_html_(self): if OPTIONS["display_style"] == "text": return f"
{escape(repr(self))}
" return formatting_html.dataset_repr(self) def info(self, buf=None) -> None: """ Concise summary of a Dataset variables and attributes. Parameters ---------- buf : file-like, default: sys.stdout writable buffer See Also -------- pandas.DataFrame.assign ncdump: netCDF's ncdump """ if buf is None: # pragma: no cover buf = sys.stdout lines = [] lines.append("xarray.Dataset {") lines.append("dimensions:") for name, size in self.dims.items(): lines.append(f"\t{name} = {size} ;") lines.append("\nvariables:") for name, da in self.variables.items(): dims = ", ".join(da.dims) lines.append(f"\t{da.dtype} {name}({dims}) ;") for k, v in da.attrs.items(): lines.append(f"\t\t{name}:{k} = {v} ;") lines.append("\n// global attributes:") for k, v in self.attrs.items(): lines.append(f"\t:{k} = {v} ;") lines.append("}") buf.write("\n".join(lines)) @property def chunks(self) -> Mapping[Hashable, Tuple[int, ...]]: """Block dimensions for this dataset's data or None if it's not a dask array. """ chunks: Dict[Hashable, Tuple[int, ...]] = {} for v in self.variables.values(): if v.chunks is not None: for dim, c in zip(v.dims, v.chunks): if dim in chunks and c != chunks[dim]: raise ValueError( f"Object has inconsistent chunks along dimension {dim}. " "This can be fixed by calling unify_chunks()." ) chunks[dim] = c return Frozen(SortedKeysDict(chunks)) def chunk( self, chunks: Union[ None, Number, str, Mapping[Hashable, Union[None, Number, str, Tuple[Number, ...]]], ] = None, name_prefix: str = "xarray-", token: str = None, lock: bool = False, ) -> "Dataset": """Coerce all arrays in this dataset into dask arrays with the given chunks. Non-dask arrays in this dataset will be converted to dask arrays. Dask arrays will be rechunked to the given chunk sizes. If neither chunks is not provided for one or more dimensions, chunk sizes along that dimension will not be updated; non-dask arrays will be converted into dask arrays with a single block. Parameters ---------- chunks : int, 'auto' or mapping, optional Chunk sizes along each dimension, e.g., ``5`` or ``{"x": 5, "y": 5}``. name_prefix : str, optional Prefix for the name of any new dask arrays. token : str, optional Token uniquely identifying this dataset. lock : optional Passed on to :py:func:`dask.array.from_array`, if the array is not already as dask array. Returns ------- chunked : xarray.Dataset """ if isinstance(chunks, (Number, str)): chunks = dict.fromkeys(self.dims, chunks) if chunks is not None: bad_dims = chunks.keys() - self.dims.keys() if bad_dims: raise ValueError( "some chunks keys are not dimensions on this " "object: %s" % bad_dims ) variables = { k: _maybe_chunk(k, v, chunks, token, lock, name_prefix) for k, v in self.variables.items() } return self._replace(variables) def _validate_indexers( self, indexers: Mapping[Hashable, Any], missing_dims: str = "raise" ) -> Iterator[Tuple[Hashable, Union[int, slice, np.ndarray, Variable]]]: """Here we make sure + indexer has a valid keys + indexer is in a valid data type + string indexers are cast to the appropriate date type if the associated index is a DatetimeIndex or CFTimeIndex """ from .dataarray import DataArray indexers = drop_dims_from_indexers(indexers, self.dims, missing_dims) # all indexers should be int, slice, np.ndarrays, or Variable for k, v in indexers.items(): if isinstance(v, (int, slice, Variable)): yield k, v elif isinstance(v, DataArray): yield k, v.variable elif isinstance(v, tuple): yield k, as_variable(v) elif isinstance(v, Dataset): raise TypeError("cannot use a Dataset as an indexer") elif isinstance(v, Sequence) and len(v) == 0: yield k, np.empty((0,), dtype="int64") else: v = np.asarray(v) if v.dtype.kind in "US": index = self.indexes[k] if isinstance(index, pd.DatetimeIndex): v = v.astype("datetime64[ns]") elif isinstance(index, xr.CFTimeIndex): v = _parse_array_of_cftime_strings(v, index.date_type) if v.ndim > 1: raise IndexError( "Unlabeled multi-dimensional array cannot be " "used for indexing: {}".format(k) ) yield k, v def _validate_interp_indexers( self, indexers: Mapping[Hashable, Any] ) -> Iterator[Tuple[Hashable, Variable]]: """Variant of _validate_indexers to be used for interpolation""" for k, v in self._validate_indexers(indexers): if isinstance(v, Variable): if v.ndim == 1: yield k, v.to_index_variable() else: yield k, v elif isinstance(v, int): yield k, Variable((), v) elif isinstance(v, np.ndarray): if v.ndim == 0: yield k, Variable((), v) elif v.ndim == 1: yield k, IndexVariable((k,), v) else: raise AssertionError() # Already tested by _validate_indexers else: raise TypeError(type(v)) def _get_indexers_coords_and_indexes(self, indexers): """Extract coordinates and indexes from indexers. Only coordinate with a name different from any of self.variables will be attached. """ from .dataarray import DataArray coords_list = [] for k, v in indexers.items(): if isinstance(v, DataArray): if v.dtype.kind == "b": if v.ndim != 1: # we only support 1-d boolean array raise ValueError( "{:d}d-boolean array is used for indexing along " "dimension {!r}, but only 1d boolean arrays are " "supported.".format(v.ndim, k) ) # Make sure in case of boolean DataArray, its # coordinate also should be indexed. v_coords = v[v.values.nonzero()[0]].coords else: v_coords = v.coords coords_list.append(v_coords) # we don't need to call align() explicitly or check indexes for # alignment, because merge_variables already checks for exact alignment # between dimension coordinates coords, indexes = merge_coordinates_without_align(coords_list) assert_coordinate_consistent(self, coords) # silently drop the conflicted variables. attached_coords = {k: v for k, v in coords.items() if k not in self._variables} attached_indexes = { k: v for k, v in indexes.items() if k not in self._variables } return attached_coords, attached_indexes def isel( self, indexers: Mapping[Hashable, Any] = None, drop: bool = False, missing_dims: str = "raise", **indexers_kwargs: Any, ) -> "Dataset": """Returns a new dataset with each array indexed along the specified dimension(s). This method selects values from each array using its `__getitem__` method, except this method does not require knowing the order of each array's dimensions. Parameters ---------- indexers : dict, optional A dict with keys matching dimensions and values given by integers, slice objects or arrays. indexer can be a integer, slice, array-like or DataArray. If DataArrays are passed as indexers, xarray-style indexing will be carried out. See :ref:`indexing` for the details. One of indexers or indexers_kwargs must be provided. drop : bool, optional If ``drop=True``, drop coordinates variables indexed by integers instead of making them scalar. missing_dims : {"raise", "warn", "ignore"}, default: "raise" What to do if dimensions that should be selected from are not present in the Dataset: - "raise": raise an exception - "warning": raise a warning, and ignore the missing dimensions - "ignore": ignore the missing dimensions **indexers_kwargs : {dim: indexer, ...}, optional The keyword arguments form of ``indexers``. One of indexers or indexers_kwargs must be provided. Returns ------- obj : Dataset A new Dataset with the same contents as this dataset, except each array and dimension is indexed by the appropriate indexers. If indexer DataArrays have coordinates that do not conflict with this object, then these coordinates will be attached. In general, each array's data will be a view of the array's data in this dataset, unless vectorized indexing was triggered by using an array indexer, in which case the data will be a copy. See Also -------- Dataset.sel DataArray.isel """ indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "isel") if any(is_fancy_indexer(idx) for idx in indexers.values()): return self._isel_fancy(indexers, drop=drop, missing_dims=missing_dims) # Much faster algorithm for when all indexers are ints, slices, one-dimensional # lists, or zero or one-dimensional np.ndarray's indexers = drop_dims_from_indexers(indexers, self.dims, missing_dims) variables = {} dims: Dict[Hashable, Tuple[int, ...]] = {} coord_names = self._coord_names.copy() indexes = self._indexes.copy() if self._indexes is not None else None for var_name, var_value in self._variables.items(): var_indexers = {k: v for k, v in indexers.items() if k in var_value.dims} if var_indexers: var_value = var_value.isel(var_indexers) if drop and var_value.ndim == 0 and var_name in coord_names: coord_names.remove(var_name) if indexes: indexes.pop(var_name, None) continue if indexes and var_name in indexes: if var_value.ndim == 1: indexes[var_name] = var_value.to_index() else: del indexes[var_name] variables[var_name] = var_value dims.update(zip(var_value.dims, var_value.shape)) return self._construct_direct( variables=variables, coord_names=coord_names, dims=dims, attrs=self._attrs, indexes=indexes, encoding=self._encoding, file_obj=self._file_obj, ) def _isel_fancy( self, indexers: Mapping[Hashable, Any], *, drop: bool, missing_dims: str = "raise", ) -> "Dataset": # Note: we need to preserve the original indexers variable in order to merge the # coords below indexers_list = list(self._validate_indexers(indexers, missing_dims)) variables: Dict[Hashable, Variable] = {} indexes: Dict[Hashable, pd.Index] = {} for name, var in self.variables.items(): var_indexers = {k: v for k, v in indexers_list if k in var.dims} if drop and name in var_indexers: continue # drop this variable if name in self.indexes: new_var, new_index = isel_variable_and_index( name, var, self.indexes[name], var_indexers ) if new_index is not None: indexes[name] = new_index elif var_indexers: new_var = var.isel(indexers=var_indexers) else: new_var = var.copy(deep=False) variables[name] = new_var coord_names = self._coord_names & variables.keys() selected = self._replace_with_new_dims(variables, coord_names, indexes) # Extract coordinates from indexers coord_vars, new_indexes = selected._get_indexers_coords_and_indexes(indexers) variables.update(coord_vars) indexes.update(new_indexes) coord_names = self._coord_names & variables.keys() | coord_vars.keys() return self._replace_with_new_dims(variables, coord_names, indexes=indexes) def sel( self, indexers: Mapping[Hashable, Any] = None, method: str = None, tolerance: Number = None, drop: bool = False, **indexers_kwargs: Any, ) -> "Dataset": """Returns a new dataset with each array indexed by tick labels along the specified dimension(s). In contrast to `Dataset.isel`, indexers for this method should use labels instead of integers. Under the hood, this method is powered by using pandas's powerful Index objects. This makes label based indexing essentially just as fast as using integer indexing. It also means this method uses pandas's (well documented) logic for indexing. This means you can use string shortcuts for datetime indexes (e.g., '2000-01' to select all values in January 2000). It also means that slices are treated as inclusive of both the start and stop values, unlike normal Python indexing. Parameters ---------- indexers : dict, optional A dict with keys matching dimensions and values given by scalars, slices or arrays of tick labels. For dimensions with multi-index, the indexer may also be a dict-like object with keys matching index level names. If DataArrays are passed as indexers, xarray-style indexing will be carried out. See :ref:`indexing` for the details. One of indexers or indexers_kwargs must be provided. method : {None, "nearest", "pad", "ffill", "backfill", "bfill"}, optional Method to use for inexact matches: * None (default): only exact matches * pad / ffill: propagate last valid index value forward * backfill / bfill: propagate next valid index value backward * nearest: use nearest valid index value tolerance : optional Maximum distance between original and new labels for inexact matches. The values of the index at the matching locations must satisfy the equation ``abs(index[indexer] - target) <= tolerance``. drop : bool, optional If ``drop=True``, drop coordinates variables in `indexers` instead of making them scalar. **indexers_kwargs : {dim: indexer, ...}, optional The keyword arguments form of ``indexers``. One of indexers or indexers_kwargs must be provided. Returns ------- obj : Dataset A new Dataset with the same contents as this dataset, except each variable and dimension is indexed by the appropriate indexers. If indexer DataArrays have coordinates that do not conflict with this object, then these coordinates will be attached. In general, each array's data will be a view of the array's data in this dataset, unless vectorized indexing was triggered by using an array indexer, in which case the data will be a copy. See Also -------- Dataset.isel DataArray.sel """ indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "sel") pos_indexers, new_indexes = remap_label_indexers( self, indexers=indexers, method=method, tolerance=tolerance ) result = self.isel(indexers=pos_indexers, drop=drop) return result._overwrite_indexes(new_indexes) def head( self, indexers: Union[Mapping[Hashable, int], int] = None, **indexers_kwargs: Any, ) -> "Dataset": """Returns a new dataset with the first `n` values of each array for the specified dimension(s). Parameters ---------- indexers : dict or int, default: 5 A dict with keys matching dimensions and integer values `n` or a single integer `n` applied over all dimensions. One of indexers or indexers_kwargs must be provided. **indexers_kwargs : {dim: n, ...}, optional The keyword arguments form of ``indexers``. One of indexers or indexers_kwargs must be provided. See Also -------- Dataset.tail Dataset.thin DataArray.head """ if not indexers_kwargs: if indexers is None: indexers = 5 if not isinstance(indexers, int) and not is_dict_like(indexers): raise TypeError("indexers must be either dict-like or a single integer") if isinstance(indexers, int): indexers = {dim: indexers for dim in self.dims} indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "head") for k, v in indexers.items(): if not isinstance(v, int): raise TypeError( "expected integer type indexer for " "dimension %r, found %r" % (k, type(v)) ) elif v < 0: raise ValueError( "expected positive integer as indexer " "for dimension %r, found %s" % (k, v) ) indexers_slices = {k: slice(val) for k, val in indexers.items()} return self.isel(indexers_slices) def tail( self, indexers: Union[Mapping[Hashable, int], int] = None, **indexers_kwargs: Any, ) -> "Dataset": """Returns a new dataset with the last `n` values of each array for the specified dimension(s). Parameters ---------- indexers : dict or int, default: 5 A dict with keys matching dimensions and integer values `n` or a single integer `n` applied over all dimensions. One of indexers or indexers_kwargs must be provided. **indexers_kwargs : {dim: n, ...}, optional The keyword arguments form of ``indexers``. One of indexers or indexers_kwargs must be provided. See Also -------- Dataset.head Dataset.thin DataArray.tail """ if not indexers_kwargs: if indexers is None: indexers = 5 if not isinstance(indexers, int) and not is_dict_like(indexers): raise TypeError("indexers must be either dict-like or a single integer") if isinstance(indexers, int): indexers = {dim: indexers for dim in self.dims} indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "tail") for k, v in indexers.items(): if not isinstance(v, int): raise TypeError( "expected integer type indexer for " "dimension %r, found %r" % (k, type(v)) ) elif v < 0: raise ValueError( "expected positive integer as indexer " "for dimension %r, found %s" % (k, v) ) indexers_slices = { k: slice(-val, None) if val != 0 else slice(val) for k, val in indexers.items() } return self.isel(indexers_slices) def thin( self, indexers: Union[Mapping[Hashable, int], int] = None, **indexers_kwargs: Any, ) -> "Dataset": """Returns a new dataset with each array indexed along every `n`-th value for the specified dimension(s) Parameters ---------- indexers : dict or int A dict with keys matching dimensions and integer values `n` or a single integer `n` applied over all dimensions. One of indexers or indexers_kwargs must be provided. **indexers_kwargs : {dim: n, ...}, optional The keyword arguments form of ``indexers``. One of indexers or indexers_kwargs must be provided. See Also -------- Dataset.head Dataset.tail DataArray.thin """ if ( not indexers_kwargs and not isinstance(indexers, int) and not is_dict_like(indexers) ): raise TypeError("indexers must be either dict-like or a single integer") if isinstance(indexers, int): indexers = {dim: indexers for dim in self.dims} indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "thin") for k, v in indexers.items(): if not isinstance(v, int): raise TypeError( "expected integer type indexer for " "dimension %r, found %r" % (k, type(v)) ) elif v < 0: raise ValueError( "expected positive integer as indexer " "for dimension %r, found %s" % (k, v) ) elif v == 0: raise ValueError("step cannot be zero") indexers_slices = {k: slice(None, None, val) for k, val in indexers.items()} return self.isel(indexers_slices) def broadcast_like( self, other: Union["Dataset", "DataArray"], exclude: Iterable[Hashable] = None ) -> "Dataset": """Broadcast this DataArray against another Dataset or DataArray. This is equivalent to xr.broadcast(other, self)[1] Parameters ---------- other : Dataset or DataArray Object against which to broadcast this array. exclude : iterable of hashable, optional Dimensions that must not be broadcasted """ if exclude is None: exclude = set() else: exclude = set(exclude) args = align(other, self, join="outer", copy=False, exclude=exclude) dims_map, common_coords = _get_broadcast_dims_map_common_coords(args, exclude) return _broadcast_helper(args[1], exclude, dims_map, common_coords) def reindex_like( self, other: Union["Dataset", "DataArray"], method: str = None, tolerance: Number = None, copy: bool = True, fill_value: Any = dtypes.NA, ) -> "Dataset": """Conform this object onto the indexes of another object, filling in missing values with ``fill_value``. The default fill value is NaN. Parameters ---------- other : Dataset or DataArray Object with an 'indexes' attribute giving a mapping from dimension names to pandas.Index objects, which provides coordinates upon which to index the variables in this dataset. The indexes on this other object need not be the same as the indexes on this dataset. Any mis-matched index values will be filled in with NaN, and any mis-matched dimension names will simply be ignored. method : {None, "nearest", "pad", "ffill", "backfill", "bfill"}, optional Method to use for filling index values from other not found in this dataset: * None (default): don't fill gaps * pad / ffill: propagate last valid index value forward * backfill / bfill: propagate next valid index value backward * nearest: use nearest valid index value tolerance : optional Maximum distance between original and new labels for inexact matches. The values of the index at the matching locations must satisfy the equation ``abs(index[indexer] - target) <= tolerance``. copy : bool, optional If ``copy=True``, data in the return value 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, a new xarray object is always returned. fill_value : scalar or dict-like, optional Value to use for newly missing values. If a dict-like maps variable names to fill values. Returns ------- reindexed : Dataset Another dataset, with this dataset's data but coordinates from the other object. See Also -------- Dataset.reindex align """ indexers = alignment.reindex_like_indexers(self, other) return self.reindex( indexers=indexers, method=method, copy=copy, fill_value=fill_value, tolerance=tolerance, ) def reindex( self, indexers: Mapping[Hashable, Any] = None, method: str = None, tolerance: Number = None, copy: bool = True, fill_value: Any = dtypes.NA, **indexers_kwargs: Any, ) -> "Dataset": """Conform this object onto a new set of indexes, filling in missing values with ``fill_value``. The default fill value is NaN. Parameters ---------- indexers : dict, optional Dictionary with keys given by dimension names and values given by arrays of coordinates tick labels. Any mis-matched coordinate values will be filled in with NaN, and any mis-matched dimension names will simply be ignored. One of indexers or indexers_kwargs must be provided. method : {None, "nearest", "pad", "ffill", "backfill", "bfill"}, optional Method to use for filling index values in ``indexers`` not found in this dataset: * None (default): don't fill gaps * pad / ffill: propagate last valid index value forward * backfill / bfill: propagate next valid index value backward * nearest: use nearest valid index value tolerance : optional Maximum distance between original and new labels for inexact matches. The values of the index at the matching locations must satisfy the equation ``abs(index[indexer] - target) <= tolerance``. copy : bool, optional If ``copy=True``, data in the return value 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, a new xarray object is always returned. fill_value : scalar or dict-like, optional Value to use for newly missing values. If a dict-like, maps variable names (including coordinates) to fill values. sparse : bool, default: False use sparse-array. **indexers_kwargs : {dim: indexer, ...}, optional Keyword arguments in the same form as ``indexers``. One of indexers or indexers_kwargs must be provided. Returns ------- reindexed : Dataset Another dataset, with this dataset's data but replaced coordinates. See Also -------- Dataset.reindex_like align pandas.Index.get_indexer Examples -------- Create a dataset with some fictional data. >>> import xarray as xr >>> import pandas as pd >>> x = xr.Dataset( ... { ... "temperature": ("station", 20 * np.random.rand(4)), ... "pressure": ("station", 500 * np.random.rand(4)), ... }, ... coords={"station": ["boston", "nyc", "seattle", "denver"]}, ... ) >>> x Dimensions: (station: 4) Coordinates: * station (station) >> x.indexes station: Index(['boston', 'nyc', 'seattle', 'denver'], dtype='object', name='station') Create a new index and reindex the dataset. By default values in the new index that do not have corresponding records in the dataset are assigned `NaN`. >>> new_index = ["boston", "austin", "seattle", "lincoln"] >>> x.reindex({"station": new_index}) Dimensions: (station: 4) Coordinates: * station (station) object 'boston' 'austin' 'seattle' 'lincoln' Data variables: temperature (station) float64 10.98 nan 12.06 nan pressure (station) float64 211.8 nan 218.8 nan We can fill in the missing values by passing a value to the keyword `fill_value`. >>> x.reindex({"station": new_index}, fill_value=0) Dimensions: (station: 4) Coordinates: * station (station) object 'boston' 'austin' 'seattle' 'lincoln' Data variables: temperature (station) float64 10.98 0.0 12.06 0.0 pressure (station) float64 211.8 0.0 218.8 0.0 We can also use different fill values for each variable. >>> x.reindex( ... {"station": new_index}, fill_value={"temperature": 0, "pressure": 100} ... ) Dimensions: (station: 4) Coordinates: * station (station) object 'boston' 'austin' 'seattle' 'lincoln' Data variables: temperature (station) float64 10.98 0.0 12.06 0.0 pressure (station) float64 211.8 100.0 218.8 100.0 Because the index is not monotonically increasing or decreasing, we cannot use arguments to the keyword method to fill the `NaN` values. >>> x.reindex({"station": new_index}, method="nearest") Traceback (most recent call last): ... raise ValueError('index must be monotonic increasing or decreasing') ValueError: index must be monotonic increasing or decreasing To further illustrate the filling functionality in reindex, we will create a dataset with a monotonically increasing index (for example, a sequence of dates). >>> x2 = xr.Dataset( ... { ... "temperature": ( ... "time", ... [15.57, 12.77, np.nan, 0.3081, 16.59, 15.12], ... ), ... "pressure": ("time", 500 * np.random.rand(6)), ... }, ... coords={"time": pd.date_range("01/01/2019", periods=6, freq="D")}, ... ) >>> x2 Dimensions: (time: 6) Coordinates: * time (time) datetime64[ns] 2019-01-01 2019-01-02 ... 2019-01-06 Data variables: temperature (time) float64 15.57 12.77 nan 0.3081 16.59 15.12 pressure (time) float64 481.8 191.7 395.9 264.4 284.0 462.8 Suppose we decide to expand the dataset to cover a wider date range. >>> time_index2 = pd.date_range("12/29/2018", periods=10, freq="D") >>> x2.reindex({"time": time_index2}) Dimensions: (time: 10) Coordinates: * time (time) datetime64[ns] 2018-12-29 2018-12-30 ... 2019-01-07 Data variables: temperature (time) float64 nan nan nan 15.57 ... 0.3081 16.59 15.12 nan pressure (time) float64 nan nan nan 481.8 ... 264.4 284.0 462.8 nan The index entries that did not have a value in the original data frame (for example, `2018-12-29`) are by default filled with NaN. If desired, we can fill in the missing values using one of several options. For example, to back-propagate the last valid value to fill the `NaN` values, pass `bfill` as an argument to the `method` keyword. >>> x3 = x2.reindex({"time": time_index2}, method="bfill") >>> x3 Dimensions: (time: 10) Coordinates: * time (time) datetime64[ns] 2018-12-29 2018-12-30 ... 2019-01-07 Data variables: temperature (time) float64 15.57 15.57 15.57 15.57 ... 16.59 15.12 nan pressure (time) float64 481.8 481.8 481.8 481.8 ... 284.0 462.8 nan Please note that the `NaN` value present in the original dataset (at index value `2019-01-03`) will not be filled by any of the value propagation schemes. >>> x2.where(x2.temperature.isnull(), drop=True) Dimensions: (time: 1) Coordinates: * time (time) datetime64[ns] 2019-01-03 Data variables: temperature (time) float64 nan pressure (time) float64 395.9 >>> x3.where(x3.temperature.isnull(), drop=True) Dimensions: (time: 2) Coordinates: * time (time) datetime64[ns] 2019-01-03 2019-01-07 Data variables: temperature (time) float64 nan nan pressure (time) float64 395.9 nan This is because filling while reindexing does not look at dataset values, but only compares the original and desired indexes. If you do want to fill in the `NaN` values present in the original dataset, use the :py:meth:`~Dataset.fillna()` method. """ return self._reindex( indexers, method, tolerance, copy, fill_value, sparse=False, **indexers_kwargs, ) def _reindex( self, indexers: Mapping[Hashable, Any] = None, method: str = None, tolerance: Number = None, copy: bool = True, fill_value: Any = dtypes.NA, sparse: bool = False, **indexers_kwargs: Any, ) -> "Dataset": """ same to _reindex but support sparse option """ indexers = utils.either_dict_or_kwargs(indexers, indexers_kwargs, "reindex") bad_dims = [d for d in indexers if d not in self.dims] if bad_dims: raise ValueError("invalid reindex dimensions: %s" % bad_dims) variables, indexes = alignment.reindex_variables( self.variables, self.sizes, self.indexes, indexers, method, tolerance, copy=copy, fill_value=fill_value, sparse=sparse, ) coord_names = set(self._coord_names) coord_names.update(indexers) return self._replace_with_new_dims(variables, coord_names, indexes=indexes) def interp( self, coords: Mapping[Hashable, Any] = None, method: str = "linear", assume_sorted: bool = False, kwargs: Mapping[str, Any] = None, **coords_kwargs: Any, ) -> "Dataset": """Multidimensional interpolation of Dataset. Parameters ---------- coords : dict, optional Mapping from dimension names to the new coordinates. New coordinate can be a scalar, array-like or DataArray. If DataArrays are passed as new coordinates, their dimensions are used for the broadcasting. Missing values are skipped. method : str, optional {"linear", "nearest"} for multidimensional array, {"linear", "nearest", "zero", "slinear", "quadratic", "cubic"} for 1-dimensional array. "linear" is used by default. assume_sorted : bool, optional If False, values of coordinates that are interpolated over can be in any order and they are sorted first. If True, interpolated coordinates are assumed to be an array of monotonically increasing values. kwargs: dict, optional Additional keyword arguments passed to scipy's interpolator. Valid options and their behavior depend on if 1-dimensional or multi-dimensional interpolation is used. **coords_kwargs : {dim: coordinate, ...}, optional The keyword arguments form of ``coords``. One of coords or coords_kwargs must be provided. Returns ------- interpolated : Dataset New dataset on the new coordinates. Notes ----- scipy is required. See Also -------- scipy.interpolate.interp1d scipy.interpolate.interpn Examples -------- >>> ds = xr.Dataset( ... data_vars={ ... "a": ("x", [5, 7, 4]), ... "b": ( ... ("x", "y"), ... [[1, 4, 2, 9], [2, 7, 6, np.nan], [6, np.nan, 5, 8]], ... ), ... }, ... coords={"x": [0, 1, 2], "y": [10, 12, 14, 16]}, ... ) >>> ds Dimensions: (x: 3, y: 4) Coordinates: * x (x) int64 0 1 2 * y (y) int64 10 12 14 16 Data variables: a (x) int64 5 7 4 b (x, y) float64 1.0 4.0 2.0 9.0 2.0 7.0 6.0 nan 6.0 nan 5.0 8.0 1D interpolation with the default method (linear): >>> ds.interp(x=[0, 0.75, 1.25, 1.75]) Dimensions: (x: 4, y: 4) Coordinates: * y (y) int64 10 12 14 16 * x (x) float64 0.0 0.75 1.25 1.75 Data variables: a (x) float64 5.0 6.5 6.25 4.75 b (x, y) float64 1.0 4.0 2.0 nan 1.75 6.25 ... nan 5.0 nan 5.25 nan 1D interpolation with a different method: >>> ds.interp(x=[0, 0.75, 1.25, 1.75], method="nearest") Dimensions: (x: 4, y: 4) Coordinates: * y (y) int64 10 12 14 16 * x (x) float64 0.0 0.75 1.25 1.75 Data variables: a (x) float64 5.0 7.0 7.0 4.0 b (x, y) float64 1.0 4.0 2.0 9.0 2.0 7.0 ... 6.0 nan 6.0 nan 5.0 8.0 1D extrapolation: >>> ds.interp( ... x=[1, 1.5, 2.5, 3.5], ... method="linear", ... kwargs={"fill_value": "extrapolate"}, ... ) Dimensions: (x: 4, y: 4) Coordinates: * y (y) int64 10 12 14 16 * x (x) float64 1.0 1.5 2.5 3.5 Data variables: a (x) float64 7.0 5.5 2.5 -0.5 b (x, y) float64 2.0 7.0 6.0 nan 4.0 nan ... 4.5 nan 12.0 nan 3.5 nan 2D interpolation: >>> ds.interp(x=[0, 0.75, 1.25, 1.75], y=[11, 13, 15], method="linear") Dimensions: (x: 4, y: 3) Coordinates: * x (x) float64 0.0 0.75 1.25 1.75 * y (y) int64 11 13 15 Data variables: a (x) float64 5.0 6.5 6.25 4.75 b (x, y) float64 2.5 3.0 nan 4.0 5.625 nan nan nan nan nan nan nan """ from . import missing if kwargs is None: kwargs = {} coords = either_dict_or_kwargs(coords, coords_kwargs, "interp") indexers = dict(self._validate_interp_indexers(coords)) if coords: # This avoids broadcasting over coordinates that are both in # the original array AND in the indexing array. It essentially # forces interpolation along the shared coordinates. sdims = ( set(self.dims) .intersection(*[set(nx.dims) for nx in indexers.values()]) .difference(coords.keys()) ) indexers.update({d: self.variables[d] for d in sdims}) obj = self if assume_sorted else self.sortby([k for k in coords]) def maybe_variable(obj, k): # workaround to get variable for dimension without coordinate. try: return obj._variables[k] except KeyError: return as_variable((k, range(obj.dims[k]))) def _validate_interp_indexer(x, new_x): # In the case of datetimes, the restrictions placed on indexers # used with interp are stronger than those which are placed on # isel, so we need an additional check after _validate_indexers. if _contains_datetime_like_objects( x ) and not _contains_datetime_like_objects(new_x): raise TypeError( "When interpolating over a datetime-like " "coordinate, the coordinates to " "interpolate to must be either datetime " "strings or datetimes. " "Instead got\n{}".format(new_x) ) return x, new_x variables: Dict[Hashable, Variable] = {} for name, var in obj._variables.items(): if name in indexers: continue if var.dtype.kind in "uifc": var_indexers = { k: _validate_interp_indexer(maybe_variable(obj, k), v) for k, v in indexers.items() if k in var.dims } variables[name] = missing.interp(var, var_indexers, method, **kwargs) elif all(d not in indexers for d in var.dims): # keep unrelated object array variables[name] = var coord_names = obj._coord_names & variables.keys() indexes = {k: v for k, v in obj.indexes.items() if k not in indexers} selected = self._replace_with_new_dims( variables.copy(), coord_names, indexes=indexes ) # attach indexer as coordinate variables.update(indexers) for k, v in indexers.items(): assert isinstance(v, Variable) if v.dims == (k,): indexes[k] = v.to_index() # Extract coordinates from indexers coord_vars, new_indexes = selected._get_indexers_coords_and_indexes(coords) variables.update(coord_vars) indexes.update(new_indexes) coord_names = obj._coord_names & variables.keys() | coord_vars.keys() return self._replace_with_new_dims(variables, coord_names, indexes=indexes) def interp_like( self, other: Union["Dataset", "DataArray"], method: str = "linear", assume_sorted: bool = False, kwargs: Mapping[str, Any] = None, ) -> "Dataset": """Interpolate this object onto the coordinates of another object, filling the out of range values with NaN. Parameters ---------- other : Dataset or DataArray Object with an 'indexes' attribute giving a mapping from dimension names to an 1d array-like, which provides coordinates upon which to index the variables in this dataset. Missing values are skipped. method : str, optional {"linear", "nearest"} for multidimensional array, {"linear", "nearest", "zero", "slinear", "quadratic", "cubic"} for 1-dimensional array. 'linear' is used by default. assume_sorted : bool, optional If False, values of coordinates that are interpolated over can be in any order and they are sorted first. If True, interpolated coordinates are assumed to be an array of monotonically increasing values. kwargs: dict, optional Additional keyword passed to scipy's interpolator. Returns ------- interpolated : Dataset Another dataset by interpolating this dataset's data along the coordinates of the other object. Notes ----- scipy is required. If the dataset has object-type coordinates, reindex is used for these coordinates instead of the interpolation. See Also -------- Dataset.interp Dataset.reindex_like """ if kwargs is None: kwargs = {} coords = alignment.reindex_like_indexers(self, other) numeric_coords: Dict[Hashable, pd.Index] = {} object_coords: Dict[Hashable, pd.Index] = {} for k, v in coords.items(): if v.dtype.kind in "uifcMm": numeric_coords[k] = v else: object_coords[k] = v ds = self if object_coords: # We do not support interpolation along object coordinate. # reindex instead. ds = self.reindex(object_coords) return ds.interp(numeric_coords, method, assume_sorted, kwargs) # Helper methods for rename() def _rename_vars(self, name_dict, dims_dict): variables = {} coord_names = set() for k, v in self.variables.items(): var = v.copy(deep=False) var.dims = tuple(dims_dict.get(dim, dim) for dim in v.dims) name = name_dict.get(k, k) if name in variables: raise ValueError(f"the new name {name!r} conflicts") variables[name] = var if k in self._coord_names: coord_names.add(name) return variables, coord_names def _rename_dims(self, name_dict): return {name_dict.get(k, k): v for k, v in self.dims.items()} def _rename_indexes(self, name_dict, dims_set): if self._indexes is None: return None indexes = {} for k, v in self.indexes.items(): new_name = name_dict.get(k, k) if new_name not in dims_set: continue if isinstance(v, pd.MultiIndex): new_names = [name_dict.get(k, k) for k in v.names] index = v.rename(names=new_names) else: index = v.rename(new_name) indexes[new_name] = index return indexes def _rename_all(self, name_dict, dims_dict): variables, coord_names = self._rename_vars(name_dict, dims_dict) dims = self._rename_dims(dims_dict) indexes = self._rename_indexes(name_dict, dims.keys()) return variables, coord_names, dims, indexes def rename( self, name_dict: Mapping[Hashable, Hashable] = None, **names: Hashable, ) -> "Dataset": """Returns a new object with renamed variables and dimensions. Parameters ---------- name_dict : dict-like, optional Dictionary whose keys are current variable or dimension names and whose values are the desired names. **names : optional Keyword form of ``name_dict``. One of name_dict or names must be provided. Returns ------- renamed : Dataset Dataset with renamed variables and dimensions. See Also -------- Dataset.swap_dims Dataset.rename_vars Dataset.rename_dims DataArray.rename """ name_dict = either_dict_or_kwargs(name_dict, names, "rename") for k in name_dict.keys(): if k not in self and k not in self.dims: raise ValueError( "cannot rename %r because it is not a " "variable or dimension in this dataset" % k ) variables, coord_names, dims, indexes = self._rename_all( name_dict=name_dict, dims_dict=name_dict ) assert_unique_multiindex_level_names(variables) return self._replace(variables, coord_names, dims=dims, indexes=indexes) def rename_dims( self, dims_dict: Mapping[Hashable, Hashable] = None, **dims: Hashable ) -> "Dataset": """Returns a new object with renamed dimensions only. Parameters ---------- dims_dict : dict-like, optional Dictionary whose keys are current dimension names and whose values are the desired names. The desired names must not be the name of an existing dimension or Variable in the Dataset. **dims : optional Keyword form of ``dims_dict``. One of dims_dict or dims must be provided. Returns ------- renamed : Dataset Dataset with renamed dimensions. See Also -------- Dataset.swap_dims Dataset.rename Dataset.rename_vars DataArray.rename """ dims_dict = either_dict_or_kwargs(dims_dict, dims, "rename_dims") for k, v in dims_dict.items(): if k not in self.dims: raise ValueError( "cannot rename %r because it is not a " "dimension in this dataset" % k ) if v in self.dims or v in self: raise ValueError( f"Cannot rename {k} to {v} because {v} already exists. " "Try using swap_dims instead." ) variables, coord_names, sizes, indexes = self._rename_all( name_dict={}, dims_dict=dims_dict ) return self._replace(variables, coord_names, dims=sizes, indexes=indexes) def rename_vars( self, name_dict: Mapping[Hashable, Hashable] = None, **names: Hashable ) -> "Dataset": """Returns a new object with renamed variables including coordinates Parameters ---------- name_dict : dict-like, optional Dictionary whose keys are current variable or coordinate names and whose values are the desired names. **names : optional Keyword form of ``name_dict``. One of name_dict or names must be provided. Returns ------- renamed : Dataset Dataset with renamed variables including coordinates See Also -------- Dataset.swap_dims Dataset.rename Dataset.rename_dims DataArray.rename """ name_dict = either_dict_or_kwargs(name_dict, names, "rename_vars") for k in name_dict: if k not in self: raise ValueError( "cannot rename %r because it is not a " "variable or coordinate in this dataset" % k ) variables, coord_names, dims, indexes = self._rename_all( name_dict=name_dict, dims_dict={} ) return self._replace(variables, coord_names, dims=dims, indexes=indexes) def swap_dims( self, dims_dict: Mapping[Hashable, Hashable], inplace: bool = None ) -> "Dataset": """Returns a new object with swapped dimensions. Parameters ---------- dims_dict : dict-like Dictionary whose keys are current dimension names and whose values are new names. Returns ------- swapped : Dataset Dataset with swapped dimensions. Examples -------- >>> ds = xr.Dataset( ... data_vars={"a": ("x", [5, 7]), "b": ("x", [0.1, 2.4])}, ... coords={"x": ["a", "b"], "y": ("x", [0, 1])}, ... ) >>> ds Dimensions: (x: 2) Coordinates: * x (x) >> ds.swap_dims({"x": "y"}) Dimensions: (y: 2) Coordinates: x (y) >> ds.swap_dims({"x": "z"}) Dimensions: (z: 2) Coordinates: x (z) "Dataset": """Return a new object with an additional axis (or axes) inserted at the corresponding position in the array shape. The new object is a view into the underlying array, not a copy. If dim is already a scalar coordinate, it will be promoted to a 1D coordinate consisting of a single value. Parameters ---------- dim : hashable, sequence of hashable, mapping, or None Dimensions to include on the new variable. If provided as hashable or sequence of hashable, then dimensions are inserted with length 1. If provided as a mapping, then the keys are the new dimensions and the values are either integers (giving the length of the new dimensions) or array-like (giving the coordinates of the new dimensions). axis : int, sequence of int, or None Axis position(s) where new axis is to be inserted (position(s) on the result array). If a list (or tuple) of integers is passed, multiple axes are inserted. In this case, dim arguments should be same length list. If axis=None is passed, all the axes will be inserted to the start of the result array. **dim_kwargs : int or sequence or ndarray The keywords are arbitrary dimensions being inserted and the values are either the lengths of the new dims (if int is given), or their coordinates. Note, this is an alternative to passing a dict to the dim kwarg and will only be used if dim is None. Returns ------- expanded : same type as caller This object, but with an additional dimension(s). """ if dim is None: pass elif isinstance(dim, Mapping): # We're later going to modify dim in place; don't tamper with # the input dim = dict(dim) elif isinstance(dim, int): raise TypeError( "dim should be hashable or sequence of hashables or mapping" ) elif isinstance(dim, str) or not isinstance(dim, Sequence): dim = {dim: 1} elif isinstance(dim, Sequence): if len(dim) != len(set(dim)): raise ValueError("dims should not contain duplicate values.") dim = {d: 1 for d in dim} dim = either_dict_or_kwargs(dim, dim_kwargs, "expand_dims") assert isinstance(dim, MutableMapping) if axis is None: axis = list(range(len(dim))) elif not isinstance(axis, Sequence): axis = [axis] if len(dim) != len(axis): raise ValueError("lengths of dim and axis should be identical.") for d in dim: if d in self.dims: raise ValueError(f"Dimension {d} already exists.") if d in self._variables and not utils.is_scalar(self._variables[d]): raise ValueError( "{dim} already exists as coordinate or" " variable name.".format(dim=d) ) variables: Dict[Hashable, Variable] = {} coord_names = self._coord_names.copy() # If dim is a dict, then ensure that the values are either integers # or iterables. for k, v in dim.items(): if hasattr(v, "__iter__"): # If the value for the new dimension is an iterable, then # save the coordinates to the variables dict, and set the # value within the dim dict to the length of the iterable # for later use. variables[k] = xr.IndexVariable((k,), v) coord_names.add(k) dim[k] = variables[k].size elif isinstance(v, int): pass # Do nothing if the dimensions value is just an int else: raise TypeError( "The value of new dimension {k} must be " "an iterable or an int".format(k=k) ) for k, v in self._variables.items(): if k not in dim: if k in coord_names: # Do not change coordinates variables[k] = v else: result_ndim = len(v.dims) + len(axis) for a in axis: if a < -result_ndim or result_ndim - 1 < a: raise IndexError( f"Axis {a} of variable {k} is out of bounds of the " f"expanded dimension size {result_ndim}" ) axis_pos = [a if a >= 0 else result_ndim + a for a in axis] if len(axis_pos) != len(set(axis_pos)): raise ValueError("axis should not contain duplicate values") # We need to sort them to make sure `axis` equals to the # axis positions of the result array. zip_axis_dim = sorted(zip(axis_pos, dim.items())) all_dims = list(zip(v.dims, v.shape)) for d, c in zip_axis_dim: all_dims.insert(d, c) variables[k] = v.set_dims(dict(all_dims)) else: # If dims includes a label of a non-dimension coordinate, # it will be promoted to a 1D coordinate with a single value. variables[k] = v.set_dims(k).to_index_variable() new_dims = self._dims.copy() new_dims.update(dim) return self._replace_vars_and_dims( variables, dims=new_dims, coord_names=coord_names ) def set_index( self, indexes: Mapping[Hashable, Union[Hashable, Sequence[Hashable]]] = None, append: bool = False, inplace: bool = None, **indexes_kwargs: Union[Hashable, Sequence[Hashable]], ) -> "Dataset": """Set Dataset (multi-)indexes using one or more existing coordinates or variables. Parameters ---------- indexes : {dim: index, ...} Mapping from names matching dimensions and values given by (lists of) the names of existing coordinates or variables to set as new (multi-)index. append : bool, optional If True, append the supplied index(es) to the existing index(es). Otherwise replace the existing index(es) (default). **indexes_kwargs : optional The keyword arguments form of ``indexes``. One of indexes or indexes_kwargs must be provided. Returns ------- obj : Dataset Another dataset, with this dataset's data but replaced coordinates. Examples -------- >>> arr = xr.DataArray( ... data=np.ones((2, 3)), ... dims=["x", "y"], ... coords={"x": range(2), "y": range(3), "a": ("x", [3, 4])}, ... ) >>> ds = xr.Dataset({"v": arr}) >>> ds Dimensions: (x: 2, y: 3) Coordinates: * x (x) int64 0 1 * y (y) int64 0 1 2 a (x) int64 3 4 Data variables: v (x, y) float64 1.0 1.0 1.0 1.0 1.0 1.0 >>> ds.set_index(x="a") Dimensions: (x: 2, y: 3) Coordinates: * x (x) int64 3 4 * y (y) int64 0 1 2 Data variables: v (x, y) float64 1.0 1.0 1.0 1.0 1.0 1.0 See Also -------- Dataset.reset_index Dataset.swap_dims """ _check_inplace(inplace) indexes = either_dict_or_kwargs(indexes, indexes_kwargs, "set_index") variables, coord_names = merge_indexes( indexes, self._variables, self._coord_names, append=append ) return self._replace_vars_and_dims(variables, coord_names=coord_names) def reset_index( self, dims_or_levels: Union[Hashable, Sequence[Hashable]], drop: bool = False, inplace: bool = None, ) -> "Dataset": """Reset the specified index(es) or multi-index level(s). Parameters ---------- dims_or_levels : str or list Name(s) of the dimension(s) and/or multi-index level(s) that will be reset. drop : bool, optional If True, remove the specified indexes and/or multi-index levels instead of extracting them as new coordinates (default: False). Returns ------- obj : Dataset Another dataset, with this dataset's data but replaced coordinates. See Also -------- Dataset.set_index """ _check_inplace(inplace) variables, coord_names = split_indexes( dims_or_levels, self._variables, self._coord_names, cast(Mapping[Hashable, Hashable], self._level_coords), drop=drop, ) return self._replace_vars_and_dims(variables, coord_names=coord_names) def reorder_levels( self, dim_order: Mapping[Hashable, Sequence[int]] = None, inplace: bool = None, **dim_order_kwargs: Sequence[int], ) -> "Dataset": """Rearrange index levels using input order. Parameters ---------- dim_order : optional Mapping from names matching dimensions and values given by lists representing new level orders. Every given dimension must have a multi-index. **dim_order_kwargs : optional The keyword arguments form of ``dim_order``. One of dim_order or dim_order_kwargs must be provided. Returns ------- obj : Dataset Another dataset, with this dataset's data but replaced coordinates. """ _check_inplace(inplace) dim_order = either_dict_or_kwargs(dim_order, dim_order_kwargs, "reorder_levels") variables = self._variables.copy() indexes = dict(self.indexes) for dim, order in dim_order.items(): coord = self._variables[dim] index = self.indexes[dim] if not isinstance(index, pd.MultiIndex): raise ValueError(f"coordinate {dim} has no MultiIndex") new_index = index.reorder_levels(order) variables[dim] = IndexVariable(coord.dims, new_index) indexes[dim] = new_index return self._replace(variables, indexes=indexes) def _stack_once(self, dims, new_dim): if ... in dims: dims = list(infix_dims(dims, self.dims)) variables = {} for name, var in self.variables.items(): if name not in dims: if any(d in var.dims for d in dims): add_dims = [d for d in dims if d not in var.dims] vdims = list(var.dims) + add_dims shape = [self.dims[d] for d in vdims] exp_var = var.set_dims(vdims, shape) stacked_var = exp_var.stack(**{new_dim: dims}) variables[name] = stacked_var else: variables[name] = var.copy(deep=False) # consider dropping levels that are unused? levels = [self.get_index(dim) for dim in dims] idx = utils.multiindex_from_product_levels(levels, names=dims) variables[new_dim] = IndexVariable(new_dim, idx) coord_names = set(self._coord_names) - set(dims) | {new_dim} indexes = {k: v for k, v in self.indexes.items() if k not in dims} indexes[new_dim] = idx return self._replace_with_new_dims( variables, coord_names=coord_names, indexes=indexes ) def stack( self, dimensions: Mapping[Hashable, Sequence[Hashable]] = None, **dimensions_kwargs: Sequence[Hashable], ) -> "Dataset": """ Stack any number of existing dimensions into a single new dimension. New dimensions will be added at the end, and the corresponding coordinate variables will be combined into a MultiIndex. Parameters ---------- dimensions : mapping of hashable to sequence of hashable Mapping of the form `new_name=(dim1, dim2, ...)`. Names of new dimensions, and the existing dimensions that they replace. An ellipsis (`...`) will be replaced by all unlisted dimensions. Passing a list containing an ellipsis (`stacked_dim=[...]`) will stack over all dimensions. **dimensions_kwargs The keyword arguments form of ``dimensions``. One of dimensions or dimensions_kwargs must be provided. Returns ------- stacked : Dataset Dataset with stacked data. See also -------- Dataset.unstack """ dimensions = either_dict_or_kwargs(dimensions, dimensions_kwargs, "stack") result = self for new_dim, dims in dimensions.items(): result = result._stack_once(dims, new_dim) return result def to_stacked_array( self, new_dim: Hashable, sample_dims: Sequence[Hashable], variable_dim: str = "variable", name: Hashable = None, ) -> "DataArray": """Combine variables of differing dimensionality into a DataArray without broadcasting. This method is similar to Dataset.to_array but does not broadcast the variables. Parameters ---------- new_dim : hashable Name of the new stacked coordinate sample_dims : sequence of hashable Dimensions that **will not** be stacked. Each array in the dataset must share these dimensions. For machine learning applications, these define the dimensions over which samples are drawn. variable_dim : str, optional Name of the level in the stacked coordinate which corresponds to the variables. name : str, optional Name of the new data array. Returns ------- stacked : DataArray DataArray with the specified dimensions and data variables stacked together. The stacked coordinate is named ``new_dim`` and represented by a MultiIndex object with a level containing the data variable names. The name of this level is controlled using the ``variable_dim`` argument. See Also -------- Dataset.to_array Dataset.stack DataArray.to_unstacked_dataset Examples -------- >>> data = xr.Dataset( ... data_vars={ ... "a": (("x", "y"), [[0, 1, 2], [3, 4, 5]]), ... "b": ("x", [6, 7]), ... }, ... coords={"y": ["u", "v", "w"]}, ... ) >>> data Dimensions: (x: 2, y: 3) Coordinates: * y (y) >> data.to_stacked_array("z", sample_dims=["x"]) array([[0, 1, 2, 6], [3, 4, 5, 7]]) Coordinates: * z (z) MultiIndex - variable (z) object 'a' 'a' 'a' 'b' - y (z) object 'u' 'v' 'w' nan Dimensions without coordinates: x """ stacking_dims = tuple(dim for dim in self.dims if dim not in sample_dims) for variable in self: dims = self[variable].dims dims_include_sample_dims = set(sample_dims) <= set(dims) if not dims_include_sample_dims: raise ValueError( "All variables in the dataset must contain the " "dimensions {}.".format(dims) ) def ensure_stackable(val): assign_coords = {variable_dim: val.name} for dim in stacking_dims: if dim not in val.dims: assign_coords[dim] = None expand_dims = set(stacking_dims).difference(set(val.dims)) expand_dims.add(variable_dim) # must be list for .expand_dims expand_dims = list(expand_dims) return ( val.assign_coords(**assign_coords) .expand_dims(expand_dims) .stack({new_dim: (variable_dim,) + stacking_dims}) ) # concatenate the arrays stackable_vars = [ensure_stackable(self[key]) for key in self.data_vars] data_array = xr.concat(stackable_vars, dim=new_dim) # coerce the levels of the MultiIndex to have the same type as the # input dimensions. This code is messy, so it might be better to just # input a dummy value for the singleton dimension. idx = data_array.indexes[new_dim] levels = [idx.levels[0]] + [ level.astype(self[level.name].dtype) for level in idx.levels[1:] ] new_idx = idx.set_levels(levels) data_array[new_dim] = IndexVariable(new_dim, new_idx) if name is not None: data_array.name = name return data_array def _unstack_once(self, dim: Hashable, fill_value, sparse) -> "Dataset": index = self.get_index(dim) index = remove_unused_levels_categories(index) full_idx = pd.MultiIndex.from_product(index.levels, names=index.names) # take a shortcut in case the MultiIndex was not modified. if index.equals(full_idx): obj = self else: obj = self._reindex( {dim: full_idx}, copy=False, fill_value=fill_value, sparse=sparse ) new_dim_names = index.names new_dim_sizes = [lev.size for lev in index.levels] variables: Dict[Hashable, Variable] = {} indexes = {k: v for k, v in self.indexes.items() if k != dim} for name, var in obj.variables.items(): if name != dim: if dim in var.dims: new_dims = dict(zip(new_dim_names, new_dim_sizes)) variables[name] = var.unstack({dim: new_dims}) else: variables[name] = var for name, lev in zip(new_dim_names, index.levels): variables[name] = IndexVariable(name, lev) indexes[name] = lev coord_names = set(self._coord_names) - {dim} | set(new_dim_names) return self._replace_with_new_dims( variables, coord_names=coord_names, indexes=indexes ) def unstack( self, dim: Union[Hashable, Iterable[Hashable]] = None, fill_value: Any = dtypes.NA, sparse: bool = False, ) -> "Dataset": """ Unstack existing dimensions corresponding to MultiIndexes into multiple new dimensions. New dimensions will be added at the end. Parameters ---------- dim : hashable or iterable of hashable, optional Dimension(s) over which to unstack. By default unstacks all MultiIndexes. fill_value : scalar or dict-like, default: nan value to be filled. If a dict-like, maps variable names to fill values. If not provided or if the dict-like does not contain all variables, the dtype's NA value will be used. sparse : bool, default: False use sparse-array if True Returns ------- unstacked : Dataset Dataset with unstacked data. See also -------- Dataset.stack """ if dim is None: dims = [ d for d in self.dims if isinstance(self.get_index(d), pd.MultiIndex) ] else: if isinstance(dim, str) or not isinstance(dim, Iterable): dims = [dim] else: dims = list(dim) missing_dims = [d for d in dims if d not in self.dims] if missing_dims: raise ValueError( "Dataset does not contain the dimensions: %s" % missing_dims ) non_multi_dims = [ d for d in dims if not isinstance(self.get_index(d), pd.MultiIndex) ] if non_multi_dims: raise ValueError( "cannot unstack dimensions that do not " "have a MultiIndex: %s" % non_multi_dims ) result = self.copy(deep=False) for dim in dims: result = result._unstack_once(dim, fill_value, sparse) return result def update(self, other: "CoercibleMapping", inplace: bool = None) -> "Dataset": """Update this dataset's variables with those from another dataset. Parameters ---------- other : Dataset or mapping Variables with which to update this dataset. One of: - Dataset - mapping {var name: DataArray} - mapping {var name: Variable} - mapping {var name: (dimension name, array-like)} - mapping {var name: (tuple of dimension names, array-like)} Returns ------- updated : Dataset Updated dataset. Raises ------ ValueError If any dimensions would have inconsistent sizes in the updated dataset. """ _check_inplace(inplace) merge_result = dataset_update_method(self, other) return self._replace(inplace=True, **merge_result._asdict()) def merge( self, other: Union["CoercibleMapping", "DataArray"], inplace: bool = None, overwrite_vars: Union[Hashable, Iterable[Hashable]] = frozenset(), compat: str = "no_conflicts", join: str = "outer", fill_value: Any = dtypes.NA, ) -> "Dataset": """Merge the arrays of two datasets into a single dataset. This method generally does not allow for overriding data, with the exception of attributes, which are ignored on the second dataset. Variables with the same name are checked for conflicts via the equals or identical methods. Parameters ---------- other : Dataset or mapping Dataset or variables to merge with this dataset. overwrite_vars : hashable or iterable of hashable, optional If provided, update variables of these name(s) without checking for conflicts in this dataset. compat : {"broadcast_equals", "equals", "identical", \ "no_conflicts"}, optional String indicating how to compare variables of the same name for potential conflicts: - 'broadcast_equals': all values must be equal when variables are broadcast against each other to ensure common dimensions. - 'equals': all values and dimensions must be the same. - 'identical': all values, dimensions and attributes must be the same. - 'no_conflicts': only values which are not null in both datasets must be equal. The returned dataset then contains the combination of all non-null values. join : {"outer", "inner", "left", "right", "exact"}, optional Method for joining ``self`` and ``other`` along shared dimensions: - 'outer': use the union of the indexes - 'inner': use the intersection of the indexes - 'left': use indexes from ``self`` - 'right': use indexes from ``other`` - 'exact': error instead of aligning non-equal indexes fill_value : scalar or dict-like, optional Value to use for newly missing values. If a dict-like, maps variable names (including coordinates) to fill values. Returns ------- merged : Dataset Merged dataset. Raises ------ MergeError If any variables conflict (see ``compat``). """ _check_inplace(inplace) other = other.to_dataset() if isinstance(other, xr.DataArray) else other merge_result = dataset_merge_method( self, other, overwrite_vars=overwrite_vars, compat=compat, join=join, fill_value=fill_value, ) return self._replace(**merge_result._asdict()) def _assert_all_in_dataset( self, names: Iterable[Hashable], virtual_okay: bool = False ) -> None: bad_names = set(names) - set(self._variables) if virtual_okay: bad_names -= self.virtual_variables if bad_names: raise ValueError( "One or more of the specified variables " "cannot be found in this dataset" ) def drop_vars( self, names: Union[Hashable, Iterable[Hashable]], *, errors: str = "raise" ) -> "Dataset": """Drop variables from this dataset. Parameters ---------- names : hashable or iterable of hashable Name(s) of variables to drop. errors : {"raise", "ignore"}, optional If 'raise' (default), raises a ValueError error if any of the variable passed are not in the dataset. If 'ignore', any given names that are in the dataset are dropped and no error is raised. Returns ------- dropped : Dataset """ # the Iterable check is required for mypy if is_scalar(names) or not isinstance(names, Iterable): names = {names} else: names = set(names) if errors == "raise": self._assert_all_in_dataset(names) variables = {k: v for k, v in self._variables.items() if k not in names} coord_names = {k for k in self._coord_names if k in variables} indexes = {k: v for k, v in self.indexes.items() if k not in names} return self._replace_with_new_dims( variables, coord_names=coord_names, indexes=indexes ) def drop(self, labels=None, dim=None, *, errors="raise", **labels_kwargs): """Backward compatible method based on `drop_vars` and `drop_sel` Using either `drop_vars` or `drop_sel` is encouraged See Also -------- Dataset.drop_vars Dataset.drop_sel """ if errors not in ["raise", "ignore"]: raise ValueError('errors must be either "raise" or "ignore"') if is_dict_like(labels) and not isinstance(labels, dict): warnings.warn( "dropping coordinates using `drop` is be deprecated; use drop_vars.", FutureWarning, stacklevel=2, ) return self.drop_vars(labels, errors=errors) if labels_kwargs or isinstance(labels, dict): if dim is not None: raise ValueError("cannot specify dim and dict-like arguments.") labels = either_dict_or_kwargs(labels, labels_kwargs, "drop") if dim is None and (is_scalar(labels) or isinstance(labels, Iterable)): warnings.warn( "dropping variables using `drop` will be deprecated; using drop_vars is encouraged.", PendingDeprecationWarning, stacklevel=2, ) return self.drop_vars(labels, errors=errors) if dim is not None: warnings.warn( "dropping labels using list-like labels is deprecated; using " "dict-like arguments with `drop_sel`, e.g. `ds.drop_sel(dim=[labels]).", DeprecationWarning, stacklevel=2, ) return self.drop_sel({dim: labels}, errors=errors, **labels_kwargs) warnings.warn( "dropping labels using `drop` will be deprecated; using drop_sel is encouraged.", PendingDeprecationWarning, stacklevel=2, ) return self.drop_sel(labels, errors=errors) def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs): """Drop index labels from this dataset. Parameters ---------- labels : mapping of hashable to Any Index labels to drop errors : {"raise", "ignore"}, optional If 'raise' (default), raises a ValueError error if any of the index labels passed are not in the dataset. If 'ignore', any given labels that are in the dataset are dropped and no error is raised. **labels_kwargs : {dim: label, ...}, optional The keyword arguments form of ``dim`` and ``labels`` Returns ------- dropped : Dataset Examples -------- >>> data = np.random.randn(2, 3) >>> labels = ["a", "b", "c"] >>> ds = xr.Dataset({"A": (["x", "y"], data), "y": labels}) >>> ds.drop_sel(y=["a", "c"]) Dimensions: (x: 2, y: 1) Coordinates: * y (y) >> ds.drop_sel(y="b") Dimensions: (x: 2, y: 2) Coordinates: * y (y) "Dataset": """Drop dimensions and associated variables from this dataset. Parameters ---------- drop_dims : hashable or iterable of hashable Dimension or dimensions to drop. errors : {"raise", "ignore"}, optional If 'raise' (default), raises a ValueError error if any of the dimensions passed are not in the dataset. If 'ignore', any given labels that are in the dataset are dropped and no error is raised. Returns ------- obj : Dataset The dataset without the given dimensions (or any variables containing those dimensions) errors : {"raise", "ignore"}, optional If 'raise' (default), raises a ValueError error if any of the dimensions passed are not in the dataset. If 'ignore', any given dimensions that are in the dataset are dropped and no error is raised. """ if errors not in ["raise", "ignore"]: raise ValueError('errors must be either "raise" or "ignore"') if isinstance(drop_dims, str) or not isinstance(drop_dims, Iterable): drop_dims = {drop_dims} else: drop_dims = set(drop_dims) if errors == "raise": missing_dims = drop_dims - set(self.dims) if missing_dims: raise ValueError( "Dataset does not contain the dimensions: %s" % missing_dims ) drop_vars = {k for k, v in self._variables.items() if set(v.dims) & drop_dims} return self.drop_vars(drop_vars) def transpose(self, *dims: Hashable) -> "Dataset": """Return a new Dataset object with all array dimensions transposed. Although the order of dimensions on each array will change, the dataset dimensions themselves will remain in fixed (sorted) order. Parameters ---------- *dims : hashable, optional By default, reverse the dimensions on each array. Otherwise, reorder the dimensions to this order. Returns ------- transposed : Dataset Each array in the dataset (including) coordinates will be transposed to the given order. Notes ----- This operation returns a view of each array's data. It is lazy for dask-backed DataArrays but not for numpy-backed DataArrays -- the data will be fully loaded into memory. See Also -------- numpy.transpose DataArray.transpose """ if dims: if set(dims) ^ set(self.dims) and ... not in dims: raise ValueError( "arguments to transpose (%s) must be " "permuted dataset dimensions (%s)" % (dims, tuple(self.dims)) ) ds = self.copy() for name, var in self._variables.items(): var_dims = tuple(dim for dim in dims if dim in (var.dims + (...,))) ds._variables[name] = var.transpose(*var_dims) return ds def dropna( self, dim: Hashable, how: str = "any", thresh: int = None, subset: Iterable[Hashable] = None, ): """Returns a new dataset with dropped labels for missing values along the provided dimension. Parameters ---------- dim : hashable Dimension along which to drop missing values. Dropping along multiple dimensions simultaneously is not yet supported. how : {"any", "all"}, default: "any" * any : if any NA values are present, drop that label * all : if all values are NA, drop that label thresh : int, default: None If supplied, require this many non-NA values. subset : iterable of hashable, optional Which variables to check for missing values. By default, all variables in the dataset are checked. Returns ------- Dataset """ # TODO: consider supporting multiple dimensions? Or not, given that # there are some ugly edge cases, e.g., pandas's dropna differs # depending on the order of the supplied axes. if dim not in self.dims: raise ValueError("%s must be a single dataset dimension" % dim) if subset is None: subset = iter(self.data_vars) count = np.zeros(self.dims[dim], dtype=np.int64) size = 0 for k in subset: array = self._variables[k] if dim in array.dims: dims = [d for d in array.dims if d != dim] count += np.asarray(array.count(dims)) # type: ignore size += np.prod([self.dims[d] for d in dims]) if thresh is not None: mask = count >= thresh elif how == "any": mask = count == size elif how == "all": mask = count > 0 elif how is not None: raise ValueError("invalid how option: %s" % how) else: raise TypeError("must specify how or thresh") return self.isel({dim: mask}) def fillna(self, value: Any) -> "Dataset": """Fill missing values in this object. This operation follows the normal broadcasting and alignment rules that xarray uses for binary arithmetic, except the result is aligned to this object (``join='left'``) instead of aligned to the intersection of index coordinates (``join='inner'``). Parameters ---------- value : scalar, ndarray, DataArray, dict or Dataset Used to fill all matching missing values in this dataset's data variables. Scalars, ndarrays or DataArrays arguments are used to fill all data with aligned coordinates (for DataArrays). Dictionaries or datasets match data variables and then align coordinates if necessary. Returns ------- Dataset Examples -------- >>> import numpy as np >>> import xarray as xr >>> ds = xr.Dataset( ... { ... "A": ("x", [np.nan, 2, np.nan, 0]), ... "B": ("x", [3, 4, np.nan, 1]), ... "C": ("x", [np.nan, np.nan, np.nan, 5]), ... "D": ("x", [np.nan, 3, np.nan, 4]), ... }, ... coords={"x": [0, 1, 2, 3]}, ... ) >>> ds Dimensions: (x: 4) Coordinates: * x (x) int64 0 1 2 3 Data variables: A (x) float64 nan 2.0 nan 0.0 B (x) float64 3.0 4.0 nan 1.0 C (x) float64 nan nan nan 5.0 D (x) float64 nan 3.0 nan 4.0 Replace all `NaN` values with 0s. >>> ds.fillna(0) Dimensions: (x: 4) Coordinates: * x (x) int64 0 1 2 3 Data variables: A (x) float64 0.0 2.0 0.0 0.0 B (x) float64 3.0 4.0 0.0 1.0 C (x) float64 0.0 0.0 0.0 5.0 D (x) float64 0.0 3.0 0.0 4.0 Replace all `NaN` elements in column ‘A’, ‘B’, ‘C’, and ‘D’, with 0, 1, 2, and 3 respectively. >>> values = {"A": 0, "B": 1, "C": 2, "D": 3} >>> ds.fillna(value=values) Dimensions: (x: 4) Coordinates: * x (x) int64 0 1 2 3 Data variables: A (x) float64 0.0 2.0 0.0 0.0 B (x) float64 3.0 4.0 1.0 1.0 C (x) float64 2.0 2.0 2.0 5.0 D (x) float64 3.0 3.0 3.0 4.0 """ if utils.is_dict_like(value): value_keys = getattr(value, "data_vars", value).keys() if not set(value_keys) <= set(self.data_vars.keys()): raise ValueError( "all variables in the argument to `fillna` " "must be contained in the original dataset" ) out = ops.fillna(self, value) return out def interpolate_na( self, dim: Hashable = None, method: str = "linear", limit: int = None, use_coordinate: Union[bool, Hashable] = True, max_gap: Union[ int, float, str, pd.Timedelta, np.timedelta64, datetime.timedelta ] = None, **kwargs: Any, ) -> "Dataset": """Fill in NaNs by interpolating according to different methods. Parameters ---------- dim : str Specifies the dimension along which to interpolate. method : str, optional String indicating which method to use for interpolation: - 'linear': linear interpolation (Default). Additional keyword arguments are passed to :py:func:`numpy.interp` - 'nearest', 'zero', 'slinear', 'quadratic', 'cubic', 'polynomial': are passed to :py:func:`scipy.interpolate.interp1d`. If ``method='polynomial'``, the ``order`` keyword argument must also be provided. - 'barycentric', 'krog', 'pchip', 'spline', 'akima': use their respective :py:class:`scipy.interpolate` classes. use_coordinate : bool, str, default: True Specifies which index to use as the x values in the interpolation formulated as `y = f(x)`. If False, values are treated as if eqaully-spaced along ``dim``. If True, the IndexVariable `dim` is used. If ``use_coordinate`` is a string, it specifies the name of a coordinate variariable to use as the index. limit : int, default: None Maximum number of consecutive NaNs to fill. Must be greater than 0 or None for no limit. This filling is done regardless of the size of the gap in the data. To only interpolate over gaps less than a given length, see ``max_gap``. max_gap : int, float, str, pandas.Timedelta, numpy.timedelta64, datetime.timedelta, default: None Maximum size of gap, a continuous sequence of NaNs, that will be filled. Use None for no limit. When interpolating along a datetime64 dimension and ``use_coordinate=True``, ``max_gap`` can be one of the following: - a string that is valid input for pandas.to_timedelta - a :py:class:`numpy.timedelta64` object - a :py:class:`pandas.Timedelta` object - a :py:class:`datetime.timedelta` object Otherwise, ``max_gap`` must be an int or a float. Use of ``max_gap`` with unlabeled dimensions has not been implemented yet. Gap length is defined as the difference between coordinate values at the first data point after a gap and the last value before a gap. For gaps at the beginning (end), gap length is defined as the difference between coordinate values at the first (last) valid data point and the first (last) NaN. For example, consider:: array([nan, nan, nan, 1., nan, nan, 4., nan, nan]) Coordinates: * x (x) int64 0 1 2 3 4 5 6 7 8 The gap lengths are 3-0 = 3; 6-3 = 3; and 8-6 = 2 respectively kwargs : dict, optional parameters passed verbatim to the underlying interpolation function Returns ------- interpolated: Dataset Filled in Dataset. See also -------- numpy.interp scipy.interpolate Examples -------- >>> ds = xr.Dataset( ... { ... "A": ("x", [np.nan, 2, 3, np.nan, 0]), ... "B": ("x", [3, 4, np.nan, 1, 7]), ... "C": ("x", [np.nan, np.nan, np.nan, 5, 0]), ... "D": ("x", [np.nan, 3, np.nan, -1, 4]), ... }, ... coords={"x": [0, 1, 2, 3, 4]}, ... ) >>> ds Dimensions: (x: 5) Coordinates: * x (x) int64 0 1 2 3 4 Data variables: A (x) float64 nan 2.0 3.0 nan 0.0 B (x) float64 3.0 4.0 nan 1.0 7.0 C (x) float64 nan nan nan 5.0 0.0 D (x) float64 nan 3.0 nan -1.0 4.0 >>> ds.interpolate_na(dim="x", method="linear") Dimensions: (x: 5) Coordinates: * x (x) int64 0 1 2 3 4 Data variables: A (x) float64 nan 2.0 3.0 1.5 0.0 B (x) float64 3.0 4.0 2.5 1.0 7.0 C (x) float64 nan nan nan 5.0 0.0 D (x) float64 nan 3.0 1.0 -1.0 4.0 >>> ds.interpolate_na(dim="x", method="linear", fill_value="extrapolate") Dimensions: (x: 5) Coordinates: * x (x) int64 0 1 2 3 4 Data variables: A (x) float64 1.0 2.0 3.0 1.5 0.0 B (x) float64 3.0 4.0 2.5 1.0 7.0 C (x) float64 20.0 15.0 10.0 5.0 0.0 D (x) float64 5.0 3.0 1.0 -1.0 4.0 """ from .missing import _apply_over_vars_with_dim, interp_na new = _apply_over_vars_with_dim( interp_na, self, dim=dim, method=method, limit=limit, use_coordinate=use_coordinate, max_gap=max_gap, **kwargs, ) return new def ffill(self, dim: Hashable, limit: int = None) -> "Dataset": """Fill NaN values by propogating values forward *Requires bottleneck.* Parameters ---------- dim : Hashable Specifies the dimension along which to propagate values when filling. limit : int, default: None The maximum number of consecutive NaN values to forward fill. In other words, if there is a gap with more than this number of consecutive NaNs, it will only be partially filled. Must be greater than 0 or None for no limit. Returns ------- Dataset """ from .missing import _apply_over_vars_with_dim, ffill new = _apply_over_vars_with_dim(ffill, self, dim=dim, limit=limit) return new def bfill(self, dim: Hashable, limit: int = None) -> "Dataset": """Fill NaN values by propogating values backward *Requires bottleneck.* Parameters ---------- dim : str Specifies the dimension along which to propagate values when filling. limit : int, default: None The maximum number of consecutive NaN values to backward fill. In other words, if there is a gap with more than this number of consecutive NaNs, it will only be partially filled. Must be greater than 0 or None for no limit. Returns ------- Dataset """ from .missing import _apply_over_vars_with_dim, bfill new = _apply_over_vars_with_dim(bfill, self, dim=dim, limit=limit) return new def combine_first(self, other: "Dataset") -> "Dataset": """Combine two Datasets, default to data_vars of self. The new coordinates follow the normal broadcasting and alignment rules of ``join='outer'``. Vacant cells in the expanded coordinates are filled with np.nan. Parameters ---------- other : Dataset Used to fill all matching missing values in this array. Returns ------- Dataset """ out = ops.fillna(self, other, join="outer", dataset_join="outer") return out def reduce( self, func: Callable, dim: Union[Hashable, Iterable[Hashable]] = None, keep_attrs: bool = None, keepdims: bool = False, numeric_only: bool = False, **kwargs: Any, ) -> "Dataset": """Reduce this dataset by applying `func` along some dimension(s). Parameters ---------- func : callable Function which can be called in the form `f(x, axis=axis, **kwargs)` to return the result of reducing an np.ndarray over an integer valued axis. dim : str or sequence of str, optional Dimension(s) over which to apply `func`. By default `func` is applied over all dimensions. keep_attrs : bool, optional If True, the dataset's attributes (`attrs`) will be copied from the original object to the new one. If False (default), the new object will be returned without attributes. keepdims : bool, default: False If True, the dimensions which are reduced are left in the result as dimensions of size one. Coordinates that use these dimensions are removed. numeric_only : bool, optional If True, only apply ``func`` to variables with a numeric dtype. **kwargs : Any Additional keyword arguments passed on to ``func``. Returns ------- reduced : Dataset Dataset with this object's DataArrays replaced with new DataArrays of summarized data and the indicated dimension(s) removed. """ if dim is None or dim is ...: dims = set(self.dims) elif isinstance(dim, str) or not isinstance(dim, Iterable): dims = {dim} else: dims = set(dim) missing_dimensions = [d for d in dims if d not in self.dims] if missing_dimensions: raise ValueError( "Dataset does not contain the dimensions: %s" % missing_dimensions ) if keep_attrs is None: keep_attrs = _get_keep_attrs(default=False) variables: Dict[Hashable, Variable] = {} for name, var in self._variables.items(): reduce_dims = [d for d in var.dims if d in dims] if name in self.coords: if not reduce_dims: variables[name] = var else: if ( not numeric_only or np.issubdtype(var.dtype, np.number) or (var.dtype == np.bool_) ): if len(reduce_dims) == 1: # unpack dimensions for the benefit of functions # like np.argmin which can't handle tuple arguments (reduce_dims,) = reduce_dims elif len(reduce_dims) == var.ndim: # prefer to aggregate over axis=None rather than # axis=(0, 1) if they will be equivalent, because # the former is often more efficient reduce_dims = None # type: ignore variables[name] = var.reduce( func, dim=reduce_dims, keep_attrs=keep_attrs, keepdims=keepdims, **kwargs, ) coord_names = {k for k in self.coords if k in variables} indexes = {k: v for k, v in self.indexes.items() if k in variables} attrs = self.attrs if keep_attrs else None return self._replace_with_new_dims( variables, coord_names=coord_names, attrs=attrs, indexes=indexes ) def map( self, func: Callable, keep_attrs: bool = None, args: Iterable[Any] = (), **kwargs: Any, ) -> "Dataset": """Apply a function to each variable in this dataset Parameters ---------- func : callable Function which can be called in the form `func(x, *args, **kwargs)` to transform each DataArray `x` in this dataset into another DataArray. keep_attrs : bool, optional If True, the dataset's attributes (`attrs`) will be copied from the original object to the new one. If False, the new object will be returned without attributes. args : tuple, optional Positional arguments passed on to `func`. **kwargs : Any Keyword arguments passed on to `func`. Returns ------- applied : Dataset Resulting dataset from applying ``func`` to each data variable. Examples -------- >>> da = xr.DataArray(np.random.randn(2, 3)) >>> ds = xr.Dataset({"foo": da, "bar": ("x", [-1, 2])}) >>> ds Dimensions: (dim_0: 2, dim_1: 3, x: 2) Dimensions without coordinates: dim_0, dim_1, x Data variables: foo (dim_0, dim_1) float64 1.764 0.4002 0.9787 2.241 1.868 -0.9773 bar (x) int64 -1 2 >>> ds.map(np.fabs) Dimensions: (dim_0: 2, dim_1: 3, x: 2) Dimensions without coordinates: dim_0, dim_1, x Data variables: foo (dim_0, dim_1) float64 1.764 0.4002 0.9787 2.241 1.868 0.9773 bar (x) float64 1.0 2.0 """ if keep_attrs is None: keep_attrs = _get_keep_attrs(default=False) variables = { k: maybe_wrap_array(v, func(v, *args, **kwargs)) for k, v in self.data_vars.items() } if keep_attrs: for k, v in variables.items(): v._copy_attrs_from(self.data_vars[k]) attrs = self.attrs if keep_attrs else None return type(self)(variables, attrs=attrs) def apply( self, func: Callable, keep_attrs: bool = None, args: Iterable[Any] = (), **kwargs: Any, ) -> "Dataset": """ Backward compatible implementation of ``map`` See Also -------- Dataset.map """ warnings.warn( "Dataset.apply may be deprecated in the future. Using Dataset.map is encouraged", PendingDeprecationWarning, stacklevel=2, ) return self.map(func, keep_attrs, args, **kwargs) def assign( self, variables: Mapping[Hashable, Any] = None, **variables_kwargs: Hashable ) -> "Dataset": """Assign new data variables to a Dataset, returning a new object with all the original variables in addition to the new ones. Parameters ---------- variables : mapping of hashable to Any Mapping from variables names to the new values. If the new values are callable, they are computed on the Dataset and assigned to new data variables. If the values are not callable, (e.g. a DataArray, scalar, or array), they are simply assigned. **variables_kwargs The keyword arguments form of ``variables``. One of variables or variables_kwargs must be provided. Returns ------- ds : Dataset A new Dataset with the new variables in addition to all the existing variables. Notes ----- Since ``kwargs`` is a dictionary, the order of your arguments may not be preserved, and so the order of the new variables is not well defined. Assigning multiple variables within the same ``assign`` is possible, but you cannot reference other variables created within the same ``assign`` call. See Also -------- pandas.DataFrame.assign Examples -------- >>> x = xr.Dataset( ... { ... "temperature_c": ( ... ("lat", "lon"), ... 20 * np.random.rand(4).reshape(2, 2), ... ), ... "precipitation": (("lat", "lon"), np.random.rand(4).reshape(2, 2)), ... }, ... coords={"lat": [10, 20], "lon": [150, 160]}, ... ) >>> x Dimensions: (lat: 2, lon: 2) Coordinates: * lat (lat) int64 10 20 * lon (lon) int64 150 160 Data variables: temperature_c (lat, lon) float64 10.98 14.3 12.06 10.9 precipitation (lat, lon) float64 0.4237 0.6459 0.4376 0.8918 Where the value is a callable, evaluated on dataset: >>> x.assign(temperature_f=lambda x: x.temperature_c * 9 / 5 + 32) Dimensions: (lat: 2, lon: 2) Coordinates: * lat (lat) int64 10 20 * lon (lon) int64 150 160 Data variables: temperature_c (lat, lon) float64 10.98 14.3 12.06 10.9 precipitation (lat, lon) float64 0.4237 0.6459 0.4376 0.8918 temperature_f (lat, lon) float64 51.76 57.75 53.7 51.62 Alternatively, the same behavior can be achieved by directly referencing an existing dataarray: >>> x.assign(temperature_f=x["temperature_c"] * 9 / 5 + 32) Dimensions: (lat: 2, lon: 2) Coordinates: * lat (lat) int64 10 20 * lon (lon) int64 150 160 Data variables: temperature_c (lat, lon) float64 10.98 14.3 12.06 10.9 precipitation (lat, lon) float64 0.4237 0.6459 0.4376 0.8918 temperature_f (lat, lon) float64 51.76 57.75 53.7 51.62 """ variables = either_dict_or_kwargs(variables, variables_kwargs, "assign") data = self.copy() # do all calculations first... results = data._calc_assign_results(variables) # ... and then assign data.update(results) return data def to_array(self, dim="variable", name=None): """Convert this dataset into an xarray.DataArray The data variables of this dataset will be broadcast against each other and stacked along the first axis of the new array. All coordinates of this dataset will remain coordinates. Parameters ---------- dim : str, optional Name of the new dimension. name : str, optional Name of the new data array. Returns ------- array : xarray.DataArray """ from .dataarray import DataArray data_vars = [self.variables[k] for k in self.data_vars] broadcast_vars = broadcast_variables(*data_vars) data = duck_array_ops.stack([b.data for b in broadcast_vars], axis=0) coords = dict(self.coords) coords[dim] = list(self.data_vars) indexes = propagate_indexes(self._indexes) dims = (dim,) + broadcast_vars[0].dims return DataArray( data, coords, dims, attrs=self.attrs, name=name, indexes=indexes ) def _normalize_dim_order( self, dim_order: List[Hashable] = None ) -> Dict[Hashable, int]: """ Check the validity of the provided dimensions if any and return the mapping between dimension name and their size. Parameters ---------- dim_order Dimension order to validate (default to the alphabetical order if None). Returns ------- result Validated dimensions mapping. """ if dim_order is None: dim_order = list(self.dims) elif set(dim_order) != set(self.dims): raise ValueError( "dim_order {} does not match the set of dimensions of this " "Dataset: {}".format(dim_order, list(self.dims)) ) ordered_dims = {k: self.dims[k] for k in dim_order} return ordered_dims def _to_dataframe(self, ordered_dims: Mapping[Hashable, int]): columns = [k for k in self.variables if k not in self.dims] data = [ self._variables[k].set_dims(ordered_dims).values.reshape(-1) for k in columns ] index = self.coords.to_index([*ordered_dims]) return pd.DataFrame(dict(zip(columns, data)), index=index) def to_dataframe(self, dim_order: List[Hashable] = None) -> pd.DataFrame: """Convert this dataset into a pandas.DataFrame. Non-index variables in this dataset form the columns of the DataFrame. The DataFrame is indexed by the Cartesian product of this dataset's indices. Parameters ---------- dim_order Hierarchical dimension order for the resulting dataframe. All arrays are transposed to this order and then written out as flat vectors in contiguous order, so the last dimension in this list will be contiguous in the resulting DataFrame. This has a major influence on which operations are efficient on the resulting dataframe. If provided, must include all dimensions of this dataset. By default, dimensions are sorted alphabetically. Returns ------- result Dataset as a pandas DataFrame. """ ordered_dims = self._normalize_dim_order(dim_order=dim_order) return self._to_dataframe(ordered_dims=ordered_dims) def _set_sparse_data_from_dataframe( self, idx: pd.Index, arrays: List[Tuple[Hashable, np.ndarray]], dims: tuple ) -> None: from sparse import COO if isinstance(idx, pd.MultiIndex): coords = np.stack([np.asarray(code) for code in idx.codes], axis=0) is_sorted = idx.is_lexsorted() shape = tuple(lev.size for lev in idx.levels) else: coords = np.arange(idx.size).reshape(1, -1) is_sorted = True shape = (idx.size,) for name, values in arrays: # In virtually all real use cases, the sparse array will now have # missing values and needs a fill_value. For consistency, don't # special case the rare exceptions (e.g., dtype=int without a # MultiIndex). dtype, fill_value = dtypes.maybe_promote(values.dtype) values = np.asarray(values, dtype=dtype) data = COO( coords, values, shape, has_duplicates=False, sorted=is_sorted, fill_value=fill_value, ) self[name] = (dims, data) def _set_numpy_data_from_dataframe( self, idx: pd.Index, arrays: List[Tuple[Hashable, np.ndarray]], dims: tuple ) -> None: if not isinstance(idx, pd.MultiIndex): for name, values in arrays: self[name] = (dims, values) return shape = tuple(lev.size for lev in idx.levels) indexer = tuple(idx.codes) # We already verified that the MultiIndex has all unique values, so # there are missing values if and only if the size of output arrays is # larger that the index. missing_values = np.prod(shape) > idx.shape[0] for name, values in arrays: # NumPy indexing is much faster than using DataFrame.reindex() to # fill in missing values: # https://stackoverflow.com/a/35049899/809705 if missing_values: dtype, fill_value = dtypes.maybe_promote(values.dtype) data = np.full(shape, fill_value, dtype) else: # If there are no missing values, keep the existing dtype # instead of promoting to support NA, e.g., keep integer # columns as integers. # TODO: consider removing this special case, which doesn't # exist for sparse=True. data = np.zeros(shape, values.dtype) data[indexer] = values self[name] = (dims, data) @classmethod def from_dataframe(cls, dataframe: pd.DataFrame, sparse: bool = False) -> "Dataset": """Convert a pandas.DataFrame into an xarray.Dataset Each column will be converted into an independent variable in the Dataset. If the dataframe's index is a MultiIndex, it will be expanded into a tensor product of one-dimensional indices (filling in missing values with NaN). This method will produce a Dataset very similar to that on which the 'to_dataframe' method was called, except with possibly redundant dimensions (since all dataset variables will have the same dimensionality) Parameters ---------- dataframe : DataFrame DataFrame from which to copy data and indices. sparse : bool, default: False If true, create a sparse arrays instead of dense numpy arrays. This can potentially save a large amount of memory if the DataFrame has a MultiIndex. Requires the sparse package (sparse.pydata.org). Returns ------- New Dataset. See also -------- xarray.DataArray.from_series pandas.DataFrame.to_xarray """ # TODO: Add an option to remove dimensions along which the variables # are constant, to enable consistent serialization to/from a dataframe, # even if some variables have different dimensionality. if not dataframe.columns.is_unique: raise ValueError("cannot convert DataFrame with non-unique columns") idx = remove_unused_levels_categories(dataframe.index) if isinstance(idx, pd.MultiIndex) and not idx.is_unique: raise ValueError( "cannot convert a DataFrame with a non-unique MultiIndex into xarray" ) # Cast to a NumPy array first, in case the Series is a pandas Extension # array (which doesn't have a valid NumPy dtype) # TODO: allow users to control how this casting happens, e.g., by # forwarding arguments to pandas.Series.to_numpy? arrays = [(k, np.asarray(v)) for k, v in dataframe.items()] obj = cls() if isinstance(idx, pd.MultiIndex): dims = tuple( name if name is not None else "level_%i" % n for n, name in enumerate(idx.names) ) for dim, lev in zip(dims, idx.levels): obj[dim] = (dim, lev) else: index_name = idx.name if idx.name is not None else "index" dims = (index_name,) obj[index_name] = (dims, idx) if sparse: obj._set_sparse_data_from_dataframe(idx, arrays, dims) else: obj._set_numpy_data_from_dataframe(idx, arrays, dims) return obj def to_dask_dataframe(self, dim_order=None, set_index=False): """ Convert this dataset into a dask.dataframe.DataFrame. The dimensions, coordinates and data variables in this dataset form the columns of the DataFrame. Parameters ---------- dim_order : list, optional Hierarchical dimension order for the resulting dataframe. All arrays are transposed to this order and then written out as flat vectors in contiguous order, so the last dimension in this list will be contiguous in the resulting DataFrame. This has a major influence on which operations are efficient on the resulting dask dataframe. If provided, must include all dimensions of this dataset. By default, dimensions are sorted alphabetically. set_index : bool, optional If set_index=True, the dask DataFrame is indexed by this dataset's coordinate. Since dask DataFrames do not support multi-indexes, set_index only works if the dataset only contains one dimension. Returns ------- dask.dataframe.DataFrame """ import dask.array as da import dask.dataframe as dd ordered_dims = self._normalize_dim_order(dim_order=dim_order) columns = list(ordered_dims) columns.extend(k for k in self.coords if k not in self.dims) columns.extend(self.data_vars) series_list = [] for name in columns: try: var = self.variables[name] except KeyError: # dimension without a matching coordinate size = self.dims[name] data = da.arange(size, chunks=size, dtype=np.int64) var = Variable((name,), data) # IndexVariable objects have a dummy .chunk() method if isinstance(var, IndexVariable): var = var.to_base_variable() dask_array = var.set_dims(ordered_dims).chunk(self.chunks).data series = dd.from_array(dask_array.reshape(-1), columns=[name]) series_list.append(series) df = dd.concat(series_list, axis=1) if set_index: dim_order = [*ordered_dims] if len(dim_order) == 1: (dim,) = dim_order df = df.set_index(dim) else: # triggers an error about multi-indexes, even if only one # dimension is passed df = df.set_index(dim_order) return df def to_dict(self, data=True): """ Convert this dataset to a dictionary following xarray naming conventions. Converts all variables and attributes to native Python objects Useful for converting to json. To avoid datetime incompatibility use decode_times=False kwarg in xarrray.open_dataset. Parameters ---------- data : bool, optional Whether to include the actual data in the dictionary. When set to False, returns just the schema. See also -------- Dataset.from_dict """ d = { "coords": {}, "attrs": decode_numpy_dict_values(self.attrs), "dims": dict(self.dims), "data_vars": {}, } for k in self.coords: d["coords"].update({k: self[k].variable.to_dict(data=data)}) for k in self.data_vars: d["data_vars"].update({k: self[k].variable.to_dict(data=data)}) return d @classmethod def from_dict(cls, d): """ Convert a dictionary into an xarray.Dataset. Input dict can take several forms: .. code:: python d = { "t": {"dims": ("t"), "data": t}, "a": {"dims": ("t"), "data": x}, "b": {"dims": ("t"), "data": y}, } d = { "coords": {"t": {"dims": "t", "data": t, "attrs": {"units": "s"}}}, "attrs": {"title": "air temperature"}, "dims": "t", "data_vars": { "a": {"dims": "t", "data": x}, "b": {"dims": "t", "data": y}, }, } where "t" is the name of the dimesion, "a" and "b" are names of data variables and t, x, and y are lists, numpy.arrays or pandas objects. Parameters ---------- d : dict-like Mapping with a minimum structure of ``{"var_0": {"dims": [..], "data": [..]}, \ ...}`` Returns ------- obj : xarray.Dataset See also -------- Dataset.to_dict DataArray.from_dict """ if not {"coords", "data_vars"}.issubset(set(d)): variables = d.items() else: import itertools variables = itertools.chain( d.get("coords", {}).items(), d.get("data_vars", {}).items() ) try: variable_dict = { k: (v["dims"], v["data"], v.get("attrs")) for k, v in variables } except KeyError as e: raise ValueError( "cannot convert dict without the key " "'{dims_data}'".format(dims_data=str(e.args[0])) ) obj = cls(variable_dict) # what if coords aren't dims? coords = set(d.get("coords", {})) - set(d.get("dims", {})) obj = obj.set_coords(coords) obj.attrs.update(d.get("attrs", {})) return obj @staticmethod def _unary_op(f): @functools.wraps(f) def func(self, *args, **kwargs): variables = {} keep_attrs = kwargs.pop("keep_attrs", None) if keep_attrs is None: keep_attrs = _get_keep_attrs(default=True) for k, v in self._variables.items(): if k in self._coord_names: variables[k] = v else: variables[k] = f(v, *args, **kwargs) if keep_attrs: variables[k].attrs = v._attrs attrs = self._attrs if keep_attrs else None return self._replace_with_new_dims(variables, attrs=attrs) return func @staticmethod def _binary_op(f, reflexive=False, join=None): @functools.wraps(f) def func(self, other): from .dataarray import DataArray if isinstance(other, groupby.GroupBy): return NotImplemented align_type = OPTIONS["arithmetic_join"] if join is None else join if isinstance(other, (DataArray, Dataset)): self, other = align(self, other, join=align_type, copy=False) g = f if not reflexive else lambda x, y: f(y, x) ds = self._calculate_binary_op(g, other, join=align_type) return ds return func @staticmethod def _inplace_binary_op(f): @functools.wraps(f) def func(self, other): from .dataarray import DataArray if isinstance(other, groupby.GroupBy): raise TypeError( "in-place operations between a Dataset and " "a grouped object are not permitted" ) # we don't actually modify arrays in-place with in-place Dataset # arithmetic -- this lets us automatically align things if isinstance(other, (DataArray, Dataset)): other = other.reindex_like(self, copy=False) g = ops.inplace_to_noninplace_op(f) ds = self._calculate_binary_op(g, other, inplace=True) self._replace_with_new_dims( ds._variables, ds._coord_names, attrs=ds._attrs, indexes=ds._indexes, inplace=True, ) return self return func def _calculate_binary_op(self, f, other, join="inner", inplace=False): def apply_over_both(lhs_data_vars, rhs_data_vars, lhs_vars, rhs_vars): if inplace and set(lhs_data_vars) != set(rhs_data_vars): raise ValueError( "datasets must have the same data variables " "for in-place arithmetic operations: %s, %s" % (list(lhs_data_vars), list(rhs_data_vars)) ) dest_vars = {} for k in lhs_data_vars: if k in rhs_data_vars: dest_vars[k] = f(lhs_vars[k], rhs_vars[k]) elif join in ["left", "outer"]: dest_vars[k] = f(lhs_vars[k], np.nan) for k in rhs_data_vars: if k not in dest_vars and join in ["right", "outer"]: dest_vars[k] = f(rhs_vars[k], np.nan) return dest_vars if utils.is_dict_like(other) and not isinstance(other, Dataset): # can't use our shortcut of doing the binary operation with # Variable objects, so apply over our data vars instead. new_data_vars = apply_over_both( self.data_vars, other, self.data_vars, other ) return Dataset(new_data_vars) other_coords = getattr(other, "coords", None) ds = self.coords.merge(other_coords) if isinstance(other, Dataset): new_vars = apply_over_both( self.data_vars, other.data_vars, self.variables, other.variables ) else: other_variable = getattr(other, "variable", other) new_vars = {k: f(self.variables[k], other_variable) for k in self.data_vars} ds._variables.update(new_vars) ds._dims = calculate_dimensions(ds._variables) return ds def _copy_attrs_from(self, other): self.attrs = other.attrs for v in other.variables: if v in self.variables: self.variables[v].attrs = other.variables[v].attrs def diff(self, dim, n=1, label="upper"): """Calculate the n-th order discrete difference along given axis. Parameters ---------- dim : str Dimension over which to calculate the finite difference. n : int, optional The number of times values are differenced. label : str, optional The new coordinate in dimension ``dim`` will have the values of either the minuend's or subtrahend's coordinate for values 'upper' and 'lower', respectively. Other values are not supported. Returns ------- difference : same type as caller The n-th order finite difference of this object. .. note:: `n` matches numpy's behavior and is different from pandas' first argument named `periods`. Examples -------- >>> ds = xr.Dataset({"foo": ("x", [5, 5, 6, 6])}) >>> ds.diff("x") Dimensions: (x: 3) Dimensions without coordinates: x Data variables: foo (x) int64 0 1 0 >>> ds.diff("x", 2) Dimensions: (x: 2) Dimensions without coordinates: x Data variables: foo (x) int64 1 -1 See Also -------- Dataset.differentiate """ if n == 0: return self if n < 0: raise ValueError(f"order `n` must be non-negative but got {n}") # prepare slices kwargs_start = {dim: slice(None, -1)} kwargs_end = {dim: slice(1, None)} # prepare new coordinate if label == "upper": kwargs_new = kwargs_end elif label == "lower": kwargs_new = kwargs_start else: raise ValueError("The 'label' argument has to be either 'upper' or 'lower'") variables = {} for name, var in self.variables.items(): if dim in var.dims: if name in self.data_vars: variables[name] = var.isel(**kwargs_end) - var.isel(**kwargs_start) else: variables[name] = var.isel(**kwargs_new) else: variables[name] = var indexes = dict(self.indexes) if dim in indexes: indexes[dim] = indexes[dim][kwargs_new[dim]] difference = self._replace_with_new_dims(variables, indexes=indexes) if n > 1: return difference.diff(dim, n - 1) else: return difference def shift(self, shifts=None, fill_value=dtypes.NA, **shifts_kwargs): """Shift this dataset by an offset along one or more dimensions. Only data variables are moved; coordinates stay in place. This is consistent with the behavior of ``shift`` in pandas. Parameters ---------- shifts : mapping of hashable to int Integer offset to shift along each of the given dimensions. Positive offsets shift to the right; negative offsets shift to the left. fill_value : scalar or dict-like, optional Value to use for newly missing values. If a dict-like, maps variable names (including coordinates) to fill values. **shifts_kwargs The keyword arguments form of ``shifts``. One of shifts or shifts_kwargs must be provided. Returns ------- shifted : Dataset Dataset with the same coordinates and attributes but shifted data variables. See also -------- roll Examples -------- >>> ds = xr.Dataset({"foo": ("x", list("abcde"))}) >>> ds.shift(x=2) Dimensions: (x: 5) Dimensions without coordinates: x Data variables: foo (x) object nan nan 'a' 'b' 'c' """ shifts = either_dict_or_kwargs(shifts, shifts_kwargs, "shift") invalid = [k for k in shifts if k not in self.dims] if invalid: raise ValueError("dimensions %r do not exist" % invalid) variables = {} for name, var in self.variables.items(): if name in self.data_vars: fill_value_ = ( fill_value.get(name, dtypes.NA) if isinstance(fill_value, dict) else fill_value ) var_shifts = {k: v for k, v in shifts.items() if k in var.dims} variables[name] = var.shift(fill_value=fill_value_, shifts=var_shifts) else: variables[name] = var return self._replace(variables) def roll(self, shifts=None, roll_coords=None, **shifts_kwargs): """Roll this dataset by an offset along one or more dimensions. Unlike shift, roll may rotate all variables, including coordinates if specified. The direction of rotation is consistent with :py:func:`numpy.roll`. Parameters ---------- shifts : dict, optional A dict with keys matching dimensions and values given by integers to rotate each of the given dimensions. Positive offsets roll to the right; negative offsets roll to the left. roll_coords : bool Indicates whether to roll the coordinates by the offset The current default of roll_coords (None, equivalent to True) is deprecated and will change to False in a future version. Explicitly pass roll_coords to silence the warning. **shifts_kwargs : {dim: offset, ...}, optional The keyword arguments form of ``shifts``. One of shifts or shifts_kwargs must be provided. Returns ------- rolled : Dataset Dataset with the same coordinates and attributes but rolled variables. See also -------- shift Examples -------- >>> ds = xr.Dataset({"foo": ("x", list("abcde"))}) >>> ds.roll(x=2) Dimensions: (x: 5) Dimensions without coordinates: x Data variables: foo (x) >> ds = xr.Dataset( ... {"a": (("x", "y"), [[0.7, 4.2, 9.4, 1.5], [6.5, 7.3, 2.6, 1.9]])}, ... coords={"x": [7, 9], "y": [1, 1.5, 2, 2.5]}, ... ) >>> ds.quantile(0) # or ds.quantile(0, dim=...) Dimensions: () Coordinates: quantile float64 0.0 Data variables: a float64 0.7 >>> ds.quantile(0, dim="x") Dimensions: (y: 4) Coordinates: * y (y) float64 1.0 1.5 2.0 2.5 quantile float64 0.0 Data variables: a (y) float64 0.7 4.2 2.6 1.5 >>> ds.quantile([0, 0.5, 1]) Dimensions: (quantile: 3) Coordinates: * quantile (quantile) float64 0.0 0.5 1.0 Data variables: a (quantile) float64 0.7 3.4 9.4 >>> ds.quantile([0, 0.5, 1], dim="x") Dimensions: (quantile: 3, y: 4) Coordinates: * y (y) float64 1.0 1.5 2.0 2.5 * quantile (quantile) float64 0.0 0.5 1.0 Data variables: a (quantile, y) float64 0.7 4.2 2.6 1.5 3.6 ... 1.7 6.5 7.3 9.4 1.9 """ if isinstance(dim, str): dims = {dim} elif dim in [None, ...]: dims = set(self.dims) else: dims = set(dim) _assert_empty( [d for d in dims if d not in self.dims], "Dataset does not contain the dimensions: %s", ) q = np.asarray(q, dtype=np.float64) variables = {} for name, var in self.variables.items(): reduce_dims = [d for d in var.dims if d in dims] if reduce_dims or not var.dims: if name not in self.coords: if ( not numeric_only or np.issubdtype(var.dtype, np.number) or var.dtype == np.bool_ ): if len(reduce_dims) == var.ndim: # prefer to aggregate over axis=None rather than # axis=(0, 1) if they will be equivalent, because # the former is often more efficient reduce_dims = None variables[name] = var.quantile( q, dim=reduce_dims, interpolation=interpolation, keep_attrs=keep_attrs, skipna=skipna, ) else: variables[name] = var # construct the new dataset coord_names = {k for k in self.coords if k in variables} indexes = {k: v for k, v in self.indexes.items() if k in variables} if keep_attrs is None: keep_attrs = _get_keep_attrs(default=False) attrs = self.attrs if keep_attrs else None new = self._replace_with_new_dims( variables, coord_names=coord_names, attrs=attrs, indexes=indexes ) return new.assign_coords(quantile=q) def rank(self, dim, pct=False, keep_attrs=None): """Ranks the data. Equal values are assigned a rank that is the average of the ranks that would have been otherwise assigned to all of the values within that set. Ranks begin at 1, not 0. If pct is True, computes percentage ranks. NaNs in the input array are returned as NaNs. The `bottleneck` library is required. Parameters ---------- dim : str Dimension over which to compute rank. pct : bool, optional If True, compute percentage ranks, otherwise compute integer ranks. keep_attrs : bool, optional If True, the dataset's attributes (`attrs`) will be copied from the original object to the new one. If False (default), the new object will be returned without attributes. Returns ------- ranked : Dataset Variables that do not depend on `dim` are dropped. """ if dim not in self.dims: raise ValueError("Dataset does not contain the dimension: %s" % dim) variables = {} for name, var in self.variables.items(): if name in self.data_vars: if dim in var.dims: variables[name] = var.rank(dim, pct=pct) else: variables[name] = var coord_names = set(self.coords) if keep_attrs is None: keep_attrs = _get_keep_attrs(default=False) attrs = self.attrs if keep_attrs else None return self._replace(variables, coord_names, attrs=attrs) def differentiate(self, coord, edge_order=1, datetime_unit=None): """ Differentiate with the second order accurate central differences. .. note:: This feature is limited to simple cartesian geometry, i.e. coord must be one dimensional. Parameters ---------- coord : str The coordinate to be used to compute the gradient. edge_order : {1, 2}, default: 1 N-th order accurate differences at the boundaries. datetime_unit : None or {"Y", "M", "W", "D", "h", "m", "s", "ms", \ "us", "ns", "ps", "fs", "as"}, default: None Unit to compute gradient. Only valid for datetime coordinate. Returns ------- differentiated: Dataset See also -------- numpy.gradient: corresponding numpy function """ from .variable import Variable if coord not in self.variables and coord not in self.dims: raise ValueError(f"Coordinate {coord} does not exist.") coord_var = self[coord].variable if coord_var.ndim != 1: raise ValueError( "Coordinate {} must be 1 dimensional but is {}" " dimensional".format(coord, coord_var.ndim) ) dim = coord_var.dims[0] if _contains_datetime_like_objects(coord_var): if coord_var.dtype.kind in "mM" and datetime_unit is None: datetime_unit, _ = np.datetime_data(coord_var.dtype) elif datetime_unit is None: datetime_unit = "s" # Default to seconds for cftime objects coord_var = coord_var._to_numeric(datetime_unit=datetime_unit) variables = {} for k, v in self.variables.items(): if k in self.data_vars and dim in v.dims and k not in self.coords: if _contains_datetime_like_objects(v): v = v._to_numeric(datetime_unit=datetime_unit) grad = duck_array_ops.gradient( v.data, coord_var, edge_order=edge_order, axis=v.get_axis_num(dim) ) variables[k] = Variable(v.dims, grad) else: variables[k] = v return self._replace(variables) def integrate(self, coord, datetime_unit=None): """ integrate the array with the trapezoidal rule. .. note:: This feature is limited to simple cartesian geometry, i.e. coord must be one dimensional. Parameters ---------- coord: str, or sequence of str Coordinate(s) used for the integration. datetime_unit : {"Y", "M", "W", "D", "h", "m", "s", "ms", "us", "ns", \ "ps", "fs", "as"}, optional Can be specify the unit if datetime coordinate is used. Returns ------- integrated : Dataset See also -------- DataArray.integrate numpy.trapz: corresponding numpy function Examples -------- >>> ds = xr.Dataset( ... data_vars={"a": ("x", [5, 5, 6, 6]), "b": ("x", [1, 2, 1, 0])}, ... coords={"x": [0, 1, 2, 3], "y": ("x", [1, 7, 3, 5])}, ... ) >>> ds Dimensions: (x: 4) Coordinates: * x (x) int64 0 1 2 3 y (x) int64 1 7 3 5 Data variables: a (x) int64 5 5 6 6 b (x) int64 1 2 1 0 >>> ds.integrate("x") Dimensions: () Data variables: a float64 16.5 b float64 3.5 >>> ds.integrate("y") Dimensions: () Data variables: a float64 20.0 b float64 4.0 """ if not isinstance(coord, (list, tuple)): coord = (coord,) result = self for c in coord: result = result._integrate_one(c, datetime_unit=datetime_unit) return result def _integrate_one(self, coord, datetime_unit=None): from .variable import Variable if coord not in self.variables and coord not in self.dims: raise ValueError(f"Coordinate {coord} does not exist.") coord_var = self[coord].variable if coord_var.ndim != 1: raise ValueError( "Coordinate {} must be 1 dimensional but is {}" " dimensional".format(coord, coord_var.ndim) ) dim = coord_var.dims[0] if _contains_datetime_like_objects(coord_var): if coord_var.dtype.kind in "mM" and datetime_unit is None: datetime_unit, _ = np.datetime_data(coord_var.dtype) elif datetime_unit is None: datetime_unit = "s" # Default to seconds for cftime objects coord_var = coord_var._replace( data=datetime_to_numeric(coord_var.data, datetime_unit=datetime_unit) ) variables = {} coord_names = set() for k, v in self.variables.items(): if k in self.coords: if dim not in v.dims: variables[k] = v coord_names.add(k) else: if k in self.data_vars and dim in v.dims: if _contains_datetime_like_objects(v): v = datetime_to_numeric(v, datetime_unit=datetime_unit) integ = duck_array_ops.trapz( v.data, coord_var.data, axis=v.get_axis_num(dim) ) v_dims = list(v.dims) v_dims.remove(dim) variables[k] = Variable(v_dims, integ) else: variables[k] = v indexes = {k: v for k, v in self.indexes.items() if k in variables} return self._replace_with_new_dims( variables, coord_names=coord_names, indexes=indexes ) @property def real(self): return self.map(lambda x: x.real, keep_attrs=True) @property def imag(self): return self.map(lambda x: x.imag, keep_attrs=True) plot = utils.UncachedAccessor(_Dataset_PlotMethods) def filter_by_attrs(self, **kwargs): """Returns a ``Dataset`` with variables that match specific conditions. Can pass in ``key=value`` or ``key=callable``. A Dataset is returned containing only the variables for which all the filter tests pass. These tests are either ``key=value`` for which the attribute ``key`` has the exact value ``value`` or the callable passed into ``key=callable`` returns True. The callable will be passed a single value, either the value of the attribute ``key`` or ``None`` if the DataArray does not have an attribute with the name ``key``. Parameters ---------- **kwargs key : str Attribute name. value : callable or obj If value is a callable, it should return a boolean in the form of bool = func(attr) where attr is da.attrs[key]. Otherwise, value will be compared to the each DataArray's attrs[key]. Returns ------- new : Dataset New dataset with variables filtered by attribute. Examples -------- >>> # Create an example dataset: ... >>> import numpy as np >>> import pandas as pd >>> import xarray as xr >>> temp = 15 + 8 * np.random.randn(2, 2, 3) >>> precip = 10 * np.random.rand(2, 2, 3) >>> lon = [[-99.83, -99.32], [-99.79, -99.23]] >>> lat = [[42.25, 42.21], [42.63, 42.59]] >>> dims = ["x", "y", "time"] >>> temp_attr = dict(standard_name="air_potential_temperature") >>> precip_attr = dict(standard_name="convective_precipitation_flux") >>> ds = xr.Dataset( ... { ... "temperature": (dims, temp, temp_attr), ... "precipitation": (dims, precip, precip_attr), ... }, ... coords={ ... "lon": (["x", "y"], lon), ... "lat": (["x", "y"], lat), ... "time": pd.date_range("2014-09-06", periods=3), ... "reference_time": pd.Timestamp("2014-09-05"), ... }, ... ) >>> # Get variables matching a specific standard_name. >>> ds.filter_by_attrs(standard_name="convective_precipitation_flux") Dimensions: (time: 3, x: 2, y: 2) Coordinates: lon (x, y) float64 -99.83 -99.32 -99.79 -99.23 lat (x, y) float64 42.25 42.21 42.63 42.59 * time (time) datetime64[ns] 2014-09-06 2014-09-07 2014-09-08 reference_time datetime64[ns] 2014-09-05 Dimensions without coordinates: x, y Data variables: precipitation (x, y, time) float64 5.68 9.256 0.7104 ... 7.992 4.615 7.805 >>> # Get all variables that have a standard_name attribute. >>> standard_name = lambda v: v is not None >>> ds.filter_by_attrs(standard_name=standard_name) Dimensions: (time: 3, x: 2, y: 2) Coordinates: lon (x, y) float64 -99.83 -99.32 -99.79 -99.23 lat (x, y) float64 42.25 42.21 42.63 42.59 * time (time) datetime64[ns] 2014-09-06 2014-09-07 2014-09-08 reference_time datetime64[ns] 2014-09-05 Dimensions without coordinates: x, y Data variables: temperature (x, y, time) float64 29.11 18.2 22.83 ... 18.28 16.15 26.63 precipitation (x, y, time) float64 5.68 9.256 0.7104 ... 7.992 4.615 7.805 """ selection = [] for var_name, variable in self.variables.items(): has_value_flag = False for attr_name, pattern in kwargs.items(): attr_value = variable.attrs.get(attr_name) if (callable(pattern) and pattern(attr_value)) or attr_value == pattern: has_value_flag = True else: has_value_flag = False break if has_value_flag is True: selection.append(var_name) return self[selection] def unify_chunks(self) -> "Dataset": """Unify chunk size along all chunked dimensions of this Dataset. Returns ------- Dataset with consistent chunk sizes for all dask-array variables See Also -------- dask.array.core.unify_chunks """ try: self.chunks except ValueError: # "inconsistent chunks" pass else: # No variables with dask backend, or all chunks are already aligned return self.copy() # import dask is placed after the quick exit test above to allow # running this method if dask isn't installed and there are no chunks import dask.array ds = self.copy() dims_pos_map = {dim: index for index, dim in enumerate(ds.dims)} dask_array_names = [] dask_unify_args = [] for name, variable in ds.variables.items(): if isinstance(variable.data, dask.array.Array): dims_tuple = [dims_pos_map[dim] for dim in variable.dims] dask_array_names.append(name) dask_unify_args.append(variable.data) dask_unify_args.append(dims_tuple) _, rechunked_arrays = dask.array.core.unify_chunks(*dask_unify_args) for name, new_array in zip(dask_array_names, rechunked_arrays): ds.variables[name]._data = new_array return ds def map_blocks( self, func: "Callable[..., T_DSorDA]", args: Sequence[Any] = (), kwargs: Mapping[str, Any] = None, template: Union["DataArray", "Dataset"] = None, ) -> "T_DSorDA": """ Apply a function to each block of this Dataset. .. warning:: This method is experimental and its signature may change. Parameters ---------- func : callable User-provided function that accepts a Dataset as its first parameter. The function will receive a subset or 'block' of this Dataset (see below), corresponding to one chunk along each chunked dimension. ``func`` will be executed as ``func(subset_dataset, *subset_args, **kwargs)``. This function must return either a single DataArray or a single Dataset. This function cannot add a new chunked dimension. args : sequence Passed to func after unpacking and subsetting any xarray objects by blocks. xarray objects in args must be aligned with obj, otherwise an error is raised. kwargs : mapping Passed verbatim to func after unpacking. xarray objects, if any, will not be subset to blocks. Passing dask collections in kwargs is not allowed. template : DataArray or Dataset, optional xarray object representing the final result after compute is called. If not provided, the function will be first run on mocked-up data, that looks like this object but has sizes 0, to determine properties of the returned object such as dtype, variable names, attributes, new dimensions and new indexes (if any). ``template`` must be provided if the function changes the size of existing dimensions. When provided, ``attrs`` on variables in `template` are copied over to the result. Any ``attrs`` set by ``func`` will be ignored. Returns ------- A single DataArray or Dataset with dask backend, reassembled from the outputs of the function. Notes ----- This function is designed for when ``func`` needs to manipulate a whole xarray object subset to each block. In the more common case where ``func`` can work on numpy arrays, it is recommended to use ``apply_ufunc``. If none of the variables in this object is backed by dask arrays, calling this function is equivalent to calling ``func(obj, *args, **kwargs)``. See Also -------- dask.array.map_blocks, xarray.apply_ufunc, xarray.Dataset.map_blocks, xarray.DataArray.map_blocks Examples -------- Calculate an anomaly from climatology using ``.groupby()``. Using ``xr.map_blocks()`` allows for parallel operations with knowledge of ``xarray``, its indices, and its methods like ``.groupby()``. >>> def calculate_anomaly(da, groupby_type="time.month"): ... gb = da.groupby(groupby_type) ... clim = gb.mean(dim="time") ... return gb - clim ... >>> time = xr.cftime_range("1990-01", "1992-01", freq="M") >>> month = xr.DataArray(time.month, coords={"time": time}, dims=["time"]) >>> np.random.seed(123) >>> array = xr.DataArray( ... np.random.rand(len(time)), ... dims=["time"], ... coords={"time": time, "month": month}, ... ).chunk() >>> ds = xr.Dataset({"a": array}) >>> ds.map_blocks(calculate_anomaly, template=ds).compute() Dimensions: (time: 24) Coordinates: * time (time) object 1990-01-31 00:00:00 ... 1991-12-31 00:00:00 month (time) int64 1 2 3 4 5 6 7 8 9 10 11 12 1 2 3 4 5 6 7 8 9 10 11 12 Data variables: a (time) float64 0.1289 0.1132 -0.0856 ... 0.2287 0.1906 -0.05901 Note that one must explicitly use ``args=[]`` and ``kwargs={}`` to pass arguments to the function being applied in ``xr.map_blocks()``: >>> ds.map_blocks( ... calculate_anomaly, ... kwargs={"groupby_type": "time.year"}, ... template=ds, ... ) Dimensions: (time: 24) Coordinates: * time (time) object 1990-01-31 00:00:00 ... 1991-12-31 00:00:00 month (time) int64 dask.array Data variables: a (time) float64 dask.array """ from .parallel import map_blocks return map_blocks(func, self, args, kwargs, template) def polyfit( self, dim: Hashable, deg: int, skipna: bool = None, rcond: float = None, w: Union[Hashable, Any] = None, full: bool = False, cov: Union[bool, str] = False, ): """ Least squares polynomial fit. This replicates the behaviour of `numpy.polyfit` but differs by skipping invalid values when `skipna = True`. Parameters ---------- dim : hashable Coordinate along which to fit the polynomials. deg : int Degree of the fitting polynomial. skipna : bool, optional If True, removes all invalid values before fitting each 1D slices of the array. Default is True if data is stored in a dask.array or if there is any invalid values, False otherwise. rcond : float, optional Relative condition number to the fit. w : hashable or Any, optional Weights to apply to the y-coordinate of the sample points. Can be an array-like object or the name of a coordinate in the dataset. full : bool, optional Whether to return the residuals, matrix rank and singular values in addition to the coefficients. cov : bool or str, optional Whether to return to the covariance matrix in addition to the coefficients. The matrix is not scaled if `cov='unscaled'`. Returns ------- polyfit_results : Dataset A single dataset which contains (for each "var" in the input dataset): [var]_polyfit_coefficients The coefficients of the best fit for each variable in this dataset. [var]_polyfit_residuals The residuals of the least-square computation for each variable (only included if `full=True`) When the matrix rank is deficient, np.nan is returned. [dim]_matrix_rank The effective rank of the scaled Vandermonde coefficient matrix (only included if `full=True`) The rank is computed ignoring the NaN values that might be skipped. [dim]_singular_values The singular values of the scaled Vandermonde coefficient matrix (only included if `full=True`) [var]_polyfit_covariance The covariance matrix of the polynomial coefficient estimates (only included if `full=False` and `cov=True`) Warns ----- RankWarning The rank of the coefficient matrix in the least-squares fit is deficient. The warning is not raised with in-memory (not dask) data and `full=True`. See also -------- numpy.polyfit """ variables = {} skipna_da = skipna x = get_clean_interp_index(self, dim, strict=False) xname = "{}_".format(self[dim].name) order = int(deg) + 1 lhs = np.vander(x, order) if rcond is None: rcond = x.shape[0] * np.core.finfo(x.dtype).eps # Weights: if w is not None: if isinstance(w, Hashable): w = self.coords[w] w = np.asarray(w) if w.ndim != 1: raise TypeError("Expected a 1-d array for weights.") if w.shape[0] != lhs.shape[0]: raise TypeError("Expected w and {} to have the same length".format(dim)) lhs *= w[:, np.newaxis] # Scaling scale = np.sqrt((lhs * lhs).sum(axis=0)) lhs /= scale degree_dim = utils.get_temp_dimname(self.dims, "degree") rank = np.linalg.matrix_rank(lhs) if full: rank = xr.DataArray(rank, name=xname + "matrix_rank") variables[rank.name] = rank sing = np.linalg.svd(lhs, compute_uv=False) sing = xr.DataArray( sing, dims=(degree_dim,), coords={degree_dim: np.arange(rank - 1, -1, -1)}, name=xname + "singular_values", ) variables[sing.name] = sing for name, da in self.data_vars.items(): if dim not in da.dims: continue if is_duck_dask_array(da.data) and ( rank != order or full or skipna is None ): # Current algorithm with dask and skipna=False neither supports # deficient ranks nor does it output the "full" info (issue dask/dask#6516) skipna_da = True elif skipna is None: skipna_da = np.any(da.isnull()) dims_to_stack = [dimname for dimname in da.dims if dimname != dim] stacked_coords: Dict[Hashable, DataArray] = {} if dims_to_stack: stacked_dim = utils.get_temp_dimname(dims_to_stack, "stacked") rhs = da.transpose(dim, *dims_to_stack).stack( {stacked_dim: dims_to_stack} ) stacked_coords = {stacked_dim: rhs[stacked_dim]} scale_da = scale[:, np.newaxis] else: rhs = da scale_da = scale if w is not None: rhs *= w[:, np.newaxis] with warnings.catch_warnings(): if full: # Copy np.polyfit behavior warnings.simplefilter("ignore", np.RankWarning) else: # Raise only once per variable warnings.simplefilter("once", np.RankWarning) coeffs, residuals = duck_array_ops.least_squares( lhs, rhs.data, rcond=rcond, skipna=skipna_da ) if isinstance(name, str): name = "{}_".format(name) else: # Thus a ReprObject => polyfit was called on a DataArray name = "" coeffs = xr.DataArray( coeffs / scale_da, dims=[degree_dim] + list(stacked_coords.keys()), coords={degree_dim: np.arange(order)[::-1], **stacked_coords}, name=name + "polyfit_coefficients", ) if dims_to_stack: coeffs = coeffs.unstack(stacked_dim) variables[coeffs.name] = coeffs if full or (cov is True): residuals = xr.DataArray( residuals if dims_to_stack else residuals.squeeze(), dims=list(stacked_coords.keys()), coords=stacked_coords, name=name + "polyfit_residuals", ) if dims_to_stack: residuals = residuals.unstack(stacked_dim) variables[residuals.name] = residuals if cov: Vbase = np.linalg.inv(np.dot(lhs.T, lhs)) Vbase /= np.outer(scale, scale) if cov == "unscaled": fac = 1 else: if x.shape[0] <= order: raise ValueError( "The number of data points must exceed order to scale the covariance matrix." ) fac = residuals / (x.shape[0] - order) covariance = xr.DataArray(Vbase, dims=("cov_i", "cov_j")) * fac variables[name + "polyfit_covariance"] = covariance return Dataset(data_vars=variables, attrs=self.attrs.copy()) def pad( self, pad_width: Mapping[Hashable, Union[int, Tuple[int, int]]] = None, mode: str = "constant", stat_length: Union[ int, Tuple[int, int], Mapping[Hashable, Tuple[int, int]] ] = None, constant_values: Union[ int, Tuple[int, int], Mapping[Hashable, Tuple[int, int]] ] = None, end_values: Union[ int, Tuple[int, int], Mapping[Hashable, Tuple[int, int]] ] = None, reflect_type: str = None, **pad_width_kwargs: Any, ) -> "Dataset": """Pad this dataset along one or more dimensions. .. warning:: This function is experimental and its behaviour is likely to change especially regarding padding of dimension coordinates (or IndexVariables). When using one of the modes ("edge", "reflect", "symmetric", "wrap"), coordinates will be padded with the same mode, otherwise coordinates are padded using the "constant" mode with fill_value dtypes.NA. Parameters ---------- pad_width : mapping of hashable to tuple of int Mapping with the form of {dim: (pad_before, pad_after)} describing the number of values padded along each dimension. {dim: pad} is a shortcut for pad_before = pad_after = pad mode : str, default: "constant" One of the following string values (taken from numpy docs). 'constant' (default) Pads with a constant value. 'edge' Pads with the edge values of array. 'linear_ramp' Pads with the linear ramp between end_value and the array edge value. 'maximum' Pads with the maximum value of all or part of the vector along each axis. 'mean' Pads with the mean value of all or part of the vector along each axis. 'median' Pads with the median value of all or part of the vector along each axis. 'minimum' Pads with the minimum value of all or part of the vector along each axis. 'reflect' Pads with the reflection of the vector mirrored on the first and last values of the vector along each axis. 'symmetric' Pads with the reflection of the vector mirrored along the edge of the array. 'wrap' Pads with the wrap of the vector along the axis. The first values are used to pad the end and the end values are used to pad the beginning. stat_length : int, tuple or mapping of hashable to tuple, default: None Used in 'maximum', 'mean', 'median', and 'minimum'. Number of values at edge of each axis used to calculate the statistic value. {dim_1: (before_1, after_1), ... dim_N: (before_N, after_N)} unique statistic lengths along each dimension. ((before, after),) yields same before and after statistic lengths for each dimension. (stat_length,) or int is a shortcut for before = after = statistic length for all axes. Default is ``None``, to use the entire axis. constant_values : scalar, tuple or mapping of hashable to tuple, default: 0 Used in 'constant'. The values to set the padded values for each axis. ``{dim_1: (before_1, after_1), ... dim_N: (before_N, after_N)}`` unique pad constants along each dimension. ``((before, after),)`` yields same before and after constants for each dimension. ``(constant,)`` or ``constant`` is a shortcut for ``before = after = constant`` for all dimensions. Default is 0. end_values : scalar, tuple or mapping of hashable to tuple, default: 0 Used in 'linear_ramp'. The values used for the ending value of the linear_ramp and that will form the edge of the padded array. ``{dim_1: (before_1, after_1), ... dim_N: (before_N, after_N)}`` unique end values along each dimension. ``((before, after),)`` yields same before and after end values for each axis. ``(constant,)`` or ``constant`` is a shortcut for ``before = after = constant`` for all axes. Default is 0. reflect_type : {"even", "odd"}, optional Used in "reflect", and "symmetric". The "even" style is the default with an unaltered reflection around the edge value. For the "odd" style, the extended part of the array is created by subtracting the reflected values from two times the edge value. **pad_width_kwargs The keyword arguments form of ``pad_width``. One of ``pad_width`` or ``pad_width_kwargs`` must be provided. Returns ------- padded : Dataset Dataset with the padded coordinates and data. See also -------- Dataset.shift, Dataset.roll, Dataset.bfill, Dataset.ffill, numpy.pad, dask.array.pad Notes ----- By default when ``mode="constant"`` and ``constant_values=None``, integer types will be promoted to ``float`` and padded with ``np.nan``. To avoid type promotion specify ``constant_values=np.nan`` Examples -------- >>> ds = xr.Dataset({"foo": ("x", range(5))}) >>> ds.pad(x=(1, 2)) Dimensions: (x: 8) Dimensions without coordinates: x Data variables: foo (x) float64 nan 0.0 1.0 2.0 3.0 4.0 nan nan """ pad_width = either_dict_or_kwargs(pad_width, pad_width_kwargs, "pad") if mode in ("edge", "reflect", "symmetric", "wrap"): coord_pad_mode = mode coord_pad_options = { "stat_length": stat_length, "constant_values": constant_values, "end_values": end_values, "reflect_type": reflect_type, } else: coord_pad_mode = "constant" coord_pad_options = {} variables = {} for name, var in self.variables.items(): var_pad_width = {k: v for k, v in pad_width.items() if k in var.dims} if not var_pad_width: variables[name] = var elif name in self.data_vars: variables[name] = var.pad( pad_width=var_pad_width, mode=mode, stat_length=stat_length, constant_values=constant_values, end_values=end_values, reflect_type=reflect_type, ) else: variables[name] = var.pad( pad_width=var_pad_width, mode=coord_pad_mode, **coord_pad_options, # type: ignore ) return self._replace_vars_and_dims(variables) def idxmin( self, dim: Hashable = None, skipna: bool = None, fill_value: Any = dtypes.NA, keep_attrs: bool = None, ) -> "Dataset": """Return the coordinate label of the minimum value along a dimension. Returns a new `Dataset` named after the dimension with the values of the coordinate labels along that dimension corresponding to minimum values along that dimension. In comparison to :py:meth:`~Dataset.argmin`, this returns the coordinate label while :py:meth:`~Dataset.argmin` returns the index. Parameters ---------- dim : str, optional Dimension over which to apply `idxmin`. This is optional for 1D variables, but required for variables with 2 or more dimensions. skipna : bool or None, default: None If True, skip missing values (as marked by NaN). By default, only skips missing values for ``float``, ``complex``, and ``object`` dtypes; other dtypes either do not have a sentinel missing value (``int``) or ``skipna=True`` has not been implemented (``datetime64`` or ``timedelta64``). fill_value : Any, default: NaN Value to be filled in case all of the values along a dimension are null. By default this is NaN. The fill value and result are automatically converted to a compatible dtype if possible. Ignored if ``skipna`` is False. keep_attrs : bool, default: False If True, the attributes (``attrs``) will be copied from the original object to the new one. If False (default), the new object will be returned without attributes. Returns ------- reduced : Dataset New `Dataset` object with `idxmin` applied to its data and the indicated dimension removed. See also -------- DataArray.idxmin, Dataset.idxmax, Dataset.min, Dataset.argmin Examples -------- >>> array1 = xr.DataArray( ... [0, 2, 1, 0, -2], dims="x", coords={"x": ["a", "b", "c", "d", "e"]} ... ) >>> array2 = xr.DataArray( ... [ ... [2.0, 1.0, 2.0, 0.0, -2.0], ... [-4.0, np.NaN, 2.0, np.NaN, -2.0], ... [np.NaN, np.NaN, 1.0, np.NaN, np.NaN], ... ], ... dims=["y", "x"], ... coords={"y": [-1, 0, 1], "x": ["a", "b", "c", "d", "e"]}, ... ) >>> ds = xr.Dataset({"int": array1, "float": array2}) >>> ds.min(dim="x") Dimensions: (y: 3) Coordinates: * y (y) int64 -1 0 1 Data variables: int int64 -2 float (y) float64 -2.0 -4.0 1.0 >>> ds.argmin(dim="x") Dimensions: (y: 3) Coordinates: * y (y) int64 -1 0 1 Data variables: int int64 4 float (y) int64 4 0 2 >>> ds.idxmin(dim="x") Dimensions: (y: 3) Coordinates: * y (y) int64 -1 0 1 Data variables: int "Dataset": """Return the coordinate label of the maximum value along a dimension. Returns a new `Dataset` named after the dimension with the values of the coordinate labels along that dimension corresponding to maximum values along that dimension. In comparison to :py:meth:`~Dataset.argmax`, this returns the coordinate label while :py:meth:`~Dataset.argmax` returns the index. Parameters ---------- dim : str, optional Dimension over which to apply `idxmax`. This is optional for 1D variables, but required for variables with 2 or more dimensions. skipna : bool or None, default: None If True, skip missing values (as marked by NaN). By default, only skips missing values for ``float``, ``complex``, and ``object`` dtypes; other dtypes either do not have a sentinel missing value (``int``) or ``skipna=True`` has not been implemented (``datetime64`` or ``timedelta64``). fill_value : Any, default: NaN Value to be filled in case all of the values along a dimension are null. By default this is NaN. The fill value and result are automatically converted to a compatible dtype if possible. Ignored if ``skipna`` is False. keep_attrs : bool, default: False If True, the attributes (``attrs``) will be copied from the original object to the new one. If False (default), the new object will be returned without attributes. Returns ------- reduced : Dataset New `Dataset` object with `idxmax` applied to its data and the indicated dimension removed. See also -------- DataArray.idxmax, Dataset.idxmin, Dataset.max, Dataset.argmax Examples -------- >>> array1 = xr.DataArray( ... [0, 2, 1, 0, -2], dims="x", coords={"x": ["a", "b", "c", "d", "e"]} ... ) >>> array2 = xr.DataArray( ... [ ... [2.0, 1.0, 2.0, 0.0, -2.0], ... [-4.0, np.NaN, 2.0, np.NaN, -2.0], ... [np.NaN, np.NaN, 1.0, np.NaN, np.NaN], ... ], ... dims=["y", "x"], ... coords={"y": [-1, 0, 1], "x": ["a", "b", "c", "d", "e"]}, ... ) >>> ds = xr.Dataset({"int": array1, "float": array2}) >>> ds.max(dim="x") Dimensions: (y: 3) Coordinates: * y (y) int64 -1 0 1 Data variables: int int64 2 float (y) float64 2.0 2.0 1.0 >>> ds.argmax(dim="x") Dimensions: (y: 3) Coordinates: * y (y) int64 -1 0 1 Data variables: int int64 1 float (y) int64 0 2 2 >>> ds.idxmax(dim="x") Dimensions: (y: 3) Coordinates: * y (y) int64 -1 0 1 Data variables: int