Sparse n-dimensional arrays for the PyData ecosystem with multiple backend implementations
—
Functions for reshaping, indexing, slicing, and reorganizing sparse arrays while maintaining sparsity structure efficiently. These operations provide flexible ways to transform and access sparse array data.
Functions for changing array dimensions and organization.
def reshape(a, shape):
"""
Reshape sparse array to new shape.
Returns array with same data but different shape. Total number of
elements must remain constant.
Parameters:
- a: sparse array, input array to reshape
- shape: tuple, new shape for the array
Returns:
Sparse array with specified shape
"""
def squeeze(a, axis=None):
"""
Remove single-dimensional entries from array shape.
Parameters:
- a: sparse array, input array
- axis: int or tuple, specific axes to squeeze (None for all size-1 axes)
Returns:
Sparse array with specified single-dimensional entries removed
"""
def expand_dims(a, axis):
"""
Expand array dimensions by inserting new axes.
Parameters:
- a: sparse array, input array
- axis: int or tuple, position(s) of new axes
Returns:
Sparse array with expanded dimensions
"""Functions for moving and permuting array axes.
def moveaxis(a, source, destination):
"""
Move array axes to new positions.
Parameters:
- a: sparse array, input array
- source: int or sequence, original positions of axes to move
- destination: int or sequence, destination positions for moved axes
Returns:
Sparse array with moved axes
"""
def permute_dims(a, axes):
"""
Permute array dimensions according to given axes.
Parameters:
- a: sparse array, input array
- axes: tuple, permutation of axes (0, 1, ..., ndim-1)
Returns:
Sparse array with permuted dimensions
"""
def transpose(a, axes=None):
"""
Transpose sparse array by reversing or permuting axes.
Parameters:
- a: sparse array, input array
- axes: tuple, permutation of axes (None reverses all axes)
Returns:
Sparse array with transposed axes
"""Functions for broadcasting arrays to compatible shapes.
def broadcast_to(array, shape):
"""
Broadcast array to specified shape.
Parameters:
- array: sparse array, input array to broadcast
- shape: tuple, target shape for broadcasting
Returns:
Sparse array broadcast to target shape
"""
def broadcast_arrays(*arrays):
"""
Broadcast multiple arrays to common shape.
Parameters:
- arrays: sparse arrays, input arrays to broadcast together
Returns:
Tuple of sparse arrays broadcast to common shape
"""Functions for joining and stacking arrays.
def concatenate(arrays, axis=0):
"""
Join arrays along existing axis.
Parameters:
- arrays: sequence of sparse arrays with same shape except along axis
- axis: int, axis along which to concatenate
Returns:
Sparse array formed by concatenating input arrays
"""
def concat(arrays, axis=0):
"""
Alias for concatenate - join arrays along existing axis.
Parameters:
- arrays: sequence of sparse arrays
- axis: int, axis along which to concatenate
Returns:
Sparse array formed by concatenating input arrays
"""
def stack(arrays, axis=0):
"""
Join arrays along new axis.
Parameters:
- arrays: sequence of sparse arrays with same shape
- axis: int, axis along which to stack (creates new dimension)
Returns:
Sparse array formed by stacking input arrays along new axis
"""Functions for selecting and extracting array elements.
def take(a, indices, axis=None):
"""
Take elements from array along specified axis.
Parameters:
- a: sparse array, input array
- indices: array-like, indices of elements to take
- axis: int, axis along which to take elements (None flattens first)
Returns:
Sparse array with selected elements
"""
def compress(condition, a, axis=None):
"""
Select array elements using boolean mask.
Parameters:
- condition: array-like of bools, selection mask
- a: sparse array, input array
- axis: int, axis along which to apply mask (None flattens first)
Returns:
Sparse array with elements where condition is True
"""
def where(condition, x=None, y=None):
"""
Select elements from arrays based on condition.
If x and y provided: return elements from x where condition is True,
elements from y where condition is False.
If only condition provided: return indices where condition is True.
Parameters:
- condition: sparse boolean array, selection condition
- x: sparse array, values where condition is True (optional)
- y: sparse array, values where condition is False (optional)
Returns:
Selected elements or indices based on condition
"""
def nonzero(a):
"""
Find indices of non-zero elements.
Parameters:
- a: sparse array, input array
Returns:
Tuple of arrays containing indices of non-zero elements
"""
def argwhere(a):
"""
Find indices where condition is True.
Parameters:
- a: sparse array, input array (typically boolean)
Returns:
Array of indices where input is non-zero/True
"""Functions for working with array diagonals.
def diagonal(a, offset=0, axis1=0, axis2=1):
"""
Extract diagonal elements from array.
Parameters:
- a: sparse array, input array (at least 2-D)
- offset: int, diagonal offset (0=main, >0=upper, <0=lower)
- axis1: int, first axis of 2-D sub-arrays
- axis2: int, second axis of 2-D sub-arrays
Returns:
Sparse array containing diagonal elements
"""
def diagonalize(v):
"""
Create diagonal matrix from 1-D array.
Parameters:
- v: sparse 1-D array, diagonal elements
Returns:
Sparse 2-D array with v on the diagonal
"""Functions for sorting array elements and finding values.
def sort(a, axis=-1):
"""
Sort array elements along specified axis.
Parameters:
- a: sparse array, input array to sort
- axis: int, axis along which to sort (None flattens first)
Returns:
Sparse array with sorted elements
"""
def argmax(a, axis=None, keepdims=False):
"""
Find indices of maximum values along axis.
Parameters:
- a: sparse array, input array
- axis: int, axis along which to find max (None for global)
- keepdims: bool, whether to preserve dimensions
Returns:
Array of indices of maximum values
"""
def argmin(a, axis=None, keepdims=False):
"""
Find indices of minimum values along axis.
Parameters:
- a: sparse array, input array
- axis: int, axis along which to find min (None for global)
- keepdims: bool, whether to preserve dimensions
Returns:
Array of indices of minimum values
"""
def unique_values(x):
"""
Find unique elements in array.
Parameters:
- x: sparse array, input array
Returns:
Sparse array containing unique elements
"""
def unique_counts(x):
"""
Find unique elements and their counts.
Parameters:
- x: sparse array, input array
Returns:
Tuple of (unique_elements, counts) arrays
"""Functions for flipping, rolling, and padding arrays.
def flip(a, axis=None):
"""
Reverse array elements along specified axes.
Parameters:
- a: sparse array, input array
- axis: int or tuple, axes along which to flip (None for all)
Returns:
Sparse array with reversed elements
"""
def roll(a, shift, axis=None):
"""
Roll array elements along specified axis.
Parameters:
- a: sparse array, input array
- shift: int or tuple, number of positions to roll
- axis: int or tuple, axes along which to roll (None for flattened)
Returns:
Sparse array with rolled elements
"""
def pad(array, pad_width, mode='constant', constant_values=0):
"""
Pad array with values.
Parameters:
- array: sparse array, input array to pad
- pad_width: sequence, padding widths for each dimension
- mode: str, padding mode ('constant', 'edge', 'wrap', etc.)
- constant_values: scalar, value for constant padding
Returns:
Sparse array with padding applied
"""
def clip(a, min=None, max=None):
"""
Clip array values to specified range.
Parameters:
- a: sparse array, input array
- min: scalar, minimum value (None for no lower bound)
- max: scalar, maximum value (None for no upper bound)
Returns:
Sparse array with clipped values
"""Functions for extracting triangular parts of matrices.
def tril(m, k=0):
"""
Extract lower triangular part of matrix.
Parameters:
- m: sparse array, input matrix (2-D)
- k: int, diagonal offset (0=include main diagonal)
Returns:
Sparse array with upper triangle zeroed out
"""
def triu(m, k=0):
"""
Extract upper triangular part of matrix.
Parameters:
- m: sparse array, input matrix (2-D)
- k: int, diagonal offset (0=include main diagonal)
Returns:
Sparse array with lower triangle zeroed out
"""import sparse
import numpy as np
# Create test array
original = sparse.random((6, 8), density=0.2)
print(f"Original shape: {original.shape}")
# Reshape operations
reshaped = sparse.reshape(original, (4, 12))
print(f"Reshaped to: {reshaped.shape}")
flattened = sparse.reshape(original, (-1,)) # -1 infers dimension
print(f"Flattened shape: {flattened.shape}")
# Dimension manipulation
expanded = sparse.expand_dims(original, axis=0) # Add batch dimension
print(f"Expanded shape: {expanded.shape}") # (1, 6, 8)
squeezed = sparse.squeeze(expanded, axis=0) # Remove batch dimension
print(f"Squeezed shape: {squeezed.shape}") # (6, 8)# Create 3D array
array_3d = sparse.random((3, 4, 5), density=0.1)
print(f"Original shape: {array_3d.shape}")
# Various transpose operations
transposed = array_3d.transpose() # Reverse all axes
print(f"Full transpose: {transposed.shape}") # (5, 4, 3)
partial_transpose = array_3d.transpose((0, 2, 1)) # Swap last two axes
print(f"Partial transpose: {partial_transpose.shape}") # (3, 5, 4)
# Move specific axes
moved = sparse.moveaxis(array_3d, 0, -1) # Move first to last
print(f"Moved axes: {moved.shape}") # (4, 5, 3)
# Permute dimensions
permuted = sparse.permute_dims(array_3d, (2, 0, 1))
print(f"Permuted: {permuted.shape}") # (5, 3, 4)# Create arrays for combination
a1 = sparse.random((3, 4), density=0.2)
a2 = sparse.random((3, 4), density=0.2)
a3 = sparse.random((3, 4), density=0.2)
# Concatenation along existing axis
concat_result = sparse.concatenate([a1, a2, a3], axis=0)
print(f"Concatenated shape: {concat_result.shape}") # (9, 4)
concat_axis1 = sparse.concatenate([a1, a2], axis=1)
print(f"Concatenated axis 1: {concat_axis1.shape}") # (3, 8)
# Stacking along new axis
stacked = sparse.stack([a1, a2, a3], axis=0)
print(f"Stacked shape: {stacked.shape}") # (3, 3, 4)
stacked_end = sparse.stack([a1, a2], axis=-1)
print(f"Stacked at end: {stacked_end.shape}") # (3, 4, 2)# Create test array with known pattern
test_array = sparse.COO.from_numpy(
np.array([[1, 0, 3, 0], [5, 2, 0, 4], [0, 0, 6, 1]])
)
# Take specific elements
taken = sparse.take(test_array, [0, 2], axis=0) # Take rows 0 and 2
print(f"Taken rows shape: {taken.shape}") # (2, 4)
taken_cols = sparse.take(test_array, [1, 3], axis=1) # Take columns 1 and 3
print(f"Taken columns shape: {taken_cols.shape}") # (3, 2)
# Boolean masking
condition = sparse.greater(test_array, 2)
compressed = sparse.compress(condition.todense().any(axis=1), test_array, axis=0)
print(f"Compressed shape: {compressed.shape}")
# Find non-zero locations
nz_coords = sparse.nonzero(test_array)
print(f"Non-zero coordinates: {len(nz_coords)} arrays")
print(f"Number of non-zeros: {len(nz_coords[0])}")# Create arrays for conditional selection
x = sparse.random((5, 5), density=0.3)
y = sparse.ones((5, 5))
condition = sparse.greater(x, 0.5)
# Select elements based on condition
result = sparse.where(condition, x, y) # x where condition, else y
print(f"Conditional result nnz: {result.nnz}")
# Just get indices where condition is true
indices = sparse.where(condition)
print(f"True condition indices: {len(indices[0])} locations")
# Using argwhere for coordinate format
coord_indices = sparse.argwhere(condition)
print(f"Coordinate indices shape: {coord_indices.shape}")# Create matrix for diagonal operations
matrix = sparse.random((6, 6), density=0.15)
# Extract diagonals
main_diag = sparse.diagonal(matrix)
print(f"Main diagonal nnz: {main_diag.nnz}")
upper_diag = sparse.diagonal(matrix, offset=1) # Super-diagonal
lower_diag = sparse.diagonal(matrix, offset=-1) # Sub-diagonal
# Create diagonal matrix from vector
vector = sparse.COO.from_numpy(np.array([1, 2, 3, 4]))
diag_matrix = sparse.diagonalize(vector)
print(f"Diagonal matrix shape: {diag_matrix.shape}") # (4, 4)
print(f"Diagonal matrix nnz: {diag_matrix.nnz}") # 4# Create test matrix
square_matrix = sparse.random((5, 5), density=0.3)
# Extract triangular parts
lower_tri = sparse.tril(square_matrix) # Include main diagonal
upper_tri = sparse.triu(square_matrix) # Include main diagonal
strict_lower = sparse.tril(square_matrix, k=-1) # Exclude main diagonal
strict_upper = sparse.triu(square_matrix, k=1) # Exclude main diagonal
print(f"Original nnz: {square_matrix.nnz}")
print(f"Lower triangle nnz: {lower_tri.nnz}")
print(f"Upper triangle nnz: {upper_tri.nnz}")
print(f"Strict lower nnz: {strict_lower.nnz}")# Array flipping and rolling
array = sparse.COO.from_numpy(np.array([[1, 2, 3], [4, 5, 6]]))
# Flip operations
flipped_lr = sparse.flip(array, axis=1) # Flip left-right
flipped_ud = sparse.flip(array, axis=0) # Flip up-down
flipped_both = sparse.flip(array) # Flip all axes
# Roll operations
rolled = sparse.roll(array, shift=1, axis=1) # Roll columns right
print(f"Original vs rolled:")
print(array.todense())
print(rolled.todense())
# Clipping values
clipped = sparse.clip(array, min=2, max=5) # Clip to range [2,5]
print(f"Clipped array nnz: {clipped.nnz}")# Efficient operations for large sparse arrays
large_sparse = sparse.random((1000, 1000), density=0.01)
# Use views when possible (transpose, reshape with compatible shapes)
transposed_view = large_sparse.T # No data copying
reshaped_view = sparse.reshape(large_sparse, (2000, 500)) # Efficient reshape
# Batch operations
arrays_to_stack = [sparse.random((100, 100), density=0.05) for _ in range(10)]
batch_result = sparse.stack(arrays_to_stack, axis=0) # Shape: (10, 100, 100)
print(f"Batch operation result shape: {batch_result.shape}")Install with Tessl CLI
npx tessl i tessl/pypi-sparse