import operator import warnings from collections.abc import Iterable from functools import reduce from typing import Any, NamedTuple import numba import numpy as np from .._sparse_array import SparseArray from .._utils import ( can_store, check_consistent_fill_value, check_zero_fill_value, is_unsigned_dtype, isscalar, normalize_axis, ) def asCOO(x, name="asCOO", check=True): """ Convert the input to :obj:`COO`. Passes through :obj:`COO` objects as-is. Parameters ---------- x : Union[SparseArray, scipy.sparse.spmatrix, numpy.ndarray] The input array to convert. name : str, optional The name of the operation to use in the exception. check : bool, optional Whether to check for a dense input. Returns ------- COO The converted :obj:`COO` array. Raises ------ ValueError If ``check`` is true and a dense input is supplied. """ from .._common import _is_sparse from .core import COO if check and not _is_sparse(x): raise ValueError(f"Performing this operation would produce a dense result: {name}") if not isinstance(x, COO): x = COO(x) return x def linear_loc(coords, shape): if shape == () and len(coords) == 0: # `np.ravel_multi_index` is not aware of arrays, so cannot produce a # sensible result here (https://github.com/numpy/numpy/issues/15690). # Since `coords` is an array and not a sequence, we know the correct # dimensions. return np.zeros(coords.shape[1:], dtype=np.intp) return np.ravel_multi_index(coords, shape) def kron(a, b): """Kronecker product of 2 sparse arrays. Parameters ---------- a, b : SparseArray, scipy.sparse.spmatrix, or np.ndarray The arrays over which to compute the Kronecker product. Returns ------- res : COO The kronecker product Raises ------ ValueError If all arguments are dense or arguments have nonzero fill-values. Examples -------- >>> from sparse import eye >>> a = eye(3, dtype="i8") >>> b = np.array([1, 2, 3], dtype="i8") >>> res = kron(a, b) >>> res.todense() # doctest: +SKIP array([[1, 2, 3, 0, 0, 0, 0, 0, 0], [0, 0, 0, 1, 2, 3, 0, 0, 0], [0, 0, 0, 0, 0, 0, 1, 2, 3]], dtype=int64) """ from .._common import _is_sparse from .._umath import _cartesian_product from .core import COO check_zero_fill_value(a, b) a_sparse = _is_sparse(a) b_sparse = _is_sparse(b) a_ndim = np.ndim(a) b_ndim = np.ndim(b) if not (a_sparse or b_sparse): raise ValueError("Performing this operation would produce a dense result: kron") if a_ndim == 0 or b_ndim == 0: return a * b a = asCOO(a, check=False) b = asCOO(b, check=False) # Match dimensions max_dim = max(a.ndim, b.ndim) a = a.reshape((1,) * (max_dim - a.ndim) + a.shape) b = b.reshape((1,) * (max_dim - b.ndim) + b.shape) a_idx, b_idx = _cartesian_product(np.arange(a.nnz), np.arange(b.nnz)) a_expanded_coords = a.coords[:, a_idx] b_expanded_coords = b.coords[:, b_idx] o_coords = a_expanded_coords * np.asarray(b.shape)[:, None] + b_expanded_coords o_data = a.data[a_idx] * b.data[b_idx] o_shape = tuple(i * j for i, j in zip(a.shape, b.shape, strict=True)) return COO(o_coords, o_data, shape=o_shape, has_duplicates=False) def concatenate(arrays, axis=0): """ Concatenate the input arrays along the given dimension. Parameters ---------- arrays : Iterable[SparseArray] The input arrays to concatenate. axis : int, optional The axis along which to concatenate the input arrays. The default is zero. Returns ------- COO The output concatenated array. Raises ------ ValueError If all elements of :code:`arrays` don't have the same fill-value. See Also -------- numpy.concatenate : NumPy equivalent function """ from .core import COO check_consistent_fill_value(arrays) if axis is None: axis = 0 arrays = [x.flatten() for x in arrays] arrays = [x if isinstance(x, COO) else COO(x) for x in arrays] axis = normalize_axis(axis, arrays[0].ndim) assert all(x.shape[ax] == arrays[0].shape[ax] for x in arrays for ax in set(range(arrays[0].ndim)) - {axis}) nnz = 0 dim = sum(x.shape[axis] for x in arrays) shape = list(arrays[0].shape) shape[axis] = dim data = np.concatenate([x.data for x in arrays]) coords = np.concatenate([x.coords for x in arrays], axis=1) if not can_store(coords.dtype, max(shape)): coords = coords.astype(np.min_scalar_type(max(shape))) dim = 0 for x in arrays: if dim: coords[axis, nnz : x.nnz + nnz] += dim dim += x.shape[axis] nnz += x.nnz return COO( coords, data, shape=shape, has_duplicates=False, sorted=(axis == 0), fill_value=arrays[0].fill_value, ) def stack(arrays, axis=0): """ Stack the input arrays along the given dimension. Parameters ---------- arrays : Iterable[SparseArray] The input arrays to stack. axis : int, optional The axis along which to stack the input arrays. Returns ------- COO The output stacked array. Raises ------ ValueError If all elements of :code:`arrays` don't have the same fill-value. See Also -------- numpy.stack : NumPy equivalent function """ from .core import COO check_consistent_fill_value(arrays) assert len({x.shape for x in arrays}) == 1 arrays = [x if isinstance(x, COO) else COO(x) for x in arrays] axis = normalize_axis(axis, arrays[0].ndim + 1) data = np.concatenate([x.data for x in arrays]) coords = np.concatenate([x.coords for x in arrays], axis=1) shape = list(arrays[0].shape) shape.insert(axis, len(arrays)) nnz = 0 new = np.empty(shape=(coords.shape[1],), dtype=np.intp) for dim, x in enumerate(arrays): new[nnz : x.nnz + nnz] = dim nnz += x.nnz coords = [coords[i] for i in range(coords.shape[0])] coords.insert(axis, new) coords = np.stack(coords, axis=0) return COO( coords, data, shape=shape, has_duplicates=False, sorted=(axis == 0), fill_value=arrays[0].fill_value, ) def triu(x, k=0): """ Returns an array with all elements below the k-th diagonal set to zero. Parameters ---------- x : COO The input array. k : int, optional The diagonal below which elements are set to zero. The default is zero, which corresponds to the main diagonal. Returns ------- COO The output upper-triangular matrix. Raises ------ ValueError If :code:`x` doesn't have zero fill-values. See Also -------- numpy.triu : NumPy equivalent function """ from .core import COO check_zero_fill_value(x) if not x.ndim >= 2: raise NotImplementedError("sparse.triu is not implemented for scalars or 1-D arrays.") mask = x.coords[-2] + k <= x.coords[-1] coords = x.coords[:, mask] data = x.data[mask] return COO(coords, data, shape=x.shape, has_duplicates=False, sorted=True) def tril(x, k=0): """ Returns an array with all elements above the k-th diagonal set to zero. Parameters ---------- x : COO The input array. k : int, optional The diagonal above which elements are set to zero. The default is zero, which corresponds to the main diagonal. Returns ------- COO The output lower-triangular matrix. Raises ------ ValueError If :code:`x` doesn't have zero fill-values. See Also -------- numpy.tril : NumPy equivalent function """ from .core import COO check_zero_fill_value(x) if not x.ndim >= 2: raise NotImplementedError("sparse.tril is not implemented for scalars or 1-D arrays.") mask = x.coords[-2] + k >= x.coords[-1] coords = x.coords[:, mask] data = x.data[mask] return COO(coords, data, shape=x.shape, has_duplicates=False, sorted=True) def nansum(x, axis=None, keepdims=False, dtype=None, out=None): """ Performs a ``NaN`` skipping sum operation along the given axes. Uses all axes by default. Parameters ---------- x : SparseArray The array to perform the reduction on. axis : Union[int, Iterable[int]], optional The axes along which to sum. Uses all axes by default. keepdims : bool, optional Whether or not to keep the dimensions of the original array. dtype : numpy.dtype The data type of the output array. Returns ------- COO The reduced output sparse array. See Also -------- :obj:`COO.sum` : Function without ``NaN`` skipping. numpy.nansum : Equivalent Numpy function. """ assert out is None x = asCOO(x, name="nansum") return nanreduce(x, np.add, axis=axis, keepdims=keepdims, dtype=dtype) def nanmean(x, axis=None, keepdims=False, dtype=None, out=None): """ Performs a ``NaN`` skipping mean operation along the given axes. Uses all axes by default. Parameters ---------- x : SparseArray The array to perform the reduction on. axis : Union[int, Iterable[int]], optional The axes along which to compute the mean. Uses all axes by default. keepdims : bool, optional Whether or not to keep the dimensions of the original array. dtype : numpy.dtype The data type of the output array. Returns ------- COO The reduced output sparse array. See Also -------- :obj:`COO.mean` : Function without ``NaN`` skipping. numpy.nanmean : Equivalent Numpy function. """ assert out is None x = asCOO(x, name="nanmean") if not (np.issubdtype(x.dtype, np.floating) or np.issubdtype(x.dtype, np.complexfloating)): return x.mean(axis=axis, keepdims=keepdims, dtype=dtype) mask = np.isnan(x) x2 = where(mask, 0, x) # Count the number non-nan elements along axis nancount = mask.sum(axis=axis, dtype="i8", keepdims=keepdims) if axis is None: axis = tuple(range(x.ndim)) elif not isinstance(axis, tuple): axis = (axis,) den = reduce(operator.mul, (x.shape[i] for i in axis), 1) den -= nancount if (den == 0).any(): warnings.warn("Mean of empty slice", RuntimeWarning, stacklevel=1) num = np.sum(x2, axis=axis, dtype=dtype, keepdims=keepdims) with np.errstate(invalid="ignore", divide="ignore"): if num.ndim: return np.true_divide(num, den, casting="unsafe") return (num / den).astype(dtype if dtype is not None else x.dtype) def nanmax(x, axis=None, keepdims=False, dtype=None, out=None): """ Maximize along the given axes, skipping ``NaN`` values. Uses all axes by default. Parameters ---------- x : SparseArray The array to perform the reduction on. axis : Union[int, Iterable[int]], optional The axes along which to maximize. Uses all axes by default. keepdims : bool, optional Whether or not to keep the dimensions of the original array. dtype : numpy.dtype The data type of the output array. Returns ------- COO The reduced output sparse array. See Also -------- :obj:`COO.max` : Function without ``NaN`` skipping. numpy.nanmax : Equivalent Numpy function. """ assert out is None x = asCOO(x, name="nanmax") ar = x.reduce(np.fmax, axis=axis, keepdims=keepdims, dtype=dtype) if (isscalar(ar) and np.isnan(ar)) or np.isnan(ar.data).any(): warnings.warn("All-NaN slice encountered", RuntimeWarning, stacklevel=1) return ar def nanmin(x, axis=None, keepdims=False, dtype=None, out=None): """ Minimize along the given axes, skipping ``NaN`` values. Uses all axes by default. Parameters ---------- x : SparseArray The array to perform the reduction on. axis : Union[int, Iterable[int]], optional The axes along which to minimize. Uses all axes by default. keepdims : bool, optional Whether or not to keep the dimensions of the original array. dtype : numpy.dtype The data type of the output array. Returns ------- COO The reduced output sparse array. See Also -------- :obj:`COO.min` : Function without ``NaN`` skipping. numpy.nanmin : Equivalent Numpy function. """ assert out is None x = asCOO(x, name="nanmin") ar = x.reduce(np.fmin, axis=axis, keepdims=keepdims, dtype=dtype) if (isscalar(ar) and np.isnan(ar)) or np.isnan(ar.data).any(): warnings.warn("All-NaN slice encountered", RuntimeWarning, stacklevel=1) return ar def nanprod(x, axis=None, keepdims=False, dtype=None, out=None): """ Performs a product operation along the given axes, skipping ``NaN`` values. Uses all axes by default. Parameters ---------- x : SparseArray The array to perform the reduction on. axis : Union[int, Iterable[int]], optional The axes along which to multiply. Uses all axes by default. keepdims : bool, optional Whether or not to keep the dimensions of the original array. dtype : numpy.dtype The data type of the output array. Returns ------- COO The reduced output sparse array. See Also -------- :obj:`COO.prod` : Function without ``NaN`` skipping. numpy.nanprod : Equivalent Numpy function. """ assert out is None x = asCOO(x) return nanreduce(x, np.multiply, axis=axis, keepdims=keepdims, dtype=dtype) def where(condition, x=None, y=None): """ Select values from either ``x`` or ``y`` depending on ``condition``. If ``x`` and ``y`` are not given, returns indices where ``condition`` is nonzero. Performs the equivalent of :obj:`numpy.where`. Parameters ---------- condition : SparseArray The condition based on which to select values from either ``x`` or ``y``. x : SparseArray, optional The array to select values from if ``condition`` is nonzero. y : SparseArray, optional The array to select values from if ``condition`` is zero. Returns ------- COO The output array with selected values if ``x`` and ``y`` are given; else where the array is nonzero. Raises ------ ValueError If the operation would produce a dense result; or exactly one of ``x`` and ``y`` are given. See Also -------- numpy.where : Equivalent Numpy function. """ from .._umath import elemwise x_given = x is not None y_given = y is not None if not (x_given or y_given): check_zero_fill_value(condition) condition = asCOO(condition, name=str(np.where)) return tuple(condition.coords) if x_given != y_given: raise ValueError("either both or neither of x and y should be given") return elemwise(np.where, condition, x, y) def argwhere(a): """ Find the indices of array elements that are non-zero, grouped by element. Parameters ---------- a : array_like Input data. Returns ------- index_array : numpy.ndarray See Also -------- :obj:`where`, :obj:`COO.nonzero` Examples -------- >>> import sparse >>> x = sparse.COO(np.arange(6).reshape((2, 3))) >>> sparse.argwhere(x > 1) array([[0, 2], [1, 0], [1, 1], [1, 2]]) """ return np.transpose(a.nonzero()) def argmax(x, /, *, axis=None, keepdims=False): """ Returns the indices of the maximum values along a specified axis. When the maximum value occurs multiple times, only the indices corresponding to the first occurrence are returned. Parameters ---------- x : SparseArray Input array. The fill value must be ``0.0`` and all non-zero values must be greater than ``0.0``. axis : int, optional Axis along which to search. If ``None``, the function must return the index of the maximum value of the flattened array. Default: ``None``. keepdims : bool, optional If ``True``, the reduced axes (dimensions) must be included in the result as singleton dimensions, and, accordingly, the result must be compatible with the input array. Otherwise, if ``False``, the reduced axes (dimensions) must not be included in the result. Default: ``False``. Returns ------- out : numpy.ndarray If ``axis`` is ``None``, a zero-dimensional array containing the index of the first occurrence of the maximum value. Otherwise, a non-zero-dimensional array containing the indices of the maximum values. """ return _arg_minmax_common(x, axis=axis, keepdims=keepdims, mode="max") def argmin(x, /, *, axis=None, keepdims=False): """ Returns the indices of the minimum values along a specified axis. When the minimum value occurs multiple times, only the indices corresponding to the first occurrence are returned. Parameters ---------- x : SparseArray Input array. The fill value must be ``0.0`` and all non-zero values must be less than ``0.0``. axis : int, optional Axis along which to search. If ``None``, the function must return the index of the minimum value of the flattened array. Default: ``None``. keepdims : bool, optional If ``True``, the reduced axes (dimensions) must be included in the result as singleton dimensions, and, accordingly, the result must be compatible with the input array. Otherwise, if ``False``, the reduced axes (dimensions) must not be included in the result. Default: ``False``. Returns ------- out : numpy.ndarray If ``axis`` is ``None``, a zero-dimensional array containing the index of the first occurrence of the minimum value. Otherwise, a non-zero-dimensional array containing the indices of the minimum values. """ return _arg_minmax_common(x, axis=axis, keepdims=keepdims, mode="min") def _replace_nan(array, value): """ Replaces ``NaN``s in ``array`` with ``value``. Parameters ---------- array : COO The input array. value : numpy.number The values to replace ``NaN`` with. Returns ------- COO A copy of ``array`` with the ``NaN``s replaced. """ if not np.issubdtype(array.dtype, np.floating): return array return where(np.isnan(array), value, array) def nanreduce(x, method, identity=None, axis=None, keepdims=False, **kwargs): """ Performs an ``NaN`` skipping reduction on this array. See the documentation on :obj:`COO.reduce` for examples. Parameters ---------- x : COO The array to reduce. method : numpy.ufunc The method to use for performing the reduction. identity : numpy.number The identity value for this reduction. Inferred from ``method`` if not given. Note that some ``ufunc`` objects don't have this, so it may be necessary to give it. axis : Union[int, Iterable[int]], optional The axes along which to perform the reduction. Uses all axes by default. keepdims : bool, optional Whether or not to keep the dimensions of the original array. **kwargs : dict Any extra arguments to pass to the reduction operation. Returns ------- COO The result of the reduction operation. Raises ------ ValueError If reducing an all-zero axis would produce a nonzero result. See Also -------- COO.reduce : Similar method without ``NaN`` skipping functionality. """ arr = _replace_nan(x, method.identity if identity is None else identity) return arr.reduce(method, axis, keepdims, **kwargs) def roll(a, shift, axis=None): """ Shifts elements of an array along specified axis. Elements that roll beyond the last position are circulated and re-introduced at the first. Parameters ---------- a : COO Input array shift : int or tuple of ints Number of index positions that elements are shifted. If a tuple is provided, then axis must be a tuple of the same size, and each of the given axes is shifted by the corresponding number. If an int while axis is a tuple of ints, then broadcasting is used so the same shift is applied to all axes. axis : int or tuple of ints, optional Axis or tuple specifying multiple axes. By default, the array is flattened before shifting, after which the original shape is restored. Returns ------- res : ndarray Output array, with the same shape as a. """ from .core import COO, as_coo a = as_coo(a) # roll flattened array if axis is None: return roll(a.reshape((-1,)), shift, 0).reshape(a.shape) # roll across specified axis # parse axis input, wrap in tuple axis = normalize_axis(axis, a.ndim) if not isinstance(axis, tuple): axis = (axis,) # make shift iterable if not isinstance(shift, Iterable): shift = (shift,) elif np.ndim(shift) > 1: raise ValueError("'shift' and 'axis' must be integers or 1D sequences.") # handle broadcasting if len(shift) == 1: shift = np.full(len(axis), shift) # check if dimensions are consistent if len(axis) != len(shift): raise ValueError("If 'shift' is a 1D sequence, 'axis' must have equal length.") if not can_store(a.coords.dtype, max(a.shape + shift)): raise ValueError( f"cannot roll with coords.dtype {a.coords.dtype} and shift {shift}. Try casting coords to a larger dtype." ) # shift elements coords, data = np.copy(a.coords), np.copy(a.data) try: for sh, ax in zip(shift, axis, strict=True): coords[ax] += sh coords[ax] %= a.shape[ax] except TypeError as e: if is_unsigned_dtype(coords.dtype): raise ValueError( f"rolling with coords.dtype as {coords.dtype} is not safe. Try using a signed dtype." ) from e return COO( coords, data=data, shape=a.shape, has_duplicates=False, fill_value=a.fill_value, ) def diagonal(a, offset=0, axis1=0, axis2=1): """ Extract diagonal from a COO array. The equivalent of :obj:`numpy.diagonal`. Parameters ---------- a : COO The array to perform the operation on. offset : int, optional Offset of the diagonal from the main diagonal. Defaults to main diagonal (0). axis1 : int, optional First axis from which the diagonals should be taken. Defaults to first axis (0). axis2 : int, optional Second axis from which the diagonals should be taken. Defaults to second axis (1). Examples -------- >>> import sparse >>> x = sparse.as_coo(np.arange(9).reshape(3, 3)) >>> sparse.diagonal(x).todense() array([0, 4, 8]) >>> sparse.diagonal(x, offset=1).todense() array([1, 5]) >>> x = sparse.as_coo(np.arange(12).reshape((2, 3, 2))) >>> x_diag = sparse.diagonal(x, axis1=0, axis2=2) >>> x_diag.shape (3, 2) >>> x_diag.todense() array([[ 0, 7], [ 2, 9], [ 4, 11]]) Returns ------- out: COO The result of the operation. Raises ------ ValueError If a.shape[axis1] != a.shape[axis2] See Also -------- :obj:`numpy.diagonal` : NumPy equivalent function """ from .core import COO if a.shape[axis1] != a.shape[axis2]: raise ValueError("a.shape[axis1] != a.shape[axis2]") diag_axes = [axis for axis in range(len(a.shape)) if axis != axis1 and axis != axis2] + [axis1] diag_shape = [a.shape[axis] for axis in diag_axes] diag_shape[-1] -= abs(offset) diag_idx = _diagonal_idx(a.coords, axis1, axis2, offset) diag_coords = [a.coords[axis][diag_idx] for axis in diag_axes] diag_data = a.data[diag_idx] return COO(diag_coords, diag_data, diag_shape) def diagonalize(a, axis=0): """ Diagonalize a COO array. The new dimension is appended at the end. .. WARNING:: :obj:`diagonalize` is not :obj:`numpy` compatible as there is no direct :obj:`numpy` equivalent. The API may change in the future. Parameters ---------- a : Union[COO, np.ndarray, scipy.sparse.spmatrix] The array to diagonalize. axis : int, optional The axis to diagonalize. Defaults to first axis (0). Examples -------- >>> import sparse >>> x = sparse.as_coo(np.arange(1, 4)) >>> sparse.diagonalize(x).todense() array([[1, 0, 0], [0, 2, 0], [0, 0, 3]]) >>> x = sparse.as_coo(np.arange(24).reshape((2, 3, 4))) >>> x_diag = sparse.diagonalize(x, axis=1) >>> x_diag.shape (2, 3, 4, 3) :obj:`diagonalize` is the inverse of :obj:`diagonal` >>> a = sparse.random((3, 3, 3, 3, 3), density=0.3) >>> a_diag = sparse.diagonalize(a, axis=2) >>> (sparse.diagonal(a_diag, axis1=2, axis2=5) == a.transpose([0, 1, 3, 4, 2])).all() True Returns ------- out: COO The result of the operation. See Also -------- :obj:`numpy.diag` : NumPy equivalent for 1D array """ from .core import COO, as_coo a = as_coo(a) diag_shape = a.shape + (a.shape[axis],) diag_coords = np.vstack([a.coords, a.coords[axis]]) return COO(diag_coords, a.data, diag_shape) def isposinf(x, out=None): """ Test element-wise for positive infinity, return result as sparse ``bool`` array. Parameters ---------- x Input out, optional Output array Examples -------- >>> import sparse >>> x = sparse.as_coo(np.array([np.inf])) >>> sparse.isposinf(x).todense() array([ True]) See Also -------- numpy.isposinf : The NumPy equivalent """ from .core import elemwise return elemwise(lambda x, out=None, dtype=None: np.isposinf(x, out=out), x, out=out) def isneginf(x, out=None): """ Test element-wise for negative infinity, return result as sparse ``bool`` array. Parameters ---------- x Input out, optional Output array Examples -------- >>> import sparse >>> x = sparse.as_coo(np.array([-np.inf])) >>> sparse.isneginf(x).todense() array([ True]) See Also -------- numpy.isneginf : The NumPy equivalent """ from .core import elemwise return elemwise(lambda x, out=None, dtype=None: np.isneginf(x, out=out), x, out=out) def result_type(*arrays_and_dtypes): """Returns the type that results from applying the NumPy type promotion rules to the arguments. See Also -------- numpy.result_type : The NumPy equivalent """ return np.result_type(*(_as_result_type_arg(x) for x in arrays_and_dtypes)) def _as_result_type_arg(x): if not isinstance(x, SparseArray): return x if x.ndim > 0: return x.dtype # 0-dimensional arrays give different result_type outputs than their dtypes return x.todense() @numba.jit(nopython=True, nogil=True) def _diagonal_idx(coordlist, axis1, axis2, offset): """ Utility function that returns all indices that correspond to a diagonal element. Parameters ---------- coordlist : list of lists Coordinate indices. axis1, axis2 : int The axes of the diagonal. offset : int Offset of the diagonal from the main diagonal. Defaults to main diagonal (0). """ return np.array([i for i in range(len(coordlist[axis1])) if coordlist[axis1][i] + offset == coordlist[axis2][i]]) def clip(a, a_min=None, a_max=None, out=None): """ Clip (limit) the values in the array. Return an array whose values are limited to ``[min, max]``. One of min or max must be given. Parameters ---------- a a_min : scalar or `SparseArray` or `None` Minimum value. If `None`, clipping is not performed on lower interval edge. a_max : scalar or `SparseArray` or `None` Maximum value. If `None`, clipping is not performed on upper interval edge. out : SparseArray, optional If provided, the results will be placed in this array. It may be the input array for in-place clipping. `out` must be of the right shape to hold the output. Its type is preserved. Returns ------- clipped_array : SparseArray An array with the elements of `self`, but where values < `min` are replaced with `min`, and those > `max` with `max`. Examples -------- >>> import sparse >>> x = sparse.COO.from_numpy([0, 0, 0, 1, 2, 3]) >>> sparse.clip(x, a_min=1).todense() # doctest: +NORMALIZE_WHITESPACE array([1, 1, 1, 1, 2, 3]) >>> sparse.clip(x, a_max=1).todense() # doctest: +NORMALIZE_WHITESPACE array([0, 0, 0, 1, 1, 1]) >>> sparse.clip(x, a_min=1, a_max=2).todense() # doctest: +NORMALIZE_WHITESPACE array([1, 1, 1, 1, 2, 2]) See Also -------- numpy.clip : Equivalent NumPy function """ a = asCOO(a, name="clip") return a.clip(a_min, a_max) def expand_dims(x, /, *, axis=0): """ Expands the shape of an array by inserting a new axis (dimension) of size one at the position specified by ``axis``. Parameters ---------- a : COO Input COO array. axis : int Position in the expanded axes where the new axis is placed. Returns ------- result : COO An expanded output COO array having the same data type as ``x``. Examples -------- >>> import sparse >>> x = sparse.COO.from_numpy([[1, 0, 0, 0, 2, -3]]) >>> x.shape (1, 6) >>> y1 = sparse.expand_dims(x, axis=1) >>> y1.shape (1, 1, 6) >>> y2 = sparse.expand_dims(x, axis=2) >>> y2.shape (1, 6, 1) """ x = _validate_coo_input(x) if not isinstance(axis, int): raise IndexError(f"Invalid axis position: {axis}") axis = normalize_axis(axis, x.ndim + 1) new_coords = np.insert(x.coords, obj=axis, values=np.zeros(x.nnz, dtype=np.intp), axis=0) new_shape = list(x.shape) new_shape.insert(axis, 1) new_shape = tuple(new_shape) from .core import COO return COO( new_coords, x.data, shape=new_shape, fill_value=x.fill_value, ) def flip(x, /, *, axis=None): """ Reverses the order of elements in an array along the given axis. The shape of the array is preserved. Parameters ---------- a : COO Input COO array. axis : int or tuple of ints, optional Axis (or axes) along which to flip. If ``axis`` is ``None``, the function must flip all input array axes. If ``axis`` is negative, the function must count from the last dimension. If provided more than one axis, the function must flip only the specified axes. Default: ``None``. Returns ------- result : COO An output array having the same data type and shape as ``x`` and whose elements, relative to ``x``, are reordered. """ x = _validate_coo_input(x) if axis is None: axis = range(x.ndim) if not isinstance(axis, Iterable): axis = (axis,) new_coords = x.coords.copy() for ax in axis: new_coords[ax, :] = x.shape[ax] - 1 - x.coords[ax, :] from .core import COO return COO( new_coords, x.data, shape=x.shape, fill_value=x.fill_value, ) # Array API set functions class UniqueCountsResult(NamedTuple): values: np.ndarray counts: np.ndarray def unique_counts(x, /): """ Returns the unique elements of an input array `x`, and the corresponding counts for each unique element in `x`. Parameters ---------- x : COO Input COO array. It will be flattened if it is not already 1-D. Returns ------- out : namedtuple The result containing: * values - The unique elements of an input array. * counts - The corresponding counts for each unique element. Raises ------ ValueError If the input array is in a different format than COO. Examples -------- >>> import sparse >>> x = sparse.COO.from_numpy([1, 0, 2, 1, 2, -3]) >>> sparse.unique_counts(x) UniqueCountsResult(values=array([-3, 0, 1, 2]), counts=array([1, 1, 2, 2])) """ x = _validate_coo_input(x) x = x.flatten() values, counts = np.unique(x.data, return_counts=True) if x.nnz < x.size: values = np.concatenate([[x.fill_value], values]) counts = np.concatenate([[x.size - x.nnz], counts]) sorted_indices = np.argsort(values) values[sorted_indices] = values.copy() counts[sorted_indices] = counts.copy() return UniqueCountsResult(values, counts) def unique_values(x, /): """ Returns the unique elements of an input array `x`. Parameters ---------- x : COO Input COO array. It will be flattened if it is not already 1-D. Returns ------- out : ndarray The unique elements of an input array. Raises ------ ValueError If the input array is in a different format than COO. Examples -------- >>> import sparse >>> x = sparse.COO.from_numpy([1, 0, 2, 1, 2, -3]) >>> sparse.unique_values(x) array([-3, 0, 1, 2]) """ x = _validate_coo_input(x) x = x.flatten() values = np.unique(x.data) if x.nnz < x.size: values = np.sort(np.concatenate([[x.fill_value], values])) return values def sort(x, /, *, axis=-1, descending=False, stable=False): """ Returns a sorted copy of an input array ``x``. Parameters ---------- x : SparseArray Input array. Should have a real-valued data type. axis : int Axis along which to sort. If set to ``-1``, the function must sort along the last axis. Default: ``-1``. descending : bool Sort order. If ``True``, the array must be sorted in descending order (by value). If ``False``, the array must be sorted in ascending order (by value). Default: ``False``. stable : bool Whether the sort is stable. Only ``False`` is supported currently. Returns ------- out : COO A sorted array. Raises ------ ValueError If the input array isn't and can't be converted to COO format. Examples -------- >>> import sparse >>> x = sparse.COO.from_numpy([1, 0, 2, 0, 2, -3]) >>> sparse.sort(x).todense() array([-3, 0, 0, 1, 2, 2]) >>> sparse.sort(x, descending=True).todense() array([ 2, 2, 1, 0, 0, -3]) """ from .._common import moveaxis from .core import COO x = _validate_coo_input(x) if stable: raise ValueError("`stable=True` isn't currently supported.") original_ndim = x.ndim if x.ndim == 1: x = x[None, :] axis = -1 x = moveaxis(x, source=axis, destination=-1) x_shape = x.shape x = x.reshape((-1, x_shape[-1])) new_coords, new_data = _sort_coo(x.coords, x.data, x.fill_value, sort_axis_len=x_shape[-1], descending=descending) x = COO(new_coords, new_data, x.shape, has_duplicates=False, sorted=True, fill_value=x.fill_value) x = x.reshape(x_shape[:-1] + (x_shape[-1],)) x = moveaxis(x, source=-1, destination=axis) return x if original_ndim == x.ndim else x.squeeze() def take(x, indices, /, *, axis=None): """ Returns elements of an array along an axis. Parameters ---------- x : SparseArray Input array. indices : ndarray Array indices. The array must be one-dimensional and have an integer data type. axis : int Axis over which to select values. If ``axis`` is negative, the function must determine the axis along which to select values by counting from the last dimension. For ``None``, the flattened input array is used. Default: ``None``. Returns ------- out : COO A COO array with requested indices. Raises ------ ValueError If the input array isn't and can't be converted to COO format. """ x = _validate_coo_input(x) if axis is None: x = x.flatten() return x[indices] axis = normalize_axis(axis, x.ndim) full_index = (slice(None),) * axis + (indices, ...) return x[full_index] def _validate_coo_input(x: Any): from .._common import _is_scipy_sparse_obj from .core import COO if _is_scipy_sparse_obj(x): x = COO.from_scipy_sparse(x) elif not isinstance(x, SparseArray): raise ValueError(f"Input must be an instance of SparseArray, but it's {type(x)}.") elif not isinstance(x, COO): x = x.asformat(COO) return x @numba.jit(nopython=True, nogil=True) def _sort_coo( coords: np.ndarray, data: np.ndarray, fill_value: float, sort_axis_len: int, descending: bool ) -> tuple[np.ndarray, np.ndarray]: assert coords.shape[0] == 2 group_coords = coords[0, :] sort_coords = coords[1, :] data = data.copy() result_indices = np.empty_like(sort_coords) # We iterate through all groups and sort each one of them. # first and last index of a group is tracked. prev_group = -1 group_first_idx = -1 group_last_idx = -1 # We add `-1` sentinel to know when the last group ends for idx, group in enumerate(np.append(group_coords, -1)): if group == prev_group: continue if prev_group != -1: group_last_idx = idx group_slice = slice(group_first_idx, group_last_idx) group_size = group_last_idx - group_first_idx # SORT VALUES if group_size > 1: # np.sort in numba doesn't support `np.sort`'s arguments so `stable` # keyword can't be supported. # https://numba.pydata.org/numba-doc/latest/reference/numpysupported.html#other-methods data[group_slice] = np.sort(data[group_slice]) if descending: data[group_slice] = data[group_slice][::-1] # SORT INDICES fill_value_count = sort_axis_len - group_size indices = np.arange(group_size) # find a place where fill_value would be for pos in range(group_size): if (not descending and fill_value < data[group_slice][pos]) or ( descending and fill_value > data[group_slice][pos] ): indices[pos:] += fill_value_count break result_indices[group_first_idx:group_last_idx] = indices prev_group = group group_first_idx = idx return np.vstack((group_coords, result_indices)), data @numba.jit(nopython=True, nogil=True) def _compute_minmax_args( coords: np.ndarray, data: np.ndarray, reduce_size: int, fill_value: float, max_mode_flag: bool, ) -> tuple[np.ndarray, np.ndarray]: assert coords.shape[0] == 2 reduce_coords = coords[0, :] index_coords = coords[1, :] result_indices = np.unique(index_coords) result_data = [] # we iterate through each trace for result_index in np.nditer(result_indices): mask = index_coords == result_index masked_reduce_coords = reduce_coords[mask] masked_data = data[mask] compared_data = operator.gt(masked_data, fill_value) if max_mode_flag else operator.lt(masked_data, fill_value) if np.any(compared_data) or len(masked_data) == reduce_size: # best value is a non-fill value best_arg = np.argmax(masked_data) if max_mode_flag else np.argmin(masked_data) result_data.append(masked_reduce_coords[best_arg]) else: # best value is a fill value, find the first occurrence of it current_coord = np.array(-1, dtype=coords.dtype) found = False for idx, new_coord in enumerate(np.nditer(np.sort(masked_reduce_coords))): # there is at least one fill value between consecutive non-fill values if new_coord - current_coord > 1: result_data.append(idx) found = True break current_coord = new_coord # get the first fill value after all non-fill values if not found: result_data.append(current_coord + 1) return (result_indices, np.array(result_data, dtype=np.intp)) def _arg_minmax_common( x: SparseArray, axis: int | None, keepdims: bool, mode: str, ): """ Internal implementation for argmax and argmin functions. """ assert mode in ("max", "min") max_mode_flag = mode == "max" x = _validate_coo_input(x) if not isinstance(axis, int | type(None)): raise ValueError(f"`axis` must be `int` or `None`, but it's: {type(axis)}.") if isinstance(axis, int) and axis >= x.ndim: raise ValueError(f"`axis={axis}` is out of bounds for array of dimension {x.ndim}.") if x.ndim == 0: raise ValueError("Input array must be at least 1-D, but it's 0-D.") # If `axis` is None then we need to flatten the input array and memorize # the original dimensionality for the final reshape operation. axis_none_original_ndim: int | None = None if axis is None: axis_none_original_ndim = x.ndim x = x.reshape(-1)[:, None] axis = 0 # A 1-D array must have one more singleton dimension. if axis == 0 and x.ndim == 1: x = x[:, None] # We need to move `axis` to the front. new_transpose = list(range(x.ndim)) new_transpose.insert(0, new_transpose.pop(axis)) new_transpose = tuple(new_transpose) # And reshape it to 2-D (reduce axis, the rest of axes flattened) new_shape = list(x.shape) new_shape.insert(0, new_shape.pop(axis)) new_shape = tuple(new_shape) x = x.transpose(new_transpose) x = x.reshape((new_shape[0], np.prod(new_shape[1:]))) # Compute max/min arguments result_indices, result_data = _compute_minmax_args( x.coords.copy(), x.data.copy(), reduce_size=x.shape[0], fill_value=x.fill_value, max_mode_flag=max_mode_flag, ) from .core import COO result = COO(result_indices, result_data, shape=(x.shape[1],), fill_value=0, prune=True) # Let's reshape the result to the original shape. result = result.reshape((1, *new_shape[1:])) new_transpose = list(range(result.ndim)) new_transpose.insert(axis, new_transpose.pop(0)) result = result.transpose(new_transpose) # If `axis=None` we need to reshape flattened array into original dimensionality. if axis_none_original_ndim is not None: result = result.reshape([1 for _ in range(axis_none_original_ndim)]) return result if keepdims else result.squeeze() def matrix_transpose(x, /): """ Transposes a matrix or a stack of matrices. Parameters ---------- x : SparseArray Input array. Returns ------- out : COO Transposed COO array. Raises ------ ValueError If the input array isn't and can't be converted to COO format, or if ``x.ndim < 2``. """ if hasattr(x, "ndim") and x.ndim < 2: raise ValueError("`x.ndim >= 2` must hold.") x = _validate_coo_input(x) transpose_axes = list(range(x.ndim)) transpose_axes[-2:] = transpose_axes[-2:][::-1] return x.transpose(transpose_axes)