scirs2_sparse/
dok_array.rs

1// Dictionary of Keys (DOK) Array implementation
2//
3// This module provides the DOK (Dictionary of Keys) array format,
4// which is efficient for incremental construction of sparse arrays.
5
6use scirs2_core::ndarray::{Array1, Array2, ArrayView1};
7use scirs2_core::numeric::{Float, SparseElement};
8use std::any::Any;
9use std::collections::HashMap;
10use std::fmt::Debug;
11use std::ops::{Add, Div, Mul, Sub};
12
13use crate::coo_array::CooArray;
14use crate::error::{SparseError, SparseResult};
15use crate::lil_array::LilArray;
16use crate::sparray::{SparseArray, SparseSum};
17
18/// DOK Array format
19///
20/// The DOK (Dictionary of Keys) format stores a sparse array in a dictionary (HashMap)
21/// mapping (row, col) coordinate tuples to values.
22///
23/// # Notes
24///
25/// - Efficient for incremental construction (setting elements one by one)
26/// - Fast random access to individual elements (get/set)
27/// - Slow operations that require iterating over all elements
28/// - Slow arithmetic operations
29/// - Not suitable for large-scale computational operations
30///
31#[derive(Clone)]
32pub struct DokArray<T>
33where
34    T: SparseElement + Div<Output = T> + 'static,
35{
36    /// Dictionary mapping (row, col) to value
37    data: HashMap<(usize, usize), T>,
38    /// Shape of the sparse array
39    shape: (usize, usize),
40}
41
42impl<T> DokArray<T>
43where
44    T: SparseElement + Div<Output = T> + 'static,
45{
46    /// Creates a new DOK array with the given shape
47    ///
48    /// # Arguments
49    /// * `shape` - Shape of the sparse array (rows, cols)
50    ///
51    /// # Returns
52    /// A new empty `DokArray`
53    pub fn new(shape: (usize, usize)) -> Self {
54        Self {
55            data: HashMap::new(),
56            shape,
57        }
58    }
59
60    /// Creates a DOK array from triplet format (COO-like)
61    ///
62    /// # Arguments
63    /// * `rows` - Row indices
64    /// * `cols` - Column indices
65    /// * `data` - Values
66    /// * `shape` - Shape of the sparse array
67    ///
68    /// # Returns
69    /// A new `DokArray`
70    ///
71    /// # Errors
72    /// Returns an error if the data is not consistent
73    pub fn from_triplets(
74        rows: &[usize],
75        cols: &[usize],
76        data: &[T],
77        shape: (usize, usize),
78    ) -> SparseResult<Self> {
79        if rows.len() != cols.len() || rows.len() != data.len() {
80            return Err(SparseError::InconsistentData {
81                reason: "rows, cols, and data must have the same length".to_string(),
82            });
83        }
84
85        let mut dok = Self::new(shape);
86        for i in 0..rows.len() {
87            if rows[i] >= shape.0 || cols[i] >= shape.1 {
88                return Err(SparseError::IndexOutOfBounds {
89                    index: (rows[i], cols[i]),
90                    shape,
91                });
92            }
93            // Only set non-zero values
94            if !SparseElement::is_zero(&data[i]) {
95                dok.data.insert((rows[i], cols[i]), data[i]);
96            }
97        }
98
99        Ok(dok)
100    }
101
102    /// Returns a reference to the internal HashMap
103    pub fn get_data(&self) -> &HashMap<(usize, usize), T> {
104        &self.data
105    }
106
107    /// Returns the triplet representation (row indices, column indices, data)
108    pub fn to_triplets(&self) -> (Array1<usize>, Array1<usize>, Array1<T>)
109    where
110        T: Float + PartialOrd,
111    {
112        let nnz = self.nnz();
113        let mut row_indices = Vec::with_capacity(nnz);
114        let mut col_indices = Vec::with_capacity(nnz);
115        let mut values = Vec::with_capacity(nnz);
116
117        // Sort by row, then column for deterministic output
118        let mut entries: Vec<_> = self.data.iter().collect();
119        entries.sort_by_key(|(&(row, col), _)| (row, col));
120
121        for (&(row, col), &value) in entries {
122            row_indices.push(row);
123            col_indices.push(col);
124            values.push(value);
125        }
126
127        (
128            Array1::from_vec(row_indices),
129            Array1::from_vec(col_indices),
130            Array1::from_vec(values),
131        )
132    }
133
134    /// Creates a DOK array from a dense ndarray
135    ///
136    /// # Arguments
137    /// * `array` - Dense ndarray
138    ///
139    /// # Returns
140    /// A new `DokArray` containing non-zero elements from the input array
141    pub fn from_array(array: &Array2<T>) -> Self {
142        let shape = (array.shape()[0], array.shape()[1]);
143        let mut dok = Self::new(shape);
144
145        for ((i, j), &value) in array.indexed_iter() {
146            if !SparseElement::is_zero(&value) {
147                dok.data.insert((i, j), value);
148            }
149        }
150
151        dok
152    }
153}
154
155impl<T> SparseArray<T> for DokArray<T>
156where
157    T: SparseElement + Div<Output = T> + Float + PartialOrd + 'static,
158{
159    fn shape(&self) -> (usize, usize) {
160        self.shape
161    }
162
163    fn nnz(&self) -> usize {
164        self.data.len()
165    }
166
167    fn dtype(&self) -> &str {
168        std::any::type_name::<T>()
169    }
170
171    fn to_array(&self) -> Array2<T> {
172        let (rows, cols) = self.shape;
173        let mut result = Array2::zeros((rows, cols));
174
175        for (&(row, col), &value) in &self.data {
176            result[[row, col]] = value;
177        }
178
179        result
180    }
181
182    fn toarray(&self) -> Array2<T> {
183        self.to_array()
184    }
185
186    fn to_coo(&self) -> SparseResult<Box<dyn SparseArray<T>>> {
187        let (row_indices, col_indices, data) = self.to_triplets();
188        CooArray::new(data, row_indices, col_indices, self.shape, true)
189            .map(|array| Box::new(array) as Box<dyn SparseArray<T>>)
190    }
191
192    fn to_csr(&self) -> SparseResult<Box<dyn SparseArray<T>>> {
193        // First convert to COO, then to CSR
194        match self.to_coo() {
195            Ok(coo) => coo.to_csr(),
196            Err(e) => Err(e),
197        }
198    }
199
200    fn to_csc(&self) -> SparseResult<Box<dyn SparseArray<T>>> {
201        // First convert to COO, then to CSC
202        match self.to_coo() {
203            Ok(coo) => coo.to_csc(),
204            Err(e) => Err(e),
205        }
206    }
207
208    fn to_dok(&self) -> SparseResult<Box<dyn SparseArray<T>>> {
209        // We're already a DOK array
210        Ok(Box::new(self.clone()))
211    }
212
213    fn to_lil(&self) -> SparseResult<Box<dyn SparseArray<T>>> {
214        let (rows_arr, cols_arr, vals_arr) = self.to_triplets();
215        let rows_slice = rows_arr
216            .as_slice()
217            .ok_or_else(|| SparseError::ValueError("non-contiguous row indices".to_string()))?;
218        let cols_slice = cols_arr
219            .as_slice()
220            .ok_or_else(|| SparseError::ValueError("non-contiguous col indices".to_string()))?;
221        let vals_slice = vals_arr
222            .as_slice()
223            .ok_or_else(|| SparseError::ValueError("non-contiguous values".to_string()))?;
224        let lil = LilArray::from_triplets(rows_slice, cols_slice, vals_slice, self.shape)?;
225        Ok(Box::new(lil))
226    }
227
228    fn to_dia(&self) -> SparseResult<Box<dyn SparseArray<T>>> {
229        self.to_csr()?.to_dia()
230    }
231
232    fn to_bsr(&self) -> SparseResult<Box<dyn SparseArray<T>>> {
233        self.to_csr()?.to_bsr()
234    }
235
236    fn add(&self, other: &dyn SparseArray<T>) -> SparseResult<Box<dyn SparseArray<T>>> {
237        if self.shape() != other.shape() {
238            return Err(SparseError::DimensionMismatch {
239                expected: self.shape().0,
240                found: other.shape().0,
241            });
242        }
243
244        let mut result = self.clone();
245        let other_array = other.to_array();
246
247        // Add existing values from self
248        for (&(row, col), &value) in &self.data {
249            result.set(row, col, value + other_array[[row, col]])?;
250        }
251
252        // Add values from other that aren't in self
253        for ((row, col), &value) in other_array.indexed_iter() {
254            if !self.data.contains_key(&(row, col)) && !SparseElement::is_zero(&value) {
255                result.set(row, col, value)?;
256            }
257        }
258
259        Ok(Box::new(result))
260    }
261
262    fn sub(&self, other: &dyn SparseArray<T>) -> SparseResult<Box<dyn SparseArray<T>>> {
263        if self.shape() != other.shape() {
264            return Err(SparseError::DimensionMismatch {
265                expected: self.shape().0,
266                found: other.shape().0,
267            });
268        }
269
270        let mut result = self.clone();
271        let other_array = other.to_array();
272
273        // Subtract existing values from self
274        for (&(row, col), &value) in &self.data {
275            result.set(row, col, value - other_array[[row, col]])?;
276        }
277
278        // Subtract values from other that aren't in self
279        for ((row, col), &value) in other_array.indexed_iter() {
280            if !self.data.contains_key(&(row, col)) && !SparseElement::is_zero(&value) {
281                result.set(row, col, -value)?;
282            }
283        }
284
285        Ok(Box::new(result))
286    }
287
288    fn mul(&self, other: &dyn SparseArray<T>) -> SparseResult<Box<dyn SparseArray<T>>> {
289        if self.shape() != other.shape() {
290            return Err(SparseError::DimensionMismatch {
291                expected: self.shape().0,
292                found: other.shape().0,
293            });
294        }
295
296        let mut result = DokArray::new(self.shape());
297        let other_array = other.to_array();
298
299        // Only need to process entries in self
300        // since a*0 = 0 for any a
301        for (&(row, col), &value) in &self.data {
302            let product = value * other_array[[row, col]];
303            if !SparseElement::is_zero(&product) {
304                result.set(row, col, product)?;
305            }
306        }
307
308        Ok(Box::new(result))
309    }
310
311    fn div(&self, other: &dyn SparseArray<T>) -> SparseResult<Box<dyn SparseArray<T>>> {
312        if self.shape() != other.shape() {
313            return Err(SparseError::DimensionMismatch {
314                expected: self.shape().0,
315                found: other.shape().0,
316            });
317        }
318
319        let mut result = DokArray::new(self.shape());
320        let other_array = other.to_array();
321
322        for (&(row, col), &value) in &self.data {
323            let divisor = other_array[[row, col]];
324            if SparseElement::is_zero(&divisor) {
325                return Err(SparseError::ComputationError(
326                    "Division by zero".to_string(),
327                ));
328            }
329
330            let quotient = value / divisor;
331            if !SparseElement::is_zero(&quotient) {
332                result.set(row, col, quotient)?;
333            }
334        }
335
336        Ok(Box::new(result))
337    }
338
339    fn dot(&self, other: &dyn SparseArray<T>) -> SparseResult<Box<dyn SparseArray<T>>> {
340        let (_m, n) = self.shape();
341        let (p, q) = other.shape();
342
343        if n != p {
344            return Err(SparseError::DimensionMismatch {
345                expected: n,
346                found: p,
347            });
348        }
349
350        // Convert to CSR for efficient matrix multiplication
351        let csr_self = self.to_csr()?;
352        let csr_other = other.to_csr()?;
353
354        csr_self.dot(&*csr_other)
355    }
356
357    fn dot_vector(&self, other: &ArrayView1<T>) -> SparseResult<Array1<T>> {
358        let (m, n) = self.shape();
359        if n != other.len() {
360            return Err(SparseError::DimensionMismatch {
361                expected: n,
362                found: other.len(),
363            });
364        }
365
366        let mut result = Array1::zeros(m);
367
368        for (&(row, col), &value) in &self.data {
369            result[row] = result[row] + value * other[col];
370        }
371
372        Ok(result)
373    }
374
375    fn transpose(&self) -> SparseResult<Box<dyn SparseArray<T>>> {
376        let (rows, cols) = self.shape;
377        let mut result = DokArray::new((cols, rows));
378
379        for (&(row, col), &value) in &self.data {
380            result.set(col, row, value)?;
381        }
382
383        Ok(Box::new(result))
384    }
385
386    fn copy(&self) -> Box<dyn SparseArray<T>> {
387        Box::new(self.clone())
388    }
389
390    fn get(&self, i: usize, j: usize) -> T {
391        if i >= self.shape.0 || j >= self.shape.1 {
392            return T::sparse_zero();
393        }
394
395        *self.data.get(&(i, j)).unwrap_or(&T::sparse_zero())
396    }
397
398    fn set(&mut self, i: usize, j: usize, value: T) -> SparseResult<()> {
399        if i >= self.shape.0 || j >= self.shape.1 {
400            return Err(SparseError::IndexOutOfBounds {
401                index: (i, j),
402                shape: self.shape,
403            });
404        }
405
406        if SparseElement::is_zero(&value) {
407            // Remove zero entries
408            self.data.remove(&(i, j));
409        } else {
410            // Set non-zero value
411            self.data.insert((i, j), value);
412        }
413
414        Ok(())
415    }
416
417    fn eliminate_zeros(&mut self) {
418        // DOK format already doesn't store zeros, but just in case
419        self.data
420            .retain(|_, &mut value| !SparseElement::is_zero(&value));
421    }
422
423    fn sort_indices(&mut self) {
424        // No-op for DOK format since it's a HashMap
425    }
426
427    fn sorted_indices(&self) -> Box<dyn SparseArray<T>> {
428        // DOK doesn't have the concept of sorted indices
429        self.copy()
430    }
431
432    fn has_sorted_indices(&self) -> bool {
433        true // DOK format doesn't have the concept of sorted indices
434    }
435
436    fn sum(&self, axis: Option<usize>) -> SparseResult<SparseSum<T>> {
437        match axis {
438            None => {
439                // Sum all elements
440                let mut sum = T::sparse_zero();
441                for &value in self.data.values() {
442                    sum = sum + value;
443                }
444                Ok(SparseSum::Scalar(sum))
445            }
446            Some(0) => {
447                // Sum along rows
448                let (_, cols) = self.shape();
449                let mut result = DokArray::new((1, cols));
450
451                for (&(_row, col), &value) in &self.data {
452                    let current = result.get(0, col);
453                    result.set(0, col, current + value)?;
454                }
455
456                Ok(SparseSum::SparseArray(Box::new(result)))
457            }
458            Some(1) => {
459                // Sum along columns
460                let (rows, _) = self.shape();
461                let mut result = DokArray::new((rows, 1));
462
463                for (&(row, col), &value) in &self.data {
464                    let current = result.get(row, 0);
465                    result.set(row, 0, current + value)?;
466                }
467
468                Ok(SparseSum::SparseArray(Box::new(result)))
469            }
470            _ => Err(SparseError::InvalidAxis),
471        }
472    }
473
474    fn max(&self) -> T {
475        if self.data.is_empty() {
476            return T::nan();
477        }
478
479        self.data
480            .values()
481            .fold(T::neg_infinity(), |acc, &x| acc.max(x))
482    }
483
484    fn min(&self) -> T {
485        if self.data.is_empty() {
486            return T::nan();
487        }
488
489        self.data
490            .values()
491            .fold(T::sparse_zero(), |acc, &x| acc.min(x))
492    }
493
494    fn find(&self) -> (Array1<usize>, Array1<usize>, Array1<T>) {
495        self.to_triplets()
496    }
497
498    fn slice(
499        &self,
500        row_range: (usize, usize),
501        col_range: (usize, usize),
502    ) -> SparseResult<Box<dyn SparseArray<T>>> {
503        let (start_row, end_row) = row_range;
504        let (start_col, end_col) = col_range;
505        let (rows, cols) = self.shape;
506
507        if start_row >= rows
508            || end_row > rows
509            || start_col >= cols
510            || end_col > cols
511            || start_row >= end_row
512            || start_col >= end_col
513        {
514            return Err(SparseError::InvalidSliceRange);
515        }
516
517        let sliceshape = (end_row - start_row, end_col - start_col);
518        let mut result = DokArray::new(sliceshape);
519
520        for (&(row, col), &value) in &self.data {
521            if row >= start_row && row < end_row && col >= start_col && col < end_col {
522                result.set(row - start_row, col - start_col, value)?;
523            }
524        }
525
526        Ok(Box::new(result))
527    }
528
529    fn as_any(&self) -> &dyn Any {
530        self
531    }
532}
533
534#[cfg(test)]
535mod tests {
536    use super::*;
537    use scirs2_core::ndarray::Array;
538
539    #[test]
540    fn test_dok_array_dtype_reflects_actual_element_type() {
541        let dok_f64 = DokArray::<f64>::new((2, 2));
542        let dok_f32 = DokArray::<f32>::new((2, 2));
543
544        // Previously `dtype()` always returned the literal string "float"
545        // regardless of the actual generic element type.
546        assert_eq!(dok_f64.dtype(), "f64");
547        assert_eq!(dok_f32.dtype(), "f32");
548        assert_ne!(dok_f64.dtype(), "float");
549        assert_ne!(dok_f32.dtype(), dok_f64.dtype());
550    }
551
552    #[test]
553    fn test_dok_array_create_and_access() {
554        // Create a 3x3 sparse array
555        let mut array = DokArray::<f64>::new((3, 3));
556
557        // Set some values
558        array
559            .set(0, 0, 1.0)
560            .expect("Test: failed to set array element");
561        array
562            .set(0, 2, 2.0)
563            .expect("Test: failed to set array element");
564        array
565            .set(1, 2, 3.0)
566            .expect("Test: failed to set array element");
567        array
568            .set(2, 0, 4.0)
569            .expect("Test: failed to set array element");
570        array
571            .set(2, 1, 5.0)
572            .expect("Test: failed to set array element");
573
574        assert_eq!(array.nnz(), 5);
575
576        // Access values
577        assert_eq!(array.get(0, 0), 1.0);
578        assert_eq!(array.get(0, 1), 0.0); // Zero entry
579        assert_eq!(array.get(0, 2), 2.0);
580        assert_eq!(array.get(1, 2), 3.0);
581        assert_eq!(array.get(2, 0), 4.0);
582        assert_eq!(array.get(2, 1), 5.0);
583
584        // Set a value to zero should remove it
585        array
586            .set(0, 0, 0.0)
587            .expect("Test: failed to set array element");
588        assert_eq!(array.nnz(), 4);
589        assert_eq!(array.get(0, 0), 0.0);
590
591        // Out of bounds access should return zero
592        assert_eq!(array.get(3, 0), 0.0);
593        assert_eq!(array.get(0, 3), 0.0);
594    }
595
596    #[test]
597    fn test_dok_array_from_triplets() {
598        let rows = vec![0, 0, 1, 2, 2];
599        let cols = vec![0, 2, 2, 0, 1];
600        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
601
602        let array = DokArray::from_triplets(&rows, &cols, &data, (3, 3))
603            .expect("Test: failed to create DokArray from triplets");
604
605        assert_eq!(array.nnz(), 5);
606        assert_eq!(array.get(0, 0), 1.0);
607        assert_eq!(array.get(0, 2), 2.0);
608        assert_eq!(array.get(1, 2), 3.0);
609        assert_eq!(array.get(2, 0), 4.0);
610        assert_eq!(array.get(2, 1), 5.0);
611    }
612
613    #[test]
614    fn test_dok_array_to_array() {
615        let rows = vec![0, 0, 1, 2, 2];
616        let cols = vec![0, 2, 2, 0, 1];
617        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
618
619        let array = DokArray::from_triplets(&rows, &cols, &data, (3, 3))
620            .expect("Test: failed to create DokArray from triplets");
621        let dense = array.to_array();
622
623        let expected =
624            Array::from_shape_vec((3, 3), vec![1.0, 0.0, 2.0, 0.0, 0.0, 3.0, 4.0, 5.0, 0.0])
625                .expect("Test: failed to create array from shape vec");
626
627        assert_eq!(dense, expected);
628    }
629
630    #[test]
631    fn test_dok_array_from_array() {
632        let dense =
633            Array::from_shape_vec((3, 3), vec![1.0, 0.0, 2.0, 0.0, 0.0, 3.0, 4.0, 5.0, 0.0])
634                .expect("Test: failed to create array from shape vec");
635
636        let array = DokArray::from_array(&dense);
637
638        assert_eq!(array.nnz(), 5);
639        assert_eq!(array.get(0, 0), 1.0);
640        assert_eq!(array.get(0, 2), 2.0);
641        assert_eq!(array.get(1, 2), 3.0);
642        assert_eq!(array.get(2, 0), 4.0);
643        assert_eq!(array.get(2, 1), 5.0);
644    }
645
646    #[test]
647    fn test_dok_array_add() {
648        let mut array1 = DokArray::<f64>::new((2, 2));
649        array1
650            .set(0, 0, 1.0)
651            .expect("Test: failed to set array element");
652        array1
653            .set(0, 1, 2.0)
654            .expect("Test: failed to set array element");
655        array1
656            .set(1, 0, 3.0)
657            .expect("Test: failed to set array element");
658
659        let mut array2 = DokArray::<f64>::new((2, 2));
660        array2
661            .set(0, 0, 4.0)
662            .expect("Test: failed to set array element");
663        array2
664            .set(1, 1, 5.0)
665            .expect("Test: failed to set array element");
666
667        let result = array1.add(&array2).expect("Test: array addition failed");
668        let dense_result = result.to_array();
669
670        assert_eq!(dense_result[[0, 0]], 5.0);
671        assert_eq!(dense_result[[0, 1]], 2.0);
672        assert_eq!(dense_result[[1, 0]], 3.0);
673        assert_eq!(dense_result[[1, 1]], 5.0);
674    }
675
676    #[test]
677    fn test_dok_array_mul() {
678        let mut array1 = DokArray::<f64>::new((2, 2));
679        array1
680            .set(0, 0, 1.0)
681            .expect("Test: failed to set array element");
682        array1
683            .set(0, 1, 2.0)
684            .expect("Test: failed to set array element");
685        array1
686            .set(1, 0, 3.0)
687            .expect("Test: failed to set array element");
688        array1
689            .set(1, 1, 4.0)
690            .expect("Test: failed to set array element");
691
692        let mut array2 = DokArray::<f64>::new((2, 2));
693        array2
694            .set(0, 0, 5.0)
695            .expect("Test: failed to set array element");
696        array2
697            .set(0, 1, 6.0)
698            .expect("Test: failed to set array element");
699        array2
700            .set(1, 0, 7.0)
701            .expect("Test: failed to set array element");
702        array2
703            .set(1, 1, 8.0)
704            .expect("Test: failed to set array element");
705
706        // Element-wise multiplication
707        let result = array1
708            .mul(&array2)
709            .expect("Test: array multiplication failed");
710        let dense_result = result.to_array();
711
712        assert_eq!(dense_result[[0, 0]], 5.0);
713        assert_eq!(dense_result[[0, 1]], 12.0);
714        assert_eq!(dense_result[[1, 0]], 21.0);
715        assert_eq!(dense_result[[1, 1]], 32.0);
716    }
717
718    #[test]
719    fn test_dok_array_dot() {
720        let mut array1 = DokArray::<f64>::new((2, 2));
721        array1
722            .set(0, 0, 1.0)
723            .expect("Test: failed to set array element");
724        array1
725            .set(0, 1, 2.0)
726            .expect("Test: failed to set array element");
727        array1
728            .set(1, 0, 3.0)
729            .expect("Test: failed to set array element");
730        array1
731            .set(1, 1, 4.0)
732            .expect("Test: failed to set array element");
733
734        let mut array2 = DokArray::<f64>::new((2, 2));
735        array2
736            .set(0, 0, 5.0)
737            .expect("Test: failed to set array element");
738        array2
739            .set(0, 1, 6.0)
740            .expect("Test: failed to set array element");
741        array2
742            .set(1, 0, 7.0)
743            .expect("Test: failed to set array element");
744        array2
745            .set(1, 1, 8.0)
746            .expect("Test: failed to set array element");
747
748        // Matrix multiplication
749        let result = array1.dot(&array2).expect("Test: array dot product failed");
750        let dense_result = result.to_array();
751
752        // [1 2] [5 6] = [1*5 + 2*7, 1*6 + 2*8] = [19, 22]
753        // [3 4] [7 8]   [3*5 + 4*7, 3*6 + 4*8]   [43, 50]
754        assert_eq!(dense_result[[0, 0]], 19.0);
755        assert_eq!(dense_result[[0, 1]], 22.0);
756        assert_eq!(dense_result[[1, 0]], 43.0);
757        assert_eq!(dense_result[[1, 1]], 50.0);
758    }
759
760    #[test]
761    fn test_dok_array_transpose() {
762        let mut array = DokArray::<f64>::new((2, 3));
763        array
764            .set(0, 0, 1.0)
765            .expect("Test: failed to set array element");
766        array
767            .set(0, 1, 2.0)
768            .expect("Test: failed to set array element");
769        array
770            .set(0, 2, 3.0)
771            .expect("Test: failed to set array element");
772        array
773            .set(1, 0, 4.0)
774            .expect("Test: failed to set array element");
775        array
776            .set(1, 1, 5.0)
777            .expect("Test: failed to set array element");
778        array
779            .set(1, 2, 6.0)
780            .expect("Test: failed to set array element");
781
782        let transposed = array.transpose().expect("Test: array transpose failed");
783
784        assert_eq!(transposed.shape(), (3, 2));
785        assert_eq!(transposed.get(0, 0), 1.0);
786        assert_eq!(transposed.get(1, 0), 2.0);
787        assert_eq!(transposed.get(2, 0), 3.0);
788        assert_eq!(transposed.get(0, 1), 4.0);
789        assert_eq!(transposed.get(1, 1), 5.0);
790        assert_eq!(transposed.get(2, 1), 6.0);
791    }
792
793    #[test]
794    fn test_dok_array_slice() {
795        let mut array = DokArray::<f64>::new((3, 3));
796        array
797            .set(0, 0, 1.0)
798            .expect("Test: failed to set array element");
799        array
800            .set(0, 1, 2.0)
801            .expect("Test: failed to set array element");
802        array
803            .set(0, 2, 3.0)
804            .expect("Test: failed to set array element");
805        array
806            .set(1, 0, 4.0)
807            .expect("Test: failed to set array element");
808        array
809            .set(1, 1, 5.0)
810            .expect("Test: failed to set array element");
811        array
812            .set(1, 2, 6.0)
813            .expect("Test: failed to set array element");
814        array
815            .set(2, 0, 7.0)
816            .expect("Test: failed to set array element");
817        array
818            .set(2, 1, 8.0)
819            .expect("Test: failed to set array element");
820        array
821            .set(2, 2, 9.0)
822            .expect("Test: failed to set array element");
823
824        let slice = array
825            .slice((0, 2), (1, 3))
826            .expect("Test: array slice failed");
827
828        assert_eq!(slice.shape(), (2, 2));
829        assert_eq!(slice.get(0, 0), 2.0);
830        assert_eq!(slice.get(0, 1), 3.0);
831        assert_eq!(slice.get(1, 0), 5.0);
832        assert_eq!(slice.get(1, 1), 6.0);
833    }
834
835    #[test]
836    fn test_dok_array_sum() {
837        let mut array = DokArray::<f64>::new((2, 3));
838        array
839            .set(0, 0, 1.0)
840            .expect("Test: failed to set array element");
841        array
842            .set(0, 1, 2.0)
843            .expect("Test: failed to set array element");
844        array
845            .set(0, 2, 3.0)
846            .expect("Test: failed to set array element");
847        array
848            .set(1, 0, 4.0)
849            .expect("Test: failed to set array element");
850        array
851            .set(1, 1, 5.0)
852            .expect("Test: failed to set array element");
853        array
854            .set(1, 2, 6.0)
855            .expect("Test: failed to set array element");
856
857        // Sum all elements
858        match array.sum(None).expect("Test: array sum failed") {
859            SparseSum::Scalar(sum) => assert_eq!(sum, 21.0),
860            _ => panic!("Expected scalar sum"),
861        }
862
863        // Sum along rows (axis 0)
864        match array.sum(Some(0)).expect("Test: array sum failed") {
865            SparseSum::SparseArray(sum_array) => {
866                assert_eq!(sum_array.shape(), (1, 3));
867                assert_eq!(sum_array.get(0, 0), 5.0);
868                assert_eq!(sum_array.get(0, 1), 7.0);
869                assert_eq!(sum_array.get(0, 2), 9.0);
870            }
871            _ => panic!("Expected sparse array"),
872        }
873
874        // Sum along columns (axis 1)
875        match array.sum(Some(1)).expect("Test: array sum failed") {
876            SparseSum::SparseArray(sum_array) => {
877                assert_eq!(sum_array.shape(), (2, 1));
878                assert_eq!(sum_array.get(0, 0), 6.0);
879                assert_eq!(sum_array.get(1, 0), 15.0);
880            }
881            _ => panic!("Expected sparse array"),
882        }
883    }
884}
scirs2_sparse/dok_array.rs

scirs2_sparse/
dok_array.rs
