oxiphysics_gpu/
sparse_gpu.rs

1// Copyright 2026 COOLJAPAN OU (Team KitaSan)
2// SPDX-License-Identifier: Apache-2.0
3
4//! GPU-ready sparse matrix formats and iterative solvers with data-oriented layouts.
5//!
6//! Provides CSR, ELLPACK, Hybrid (ELL+COO), and Block-CSR formats along with
7//! iterative solvers (CG, BiCGSTAB, preconditioned CG) suitable for GPU offload.
8
9// ---------------------------------------------------------------------------
10// Sparse vector operations
11// ---------------------------------------------------------------------------
12
13/// Dot product of two vectors.
14pub fn dot(x: &[f64], y: &[f64]) -> f64 {
15    x.iter().zip(y.iter()).map(|(a, b)| a * b).sum()
16}
17
18/// `y + alpha * x` (AXPY).
19pub fn axpy(alpha: f64, x: &[f64], y: &[f64]) -> Vec<f64> {
20    y.iter()
21        .zip(x.iter())
22        .map(|(yi, xi)| yi + alpha * xi)
23        .collect()
24}
25
26/// Euclidean norm of a vector.
27pub fn norm2(x: &[f64]) -> f64 {
28    dot(x, x).sqrt()
29}
30
31/// Scale every element of `x` by `s`.
32pub fn scale_vec(x: &[f64], s: f64) -> Vec<f64> {
33    x.iter().map(|v| v * s).collect()
34}
35
36// ---------------------------------------------------------------------------
37// SparseTriplet – coordinate (COO) format for assembly
38// ---------------------------------------------------------------------------
39
40/// Coordinate (COO) format sparse matrix for incremental assembly.
41pub struct SparseTriplet {
42    /// Row indices of non-zero entries.
43    pub rows: Vec<usize>,
44    /// Column indices of non-zero entries.
45    pub cols: Vec<usize>,
46    /// Values of non-zero entries.
47    pub vals: Vec<f64>,
48}
49
50impl SparseTriplet {
51    /// Create an empty triplet store.
52    pub fn new() -> Self {
53        Self {
54            rows: Vec::new(),
55            cols: Vec::new(),
56            vals: Vec::new(),
57        }
58    }
59
60    /// Push a single entry `(row, col, val)`.
61    pub fn add(&mut self, row: usize, col: usize, val: f64) {
62        self.rows.push(row);
63        self.cols.push(col);
64        self.vals.push(val);
65    }
66
67    /// Convert to [`CsrMatrix`], sorting by (row, col) and summing duplicates.
68    pub fn to_csr(&self, n_rows: usize, n_cols: usize) -> CsrMatrix {
69        // Sort indices by (row, col)
70        let mut order: Vec<usize> = (0..self.rows.len()).collect();
71        order.sort_by_key(|&i| (self.rows[i], self.cols[i]));
72
73        // Accumulate into (row, col, val) triples, summing duplicates
74        let mut entries: Vec<(usize, usize, f64)> = Vec::new();
75        for &i in &order {
76            let r = self.rows[i];
77            let c = self.cols[i];
78            let v = self.vals[i];
79            if let Some(last) = entries.last_mut()
80                && last.0 == r
81                && last.1 == c
82            {
83                last.2 += v;
84                continue;
85            }
86            entries.push((r, c, v));
87        }
88
89        // Build CSR
90        let nnz = entries.len();
91        let mut row_ptr = vec![0usize; n_rows + 1];
92        let mut col_idx = Vec::with_capacity(nnz);
93        let mut values = Vec::with_capacity(nnz);
94
95        for &(r, c, v) in &entries {
96            row_ptr[r + 1] += 1;
97            col_idx.push(c);
98            values.push(v);
99        }
100        for i in 0..n_rows {
101            row_ptr[i + 1] += row_ptr[i];
102        }
103
104        CsrMatrix {
105            n_rows,
106            n_cols,
107            row_ptr,
108            col_idx,
109            values,
110        }
111    }
112}
113
114impl Default for SparseTriplet {
115    fn default() -> Self {
116        Self::new()
117    }
118}
119
120// ---------------------------------------------------------------------------
121// CsrMatrix – Compressed Sparse Row
122// ---------------------------------------------------------------------------
123
124/// Compressed Sparse Row (CSR) matrix stored as plain `f64` arrays.
125pub struct CsrMatrix {
126    /// Number of rows.
127    pub n_rows: usize,
128    /// Number of columns.
129    pub n_cols: usize,
130    /// Row start indices, length `n_rows + 1`.
131    pub row_ptr: Vec<usize>,
132    /// Column indices of non-zeros.
133    pub col_idx: Vec<usize>,
134    /// Non-zero values.
135    pub values: Vec<f64>,
136}
137
138impl CsrMatrix {
139    /// Create an empty (all-zero) CSR matrix.
140    pub fn new(n_rows: usize, n_cols: usize) -> Self {
141        Self {
142            n_rows,
143            n_cols,
144            row_ptr: vec![0; n_rows + 1],
145            col_idx: Vec::new(),
146            values: Vec::new(),
147        }
148    }
149
150    /// Build from a dense row-major matrix (rows of length `n_cols`).
151    pub fn from_dense(m: &[Vec<f64>]) -> Self {
152        let n_rows = m.len();
153        let n_cols = if n_rows > 0 { m[0].len() } else { 0 };
154        let mut row_ptr = vec![0usize; n_rows + 1];
155        let mut col_idx = Vec::new();
156        let mut values = Vec::new();
157        for (r, row) in m.iter().enumerate() {
158            for (c, &v) in row.iter().enumerate() {
159                if v != 0.0 {
160                    col_idx.push(c);
161                    values.push(v);
162                }
163            }
164            row_ptr[r + 1] = col_idx.len();
165        }
166        Self {
167            n_rows,
168            n_cols,
169            row_ptr,
170            col_idx,
171            values,
172        }
173    }
174
175    /// Number of stored non-zeros.
176    pub fn nnz(&self) -> usize {
177        self.values.len()
178    }
179
180    /// Sparse matrix–vector product `y = A * x`.
181    pub fn spmv(&self, x: &[f64]) -> Vec<f64> {
182        let mut y = vec![0.0f64; self.n_rows];
183        for (r, yr) in y.iter_mut().enumerate() {
184            let start = self.row_ptr[r];
185            let end = self.row_ptr[r + 1];
186            let mut sum = 0.0;
187            for k in start..end {
188                sum += self.values[k] * x[self.col_idx[k]];
189            }
190            *yr = sum;
191        }
192        y
193    }
194
195    /// Get the value at `(row, col)`, returning `0.0` if not stored.
196    pub fn get(&self, row: usize, col: usize) -> f64 {
197        let start = self.row_ptr[row];
198        let end = self.row_ptr[row + 1];
199        for k in start..end {
200            if self.col_idx[k] == col {
201                return self.values[k];
202            }
203        }
204        0.0
205    }
206
207    /// Return `A^T` as a new `CsrMatrix`.
208    pub fn transpose(&self) -> CsrMatrix {
209        // Count non-zeros per column (which become rows of A^T)
210        let mut row_ptr = vec![0usize; self.n_cols + 1];
211        for &c in &self.col_idx {
212            row_ptr[c + 1] += 1;
213        }
214        for i in 0..self.n_cols {
215            row_ptr[i + 1] += row_ptr[i];
216        }
217
218        let nnz = self.values.len();
219        let mut col_idx = vec![0usize; nnz];
220        let mut values = vec![0.0f64; nnz];
221        let mut pos = row_ptr[..self.n_cols].to_vec();
222
223        for r in 0..self.n_rows {
224            let start = self.row_ptr[r];
225            let end = self.row_ptr[r + 1];
226            for k in start..end {
227                let c = self.col_idx[k];
228                let dest = pos[c];
229                col_idx[dest] = r;
230                values[dest] = self.values[k];
231                pos[c] += 1;
232            }
233        }
234
235        CsrMatrix {
236            n_rows: self.n_cols,
237            n_cols: self.n_rows,
238            row_ptr,
239            col_idx,
240            values,
241        }
242    }
243
244    /// Add two CSR matrices of identical shape.
245    pub fn add(&self, other: &CsrMatrix) -> CsrMatrix {
246        assert_eq!(self.n_rows, other.n_rows);
247        assert_eq!(self.n_cols, other.n_cols);
248        // Assemble via triplet
249        let mut trip = SparseTriplet::new();
250        for r in 0..self.n_rows {
251            for k in self.row_ptr[r]..self.row_ptr[r + 1] {
252                trip.add(r, self.col_idx[k], self.values[k]);
253            }
254            for k in other.row_ptr[r]..other.row_ptr[r + 1] {
255                trip.add(r, other.col_idx[k], other.values[k]);
256            }
257        }
258        trip.to_csr(self.n_rows, self.n_cols)
259    }
260
261    /// Return a new matrix with every value multiplied by `s`.
262    pub fn scale(&self, s: f64) -> CsrMatrix {
263        CsrMatrix {
264            n_rows: self.n_rows,
265            n_cols: self.n_cols,
266            row_ptr: self.row_ptr.clone(),
267            col_idx: self.col_idx.clone(),
268            values: self.values.iter().map(|v| v * s).collect(),
269        }
270    }
271}
272
273// ---------------------------------------------------------------------------
274// EllMatrix – ELLPACK / ITPACK format (GPU-friendly padded layout)
275// ---------------------------------------------------------------------------
276
277/// ELLPACK-format sparse matrix: rows padded to `max_nnz_per_row`.
278pub struct EllMatrix {
279    /// Number of rows.
280    pub n_rows: usize,
281    /// Number of columns.
282    pub n_cols: usize,
283    /// Maximum number of non-zeros per row (padding width).
284    pub max_nnz_per_row: usize,
285    /// Column indices, row-major `[n_rows × max_nnz_per_row]`.
286    pub col_idx: Vec<usize>,
287    /// Values, row-major `[n_rows × max_nnz_per_row]`.
288    pub values: Vec<f64>,
289}
290
291impl EllMatrix {
292    /// Convert a [`CsrMatrix`] to ELLPACK.
293    pub fn from_csr(csr: &CsrMatrix) -> Self {
294        let n_rows = csr.n_rows;
295        let n_cols = csr.n_cols;
296        let max_nnz_per_row = (0..n_rows)
297            .map(|r| csr.row_ptr[r + 1] - csr.row_ptr[r])
298            .max()
299            .unwrap_or(0);
300
301        let size = n_rows * max_nnz_per_row;
302        let mut col_idx = vec![0usize; size];
303        let mut values = vec![0.0f64; size];
304
305        for r in 0..n_rows {
306            let start = csr.row_ptr[r];
307            let end = csr.row_ptr[r + 1];
308            for (j, k) in (start..end).enumerate() {
309                col_idx[r * max_nnz_per_row + j] = csr.col_idx[k];
310                values[r * max_nnz_per_row + j] = csr.values[k];
311            }
312        }
313
314        Self {
315            n_rows,
316            n_cols,
317            max_nnz_per_row,
318            col_idx,
319            values,
320        }
321    }
322
323    /// Sparse matrix–vector product `y = A * x`.
324    pub fn spmv(&self, x: &[f64]) -> Vec<f64> {
325        let mut y = vec![0.0f64; self.n_rows];
326        for (r, yr) in y.iter_mut().enumerate() {
327            let mut sum = 0.0;
328            for j in 0..self.max_nnz_per_row {
329                let v = self.values[r * self.max_nnz_per_row + j];
330                if v != 0.0 {
331                    let c = self.col_idx[r * self.max_nnz_per_row + j];
332                    sum += v * x[c];
333                }
334            }
335            *yr = sum;
336        }
337        y
338    }
339}
340
341// ---------------------------------------------------------------------------
342// HybridMatrix – ELL + COO for irregular sparsity
343// ---------------------------------------------------------------------------
344
345/// Hybrid ELL+COO matrix: regular rows stored in ELL, overflow in COO.
346pub struct HybridMatrix {
347    /// Regular (padded) portion stored in ELLPACK format.
348    pub ell: EllMatrix,
349    /// Row indices of COO overflow entries.
350    pub coo_row: Vec<usize>,
351    /// Column indices of COO overflow entries.
352    pub coo_col: Vec<usize>,
353    /// Values of COO overflow entries.
354    pub coo_val: Vec<f64>,
355}
356
357impl HybridMatrix {
358    /// Sparse matrix–vector product `y = (ELL + COO) * x`.
359    pub fn spmv(&self, x: &[f64]) -> Vec<f64> {
360        let mut y = self.ell.spmv(x);
361        for k in 0..self.coo_val.len() {
362            y[self.coo_row[k]] += self.coo_val[k] * x[self.coo_col[k]];
363        }
364        y
365    }
366}
367
368// ---------------------------------------------------------------------------
369// BlockCsrMatrix – block sparse row for FEM
370// ---------------------------------------------------------------------------
371
372/// Block-CSR matrix where every stored entry is a `block_size × block_size` dense tile.
373pub struct BlockCsrMatrix {
374    /// Size of each dense square block.
375    pub block_size: usize,
376    /// Number of block rows.
377    pub n_block_rows: usize,
378    /// Number of block columns.
379    pub n_block_cols: usize,
380    /// Block row start indices, length `n_block_rows + 1`.
381    pub row_ptr: Vec<usize>,
382    /// Column (block) indices.
383    pub col_idx: Vec<usize>,
384    /// Dense tiles, each of length `block_size * block_size`.
385    pub blocks: Vec<Vec<f64>>,
386}
387
388impl BlockCsrMatrix {
389    /// Sparse matrix–vector product treating the matrix as `(n_block_rows * block_size)` × `(n_block_cols * block_size)`.
390    pub fn spmv_block(&self, x: &[f64]) -> Vec<f64> {
391        let bs = self.block_size;
392        let n = self.n_block_rows * bs;
393        let mut y = vec![0.0f64; n];
394        for br in 0..self.n_block_rows {
395            let row_start = self.row_ptr[br];
396            let row_end = self.row_ptr[br + 1];
397            for k in row_start..row_end {
398                let bc = self.col_idx[k];
399                let blk = &self.blocks[k];
400                // Multiply dense block into y
401                for i in 0..bs {
402                    let mut s = 0.0;
403                    for j in 0..bs {
404                        s += blk[i * bs + j] * x[bc * bs + j];
405                    }
406                    y[br * bs + i] += s;
407                }
408            }
409        }
410        y
411    }
412
413    /// Convert to a flat [`CsrMatrix`].
414    pub fn to_csr(&self) -> CsrMatrix {
415        let bs = self.block_size;
416        let n_rows = self.n_block_rows * bs;
417        let n_cols = self.n_block_cols * bs;
418        let mut trip = SparseTriplet::new();
419        for br in 0..self.n_block_rows {
420            for k in self.row_ptr[br]..self.row_ptr[br + 1] {
421                let bc = self.col_idx[k];
422                let blk = &self.blocks[k];
423                for i in 0..bs {
424                    for j in 0..bs {
425                        let v = blk[i * bs + j];
426                        if v != 0.0 {
427                            trip.add(br * bs + i, bc * bs + j, v);
428                        }
429                    }
430                }
431            }
432        }
433        trip.to_csr(n_rows, n_cols)
434    }
435}
436
437// ---------------------------------------------------------------------------
438// Iterative solvers
439// ---------------------------------------------------------------------------
440
441/// Conjugate Gradient solver for symmetric positive-definite systems `A x = b`.
442///
443/// Returns `(solution, iterations_used)`.
444pub fn cg_solve(
445    a: &CsrMatrix,
446    b: &[f64],
447    x0: &[f64],
448    max_iter: usize,
449    tol: f64,
450) -> (Vec<f64>, usize) {
451    let n = b.len();
452    let mut x = x0.to_vec();
453    // r = b - A*x
454    let ax = a.spmv(&x);
455    let mut r: Vec<f64> = (0..n).map(|i| b[i] - ax[i]).collect();
456    let mut p = r.clone();
457    let mut rs_old = dot(&r, &r);
458
459    for iter in 0..max_iter {
460        if rs_old.sqrt() < tol {
461            return (x, iter);
462        }
463        let ap = a.spmv(&p);
464        let alpha = rs_old / dot(&p, &ap);
465        x = axpy(alpha, &p, &x);
466        r = axpy(-alpha, &ap, &r);
467        let rs_new = dot(&r, &r);
468        let beta = rs_new / rs_old;
469        p = axpy(beta, &p, &r);
470        rs_old = rs_new;
471    }
472    (x, max_iter)
473}
474
475/// BiCGSTAB solver for general (possibly non-symmetric) systems `A x = b`.
476///
477/// Returns `(solution, iterations_used)`.
478pub fn bicgstab_solve(
479    a: &CsrMatrix,
480    b: &[f64],
481    x0: &[f64],
482    max_iter: usize,
483    tol: f64,
484) -> (Vec<f64>, usize) {
485    let n = b.len();
486    let mut x = x0.to_vec();
487    let ax = a.spmv(&x);
488    let mut r: Vec<f64> = (0..n).map(|i| b[i] - ax[i]).collect();
489    let r_hat = r.clone();
490    let mut rho = 1.0_f64;
491    let mut alpha_s = 1.0_f64;
492    let mut omega = 1.0_f64;
493    let mut v = vec![0.0f64; n];
494    let mut p = vec![0.0f64; n];
495
496    for iter in 0..max_iter {
497        if norm2(&r) < tol {
498            return (x, iter);
499        }
500        let rho_new = dot(&r_hat, &r);
501        let beta = (rho_new / rho) * (alpha_s / omega);
502        p = axpy(beta, &p, &axpy(-beta * omega, &v, &r));
503        v = a.spmv(&p);
504        let denom = dot(&r_hat, &v);
505        if denom.abs() < 1e-300 {
506            return (x, iter);
507        }
508        alpha_s = rho_new / denom;
509        let s: Vec<f64> = axpy(-alpha_s, &v, &r);
510        if norm2(&s) < tol {
511            x = axpy(alpha_s, &p, &x);
512            return (x, iter + 1);
513        }
514        let t = a.spmv(&s);
515        let tt = dot(&t, &t);
516        omega = if tt.abs() < 1e-300 {
517            0.0
518        } else {
519            dot(&t, &s) / tt
520        };
521        x = axpy(omega, &s, &axpy(alpha_s, &p, &x));
522        r = axpy(-omega, &t, &s);
523        rho = rho_new;
524    }
525    (x, max_iter)
526}
527
528/// Conjugate Gradient with Jacobi (diagonal) preconditioner for `A x = b` (SPD).
529///
530/// Returns `(solution, iterations_used)`.
531pub fn jacobi_preconditioned_cg(
532    a: &CsrMatrix,
533    b: &[f64],
534    max_iter: usize,
535    tol: f64,
536) -> (Vec<f64>, usize) {
537    let n = b.len();
538    // Build inverse diagonal preconditioner M^{-1}
539    let mut m_inv = vec![1.0f64; n];
540    for (r, m) in m_inv.iter_mut().enumerate() {
541        let d = a.get(r, r);
542        if d.abs() > 1e-300 {
543            *m = 1.0 / d;
544        }
545    }
546
547    let mut x = vec![0.0f64; n];
548    let ax = a.spmv(&x);
549    let mut r: Vec<f64> = b
550        .iter()
551        .zip(ax.iter())
552        .map(|(&bi, &axi)| bi - axi)
553        .collect();
554    // z = M^{-1} r
555    let z: Vec<f64> = (0..n).map(|i| m_inv[i] * r[i]).collect();
556    let mut p = z.clone();
557    let mut rz_old = dot(&r, &z);
558
559    for iter in 0..max_iter {
560        if norm2(&r) < tol {
561            return (x, iter);
562        }
563        let ap = a.spmv(&p);
564        let alpha = rz_old / dot(&p, &ap);
565        x = axpy(alpha, &p, &x);
566        r = axpy(-alpha, &ap, &r);
567        let z_new: Vec<f64> = (0..n).map(|i| m_inv[i] * r[i]).collect();
568        let rz_new = dot(&r, &z_new);
569        let beta = rz_new / rz_old;
570        p = axpy(beta, &p, &z_new);
571        rz_old = rz_new;
572    }
573    (x, max_iter)
574}
575
576// ---------------------------------------------------------------------------
577// GPU simulation utilities
578// ---------------------------------------------------------------------------
579
580/// Estimate SpMV throughput in GFLOPS given matrix dimensions and nnz count.
581///
582/// Uses a simple roofline model: bandwidth-bound at 100 GB/s with 12 bytes/flop.
583pub fn simulate_spmv_throughput(n: usize, nnz: usize) -> f64 {
584    // Memory traffic estimate: col_idx (8 bytes) + values (8 bytes) + x access (~random, 8 bytes)
585    // = 24 bytes per nnz, plus row_ptr read (n * 8 bytes)
586    let _ = n; // size used for context but dominated by nnz traffic
587    let bytes_transferred = (nnz * 24) as f64;
588    let bandwidth_gb_s = 100.0_f64; // typical GPU HBM bandwidth
589    let time_s = bytes_transferred / (bandwidth_gb_s * 1e9);
590    let flops = 2.0 * nnz as f64; // one multiply + one add per nnz
591    flops / time_s / 1e9 // GFLOPS
592}
593
594/// Choose the ELLPACK row width (max non-zeros per row) to minimize padding waste.
595///
596/// Uses the 75th-percentile of the per-row nnz distribution.
597pub fn optimal_ell_row_width(nnz_distribution: &[usize]) -> usize {
598    if nnz_distribution.is_empty() {
599        return 0;
600    }
601    let mut sorted = nnz_distribution.to_vec();
602    sorted.sort_unstable();
603    let idx = (sorted.len() * 3) / 4; // 75th percentile index
604    sorted[idx]
605}
606
607// ---------------------------------------------------------------------------
608// SpMV – segmented (row-parallel) variant
609// ---------------------------------------------------------------------------
610
611/// Segmented SpMV: processes each row independently to prepare for
612/// GPU-style parallel execution. Functionally identical to `CsrMatrix::spmv`
613/// but structured for row-parallel dispatch.
614pub fn spmv_segmented(a: &CsrMatrix, x: &[f64]) -> Vec<f64> {
615    let mut y = vec![0.0_f64; a.n_rows];
616    for (r, yr) in y.iter_mut().enumerate() {
617        let start = a.row_ptr[r];
618        let end = a.row_ptr[r + 1];
619        let mut acc = 0.0_f64;
620        for k in start..end {
621            acc += a.values[k] * x[a.col_idx[k]];
622        }
623        *yr = acc;
624    }
625    y
626}
627
628// ---------------------------------------------------------------------------
629// Sparse matrix assembly helpers
630// ---------------------------------------------------------------------------
631
632/// Assemble a 1D Laplacian matrix of size `n × n` (tridiagonal: 2 on diag, -1 off-diag).
633pub fn assemble_1d_laplacian(n: usize) -> CsrMatrix {
634    let mut trip = SparseTriplet::new();
635    for i in 0..n {
636        trip.add(i, i, 2.0);
637        if i > 0 {
638            trip.add(i, i - 1, -1.0);
639        }
640        if i + 1 < n {
641            trip.add(i, i + 1, -1.0);
642        }
643    }
644    trip.to_csr(n, n)
645}
646
647// ---------------------------------------------------------------------------
648// CSR-to-ELL conversion (alternative entry point)
649// ---------------------------------------------------------------------------
650
651/// Convert a CSR matrix to ELLPACK format (convenience wrapper).
652pub fn csr_to_ell(csr: &CsrMatrix) -> EllMatrix {
653    EllMatrix::from_csr(csr)
654}
655
656// ---------------------------------------------------------------------------
657// Extract diagonal
658// ---------------------------------------------------------------------------
659
660/// Extract the main diagonal of a CSR matrix.
661pub fn extract_diagonal(a: &CsrMatrix) -> Vec<f64> {
662    let n = a.n_rows.min(a.n_cols);
663    let mut diag = vec![0.0_f64; n];
664    for (r, d) in diag.iter_mut().enumerate() {
665        *d = a.get(r, r);
666    }
667    diag
668}
669
670// ---------------------------------------------------------------------------
671// Per-row nnz distribution
672// ---------------------------------------------------------------------------
673
674/// Compute the number of non-zeros per row for a CSR matrix.
675pub fn compute_nnz_per_row(a: &CsrMatrix) -> Vec<usize> {
676    (0..a.n_rows)
677        .map(|r| a.row_ptr[r + 1] - a.row_ptr[r])
678        .collect()
679}
680
681// ---------------------------------------------------------------------------
682// Frobenius norm
683// ---------------------------------------------------------------------------
684
685/// Compute the Frobenius norm of a sparse matrix: `sqrt(sum(a_ij^2))`.
686pub fn frobenius_norm(a: &CsrMatrix) -> f64 {
687    let sum_sq: f64 = a.values.iter().map(|v| v * v).sum();
688    sum_sq.sqrt()
689}
690
691// ---------------------------------------------------------------------------
692// Sparse triangular solves
693// ---------------------------------------------------------------------------
694
695/// Forward-substitution solve `L x = b` where `L` is lower-triangular (CSR).
696///
697/// Assumes `L` has non-zero diagonal entries. The diagonal entry for row `i`
698/// is taken as `L.get(i, i)`.
699pub fn sparse_lower_triangular_solve(l: &CsrMatrix, b: &[f64]) -> Vec<f64> {
700    let n = b.len();
701    let mut x = vec![0.0_f64; n];
702    for i in 0..n {
703        let mut sum = b[i];
704        let start = l.row_ptr[i];
705        let end = l.row_ptr[i + 1];
706        let mut diag = 1.0_f64;
707        for k in start..end {
708            let c = l.col_idx[k];
709            if c < i {
710                sum -= l.values[k] * x[c];
711            } else if c == i {
712                diag = l.values[k];
713            }
714        }
715        x[i] = sum / diag;
716    }
717    x
718}
719
720/// Back-substitution solve `U x = b` where `U` is upper-triangular (CSR).
721///
722/// Assumes `U` has non-zero diagonal entries.
723pub fn sparse_upper_triangular_solve(u: &CsrMatrix, b: &[f64]) -> Vec<f64> {
724    let n = b.len();
725    let mut x = vec![0.0_f64; n];
726    for i in (0..n).rev() {
727        let mut sum = b[i];
728        let start = u.row_ptr[i];
729        let end = u.row_ptr[i + 1];
730        let mut diag = 1.0_f64;
731        for k in start..end {
732            let c = u.col_idx[k];
733            if c > i {
734                sum -= u.values[k] * x[c];
735            } else if c == i {
736                diag = u.values[k];
737            }
738        }
739        x[i] = sum / diag;
740    }
741    x
742}
743
744// ---------------------------------------------------------------------------
745// Tests
746// ---------------------------------------------------------------------------
747
748#[cfg(test)]
749mod tests {
750    use super::*;
751
752    #[test]
753    fn test_csr_from_dense_nnz() {
754        let m = vec![
755            vec![1.0, 0.0, 2.0],
756            vec![0.0, 3.0, 0.0],
757            vec![4.0, 5.0, 6.0],
758        ];
759        let csr = CsrMatrix::from_dense(&m);
760        assert_eq!(csr.n_rows, 3);
761        assert_eq!(csr.n_cols, 3);
762        assert_eq!(csr.nnz(), 6);
763    }
764
765    #[test]
766    fn test_csr_spmv_identity() {
767        // 3x3 identity
768        let m = vec![
769            vec![1.0, 0.0, 0.0],
770            vec![0.0, 1.0, 0.0],
771            vec![0.0, 0.0, 1.0],
772        ];
773        let csr = CsrMatrix::from_dense(&m);
774        let x = vec![3.0, 7.0, -2.0];
775        let y = csr.spmv(&x);
776        assert_eq!(y, x);
777    }
778
779    #[test]
780    fn test_csr_spmv_known_3x3() {
781        // A = [[2,1,0],[1,3,1],[0,1,2]]
782        let m = vec![
783            vec![2.0, 1.0, 0.0],
784            vec![1.0, 3.0, 1.0],
785            vec![0.0, 1.0, 2.0],
786        ];
787        let csr = CsrMatrix::from_dense(&m);
788        let x = vec![1.0, 2.0, 3.0];
789        let y = csr.spmv(&x);
790        // y[0] = 2+2 = 4, y[1] = 1+6+3 = 10, y[2] = 2+6 = 8
791        assert!((y[0] - 4.0).abs() < 1e-12);
792        assert!((y[1] - 10.0).abs() < 1e-12);
793        assert!((y[2] - 8.0).abs() < 1e-12);
794    }
795
796    #[test]
797    fn test_cg_solve_diagonal_spd() {
798        // A = diag(1,2,3,4), b = [1,2,3,4], exact solution x = [1,1,1,1]
799        let m = vec![
800            vec![1.0, 0.0, 0.0, 0.0],
801            vec![0.0, 2.0, 0.0, 0.0],
802            vec![0.0, 0.0, 3.0, 0.0],
803            vec![0.0, 0.0, 0.0, 4.0],
804        ];
805        let a = CsrMatrix::from_dense(&m);
806        let b = vec![1.0, 2.0, 3.0, 4.0];
807        let x0 = vec![0.0; 4];
808        let (x, _iters) = cg_solve(&a, &b, &x0, 100, 1e-12);
809        for v in &x {
810            assert!((v - 1.0).abs() < 1e-10, "x value {v} not close to 1.0");
811        }
812    }
813
814    #[test]
815    fn test_sparse_triplet_to_csr_duplicate_sum() {
816        let mut trip = SparseTriplet::new();
817        trip.add(0, 0, 1.0);
818        trip.add(0, 0, 2.0); // duplicate → should sum to 3.0
819        trip.add(1, 1, 5.0);
820        let csr = trip.to_csr(2, 2);
821        assert!((csr.get(0, 0) - 3.0).abs() < 1e-12);
822        assert!((csr.get(1, 1) - 5.0).abs() < 1e-12);
823        assert_eq!(csr.nnz(), 2);
824    }
825
826    #[test]
827    fn test_ell_spmv_matches_csr() {
828        let m = vec![
829            vec![2.0, 1.0, 0.0],
830            vec![1.0, 3.0, 1.0],
831            vec![0.0, 1.0, 2.0],
832        ];
833        let csr = CsrMatrix::from_dense(&m);
834        let ell = EllMatrix::from_csr(&csr);
835        let x = vec![1.0, -1.0, 2.0];
836        let y_csr = csr.spmv(&x);
837        let y_ell = ell.spmv(&x);
838        for (a, b) in y_csr.iter().zip(y_ell.iter()) {
839            assert!((a - b).abs() < 1e-12, "ELL mismatch: {a} vs {b}");
840        }
841    }
842
843    #[test]
844    fn test_csr_transpose() {
845        // A = [[1,2,3],[4,5,6]]  →  A^T = [[1,4],[2,5],[3,6]]
846        let m = vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]];
847        let csr = CsrMatrix::from_dense(&m);
848        let at = csr.transpose();
849        assert_eq!(at.n_rows, 3);
850        assert_eq!(at.n_cols, 2);
851        assert!((at.get(0, 0) - 1.0).abs() < 1e-12);
852        assert!((at.get(0, 1) - 4.0).abs() < 1e-12);
853        assert!((at.get(1, 0) - 2.0).abs() < 1e-12);
854        assert!((at.get(1, 1) - 5.0).abs() < 1e-12);
855        assert!((at.get(2, 0) - 3.0).abs() < 1e-12);
856        assert!((at.get(2, 1) - 6.0).abs() < 1e-12);
857    }
858
859    // ── spmv_segmented ─────────────────────────────────────────────────────
860
861    #[test]
862    fn test_spmv_segmented_identity() {
863        let m = vec![
864            vec![1.0, 0.0, 0.0],
865            vec![0.0, 1.0, 0.0],
866            vec![0.0, 0.0, 1.0],
867        ];
868        let csr = CsrMatrix::from_dense(&m);
869        let x = vec![3.0, 7.0, -2.0];
870        let y = spmv_segmented(&csr, &x);
871        for (a, b) in y.iter().zip(x.iter()) {
872            assert!((a - b).abs() < 1e-12);
873        }
874    }
875
876    #[test]
877    fn test_spmv_segmented_matches_csr() {
878        let m = vec![
879            vec![2.0, 1.0, 0.0],
880            vec![1.0, 3.0, 1.0],
881            vec![0.0, 1.0, 2.0],
882        ];
883        let csr = CsrMatrix::from_dense(&m);
884        let x = vec![1.0, -1.0, 2.0];
885        let y_std = csr.spmv(&x);
886        let y_seg = spmv_segmented(&csr, &x);
887        for (a, b) in y_std.iter().zip(y_seg.iter()) {
888            assert!((a - b).abs() < 1e-12);
889        }
890    }
891
892    // ── assemble_1d_laplacian ──────────────────────────────────────────────
893
894    #[test]
895    fn test_assemble_1d_laplacian_3x3() {
896        let l = assemble_1d_laplacian(3);
897        // Expected: [[2,-1,0],[-1,2,-1],[0,-1,2]]
898        assert!((l.get(0, 0) - 2.0).abs() < 1e-12);
899        assert!((l.get(0, 1) - (-1.0)).abs() < 1e-12);
900        assert!((l.get(0, 2)).abs() < 1e-12);
901        assert!((l.get(1, 0) - (-1.0)).abs() < 1e-12);
902        assert!((l.get(1, 1) - 2.0).abs() < 1e-12);
903        assert!((l.get(1, 2) - (-1.0)).abs() < 1e-12);
904        assert!((l.get(2, 2) - 2.0).abs() < 1e-12);
905    }
906
907    #[test]
908    fn test_assemble_1d_laplacian_spd() {
909        // 1D Laplacian is symmetric positive definite
910        let n = 5;
911        let l = assemble_1d_laplacian(n);
912        // Check symmetry
913        for i in 0..n {
914            for j in 0..n {
915                assert!((l.get(i, j) - l.get(j, i)).abs() < 1e-12);
916            }
917        }
918        // CG should converge for SPD systems
919        let b = vec![1.0; n];
920        let x0 = vec![0.0; n];
921        let (x, iters) = cg_solve(&l, &b, &x0, 200, 1e-10);
922        assert!(iters < 200);
923        // Verify Ax ≈ b
924        let ax = l.spmv(&x);
925        for i in 0..n {
926            assert!((ax[i] - b[i]).abs() < 1e-8);
927        }
928    }
929
930    // ── csr_to_ell ─────────────────────────────────────────────────────────
931
932    #[test]
933    fn test_csr_to_ell_spmv() {
934        let m = vec![
935            vec![5.0, 0.0, 1.0, 0.0],
936            vec![0.0, 3.0, 0.0, 2.0],
937            vec![1.0, 0.0, 4.0, 0.0],
938            vec![0.0, 0.0, 0.0, 6.0],
939        ];
940        let csr = CsrMatrix::from_dense(&m);
941        let ell = csr_to_ell(&csr);
942        let x = vec![1.0, 2.0, 3.0, 4.0];
943        let y_csr = csr.spmv(&x);
944        let y_ell = ell.spmv(&x);
945        for (a, b) in y_csr.iter().zip(y_ell.iter()) {
946            assert!((a - b).abs() < 1e-12, "mismatch: {a} vs {b}");
947        }
948    }
949
950    #[test]
951    fn test_csr_to_ell_max_nnz() {
952        let m = vec![
953            vec![1.0, 0.0, 0.0],
954            vec![1.0, 2.0, 3.0], // 3 nnz → max
955            vec![0.0, 0.0, 1.0],
956        ];
957        let csr = CsrMatrix::from_dense(&m);
958        let ell = csr_to_ell(&csr);
959        assert_eq!(ell.max_nnz_per_row, 3);
960    }
961
962    // ── BlockCsrMatrix additional tests ────────────────────────────────────
963
964    #[test]
965    fn test_block_csr_spmv_2x2() {
966        // Single 2x2 block [[1,2],[3,4]] at block position (0,0)
967        let bcsr = BlockCsrMatrix {
968            block_size: 2,
969            n_block_rows: 1,
970            n_block_cols: 1,
971            row_ptr: vec![0, 1],
972            col_idx: vec![0],
973            blocks: vec![vec![1.0, 2.0, 3.0, 4.0]],
974        };
975        let x = vec![1.0, 1.0];
976        let y = bcsr.spmv_block(&x);
977        assert!((y[0] - 3.0).abs() < 1e-12); // 1+2
978        assert!((y[1] - 7.0).abs() < 1e-12); // 3+4
979    }
980
981    #[test]
982    fn test_block_csr_to_csr_roundtrip() {
983        let bcsr = BlockCsrMatrix {
984            block_size: 2,
985            n_block_rows: 2,
986            n_block_cols: 2,
987            row_ptr: vec![0, 1, 2],
988            col_idx: vec![0, 1],
989            blocks: vec![vec![1.0, 2.0, 3.0, 4.0], vec![5.0, 6.0, 7.0, 8.0]],
990        };
991        let csr = bcsr.to_csr();
992        let x = vec![1.0, 1.0, 1.0, 1.0];
993        let y_block = bcsr.spmv_block(&x);
994        let y_csr = csr.spmv(&x);
995        for (a, b) in y_block.iter().zip(y_csr.iter()) {
996            assert!((a - b).abs() < 1e-12);
997        }
998    }
999
1000    // ── sparse_lower_triangular_solve ──────────────────────────────────────
1001
1002    #[test]
1003    fn test_lower_tri_solve_identity() {
1004        let m = vec![
1005            vec![1.0, 0.0, 0.0],
1006            vec![0.0, 1.0, 0.0],
1007            vec![0.0, 0.0, 1.0],
1008        ];
1009        let l = CsrMatrix::from_dense(&m);
1010        let b = vec![3.0, 7.0, -2.0];
1011        let x = sparse_lower_triangular_solve(&l, &b);
1012        for (a, bv) in x.iter().zip(b.iter()) {
1013            assert!((a - bv).abs() < 1e-12);
1014        }
1015    }
1016
1017    #[test]
1018    fn test_lower_tri_solve_3x3() {
1019        // L = [[2,0,0],[1,3,0],[4,2,5]]
1020        let m = vec![
1021            vec![2.0, 0.0, 0.0],
1022            vec![1.0, 3.0, 0.0],
1023            vec![4.0, 2.0, 5.0],
1024        ];
1025        let l = CsrMatrix::from_dense(&m);
1026        let b = vec![4.0, 7.0, 26.0];
1027        let x = sparse_lower_triangular_solve(&l, &b);
1028        // x[0] = 4/2 = 2
1029        // x[1] = (7 - 1*2)/3 = 5/3
1030        // x[2] = (26 - 4*2 - 2*(5/3))/5
1031        assert!((x[0] - 2.0).abs() < 1e-10);
1032        assert!((x[1] - 5.0 / 3.0).abs() < 1e-10);
1033        let expected_x2 = (26.0 - 8.0 - 10.0 / 3.0) / 5.0;
1034        assert!((x[2] - expected_x2).abs() < 1e-10);
1035    }
1036
1037    #[test]
1038    fn test_lower_tri_solve_verify_lx_eq_b() {
1039        let m = vec![
1040            vec![3.0, 0.0, 0.0, 0.0],
1041            vec![1.0, 2.0, 0.0, 0.0],
1042            vec![0.0, 4.0, 5.0, 0.0],
1043            vec![2.0, 0.0, 1.0, 6.0],
1044        ];
1045        let l = CsrMatrix::from_dense(&m);
1046        let b = vec![9.0, 8.0, 22.0, 29.0];
1047        let x = sparse_lower_triangular_solve(&l, &b);
1048        // Verify L*x = b
1049        let lx = l.spmv(&x);
1050        for i in 0..4 {
1051            assert!(
1052                (lx[i] - b[i]).abs() < 1e-10,
1053                "row {i}: {} vs {}",
1054                lx[i],
1055                b[i]
1056            );
1057        }
1058    }
1059
1060    // ── sparse_upper_triangular_solve ──────────────────────────────────────
1061
1062    #[test]
1063    fn test_upper_tri_solve_identity() {
1064        let m = vec![
1065            vec![1.0, 0.0, 0.0],
1066            vec![0.0, 1.0, 0.0],
1067            vec![0.0, 0.0, 1.0],
1068        ];
1069        let u = CsrMatrix::from_dense(&m);
1070        let b = vec![3.0, 7.0, -2.0];
1071        let x = sparse_upper_triangular_solve(&u, &b);
1072        for (a, bv) in x.iter().zip(b.iter()) {
1073            assert!((a - bv).abs() < 1e-12);
1074        }
1075    }
1076
1077    #[test]
1078    fn test_upper_tri_solve_verify_ux_eq_b() {
1079        // U = [[2,1,3],[0,4,2],[0,0,5]]
1080        let m = vec![
1081            vec![2.0, 1.0, 3.0],
1082            vec![0.0, 4.0, 2.0],
1083            vec![0.0, 0.0, 5.0],
1084        ];
1085        let u = CsrMatrix::from_dense(&m);
1086        let b = vec![13.0, 14.0, 10.0];
1087        let x = sparse_upper_triangular_solve(&u, &b);
1088        let ux = u.spmv(&x);
1089        for i in 0..3 {
1090            assert!(
1091                (ux[i] - b[i]).abs() < 1e-10,
1092                "row {i}: {} vs {}",
1093                ux[i],
1094                b[i]
1095            );
1096        }
1097    }
1098
1099    // ── CsrMatrix::add & scale ─────────────────────────────────────────────
1100
1101    #[test]
1102    fn test_csr_add() {
1103        let m1 = vec![vec![1.0, 0.0], vec![0.0, 2.0]];
1104        let m2 = vec![vec![0.0, 3.0], vec![4.0, 0.0]];
1105        let a = CsrMatrix::from_dense(&m1);
1106        let b = CsrMatrix::from_dense(&m2);
1107        let c = a.add(&b);
1108        assert!((c.get(0, 0) - 1.0).abs() < 1e-12);
1109        assert!((c.get(0, 1) - 3.0).abs() < 1e-12);
1110        assert!((c.get(1, 0) - 4.0).abs() < 1e-12);
1111        assert!((c.get(1, 1) - 2.0).abs() < 1e-12);
1112    }
1113
1114    #[test]
1115    fn test_csr_scale() {
1116        let m = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
1117        let a = CsrMatrix::from_dense(&m);
1118        let b = a.scale(2.0);
1119        assert!((b.get(0, 0) - 2.0).abs() < 1e-12);
1120        assert!((b.get(0, 1) - 4.0).abs() < 1e-12);
1121        assert!((b.get(1, 0) - 6.0).abs() < 1e-12);
1122        assert!((b.get(1, 1) - 8.0).abs() < 1e-12);
1123    }
1124
1125    // ── BiCGSTAB ───────────────────────────────────────────────────────────
1126
1127    #[test]
1128    fn test_bicgstab_diagonal() {
1129        let m = vec![
1130            vec![2.0, 0.0, 0.0],
1131            vec![0.0, 3.0, 0.0],
1132            vec![0.0, 0.0, 4.0],
1133        ];
1134        let a = CsrMatrix::from_dense(&m);
1135        let b = vec![4.0, 9.0, 16.0];
1136        let x0 = vec![0.0; 3];
1137        let (x, _iters) = bicgstab_solve(&a, &b, &x0, 100, 1e-10);
1138        assert!((x[0] - 2.0).abs() < 1e-8);
1139        assert!((x[1] - 3.0).abs() < 1e-8);
1140        assert!((x[2] - 4.0).abs() < 1e-8);
1141    }
1142
1143    #[test]
1144    fn test_bicgstab_nonsymmetric() {
1145        // Non-symmetric but diagonally dominant
1146        let m = vec![
1147            vec![4.0, 1.0, 0.0],
1148            vec![0.0, 3.0, 1.0],
1149            vec![0.0, 0.0, 5.0],
1150        ];
1151        let a = CsrMatrix::from_dense(&m);
1152        let b = vec![5.0, 4.0, 5.0];
1153        let x0 = vec![0.0; 3];
1154        let (x, _iters) = bicgstab_solve(&a, &b, &x0, 200, 1e-10);
1155        // Verify Ax ≈ b
1156        let ax = a.spmv(&x);
1157        for i in 0..3 {
1158            assert!(
1159                (ax[i] - b[i]).abs() < 1e-6,
1160                "row {i}: {} vs {}",
1161                ax[i],
1162                b[i]
1163            );
1164        }
1165    }
1166
1167    // ── Jacobi preconditioned CG ───────────────────────────────────────────
1168
1169    #[test]
1170    fn test_jacobi_pcg_laplacian() {
1171        let n = 8;
1172        let a = assemble_1d_laplacian(n);
1173        let b: Vec<f64> = (0..n).map(|i| (i as f64 + 1.0).sin()).collect();
1174        let (x, iters) = jacobi_preconditioned_cg(&a, &b, 500, 1e-10);
1175        assert!(iters < 500, "PCG should converge, used {iters} iterations");
1176        let ax = a.spmv(&x);
1177        for i in 0..n {
1178            assert!(
1179                (ax[i] - b[i]).abs() < 1e-6,
1180                "row {i}: {} vs {}",
1181                ax[i],
1182                b[i]
1183            );
1184        }
1185    }
1186
1187    // ── GPU simulation utilities ───────────────────────────────────────────
1188
1189    #[test]
1190    fn test_simulate_spmv_throughput() {
1191        let gflops = simulate_spmv_throughput(100, 1000);
1192        assert!(gflops > 0.0);
1193        // With 100 GB/s bandwidth and 24 bytes/nnz, bigger nnz should give higher GFLOPS
1194        let gflops2 = simulate_spmv_throughput(1000, 10000);
1195        assert!((gflops2 - gflops).abs() < 1.0); // roofline is flat w.r.t. nnz ratio
1196    }
1197
1198    #[test]
1199    fn test_optimal_ell_row_width_empty() {
1200        assert_eq!(optimal_ell_row_width(&[]), 0);
1201    }
1202
1203    #[test]
1204    fn test_optimal_ell_row_width_uniform() {
1205        // All rows have 5 nnz
1206        let dist = vec![5; 10];
1207        assert_eq!(optimal_ell_row_width(&dist), 5);
1208    }
1209
1210    #[test]
1211    fn test_optimal_ell_row_width_varied() {
1212        let dist = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
1213        let w = optimal_ell_row_width(&dist);
1214        // 75th percentile index = (10 * 3) / 4 = 7 → sorted[7] = 8
1215        assert_eq!(w, 8);
1216    }
1217
1218    // ── SparseTriplet default ──────────────────────────────────────────────
1219
1220    #[test]
1221    fn test_sparse_triplet_default() {
1222        let t = SparseTriplet::default();
1223        assert!(t.rows.is_empty());
1224        assert!(t.cols.is_empty());
1225        assert!(t.vals.is_empty());
1226    }
1227
1228    // ── vector operations ──────────────────────────────────────────────────
1229
1230    #[test]
1231    fn test_dot_product() {
1232        let x = vec![1.0, 2.0, 3.0];
1233        let y = vec![4.0, 5.0, 6.0];
1234        assert!((dot(&x, &y) - 32.0).abs() < 1e-12);
1235    }
1236
1237    #[test]
1238    fn test_axpy() {
1239        let x = vec![1.0, 2.0, 3.0];
1240        let y = vec![10.0, 20.0, 30.0];
1241        let z = axpy(2.0, &x, &y);
1242        assert_eq!(z, vec![12.0, 24.0, 36.0]);
1243    }
1244
1245    #[test]
1246    fn test_norm2() {
1247        let x = vec![3.0, 4.0];
1248        assert!((norm2(&x) - 5.0).abs() < 1e-12);
1249    }
1250
1251    #[test]
1252    fn test_scale_vec() {
1253        let x = vec![1.0, 2.0, 3.0];
1254        let s = scale_vec(&x, 3.0);
1255        assert_eq!(s, vec![3.0, 6.0, 9.0]);
1256    }
1257
1258    // ── HybridMatrix ──────────────────────────────────────────────────────
1259
1260    #[test]
1261    fn test_hybrid_spmv() {
1262        let m = vec![
1263            vec![1.0, 2.0, 0.0],
1264            vec![0.0, 3.0, 0.0],
1265            vec![0.0, 0.0, 4.0],
1266        ];
1267        let csr = CsrMatrix::from_dense(&m);
1268        let ell = EllMatrix::from_csr(&csr);
1269        let hybrid = HybridMatrix {
1270            ell,
1271            coo_row: vec![],
1272            coo_col: vec![],
1273            coo_val: vec![],
1274        };
1275        let x = vec![1.0, 1.0, 1.0];
1276        let y = hybrid.spmv(&x);
1277        assert!((y[0] - 3.0).abs() < 1e-12);
1278        assert!((y[1] - 3.0).abs() < 1e-12);
1279        assert!((y[2] - 4.0).abs() < 1e-12);
1280    }
1281
1282    #[test]
1283    fn test_hybrid_spmv_with_coo() {
1284        // ELL part: identity, COO part: adds off-diagonal
1285        let m = vec![vec![1.0, 0.0], vec![0.0, 1.0]];
1286        let csr = CsrMatrix::from_dense(&m);
1287        let ell = EllMatrix::from_csr(&csr);
1288        let hybrid = HybridMatrix {
1289            ell,
1290            coo_row: vec![0],
1291            coo_col: vec![1],
1292            coo_val: vec![5.0],
1293        };
1294        let x = vec![1.0, 2.0];
1295        let y = hybrid.spmv(&x);
1296        assert!((y[0] - 11.0).abs() < 1e-12); // 1*1 + 5*2
1297        assert!((y[1] - 2.0).abs() < 1e-12);
1298    }
1299
1300    // ── CsrMatrix from empty ──────────────────────────────────────────────
1301
1302    #[test]
1303    fn test_csr_empty() {
1304        let csr = CsrMatrix::new(3, 3);
1305        assert_eq!(csr.nnz(), 0);
1306        let y = csr.spmv(&[1.0, 2.0, 3.0]);
1307        assert_eq!(y, vec![0.0, 0.0, 0.0]);
1308    }
1309
1310    // ── extract_diagonal ──────────────────────────────────────────────────
1311
1312    #[test]
1313    fn test_extract_diagonal() {
1314        let m = vec![
1315            vec![5.0, 1.0, 0.0],
1316            vec![0.0, 3.0, 2.0],
1317            vec![0.0, 0.0, 7.0],
1318        ];
1319        let csr = CsrMatrix::from_dense(&m);
1320        let diag = extract_diagonal(&csr);
1321        assert!((diag[0] - 5.0).abs() < 1e-12);
1322        assert!((diag[1] - 3.0).abs() < 1e-12);
1323        assert!((diag[2] - 7.0).abs() < 1e-12);
1324    }
1325
1326    #[test]
1327    fn test_extract_diagonal_missing() {
1328        // Matrix with zero diagonal
1329        let m = vec![vec![0.0, 1.0], vec![1.0, 0.0]];
1330        let csr = CsrMatrix::from_dense(&m);
1331        let diag = extract_diagonal(&csr);
1332        assert!((diag[0]).abs() < 1e-12);
1333        assert!((diag[1]).abs() < 1e-12);
1334    }
1335
1336    // ── compute_nnz_per_row ────────────────────────────────────────────────
1337
1338    #[test]
1339    fn test_compute_nnz_per_row() {
1340        let m = vec![
1341            vec![1.0, 2.0, 0.0],
1342            vec![0.0, 3.0, 0.0],
1343            vec![4.0, 5.0, 6.0],
1344        ];
1345        let csr = CsrMatrix::from_dense(&m);
1346        let nnz = compute_nnz_per_row(&csr);
1347        assert_eq!(nnz, vec![2, 1, 3]);
1348    }
1349
1350    // ── frobenius_norm ─────────────────────────────────────────────────────
1351
1352    #[test]
1353    fn test_frobenius_norm() {
1354        // Identity 3x3 → frobenius = sqrt(3)
1355        let m = vec![
1356            vec![1.0, 0.0, 0.0],
1357            vec![0.0, 1.0, 0.0],
1358            vec![0.0, 0.0, 1.0],
1359        ];
1360        let csr = CsrMatrix::from_dense(&m);
1361        let f = frobenius_norm(&csr);
1362        assert!((f - 3.0_f64.sqrt()).abs() < 1e-12);
1363    }
1364
1365    #[test]
1366    fn test_frobenius_norm_known() {
1367        let m = vec![vec![3.0, 4.0], vec![0.0, 0.0]];
1368        let csr = CsrMatrix::from_dense(&m);
1369        let f = frobenius_norm(&csr);
1370        // sqrt(9 + 16) = 5
1371        assert!((f - 5.0).abs() < 1e-12);
1372    }
1373}
oxiphysics_gpu/sparse_gpu.rs

oxiphysics_gpu/
sparse_gpu.rs
