0
# Array Manipulation Functions
1

2
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.
3

4
## Capabilities
5

6
### Shape Manipulation
7

8
Functions for changing array dimensions and organization.
9

10
```python { .api }
11
def reshape(a, shape):
12
    """
13
    Reshape sparse array to new shape.
14
    
15
    Returns array with same data but different shape. Total number of 
16
    elements must remain constant.
17
    
18
    Parameters:
19
    - a: sparse array, input array to reshape
20
    - shape: tuple, new shape for the array
21
    
22
    Returns:
23
    Sparse array with specified shape
24
    """
25

26
def squeeze(a, axis=None):
27
    """
28
    Remove single-dimensional entries from array shape.
29
    
30
    Parameters:
31
    - a: sparse array, input array
32
    - axis: int or tuple, specific axes to squeeze (None for all size-1 axes)
33
    
34
    Returns:
35
    Sparse array with specified single-dimensional entries removed
36
    """
37

38
def expand_dims(a, axis):
39
    """
40
    Expand array dimensions by inserting new axes.
41
    
42
    Parameters:
43
    - a: sparse array, input array
44
    - axis: int or tuple, position(s) of new axes
45
    
46
    Returns:
47
    Sparse array with expanded dimensions
48
    """
49
```
50

51
### Axis Operations
52

53
Functions for moving and permuting array axes.
54

55
```python { .api }
56
def moveaxis(a, source, destination):
57
    """
58
    Move array axes to new positions.
59
    
60
    Parameters:
61
    - a: sparse array, input array
62
    - source: int or sequence, original positions of axes to move
63
    - destination: int or sequence, destination positions for moved axes
64
    
65
    Returns:
66
    Sparse array with moved axes
67
    """
68

69
def permute_dims(a, axes):
70
    """
71
    Permute array dimensions according to given axes.
72
    
73
    Parameters:
74
    - a: sparse array, input array
75
    - axes: tuple, permutation of axes (0, 1, ..., ndim-1)
76
    
77
    Returns:
78
    Sparse array with permuted dimensions
79
    """
80

81
def transpose(a, axes=None):
82
    """
83
    Transpose sparse array by reversing or permuting axes.
84
    
85
    Parameters:
86
    - a: sparse array, input array
87
    - axes: tuple, permutation of axes (None reverses all axes)
88
    
89
    Returns:
90
    Sparse array with transposed axes
91
    """
92
```
93

94
### Broadcasting Operations
95

96
Functions for broadcasting arrays to compatible shapes.
97

98
```python { .api }
99
def broadcast_to(array, shape):
100
    """
101
    Broadcast array to specified shape.
102
    
103
    Parameters:
104
    - array: sparse array, input array to broadcast
105
    - shape: tuple, target shape for broadcasting
106
    
107
    Returns:
108
    Sparse array broadcast to target shape
109
    """
110

111
def broadcast_arrays(*arrays):
112
    """
113
    Broadcast multiple arrays to common shape.
114
    
115
    Parameters:
116
    - arrays: sparse arrays, input arrays to broadcast together
117
    
118
    Returns:
119
    Tuple of sparse arrays broadcast to common shape
120
    """
121
```
122

123
### Array Combination
124

125
Functions for joining and stacking arrays.
126

127
```python { .api }
128
def concatenate(arrays, axis=0):
129
    """
130
    Join arrays along existing axis.
131
    
132
    Parameters:
133
    - arrays: sequence of sparse arrays with same shape except along axis
134
    - axis: int, axis along which to concatenate
135
    
136
    Returns:
137
    Sparse array formed by concatenating input arrays
138
    """
139

140
def concat(arrays, axis=0):
141
    """
142
    Alias for concatenate - join arrays along existing axis.
143
    
144
    Parameters:
145
    - arrays: sequence of sparse arrays
146
    - axis: int, axis along which to concatenate
147
    
148
    Returns:
149
    Sparse array formed by concatenating input arrays
150
    """
151

152
def stack(arrays, axis=0):
153
    """
154
    Join arrays along new axis.
155
    
156
    Parameters:
157
    - arrays: sequence of sparse arrays with same shape
158
    - axis: int, axis along which to stack (creates new dimension)
159
    
160
    Returns:
161
    Sparse array formed by stacking input arrays along new axis
162
    """
163
```
164

165
### Indexing and Selection
166

167
Functions for selecting and extracting array elements.
168

169
```python { .api }
170
def take(a, indices, axis=None):
171
    """
172
    Take elements from array along specified axis.
173
    
174
    Parameters:
175
    - a: sparse array, input array
176
    - indices: array-like, indices of elements to take
177
    - axis: int, axis along which to take elements (None flattens first)
178
    
179
    Returns:
180
    Sparse array with selected elements
181
    """
182

183
def compress(condition, a, axis=None):
184
    """
185
    Select array elements using boolean mask.
186
    
187
    Parameters:
188
    - condition: array-like of bools, selection mask
189
    - a: sparse array, input array
190
    - axis: int, axis along which to apply mask (None flattens first)
191
    
192
    Returns:
193
    Sparse array with elements where condition is True
194
    """
195

196
def where(condition, x=None, y=None):
197
    """
198
    Select elements from arrays based on condition.
199
    
200
    If x and y provided: return elements from x where condition is True,
201
    elements from y where condition is False.
202
    If only condition provided: return indices where condition is True.
203
    
204
    Parameters:
205
    - condition: sparse boolean array, selection condition
206
    - x: sparse array, values where condition is True (optional)
207
    - y: sparse array, values where condition is False (optional)
208
    
209
    Returns:
210
    Selected elements or indices based on condition
211
    """
212

213
def nonzero(a):
214
    """
215
    Find indices of non-zero elements.
216
    
217
    Parameters:
218
    - a: sparse array, input array
219
    
220
    Returns:
221
    Tuple of arrays containing indices of non-zero elements
222
    """
223

224
def argwhere(a):
225
    """
226
    Find indices where condition is True.
227
    
228
    Parameters:
229
    - a: sparse array, input array (typically boolean)
230
    
231
    Returns:
232
    Array of indices where input is non-zero/True
233
    """
234
```
235

236
### Diagonal Operations
237

238
Functions for working with array diagonals.
239

240
```python { .api }
241
def diagonal(a, offset=0, axis1=0, axis2=1):
242
    """
243
    Extract diagonal elements from array.
244
    
245
    Parameters:
246
    - a: sparse array, input array (at least 2-D)
247
    - offset: int, diagonal offset (0=main, >0=upper, <0=lower)
248
    - axis1: int, first axis of 2-D sub-arrays
249
    - axis2: int, second axis of 2-D sub-arrays  
250
    
251
    Returns:
252
    Sparse array containing diagonal elements
253
    """
254

255
def diagonalize(v):
256
    """
257
    Create diagonal matrix from 1-D array.
258
    
259
    Parameters:
260
    - v: sparse 1-D array, diagonal elements
261
    
262
    Returns:
263
    Sparse 2-D array with v on the diagonal
264
    """
265
```
266

267
### Sorting and Searching
268

269
Functions for sorting array elements and finding values.
270

271
```python { .api }
272
def sort(a, axis=-1):
273
    """
274
    Sort array elements along specified axis.
275
    
276
    Parameters:
277
    - a: sparse array, input array to sort
278
    - axis: int, axis along which to sort (None flattens first)
279
    
280
    Returns:
281
    Sparse array with sorted elements
282
    """
283

284
def argmax(a, axis=None, keepdims=False):
285
    """
286
    Find indices of maximum values along axis.
287
    
288
    Parameters:
289
    - a: sparse array, input array
290
    - axis: int, axis along which to find max (None for global)
291
    - keepdims: bool, whether to preserve dimensions
292
    
293
    Returns:
294
    Array of indices of maximum values
295
    """
296

297
def argmin(a, axis=None, keepdims=False):
298
    """
299
    Find indices of minimum values along axis.
300
    
301
    Parameters:
302
    - a: sparse array, input array
303
    - axis: int, axis along which to find min (None for global)
304
    - keepdims: bool, whether to preserve dimensions
305
    
306
    Returns:
307
    Array of indices of minimum values
308
    """
309

310
def unique_values(x):
311
    """
312
    Find unique elements in array.
313
    
314
    Parameters:
315
    - x: sparse array, input array
316
    
317
    Returns:
318
    Sparse array containing unique elements
319
    """
320

321
def unique_counts(x):
322
    """
323
    Find unique elements and their counts.
324
    
325
    Parameters:
326
    - x: sparse array, input array
327
    
328
    Returns:
329
    Tuple of (unique_elements, counts) arrays
330
    """
331
```
332

333
### Array Manipulation Utilities
334

335
Functions for flipping, rolling, and padding arrays.
336

337
```python { .api }
338
def flip(a, axis=None):
339
    """
340
    Reverse array elements along specified axes.
341
    
342
    Parameters:
343
    - a: sparse array, input array
344
    - axis: int or tuple, axes along which to flip (None for all)
345
    
346
    Returns:
347
    Sparse array with reversed elements
348
    """
349

350
def roll(a, shift, axis=None):
351
    """
352
    Roll array elements along specified axis.
353
    
354
    Parameters:
355
    - a: sparse array, input array
356
    - shift: int or tuple, number of positions to roll
357
    - axis: int or tuple, axes along which to roll (None for flattened)
358
    
359
    Returns:
360
    Sparse array with rolled elements
361
    """
362

363
def pad(array, pad_width, mode='constant', constant_values=0):
364
    """
365
    Pad array with values.
366
    
367
    Parameters:
368
    - array: sparse array, input array to pad
369
    - pad_width: sequence, padding widths for each dimension
370
    - mode: str, padding mode ('constant', 'edge', 'wrap', etc.)
371
    - constant_values: scalar, value for constant padding
372
    
373
    Returns:
374
    Sparse array with padding applied
375
    """
376

377
def clip(a, min=None, max=None):
378
    """
379
    Clip array values to specified range.
380
    
381
    Parameters:
382
    - a: sparse array, input array
383
    - min: scalar, minimum value (None for no lower bound)
384
    - max: scalar, maximum value (None for no upper bound)
385
    
386
    Returns:
387
    Sparse array with clipped values
388
    """
389
```
390

391
### Triangular Matrix Operations
392

393
Functions for extracting triangular parts of matrices.
394

395
```python { .api }
396
def tril(m, k=0):
397
    """
398
    Extract lower triangular part of matrix.
399
    
400
    Parameters:
401
    - m: sparse array, input matrix (2-D)
402
    - k: int, diagonal offset (0=include main diagonal)
403
    
404
    Returns:
405
    Sparse array with upper triangle zeroed out
406
    """
407

408
def triu(m, k=0):
409
    """
410
    Extract upper triangular part of matrix.
411
    
412
    Parameters:
413
    - m: sparse array, input matrix (2-D)
414
    - k: int, diagonal offset (0=include main diagonal)
415
    
416
    Returns:
417
    Sparse array with lower triangle zeroed out
418
    """
419
```
420

421
## Usage Examples
422

423
### Shape Manipulation
424

425
```python
426
import sparse
427
import numpy as np
428

429
# Create test array
430
original = sparse.random((6, 8), density=0.2)
431
print(f"Original shape: {original.shape}")
432

433
# Reshape operations
434
reshaped = sparse.reshape(original, (4, 12))
435
print(f"Reshaped to: {reshaped.shape}")
436

437
flattened = sparse.reshape(original, (-1,))  # -1 infers dimension
438
print(f"Flattened shape: {flattened.shape}")
439

440
# Dimension manipulation
441
expanded = sparse.expand_dims(original, axis=0)  # Add batch dimension
442
print(f"Expanded shape: {expanded.shape}")       # (1, 6, 8)
443

444
squeezed = sparse.squeeze(expanded, axis=0)      # Remove batch dimension
445
print(f"Squeezed shape: {squeezed.shape}")      # (6, 8)
446
```
447

448
### Array Transposition and Axis Operations
449

450
```python
451
# Create 3D array
452
array_3d = sparse.random((3, 4, 5), density=0.1)
453
print(f"Original shape: {array_3d.shape}")
454

455
# Various transpose operations
456
transposed = array_3d.transpose()                    # Reverse all axes
457
print(f"Full transpose: {transposed.shape}")        # (5, 4, 3)
458

459
partial_transpose = array_3d.transpose((0, 2, 1))   # Swap last two axes
460
print(f"Partial transpose: {partial_transpose.shape}")  # (3, 5, 4)
461

462
# Move specific axes
463
moved = sparse.moveaxis(array_3d, 0, -1)            # Move first to last
464
print(f"Moved axes: {moved.shape}")                 # (4, 5, 3)
465

466
# Permute dimensions
467
permuted = sparse.permute_dims(array_3d, (2, 0, 1))
468
print(f"Permuted: {permuted.shape}")                # (5, 3, 4)
469
```
470

471
### Array Combination
472

473
```python
474
# Create arrays for combination
475
a1 = sparse.random((3, 4), density=0.2)
476
a2 = sparse.random((3, 4), density=0.2)  
477
a3 = sparse.random((3, 4), density=0.2)
478

479
# Concatenation along existing axis
480
concat_result = sparse.concatenate([a1, a2, a3], axis=0)
481
print(f"Concatenated shape: {concat_result.shape}")  # (9, 4)
482

483
concat_axis1 = sparse.concatenate([a1, a2], axis=1)
484
print(f"Concatenated axis 1: {concat_axis1.shape}")  # (3, 8)
485

486
# Stacking along new axis
487
stacked = sparse.stack([a1, a2, a3], axis=0)
488
print(f"Stacked shape: {stacked.shape}")             # (3, 3, 4)
489

490
stacked_end = sparse.stack([a1, a2], axis=-1)
491
print(f"Stacked at end: {stacked_end.shape}")        # (3, 4, 2)
492
```
493

494
### Indexing and Selection
495

496
```python
497
# Create test array with known pattern
498
test_array = sparse.COO.from_numpy(
499
    np.array([[1, 0, 3, 0], [5, 2, 0, 4], [0, 0, 6, 1]])
500
)
501

502
# Take specific elements
503
taken = sparse.take(test_array, [0, 2], axis=0)      # Take rows 0 and 2
504
print(f"Taken rows shape: {taken.shape}")           # (2, 4)
505

506
taken_cols = sparse.take(test_array, [1, 3], axis=1) # Take columns 1 and 3
507
print(f"Taken columns shape: {taken_cols.shape}")   # (3, 2)
508

509
# Boolean masking
510
condition = sparse.greater(test_array, 2)
511
compressed = sparse.compress(condition.todense().any(axis=1), test_array, axis=0)
512
print(f"Compressed shape: {compressed.shape}")
513

514
# Find non-zero locations
515
nz_coords = sparse.nonzero(test_array)
516
print(f"Non-zero coordinates: {len(nz_coords)} arrays")
517
print(f"Number of non-zeros: {len(nz_coords[0])}")
518
```
519

520
### Conditional Selection
521

522
```python
523
# Create arrays for conditional selection
524
x = sparse.random((5, 5), density=0.3)
525
y = sparse.ones((5, 5))
526
condition = sparse.greater(x, 0.5)
527

528
# Select elements based on condition
529
result = sparse.where(condition, x, y)  # x where condition, else y
530
print(f"Conditional result nnz: {result.nnz}")
531

532
# Just get indices where condition is true
533
indices = sparse.where(condition)
534
print(f"True condition indices: {len(indices[0])} locations")
535

536
# Using argwhere for coordinate format
537
coord_indices = sparse.argwhere(condition)
538
print(f"Coordinate indices shape: {coord_indices.shape}")
539
```
540

541
### Diagonal Operations
542

543
```python
544
# Create matrix for diagonal operations
545
matrix = sparse.random((6, 6), density=0.15)
546

547
# Extract diagonals
548
main_diag = sparse.diagonal(matrix)
549
print(f"Main diagonal nnz: {main_diag.nnz}")
550

551
upper_diag = sparse.diagonal(matrix, offset=1)  # Super-diagonal
552
lower_diag = sparse.diagonal(matrix, offset=-1) # Sub-diagonal
553

554
# Create diagonal matrix from vector
555
vector = sparse.COO.from_numpy(np.array([1, 2, 3, 4]))
556
diag_matrix = sparse.diagonalize(vector)
557
print(f"Diagonal matrix shape: {diag_matrix.shape}")  # (4, 4)
558
print(f"Diagonal matrix nnz: {diag_matrix.nnz}")     # 4
559
```
560

561
### Triangular Operations
562

563
```python
564
# Create test matrix
565
square_matrix = sparse.random((5, 5), density=0.3)
566

567
# Extract triangular parts
568
lower_tri = sparse.tril(square_matrix)              # Include main diagonal
569
upper_tri = sparse.triu(square_matrix)              # Include main diagonal
570

571
strict_lower = sparse.tril(square_matrix, k=-1)     # Exclude main diagonal
572
strict_upper = sparse.triu(square_matrix, k=1)      # Exclude main diagonal
573

574
print(f"Original nnz: {square_matrix.nnz}")
575
print(f"Lower triangle nnz: {lower_tri.nnz}")
576
print(f"Upper triangle nnz: {upper_tri.nnz}")
577
print(f"Strict lower nnz: {strict_lower.nnz}")
578
```
579

580
### Array Manipulation Utilities
581

582
```python
583
# Array flipping and rolling
584
array = sparse.COO.from_numpy(np.array([[1, 2, 3], [4, 5, 6]]))
585

586
# Flip operations
587
flipped_lr = sparse.flip(array, axis=1)    # Flip left-right
588
flipped_ud = sparse.flip(array, axis=0)    # Flip up-down
589
flipped_both = sparse.flip(array)          # Flip all axes
590

591
# Roll operations
592
rolled = sparse.roll(array, shift=1, axis=1)  # Roll columns right
593
print(f"Original vs rolled:")
594
print(array.todense())
595
print(rolled.todense())
596

597
# Clipping values
598
clipped = sparse.clip(array, min=2, max=5)     # Clip to range [2,5]
599
print(f"Clipped array nnz: {clipped.nnz}")
600
```
601

602
### Performance-Optimized Manipulation
603

604
```python
605
# Efficient operations for large sparse arrays
606
large_sparse = sparse.random((1000, 1000), density=0.01)
607

608
# Use views when possible (transpose, reshape with compatible shapes)
609
transposed_view = large_sparse.T  # No data copying
610
reshaped_view = sparse.reshape(large_sparse, (2000, 500))  # Efficient reshape
611

612
# Batch operations
613
arrays_to_stack = [sparse.random((100, 100), density=0.05) for _ in range(10)]
614
batch_result = sparse.stack(arrays_to_stack, axis=0)  # Shape: (10, 100, 100)
615
print(f"Batch operation result shape: {batch_result.shape}")
616
```
617

618
## Memory and Performance Considerations
619

620
- **Shape operations**: Most operations preserve sparsity structure
621
- **Concatenation**: May require coordinate recomputation for optimal storage
622
- **Indexing**: Advanced indexing may increase result density
623
- **Broadcasting**: Can significantly increase memory usage with sparse arrays
624
- **Views vs. Copies**: Transpose and compatible reshapes create views when possible