import sys import types import numpy as np import pandas as pd from .. import util from ..dimension import Dimension, asdim, dimension_name from ..element import Element from ..ndmapping import NdMapping, item_check, sorted_context from .grid import GridInterface from .interface import DataError, Interface from .util import dask_array_module, finite_range def is_cupy(array): if 'cupy' not in sys.modules: return False from cupy import ndarray return isinstance(array, ndarray) class XArrayInterface(GridInterface): types = () datatype = 'xarray' @classmethod def loaded(cls): return 'xarray' in sys.modules @classmethod def applies(cls, obj): if not cls.loaded(): return False import xarray as xr return isinstance(obj, (xr.Dataset, xr.DataArray)) @classmethod def dimension_type(cls, dataset, dim): name = dataset.get_dimension(dim, strict=True).name if cls.packed(dataset) and name in dataset.vdims: return dataset.data.dtype.type return dataset.data[name].dtype.type @classmethod def dtype(cls, dataset, dim): name = dataset.get_dimension(dim, strict=True).name if cls.packed(dataset) and name in dataset.vdims: return dataset.data.dtype return dataset.data[name].dtype @classmethod def packed(cls, dataset): import xarray as xr return isinstance(dataset.data, xr.DataArray) @classmethod def shape(cls, dataset, gridded=False): if cls.packed(dataset): array = dataset.data[..., 0] else: array = dataset.data[dataset.vdims[0].name] if not gridded: return (np.prod(array.shape, dtype=np.intp), len(dataset.dimensions())) shape_map = dict(zip(array.dims, array.shape)) return tuple(shape_map.get(kd.name, np.nan) for kd in dataset.kdims[::-1]) @classmethod def init(cls, eltype, data, kdims, vdims): import xarray as xr element_params = eltype.param.objects() kdim_param = element_params['kdims'] vdim_param = element_params['vdims'] def retrieve_unit_and_label(dim): if isinstance(dim, Dimension): return dim dim = asdim(dim) coord = data[dim.name] unit = coord.attrs.get('units') if dim.unit is None else dim.unit if isinstance(unit, tuple): unit = unit[0] if isinstance(coord.attrs.get("long_name"), str): spec = (dim.name, coord.attrs['long_name']) else: spec = (dim.name, dim.label) nodata = coord.attrs.get('NODATA') return dim.clone(spec, unit=unit, nodata=nodata) packed = False if isinstance(data, xr.DataArray): kdim_len = len(kdim_param.default) if kdims is None else len(kdims) vdim_len = len(vdim_param.default) if vdims is None else len(vdims) if vdim_len > 1 and kdim_len == len(data.dims)-1 and data.shape[-1] == vdim_len: packed = True elif vdims: vdim = vdims[0] elif data.name: vdim = Dimension(data.name) vdim.unit = data.attrs.get('units') vdim.nodata = data.attrs.get('NODATA') label = data.attrs.get('long_name') if isinstance(label, str): vdim.label = label elif len(vdim_param.default) == 1: vdim = asdim(vdim_param.default[0]) if vdim.name in data.dims: raise DataError("xarray DataArray does not define a name, " "and the default of '%s' clashes with a " "coordinate dimension. Give the DataArray " "a name or supply an explicit value dimension." % vdim.name, cls) else: raise DataError("xarray DataArray does not define a name " "and %s does not define a default value " "dimension. Give the DataArray a name or " "supply an explicit vdim." % eltype.__name__, cls) if not packed: if vdim in data.dims: data = data.to_dataset(vdim.name) vdims = [asdim(vd) for vd in data.data_vars] else: vdims = [vdim] data = data.to_dataset(name=vdim.name) if not isinstance(data, (xr.Dataset, xr.DataArray)): if kdims is None: kdims = kdim_param.default if vdims is None: vdims = vdim_param.default kdims = [asdim(kd) for kd in kdims] vdims = [asdim(vd) for vd in vdims] if isinstance(data, np.ndarray) and data.ndim == 2 and data.shape[1] == len(kdims+vdims): data = tuple(data) ndims = len(kdims) if isinstance(data, tuple): dimensions = [d.name for d in kdims+vdims] if (len(data) != len(dimensions) and len(data) == (ndims+1) and len(data[-1].shape) == (ndims+1)): value_array = data[-1] data = {d: v for d, v in zip(dimensions, data[:-1])} packed = True else: data = {d: v for d, v in zip(dimensions, data)} elif isinstance(data, (list, np.ndarray)) and len(data) == 0: dimensions = [d.name for d in kdims + vdims] data = {d: np.array([]) for d in dimensions[:ndims]} data.update({d: np.empty((0,) * ndims) for d in dimensions[ndims:]}) if not isinstance(data, dict): raise TypeError('XArrayInterface could not interpret data type') data = {d: np.asarray(values) if d in kdims else values for d, values in data.items()} coord_dims = [data[kd.name].ndim for kd in kdims] dims = tuple('dim_%d' % i for i in range(max(coord_dims)))[::-1] coords = {} for kd in kdims: coord_vals = data[kd.name] if coord_vals.ndim > 1: coord = (dims[:coord_vals.ndim], coord_vals) else: coord = coord_vals coords[kd.name] = coord xr_kwargs = {'dims': dims if max(coord_dims) > 1 else list(coords)[::-1]} if packed: xr_kwargs['dims'] = list(coords)[::-1] + ['band'] coords['band'] = list(range(len(vdims))) data = xr.DataArray(value_array, coords=coords, **xr_kwargs) else: arrays = {} for vdim in vdims: arr = data[vdim.name] if not isinstance(arr, xr.DataArray): arr = xr.DataArray(arr, coords=coords, **xr_kwargs) arrays[vdim.name] = arr data = xr.Dataset(arrays) else: # Started to warn in xarray 2023.12.0: # The return type of `Dataset.dims` will be changed to return a # set of dimension names in future, in order to be more consistent # with `DataArray.dims`. To access a mapping from dimension names to # lengths, please use `Dataset.sizes`. data_info = data.sizes if hasattr(data, "sizes") else data.dims if not data.coords: data = data.assign_coords(**{k: range(v) for k, v in data_info.items()}) if vdims is None: vdims = list(data.data_vars) if kdims is None: xrdims = list(data_info) xrcoords = list(data.coords) kdims = [name for name in data.indexes.keys() if isinstance(data[name].data, np.ndarray)] kdims = sorted(kdims, key=lambda x: (xrcoords.index(x) if x in xrcoords else float('inf'), x)) if packed: kdims = kdims[:-1] elif set(xrdims) != set(kdims): virtual_dims = [xd for xd in xrdims if xd not in kdims] for c in data.coords: if c not in kdims and set(data[c].dims) == set(virtual_dims): kdims.append(c) kdims = [retrieve_unit_and_label(kd) for kd in kdims] vdims = [retrieve_unit_and_label(vd) for vd in vdims] not_found = [] for d in kdims: if not any(d.name == k or (isinstance(v, xr.DataArray) and d.name in v.dims) for k, v in data.coords.items()): not_found.append(d) if not isinstance(data, (xr.Dataset, xr.DataArray)): raise TypeError('Data must be be an xarray Dataset type.') elif not_found: raise DataError("xarray Dataset must define coordinates " "for all defined kdims, %s coordinates not found." % not_found, cls) for vdim in vdims: if packed: continue da = data[vdim.name] # Do not enforce validation for irregular arrays since they # not need to be canonicalized if any(len(da.coords[c].shape) > 1 for c in da.coords): continue undeclared = [] for c in da.coords: if c in kdims or len(da[c].shape) != 1 or da[c].shape[0] <= 1: # Skip if coord is declared, represents irregular coordinates or is constant continue elif all(d in kdims for d in da[c].dims): continue # Skip if coord is alias for another dimension elif any(all(d in da[kd.name].dims for d in da[c].dims) for kd in kdims): # Skip if all the dims on the coord are present on another coord continue undeclared.append(c) if undeclared and eltype.param.kdims.bounds[1] not in (0, None): raise DataError( f'The coordinates on the {vdim.name!r} DataArray do not match the ' 'provided key dimensions (kdims). The following coords ' f'were left unspecified: {undeclared!r}. If you are requesting a ' 'lower dimensional view such as a histogram cast ' 'the xarray to a columnar format using the .to_dataframe ' 'or .to_dask_dataframe methods before providing it to ' 'HoloViews.') return data, {'kdims': kdims, 'vdims': vdims}, {} @classmethod def validate(cls, dataset, vdims=True): import xarray as xr if isinstance(dataset.data, xr.Dataset): Interface.validate(dataset, vdims) else: not_found = [kd.name for kd in dataset.kdims if kd.name not in dataset.data.coords] if not_found: raise DataError("Supplied data does not contain specified " "dimensions, the following dimensions were " "not found: %s" % repr(not_found), cls) # Check whether irregular (i.e. multi-dimensional) coordinate # array dimensionality matches irregular = [] for kd in dataset.kdims: if cls.irregular(dataset, kd): irregular.append((kd, dataset.data[kd.name].dims)) if irregular: nonmatching = [f'{kd}: {dims}' for kd, dims in irregular[1:] if set(dims) != set(irregular[0][1])] if nonmatching: nonmatching = ['{}: {}'.format(*irregular[0])] + nonmatching raise DataError("The dimensions of coordinate arrays " "on irregular data must match. The " "following kdims were found to have " "non-matching array dimensions:\n\n%s" % ('\n'.join(nonmatching)), cls) @classmethod def compute(cls, dataset): return dataset.clone(dataset.data.compute()) @classmethod def persist(cls, dataset): return dataset.clone(dataset.data.persist()) @classmethod def range(cls, dataset, dimension): dimension = dataset.get_dimension(dimension, strict=True) dim = dimension.name edges = dataset._binned and dimension in dataset.kdims if edges: data = cls.coords(dataset, dim, edges=True) else: if cls.packed(dataset) and dim in dataset.vdims: data = dataset.data.values[..., dataset.vdims.index(dim)] else: data = dataset.data[dim] if dimension.nodata is not None: data = cls.replace_value(data, dimension.nodata) if not len(data): dmin, dmax = np.nan, np.nan elif data.dtype.kind == 'M' or not edges: dmin, dmax = data.min(), data.max() if not edges: dmin, dmax = dmin.data, dmax.data else: dmin, dmax = np.nanmin(data), np.nanmax(data) da = dask_array_module() if da and isinstance(dmin, da.Array): dmin, dmax = da.compute(dmin, dmax) if isinstance(dmin, np.ndarray) and dmin.shape == (): dmin = dmin[()] if isinstance(dmax, np.ndarray) and dmax.shape == (): dmax = dmax[()] dmin = dmin if np.isscalar(dmin) or isinstance(dmin, util.datetime_types) else dmin.item() dmax = dmax if np.isscalar(dmax) or isinstance(dmax, util.datetime_types) else dmax.item() return finite_range(data, dmin, dmax) @classmethod def groupby(cls, dataset, dimensions, container_type, group_type, **kwargs): index_dims = [dataset.get_dimension(d, strict=True) for d in dimensions] element_dims = [kdim for kdim in dataset.kdims if kdim not in index_dims] invalid = [d for d in index_dims if dataset.data[d.name].ndim > 1] if invalid: if len(invalid) == 1: invalid = f"'{invalid[0]}'" raise ValueError("Cannot groupby irregularly sampled dimension(s) %s." % invalid) group_kwargs = {} if group_type != 'raw' and issubclass(group_type, Element): group_kwargs = dict(util.get_param_values(dataset), kdims=element_dims) group_kwargs.update(kwargs) drop_dim = any(d not in group_kwargs['kdims'] for d in element_dims) group_by = [d.name for d in index_dims] data = [] if len(dimensions) == 1: for k, v in dataset.data.groupby(index_dims[0].name, squeeze=False): if drop_dim: v = v.to_dataframe().reset_index() data.append((k, group_type(v, **group_kwargs))) else: unique_iters = [cls.values(dataset, d, False) for d in group_by] indexes = zip(*util.cartesian_product(unique_iters)) for k in indexes: sel = dataset.data.sel(**dict(zip(group_by, k))) if drop_dim: sel = sel.to_dataframe().reset_index() data.append((k, group_type(sel, **group_kwargs))) if issubclass(container_type, NdMapping): with item_check(False), sorted_context(False): return container_type(data, kdims=index_dims) else: return container_type(data) @classmethod def coords(cls, dataset, dimension, ordered=False, expanded=False, edges=False): import xarray as xr dim = dataset.get_dimension(dimension) dim = dimension if dim is None else dim.name irregular = cls.irregular(dataset, dim) if irregular or expanded: if irregular: data = dataset.data[dim] else: data = util.expand_grid_coords(dataset, dim) if edges: data = cls._infer_interval_breaks(data, axis=1) data = cls._infer_interval_breaks(data, axis=0) return data.values if isinstance(data, xr.DataArray) else data data = np.atleast_1d(dataset.data[dim].data) if ordered and data.shape and np.all(data[1:] < data[:-1]): data = data[::-1] shape = cls.shape(dataset, True) if dim in dataset.kdims: idx = dataset.get_dimension_index(dim) isedges = (len(shape) == dataset.ndims and len(data) == (shape[dataset.ndims-idx-1]+1)) else: isedges = False if edges and not isedges: data = cls._infer_interval_breaks(data) elif not edges and isedges: data = np.convolve(data, [0.5, 0.5], 'valid') return data.values if isinstance(data, xr.DataArray) else data @classmethod def values(cls, dataset, dim, expanded=True, flat=True, compute=True, keep_index=False): dim = dataset.get_dimension(dim, strict=True) packed = cls.packed(dataset) and dim in dataset.vdims if packed: data = dataset.data.data[..., dataset.vdims.index(dim)] else: data = dataset.data[dim.name] if not keep_index: data = data.data irregular = cls.irregular(dataset, dim) if dim in dataset.kdims else False irregular_kdims = [d for d in dataset.kdims if cls.irregular(dataset, d)] if irregular_kdims: virtual_coords = list(dataset.data[irregular_kdims[0].name].coords.dims) else: virtual_coords = [] if dim in dataset.vdims or irregular: if packed: data_coords = list(dataset.data.dims)[:-1] else: data_coords = list(dataset.data[dim.name].dims) da = dask_array_module() if compute and da and isinstance(data, da.Array): data = data.compute() if is_cupy(data): import cupy data = cupy.asnumpy(data) if not keep_index: data = cls.canonicalize(dataset, data, data_coords=data_coords, virtual_coords=virtual_coords) return data.T.flatten() if flat and not keep_index else data elif expanded: data = cls.coords(dataset, dim.name, expanded=True) return data.T.flatten() if flat else data else: if keep_index: return dataset.data[dim.name] return cls.coords(dataset, dim.name, ordered=True) @classmethod def aggregate(cls, dataset, dimensions, function, **kwargs): reduce_dims = [d.name for d in dataset.kdims if d not in dimensions] return dataset.data.reduce(function, dim=reduce_dims, **kwargs), [] @classmethod def unpack_scalar(cls, dataset, data): """ Given a dataset object and data in the appropriate format for the interface, return a simple scalar. """ if not cls.packed(dataset) and len(data.data_vars) == 1: array = data[dataset.vdims[0].name].squeeze() if len(array.shape) == 0: return array.item() return data @classmethod def ndloc(cls, dataset, indices): kdims = [d for d in dataset.kdims[::-1]] adjusted_indices = [] slice_dims = [] for kd, ind in zip(kdims, indices): if cls.irregular(dataset, kd): coords = [c for c in dataset.data.coords if c not in dataset.data.dims] dim = dataset.data[kd.name].dims[coords.index(kd.name)] shape = dataset.data[kd.name].shape[coords.index(kd.name)] coords = np.arange(shape) else: coords = cls.coords(dataset, kd, False) dim = kd.name slice_dims.append(dim) ncoords = len(coords) if np.all(coords[1:] < coords[:-1]): if np.isscalar(ind): ind = ncoords-ind-1 elif isinstance(ind, slice): start = None if ind.stop is None else ncoords-ind.stop stop = None if ind.start is None else ncoords-ind.start ind = slice(start, stop, ind.step) elif isinstance(ind, np.ndarray) and ind.dtype.kind == 'b': ind = ind[::-1] elif isinstance(ind, (np.ndarray, list)): ind = [ncoords-i-1 for i in ind] if isinstance(ind, list): ind = np.array(ind) if isinstance(ind, np.ndarray) and ind.dtype.kind == 'b': ind = np.where(ind)[0] adjusted_indices.append(ind) isel = dict(zip(slice_dims, adjusted_indices)) all_scalar = all(map(np.isscalar, indices)) if all_scalar and len(indices) == len(kdims) and len(dataset.vdims) == 1: return dataset.data[dataset.vdims[0].name].isel(**isel).values.item() # Detect if the indexing is selecting samples or slicing the array sampled = (all(isinstance(ind, np.ndarray) and ind.dtype.kind != 'b' for ind in adjusted_indices) and len(indices) == len(kdims)) if sampled or (all_scalar and len(indices) == len(kdims)): import xarray as xr if cls.packed(dataset): selected = dataset.data.isel({k: xr.DataArray(v) for k, v in isel.items()}) df = selected.to_dataframe('vdims')[['vdims']].T vdims = [vd.name for vd in dataset.vdims] return df.rename(columns={i: d for i, d in enumerate(vdims)})[vdims] if all_scalar: isel = {k: [v] for k, v in isel.items()} selected = dataset.data.isel({k: xr.DataArray(v) for k, v in isel.items()}) return selected.to_dataframe().reset_index() else: return dataset.data.isel(**isel) @classmethod def concat_dim(cls, datasets, dim, vdims): import xarray as xr return xr.concat([ds.assign_coords(**{dim.name: c}) for c, ds in datasets.items()], dim=dim.name) @classmethod def redim(cls, dataset, dimensions): renames = {k: v.name for k, v in dimensions.items()} return dataset.data.rename(renames) @classmethod def reindex(cls, dataset, kdims=None, vdims=None): dropped_kdims = [kd for kd in dataset.kdims if kd not in kdims] constant = {} for kd in dropped_kdims: vals = cls.values(dataset, kd.name, expanded=False) if len(vals) == 1: constant[kd.name] = vals[0] if len(constant) == len(dropped_kdims): dropped = dataset.data.sel(**{k: v for k, v in constant.items() if k in dataset.data.dims}) if vdims and cls.packed(dataset): return dropped.isel(**{dataset.data.dims[-1]: [dataset.vdims.index(vd) for vd in vdims]}) return dropped elif dropped_kdims: return tuple(dataset.columns(kdims+vdims).values()) return dataset.data @classmethod def sort(cls, dataset, by=None, reverse=False): if by is None: by = [] return dataset @classmethod def mask(cls, dataset, mask, mask_val=np.nan): packed = cls.packed(dataset) masked = dataset.data.copy() if packed: data_coords = list(dataset.data.dims)[:-1] mask = cls.canonicalize(dataset, mask, data_coords) try: masked.values[mask] = mask_val except ValueError: masked = masked.astype('float') masked.values[mask] = mask_val else: orig_mask = mask for vd in dataset.vdims: dims = list(dataset.data[vd.name].dims) if any(cls.irregular(dataset, kd) for kd in dataset.kdims): data_coords = dims else: data_coords = [kd.name for kd in dataset.kdims][::-1] mask = cls.canonicalize(dataset, orig_mask, data_coords) if data_coords != dims: inds = [dims.index(d) for d in data_coords] mask = mask.transpose(inds) masked[vd.name] = marr = masked[vd.name].astype('float') marr.values[mask] = mask_val return masked @classmethod def select(cls, dataset, selection_mask=None, **selection): if selection_mask is not None: return dataset.data.where(selection_mask, drop=True) validated = {} for k, v in selection.items(): dim = dataset.get_dimension(k, strict=True) if cls.irregular(dataset, dim): return GridInterface.select(dataset, selection_mask, **selection) dim = dim.name if isinstance(v, slice): v = (v.start, v.stop) if isinstance(v, set): validated[dim] = list(v) elif isinstance(v, tuple): dim_vals = dataset.data[k].values upper = None if v[1] is None else v[1]-sys.float_info.epsilon*10 v = v[0], upper if dim_vals.dtype.kind not in 'OSU' and np.all(dim_vals[1:] < dim_vals[:-1]): # If coordinates are inverted invert slice v = v[::-1] validated[dim] = slice(*v) elif isinstance(v, types.FunctionType): validated[dim] = v(dataset[k]) else: validated[dim] = v data = dataset.data.sel(**validated) # Restore constant dimensions indexed = cls.indexed(dataset, selection) dropped = dict((d.name, np.atleast_1d(data[d.name])) for d in dataset.kdims if not data[d.name].data.shape) if dropped and not indexed: data = data.expand_dims(dropped) # see https://github.com/pydata/xarray/issues/2891 # since we only expanded on dimensions of size 1 # we can monkeypatch the dataarray back to writeable. for d in data.values(): if hasattr(d.data, 'flags'): d.data.flags.writeable = True da = dask_array_module() if (indexed and len(data.data_vars) == 1 and len(data[dataset.vdims[0].name].shape) == 0): value = data[dataset.vdims[0].name] if da and isinstance(value.data, da.Array): value = value.compute() return value.item() elif indexed: values = [] for vd in dataset.vdims: value = data[vd.name] if da and isinstance(value.data, da.Array): value = value.compute() values.append(value.item()) return np.array(values) return data @classmethod def length(cls, dataset): return np.prod([len(dataset.data[d.name]) for d in dataset.kdims], dtype=np.intp) @classmethod def dframe(cls, dataset, dimensions): import xarray as xr if cls.packed(dataset): bands = {vd.name: dataset.data[..., i].drop_vars('band') for i, vd in enumerate(dataset.vdims)} data = xr.Dataset(bands) else: data = dataset.data data = data.to_dataframe().reset_index() if dimensions: return data[dimensions] return data @classmethod def sample(cls, dataset, samples=None): if samples is None: samples = [] names = [kd.name for kd in dataset.kdims] samples = [dataset.data.sel(**{k: [v] for k, v in zip(names, s)}).to_dataframe().reset_index() for s in samples] return pd.concat(samples) @classmethod def add_dimension(cls, dataset, dimension, dim_pos, values, vdim): import xarray as xr if not vdim: raise Exception("Cannot add key dimension to a dense representation.") dim = dimension_name(dimension) coords = {d.name: cls.coords(dataset, d.name) for d in dataset.kdims} arr = xr.DataArray(values, coords=coords, name=dim, dims=tuple(d.name for d in dataset.kdims[::-1])) return dataset.data.assign(**{dim: arr}) @classmethod def assign(cls, dataset, new_data): import xarray as xr data = dataset.data prev_coords = set.intersection(*[ set(var.coords) for var in data.data_vars.values() ]) coords = {} for k, v in new_data.items(): if k not in dataset.kdims: continue elif isinstance(v, xr.DataArray): coords[k] = v.rename(**({v.name: k} if v.name != k else {})) continue coord_vals = cls.coords(dataset, k) if not coord_vals.ndim > 1 and np.all(coord_vals[1:] < coord_vals[:-1]): v = v[::-1] coords[k] = (k, v) if coords: data = data.assign_coords(**coords) dims = tuple(kd.name for kd in dataset.kdims[::-1]) vars = {} for k, v in new_data.items(): if k in dataset.kdims: continue if isinstance(v, xr.DataArray): vars[k] = v else: vars[k] = (dims, cls.canonicalize(dataset, v, data_coords=dims)) if len(vars) == 1 and list(vars) == [vd.name for vd in dataset.vdims]: data = vars[dataset.vdims[0].name] used_coords = set(data.coords) else: if vars: data = data.assign(vars) used_coords = set.intersection(*[set(var.coords) for var in data.data_vars.values()]) drop_coords = set.symmetric_difference(used_coords, prev_coords) return data.drop_vars([c for c in drop_coords if c in data.coords]), list(drop_coords) Interface.register(XArrayInterface)