torsh_sparse/
scipy_sparse.rs

1//! SciPy sparse matrix interoperability
2//!
3//! This module provides conversion between ToRSh sparse tensors and SciPy sparse matrices,
4//! enabling seamless integration with Python scientific computing ecosystem.
5
6#[cfg(feature = "scipy")]
7use numpy::{PyArray1, PyReadonlyArray1};
8#[cfg(feature = "scipy")]
9use pyo3::prelude::*;
10#[cfg(feature = "scipy")]
11use pyo3::types::PyDict;
12#[cfg(feature = "scipy")]
13use pyo3::Bound;
14
15use crate::*;
16use std::collections::HashMap;
17use torsh_core::Result as TorshResult;
18
19/// Supported SciPy sparse matrix formats
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum ScipyFormat {
22    /// Compressed Sparse Row (csr_matrix)
23    Csr,
24    /// Compressed Sparse Column (csc_matrix)
25    Csc,
26    /// Coordinate format (coo_matrix)
27    Coo,
28    /// Block Sparse Row (bsr_matrix)
29    Bsr,
30    /// Diagonal format (dia_matrix)
31    Dia,
32}
33
34impl From<SparseFormat> for ScipyFormat {
35    fn from(format: SparseFormat) -> Self {
36        match format {
37            SparseFormat::Coo => ScipyFormat::Coo,
38            SparseFormat::Csr => ScipyFormat::Csr,
39            SparseFormat::Csc => ScipyFormat::Csc,
40            SparseFormat::Bsr => ScipyFormat::Bsr,
41            SparseFormat::Dia => ScipyFormat::Dia,
42            SparseFormat::Ell => ScipyFormat::Csr, // ELL -> CSR fallback
43            SparseFormat::Rle => ScipyFormat::Csr, // RLE -> CSR fallback
44            SparseFormat::Symmetric => ScipyFormat::Csr, // Symmetric -> CSR fallback
45            SparseFormat::Dsr => ScipyFormat::Csr, // DSR -> CSR fallback
46        }
47    }
48}
49
50impl From<ScipyFormat> for SparseFormat {
51    fn from(format: ScipyFormat) -> Self {
52        match format {
53            ScipyFormat::Coo => SparseFormat::Coo,
54            ScipyFormat::Csr => SparseFormat::Csr,
55            ScipyFormat::Csc => SparseFormat::Csc,
56            ScipyFormat::Bsr => SparseFormat::Bsr,
57            ScipyFormat::Dia => SparseFormat::Dia,
58        }
59    }
60}
61
62/// SciPy sparse matrix representation for data exchange
63#[derive(Debug, Clone)]
64pub struct ScipySparseData {
65    /// Matrix format
66    pub format: ScipyFormat,
67    /// Matrix shape (rows, cols)
68    pub shape: (usize, usize),
69    /// Data values
70    pub data: Vec<f64>,
71    /// Row indices (for COO and CSR) or column indices (for CSC)
72    pub indices: Vec<usize>,
73    /// Row pointers (for CSR) or column pointers (for CSC) or coordinate rows (for COO)
74    pub indptr_or_row: Vec<usize>,
75    /// Block size for BSR format
76    pub blocksize: Option<(usize, usize)>,
77    /// Number of diagonals for DIA format
78    pub diagonals: Option<Vec<i32>>,
79}
80
81impl ScipySparseData {
82    /// Create new SciPy sparse data
83    pub fn new(format: ScipyFormat, shape: (usize, usize)) -> Self {
84        Self {
85            format,
86            shape,
87            data: Vec::new(),
88            indices: Vec::new(),
89            indptr_or_row: Vec::new(),
90            blocksize: None,
91            diagonals: None,
92        }
93    }
94
95    /// Create from COO data
96    pub fn from_coo(
97        shape: (usize, usize),
98        row_indices: Vec<usize>,
99        col_indices: Vec<usize>,
100        values: Vec<f64>,
101    ) -> Self {
102        Self {
103            format: ScipyFormat::Coo,
104            shape,
105            data: values,
106            indices: col_indices,
107            indptr_or_row: row_indices,
108            blocksize: None,
109            diagonals: None,
110        }
111    }
112
113    /// Create from CSR data
114    pub fn from_csr(
115        shape: (usize, usize),
116        row_ptr: Vec<usize>,
117        col_indices: Vec<usize>,
118        values: Vec<f64>,
119    ) -> Self {
120        Self {
121            format: ScipyFormat::Csr,
122            shape,
123            data: values,
124            indices: col_indices,
125            indptr_or_row: row_ptr,
126            blocksize: None,
127            diagonals: None,
128        }
129    }
130
131    /// Create from CSC data
132    pub fn from_csc(
133        shape: (usize, usize),
134        col_ptr: Vec<usize>,
135        row_indices: Vec<usize>,
136        values: Vec<f64>,
137    ) -> Self {
138        Self {
139            format: ScipyFormat::Csc,
140            shape,
141            data: values,
142            indices: row_indices,
143            indptr_or_row: col_ptr,
144            blocksize: None,
145            diagonals: None,
146        }
147    }
148}
149
150/// SciPy sparse matrix integration utilities
151pub struct ScipySparseIntegration;
152
153impl ScipySparseIntegration {
154    /// Convert ToRSh sparse tensor to SciPy sparse data
155    pub fn to_scipy_data(sparse: &dyn SparseTensor) -> TorshResult<ScipySparseData> {
156        let shape = sparse.shape();
157        let (rows, cols) = (shape.dims()[0], shape.dims()[1]);
158
159        match sparse.format() {
160            SparseFormat::Coo => {
161                let coo = sparse.to_coo()?;
162                let triplets = coo.triplets();
163
164                let mut row_indices = Vec::new();
165                let mut col_indices = Vec::new();
166                let mut values = Vec::new();
167
168                for (row, col, val) in triplets {
169                    row_indices.push(row);
170                    col_indices.push(col);
171                    values.push(val as f64);
172                }
173
174                Ok(ScipySparseData::from_coo(
175                    (rows, cols),
176                    row_indices,
177                    col_indices,
178                    values,
179                ))
180            }
181            SparseFormat::Csr => {
182                let csr = sparse.to_csr()?;
183                let row_ptr = csr.row_ptr().to_vec();
184                let col_indices = csr.col_indices().to_vec();
185                let values = csr.values().iter().map(|&v| v as f64).collect();
186
187                Ok(ScipySparseData::from_csr(
188                    (rows, cols),
189                    row_ptr,
190                    col_indices,
191                    values,
192                ))
193            }
194            SparseFormat::Csc => {
195                let csc = sparse.to_csc()?;
196                let col_ptr = csc.col_ptr().to_vec();
197                let row_indices = csc.row_indices().to_vec();
198                let values = csc.values().iter().map(|&v| v as f64).collect();
199
200                Ok(ScipySparseData::from_csc(
201                    (rows, cols),
202                    col_ptr,
203                    row_indices,
204                    values,
205                ))
206            }
207            _ => {
208                // Convert other formats to COO first
209                let coo = sparse.to_coo()?;
210                Self::to_scipy_data(&coo)
211            }
212        }
213    }
214
215    /// Convert SciPy sparse data to ToRSh sparse tensor
216    pub fn from_scipy_data(
217        data: &ScipySparseData,
218    ) -> TorshResult<Box<dyn SparseTensor + Send + Sync>> {
219        let shape = Shape::new(vec![data.shape.0, data.shape.1]);
220
221        match data.format {
222            ScipyFormat::Coo => {
223                let mut rows = Vec::new();
224                let mut cols = Vec::new();
225                let mut values = Vec::new();
226
227                for i in 0..data.data.len() {
228                    rows.push(data.indptr_or_row[i]);
229                    cols.push(data.indices[i]);
230                    values.push(data.data[i] as f32);
231                }
232
233                let coo = CooTensor::new(rows, cols, values, shape)?;
234                Ok(Box::new(coo))
235            }
236            ScipyFormat::Csr => {
237                let row_ptr = &data.indptr_or_row;
238                let col_indices = &data.indices;
239                let values: Vec<f32> = data.data.iter().map(|&v| v as f32).collect();
240
241                // SciPy does not guarantee canonical (sorted, duplicate-free)
242                // CSR arrays, so normalise instead of asserting.
243                let csr = CsrTensor::from_unsorted_parts(
244                    row_ptr.clone(),
245                    col_indices.clone(),
246                    values,
247                    shape,
248                )?;
249
250                Ok(Box::new(csr))
251            }
252            ScipyFormat::Csc => {
253                let col_ptr = &data.indptr_or_row;
254                let row_indices = &data.indices;
255                let values: Vec<f32> = data.data.iter().map(|&v| v as f32).collect();
256
257                // SciPy does not guarantee canonical (sorted, duplicate-free)
258                // CSC arrays, so normalise instead of asserting.
259                let csc = CscTensor::from_unsorted_parts(
260                    col_ptr.clone(),
261                    row_indices.clone(),
262                    values,
263                    shape,
264                )?;
265
266                Ok(Box::new(csc))
267            }
268            _ => {
269                // Convert to COO first, then to target format
270                let coo_data = ScipySparseData {
271                    format: ScipyFormat::Coo,
272                    ..data.clone()
273                };
274                let coo = Self::from_scipy_data(&coo_data)?;
275                convert_sparse_format(coo.as_ref(), data.format.into())
276            }
277        }
278    }
279
280    /// Serialize sparse tensor to dictionary format compatible with SciPy
281    pub fn to_dict(sparse: &dyn SparseTensor) -> TorshResult<HashMap<String, Vec<f64>>> {
282        let scipy_data = Self::to_scipy_data(sparse)?;
283
284        let mut dict = HashMap::new();
285        dict.insert("data".to_string(), scipy_data.data);
286        dict.insert(
287            "indices".to_string(),
288            scipy_data.indices.iter().map(|&x| x as f64).collect(),
289        );
290        dict.insert(
291            "indptr".to_string(),
292            scipy_data.indptr_or_row.iter().map(|&x| x as f64).collect(),
293        );
294        dict.insert(
295            "shape".to_string(),
296            vec![scipy_data.shape.0 as f64, scipy_data.shape.1 as f64],
297        );
298
299        Ok(dict)
300    }
301
302    /// Generate Python code to create equivalent SciPy sparse matrix
303    pub fn to_python_code(sparse: &dyn SparseTensor, var_name: &str) -> TorshResult<String> {
304        let scipy_data = Self::to_scipy_data(sparse)?;
305        let format_name = match scipy_data.format {
306            ScipyFormat::Coo => "coo_matrix",
307            ScipyFormat::Csr => "csr_matrix",
308            ScipyFormat::Csc => "csc_matrix",
309            ScipyFormat::Bsr => "bsr_matrix",
310            ScipyFormat::Dia => "dia_matrix",
311        };
312
313        let mut code = String::new();
314        code.push_str("import numpy as np\n");
315        code.push_str("from scipy.sparse import ");
316        code.push_str(format_name);
317        code.push_str("\n\n");
318
319        match scipy_data.format {
320            ScipyFormat::Coo => {
321                code.push_str("# COO format data\n");
322                code.push_str(&format!("row = np.array({:?})\n", scipy_data.indptr_or_row));
323                code.push_str(&format!("col = np.array({:?})\n", scipy_data.indices));
324                code.push_str(&format!("data = np.array({:?})\n", scipy_data.data));
325                code.push_str(&format!("shape = {:?}\n", scipy_data.shape));
326                code.push_str(&format!(
327                    "{var_name} = {format_name}((data, (row, col)), shape=shape)\n"
328                ));
329            }
330            ScipyFormat::Csr | ScipyFormat::Csc => {
331                let ptr_name = "indptr";
332                code.push_str(&format!("# {} format data\n", format_name.to_uppercase()));
333                code.push_str(&format!("data = np.array({:?})\n", scipy_data.data));
334                code.push_str(&format!("indices = np.array({:?})\n", scipy_data.indices));
335                code.push_str(&format!(
336                    "{} = np.array({:?})\n",
337                    ptr_name, scipy_data.indptr_or_row
338                ));
339                code.push_str(&format!("shape = {:?}\n", scipy_data.shape));
340                code.push_str(&format!(
341                    "{var_name} = {format_name}((data, indices, {ptr_name}), shape=shape)\n"
342                ));
343            }
344            _ => {
345                // Fallback to COO for other formats
346                code.push_str(&format!(
347                    "# Note: {format_name} format converted to COO for compatibility\n"
348                ));
349                code.push_str(&format!("row = np.array({:?})\n", scipy_data.indptr_or_row));
350                code.push_str(&format!("col = np.array({:?})\n", scipy_data.indices));
351                code.push_str(&format!("data = np.array({:?})\n", scipy_data.data));
352                code.push_str(&format!("shape = {:?}\n", scipy_data.shape));
353                code.push_str(&format!(
354                    "{var_name} = coo_matrix((data, (row, col)), shape=shape)\n"
355                ));
356            }
357        }
358
359        Ok(code)
360    }
361}
362
363/// Python bindings for SciPy sparse integration (when scipy feature is enabled)
364#[cfg(feature = "scipy")]
365pub mod python_bindings {
366    use super::*;
367
368    /// Export sparse tensor to Python SciPy format
369    #[pyfunction]
370    pub fn torsh_to_scipy(
371        py: Python,
372        format: &str,
373        shape: (usize, usize),
374        data: Vec<f64>,
375        indices: Vec<usize>,
376        indptr: Vec<usize>,
377    ) -> PyResult<Py<PyAny>> {
378        let scipy = py.import("scipy.sparse")?;
379
380        let data_array = PyArray1::from_vec(py, data);
381        let indices_array = PyArray1::from_vec(py, indices);
382        let indptr_array = PyArray1::from_vec(py, indptr);
383
384        let args = (
385            data_array
386                .into_pyobject(py)
387                .expect("PyArray conversion should succeed"),
388            indices_array
389                .into_pyobject(py)
390                .expect("PyArray conversion should succeed"),
391            indptr_array
392                .into_pyobject(py)
393                .expect("PyArray conversion should succeed"),
394        );
395
396        let kwargs = PyDict::new(py);
397        kwargs.set_item("shape", shape)?;
398
399        let matrix_class = scipy.getattr(format)?;
400        let result = matrix_class.call(args, Some(&kwargs))?;
401
402        Ok(result.unbind())
403    }
404
405    /// Import sparse tensor from Python SciPy format
406    #[pyfunction]
407    pub fn scipy_to_torsh(
408        _py: Python,
409        scipy_matrix: &Bound<PyAny>,
410    ) -> PyResult<(String, (usize, usize), Vec<f64>, Vec<usize>, Vec<usize>)> {
411        // Get format
412        let format_attr = scipy_matrix.getattr("format")?;
413        let format: String = format_attr.extract()?;
414
415        // Get shape
416        let shape_attr = scipy_matrix.getattr("shape")?;
417        let shape: (usize, usize) = shape_attr.extract()?;
418
419        // Convert to COO format for universal handling
420        let coo_matrix = scipy_matrix.call_method0("tocoo")?;
421
422        // Extract data arrays
423        let data_attr = coo_matrix.getattr("data")?;
424        let row_attr = coo_matrix.getattr("row")?;
425        let col_attr = coo_matrix.getattr("col")?;
426
427        let data: PyReadonlyArray1<f64> = data_attr.extract()?;
428        let row: PyReadonlyArray1<i32> = row_attr.extract()?;
429        let col: PyReadonlyArray1<i32> = col_attr.extract()?;
430
431        let data_vec = data.as_slice()?.to_vec();
432        let row_vec: Vec<usize> = row.as_slice()?.iter().map(|&x| x as usize).collect();
433        let col_vec: Vec<usize> = col.as_slice()?.iter().map(|&x| x as usize).collect();
434
435        Ok((format, shape, data_vec, col_vec, row_vec))
436    }
437}
438
439/// Macro to simplify SciPy conversion
440#[macro_export]
441macro_rules! to_scipy {
442    ($sparse:expr) => {
443        ScipySparseIntegration::to_scipy_data($sparse)
444    };
445    ($sparse:expr, $format:expr) => {{
446        let scipy_data = ScipySparseIntegration::to_scipy_data($sparse)?;
447        let converted = convert_sparse_format($sparse, $format)?;
448        ScipySparseIntegration::to_scipy_data(converted.as_ref())
449    }};
450}
451
452/// Macro to simplify creation from SciPy data
453#[macro_export]
454macro_rules! from_scipy {
455    ($data:expr) => {
456        ScipySparseIntegration::from_scipy_data($data)
457    };
458}
459
460#[cfg(test)]
461mod tests {
462    use super::*;
463    use crate::coo::CooTensor;
464    use torsh_core::{DType, Shape};
465
466    #[test]
467    fn test_scipy_data_conversion() {
468        let shape = Shape::new(vec![3, 3]);
469        let mut coo = CooTensor::empty(shape.clone(), DType::F32).unwrap();
470
471        // Create a simple sparse matrix
472        coo.insert(0, 0, 1.0).unwrap();
473        coo.insert(1, 1, 2.0).unwrap();
474        coo.insert(2, 2, 3.0).unwrap();
475        coo.insert(0, 2, 4.0).unwrap();
476
477        // Convert to SciPy data
478        let scipy_data = ScipySparseIntegration::to_scipy_data(&coo).unwrap();
479
480        assert_eq!(scipy_data.format, ScipyFormat::Coo);
481        assert_eq!(scipy_data.shape, (3, 3));
482        assert_eq!(scipy_data.data.len(), 4);
483
484        // Convert back to ToRSh
485        let restored = ScipySparseIntegration::from_scipy_data(&scipy_data).unwrap();
486        assert_eq!(restored.nnz(), 4);
487        assert_eq!(restored.shape(), &shape);
488    }
489
490    #[test]
491    fn test_python_code_generation() {
492        let shape = Shape::new(vec![2, 2]);
493        let mut coo = CooTensor::empty(shape, DType::F32).unwrap();
494
495        coo.insert(0, 0, 1.0).unwrap();
496        coo.insert(1, 1, 2.0).unwrap();
497
498        let code = ScipySparseIntegration::to_python_code(&coo, "matrix").unwrap();
499
500        assert!(code.contains("import numpy as np"));
501        assert!(code.contains("from scipy.sparse import"));
502        assert!(code.contains("matrix ="));
503    }
504
505    #[test]
506    fn test_dict_conversion() {
507        let shape = Shape::new(vec![2, 2]);
508        let mut coo = CooTensor::empty(shape, DType::F32).unwrap();
509
510        coo.insert(0, 0, 1.0).unwrap();
511        coo.insert(1, 1, 2.0).unwrap();
512
513        let dict = ScipySparseIntegration::to_dict(&coo).unwrap();
514
515        assert!(dict.contains_key("data"));
516        assert!(dict.contains_key("indices"));
517        assert!(dict.contains_key("indptr"));
518        assert!(dict.contains_key("shape"));
519
520        assert_eq!(dict["shape"], vec![2.0, 2.0]);
521        assert_eq!(dict["data"].len(), 2);
522    }
523
524    #[test]
525    fn test_format_conversion() {
526        assert_eq!(ScipyFormat::from(SparseFormat::Coo), ScipyFormat::Coo);
527        assert_eq!(ScipyFormat::from(SparseFormat::Csr), ScipyFormat::Csr);
528        assert_eq!(ScipyFormat::from(SparseFormat::Csc), ScipyFormat::Csc);
529        assert_eq!(ScipyFormat::from(SparseFormat::Ell), ScipyFormat::Csr);
530
531        assert_eq!(SparseFormat::from(ScipyFormat::Coo), SparseFormat::Coo);
532        assert_eq!(SparseFormat::from(ScipyFormat::Csr), SparseFormat::Csr);
533        assert_eq!(SparseFormat::from(ScipyFormat::Csc), SparseFormat::Csc);
534    }
535
536    #[test]
537    fn test_macro_usage() {
538        let shape = Shape::new(vec![2, 2]);
539        let mut coo = CooTensor::empty(shape, DType::F32).unwrap();
540
541        coo.insert(0, 0, 1.0).unwrap();
542        coo.insert(1, 1, 2.0).unwrap();
543
544        let scipy_data = to_scipy!(&coo).unwrap();
545        assert_eq!(scipy_data.data.len(), 2);
546
547        let restored = from_scipy!(&scipy_data).unwrap();
548        assert_eq!(restored.nnz(), 2);
549    }
550}
torsh_sparse/scipy_sparse.rs

torsh_sparse/
scipy_sparse.rs
