ghostflow_core/
sparse.rs

1//! Sparse tensor operations
2//!
3//! Implements sparse tensor formats:
4//! - COO (Coordinate format)
5//! - CSR (Compressed Sparse Row)
6//! - CSC (Compressed Sparse Column)
7
8use crate::tensor::Tensor;
9use crate::error::{GhostError, Result};
10use std::collections::HashMap;
11
12/// Sparse tensor in COO (Coordinate) format
13#[derive(Debug, Clone)]
14pub struct SparseTensorCOO {
15    /// Non-zero values
16    pub values: Vec<f32>,
17    /// Row indices
18    pub rows: Vec<usize>,
19    /// Column indices
20    pub cols: Vec<usize>,
21    /// Shape of the tensor
22    pub shape: Vec<usize>,
23    /// Number of non-zero elements
24    pub nnz: usize,
25}
26
27impl SparseTensorCOO {
28    /// Create a new sparse tensor in COO format
29    pub fn new(values: Vec<f32>, rows: Vec<usize>, cols: Vec<usize>, shape: Vec<usize>) -> Result<Self> {
30        if values.len() != rows.len() || values.len() != cols.len() {
31            return Err(GhostError::InvalidShape(
32                "Values, rows, and cols must have the same length".to_string()
33            ));
34        }
35        
36        let nnz = values.len();
37        
38        Ok(SparseTensorCOO {
39            values,
40            rows,
41            cols,
42            shape,
43            nnz,
44        })
45    }
46    
47    /// Create sparse tensor from dense tensor (threshold-based)
48    pub fn from_dense(tensor: &Tensor, threshold: f32) -> Self {
49        let data = tensor.data_f32();
50        let shape = tensor.dims().to_vec();
51        
52        let mut values = Vec::new();
53        let mut rows = Vec::new();
54        let mut cols = Vec::new();
55        
56        if shape.len() == 2 {
57            let (nrows, ncols) = (shape[0], shape[1]);
58            for i in 0..nrows {
59                for j in 0..ncols {
60                    let val = data[i * ncols + j];
61                    if val.abs() > threshold {
62                        values.push(val);
63                        rows.push(i);
64                        cols.push(j);
65                    }
66                }
67            }
68        }
69        
70        let nnz = values.len();
71        SparseTensorCOO {
72            values,
73            rows,
74            cols,
75            shape,
76            nnz,
77        }
78    }
79    
80    /// Convert to dense tensor
81    pub fn to_dense(&self) -> Result<Tensor> {
82        if self.shape.len() != 2 {
83            return Err(GhostError::InvalidShape(
84                "Only 2D sparse tensors supported".to_string()
85            ));
86        }
87        
88        let (nrows, ncols) = (self.shape[0], self.shape[1]);
89        let mut data = vec![0.0f32; nrows * ncols];
90        
91        for i in 0..self.nnz {
92            let row = self.rows[i];
93            let col = self.cols[i];
94            let val = self.values[i];
95            data[row * ncols + col] = val;
96        }
97        
98        Tensor::from_slice(&data, &self.shape)
99    }
100    
101    /// Sparse matrix-vector multiplication
102    pub fn spmv(&self, vec: &Tensor) -> Result<Tensor> {
103        if self.shape.len() != 2 {
104            return Err(GhostError::InvalidShape(
105                "SpMV requires 2D sparse matrix".to_string()
106            ));
107        }
108        
109        let vec_data = vec.data_f32();
110        if vec_data.len() != self.shape[1] {
111            return Err(GhostError::InvalidShape(
112                format!("Vector length {} doesn't match matrix columns {}", vec_data.len(), self.shape[1])
113            ));
114        }
115        
116        let mut result = vec![0.0f32; self.shape[0]];
117        
118        for i in 0..self.nnz {
119            let row = self.rows[i];
120            let col = self.cols[i];
121            let val = self.values[i];
122            result[row] += val * vec_data[col];
123        }
124        
125        Tensor::from_slice(&result, &[self.shape[0]])
126    }
127    
128    /// Sparse matrix-matrix multiplication
129    pub fn spmm(&self, other: &SparseTensorCOO) -> Result<SparseTensorCOO> {
130        if self.shape[1] != other.shape[0] {
131            return Err(GhostError::InvalidShape(
132                format!("Matrix dimensions don't match: {:?} x {:?}", self.shape, other.shape)
133            ));
134        }
135        
136        // Build hash map for efficient lookup
137        let mut other_map: HashMap<(usize, usize), f32> = HashMap::new();
138        for i in 0..other.nnz {
139            other_map.insert((other.rows[i], other.cols[i]), other.values[i]);
140        }
141        
142        let mut result_map: HashMap<(usize, usize), f32> = HashMap::new();
143        
144        // Compute C = A * B
145        for i in 0..self.nnz {
146            let row = self.rows[i];
147            let k = self.cols[i];
148            let a_val = self.values[i];
149            
150            // Find all B[k, col]
151            for j in 0..other.nnz {
152                if other.rows[j] == k {
153                    let col = other.cols[j];
154                    let b_val = other.values[j];
155                    *result_map.entry((row, col)).or_insert(0.0) += a_val * b_val;
156                }
157            }
158        }
159        
160        // Convert result map to COO format
161        let mut values = Vec::new();
162        let mut rows = Vec::new();
163        let mut cols = Vec::new();
164        
165        for ((row, col), val) in result_map {
166            if val.abs() > 1e-10 {
167                values.push(val);
168                rows.push(row);
169                cols.push(col);
170            }
171        }
172        
173        SparseTensorCOO::new(values, rows, cols, vec![self.shape[0], other.shape[1]])
174    }
175    
176    /// Element-wise addition with another sparse tensor
177    pub fn add(&self, other: &SparseTensorCOO) -> Result<SparseTensorCOO> {
178        if self.shape != other.shape {
179            return Err(GhostError::InvalidShape(
180                format!("Shape mismatch: {:?} vs {:?}", self.shape, other.shape)
181            ));
182        }
183        
184        let mut result_map: HashMap<(usize, usize), f32> = HashMap::new();
185        
186        // Add values from self
187        for i in 0..self.nnz {
188            let key = (self.rows[i], self.cols[i]);
189            *result_map.entry(key).or_insert(0.0) += self.values[i];
190        }
191        
192        // Add values from other
193        for i in 0..other.nnz {
194            let key = (other.rows[i], other.cols[i]);
195            *result_map.entry(key).or_insert(0.0) += other.values[i];
196        }
197        
198        // Convert to COO
199        let mut values = Vec::new();
200        let mut rows = Vec::new();
201        let mut cols = Vec::new();
202        
203        for ((row, col), val) in result_map {
204            if val.abs() > 1e-10 {
205                values.push(val);
206                rows.push(row);
207                cols.push(col);
208            }
209        }
210        
211        SparseTensorCOO::new(values, rows, cols, self.shape.clone())
212    }
213    
214    /// Transpose the sparse matrix
215    pub fn transpose(&self) -> Result<SparseTensorCOO> {
216        if self.shape.len() != 2 {
217            return Err(GhostError::InvalidShape(
218                "Transpose only supported for 2D tensors".to_string()
219            ));
220        }
221        
222        SparseTensorCOO::new(
223            self.values.clone(),
224            self.cols.clone(),
225            self.rows.clone(),
226            vec![self.shape[1], self.shape[0]],
227        )
228    }
229    
230    /// Get sparsity ratio (fraction of zero elements)
231    pub fn sparsity(&self) -> f32 {
232        let total_elements: usize = self.shape.iter().product();
233        1.0 - (self.nnz as f32 / total_elements as f32)
234    }
235}
236
237/// Sparse tensor in CSR (Compressed Sparse Row) format
238#[derive(Debug, Clone)]
239pub struct SparseTensorCSR {
240    /// Non-zero values
241    pub values: Vec<f32>,
242    /// Column indices
243    pub col_indices: Vec<usize>,
244    /// Row pointers
245    pub row_ptrs: Vec<usize>,
246    /// Shape of the tensor
247    pub shape: Vec<usize>,
248    /// Number of non-zero elements
249    pub nnz: usize,
250}
251
252impl SparseTensorCSR {
253    /// Create a new sparse tensor in CSR format
254    pub fn new(values: Vec<f32>, col_indices: Vec<usize>, row_ptrs: Vec<usize>, shape: Vec<usize>) -> Result<Self> {
255        if values.len() != col_indices.len() {
256            return Err(GhostError::InvalidShape(
257                "Values and col_indices must have the same length".to_string()
258            ));
259        }
260        
261        let nnz = values.len();
262        
263        Ok(SparseTensorCSR {
264            values,
265            col_indices,
266            row_ptrs,
267            shape,
268            nnz,
269        })
270    }
271    
272    /// Convert from COO format
273    pub fn from_coo(coo: &SparseTensorCOO) -> Result<Self> {
274        if coo.shape.len() != 2 {
275            return Err(GhostError::InvalidShape(
276                "CSR only supports 2D tensors".to_string()
277            ));
278        }
279        
280        let nrows = coo.shape[0];
281        
282        // Sort by row, then column
283        let mut indices: Vec<usize> = (0..coo.nnz).collect();
284        indices.sort_by_key(|&i| (coo.rows[i], coo.cols[i]));
285        
286        let mut values = Vec::with_capacity(coo.nnz);
287        let mut col_indices = Vec::with_capacity(coo.nnz);
288        let mut row_ptrs = vec![0; nrows + 1];
289        
290        for &i in &indices {
291            values.push(coo.values[i]);
292            col_indices.push(coo.cols[i]);
293        }
294        
295        // Build row pointers
296        for &i in &indices {
297            row_ptrs[coo.rows[i] + 1] += 1;
298        }
299        
300        for i in 1..=nrows {
301            row_ptrs[i] += row_ptrs[i - 1];
302        }
303        
304        SparseTensorCSR::new(values, col_indices, row_ptrs, coo.shape.clone())
305    }
306    
307    /// Convert to COO format
308    pub fn to_coo(&self) -> Result<SparseTensorCOO> {
309        let mut values = Vec::with_capacity(self.nnz);
310        let mut rows = Vec::with_capacity(self.nnz);
311        let mut cols = Vec::with_capacity(self.nnz);
312        
313        for row in 0..self.shape[0] {
314            let start = self.row_ptrs[row];
315            let end = self.row_ptrs[row + 1];
316            
317            for i in start..end {
318                values.push(self.values[i]);
319                rows.push(row);
320                cols.push(self.col_indices[i]);
321            }
322        }
323        
324        SparseTensorCOO::new(values, rows, cols, self.shape.clone())
325    }
326    
327    /// Sparse matrix-vector multiplication (optimized for CSR)
328    pub fn spmv(&self, vec: &Tensor) -> Result<Tensor> {
329        let vec_data = vec.data_f32();
330        if vec_data.len() != self.shape[1] {
331            return Err(GhostError::InvalidShape(
332                format!("Vector length {} doesn't match matrix columns {}", vec_data.len(), self.shape[1])
333            ));
334        }
335        
336        let mut result = vec![0.0f32; self.shape[0]];
337        
338        for row in 0..self.shape[0] {
339            let start = self.row_ptrs[row];
340            let end = self.row_ptrs[row + 1];
341            
342            let mut sum = 0.0;
343            for i in start..end {
344                let col = self.col_indices[i];
345                let val = self.values[i];
346                sum += val * vec_data[col];
347            }
348            result[row] = sum;
349        }
350        
351        Tensor::from_slice(&result, &[self.shape[0]])
352    }
353}
354
355/// Sparse tensor in CSC (Compressed Sparse Column) format
356#[derive(Debug, Clone)]
357pub struct SparseTensorCSC {
358    /// Non-zero values
359    pub values: Vec<f32>,
360    /// Row indices
361    pub row_indices: Vec<usize>,
362    /// Column pointers
363    pub col_ptrs: Vec<usize>,
364    /// Shape of the tensor
365    pub shape: Vec<usize>,
366    /// Number of non-zero elements
367    pub nnz: usize,
368}
369
370impl SparseTensorCSC {
371    /// Create a new sparse tensor in CSC format
372    pub fn new(values: Vec<f32>, row_indices: Vec<usize>, col_ptrs: Vec<usize>, shape: Vec<usize>) -> Result<Self> {
373        if values.len() != row_indices.len() {
374            return Err(GhostError::InvalidShape(
375                "Values and row_indices must have the same length".to_string()
376            ));
377        }
378        
379        let nnz = values.len();
380        
381        Ok(SparseTensorCSC {
382            values,
383            row_indices,
384            col_ptrs,
385            shape,
386            nnz,
387        })
388    }
389    
390    /// Convert from COO format
391    pub fn from_coo(coo: &SparseTensorCOO) -> Result<Self> {
392        if coo.shape.len() != 2 {
393            return Err(GhostError::InvalidShape(
394                "CSC only supports 2D tensors".to_string()
395            ));
396        }
397        
398        let ncols = coo.shape[1];
399        
400        // Sort by column, then row
401        let mut indices: Vec<usize> = (0..coo.nnz).collect();
402        indices.sort_by_key(|&i| (coo.cols[i], coo.rows[i]));
403        
404        let mut values = Vec::with_capacity(coo.nnz);
405        let mut row_indices = Vec::with_capacity(coo.nnz);
406        let mut col_ptrs = vec![0; ncols + 1];
407        
408        for &i in &indices {
409            values.push(coo.values[i]);
410            row_indices.push(coo.rows[i]);
411        }
412        
413        // Build column pointers
414        for &i in &indices {
415            col_ptrs[coo.cols[i] + 1] += 1;
416        }
417        
418        for i in 1..=ncols {
419            col_ptrs[i] += col_ptrs[i - 1];
420        }
421        
422        SparseTensorCSC::new(values, row_indices, col_ptrs, coo.shape.clone())
423    }
424}
425
426#[cfg(test)]
427mod tests {
428    use super::*;
429    
430    #[test]
431    fn test_sparse_coo_creation() {
432        let values = vec![1.0, 2.0, 3.0];
433        let rows = vec![0, 1, 2];
434        let cols = vec![0, 1, 2];
435        let shape = vec![3, 3];
436        
437        let sparse = SparseTensorCOO::new(values, rows, cols, shape).unwrap();
438        assert_eq!(sparse.nnz, 3);
439        assert_eq!(sparse.sparsity(), 2.0 / 3.0);
440    }
441    
442    #[test]
443    fn test_sparse_from_dense() {
444        let data = vec![1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 3.0];
445        let dense = Tensor::from_slice(&data, &[3, 3]).unwrap();
446        
447        let sparse = SparseTensorCOO::from_dense(&dense, 0.5);
448        assert_eq!(sparse.nnz, 3);
449    }
450    
451    #[test]
452    fn test_sparse_to_dense() {
453        let values = vec![1.0, 2.0, 3.0];
454        let rows = vec![0, 1, 2];
455        let cols = vec![0, 1, 2];
456        let shape = vec![3, 3];
457        
458        let sparse = SparseTensorCOO::new(values, rows, cols, shape).unwrap();
459        let dense = sparse.to_dense().unwrap();
460        
461        assert_eq!(dense.dims(), &[3, 3]);
462        assert_eq!(dense.data_f32()[0], 1.0);
463        assert_eq!(dense.data_f32()[4], 2.0);
464        assert_eq!(dense.data_f32()[8], 3.0);
465    }
466    
467    #[test]
468    fn test_sparse_spmv() {
469        let values = vec![1.0, 2.0, 3.0];
470        let rows = vec![0, 1, 2];
471        let cols = vec![0, 1, 2];
472        let shape = vec![3, 3];
473        
474        let sparse = SparseTensorCOO::new(values, rows, cols, shape).unwrap();
475        let vec = Tensor::from_slice(&[1.0, 2.0, 3.0], &[3]).unwrap();
476        
477        let result = sparse.spmv(&vec).unwrap();
478        assert_eq!(result.data_f32(), vec![1.0, 4.0, 9.0]);
479    }
480    
481    #[test]
482    fn test_csr_conversion() {
483        let values = vec![1.0, 2.0, 3.0];
484        let rows = vec![0, 1, 2];
485        let cols = vec![0, 1, 2];
486        let shape = vec![3, 3];
487        
488        let coo = SparseTensorCOO::new(values, rows, cols, shape).unwrap();
489        let csr = SparseTensorCSR::from_coo(&coo).unwrap();
490        
491        assert_eq!(csr.nnz, 3);
492        assert_eq!(csr.row_ptrs.len(), 4);
493    }
494}
ghostflow_core/sparse.rs

ghostflow_core/
sparse.rs
