Skip to main content

yui_matrix/sparse/
snf.rs

1//! Sparse Smith normal form over a Euclidean ring.
2//!
3//! Mirrors [`crate::dense::snf`] (same pivoting and Euclidean elimination),
4//! but operates on hashed sparse rows with a column index, so large
5//! reducer-residuals never materialize as dense matrices. The four
6//! transformation matrices are reconstructed at the end by replaying the
7//! recorded elementary operations on sparse identities.
8
9use log::debug;
10use rustc_hash::{FxHashMap, FxHashSet};
11use yui_core::abst::{EucRing, EucRingOps};
12use yui_core::util::log::log_step_crossed;
13
14use crate::MatTrait;
15use crate::dense::snf::SnfFlags;
16use super::SpMat;
17
18// one progress line per this many pivots.
19const PIVOT_LOG_STEP: usize = 100_000;
20
21/// Result of a sparse SNF: `p * a * q = result` (diagonal), `pinv * result * qinv = a`.
22pub struct SpSnf<R>
23where R: EucRing, for<'x> &'x R: EucRingOps<R> {
24    result: SpMat<R>,
25    diag: Vec<R>,
26    p:    Option<SpMat<R>>,
27    pinv: Option<SpMat<R>>,
28    q:    Option<SpMat<R>>,
29    qinv: Option<SpMat<R>>,
30}
31
32impl<R> SpSnf<R>
33where R: EucRing, for<'x> &'x R: EucRingOps<R> {
34    pub fn result(&self) -> &SpMat<R> {
35        &self.result
36    }
37
38    pub fn p(&self) -> Option<&SpMat<R>> {
39        self.p.as_ref()
40    }
41
42    pub fn pinv(&self) -> Option<&SpMat<R>> {
43        self.pinv.as_ref()
44    }
45
46    pub fn q(&self) -> Option<&SpMat<R>> {
47        self.q.as_ref()
48    }
49
50    pub fn qinv(&self) -> Option<&SpMat<R>> {
51        self.qinv.as_ref()
52    }
53
54    pub fn rank(&self) -> usize {
55        self.diag.len()
56    }
57
58    pub fn factors(&self) -> Vec<&R> {
59        self.diag.iter().collect()
60    }
61}
62
63/// Below this core area (non-zero rows × non-zero cols), the matrix is
64/// compacted by permutation and passed to the dense SNF instead.
65const DENSE_SNF_MAX_AREA: usize = 65536; // 256 × 256
66
67/// Computes the sparse SNF of `a`, producing the transformation matrices
68/// selected by `flags = [p, pinv, q, qinv]`.
69///
70/// Small inputs take a fast pass: permute the non-zero core to the front and
71/// run [`crate::dense::snf`] on it; large ones are eliminated sparsely.
72pub fn sp_snf<R>(a: &SpMat<R>, flags: SnfFlags) -> SpSnf<R>
73where R: EucRing, for<'x> &'x R: EucRingOps<R> {
74    sp_snf_with(a, flags, DENSE_SNF_MAX_AREA)
75}
76
77fn sp_snf_with<R>(a: &SpMat<R>, flags: SnfFlags, dense_max_area: usize) -> SpSnf<R>
78where R: EucRing, for<'x> &'x R: EucRingOps<R> {
79    let (row_idx, col_idx) = nz_indices(a);
80    let (m0, n0) = (row_idx.len(), col_idx.len());
81
82    debug!("start sparse snf: {:?}, nnz: {}, core: {m0}x{n0}, flags: {:?}", a.shape(), a.iter_nz().count(), flags);
83
84    if m0 * n0 <= dense_max_area {
85        return dense_snf_in(a, flags, row_idx, col_idx);
86    }
87
88    let mut calc = SpSnfCalc::new(a, flags);
89    calc.process();
90
91    debug!("sparse snf done, rank: {}", calc.rank);
92
93    calc.into_result()
94}
95
96// Sorted non-zero row / col indices of `a`.
97fn nz_indices<R>(a: &SpMat<R>) -> (Vec<usize>, Vec<usize>)
98where R: EucRing, for<'x> &'x R: EucRingOps<R> {
99    let mut rows = vec![false; a.n_rows()];
100    let mut cols = vec![false; a.n_cols()];
101    for (i, j, _) in a.iter_nz() {
102        rows[i] = true;
103        cols[j] = true;
104    }
105    let row_idx = rows.iter().enumerate().filter_map(|(i, &b)| b.then_some(i)).collect();
106    let col_idx = cols.iter().enumerate().filter_map(|(j, &b)| b.then_some(j)).collect();
107    (row_idx, col_idx)
108}
109
110// Fast pass: gather the non-zero core `a[row_idx, col_idx]` to the front,
111// run the dense SNF on it, and lift. With `R`/`C` the gathering permutations,
112// `R·a·C = [[core, 0], [0, 0]]`, so
113// `p = diag(p_d, I)·R`, `q = C·diag(q_d, I)`, and inverses transposed-fashion.
114fn dense_snf_in<R>(a: &SpMat<R>, flags: SnfFlags, row_idx: Vec<usize>, col_idx: Vec<usize>) -> SpSnf<R>
115where R: EucRing, for<'x> &'x R: EucRingOps<R> {
116    use crate::dense::Mat;
117    use crate::dense::snf::snf_in_place;
118
119    let (m, n) = a.shape();
120    let (m0, n0) = (row_idx.len(), col_idx.len());
121
122    debug!("  snf fast-pass: dense on {m0}x{n0} core");
123
124    let row_pos: FxHashMap<usize, usize> = row_idx.iter().enumerate().map(|(k, &i)| (i, k)).collect();
125    let col_pos: FxHashMap<usize, usize> = col_idx.iter().enumerate().map(|(k, &j)| (j, k)).collect();
126
127    let mut core = Mat::zero((m0, n0));
128    for (i, j, v) in a.iter_nz() {
129        core[(row_pos[&i], col_pos[&j])] = v.clone();
130    }
131
132    let s = snf_in_place(core, flags);
133    let rank = s.rank();
134    let (res_d, [p_d, pinv_d, q_d, qinv_d]) = s.destruct();
135
136    let diag: Vec<R> = (0..rank).map(|i| res_d[(i, i)].clone()).collect();
137    let result = SpMat::from_entries((m, n), diag.iter().enumerate().map(|(i, v)| (i, i, v.clone())));
138
139    let rest_rows = || (0..m).filter(|i| !row_pos.contains_key(i));
140    let rest_cols = || (0..n).filter(|j| !col_pos.contains_key(j));
141
142    // p[k, row_idx[l]] = p_d[k, l];  p[m0+t, rest_row_t] = 1
143    let p = p_d.map(|p_d| SpMat::from_entries((m, m),
144        dense_nz(&p_d).map(|(k, l, v)| (k, row_idx[l], v))
145            .chain(rest_rows().enumerate().map(|(t, i)| (m0 + t, i, R::one())))
146    ));
147    // pinv[row_idx[k], l] = pinv_d[k, l];  pinv[rest_row_t, m0+t] = 1
148    let pinv = pinv_d.map(|pinv_d| SpMat::from_entries((m, m),
149        dense_nz(&pinv_d).map(|(k, l, v)| (row_idx[k], l, v))
150            .chain(rest_rows().enumerate().map(|(t, i)| (i, m0 + t, R::one())))
151    ));
152    // q[col_idx[k], l] = q_d[k, l];  q[rest_col_t, n0+t] = 1
153    let q = q_d.map(|q_d| SpMat::from_entries((n, n),
154        dense_nz(&q_d).map(|(k, l, v)| (col_idx[k], l, v))
155            .chain(rest_cols().enumerate().map(|(t, j)| (j, n0 + t, R::one())))
156    ));
157    // qinv[k, col_idx[l]] = qinv_d[k, l];  qinv[n0+t, rest_col_t] = 1
158    let qinv = qinv_d.map(|qinv_d| SpMat::from_entries((n, n),
159        dense_nz(&qinv_d).map(|(k, l, v)| (k, col_idx[l], v))
160            .chain(rest_cols().enumerate().map(|(t, j)| (n0 + t, j, R::one())))
161    ));
162
163    SpSnf { result, diag, p, pinv, q, qinv }
164}
165
166fn dense_nz<R>(a: &crate::dense::Mat<R>) -> impl Iterator<Item = (usize, usize, R)> + '_
167where R: EucRing, for<'x> &'x R: EucRingOps<R> {
168    let (m, n) = a.shape();
169    (0..m).flat_map(move |i| (0..n).filter_map(move |j| {
170        let v = &a[(i, j)];
171        (!v.is_zero()).then(|| (i, j, v.clone()))
172    }))
173}
174
175struct SpSnfCalc<R>
176where R: EucRing, for<'x> &'x R: EucRingOps<R> {
177    shape: (usize, usize),
178    rows: Vec<FxHashMap<usize, R>>,  // row -> (col -> value)
179    cols: Vec<FxHashSet<usize>>,     // col -> rows with a non-zero entry
180    // trans matrices, updated in realtime like the dense SnfCalc:
181    // p / qinv accumulate row-wise, pinv / q column-wise.
182    p:    Option<VecStore<R>>,
183    pinv: Option<VecStore<R>>,
184    q:    Option<VecStore<R>>,
185    qinv: Option<VecStore<R>>,
186    rank: usize,
187}
188
189impl<R> SpSnfCalc<R>
190where R: EucRing, for<'x> &'x R: EucRingOps<R> {
191    fn new(a: &SpMat<R>, flags: SnfFlags) -> Self {
192        let (m, n) = a.shape();
193        let mut rows = vec![FxHashMap::default(); m];
194        let mut cols = vec![FxHashSet::default(); n];
195
196        for (i, j, v) in a.iter_nz() {
197            rows[i].insert(j, v.clone());
198            cols[j].insert(i);
199        }
200
201        let [fp, fpinv, fq, fqinv] = flags;
202        let p    = fp.then(|| VecStore::id(m));
203        let pinv = fpinv.then(|| VecStore::id(m));
204        let q    = fq.then(|| VecStore::id(n));
205        let qinv = fqinv.then(|| VecStore::id(n));
206
207        Self { shape: (m, n), rows, cols, p, pinv, q, qinv, rank: 0 }
208    }
209
210    fn process(&mut self) {
211        self.eliminate_all();
212        debug!("  snf eliminate done: rank {}; normalize..", self.rank);
213        self.diag_normalize();
214    }
215
216    fn into_result(self) -> SpSnf<R> {
217        let diag: Vec<R> = (0..self.rank).map(|i| self.rows[i][&i].clone()).collect();
218        let result = SpMat::from_entries(self.shape, self.rows.iter().enumerate().flat_map(|(i, row)|
219            row.iter().map(move |(&j, v)| (i, j, v.clone()))
220        ));
221        let (m, n) = self.shape;
222        let p    = self.p.map(|s| s.into_spmat_rows(m));
223        let pinv = self.pinv.map(|s| s.into_spmat_cols(m));
224        let q    = self.q.map(|s| s.into_spmat_cols(n));
225        let qinv = self.qinv.map(|s| s.into_spmat_rows(n));
226        SpSnf { result, diag, p, pinv, q, qinv }
227    }
228
229    // ---- elimination ----
230
231    fn eliminate_all(&mut self) {
232        let (m, n) = self.shape;
233        let mut i = 0;
234
235        for j in 0..n {
236            if i >= m { break }
237            if self.eliminate_step(i, j) {
238                i += 1;
239                // its own message: nnz is how fill-in is watched, so the plain form won't do.
240                if log_step_crossed(i, i - 1, m, PIVOT_LOG_STEP) {
241                    let nnz: usize = self.rows.iter().map(|r| r.len()).sum();
242                    debug!("  snf progress: {i} pivots ({j}/{n} cols), nnz {nnz}");
243                }
244            }
245        }
246    }
247
248    fn eliminate_step(&mut self, i: usize, j: usize) -> bool {
249        // pivot: the row with minimal nnz among rows ≥ i having an entry in col j
250        let Some(i_p) = self.cols[j].iter()
251            .filter(|&&r| r >= i)
252            .min_by_key(|&&r| self.rows[r].len())
253            .copied()
254        else {
255            return false
256        };
257
258        if i_p > i {
259            self.swap_rows(i, i_p);
260        }
261        if j > i {
262            self.swap_cols(i, j);
263        }
264
265        let u = self.rows[i][&i].normalizing_unit();
266        if !u.is_one() {
267            self.mul_col(i, &u);
268        }
269
270        self.eliminate_at(i);
271        self.rank += 1;
272
273        true
274    }
275
276    // Clears row i and column i by Euclidean row/col operations on the pivot (i, i).
277    fn eliminate_at(&mut self, i: usize) {
278        debug_assert!(self.rows[i].contains_key(&i));
279
280        while self.rows[i].len() > 1 || self.cols[i].len() > 1 {
281            let modified = self.eliminate_col(i) | self.eliminate_row(i);
282            if !modified {
283                panic!("sparse snf: no progress at pivot {i}");
284            }
285        }
286    }
287
288    fn eliminate_col(&mut self, i: usize) -> bool {
289        let mut modified = false;
290        let targets: Vec<usize> = self.cols[i].iter().copied().filter(|&r| r != i).collect();
291
292        for i1 in targets {
293            let (Some(x), Some(y)) = (self.entry(i, i).cloned(), self.entry(i1, i).cloned()) else { continue };
294            let (d, s, t) = gcdx(&x, &y);
295            let (a, b) = (&x / &d, &y / &d);
296
297            // [ s t][rᵢ ]   [ d ]
298            // [-b a][rᵢ₁] = [ 0 ]   (at col i)
299            self.row_elem([s, t, -b, a], i, i1);
300            modified = true;
301        }
302
303        modified
304    }
305
306    fn eliminate_row(&mut self, i: usize) -> bool {
307        let mut modified = false;
308        let targets: Vec<usize> = self.rows[i].keys().copied().filter(|&c| c != i).collect();
309
310        for j1 in targets {
311            let (Some(x), Some(y)) = (self.entry(i, i).cloned(), self.entry(i, j1).cloned()) else { continue };
312            let (d, s, t) = gcdx(&x, &y);
313            let (a, b) = (&x / &d, &y / &d);
314
315            // [cᵢ cⱼ₁][s -b]   [d 0]   (at row i)
316            //         [t  a] =
317            self.col_elem([s, t, -b, a], i, j1);
318            modified = true;
319        }
320
321        modified
322    }
323
324    // ---- diagonal normalization (divisor chain d₁ | d₂ | …) ----
325
326    fn diag_normalize(&mut self) {
327        let r = self.rank;
328        if r == 0 {
329            return
330        }
331
332        // canonical unit representatives, so associates compare equal below.
333        for i in 0..r {
334            let u = self.rows[i][&i].normalizing_unit();
335            if !u.is_one() {
336                self.mul_row(i, &u);
337            }
338        }
339
340        // pre-order by divisibility: when the diagonal values are pairwise
341        // comparable (e.g. powers of H) the bubble pass below finds nothing.
342        self.sort_diag();
343
344        // bubble the divisor chain: full O(r) passes, no restart-from-zero
345        // (the dense algorithm's restart is quadratic on huge diagonals).
346        loop {
347            let mut clean = true;
348            for i in 0..r-1 {
349                if !self.diag_normalize_step(i) {
350                    clean = false;
351                }
352            }
353            if clean {
354                break
355            }
356        }
357
358        // gcd fixes may denormalize units.
359        for i in 0..r {
360            let u = self.rows[i][&i].normalizing_unit();
361            if !u.is_one() {
362                self.mul_row(i, &u);
363            }
364        }
365    }
366
367    // Sorts the diagonal (by row+col swaps) so equal values are contiguous and
368    // strictly-dividing values come first; incomparable pairs keep insertion
369    // order and are resolved by the gcd steps of the bubble pass.
370    fn sort_diag(&mut self) {
371        let r = self.rank;
372
373        let mut chain: Vec<(R, Vec<usize>)> = Vec::new();
374        for i in 0..r {
375            let v = self.rows[i][&i].clone();
376            if let Some((_, grp)) = chain.iter_mut().find(|(w, _)| w == &v) {
377                grp.push(i);
378                continue
379            }
380            let pos = chain.iter().position(|(w, _)| v.divides(w) && !w.divides(&v)).unwrap_or(chain.len());
381            chain.insert(pos, (v, vec![i]));
382        }
383
384        // target[t] = original position whose value should land at t.
385        let target: Vec<usize> = chain.into_iter().flat_map(|(_, grp)| grp).collect();
386
387        let mut pos_of: Vec<usize> = (0..r).collect(); // original -> current
388        let mut at: Vec<usize> = (0..r).collect();     // current -> original
389        for t in 0..r {
390            let cur = pos_of[target[t]];
391            if cur != t {
392                self.swap_rows(t, cur);
393                self.swap_cols(t, cur);
394                let other = at[t];
395                at.swap(t, cur);
396                pos_of[target[t]] = t;
397                pos_of[other] = cur;
398            }
399        }
400    }
401
402    fn diag_normalize_step(&mut self, i: usize) -> bool {
403        let x = self.rows[i][&i].clone();
404        let y = self.rows[i + 1][&(i + 1)].clone();
405
406        debug_assert!(!x.is_zero() && !y.is_zero());
407
408        if x.divides(&y) {
409            return true
410        }
411
412        if y.divides(&x) {
413            self.swap_rows(i, i + 1);
414            self.swap_cols(i, i + 1);
415            return false
416        }
417
418        // sx + ty = d, a = x/d, b = y/d:
419        // [1   1 ][x   ][s  -b]   [d      ]
420        // [-tb sa][   y][t   a] = [   xy/d]
421        let (d, s, t) = gcdx(&x, &y);
422        let (a, b) = (&x / &d, &y / &d);
423        let (tb, sa) = (&t * &b, &s * &a);
424
425        self.row_elem([R::one(), R::one(), -tb, sa], i, i + 1);
426        self.col_elem([s, t, -b, a], i, i + 1);
427
428        false
429    }
430
431    // ---- primitive operations (matrix + realtime trans) ----
432
433    fn entry(&self, i: usize, j: usize) -> Option<&R> {
434        self.rows[i].get(&j)
435    }
436
437    fn swap_rows(&mut self, i: usize, j: usize) {
438        let keys: FxHashSet<usize> = self.rows[i].keys().chain(self.rows[j].keys()).copied().collect();
439        self.rows.swap(i, j);
440        for k in keys {
441            let col = &mut self.cols[k];
442            let (a, b) = (col.contains(&i), col.contains(&j));
443            if a && !b {
444                col.remove(&i);
445                col.insert(j);
446            } else if b && !a {
447                col.remove(&j);
448                col.insert(i);
449            }
450        }
451        if let Some(p) = &mut self.p { p.swap(i, j); }
452        if let Some(pinv) = &mut self.pinv { pinv.swap(i, j); }
453    }
454
455    fn swap_cols(&mut self, i: usize, j: usize) {
456        let rows: Vec<usize> = self.cols[i].union(&self.cols[j]).copied().collect();
457        for r in rows {
458            let vi = self.rows[r].remove(&i);
459            let vj = self.rows[r].remove(&j);
460            if let Some(v) = vj { self.rows[r].insert(i, v); }
461            if let Some(v) = vi { self.rows[r].insert(j, v); }
462        }
463        self.cols.swap(i, j);
464        if let Some(q) = &mut self.q { q.swap(i, j); }
465        if let Some(qinv) = &mut self.qinv { qinv.swap(i, j); }
466    }
467
468    fn mul_row(&mut self, i: usize, u: &R) {
469        debug_assert!(u.is_unit());
470        for (_, v) in self.rows[i].iter_mut() {
471            *v = &*v * u;
472        }
473        if let Some(p) = &mut self.p { p.scale(i, u); }
474        if let Some(pinv) = &mut self.pinv { pinv.scale(i, &u.inv().unwrap()); }
475    }
476
477    fn mul_col(&mut self, j: usize, u: &R) {
478        debug_assert!(u.is_unit());
479        let rows: Vec<usize> = self.cols[j].iter().copied().collect();
480        for r in rows {
481            let v = self.rows[r].get_mut(&j).unwrap();
482            *v = &*v * u;
483        }
484        if let Some(q) = &mut self.q { q.scale(j, u); }
485        if let Some(qinv) = &mut self.qinv { qinv.scale(j, &u.inv().unwrap()); }
486    }
487
488    // rows (i, j) ← (a·rᵢ + b·rⱼ, c·rᵢ + d·rⱼ), det = 1.
489    fn row_elem(&mut self, comps: [R; 4], i: usize, j: usize) {
490        let [a, b, c, d] = &comps;
491        debug_assert!((a * d - b * c).is_one());
492
493        let ri = std::mem::take(&mut self.rows[i]);
494        let rj = std::mem::take(&mut self.rows[j]);
495        let keys: FxHashSet<usize> = ri.keys().chain(rj.keys()).copied().collect();
496
497        let (mut ni, mut nj) = (FxHashMap::default(), FxHashMap::default());
498        for &k in &keys {
499            let x = ri.get(&k);
500            let y = rj.get(&k);
501            let vi = lin(a, x, b, y);
502            let vj = lin(c, x, d, y);
503
504            let col = &mut self.cols[k];
505            if vi.is_zero() { col.remove(&i); } else { col.insert(i); ni.insert(k, vi); }
506            if vj.is_zero() { col.remove(&j); } else { col.insert(j); nj.insert(k, vj); }
507        }
508        self.rows[i] = ni;
509        self.rows[j] = nj;
510
511        if let Some(p) = &mut self.p { p.combine(&comps, i, j); }
512        if let Some(pinv) = &mut self.pinv { pinv.combine(&inv_comps(&comps), i, j); }
513    }
514
515    // cols (i, j) ← (a·cᵢ + b·cⱼ, c·cᵢ + d·cⱼ), det = 1.
516    fn col_elem(&mut self, comps: [R; 4], i: usize, j: usize) {
517        let [a, b, c, d] = &comps;
518        debug_assert!((a * d - b * c).is_one());
519
520        let rows: Vec<usize> = self.cols[i].union(&self.cols[j]).copied().collect();
521        for r in rows {
522            let x = self.rows[r].remove(&i);
523            let y = self.rows[r].remove(&j);
524            let vi = lin(a, x.as_ref(), b, y.as_ref());
525            let vj = lin(c, x.as_ref(), d, y.as_ref());
526
527            if vi.is_zero() { self.cols[i].remove(&r); } else { self.cols[i].insert(r); self.rows[r].insert(i, vi); }
528            if vj.is_zero() { self.cols[j].remove(&r); } else { self.cols[j].insert(r); self.rows[r].insert(j, vj); }
529        }
530
531        if let Some(q) = &mut self.q { q.combine(&comps, i, j); }
532        if let Some(qinv) = &mut self.qinv { qinv.combine(&inv_comps(&comps), i, j); }
533    }
534
535}
536
537// Inverse of an elementary [a,b;c,d] (det 1) is [d,-b;-c,a], which as a
538// combine on the opposite side takes comps [d, -c, -b, a].
539fn inv_comps<R>([a, b, c, d]: &[R; 4]) -> [R; 4]
540where R: EucRing, for<'x> &'x R: EucRingOps<R> {
541    [d.clone(), -c, -b, a.clone()]
542}
543
544// x/d is kept as the Bezout coefficient shortcut when it is a unit:
545// s = (x/d)⁻¹ satisfies s·x = d directly (avoids coefficient growth).
546fn gcdx<R>(x: &R, y: &R) -> (R, R, R)
547where R: EucRing, for<'x> &'x R: EucRingOps<R> {
548    let (d, s, t) = EucRing::gcdx(x, y);
549    let a = x / &d;
550    if a.is_unit() {
551        let s = a.inv().unwrap();
552        (d, s, R::zero())
553    } else {
554        (d, s, t)
555    }
556}
557
558fn lin<R>(a: &R, x: Option<&R>, b: &R, y: Option<&R>) -> R
559where R: EucRing, for<'x> &'x R: EucRingOps<R> {
560    match (x, y) {
561        (Some(x), Some(y)) => a * x + b * y,
562        (Some(x), None) => a * x,
563        (None, Some(y)) => b * y,
564        (None, None) => R::zero(),
565    }
566}
567
568// A sequence of sparse vectors (rows of `p`/`qinv` or columns of `pinv`/`q`)
569// supporting the replayed operations.
570struct VecStore<R> {
571    vecs: Vec<FxHashMap<usize, R>>,
572}
573
574impl<R> VecStore<R>
575where R: EucRing, for<'x> &'x R: EucRingOps<R> {
576    fn id(n: usize) -> Self {
577        let vecs = (0..n).map(|i| {
578            let mut v = FxHashMap::with_capacity_and_hasher(1, Default::default());
579            v.insert(i, R::one());
580            v
581        }).collect();
582        Self { vecs }
583    }
584
585    fn swap(&mut self, i: usize, j: usize) {
586        self.vecs.swap(i, j);
587    }
588
589    fn scale(&mut self, i: usize, u: &R) {
590        for (_, v) in self.vecs[i].iter_mut() {
591            *v = &*v * u;
592        }
593    }
594
595    // (vᵢ, vⱼ) ← (a·vᵢ + b·vⱼ, c·vᵢ + d·vⱼ)
596    fn combine(&mut self, comps: &[R; 4], i: usize, j: usize) {
597        let [a, b, c, d] = comps;
598        let vi = std::mem::take(&mut self.vecs[i]);
599        let vj = std::mem::take(&mut self.vecs[j]);
600        let keys: FxHashSet<usize> = vi.keys().chain(vj.keys()).copied().collect();
601
602        let (mut ni, mut nj) = (FxHashMap::default(), FxHashMap::default());
603        for k in keys {
604            let x = vi.get(&k);
605            let y = vj.get(&k);
606            let wi = lin(a, x, b, y);
607            let wj = lin(c, x, d, y);
608            if !wi.is_zero() { ni.insert(k, wi); }
609            if !wj.is_zero() { nj.insert(k, wj); }
610        }
611        self.vecs[i] = ni;
612        self.vecs[j] = nj;
613    }
614
615    fn into_spmat_rows(self, n_cols: usize) -> SpMat<R> {
616        let m = self.vecs.len();
617        SpMat::from_entries((m, n_cols), self.vecs.into_iter().enumerate().flat_map(|(i, row)|
618            row.into_iter().map(move |(j, v)| (i, j, v))
619        ))
620    }
621
622    fn into_spmat_cols(self, n_rows: usize) -> SpMat<R> {
623        let n = self.vecs.len();
624        SpMat::from_entries((n_rows, n), self.vecs.into_iter().enumerate().flat_map(|(j, col)|
625            col.into_iter().map(move |(i, v)| (i, j, v))
626        ))
627    }
628}
629
630#[cfg(test)]
631mod tests {
632    use super::*;
633    use num_traits::{Pow, Zero};
634    use yui_core::num::FF2;
635    use yui_core::poly::Poly;
636    use crate::dense::snf::snf;
637
638    // exercises both routes: 0 forces the sparse elimination, MAX the dense fast-pass.
639    fn check_snf<R>(a: &SpMat<R>)
640    where R: EucRing, for<'x> &'x R: EucRingOps<R> {
641        for max in [0, usize::MAX] {
642            let s = sp_snf_with(a, [true; 4], max);
643            let res = s.result();
644
645            // diagonal shape
646            for (i, j, v) in res.iter_nz() {
647                assert!(i == j || v.is_zero(), "non-diagonal entry at ({i}, {j}) (max: {max})");
648            }
649
650            // divisor chain
651            let fs = s.factors();
652            for w in fs.windows(2) {
653                assert!(w[0].divides(w[1]), "{} does not divide {} (max: {max})", w[0], w[1]);
654            }
655
656            // p * a * q = result, pinv * result * qinv = a
657            let (p, pinv, q, qinv) = (s.p().unwrap(), s.pinv().unwrap(), s.q().unwrap(), s.qinv().unwrap());
658            assert_eq!(&(&(p * a) * q), res, "p*a*q != result (max: {max})");
659            assert_eq!(&(&(pinv * res) * qinv), a, "pinv*result*qinv != a (max: {max})");
660        }
661    }
662
663    fn check_against_dense<R>(a: &SpMat<R>)
664    where R: EucRing, for<'x> &'x R: EucRingOps<R> {
665        check_snf(a);
666
667        let d = snf(&a.clone().into_dense(), [false; 4]);
668        let df: Vec<&R> = d.factors();
669        for max in [0, usize::MAX] {
670            let s = sp_snf_with(a, [false; 4], max);
671            let sf: Vec<&R> = s.factors();
672            assert_eq!(sf, df, "factors differ from dense snf (max: {max})");
673        }
674    }
675
676    #[test]
677    fn snf_int() {
678        let a: SpMat<i64> = SpMat::from_row_major((3, 3), [1, 2, 3, 4, 5, 6, 7, 8, 9]);
679        check_against_dense(&a);
680    }
681
682    #[test]
683    fn snf_int_tors() {
684        let a: SpMat<i64> = SpMat::from_row_major((5, 5), [
685            -20, -7, -27, 2, 29,
686            17, 8, 14, -4, -10,
687            13, 8, 10, -4, -6,
688            -9, -2, -14, 0, 16,
689            5, 0, 5, -1, -4
690        ]);
691        let s = sp_snf(&a, [true; 4]);
692        check_snf(&a);
693        let fs: Vec<i64> = s.factors().into_iter().cloned().collect();
694        assert_eq!(fs, vec![1, 1, 1, 2, 60]);
695    }
696
697    #[test]
698    fn snf_zero() {
699        let a: SpMat<i64> = SpMat::zero((3, 4));
700        let s = sp_snf(&a, [true; 4]);
701        assert_eq!(s.rank(), 0);
702        check_snf(&a);
703    }
704
705    #[test]
706    fn snf_empty() {
707        let a: SpMat<i64> = SpMat::zero((0, 0));
708        let s = sp_snf(&a, [true; 4]);
709        assert_eq!(s.rank(), 0);
710    }
711
712    #[test]
713    fn snf_rect() {
714        let a: SpMat<i64> = SpMat::from_row_major((2, 4), [2, 4, 6, 8, 3, 5, 7, 9]);
715        check_against_dense(&a);
716    }
717
718    #[test]
719    fn snf_poly_h_multiples() {
720        // the KhI-residual shape: all entries divisible by H, no unit pivots
721        type P = Poly<'H', FF2>;
722        let h = P::variable();
723        let o = P::zero();
724        let h2 = &h * &h;
725        let h3 = &h2 * &h;
726
727        let a: SpMat<P> = SpMat::from_row_major((3, 3), [
728            h.clone(), h2.clone(), o.clone(),
729            o.clone(), h.clone(), h3.clone(),
730            h2.clone(), o.clone(), h.clone(),
731        ]);
732        check_snf(&a);
733    }
734
735    #[test]
736    fn snf_large_monomial_diag() {
737        // 20k-entry shuffled H-power diagonal: normalization must be linear-ish,
738        // not the dense algorithm's restart-quadratic (regression for the
739        // 143k-diagonal stall on 10_162).
740        type P = Poly<'H', FF2>;
741        let n = 20_000;
742        let a: SpMat<P> = SpMat::from_entries((n, n), (0..n).map(|i|
743            (i, i, P::variable().pow((i * 7919) % 23))
744        ));
745        let s = sp_snf_with(&a, [false; 4], 0);
746        assert_eq!(s.rank(), n);
747        let fs = s.factors();
748        for w in fs.windows(2) {
749            assert!(w[0].divides(w[1]));
750        }
751    }
752
753    #[test]
754    fn snf_rand() {
755        let a: SpMat<i64> = SpMat::rand((20, 30), 0.2);
756        check_snf(&a);
757    }
758
759    #[test]
760    fn snf_rand_dense_cmp() {
761        for _ in 0..5 {
762            let a: SpMat<i64> = SpMat::rand((8, 10), 0.4);
763            check_against_dense(&a);
764        }
765    }
766}