From 29e111cc3c9a57722723be6e53ba79f6d784cc6c Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Sat, 4 Apr 2026 13:29:02 -0400 Subject: [PATCH 01/22] refactor: convert to cargo workspace with lib and py crates Split the single cdylib crate into a workspace with two members: - crates/modularsnf: pure Rust library (rlib) - crates/modularsnf-py: PyO3 Python bindings (cdylib) --- Cargo.lock | 10 +++++++++- Cargo.toml | 19 +++---------------- crates/modularsnf-py/Cargo.toml | 18 ++++++++++++++++++ crates/modularsnf/Cargo.toml | 13 +++++++++++++ 4 files changed, 43 insertions(+), 17 deletions(-) create mode 100644 crates/modularsnf-py/Cargo.toml create mode 100644 crates/modularsnf/Cargo.toml diff --git a/Cargo.lock b/Cargo.lock index b5098bd..ecf5101 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -56,9 +56,17 @@ dependencies = [ [[package]] name = "modularsnf" -version = "0.3.0" +version = "0.4.0" dependencies = [ "ndarray", +] + +[[package]] +name = "modularsnf-py" +version = "0.4.0" +dependencies = [ + "modularsnf", + "ndarray", "numpy", "pyo3", ] diff --git a/Cargo.toml b/Cargo.toml index 96c4cce..6655fcf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,19 +1,6 @@ -[package] -name = "modularsnf" -version = "0.3.0" -edition = "2021" -license = "Apache-2.0" -description = "Smith Normal Form calculations over Z/N rings" -repository = "https://github.com/events555/modularsnf" - -[lib] -name = "_rust" -crate-type = ["cdylib"] - -[dependencies] -pyo3 = { version = "0.24", features = ["abi3-py310", "extension-module"] } -numpy = "0.24" -ndarray = "0.16" +[workspace] +members = ["crates/modularsnf", "crates/modularsnf-py"] +resolver = "2" [profile.release] debug = true diff --git a/crates/modularsnf-py/Cargo.toml b/crates/modularsnf-py/Cargo.toml new file mode 100644 index 0000000..fa01c11 --- /dev/null +++ b/crates/modularsnf-py/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "modularsnf-py" +version = "0.4.0" +edition = "2021" +license = "Apache-2.0" +description = "Python bindings for modularsnf" +repository = "https://github.com/events555/modularsnf" +publish = false + +[lib] +name = "_rust" +crate-type = ["cdylib"] + +[dependencies] +modularsnf = { path = "../modularsnf" } +pyo3 = { version = "0.24", features = ["abi3-py310", "extension-module"] } +numpy = "0.24" +ndarray = "0.16" diff --git a/crates/modularsnf/Cargo.toml b/crates/modularsnf/Cargo.toml new file mode 100644 index 0000000..030d62a --- /dev/null +++ b/crates/modularsnf/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "modularsnf" +version = "0.4.0" +edition = "2021" +license = "Apache-2.0" +description = "Smith Normal Form calculations over Z/N rings" +repository = "https://github.com/events555/modularsnf" + +[lib] +name = "modularsnf" + +[dependencies] +ndarray = "0.16" From bdd846822615ff77538f7888febb6ace0819279c Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Sat, 4 Apr 2026 13:29:26 -0400 Subject: [PATCH 02/22] refactor: extract pure Rust lib crate from PyO3 bindings Move algorithm code to crates/modularsnf/src/ and strip all PyO3 dependencies. Rename RustRingZModN to RingZModN, remove *_internal method suffixes, and add a top-level smith_normal_form() entry point. --- {src => crates/modularsnf/src}/band.rs | 10 +- {src => crates/modularsnf/src}/diagonal.rs | 92 ++--------- {src => crates/modularsnf/src}/echelon.rs | 20 +-- crates/modularsnf/src/lib.rs | 40 +++++ {src => crates/modularsnf/src}/ring.rs | 168 +++++---------------- {src => crates/modularsnf/src}/snf.rs | 44 +++--- src/lib.rs | 60 -------- 7 files changed, 127 insertions(+), 307 deletions(-) rename {src => crates/modularsnf/src}/band.rs (96%) rename {src => crates/modularsnf/src}/diagonal.rs (81%) rename {src => crates/modularsnf/src}/echelon.rs (84%) create mode 100644 crates/modularsnf/src/lib.rs rename {src => crates/modularsnf/src}/ring.rs (54%) rename {src => crates/modularsnf/src}/snf.rs (91%) delete mode 100644 src/lib.rs diff --git a/src/band.rs b/crates/modularsnf/src/band.rs similarity index 96% rename from src/band.rs rename to crates/modularsnf/src/band.rs index d5522b7..7374beb 100644 --- a/src/band.rs +++ b/crates/modularsnf/src/band.rs @@ -1,10 +1,10 @@ //! Band reduction — Rust port of modularsnf/band.py. -use numpy::ndarray::{s, Array2}; +use ndarray::{s, Array2}; use crate::diagonal::matmul_mod; use crate::echelon::lemma_3_1; -use crate::ring::RustRingZModN; +use crate::ring::RingZModN; /// Positive modulo helper. #[inline] @@ -30,7 +30,7 @@ fn left_apply_block(m: &mut Array2, block: &Array2, start: usize, n_mo /// Triang step: triangulate top-right block of upper-b-banded matrix. /// Returns (B_prime, W) where W is the s2 x s2 right transform. -fn triang(b_mat: &Array2, b: usize, ring: &RustRingZModN) -> (Array2, Array2) { +fn triang(b_mat: &Array2, b: usize, ring: &RingZModN) -> (Array2, Array2) { let n_mod = ring.n(); let s1 = b / 2; let s2 = b - 1; @@ -59,7 +59,7 @@ fn triang(b_mat: &Array2, b: usize, ring: &RustRingZModN) -> (Array2, fn shift( c_mat: &Array2, b: usize, - ring: &RustRingZModN, + ring: &RingZModN, ) -> (Array2, Array2, Array2) { let n_mod = ring.n(); let s2 = b - 1; @@ -120,7 +120,7 @@ pub fn band_reduction( a: &Array2, b: usize, t_param: usize, - ring: &RustRingZModN, + ring: &RingZModN, ) -> (Array2, Array2, Array2, usize) { let n_mod = ring.n(); let n = a.nrows(); diff --git a/src/diagonal.rs b/crates/modularsnf/src/diagonal.rs similarity index 81% rename from src/diagonal.rs rename to crates/modularsnf/src/diagonal.rs index 8f6ed85..4b9d80e 100644 --- a/src/diagonal.rs +++ b/crates/modularsnf/src/diagonal.rs @@ -1,13 +1,11 @@ //! Diagonal merge for Smith Normal Form — Rust port of modularsnf/diagonal.py. //! //! Implements _merge_scalars_raw, _merge_raw, and _smith_from_diagonal_raw -//! entirely in Rust, returning results as numpy arrays to Python. +//! entirely in Rust. -use numpy::ndarray::{s, Array2}; -use numpy::{IntoPyArray, PyArray2, PyReadonlyArray2}; -use pyo3::prelude::*; +use ndarray::{s, Array2}; -use crate::ring::{mul_mod, posmod_i128, RustRingZModN}; +use crate::ring::{mul_mod, posmod_i128, RingZModN}; /// Positive modulo. #[inline] @@ -16,16 +14,16 @@ fn posmod(a: i64, n: i64) -> i64 { } /// Merge two scalar SNF entries. Returns (U, V, S) as 2x2 arrays. -fn merge_scalars(a: i64, b: i64, ring: &RustRingZModN) -> (Array2, Array2, Array2) { +fn merge_scalars(a: i64, b: i64, ring: &RingZModN) -> (Array2, Array2, Array2) { let n = ring.n(); - let (g, s, t, u, v) = ring.gcdex_internal(a, b); + let (g, s, t, u, v) = ring.gcdex(a, b); if g % n == 0 { return (Array2::eye(2), Array2::eye(2), Array2::zeros((2, 2))); } let tb = mul_mod(t, b, n); - let q_raw = ring.div_internal(tb, g).unwrap_or(0); + let q_raw = ring.div(tb, g).unwrap_or(0); let q = posmod(-q_raw, n); let u_arr = Array2::from_shape_vec( @@ -152,10 +150,10 @@ fn matmul_mod_add(a: &Array2, b: &Array2, n: i64) -> Array2 { } /// Recursive merge of two SNF blocks. Returns (U, V, S) as 2n x 2n arrays. -fn merge_raw( +pub fn merge_raw( a_arr: &Array2, b_arr: &Array2, - ring: &RustRingZModN, + ring: &RingZModN, ) -> (Array2, Array2, Array2) { let n_mod = ring.n(); let n = a_arr.nrows(); @@ -278,7 +276,7 @@ fn merge_raw( /// Bottom-up iterative diagonal SNF for power-of-two matrices. fn smith_from_diagonal_raw( diag: &Array2, - ring: &RustRingZModN, + ring: &RingZModN, ) -> (Array2, Array2, Array2) { let n_mod = ring.n(); let n = diag.nrows(); @@ -324,11 +322,11 @@ fn smith_from_diagonal_raw( blocks.into_iter().next().unwrap() } -/// Public entry for diagonal SNF from other Rust modules. +/// Compute SNF of a diagonal matrix. /// Pads to power-of-two, runs merge, crops back. -pub fn smith_from_diagonal_internal( +pub fn smith_from_diagonal( diag: &Array2, - ring: &RustRingZModN, + ring: &RingZModN, ) -> (Array2, Array2, Array2) { let n = diag.nrows(); if n <= 1 { @@ -346,69 +344,3 @@ pub fn smith_from_diagonal_internal( subblock(&s, 0, n, 0, n), ) } - -// ---- PyO3 exports ---- - -/// Compute SNF of a diagonal matrix. Takes a flat diagonal array and modulus. -/// Returns (U, V, S) as numpy arrays, cropped to original size. -#[pyfunction] -pub fn rust_smith_from_diagonal<'py>( - py: Python<'py>, - diag_data: PyReadonlyArray2<'py, i64>, - modulus: i64, -) -> PyResult<( - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, -)> { - let ring = RustRingZModN::new_internal(modulus)?; - let arr = diag_data.as_array().to_owned(); - let n = arr.nrows(); - - // Pad to power of two - let size = if n <= 1 { - n - } else { - (n as u64).next_power_of_two() as usize - }; - let mut pad = Array2::zeros((size, size)); - for i in 0..n { - for j in 0..n { - pad[[i, j]] = arr[[i, j]]; - } - } - - let (u, v, s) = smith_from_diagonal_raw(&pad, &ring); - - // Crop to original size - let u_crop = subblock(&u, 0, n, 0, n); - let v_crop = subblock(&v, 0, n, 0, n); - let s_crop = subblock(&s, 0, n, 0, n); - - Ok(( - u_crop.into_pyarray(py), - v_crop.into_pyarray(py), - s_crop.into_pyarray(py), - )) -} - -/// Merge two SNF blocks. Returns (U, V, S) as numpy arrays. -#[pyfunction] -pub fn rust_merge_smith_blocks<'py>( - py: Python<'py>, - a_data: PyReadonlyArray2<'py, i64>, - b_data: PyReadonlyArray2<'py, i64>, - modulus: i64, -) -> PyResult<( - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, -)> { - let ring = RustRingZModN::new_internal(modulus)?; - let a = a_data.as_array().to_owned(); - let b = b_data.as_array().to_owned(); - - let (u, v, s) = merge_raw(&a, &b, &ring); - - Ok((u.into_pyarray(py), v.into_pyarray(py), s.into_pyarray(py))) -} diff --git a/src/echelon.rs b/crates/modularsnf/src/echelon.rs similarity index 84% rename from src/echelon.rs rename to crates/modularsnf/src/echelon.rs index 360b462..b45aae8 100644 --- a/src/echelon.rs +++ b/crates/modularsnf/src/echelon.rs @@ -1,8 +1,8 @@ //! Row-echelon utilities — Rust port of modularsnf/echelon.py. -use numpy::ndarray::Array2; +use ndarray::Array2; -use crate::ring::{posmod_i128, RustRingZModN}; +use crate::ring::{posmod_i128, RingZModN}; /// Apply a 2x2 row transform to rows r0, r1 of both matrices. /// [s t] [row_r0] [new_r0] @@ -46,7 +46,7 @@ pub fn apply_row_2x2_pair( /// Lemma 3.1: row-echelon form via extended GCD elimination. /// Returns (U, T, rank). -pub fn lemma_3_1(a: &Array2, ring: &RustRingZModN) -> (Array2, Array2, usize) { +pub fn lemma_3_1(a: &Array2, ring: &RingZModN) -> (Array2, Array2, usize) { let n_mod = ring.n(); let n_rows = a.nrows(); let n_cols = a.ncols(); @@ -65,15 +65,15 @@ pub fn lemma_3_1(a: &Array2, ring: &RustRingZModN) -> (Array2, Array2< let a_val = t[[r, k]]; let b_val = t[[i, k]]; - if ring.is_zero_internal(b_val) { + if ring.is_zero(b_val) { continue; } - let (_, s, tv, uv, v) = ring.gcdex_internal(a_val, b_val); + let (_, s, tv, uv, v) = ring.gcdex(a_val, b_val); apply_row_2x2_pair(&mut t, &mut u, r, i, s, tv, uv, v, n_mod); } - if !ring.is_zero_internal(t[[r, k]]) { + if !ring.is_zero(t[[r, k]]) { r += 1; } } @@ -86,7 +86,7 @@ pub fn lemma_3_1(a: &Array2, ring: &RustRingZModN) -> (Array2, Array2< pub fn index1_reduce_on_columns( a: &Array2, k: usize, - ring: &RustRingZModN, + ring: &RingZModN, ) -> (Array2, Array2) { let n_mod = ring.n(); let n = a.nrows(); @@ -98,11 +98,11 @@ pub fn index1_reduce_on_columns( let sj = t[[j, j]]; for i in 0..j { let x = t[[i, j]]; - if ring.is_zero_internal(x) { + if ring.is_zero(x) { continue; } let rem = { - let b_ass = ring.gcd_internal(sj, 0); + let b_ass = ring.gcd(sj, 0); let x_val = ((x % n_mod) + n_mod) % n_mod; if b_ass == 0 { x_val @@ -111,7 +111,7 @@ pub fn index1_reduce_on_columns( } }; let diff = ((x - rem) % n_mod + n_mod) % n_mod; - let quo = ring.div_internal(diff, sj).unwrap_or(0); + let quo = ring.div(diff, sj).unwrap_or(0); let phi = ((-quo) % n_mod + n_mod) % n_mod; let cols = t.ncols(); diff --git a/crates/modularsnf/src/lib.rs b/crates/modularsnf/src/lib.rs new file mode 100644 index 0000000..24eedad --- /dev/null +++ b/crates/modularsnf/src/lib.rs @@ -0,0 +1,40 @@ +pub mod band; +pub mod diagonal; +pub mod echelon; +pub mod ring; +pub mod snf; + +pub use ring::RingZModN; + +use ndarray::Array2; + +/// Full Smith Normal Form: takes an n x m matrix and modulus, +/// returns (U, V, S) with S = U @ A @ V (mod N). +pub fn smith_normal_form( + a: &Array2, + modulus: i64, +) -> Result<(Array2, Array2, Array2), String> { + let r = RingZModN::new(modulus)?; + let n = a.nrows(); + let m = a.ncols(); + + if n == 0 || m == 0 { + let u = Array2::::eye(n); + let v = Array2::::eye(m); + return Ok((u, v, a.clone())); + } + + // Pad to square if needed + let s = n.max(m); + let mut a_pad = Array2::zeros((s, s)); + a_pad.slice_mut(ndarray::s![..n, ..m]).assign(a); + + let (u_pad, v_pad, s_pad) = snf::smith_square(&a_pad, &r); + + // Crop back + let u = u_pad.slice(ndarray::s![..n, ..n]).to_owned(); + let v = v_pad.slice(ndarray::s![..m, ..m]).to_owned(); + let s_mat = s_pad.slice(ndarray::s![..n, ..m]).to_owned(); + + Ok((u, v, s_mat)) +} diff --git a/src/ring.rs b/crates/modularsnf/src/ring.rs similarity index 54% rename from src/ring.rs rename to crates/modularsnf/src/ring.rs index 736b6b1..89c3711 100644 --- a/src/ring.rs +++ b/crates/modularsnf/src/ring.rs @@ -2,9 +2,6 @@ //! //! All operations match Storjohann's Dissertation Section 1.1. -use pyo3::exceptions::PyValueError; -use pyo3::prelude::*; - /// Extended GCD: returns (g, x, y) such that a*x + b*y = g. #[inline] fn egcd(a: i64, b: i64) -> (i64, i64, i64) { @@ -27,7 +24,7 @@ fn egcd(a: i64, b: i64) -> (i64, i64, i64) { } #[inline] -fn gcd(a: i64, b: i64) -> i64 { +fn gcd_raw(a: i64, b: i64) -> i64 { let (mut a, mut b) = (a.abs(), b.abs()); while b != 0 { let t = b; @@ -65,56 +62,53 @@ pub(crate) fn mul_mod(a: i64, b: i64, n: i64) -> i64 { posmod_i128((a as i128) * (b as i128), n) } -#[pyclass] -pub struct RustRingZModN { +pub struct RingZModN { n: i64, } -#[pymethods] -impl RustRingZModN { - #[new] - fn new(n: i64) -> PyResult { +impl RingZModN { + pub fn new(n: i64) -> Result { if n < 2 { - return Err(PyValueError::new_err("Modulus N must be >= 2")); + return Err("Modulus N must be >= 2".to_string()); } Ok(Self { n }) } - #[getter] - fn n_val(&self) -> i64 { + #[inline] + pub fn n(&self) -> i64 { self.n } - fn add(&self, a: i64, b: i64) -> i64 { + pub fn add(&self, a: i64, b: i64) -> i64 { add_mod(a, b, self.n) } - fn sub(&self, a: i64, b: i64) -> i64 { + pub fn sub(&self, a: i64, b: i64) -> i64 { sub_mod(a, b, self.n) } - fn mul(&self, a: i64, b: i64) -> i64 { + pub fn mul(&self, a: i64, b: i64) -> i64 { mul_mod(a, b, self.n) } - fn is_zero(&self, a: i64) -> bool { + pub fn is_zero(&self, a: i64) -> bool { posmod(a, self.n) == 0 } - fn gcd(&self, a: i64, b: i64) -> i64 { - gcd(posmod(a, self.n), gcd(posmod(b, self.n), self.n)) + pub fn gcd(&self, a: i64, b: i64) -> i64 { + gcd_raw(posmod(a, self.n), gcd_raw(posmod(b, self.n), self.n)) } - fn ass(&self, a: i64) -> i64 { - gcd(posmod(a, self.n), self.n) + pub fn ass(&self, a: i64) -> i64 { + gcd_raw(posmod(a, self.n), self.n) } - fn ann(&self, a: i64) -> i64 { - let g = gcd(posmod(a, self.n), self.n); + pub fn ann(&self, a: i64) -> i64 { + let g = gcd_raw(posmod(a, self.n), self.n); posmod(self.n / g, self.n) } - fn rem(&self, a: i64, b: i64) -> i64 { + pub fn rem(&self, a: i64, b: i64) -> i64 { let b_ass = self.ass(b); let a_val = posmod(a, self.n); if b_ass == 0 { @@ -124,39 +118,39 @@ impl RustRingZModN { } } - fn quo(&self, a: i64, b: i64) -> PyResult { + pub fn quo(&self, a: i64, b: i64) -> Result { let r = self.rem(a, b); let diff = self.sub(a, r); self.div(diff, b) } - fn div(&self, a: i64, b: i64) -> PyResult { + pub fn div(&self, a: i64, b: i64) -> Result { let a_val = posmod(a, self.n); let b_val = posmod(b, self.n); let (g, x, _) = egcd(b_val, self.n); if g == 0 || a_val % g != 0 { - return Err(PyValueError::new_err(format!( + return Err(format!( "{a} not divisible by {b} in Z/{}", self.n - ))); + )); } Ok(posmod_i128((x as i128) * ((a_val / g) as i128), self.n / g)) } - fn unit(&self, a: i64) -> i64 { + pub fn unit(&self, a: i64) -> i64 { let a_val = posmod(a, self.n); let (_, u, _) = egcd(a_val, self.n); posmod(u, self.n) } - fn gcdex(&self, a: i64, b: i64) -> PyResult<(i64, i64, i64, i64, i64)> { + pub fn gcdex(&self, a: i64, b: i64) -> (i64, i64, i64, i64, i64) { let a_val = posmod(a, self.n); let b_val = posmod(b, self.n); // Fast path: b is a multiple of a in Z/N - if a_val != 0 && (b_val % gcd(a_val, self.n) == 0) { + if a_val != 0 && (b_val % gcd_raw(a_val, self.n) == 0) { if let Ok(q) = self.div(b_val, a_val) { - return Ok((a_val, 1, 0, posmod(-q, self.n), 1)); + return (a_val, 1, 0, posmod(-q, self.n), 1); } } @@ -178,27 +172,23 @@ impl RustRingZModN { t0 = tmp; } - let g = r0; - let s = s0; - let t = t0; - - if g == 0 { - return Ok((0, 1, 0, 0, 1)); + if r0 == 0 { + return (0, 1, 0, 0, 1); } - let u = -(b_val / g); - let v = a_val / g; + let u = -(b_val / r0); + let v = a_val / r0; - Ok(( - posmod(g, self.n), - posmod(s, self.n), - posmod(t, self.n), + ( + posmod(r0, self.n), + posmod(s0, self.n), + posmod(t0, self.n), posmod(u, self.n), posmod(v, self.n), - )) + ) } - fn stab(&self, a: i64, b: i64, c: i64) -> PyResult { + pub fn stab(&self, a: i64, b: i64, c: i64) -> Result { let a = posmod(a, self.n); let b = posmod(b, self.n); let c = posmod(c, self.n); @@ -210,91 +200,9 @@ impl RustRingZModN { return Ok(x); } } - Err(PyValueError::new_err(format!( + Err(format!( "Stab failed for a={a}, b={b}, c={c} in Z/{}", self.n - ))) - } -} - -// Internal-only helpers for use from other Rust modules. -impl RustRingZModN { - pub fn new_internal(n: i64) -> PyResult { - if n < 2 { - return Err(pyo3::exceptions::PyValueError::new_err( - "Modulus N must be >= 2", - )); - } - Ok(Self { n }) - } - - #[inline] - pub fn n(&self) -> i64 { - self.n - } - - #[inline] - pub fn div_internal(&self, a: i64, b: i64) -> Result { - let a_val = posmod(a, self.n); - let b_val = posmod(b, self.n); - let (g, x, _) = egcd(b_val, self.n); - if g == 0 || a_val % g != 0 { - return Err(format!("{a} not divisible by {b} in Z/{}", self.n)); - } - Ok(posmod_i128((x as i128) * ((a_val / g) as i128), self.n / g)) - } - - #[inline] - pub fn gcdex_internal(&self, a: i64, b: i64) -> (i64, i64, i64, i64, i64) { - let a_val = posmod(a, self.n); - let b_val = posmod(b, self.n); - - if a_val != 0 && (b_val % gcd(a_val, self.n) == 0) { - if let Ok(q) = self.div_internal(b_val, a_val) { - return (a_val, 1, 0, posmod(-q, self.n), 1); - } - } - - let (mut r0, mut r1) = (a_val, b_val); - let (mut s0, mut s1) = (1i64, 0i64); - let (mut t0, mut t1) = (0i64, 1i64); - - while r1 != 0 { - let q = r0 / r1; - let tmp = r1; - r1 = r0 - q * r1; - r0 = tmp; - let tmp = s1; - s1 = s0 - q * s1; - s0 = tmp; - let tmp = t1; - t1 = t0 - q * t1; - t0 = tmp; - } - - if r0 == 0 { - return (0, 1, 0, 0, 1); - } - - let u = -(b_val / r0); - let v = a_val / r0; - - ( - posmod(r0, self.n), - posmod(s0, self.n), - posmod(t0, self.n), - posmod(u, self.n), - posmod(v, self.n), - ) - } - - #[inline] - pub fn gcd_internal(&self, a: i64, b: i64) -> i64 { - gcd(posmod(a, self.n), gcd(posmod(b, self.n), self.n)) - } - - #[inline] - pub fn is_zero_internal(&self, a: i64) -> bool { - posmod(a, self.n) == 0 + )) } } diff --git a/src/snf.rs b/crates/modularsnf/src/snf.rs similarity index 91% rename from src/snf.rs rename to crates/modularsnf/src/snf.rs index a27d39f..834d222 100644 --- a/src/snf.rs +++ b/crates/modularsnf/src/snf.rs @@ -1,11 +1,11 @@ //! Smith Normal Form pipeline — Rust port of modularsnf/snf.py. -use numpy::ndarray::{s, Array2}; +use ndarray::{s, Array2}; use crate::band::{band_reduction, compute_upper_bandwidth}; use crate::diagonal::matmul_mod; use crate::echelon::{apply_row_2x2_pair, index1_reduce_on_columns, lemma_3_1}; -use crate::ring::{posmod_i128, RustRingZModN}; +use crate::ring::{posmod_i128, RingZModN}; /// Positive modulo. #[inline] @@ -17,7 +17,7 @@ fn posmod(a: i64, n: i64) -> i64 { /// Returns (U, V, S) such that S = U @ A @ V. pub fn smith_square( a: &Array2, - ring: &RustRingZModN, + ring: &RingZModN, ) -> (Array2, Array2, Array2) { let n_mod = ring.n(); let n = a.nrows(); @@ -55,7 +55,7 @@ pub fn smith_square( /// Smith form for a 2-banded square matrix. fn smith_from_upper_2_banded( a: &Array2, - ring: &RustRingZModN, + ring: &RingZModN, ) -> (Array2, Array2, Array2) { let n_mod = ring.n(); let n = a.nrows(); @@ -115,7 +115,7 @@ fn smith_from_upper_2_banded( /// Step 1: split 2-banded matrix into spike + principal blocks. fn step1_split_with_spike( a: &Array2, - ring: &RustRingZModN, + ring: &RingZModN, ) -> (Array2, Array2, Array2, usize, usize) { let n_mod = ring.n(); let n = a.nrows(); @@ -130,7 +130,7 @@ fn step1_split_with_spike( let (r0, r1) = (k - 1, k); let a_val = t[[r0, k]]; let b_val = t[[r1, k]]; - let (_, s, tv, uv, v) = ring.gcdex_internal(a_val, b_val); + let (_, s, tv, uv, v) = ring.gcdex(a_val, b_val); // Note: swapped order (u,v,s,t) vs (s,t,u,v) per Python code apply_row_2x2_pair(&mut t, &mut u, r0, r1, uv, v, s, tv, n_mod); } @@ -140,7 +140,7 @@ fn step1_split_with_spike( let (r0, r1) = (k, k + 1); let a_val = t[[r0, k]]; let b_val = t[[r1, k]]; - let (_, s, tv, uv, v) = ring.gcdex_internal(a_val, b_val); + let (_, s, tv, uv, v) = ring.gcdex(a_val, b_val); apply_row_2x2_pair(&mut t, &mut u, r0, r1, s, tv, uv, v, n_mod); } @@ -153,7 +153,7 @@ fn step2_recursive_blocks( a: &Array2, n1: usize, _n2: usize, - ring: &RustRingZModN, + ring: &RingZModN, ) -> (Array2, Array2, Array2) { let n_mod = ring.n(); let n = a.nrows(); @@ -223,7 +223,7 @@ fn step3_permute(a: &Array2, n1: usize) -> (Array2, Array2, Array /// Step 4: diagonalize (n-1) x (n-1) principal block. fn step4_smith_on_n_minus_1( a: &Array2, - ring: &RustRingZModN, + ring: &RingZModN, ) -> (Array2, Array2, Array2) { let n_mod = ring.n(); let n = a.nrows(); @@ -231,8 +231,8 @@ fn step4_smith_on_n_minus_1( let b = a.slice(s![..n - 1, ..n - 1]).to_owned(); // Use diagonal SNF - use crate::diagonal::smith_from_diagonal_internal; - let (u_loc, v_loc, _s_loc) = smith_from_diagonal_internal(&b, ring); + use crate::diagonal::smith_from_diagonal; + let (u_loc, v_loc, _s_loc) = smith_from_diagonal(&b, ring); let mut u = Array2::::eye(n); let mut v = Array2::::eye(n); @@ -246,7 +246,7 @@ fn step4_smith_on_n_minus_1( /// Steps 5-8: gcd chain. fn step5_to_8_gcd_chain( a: &Array2, - ring: &RustRingZModN, + ring: &RingZModN, ) -> (Array2, Array2, Array2, usize) { let n_mod = ring.n(); let n = a.nrows(); @@ -258,7 +258,7 @@ fn step5_to_8_gcd_chain( // Find idx_k: first zero diagonal entry let mut idx_k = n - 1; for i in 0..n - 1 { - if ring.is_zero_internal(t[[i, i]]) { + if ring.is_zero(t[[i, i]]) { idx_k = i; break; } @@ -269,11 +269,11 @@ fn step5_to_8_gcd_chain( // Step 5: eliminate entries below idx_k in last column for row in idx_k + 1..n { let target_val = t[[row, last_col]]; - if ring.is_zero_internal(target_val) { + if ring.is_zero(target_val) { continue; } let pivot_val = t[[idx_k, last_col]]; - let (_, s, tv, uv, vv) = ring.gcdex_internal(pivot_val, target_val); + let (_, s, tv, uv, vv) = ring.gcdex(pivot_val, target_val); apply_row_2x2_pair(&mut t, &mut u, idx_k, row, s, tv, uv, vv, n_mod); } @@ -300,11 +300,11 @@ fn step5_to_8_gcd_chain( // stab let c = { - let target_gcd = ring.gcd_internal(a_ik, ring.gcd_internal(a_i1k, a_ii)); + let target_gcd = ring.gcd(a_ik, ring.gcd(a_i1k, a_ii)); let mut found = 0i64; for x in 0..n_mod { let candidate = posmod_i128((a_ik as i128) + (x as i128) * (a_i1k as i128), n_mod); - let current = ring.gcd_internal(candidate, a_ii); + let current = ring.gcd(candidate, a_ii); if current == target_gcd { found = x; break; @@ -317,7 +317,7 @@ fn step5_to_8_gcd_chain( let numerator = posmod_i128((c as i128) * (s_next as i128), n_mod); // quo - let a_ii_ass = ring.gcd_internal(a_ii, 0); + let a_ii_ass = ring.gcd(a_ii, 0); let num_mod = posmod(numerator, n_mod); let rem = if a_ii_ass == 0 { num_mod @@ -325,7 +325,7 @@ fn step5_to_8_gcd_chain( num_mod % a_ii_ass }; let diff = posmod(num_mod - rem, n_mod); - let q_raw = ring.div_internal(diff, a_ii).unwrap_or(0); + let q_raw = ring.div(diff, a_ii).unwrap_or(0); let q = posmod(-q_raw, n_mod); // Op 1: add c * row[i+1] to row[i] @@ -349,11 +349,11 @@ fn step5_to_8_gcd_chain( let pivot = t[[i, i]]; let target = t[[i, col_target]]; - if ring.is_zero_internal(target) { + if ring.is_zero(target) { continue; } - let (_, s, tv, uv, vv) = ring.gcdex_internal(pivot, target); + let (_, s, tv, uv, vv) = ring.gcdex(pivot, target); // Column operations for row in 0..n { @@ -388,7 +388,7 @@ fn step5_to_8_gcd_chain( fn step9_index_reduction( a: &Array2, k: usize, - ring: &RustRingZModN, + ring: &RingZModN, ) -> (Array2, Array2, Array2) { let n_mod = ring.n(); let n = a.nrows(); diff --git a/src/lib.rs b/src/lib.rs deleted file mode 100644 index 76b3170..0000000 --- a/src/lib.rs +++ /dev/null @@ -1,60 +0,0 @@ -pub mod band; -pub mod diagonal; -pub mod echelon; -pub mod ring; -pub mod snf; - -use numpy::{IntoPyArray, PyArray2, PyReadonlyArray2}; -use pyo3::prelude::*; - -/// Native Rust acceleration for modularsnf. -#[pymodule] -fn _rust(m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add_class::()?; - m.add_function(wrap_pyfunction!(diagonal::rust_smith_from_diagonal, m)?)?; - m.add_function(wrap_pyfunction!(diagonal::rust_merge_smith_blocks, m)?)?; - m.add_function(wrap_pyfunction!(rust_smith_normal_form, m)?)?; - Ok(()) -} - -/// Full Smith Normal Form: takes an n×n matrix and modulus, -/// returns (U, V, S) as numpy arrays with S = U @ A @ V (mod N). -#[pyfunction] -fn rust_smith_normal_form<'py>( - py: Python<'py>, - data: PyReadonlyArray2<'py, i64>, - modulus: i64, -) -> PyResult<( - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, -)> { - let r = ring::RustRingZModN::new_internal(modulus)?; - let a = data.as_array().to_owned(); - let n = a.nrows(); - let m = a.ncols(); - - if n == 0 || m == 0 { - let u = numpy::ndarray::Array2::::eye(n); - let v = numpy::ndarray::Array2::::eye(m); - return Ok((u.into_pyarray(py), v.into_pyarray(py), a.into_pyarray(py))); - } - - // Pad to square if needed - let s = n.max(m); - let mut a_pad = numpy::ndarray::Array2::zeros((s, s)); - a_pad.slice_mut(numpy::ndarray::s![..n, ..m]).assign(&a); - - let (u_pad, v_pad, s_pad) = snf::smith_square(&a_pad, &r); - - // Crop back - let u = u_pad.slice(numpy::ndarray::s![..n, ..n]).to_owned(); - let v = v_pad.slice(numpy::ndarray::s![..m, ..m]).to_owned(); - let s_mat = s_pad.slice(numpy::ndarray::s![..n, ..m]).to_owned(); - - Ok(( - u.into_pyarray(py), - v.into_pyarray(py), - s_mat.into_pyarray(py), - )) -} From 4b5cb61ddd9f4ad86e3df4f30a1d5d1921b233ed Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Sat, 4 Apr 2026 13:29:45 -0400 Subject: [PATCH 03/22] feat: add PyO3 wrapper crate delegating to modularsnf lib Thin bindings in crates/modularsnf-py that wrap the pure Rust API with #[pyclass]/#[pyfunction] and numpy conversions. Python API is unchanged. --- crates/modularsnf-py/src/lib.rs | 176 ++++++++++++++++++++++++++++++++ 1 file changed, 176 insertions(+) create mode 100644 crates/modularsnf-py/src/lib.rs diff --git a/crates/modularsnf-py/src/lib.rs b/crates/modularsnf-py/src/lib.rs new file mode 100644 index 0000000..4493618 --- /dev/null +++ b/crates/modularsnf-py/src/lib.rs @@ -0,0 +1,176 @@ +use modularsnf::ring::RingZModN; +use modularsnf::snf::smith_square; +use modularsnf::diagonal::{merge_raw, smith_from_diagonal}; + +use numpy::ndarray::{s, Array2}; +use numpy::{IntoPyArray, PyArray2, PyReadonlyArray2}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; + +/// Wrapper around RingZModN for Python. +#[pyclass] +#[pyo3(name = "RustRingZModN")] +struct PyRingZModN { + inner: RingZModN, +} + +#[pymethods] +impl PyRingZModN { + #[new] + fn new(n: i64) -> PyResult { + let inner = RingZModN::new(n).map_err(PyValueError::new_err)?; + Ok(Self { inner }) + } + + #[getter] + fn n_val(&self) -> i64 { + self.inner.n() + } + + fn add(&self, a: i64, b: i64) -> i64 { + self.inner.add(a, b) + } + + fn sub(&self, a: i64, b: i64) -> i64 { + self.inner.sub(a, b) + } + + fn mul(&self, a: i64, b: i64) -> i64 { + self.inner.mul(a, b) + } + + fn is_zero(&self, a: i64) -> bool { + self.inner.is_zero(a) + } + + fn gcd(&self, a: i64, b: i64) -> i64 { + self.inner.gcd(a, b) + } + + fn ass(&self, a: i64) -> i64 { + self.inner.ass(a) + } + + fn ann(&self, a: i64) -> i64 { + self.inner.ann(a) + } + + fn rem(&self, a: i64, b: i64) -> i64 { + self.inner.rem(a, b) + } + + fn quo(&self, a: i64, b: i64) -> PyResult { + self.inner.quo(a, b).map_err(PyValueError::new_err) + } + + fn div(&self, a: i64, b: i64) -> PyResult { + self.inner.div(a, b).map_err(PyValueError::new_err) + } + + fn unit(&self, a: i64) -> i64 { + self.inner.unit(a) + } + + fn gcdex(&self, a: i64, b: i64) -> PyResult<(i64, i64, i64, i64, i64)> { + Ok(self.inner.gcdex(a, b)) + } + + fn stab(&self, a: i64, b: i64, c: i64) -> PyResult { + self.inner.stab(a, b, c).map_err(PyValueError::new_err) + } +} + +/// Full Smith Normal Form: takes an n x m matrix and modulus, +/// returns (U, V, S) as numpy arrays with S = U @ A @ V (mod N). +#[pyfunction] +fn rust_smith_normal_form<'py>( + py: Python<'py>, + data: PyReadonlyArray2<'py, i64>, + modulus: i64, +) -> PyResult<( + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, +)> { + let r = RingZModN::new(modulus).map_err(PyValueError::new_err)?; + let a = data.as_array().to_owned(); + let n = a.nrows(); + let m = a.ncols(); + + if n == 0 || m == 0 { + let u = Array2::::eye(n); + let v = Array2::::eye(m); + return Ok((u.into_pyarray(py), v.into_pyarray(py), a.into_pyarray(py))); + } + + // Pad to square if needed + let sz = n.max(m); + let mut a_pad = Array2::zeros((sz, sz)); + a_pad.slice_mut(s![..n, ..m]).assign(&a); + + let (u_pad, v_pad, s_pad) = smith_square(&a_pad, &r); + + // Crop back + let u = u_pad.slice(s![..n, ..n]).to_owned(); + let v = v_pad.slice(s![..m, ..m]).to_owned(); + let s_mat = s_pad.slice(s![..n, ..m]).to_owned(); + + Ok(( + u.into_pyarray(py), + v.into_pyarray(py), + s_mat.into_pyarray(py), + )) +} + +/// Compute SNF of a diagonal matrix. +#[pyfunction] +fn rust_smith_from_diagonal<'py>( + py: Python<'py>, + diag_data: PyReadonlyArray2<'py, i64>, + modulus: i64, +) -> PyResult<( + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, +)> { + let ring = RingZModN::new(modulus).map_err(PyValueError::new_err)?; + let arr = diag_data.as_array().to_owned(); + let (u, v, s) = smith_from_diagonal(&arr, &ring); + + Ok(( + u.into_pyarray(py), + v.into_pyarray(py), + s.into_pyarray(py), + )) +} + +/// Merge two SNF blocks. +#[pyfunction] +fn rust_merge_smith_blocks<'py>( + py: Python<'py>, + a_data: PyReadonlyArray2<'py, i64>, + b_data: PyReadonlyArray2<'py, i64>, + modulus: i64, +) -> PyResult<( + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, +)> { + let ring = RingZModN::new(modulus).map_err(PyValueError::new_err)?; + let a = a_data.as_array().to_owned(); + let b = b_data.as_array().to_owned(); + + let (u, v, s) = merge_raw(&a, &b, &ring); + + Ok((u.into_pyarray(py), v.into_pyarray(py), s.into_pyarray(py))) +} + +/// Native Rust acceleration for modularsnf. +#[pymodule] +fn _rust(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_function(wrap_pyfunction!(rust_smith_from_diagonal, m)?)?; + m.add_function(wrap_pyfunction!(rust_merge_smith_blocks, m)?)?; + m.add_function(wrap_pyfunction!(rust_smith_normal_form, m)?)?; + Ok(()) +} From c4803f4358d0ecfaa5255d4d9b138b2c52c68a51 Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Sat, 4 Apr 2026 13:30:11 -0400 Subject: [PATCH 04/22] build: point maturin at py crate and bump version to 0.4.0 --- pyproject.toml | 3 ++- uv.lock | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 51c66b6..7935b0f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "modularsnf" -version = "0.3.0" +version = "0.4.0" description = "A Python module for Smith Normal Form calculations over Z/N rings" readme = "README.md" license = "Apache-2.0" @@ -29,6 +29,7 @@ testpaths = [ markers = ["perf: performance sanity checks (run with: uv run pytest -m perf)"] [tool.maturin] +manifest-path = "crates/modularsnf-py/Cargo.toml" module-name = "modularsnf._rust" compatibility = "manylinux_2_28" diff --git a/uv.lock b/uv.lock index 564c99b..e8209bf 100644 --- a/uv.lock +++ b/uv.lock @@ -180,7 +180,7 @@ wheels = [ [[package]] name = "modularsnf" -version = "0.2.0" +version = "0.4.0" source = { editable = "." } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, From bbaac9d44edccd6a68443fe1894e1bb054833f19 Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Sat, 4 Apr 2026 13:30:45 -0400 Subject: [PATCH 05/22] ci: add crates.io publish job for modularsnf lib crate --- .github/workflows/publish.yml | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index eace29a..626742f 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -131,3 +131,17 @@ jobs: name: sdist path: dist - uses: pypa/gh-action-pypi-publish@release/v1 + + publish-crates: + name: Publish To crates.io + needs: check-dist + if: github.ref == 'refs/heads/main' + runs-on: ubuntu-24.04 + environment: pypi + steps: + - uses: actions/checkout@v5 + - uses: dtolnay/rust-toolchain@stable + - name: Publish modularsnf to crates.io + run: cargo publish -p modularsnf + env: + CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }} From 15c5f505e3a84f5261cfb5d310e4111312007d17 Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Sat, 4 Apr 2026 13:33:36 -0400 Subject: [PATCH 06/22] build: use workspace version inheritance and add version parity check Define version once in workspace Cargo.toml, both crates inherit via version.workspace = true. Add check-version to justfile to ensure pyproject.toml stays in sync. --- Cargo.toml | 3 +++ crates/modularsnf-py/Cargo.toml | 2 +- crates/modularsnf/Cargo.toml | 2 +- justfile | 12 ++++++++++-- 4 files changed, 15 insertions(+), 4 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 6655fcf..ca5b9fc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,5 +2,8 @@ members = ["crates/modularsnf", "crates/modularsnf-py"] resolver = "2" +[workspace.package] +version = "0.4.0" + [profile.release] debug = true diff --git a/crates/modularsnf-py/Cargo.toml b/crates/modularsnf-py/Cargo.toml index fa01c11..60d5b37 100644 --- a/crates/modularsnf-py/Cargo.toml +++ b/crates/modularsnf-py/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "modularsnf-py" -version = "0.4.0" +version.workspace = true edition = "2021" license = "Apache-2.0" description = "Python bindings for modularsnf" diff --git a/crates/modularsnf/Cargo.toml b/crates/modularsnf/Cargo.toml index 030d62a..51f0dcd 100644 --- a/crates/modularsnf/Cargo.toml +++ b/crates/modularsnf/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "modularsnf" -version = "0.4.0" +version.workspace = true edition = "2021" license = "Apache-2.0" description = "Smith Normal Form calculations over Z/N rings" diff --git a/justfile b/justfile index 6c28f83..9d0beb7 100644 --- a/justfile +++ b/justfile @@ -1,5 +1,5 @@ -# Run all checks (lint, typecheck, test+coverage) -check: lint typecheck test +# Run all checks (lint, typecheck, test+coverage, version parity) +check: lint typecheck test check-version lint: uv run ruff check . @@ -11,6 +11,14 @@ test: uv run pytest tests/ -x --backend python uv run pytest tests/ -x --backend rust +# Verify Cargo workspace and pyproject.toml versions match +check-version: + #!/usr/bin/env bash + set -euo pipefail + cargo=$(grep -m1 '^version' Cargo.toml | sed 's/.*"\(.*\)"/\1/') + py=$(grep -m1 '^version' pyproject.toml | sed 's/.*"\(.*\)"/\1/') + if [ "$cargo" != "$py" ]; then echo "Version mismatch: Cargo=$cargo pyproject=$py"; exit 1; fi + # Build the Rust extension (debug) build: uv run maturin develop From 18f915d0c6091bc23a1509f3e275d8b9a55ae454 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 3 Jun 2026 02:52:15 +0000 Subject: [PATCH 07/22] feat: add CRT-style non-recursive SNF path Add a Chinese Remainder Theorem based Smith Normal Form that avoids the recursive divide-and-conquer pipeline entirely. Factor N = prod p^e (passed in, so factoring is amortized across calls), compute the SNF over each local ring Z/p^e via valuation-pivoted Gaussian elimination, then CRT-recombine the per-prime unimodular transforms into U, V over Z/NZ. The local elimination is flat (no recursion), produces the unimodular transforms natively, and yields the divisibility chain for free. Exposed via the new rust_crt_snf PyO3 entry point alongside the existing recursive rust_smith_normal_form, which is left untouched for comparison. https://claude.ai/code/session_013PpJ2A2YvUCUk5EVtU5rse --- crates/modularsnf-py/src/lib.rs | 35 +++++ crates/modularsnf/src/crt.rs | 267 ++++++++++++++++++++++++++++++++ crates/modularsnf/src/lib.rs | 1 + 3 files changed, 303 insertions(+) create mode 100644 crates/modularsnf/src/crt.rs diff --git a/crates/modularsnf-py/src/lib.rs b/crates/modularsnf-py/src/lib.rs index 4493618..fdbc485 100644 --- a/crates/modularsnf-py/src/lib.rs +++ b/crates/modularsnf-py/src/lib.rs @@ -1,5 +1,6 @@ use modularsnf::ring::RingZModN; use modularsnf::snf::smith_square; +use modularsnf::crt::crt_snf; use modularsnf::diagonal::{merge_raw, smith_from_diagonal}; use numpy::ndarray::{s, Array2}; @@ -165,6 +166,39 @@ fn rust_merge_smith_blocks<'py>( Ok((u.into_pyarray(py), v.into_pyarray(py), s.into_pyarray(py))) } +/// CRT-style SNF: takes an n x m matrix, modulus, and the prime +/// factorization of the modulus as (p, e) pairs (so factoring is amortized). +/// Returns (U, V, S) with S = U @ A @ V (mod N). +#[pyfunction] +fn rust_crt_snf<'py>( + py: Python<'py>, + data: PyReadonlyArray2<'py, i64>, + modulus: i64, + factors: Vec<(i64, u32)>, +) -> PyResult<( + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, +)> { + let a = data.as_array().to_owned(); + let n = a.nrows(); + let m = a.ncols(); + + if n == 0 || m == 0 { + let u = Array2::::eye(n); + let v = Array2::::eye(m); + return Ok((u.into_pyarray(py), v.into_pyarray(py), a.into_pyarray(py))); + } + + let (u, v, s_mat) = crt_snf(&a, modulus, &factors); + + Ok(( + u.into_pyarray(py), + v.into_pyarray(py), + s_mat.into_pyarray(py), + )) +} + /// Native Rust acceleration for modularsnf. #[pymodule] fn _rust(m: &Bound<'_, PyModule>) -> PyResult<()> { @@ -172,5 +206,6 @@ fn _rust(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(rust_smith_from_diagonal, m)?)?; m.add_function(wrap_pyfunction!(rust_merge_smith_blocks, m)?)?; m.add_function(wrap_pyfunction!(rust_smith_normal_form, m)?)?; + m.add_function(wrap_pyfunction!(rust_crt_snf, m)?)?; Ok(()) } diff --git a/crates/modularsnf/src/crt.rs b/crates/modularsnf/src/crt.rs new file mode 100644 index 0000000..f5420c6 --- /dev/null +++ b/crates/modularsnf/src/crt.rs @@ -0,0 +1,267 @@ +//! CRT-style Smith Normal Form over Z/NZ — non-recursive. +//! +//! Factor `N = prod p^e` (passed in, so factoring is amortized across many +//! calls with the same modulus), compute the SNF over each local ring +//! `Z/p^e` by valuation-pivoted Gaussian elimination — which yields the +//! divisibility chain for free and the unimodular transforms natively — then +//! CRT-recombine the per-prime transforms into `U, V` over `Z/NZ`. + +use ndarray::Array2; + +use crate::ring::posmod_i128; + +#[inline] +fn pmod(a: i64, n: i64) -> i64 { + posmod_i128(a as i128, n) +} + +#[inline] +fn mulmod(a: i64, b: i64, n: i64) -> i64 { + posmod_i128(a as i128 * b as i128, n) +} + +fn egcd_i128(a: i128, b: i128) -> (i128, i128, i128) { + let (mut r0, mut r1) = (a, b); + let (mut s0, mut s1) = (1i128, 0i128); + let (mut t0, mut t1) = (0i128, 1i128); + while r1 != 0 { + let q = r0 / r1; + let t = r1; + r1 = r0 - q * r1; + r0 = t; + let t = s1; + s1 = s0 - q * s1; + s0 = t; + let t = t1; + t1 = t0 - q * t1; + t0 = t; + } + (r0, s0, t0) +} + +fn inv_mod(a: i64, m: i64) -> i64 { + let (_g, x, _) = egcd_i128(pmod(a, m) as i128, m as i128); + let r = ((x % m as i128) + m as i128) % m as i128; + r as i64 +} + +/// p-adic valuation of `x` in `[0, p^cap)`, capped at `cap`; `x == 0 -> cap`. +fn pval(mut x: i64, p: i64, cap: u32) -> u32 { + if x == 0 { + return cap; + } + let mut v = 0u32; + while v < cap && x % p == 0 { + x /= p; + v += 1; + } + v +} + +/// SNF of `A` over the local ring `Z/p^e` via valuation pivoting. +/// +/// Returns `(U, V, vals)` with `U @ A @ V == diag(p^vals) (mod p^e)`, +/// `U` (n x n) and `V` (m x m) unimodular and `vals` ascending. +fn local_snf(a: &Array2, p: i64, e: u32) -> (Array2, Array2, Vec) { + let q = p.pow(e); + let n = a.nrows(); + let m = a.ncols(); + let r = n.min(m); + + let mut mat = a.mapv(|x| pmod(x, q)); + let mut u = Array2::::eye(n); + let mut v = Array2::::eye(m); + + for k in 0..r { + // Find the pivot in mat[k.., k..] with minimal p-adic valuation. + let mut best: Option<(usize, usize)> = None; + let mut bestval = e + 1; + 'search: for i in k..n { + for j in k..m { + let val = mat[[i, j]]; + if val != 0 { + let vv = pval(val, p, e); + if vv < bestval { + bestval = vv; + best = Some((i, j)); + if vv == 0 { + break 'search; + } + } + } + } + } + let (pi, pj) = match best { + Some(x) => x, + None => break, // trailing block is entirely zero + }; + + if pi != k { + for col in 0..m { + mat.swap([k, col], [pi, col]); + } + for col in 0..n { + u.swap([k, col], [pi, col]); + } + } + if pj != k { + for row in 0..n { + mat.swap([row, k], [row, pj]); + } + for row in 0..m { + v.swap([row, k], [row, pj]); + } + } + + let vv = bestval; + let pv = p.pow(vv); + + // Normalize the pivot to exactly p^vv by scaling row k by a unit. + let unit = mat[[k, k]] / pv; // exact: pivot = p^vv * unit + let uinv = inv_mod(unit, q); + for col in 0..m { + mat[[k, col]] = mulmod(mat[[k, col]], uinv, q); + } + for col in 0..n { + u[[k, col]] = mulmod(u[[k, col]], uinv, q); + } + + // Clear column k below the pivot. + for i in (k + 1)..n { + let val = mat[[i, k]]; + if val != 0 { + let c = val / pv; // exact: val(mat[i,k]) >= vv + for col in 0..m { + mat[[i, col]] = posmod_i128( + mat[[i, col]] as i128 - c as i128 * mat[[k, col]] as i128, + q, + ); + } + for col in 0..n { + u[[i, col]] = posmod_i128( + u[[i, col]] as i128 - c as i128 * u[[k, col]] as i128, + q, + ); + } + } + } + + // Clear row k to the right of the pivot. + for j in (k + 1)..m { + let val = mat[[k, j]]; + if val != 0 { + let c = val / pv; + for row in 0..n { + mat[[row, j]] = posmod_i128( + mat[[row, j]] as i128 - c as i128 * mat[[row, k]] as i128, + q, + ); + } + for row in 0..m { + v[[row, j]] = posmod_i128( + v[[row, j]] as i128 - c as i128 * v[[row, k]] as i128, + q, + ); + } + } + } + } + + let vals: Vec = (0..r).map(|i| pval(mat[[i, i]], p, e)).collect(); + (u, v, vals) +} + +/// Incremental CRT: combine `residues` mod pairwise-coprime `moduli`. +fn crt_combine(residues: &[i64], moduli: &[i64]) -> i64 { + let mut x: i128 = 0; + let mut m_acc: i128 = 1; + for (&r, &m) in residues.iter().zip(moduli) { + let m128 = m as i128; + let m_mod = ((m_acc % m128) + m128) % m128; + let (_g, s, _) = egcd_i128(m_mod, m128); + let inv = ((s % m128) + m128) % m128; + let diff = (((r as i128 - x) % m128) + m128) % m128; + let t = (diff * inv) % m128; + x += m_acc * t; + m_acc *= m128; + x = ((x % m_acc) + m_acc) % m_acc; + } + x as i64 +} + +/// CRT-style SNF over `Z/NZ`. `factors` is the prime factorization of +/// `modulus` as `(p, e)` pairs (passed in so it can be amortized). +/// +/// Returns `(U, V, S)` with `S = U @ A @ V (mod N)`, `S` diagonal with the +/// divisibility chain and `U, V` unimodular. +pub fn crt_snf( + a: &Array2, + modulus: i64, + factors: &[(i64, u32)], +) -> (Array2, Array2, Array2) { + let n = a.nrows(); + let m = a.ncols(); + let r = n.min(m); + let np = factors.len(); + + let qs: Vec = factors.iter().map(|&(p, e)| p.pow(e)).collect(); + + let mut locals: Vec<(Array2, Array2, Vec)> = + factors.iter().map(|&(p, e)| local_snf(a, p, e)).collect(); + + // Global invariant factors d_i = prod_p p^{vals_p[i]} (divides N). + let mut d = vec![0i64; r]; + for (i, di_slot) in d.iter_mut().enumerate() { + let mut di: i128 = 1; + for (pi, &(p, _)) in factors.iter().enumerate() { + di *= (p as i128).pow(locals[pi].2[i]); + } + *di_slot = (di % modulus as i128) as i64; + } + + // Unit-normalize each prime's V so all primes realize the same d_i. + for pi in 0..np { + let q = qs[pi]; + for i in 0..r { + let mut w: i128 = 1; + for (pj, &(p, _)) in factors.iter().enumerate() { + if pj != pi { + w = (w * (p as i128).pow(locals[pj].2[i])) % q as i128; + } + } + let w = w as i64; + for row in 0..m { + let val = locals[pi].1[[row, i]]; + locals[pi].1[[row, i]] = mulmod(val, w, q); + } + } + } + + // CRT-recombine U (n x n) and V (m x m) entrywise across primes. + let mut u = Array2::::zeros((n, n)); + let mut resid = vec![0i64; np]; + for ar in 0..n { + for br in 0..n { + for pi in 0..np { + resid[pi] = locals[pi].0[[ar, br]]; + } + u[[ar, br]] = crt_combine(&resid, &qs); + } + } + let mut v = Array2::::zeros((m, m)); + for ar in 0..m { + for br in 0..m { + for pi in 0..np { + resid[pi] = locals[pi].1[[ar, br]]; + } + v[[ar, br]] = crt_combine(&resid, &qs); + } + } + + let mut s = Array2::::zeros((n, m)); + for (i, &di) in d.iter().enumerate() { + s[[i, i]] = di; + } + + (u, v, s) +} diff --git a/crates/modularsnf/src/lib.rs b/crates/modularsnf/src/lib.rs index 24eedad..aa08e25 100644 --- a/crates/modularsnf/src/lib.rs +++ b/crates/modularsnf/src/lib.rs @@ -1,4 +1,5 @@ pub mod band; +pub mod crt; pub mod diagonal; pub mod echelon; pub mod ring; From 504364f5df05945ab1a6565b04c9ed4ee7b9606c Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 3 Jun 2026 02:52:21 +0000 Subject: [PATCH 08/22] chore: add CRT SNF prototype and profiling scripts Pure-Python reference implementation of the CRT-style SNF, a correctness harness validating it against the existing implementation as oracle, and a profiling script for head-to-head comparison against the recursive path. https://claude.ai/code/session_013PpJ2A2YvUCUk5EVtU5rse --- scripts/crt_snf_profile.py | 75 +++++++++++++ scripts/crt_snf_prototype.py | 207 +++++++++++++++++++++++++++++++++++ scripts/crt_snf_verify.py | 94 ++++++++++++++++ 3 files changed, 376 insertions(+) create mode 100644 scripts/crt_snf_profile.py create mode 100644 scripts/crt_snf_prototype.py create mode 100644 scripts/crt_snf_verify.py diff --git a/scripts/crt_snf_profile.py b/scripts/crt_snf_profile.py new file mode 100644 index 0000000..99a59bb --- /dev/null +++ b/scripts/crt_snf_profile.py @@ -0,0 +1,75 @@ +"""Profile CRT-style SNF prototype vs the existing implementation. + +Compares three code paths at increasing n for a composite modulus: + * existing pipeline, Rust backend (the current production fast path) + * existing pipeline, pure-Python backend + * CRT prototype (pure-Python, object dtype) +""" + +from __future__ import annotations + +import time + +import numpy as np +from crt_snf_prototype import crt_snf + +import modularsnf.diagonal as diagonal_mod +import modularsnf.ring as ring_mod +import modularsnf.snf as snf_mod +from modularsnf import smith_normal_form_mod + +# Capture real Rust hooks so we can toggle the backend. +_R = { + "ring": ring_mod._RustRing, + "diag": diagonal_mod._rust_diag, + "merge": diagonal_mod._rust_merge, + "snf": snf_mod._rust_snf, +} + + +def set_backend(rust: bool) -> None: + ring_mod._RustRing = _R["ring"] if rust else None + diagonal_mod._rust_diag = _R["diag"] if rust else None + diagonal_mod._rust_merge = _R["merge"] if rust else None + snf_mod._rust_snf = _R["snf"] if rust else None + + +def timeit(fn, *args, repeat: int = 3) -> float: + best = float("inf") + for _ in range(repeat): + t0 = time.perf_counter() + fn(*args) + best = min(best, time.perf_counter() - t0) + return best + + +def main() -> None: + N = 720 # 2^4 * 3^2 * 5 — composite with zero divisors + sizes = [4, 8, 16, 24, 32, 48, 64] + rng = np.random.default_rng(0) + + print(f"modulus N = {N} (factor = 2^4 * 3^2 * 5)") + print(f"{'n':>4} | {'rust (ms)':>12} | {'python (ms)':>12} | {'crt (ms)':>12}") + print("-" * 52) + + for n in sizes: + A = rng.integers(0, N, size=(n, n)) + Al = A.tolist() + + set_backend(True) + t_rust = timeit(lambda: smith_normal_form_mod(Al, modulus=N)) + + set_backend(False) + t_py = timeit(lambda: smith_normal_form_mod(Al, modulus=N), repeat=1) + + t_crt = timeit(lambda: crt_snf(A, N), repeat=1) + + print( + f"{n:>4} | {t_rust*1e3:>12.2f} | {t_py*1e3:>12.2f} | {t_crt*1e3:>12.2f}" + ) + + set_backend(True) + + +if __name__ == "__main__": + main() diff --git a/scripts/crt_snf_prototype.py b/scripts/crt_snf_prototype.py new file mode 100644 index 0000000..9a4011e --- /dev/null +++ b/scripts/crt_snf_prototype.py @@ -0,0 +1,207 @@ +"""Prototype: CRT-style Smith Normal Form over Z/NZ with transforms. + +Non-recursive. Factor N = prod p^e, compute SNF over each local ring Z/p^e +by valuation-pivoted Gaussian elimination (which yields the divisibility +chain for free and unimodular transforms natively), then CRT-recombine the +per-prime transforms into U, V over Z/NZ. + +This is a correctness/feasibility prototype in pure Python (object dtype to +avoid overflow), validated against the existing implementation as oracle. +""" + +from __future__ import annotations + +from math import gcd + +import numpy as np + + +def factorize(n: int) -> dict[int, int]: + factors: dict[int, int] = {} + d = 2 + while d * d <= n: + while n % d == 0: + factors[d] = factors.get(d, 0) + 1 + n //= d + d += 1 + if n > 1: + factors[n] = factors.get(n, 0) + 1 + return factors + + +def pval(x: int, p: int, cap: int) -> int: + """p-adic valuation of x, capped at cap; x == 0 -> cap.""" + x %= p**cap + if x == 0: + return cap + v = 0 + while v < cap and x % p == 0: + x //= p + v += 1 + return v + + +def egcd(a: int, b: int) -> tuple[int, int, int]: + if b == 0: + return (a, 1, 0) + g, x, y = egcd(b, a % b) + return (g, y, x - (a // b) * y) + + +def inv_mod(a: int, m: int) -> int: + g, x, _ = egcd(a % m, m) + assert g == 1, f"{a} not invertible mod {m}" + return x % m + + +def local_snf(A: np.ndarray, p: int, e: int): + """SNF of A over the local ring Z/p^e via valuation pivoting. + + Returns (U, V, diag_vals) with U @ A @ V == diag(p^vals) mod p^e, + U (n x n) and V (m x m) unimodular, vals ascending (divisibility chain). + """ + q = p**e + M = (A.astype(object)) % q + n, m = M.shape + U = np.eye(n, dtype=object) + V = np.eye(m, dtype=object) + r = min(n, m) + + for k in range(r): + # Find pivot in M[k:, k:] with minimal p-adic valuation. + best = None + bestval = e + 1 + for i in range(k, n): + for j in range(k, m): + if int(M[i, j]) % q != 0: + v = pval(int(M[i, j]), p, e) + if v < bestval: + bestval, best = v, (i, j) + if v == 0: + break + if bestval == 0: + break + if best is None: + break # entire trailing block is zero + + pi, pj = best + if pi != k: + M[[k, pi], :] = M[[pi, k], :] + U[[k, pi], :] = U[[pi, k], :] + if pj != k: + M[:, [k, pj]] = M[:, [pj, k]] + V[:, [k, pj]] = V[:, [pj, k]] + + v = bestval + pv = p**v + # Normalize pivot to exactly p^v by scaling row k by the unit inverse. + pivot = int(M[k, k]) % q + unit = pivot // pv # coprime to p + uinv = inv_mod(unit, q) + M[k, :] = (M[k, :] * uinv) % q + U[k, :] = (U[k, :] * uinv) % q + + # Clear column k below the pivot. + for i in range(k + 1, n): + if int(M[i, k]) % q != 0: + c = (int(M[i, k]) % q) // pv # exact: val(M[i,k]) >= v + M[i, :] = (M[i, :] - c * M[k, :]) % q + U[i, :] = (U[i, :] - c * U[k, :]) % q + + # Clear row k to the right of the pivot. + for j in range(k + 1, m): + if int(M[k, j]) % q != 0: + c = (int(M[k, j]) % q) // pv + M[:, j] = (M[:, j] - c * M[:, k]) % q + V[:, j] = (V[:, j] - c * V[:, k]) % q + + vals = [pval(int(M[i, i]), p, e) for i in range(r)] + return U, V, vals + + +def crt_combine(residues: list[int], moduli: list[int]) -> int: + """Combine residues mod pairwise-coprime moduli into a value mod prod.""" + x = 0 + M = 1 + for r, m in zip(residues, moduli): + # solve x ≡ x (mod M), x ≡ r (mod m) + g, s, _ = egcd(M, m) + assert g == 1 + x = (x + M * ((r - x) * s % m)) % (M * m) + M *= m + return x + + +def crt_snf(A: np.ndarray, N: int): + """CRT-style SNF over Z/NZ. Returns (U, V, S) with S = U @ A @ V mod N.""" + A = A.astype(object) % N + n, m = A.shape + r = min(n, m) + fac = factorize(N) + primes = sorted(fac) + qs = [p ** fac[p] for p in primes] + + per_prime = {} + vals = {} + for p in primes: + e = fac[p] + Up, Vp, vp = local_snf(A, p, e) + per_prime[p] = (Up, Vp) + vals[p] = vp + + # Global invariant factors d_i = prod_p p^{vals_p[i]} (divides N). + d = [] + for i in range(r): + di = 1 + for p in primes: + di *= p ** vals[p][i] + d.append(di % N) + + # Unit-normalize each prime's V so all primes realize the SAME global d_i. + for p in primes: + e = fac[p] + q = p**e + Vp = per_prime[p][1] + for i in range(r): + # w = prod_{p'!=p} p'^{vals_{p'}[i]} (a unit mod q) + w = 1 + for p2 in primes: + if p2 != p: + w *= p2 ** vals[p2][i] + w %= q + Vp[:, i] = (Vp[:, i] * w) % q + + # CRT-recombine U and V entrywise across primes. + U = np.zeros((n, n), dtype=object) + V = np.zeros((m, m), dtype=object) + for a in range(n): + for b in range(n): + U[a, b] = crt_combine( + [int(per_prime[p][0][a, b]) for p in primes], qs + ) + for a in range(m): + for b in range(m): + V[a, b] = crt_combine( + [int(per_prime[p][1][a, b]) for p in primes], qs + ) + + S = np.zeros((n, m), dtype=object) + for i in range(r): + S[i, i] = d[i] % N + return U, V, S + + +def det_mod(M: np.ndarray, N: int) -> int: + """Determinant of integer matrix mod N via fraction-free (Bareiss-ish).""" + M = M.astype(object) % N + n = M.shape[0] + if n == 0: + return 1 % N + # Plain cofactor for small n is fine for tests; use integer det via numpy on python ints. + from sympy import Matrix + + return int(Matrix(M.tolist()).det()) % N + + +def is_unimodular(M: np.ndarray, N: int) -> bool: + return gcd(det_mod(M, N) % N, N) == 1 diff --git a/scripts/crt_snf_verify.py b/scripts/crt_snf_verify.py new file mode 100644 index 0000000..328f791 --- /dev/null +++ b/scripts/crt_snf_verify.py @@ -0,0 +1,94 @@ +"""Verify the CRT-style SNF prototype against the existing implementation.""" + +from __future__ import annotations + +import random +from math import gcd + +import numpy as np +from crt_snf_prototype import crt_snf, is_unimodular + +from modularsnf import smith_normal_form_mod + + +def check_one(A: np.ndarray, N: int) -> tuple[bool, str]: + n, m = A.shape + U, V, S = crt_snf(A, N) + + # 1. Self-certifying factorization equation. + prod = (U.astype(object) @ A.astype(object) @ V.astype(object)) % N + if not np.array_equal(prod % N, S % N): + return False, "U@A@V != S mod N" + + # 2. S diagonal. + for i in range(n): + for j in range(m): + if i != j and int(S[i, j]) % N != 0: + return False, "S not diagonal" + + # 3. Divisibility chain d_i | d_{i+1} (as ideals: gcd(d_i,N) | gcd(d_{i+1},N)). + r = min(n, m) + g = [gcd(int(S[i, i]) % N, N) for i in range(r)] + for i in range(r - 1): + if g[i + 1] % g[i] != 0: + return False, f"divisibility chain broken: {g}" + + # 4. Unimodularity of U and V. + if not is_unimodular(U, N): + return False, "U not unimodular" + if not is_unimodular(V, N): + return False, "V not unimodular" + + # 5. Diagonal matches the oracle (SNF diagonal is unique up to gcd-with-N). + So, _, _ = smith_normal_form_mod(A.tolist(), modulus=N) + So = np.array(So, dtype=object) + g_oracle = [gcd(int(So[i, i]) % N, N) for i in range(r)] + if g != g_oracle: + return False, f"diagonal mismatch: mine={g} oracle={g_oracle}" + + return True, "ok" + + +def main() -> None: + random.seed(12345) + np.random.seed(12345) + + moduli = [2, 3, 4, 6, 8, 9, 12, 16, 30, 36, 60, 36, 100, 210, 720, 1024] + sizes = [1, 2, 3, 4, 5, 6] + + total = 0 + failures = 0 + for N in moduli: + for n in sizes: + for _ in range(30): + m = n # square; rectangular handled by same padding upstream + A = np.random.randint(0, N, size=(n, m)) + total += 1 + ok, msg = check_one(A, N) + if not ok: + failures += 1 + if failures <= 10: + print(f"FAIL N={N} shape={A.shape}: {msg}") + print(A) + + # Edge cases. + edge = [ + (np.zeros((3, 3), dtype=int), 12), + (np.eye(3, dtype=int), 12), + (np.array([[2, 4, 0], [6, 8, 3], [0, 3, 9]]), 36), + (np.array([[6]]), 12), + (np.full((4, 4), 6), 36), + (np.array([[4, 0], [0, 9]]), 36), + ] + for A, N in edge: + total += 1 + ok, msg = check_one(A, N) + if not ok: + failures += 1 + print(f"FAIL edge N={N} shape={A.shape}: {msg}\n{A}") + + print(f"\n{total - failures}/{total} passed, {failures} failures") + + +if __name__ == "__main__": + main() From d9f6358189ee70b0a167e70048f1ad96b741536a Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Wed, 3 Jun 2026 00:27:00 -0400 Subject: [PATCH 09/22] docs: replace PLAN.md with CRT optimization roadmap Prior 0.2.0 release-hardening plan is fully shipped (v0.4.0: int64 contract, i128 intermediates, abi3 wheels, dual-backend tests). Replace with the research-backed roadmap for optimizing the CRT SNF path (Phases 0-4). --- docs/PLAN.md | 153 +++++++++++++++++++++++++++------------------------ 1 file changed, 81 insertions(+), 72 deletions(-) diff --git a/docs/PLAN.md b/docs/PLAN.md index c765932..f220b2a 100644 --- a/docs/PLAN.md +++ b/docs/PLAN.md @@ -1,75 +1,84 @@ -# PLAN: 0.2.0 Release Hardening +# PLAN: Optimize the CRT-based Smith Normal Form path ## Goal -Ship `0.2.0` with a coherent machine-integer contract, safer Rust hot paths, -and release automation that publishes binary wheels for mainstream platforms. - -## Scope - -1. Make the public Python API explicitly `int64`-only. -2. Keep Rust as the primary accelerated path, but remove accidental claims of - arbitrary-precision support. -3. Use wider intermediates in Rust where `i64` multiplication or accumulation - can overflow before modular reduction. -4. Publish `abi3` wheels for: - - Linux `x86_64` - - Linux `aarch64` - - macOS Intel - - macOS Apple Silicon - - Windows `x86_64` -5. Clean up the profiling script so backend selection is explicit and the - reported instrumentation matches the code path being exercised. -6. Make `just check` validate the full test suite against both the pure-Python - and Rust backends without duplicating test files. -7. Finish with `just check`. - -## Work Items - -### 1. Python Contract - -- Add explicit signed-64-bit validation for: - - modulus values - - matrix entries -- Remove object-dtype fallback paths that implied arbitrary-precision support. -- Update tests to assert the new failure mode for out-of-range inputs. - -### 2. Rust Arithmetic Safety - -- Audit multiplication and accumulation sites in the Rust core. -- Use `i128` intermediates only where needed before reducing back to `i64`. -- Preserve `i64` storage and FFI types. - -### 3. Packaging And Release Automation - -- Bump package versions to `0.2.0`. -- Switch PyO3 packaging to `abi3-py310`. -- Update GitHub Actions to build and publish wheels for the release matrix. -- Keep sdist generation, accepting that source builds require Rust. - -### 4. Profiling Tooling - -- Make backend choice explicit in `scripts/profile_snf.py`. -- Ensure Python-only instrumentation is not presented as if it profiles the - Rust top-level fast path. -- Clean up style issues so it passes the repo checks. - -### 5. Validation - -- Add a pytest backend switch that can force either implementation for an - entire test session. -- Run the existing test suite twice under `just test`: - - once with `--backend python` - - once with `--backend rust` -- Keep the ring and matrix tests as exact-behavior checks. -- Keep the SNF tests focused on structural properties and oracle comparisons - rather than raw `U`, `V`, `S` tuple equality. -- Run `just check`. -- Fix lint, type, and test failures introduced by the release hardening. - -## Notes - -- The intended contract for `0.2.0` is machine integers, not arbitrary-size - Python integers. -- Using `i128` internally does not expand the public numeric contract; it only - avoids incorrect overflow in modular arithmetic on `i64` inputs. +Turn the CRT SNF path (`crates/modularsnf/src/crt.rs`, `rust_crt_snf`) into a +fully-optimized, large-`n`, multi-core implementation for the **small-modulus +(N ≤ a few hundred), large-matrix (n in the thousands)** regime. + +Baseline established (head-to-head, rust-vs-rust, branch `bench/headtohead`): +the CRT path already beats the recursive Storjohann baseline 4.4–35×. This plan +attacks its remaining bottlenecks: i128 in the hot loop, scalar elimination, +and single-threaded execution. + +## Why this design is correct (research-validated) + +- Over a chain/valuation ring `Z/p^e`, the SNF conditioner can be a permutation + matrix — valuation-pivoted Gaussian elimination, **no gcdex/Bézout** + (Storjohann diss. Prop 9.21; ESA). This is exactly what `local_snf` does. +- Local-then-recombine is independently validated (Wilkening–Yu). +- SNF/Howell over `Z/N` reduces to **matrix multiplication** — so "make it fast" + means "turn the trailing update into GEMM." + +## Library strategy (call-into vs hand-roll) + +| Component | Verdict | Rationale | +|---|---|---| +| Local elimination (Z/p^e) | **Hand-roll (Rust)** | Simple chain-ring pivot; no lib does Smith-with-transforms over Z/p^e; license-clean. | +| Modular GEMM (trailing update) | **Hand-roll thin layer + BLAS/faer** | Accumulate f64/i64, reduce mod p^e; FFLAS does this but is LGPL + field-only. | +| CRT recombination | **Hand-roll (Rust)** | Already done; O(n²·np), not the bottleneck. | +| External SNF lib | **Reference only** | FLINT `nmod_mat_howell_form` = Howell≠Smith, single-limb, **LGPL**. PARI = **GPL → incompatible** with Apache-2.0. | + +Decisive constraint: LGPL relinking is fragile in a static manylinux/abi3 wheel; +GPL is out. Hand-roll; use FFLAS-FFPACK / FLINT as algorithmic references. +Prefer **faer** (MIT/Apache) if a Rust GEMM backend is needed. + +Key arithmetic fact for our regime: exact integer GEMM stays exact while +`λ·(p−1)² < 2^53`. For p^e ≈ 256 that's λ ≈ 1.4e8 accumulations before a +reduction is needed — so reductions are essentially free and BLAS `dgemm` / +`faer` can be the inner kernel. (Use only this single-level bound; the nested +Strassen-Winograd `kWinograd` bound did not survive verification.) + +## Roadmap (phased, each verified against the head-to-head harness) + +### Phase 0 — Characterize *(in progress)* +- Env-gated timing in `crt_snf`: measure `local_snf` total vs recombination vs + unit-normalization at large n (500, 1000, 2000) and representative moduli. +- Confirm the `local_snf` i128 inner loop dominates. Record the split. + +### Phase 1 — Kill i128 in the hot loop (single-core constant win) *(next)* +- Add an i64 fast-path reduction (valid when `p^e < 2^31`, i128 fallback else), + mirroring `optimize-modulus`'s `mul_mod`. +- Apply to `local_snf` row/column clears + pivot normalization (`mulmod`). +- Verify: invariant factors + `U·A·V == S (mod N)` unchanged; measure speedup. + +### Phase 2 — Blocked right-looking LU + modular GEMM (asymptotic win) +- Restructure `local_snf` into panels; trailing update becomes one matmul + `T -= L·U_panel (mod p^e)` (Storjohann reduction-to-matmul). +- Build the f64/i64-accumulate → reduce-mod-p^e GEMM layer; evaluate `faer` + vs ndarray+BLAS vs hand-rolled as the inner kernel. + +### Phase 3 — Multi-core (rayon) +- Parallelize the trailing-update GEMM first (real scaling). +- Add across-primes parallelism (cheap, low ceiling: np = 1–3). +- **Re-characterize scaling empirically** — published speedups (e.g. 13.6×/32 + cores, FFLAS Euro-Par 2014) are over prime *fields* Z/pZ, not Z/p^e. + +### Phase 4 — GPU (deferred, literature-gated) +- No reusable GPU kernel for modular LU/SNF over Z/p(^e) was found. Revisit only + after rayon scaling is measured. + +## Open questions / risks +- Field-vs-chain-ring: parallel/GEMM benchmark evidence is Z/pZ; Z/p^e parallel + behavior is unmeasured — expect to measure it ourselves. +- Integer-vs-polynomial: Storjohann/Wilkening-Yu state results over K[x]/(p^k); + Z/p^e transfer is structurally sound extrapolation. +- Evaluate `faer` early in Phase 2 as the license-clean GEMM backend. + +## References +- Storjohann, dissertation §9.6 (Prop 9.21, 9.23): cs.uwaterloo.ca/~astorjoh/diss2up.pdf +- Storjohann, ESA (Howell/Smith reduce to matmul): cs.uwaterloo.ca/~astorjoh/esa.pdf +- Dumas/Gautier/Pernet/Sultan, parallel echelon forms (Euro-Par 2014): arxiv.org/abs/1402.3501 +- Dumas/Giorgi/Pernet, FFLAS-FFPACK (exact GEMM, delayed reduction): hal.science/hal-00018223v2 +- Wilkening–Yu, local Smith form: math.berkeley.edu/~wilken/papers/smith.pdf +- FLINT nmod_mat (Howell form over Z/N): flintlib.org/doc/nmod_mat.html From 9089563123fd881a9ae57f6cca599b550e1879b1 Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Wed, 3 Jun 2026 00:27:00 -0400 Subject: [PATCH 10/22] perf(crt): i64 fast-path reduction in local_snf hot loops (Phase 1) local_snf is 97-100% of CRT runtime (Phase 0 profiling). Its row/column clears reduced every update through posmod_i128. For prime-power moduli p^e < 2^31 the difference a-b*c fits in i64, so widen only when needed. Adds sub_mul_mod + i64 fast path in mulmod, and env-gated CRT_PROFILE phase-split instrumentation. ~1.5-2.05x faster across n=20..400, all moduli; 1.4-1.6x at n=1000-2000. Verified: invariant factors and U*A*V==S (mod N) unchanged vs dev baseline and optimize-modulus across 50 configs. --- crates/modularsnf/src/crt.rs | 70 +++++++++++++++++++++++++++--------- 1 file changed, 53 insertions(+), 17 deletions(-) diff --git a/crates/modularsnf/src/crt.rs b/crates/modularsnf/src/crt.rs index f5420c6..6d87139 100644 --- a/crates/modularsnf/src/crt.rs +++ b/crates/modularsnf/src/crt.rs @@ -15,9 +15,30 @@ fn pmod(a: i64, n: i64) -> i64 { posmod_i128(a as i128, n) } +/// Threshold below which q^2 (and a - b*c with a,b,c in [0,q)) fits in i64, +/// so the hot loops avoid i128. q < 2^31 => q^2 < 2^62, comfortably in range. +const I64_FAST_MAX: i64 = 1i64 << 31; + #[inline] fn mulmod(a: i64, b: i64, n: i64) -> i64 { - posmod_i128(a as i128 * b as i128, n) + if n < I64_FAST_MAX { + let r = (a * b) % n; + if r < 0 { r + n } else { r } + } else { + posmod_i128(a as i128 * b as i128, n) + } +} + +/// Reduce `a - b*c` into [0, q). Inputs `a, b, c` are in [0, q). +/// i64 fast path when q < 2^31; widens to i128 for larger prime-power moduli. +#[inline] +fn sub_mul_mod(a: i64, b: i64, c: i64, q: i64) -> i64 { + if q < I64_FAST_MAX { + let r = (a - b * c) % q; + if r < 0 { r + q } else { r } + } else { + posmod_i128(a as i128 - b as i128 * c as i128, q) + } } fn egcd_i128(a: i128, b: i128) -> (i128, i128, i128) { @@ -132,16 +153,10 @@ fn local_snf(a: &Array2, p: i64, e: u32) -> (Array2, Array2, Vec< if val != 0 { let c = val / pv; // exact: val(mat[i,k]) >= vv for col in 0..m { - mat[[i, col]] = posmod_i128( - mat[[i, col]] as i128 - c as i128 * mat[[k, col]] as i128, - q, - ); + mat[[i, col]] = sub_mul_mod(mat[[i, col]], c, mat[[k, col]], q); } for col in 0..n { - u[[i, col]] = posmod_i128( - u[[i, col]] as i128 - c as i128 * u[[k, col]] as i128, - q, - ); + u[[i, col]] = sub_mul_mod(u[[i, col]], c, u[[k, col]], q); } } } @@ -152,16 +167,10 @@ fn local_snf(a: &Array2, p: i64, e: u32) -> (Array2, Array2, Vec< if val != 0 { let c = val / pv; for row in 0..n { - mat[[row, j]] = posmod_i128( - mat[[row, j]] as i128 - c as i128 * mat[[row, k]] as i128, - q, - ); + mat[[row, j]] = sub_mul_mod(mat[[row, j]], c, mat[[row, k]], q); } for row in 0..m { - v[[row, j]] = posmod_i128( - v[[row, j]] as i128 - c as i128 * v[[row, k]] as i128, - q, - ); + v[[row, j]] = sub_mul_mod(v[[row, j]], c, v[[row, k]], q); } } } @@ -206,8 +215,13 @@ pub fn crt_snf( let qs: Vec = factors.iter().map(|&(p, e)| p.pow(e)).collect(); + // Phase 0 profiling: set CRT_PROFILE=1 to print a phase-time breakdown. + let prof = std::env::var("CRT_PROFILE").is_ok(); + let t_local = std::time::Instant::now(); + let mut locals: Vec<(Array2, Array2, Vec)> = factors.iter().map(|&(p, e)| local_snf(a, p, e)).collect(); + let d_local = t_local.elapsed(); // Global invariant factors d_i = prod_p p^{vals_p[i]} (divides N). let mut d = vec![0i64; r]; @@ -220,6 +234,7 @@ pub fn crt_snf( } // Unit-normalize each prime's V so all primes realize the same d_i. + let t_norm = std::time::Instant::now(); for pi in 0..np { let q = qs[pi]; for i in 0..r { @@ -237,7 +252,10 @@ pub fn crt_snf( } } + let d_norm = t_norm.elapsed(); + // CRT-recombine U (n x n) and V (m x m) entrywise across primes. + let t_recomb = std::time::Instant::now(); let mut u = Array2::::zeros((n, n)); let mut resid = vec![0i64; np]; for ar in 0..n { @@ -258,10 +276,28 @@ pub fn crt_snf( } } + let d_recomb = t_recomb.elapsed(); + let mut s = Array2::::zeros((n, m)); for (i, &di) in d.iter().enumerate() { s[[i, i]] = di; } + if prof { + let total = d_local + d_norm + d_recomb; + eprintln!( + "[CRT_PROFILE] n={n} m={m} N={modulus} np={np} | \ + local_snf={:.3}ms ({:.0}%) normalize={:.3}ms ({:.0}%) \ + recombine={:.3}ms ({:.0}%) | total={:.3}ms", + d_local.as_secs_f64() * 1e3, + 100.0 * d_local.as_secs_f64() / total.as_secs_f64(), + d_norm.as_secs_f64() * 1e3, + 100.0 * d_norm.as_secs_f64() / total.as_secs_f64(), + d_recomb.as_secs_f64() * 1e3, + 100.0 * d_recomb.as_secs_f64() / total.as_secs_f64(), + total.as_secs_f64() * 1e3, + ); + } + (u, v, s) } From b4b06e081454a45115489a020a98ae9f5ec97319 Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Wed, 3 Jun 2026 00:44:04 -0400 Subject: [PATCH 11/22] perf(crt): split local_snf into two-phase column/row elimination (Phase 2a) Restructure local_snf into Phase L (minimal-valuation column elimination -> U, upper-triangular mat) + Phase R (diagonalize -> V, reverse-order back- elimination to keep fill in unprocessed rows). This is the structural groundwork for a blocked GEMM trailing update, and already removes wasted work: the old inline row-clear updated all n rows of each column though only the pivot row was nonzero. ~1.1-1.3x faster than Phase 1. Adds a randomized cargo test (no external deps) asserting local_snf yields a valid Smith form: U*A*V == diag(p^vals) (mod q), ascending valuations, and U,V unimodular (det != 0 mod p), across 500+ cases. Verified: zero invariant- factor mismatches vs the independent dev baseline on the head-to-head sweep. --- crates/modularsnf/src/crt.rs | 142 +++++++++++++++++++++++++++++++++-- 1 file changed, 136 insertions(+), 6 deletions(-) diff --git a/crates/modularsnf/src/crt.rs b/crates/modularsnf/src/crt.rs index 6d87139..01c13ee 100644 --- a/crates/modularsnf/src/crt.rs +++ b/crates/modularsnf/src/crt.rs @@ -93,6 +93,10 @@ fn local_snf(a: &Array2, p: i64, e: u32) -> (Array2, Array2, Vec< let mut u = Array2::::eye(n); let mut v = Array2::::eye(m); + // ---- Phase L: column elimination -> U, mat becomes upper-triangular ---- + // Minimal-valuation pivoting (which sorts the diagonal by valuation, giving + // the divisibility chain). Row clears are deferred to Phase R; deferring is + // valid because column clears alone produce the upper-triangular Schur form. for k in 0..r { // Find the pivot in mat[k.., k..] with minimal p-adic valuation. let mut best: Option<(usize, usize)> = None; @@ -147,26 +151,36 @@ fn local_snf(a: &Array2, p: i64, e: u32) -> (Array2, Array2, Vec< u[[k, col]] = mulmod(u[[k, col]], uinv, q); } - // Clear column k below the pivot. + // Clear column k below the pivot (Schur update of the trailing block). for i in (k + 1)..n { let val = mat[[i, k]]; if val != 0 { let c = val / pv; // exact: val(mat[i,k]) >= vv - for col in 0..m { + for col in (k + 1)..m { mat[[i, col]] = sub_mul_mod(mat[[i, col]], c, mat[[k, col]], q); } + mat[[i, k]] = 0; for col in 0..n { u[[i, col]] = sub_mul_mod(u[[i, col]], c, u[[k, col]], q); } } } + } + + let vals: Vec = (0..r).map(|i| pval(mat[[i, i]], p, e)).collect(); - // Clear row k to the right of the pivot. + // ---- Phase R: diagonalize the upper-triangular mat -> V ---- + // Clear super-diagonal entries by column operations, processing pivots from + // last to first so that fill created above row k lands in not-yet-processed + // rows (and is cleared when those pivots are reached). col k of the upper- + // triangular mat is nonzero only in rows 0..=k. + for k in (0..r).rev() { + let pv = p.pow(vals[k]); for j in (k + 1)..m { let val = mat[[k, j]]; if val != 0 { - let c = val / pv; - for row in 0..n { + let c = val / pv; // exact: mat[k,j] divisible by pv + for row in 0..=k { mat[[row, j]] = sub_mul_mod(mat[[row, j]], c, mat[[row, k]], q); } for row in 0..m { @@ -176,7 +190,6 @@ fn local_snf(a: &Array2, p: i64, e: u32) -> (Array2, Array2, Vec< } } - let vals: Vec = (0..r).map(|i| pval(mat[[i, i]], p, e)).collect(); (u, v, vals) } @@ -301,3 +314,120 @@ pub fn crt_snf( (u, v, s) } + +#[cfg(test)] +mod tests { + use super::*; + use ndarray::Array2; + + /// Deterministic LCG so tests need no external rng dependency. + struct Lcg(u64); + impl Lcg { + fn next(&mut self) -> u64 { + self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + self.0 >> 16 + } + fn below(&mut self, n: i64) -> i64 { + (self.next() % n as u64) as i64 + } + } + + fn matmul_mod(a: &Array2, b: &Array2, q: i64) -> Array2 { + let (n, k) = (a.nrows(), a.ncols()); + let m = b.ncols(); + let mut c = Array2::::zeros((n, m)); + for i in 0..n { + for j in 0..m { + let mut acc: i128 = 0; + for t in 0..k { + acc += a[[i, t]] as i128 * b[[t, j]] as i128; + } + c[[i, j]] = posmod_i128(acc, q); + } + } + c + } + + /// True iff det(M mod p) != 0 over the field F_p (=> M unimodular mod p^e). + fn invertible_mod_p(mm: &Array2, p: i64) -> bool { + let n = mm.nrows(); + let mut a = mm.mapv(|x| pmod(x, p)); + for k in 0..n { + // find a nonzero pivot in column k at/below row k + let mut piv = None; + for i in k..n { + if a[[i, k]] != 0 { + piv = Some(i); + break; + } + } + let pr = match piv { + Some(x) => x, + None => return false, // singular mod p + }; + if pr != k { + for col in 0..n { + a.swap([k, col], [pr, col]); + } + } + let inv = inv_mod(a[[k, k]], p); + for i in (k + 1)..n { + if a[[i, k]] != 0 { + let f = pmod(a[[i, k]] * inv, p); + for col in k..n { + a[[i, col]] = pmod(a[[i, col]] - f * a[[k, col]], p); + } + } + } + } + true + } + + #[test] + fn local_snf_is_valid_smith_form() { + let cases = [(2u32, 1u32), (2, 3), (3, 2), (5, 1), (5, 2), (7, 3)]; + let mut rng = Lcg(0x1234_5678_9abc_def0); + let mut checked = 0; + for &(p, e) in &cases { + let p = p as i64; + let q = p.pow(e); + for n in [1usize, 2, 5, 9, 16, 23] { + for m in [1usize, 3, 5, 9, 16] { + for _trial in 0..6 { + let a = Array2::from_shape_fn((n, m), |_| rng.below(q)); + let (uu, vv, vals) = local_snf(&a, p, e); + + // U @ A @ V == diag(p^vals) (mod q) + let prod = matmul_mod(&matmul_mod(&uu, &a, q), &vv, q); + let r = n.min(m); + for i in 0..n { + for j in 0..m { + let expect = if i == j && i < r { + pmod(p.pow(vals[i]), q) + } else { + 0 + }; + assert_eq!( + prod[[i, j]], expect, + "U*A*V mismatch at ({i},{j}) p={p} e={e} n={n} m={m}" + ); + } + } + // valuations ascending (divisibility chain) + for i in 1..r { + assert!( + vals[i - 1] <= vals[i], + "vals not ascending: {vals:?} p={p} e={e}" + ); + } + // U (n x n) and V (m x m) unimodular + assert!(invertible_mod_p(&uu, p), "U not unimodular p={p} e={e}"); + assert!(invertible_mod_p(&vv, p), "V not unimodular p={p} e={e}"); + checked += 1; + } + } + } + } + assert!(checked > 500, "expected many cases, got {checked}"); + } +} From beb7d79554f5b9c9e1ea13a1013b2c62f528cb2f Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Wed, 3 Jun 2026 00:48:37 -0400 Subject: [PATCH 12/22] docs: record Phase 2a done + detailed Phase 2b blocked-GEMM design Capture the valuation-level-staging insight (global min-valuation pivoting defeats naive blocking; stage by p-adic level so each level is unit-pivot blocked LU) and the L21/U12 bookkeeping + deflation plan, so 2b can be executed cleanly against the cargo oracle. --- docs/PLAN.md | 37 +++++++++++++++++++++++++++++++++---- 1 file changed, 33 insertions(+), 4 deletions(-) diff --git a/docs/PLAN.md b/docs/PLAN.md index f220b2a..949d4cb 100644 --- a/docs/PLAN.md +++ b/docs/PLAN.md @@ -53,10 +53,39 @@ Strassen-Winograd `kWinograd` bound did not survive verification.) - Verify: invariant factors + `U·A·V == S (mod N)` unchanged; measure speedup. ### Phase 2 — Blocked right-looking LU + modular GEMM (asymptotic win) -- Restructure `local_snf` into panels; trailing update becomes one matmul - `T -= L·U_panel (mod p^e)` (Storjohann reduction-to-matmul). -- Build the f64/i64-accumulate → reduce-mod-p^e GEMM layer; evaluate `faer` - vs ndarray+BLAS vs hand-rolled as the inner kernel. + +**2a — two-phase split (DONE).** `local_snf` now does Phase L (minimal-valuation +column elimination → U, upper-triangular `mat`) + Phase R (reverse-order back- +elimination → V). Groundwork for blocking; also removed wasted work (old inline +row-clear touched all n rows of each column). ~1.1–1.3× over Phase 1. Gated by a +randomized cargo oracle test (`local_snf_is_valid_smith_form`, 500+ cases: +`U·A·V==diag(p^vals)`, ascending vals, `U,V` unimodular) + head-to-head vs dev. + +**2b — blocked GEMM trailing update (NEXT, designed).** Key constraint: +minimal-valuation pivoting needs the *global* min over the trailing block, which +defeats naive column-panel/lazy blocking. The blockable formulation is +**valuation-level staging**: process pivots by p-adic level ℓ = 0..e; within a +level every pivot is a unit mod p, so it degenerates to standard blocked LU. +Implementation plan for Phase L: +- Outer loop: `lev` = current global min valuation in `mat[k.., k..]`; `pl=p^lev`. +- Panel (width B≈48) of valuation-`lev` pivots. Bookkeeping per standard + right-looking blocked LU: + - factor panel columns `[pstart,pend)`: full updates to **panel rows** (so the + `U12` block `mat[pstart..pend, pend..]` is correct) and to **panel columns** + for all rows (so pivot search + the `L21` multipliers `mat[pend.., pstart..pend]` + are correct); defer non-panel trailing columns. + - **GEMM trailing update**: `mat[pend.., pend..] -= L21 @ U12 (mod p^e)`, and + the analogous `u[pend.., :] -= L21 @ u[pstart..pend, :]`. + - column-deflation: a panel column with no valuation-`lev` unit → swap out to + the active-region tail (track in V), revisit at a higher level. +- GEMM kernel: i64 accumulate with delayed reduction (small `p^e` ⇒ block + `λ·(p−1)² < 2^53` is huge ⇒ reduce rarely); i128 acc fallback for large `p^e`. + Later swap inner kernel for `faer` (MIT/Apache) or BLAS dgemm. +- RISK: exact-LU bookkeeping (L21/U12 split, panel-row vs below-panel updates, + deflation) is subtle; keep the scalar 2a `local_snf` as the cargo-test oracle + and only switch the default once the blocked version passes the full suite. +- Apply the same blocking to Phase R's V-update (triangular column reduction → + one TRMM/GEMM) once Phase L lands. ### Phase 3 — Multi-core (rayon) - Parallelize the trailing-update GEMM first (real scaling). From 71fd99177f4c36034c675f72079fa4080d6e063e Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Wed, 3 Jun 2026 01:14:19 -0400 Subject: [PATCH 13/22] perf(crt): blocked right-looking LU with GEMM trailing update (Phase 2b) Phase L now clears the valuation-0 bulk with blocked LU (panel width 48): in-place panel factorization, TRSM (L11^-1) for the deferred panel-row block, and a GEMM trailing update (mat[pend..,B..] -= L21@U12 mod q, same for U) with an i64 accumulator + delayed reduction (i128 fallback for large q). The scalar min-valuation path finishes the higher-valuation / rank-deficient tail; a final pass normalizes blocked pivots (left as units to keep mat and u row-consistent). Speedup grows with n (cache/blocking): ~1.3-1.5x over Phase 2a at n=400, and 1.6-2.6x over Phase 1 at n=1000-2000. Cumulative ~3x over the original CRT. Verified: two randomized cargo tests (600+ cases incl. multi-panel, rank- deficiency mod p forcing the blocked->scalar handoff, p|A skipping the blocked pass, rectangular) assert valid SNF; head-to-head shows zero invariant-factor mismatches vs the dev baseline and all transform_ok across the sweep. --- crates/modularsnf/src/crt.rs | 363 ++++++++++++++++++++++++++++++----- 1 file changed, 314 insertions(+), 49 deletions(-) diff --git a/crates/modularsnf/src/crt.rs b/crates/modularsnf/src/crt.rs index 01c13ee..57ac520 100644 --- a/crates/modularsnf/src/crt.rs +++ b/crates/modularsnf/src/crt.rs @@ -79,25 +79,203 @@ fn pval(mut x: i64, p: i64, cap: u32) -> u32 { v } -/// SNF of `A` over the local ring `Z/p^e` via valuation pivoting. -/// -/// Returns `(U, V, vals)` with `U @ A @ V == diag(p^vals) (mod p^e)`, -/// `U` (n x n) and `V` (m x m) unimodular and `vals` ascending. -fn local_snf(a: &Array2, p: i64, e: u32) -> (Array2, Array2, Vec) { - let q = p.pow(e); - let n = a.nrows(); - let m = a.ncols(); - let r = n.min(m); +/// Block width for the right-looking blocked LU (Phase L, unit stage). +const PANEL_B: usize = 48; +/// i64 GEMM accumulator is safe while `PANEL_B * (q-1)^2 < 2^63`. q < 2^28 +/// gives 48 * 2^56 < 2^62, so use i64 below it and i128 at/above it. +const GEMM_I64_MAX: i64 = 1i64 << 28; + +/// Minimal p-adic valuation over `mat[k.., k..]` (returns `e+1` if all zero). +fn min_valuation(mat: &Array2, k: usize, n: usize, m: usize, p: i64, e: u32) -> u32 { + let mut lev = e + 1; + for i in k..n { + for j in k..m { + let val = mat[[i, j]]; + if val != 0 { + let vv = pval(val, p, e); + if vv < lev { + if vv == 0 { + return 0; + } + lev = vv; + } + } + } + } + lev +} - let mut mat = a.mapv(|x| pmod(x, q)); - let mut u = Array2::::eye(n); - let mut v = Array2::::eye(m); +/// Apply a deferred panel update to `target` over columns `[c0, c1)`: +/// first TRSM the panel rows `[k, k+pb)` by the unit-lower-triangular `L11`, +/// then GEMM the trailing rows `[k+pb, n)` as `-= L21 @ panel (mod q)`. +/// `l11[t*pb + s]` (s < t) and `l21[i*pb + s]` are the stored multipliers. +fn apply_block_update( + target: &mut Array2, + l11: &[i64], + l21: &[i64], + k: usize, + pb: usize, + c0: usize, + c1: usize, + q: i64, +) { + if c1 <= c0 { + return; + } + let n = target.nrows(); + let pend = k + pb; + // TRSM: forward substitution against unit-lower-triangular L11. + for t in 1..pb { + for s in 0..t { + let f = l11[t * pb + s]; + if f != 0 { + for col in c0..c1 { + target[[k + t, col]] = + sub_mul_mod(target[[k + t, col]], f, target[[k + s, col]], q); + } + } + } + } + // GEMM trailing update with delayed reduction. + let trail = n - pend; + if q < GEMM_I64_MAX { + for i in 0..trail { + let row = pend + i; + for col in c0..c1 { + let mut acc = target[[row, col]]; + for s in 0..pb { + acc -= l21[i * pb + s] * target[[k + s, col]]; + } + let rr = acc % q; + target[[row, col]] = if rr < 0 { rr + q } else { rr }; + } + } + } else { + for i in 0..trail { + let row = pend + i; + for col in c0..c1 { + let mut acc = target[[row, col]] as i128; + for s in 0..pb { + acc -= l21[i * pb + s] as i128 * target[[k + s, col]] as i128; + } + target[[row, col]] = posmod_i128(acc, q); + } + } + } +} - // ---- Phase L: column elimination -> U, mat becomes upper-triangular ---- - // Minimal-valuation pivoting (which sorts the diagonal by valuation, giving - // the divisibility chain). Row clears are deferred to Phase R; deferring is - // valid because column clears alone produce the upper-triangular Schur form. - for k in 0..r { +/// Blocked right-looking LU over the unit (valuation-0) part of `mat`. +/// Eliminates leading columns whose pivot is a unit (mod p), updating `mat` +/// and `u`; returns the first pivot index it could not place (a column with no +/// unit), which the scalar path then finishes. `v` is untouched (no column +/// swaps here). Pre-condition: `min_valuation(mat, 0, ..) == 0`. +fn phase_l_blocked_unit( + mat: &mut Array2, + u: &mut Array2, + p: i64, + e: u32, + q: i64, + n: usize, + m: usize, + r: usize, +) -> usize { + let mut k = 0; + while k < r { + let pend_max = (k + PANEL_B).min(r); + + // ---- factor the panel columns [k, pend_max): unit row-pivoting ---- + let mut kk = k; + while kk < pend_max { + // find a unit (valuation 0) in column kk, rows [kk, n) + let mut prow = None; + for i in kk..n { + if mat[[i, kk]] != 0 && pval(mat[[i, kk]], p, e) == 0 { + prow = Some(i); + break; + } + } + let pr = match prow { + Some(x) => x, + None => break, // column kk has no unit -> end panel (then scalar) + }; + if pr != kk { + // full-width row swap keeps the deferred block + u consistent + for col in 0..m { + mat.swap([kk, col], [pr, col]); + } + for col in 0..n { + u.swap([kk, col], [pr, col]); + } + } + // Divide-by-pivot multipliers (the pivot is a unit). The pivot row + // is NOT normalized here; the diagonal is normalized once at the end + // so that mat and u undergo identical row operations — u's elimination + // is deferred (TRSM/GEMM via the stored L), and a normalize-then- + // eliminate / eliminate-then-normalize mismatch would corrupt it. + let uinv = inv_mod(mat[[kk, kk]], q); + for i in (kk + 1)..n { + if mat[[i, kk]] != 0 { + let c = mulmod(mat[[i, kk]], uinv, q); // = mat[i,kk] / mat[kk,kk] + for col in (kk + 1)..pend_max { + mat[[i, col]] = sub_mul_mod(mat[[i, col]], c, mat[[kk, col]], q); + } + mat[[i, kk]] = c; // store the multiplier (L) + } + } + kk += 1; + } + let pend = kk; + if pend == k { + return k; // no unit in column k -> hand the rest to the scalar path + } + let pb = pend - k; + + // ---- extract L11 (unit lower-tri) and L21 (trailing-row) multipliers ---- + let trail = n - pend; + let mut l11 = vec![0i64; pb * pb]; + for t in 0..pb { + for s in 0..t { + l11[t * pb + s] = mat[[k + t, k + s]]; + } + } + let mut l21 = vec![0i64; trail * pb]; + for i in 0..trail { + for s in 0..pb { + l21[i * pb + s] = mat[[pend + i, k + s]]; + } + } + + // ---- deferred TRSM + GEMM on mat (cols [pend_max, m)) and u (all cols) ---- + apply_block_update(mat, &l11, &l21, k, pb, pend_max, m, q); + apply_block_update(u, &l11, &l21, k, pb, 0, n, q); + + // ---- zero the L storage below the diagonal of the pivot columns ---- + for t in 0..pb { + for i in (k + t + 1)..n { + mat[[i, k + t]] = 0; + } + } + k = pend; + } + k +} + +/// Scalar Phase L (minimal-valuation column elimination), continuing from +/// `k_start`. Fully general: global min-valuation pivot + row/column swaps. +/// Produces `u` (and column swaps tracked in `v`); leaves `mat` upper-triangular. +fn phase_l_scalar( + mat: &mut Array2, + u: &mut Array2, + v: &mut Array2, + p: i64, + e: u32, + q: i64, + n: usize, + m: usize, + r: usize, + k_start: usize, +) { + for k in k_start..r { // Find the pivot in mat[k.., k..] with minimal p-adic valuation. let mut best: Option<(usize, usize)> = None; let mut bestval = e + 1; @@ -138,11 +316,10 @@ fn local_snf(a: &Array2, p: i64, e: u32) -> (Array2, Array2, Vec< } } - let vv = bestval; - let pv = p.pow(vv); + let pv = p.pow(bestval); // Normalize the pivot to exactly p^vv by scaling row k by a unit. - let unit = mat[[k, k]] / pv; // exact: pivot = p^vv * unit + let unit = mat[[k, k]] / pv; let uinv = inv_mod(unit, q); for col in 0..m { mat[[k, col]] = mulmod(mat[[k, col]], uinv, q); @@ -155,7 +332,7 @@ fn local_snf(a: &Array2, p: i64, e: u32) -> (Array2, Array2, Vec< for i in (k + 1)..n { let val = mat[[i, k]]; if val != 0 { - let c = val / pv; // exact: val(mat[i,k]) >= vv + let c = val / pv; for col in (k + 1)..m { mat[[i, col]] = sub_mul_mod(mat[[i, col]], c, mat[[k, col]], q); } @@ -166,9 +343,53 @@ fn local_snf(a: &Array2, p: i64, e: u32) -> (Array2, Array2, Vec< } } } +} + +/// SNF of `A` over the local ring `Z/p^e` via valuation pivoting. +/// +/// Returns `(U, V, vals)` with `U @ A @ V == diag(p^vals) (mod p^e)`, +/// `U` (n x n) and `V` (m x m) unimodular and `vals` ascending. +fn local_snf(a: &Array2, p: i64, e: u32) -> (Array2, Array2, Vec) { + let q = p.pow(e); + let n = a.nrows(); + let m = a.ncols(); + let r = n.min(m); + + let mut mat = a.mapv(|x| pmod(x, q)); + let mut u = Array2::::eye(n); + let mut v = Array2::::eye(m); + + // ---- Phase L: column elimination -> U, mat becomes upper-triangular ---- + // Blocked LU clears the valuation-0 bulk (trailing update = GEMM); the + // scalar path finishes the higher-valuation / column-swap tail. + let k0 = if min_valuation(&mat, 0, n, m, p, e) == 0 { + phase_l_blocked_unit(&mut mat, &mut u, p, e, q, n, m, r) + } else { + 0 + }; + phase_l_scalar(&mut mat, &mut u, &mut v, p, e, q, n, m, r, k0); let vals: Vec = (0..r).map(|i| pval(mat[[i, i]], p, e)).collect(); + // Normalize each pivot diagonal to exactly p^vals[k]: blocked pivots were + // left as units; scalar pivots are already p^vals so this no-ops for them. + for k in 0..r { + if mat[[k, k]] == 0 { + continue; // zero invariant factor (p^e ≡ 0) + } + let pvk = p.pow(vals[k]); + if mat[[k, k]] != pvk { + let unit = mat[[k, k]] / pvk; // exact: pivot = p^vals * unit + let sc = inv_mod(unit, q); + for col in 0..m { + mat[[k, col]] = mulmod(mat[[k, col]], sc, q); + } + for col in 0..n { + u[[k, col]] = mulmod(u[[k, col]], sc, q); + } + } + } + // ---- Phase R: diagonalize the upper-triangular mat -> V ---- // Clear super-diagonal entries by column operations, processing pivots from // last to first so that fill created above row k lands in not-yet-processed @@ -395,34 +616,7 @@ mod tests { for m in [1usize, 3, 5, 9, 16] { for _trial in 0..6 { let a = Array2::from_shape_fn((n, m), |_| rng.below(q)); - let (uu, vv, vals) = local_snf(&a, p, e); - - // U @ A @ V == diag(p^vals) (mod q) - let prod = matmul_mod(&matmul_mod(&uu, &a, q), &vv, q); - let r = n.min(m); - for i in 0..n { - for j in 0..m { - let expect = if i == j && i < r { - pmod(p.pow(vals[i]), q) - } else { - 0 - }; - assert_eq!( - prod[[i, j]], expect, - "U*A*V mismatch at ({i},{j}) p={p} e={e} n={n} m={m}" - ); - } - } - // valuations ascending (divisibility chain) - for i in 1..r { - assert!( - vals[i - 1] <= vals[i], - "vals not ascending: {vals:?} p={p} e={e}" - ); - } - // U (n x n) and V (m x m) unimodular - assert!(invertible_mod_p(&uu, p), "U not unimodular p={p} e={e}"); - assert!(invertible_mod_p(&vv, p), "V not unimodular p={p} e={e}"); + assert_valid_snf(&a, p, e); checked += 1; } } @@ -430,4 +624,75 @@ mod tests { } assert!(checked > 500, "expected many cases, got {checked}"); } + + /// Assert local_snf(a, p, e) is a valid Smith form: U@A@V == diag(p^vals), + /// ascending valuations, U,V unimodular. + fn assert_valid_snf(a: &Array2, p: i64, e: u32) { + let q = p.pow(e); + let (n, m) = (a.nrows(), a.ncols()); + let r = n.min(m); + let (uu, vv, vals) = local_snf(a, p, e); + let prod = matmul_mod(&matmul_mod(&uu, a, q), &vv, q); + for i in 0..n { + for j in 0..m { + let expect = if i == j && i < r { pmod(p.pow(vals[i]), q) } else { 0 }; + assert_eq!(prod[[i, j]], expect, "U*A*V at ({i},{j}) p={p} e={e} n={n} m={m}"); + } + } + for i in 1..r { + assert!(vals[i - 1] <= vals[i], "vals not ascending {vals:?} p={p} e={e}"); + } + assert!(invertible_mod_p(&uu, p), "U not unimodular p={p} e={e} n={n} m={m}"); + assert!(invertible_mod_p(&vv, p), "V not unimodular p={p} e={e} n={n} m={m}"); + } + + /// Exercise the blocked-LU paths: large n (multi-panel), rank-deficiency + /// mod p (forces the blocked -> scalar handoff), and p|A (blocked skipped). + #[test] + fn local_snf_blocked_paths() { + let cases = [(2u32, 1u32), (2, 4), (3, 2), (3, 3), (5, 2), (7, 2)]; + let mut rng = Lcg(0xdead_beef_0bad_f00d); + let mut checked = 0; + for &(p, e) in &cases { + let p = p as i64; + let q = p.pow(e); + // Larger square sizes spanning the panel width (PANEL_B = 48). + for &n in &[40usize, 49, 64, 97] { + // (1) dense random: generically full-rank mod p (single long blocked run) + let a = Array2::from_shape_fn((n, n), |_| rng.below(q)); + assert_valid_snf(&a, p, e); + checked += 1; + + // (2) low rank mod p: A = B @ C with inner dim rk << n, so the + // blocked unit pass exhausts units early and hands to scalar. + for &rk in &[1usize, 3, 7] { + let b = Array2::from_shape_fn((n, rk), |_| rng.below(q)); + let c = Array2::from_shape_fn((rk, n), |_| rng.below(q)); + let a = matmul_mod(&b, &c, q); + assert_valid_snf(&a, p, e); + checked += 1; + } + + // (3) p | A everywhere: min valuation >= 1, blocked pass skipped. + let a = Array2::from_shape_fn((n, n), |_| (rng.below(q / p.max(1)) * p) % q); + assert_valid_snf(&a, p, e); + checked += 1; + + // (4) mixed: a unit block plus a p-scaled block (multi-level). + let a = Array2::from_shape_fn((n, n), |(i, j)| { + if (i + j) % 3 == 0 { rng.below(q) } else { (rng.below(q) * p) % q } + }); + assert_valid_snf(&a, p, e); + checked += 1; + + // (5) rectangular, both orientations. + let a = Array2::from_shape_fn((n, n / 2 + 1), |_| rng.below(q)); + assert_valid_snf(&a, p, e); + let a = Array2::from_shape_fn((n / 2 + 1, n), |_| rng.below(q)); + assert_valid_snf(&a, p, e); + checked += 2; + } + } + assert!(checked > 100, "expected many cases, got {checked}"); + } } From 328f6c8993bc8fc80b3e362937585a25b307fc4f Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Wed, 3 Jun 2026 01:14:38 -0400 Subject: [PATCH 14/22] docs: mark Phase 2b done; note tiled/BLAS GEMM kernel as Phase 2c --- docs/PLAN.md | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/docs/PLAN.md b/docs/PLAN.md index 949d4cb..d01037b 100644 --- a/docs/PLAN.md +++ b/docs/PLAN.md @@ -61,7 +61,20 @@ row-clear touched all n rows of each column). ~1.1–1.3× over Phase 1. Gated b randomized cargo oracle test (`local_snf_is_valid_smith_form`, 500+ cases: `U·A·V==diag(p^vals)`, ascending vals, `U,V` unimodular) + head-to-head vs dev. -**2b — blocked GEMM trailing update (NEXT, designed).** Key constraint: +**2b — blocked GEMM trailing update (DONE).** Implemented as right-looking +blocked LU over the valuation-0 bulk (panel B=48): in-place panel factorization, +TRSM for the deferred panel rows, GEMM trailing update (i64 accumulate + delayed +reduction) on `mat` and `u`; scalar min-valuation path finishes the higher- +valuation/rank-deficient tail; final diagonal normalization. ~1.3–1.5× over 2a +at n=400, 1.6–2.6× over Phase 1 at n=1000–2000 (cumulative ~3× over original +CRT). Gated by two cargo tests (600+ cases incl. multi-panel, rank-deficiency +mod p, p|A, rectangular) + head-to-head vs dev. +NOTE: the GEMM is a naive i64 triple loop — itself not cache-tiled, which is why +n=2000 only reaches 1.58×. Swapping it for a tiled / `faer` / BLAS kernel is the +next lever (Phase 2c) before rayon. + +Original design notes (for reference): +Key constraint: minimal-valuation pivoting needs the *global* min over the trailing block, which defeats naive column-panel/lazy blocking. The blockable formulation is **valuation-level staging**: process pivots by p-adic level ℓ = 0..e; within a From 86a36d7e993d0bbde964d1771cfc1d9e231bf734 Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Wed, 3 Jun 2026 01:30:32 -0400 Subject: [PATCH 15/22] perf(crt): packed cache-friendly i64 GEMM micro-kernel (Phase 2c) The naive 2b trailing GEMM read panel rows target[k+s,col] with stride m across s -> ~pb distinct cache lines per output element. Pack U12 (post-TRSM panel rows) once into a contiguous (pb x ncols) buffer, then update each trailing row with contiguous axpy passes (autovectorizes); i128 accumulator fallback for large moduli. Tried an f64 GEMM via ndarray/matrixmultiply first: helped single-prime moduli but REGRESSED multi-prime (N=360: 0.86x) due to per-panel n*n product allocation + thin-K (pb=48). The packed i64 kernel wins across both: ~1.34x (n=1000,N=2), 1.24x (n=1000,N=360), 1.10x (n=2000) over Phase 2b. Correctness unchanged (cargo tests + head-to-head vs dev: no mismatches, all transform_ok). --- crates/modularsnf/src/crt.rs | 59 +++++++++++++++++++++++++++--------- 1 file changed, 45 insertions(+), 14 deletions(-) diff --git a/crates/modularsnf/src/crt.rs b/crates/modularsnf/src/crt.rs index 57ac520..bd913e2 100644 --- a/crates/modularsnf/src/crt.rs +++ b/crates/modularsnf/src/crt.rs @@ -81,8 +81,8 @@ fn pval(mut x: i64, p: i64, cap: u32) -> u32 { /// Block width for the right-looking blocked LU (Phase L, unit stage). const PANEL_B: usize = 48; -/// i64 GEMM accumulator is safe while `PANEL_B * (q-1)^2 < 2^63`. q < 2^28 -/// gives 48 * 2^56 < 2^62, so use i64 below it and i128 at/above it. +/// The i64 GEMM accumulator is safe while `PANEL_B * (q-1)^2 < 2^63`. q < 2^28 +/// gives 48 * 2^56 < 2^62, so use the i64 micro-kernel below it, i128 above. const GEMM_I64_MAX: i64 = 1i64 << 28; /// Minimal p-adic valuation over `mat[k.., k..]` (returns `e+1` if all zero). @@ -136,29 +136,60 @@ fn apply_block_update( } } } - // GEMM trailing update with delayed reduction. + // GEMM trailing update: trailing rows -= L21 @ U12 (mod q). + // Pack U12 (post-TRSM panel rows) into a contiguous (pb x ncols) buffer so + // the inner axpy streams contiguously instead of striding by row length. let trail = n - pend; + if trail == 0 { + return; + } + let ncols = c1 - c0; + let mut rpack = vec![0i64; pb * ncols]; + for s in 0..pb { + let base = s * ncols; + for j in 0..ncols { + rpack[base + j] = target[[k + s, c0 + j]]; + } + } if q < GEMM_I64_MAX { + let mut acc = vec![0i64; ncols]; for i in 0..trail { let row = pend + i; - for col in c0..c1 { - let mut acc = target[[row, col]]; - for s in 0..pb { - acc -= l21[i * pb + s] * target[[k + s, col]]; + for j in 0..ncols { + acc[j] = target[[row, c0 + j]]; + } + for s in 0..pb { + let l = l21[i * pb + s]; + if l != 0 { + let base = s * ncols; + for j in 0..ncols { + acc[j] -= l * rpack[base + j]; // contiguous, autovectorizes + } } - let rr = acc % q; - target[[row, col]] = if rr < 0 { rr + q } else { rr }; + } + for j in 0..ncols { + let rr = acc[j] % q; + target[[row, c0 + j]] = if rr < 0 { rr + q } else { rr }; } } } else { + let mut acc = vec![0i128; ncols]; for i in 0..trail { let row = pend + i; - for col in c0..c1 { - let mut acc = target[[row, col]] as i128; - for s in 0..pb { - acc -= l21[i * pb + s] as i128 * target[[k + s, col]] as i128; + for j in 0..ncols { + acc[j] = target[[row, c0 + j]] as i128; + } + for s in 0..pb { + let l = l21[i * pb + s] as i128; + if l != 0 { + let base = s * ncols; + for j in 0..ncols { + acc[j] -= l * rpack[base + j] as i128; + } } - target[[row, col]] = posmod_i128(acc, q); + } + for j in 0..ncols { + target[[row, c0 + j]] = posmod_i128(acc[j], q); } } } From 624e165e2b7d745ca0fa5ba4242a3b00189728f9 Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Wed, 3 Jun 2026 01:30:50 -0400 Subject: [PATCH 16/22] docs: mark Phase 2c done (packed i64 GEMM kernel) --- docs/PLAN.md | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/docs/PLAN.md b/docs/PLAN.md index d01037b..3de2678 100644 --- a/docs/PLAN.md +++ b/docs/PLAN.md @@ -69,9 +69,16 @@ valuation/rank-deficient tail; final diagonal normalization. ~1.3–1.5× over 2 at n=400, 1.6–2.6× over Phase 1 at n=1000–2000 (cumulative ~3× over original CRT). Gated by two cargo tests (600+ cases incl. multi-panel, rank-deficiency mod p, p|A, rectangular) + head-to-head vs dev. -NOTE: the GEMM is a naive i64 triple loop — itself not cache-tiled, which is why -n=2000 only reaches 1.58×. Swapping it for a tiled / `faer` / BLAS kernel is the -next lever (Phase 2c) before rayon. +**2c — packed cache-friendly GEMM micro-kernel (DONE).** The 2b GEMM read panel +rows with stride-m across the contraction index (~pb cache lines per output). +Fix: pack U12 into a contiguous (pb×ncols) buffer, then contiguous-axpy each +trailing row (autovectorizes); i128 fallback for large q. ~1.10–1.34× over 2b +across single- and multi-prime moduli. (An f64 ndarray/matrixmultiply GEMM was +tried first but regressed multi-prime — n*n product allocation + thin K=48 — so +the packed i64 kernel was kept.) Cumulative ~3.5–4× over original CRT. +NEXT: a true tiled/register-blocked or BLAS/faer kernel could push large-n +further (n=2000 still ~1.1× since the micro-kernel isn't register-tiled), but +diminishing returns vs Phase 3 (rayon). Original design notes (for reference): Key constraint: From fc02385a78d9c616d5e54cef00868b180f0d9d6b Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Sat, 6 Jun 2026 00:42:27 -0400 Subject: [PATCH 17/22] fix(snf): keep transforms unimodular for tall rank-deficient inputs Tall matrices (n > m) were padded to square and had V cropped, which could leave V non-unimodular for column-rank-deficient inputs (an all-zero column gave V = [[0]]). U A V == S still held, so the old tests missed it. Route tall matrices through the transpose, where the wide path crops U (which stays unimodular) instead of V. Add a Rust correctness oracle for smith_normal_form (U A V == S, divisibility chain, unimodularity over random inputs) plus a Storjohann-vs-CRT invariant-factor cross-check, which is what surfaced the bug. Tests use rand's seeded StdRng (dev-dependency). --- Cargo.lock | 77 +++++++++++++++ crates/modularsnf/Cargo.toml | 3 + crates/modularsnf/src/lib.rs | 182 +++++++++++++++++++++++++++++++++++ 3 files changed, 262 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index ecf5101..e9fcd04 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -14,6 +14,17 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + [[package]] name = "heck" version = "0.5.0" @@ -59,6 +70,7 @@ name = "modularsnf" version = "0.4.0" dependencies = [ "ndarray", + "rand", ] [[package]] @@ -150,6 +162,15 @@ dependencies = [ "portable-atomic", ] +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + [[package]] name = "proc-macro2" version = "1.0.106" @@ -231,6 +252,36 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rand" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +dependencies = [ + "libc", + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom", +] + [[package]] name = "rawpointer" version = "0.2.1" @@ -277,3 +328,29 @@ name = "unindent" version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "zerocopy" +version = "0.8.50" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b065d4f0e55f82fae73202e189638116a87c55ab6b8e6c2721e13dd9d854ad1" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.50" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b631b19d36a892ab55420c92dbc83ccd79274f25be714855d3074aa71cab639" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] diff --git a/crates/modularsnf/Cargo.toml b/crates/modularsnf/Cargo.toml index 51f0dcd..dc06694 100644 --- a/crates/modularsnf/Cargo.toml +++ b/crates/modularsnf/Cargo.toml @@ -11,3 +11,6 @@ name = "modularsnf" [dependencies] ndarray = "0.16" + +[dev-dependencies] +rand = "0.8" diff --git a/crates/modularsnf/src/lib.rs b/crates/modularsnf/src/lib.rs index aa08e25..0b5c7a9 100644 --- a/crates/modularsnf/src/lib.rs +++ b/crates/modularsnf/src/lib.rs @@ -25,6 +25,16 @@ pub fn smith_normal_form( return Ok((u, v, a.clone())); } + // Tall matrices (n > m) are routed through the transpose: padding columns + // and cropping V can leave V non-unimodular for column-rank-deficient + // inputs (e.g. zero columns), whereas the wide path crops U, which stays + // unimodular. From `U' Aᵀ V' = S'` we get `(V'ᵀ) A (U'ᵀ) = S'ᵀ`. + if n > m { + let at = a.t().to_owned(); + let (u_t, v_t, s_t) = smith_normal_form(&at, modulus)?; + return Ok((v_t.t().to_owned(), u_t.t().to_owned(), s_t.t().to_owned())); + } + // Pad to square if needed let s = n.max(m); let mut a_pad = Array2::zeros((s, s)); @@ -39,3 +49,175 @@ pub fn smith_normal_form( Ok((u, v, s_mat)) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::crt::crt_snf; + use crate::ring::posmod_i128; + use ndarray::Array2; + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + + fn gcd_i64(mut a: i64, mut b: i64) -> i64 { + while b != 0 { + let t = b; + b = a % b; + a = t; + } + a.abs() + } + + fn matmul_mod(a: &Array2, b: &Array2, q: i64) -> Array2 { + let (n, k) = (a.nrows(), a.ncols()); + let m = b.ncols(); + let mut c = Array2::::zeros((n, m)); + for i in 0..n { + for j in 0..m { + let mut acc: i128 = 0; + for t in 0..k { + acc += a[[i, t]] as i128 * b[[t, j]] as i128; + } + c[[i, j]] = posmod_i128(acc, q); + } + } + c + } + + /// Exact integer determinant by cofactor expansion (small n only). + fn det_i128(a: &Array2) -> i128 { + let n = a.nrows(); + let mut m: Vec> = (0..n) + .map(|i| (0..n).map(|j| a[[i, j]] as i128).collect()) + .collect(); + det_recur(&mut m) + } + + fn det_recur(m: &mut [Vec]) -> i128 { + let n = m.len(); + if n == 1 { + return m[0][0]; + } + if n == 2 { + return m[0][0] * m[1][1] - m[0][1] * m[1][0]; + } + let mut det = 0i128; + for j in 0..n { + let mut sub: Vec> = m[1..] + .iter() + .map(|row| { + row.iter() + .enumerate() + .filter(|&(c, _)| c != j) + .map(|(_, &v)| v) + .collect() + }) + .collect(); + let sign = if j % 2 == 0 { 1 } else { -1 }; + det += sign * m[0][j] * det_recur(&mut sub); + } + det + } + + /// gcd(det(M) mod N, N) == 1: M is unimodular over Z/NZ. + fn is_unimodular(m: &Array2, n: i64) -> bool { + let d = posmod_i128(det_i128(m), n); + gcd_i64(d, n) == 1 + } + + /// Normalized invariant factors: gcd(diag, N), ascending. + fn normalized_invariants(s: &Array2, n: i64) -> Vec { + let r = s.nrows().min(s.ncols()); + let mut inv: Vec = (0..r).map(|i| gcd_i64(s[[i, i]], n)).collect(); + inv.sort_unstable(); + inv + } + + fn assert_valid_snf(a: &Array2, n: i64, check_unimodular: bool) { + let (rows, cols) = (a.nrows(), a.ncols()); + let (u, v, s) = smith_normal_form(a, n).expect("smith_normal_form"); + + assert_eq!(u.dim(), (rows, rows)); + assert_eq!(v.dim(), (cols, cols)); + assert_eq!(s.dim(), (rows, cols)); + + // U A V == S (mod N) + let prod = matmul_mod(&matmul_mod(&u, a, n), &v, n); + for i in 0..rows { + for j in 0..cols { + assert_eq!( + prod[[i, j]], + posmod_i128(s[[i, j]] as i128, n), + "U A V != S at ({i},{j}) N={n} shape={rows}x{cols}" + ); + } + } + + // S diagonal. + for i in 0..rows { + for j in 0..cols { + if i != j { + assert_eq!(s[[i, j]], 0, "S not diagonal at ({i},{j})"); + } + } + } + + // Divisibility chain on ideal generators: g_i | g_{i+1}. + let inv = normalized_invariants(&s, n); + for w in inv.windows(2) { + assert_eq!(w[1] % w[0], 0, "invariants not a chain: {inv:?} N={n}"); + } + + if check_unimodular { + assert!(is_unimodular(&u, n), "U not unimodular N={n}"); + assert!(is_unimodular(&v, n), "V not unimodular N={n}"); + } + } + + #[test] + fn smith_normal_form_is_valid() { + let moduli = [2i64, 4, 6, 8, 9, 12, 30, 36, 100]; + let shapes = [(1, 1), (3, 3), (4, 6), (6, 4), (5, 1), (1, 5), (7, 7)]; + let mut rng = StdRng::seed_from_u64(0x0bad_c0de_1234_5678); + let mut checked = 0; + for &n in &moduli { + for &(rows, cols) in &shapes { + for _ in 0..4 { + let a = Array2::from_shape_fn((rows, cols), |_| rng.gen_range(0..n)); + // det helper is O(k!); only check unimodularity for small k. + assert_valid_snf(&a, n, rows.max(cols) <= 7); + checked += 1; + } + } + } + assert!(checked > 200, "expected many cases, got {checked}"); + } + + #[test] + fn smith_normal_form_matches_crt() { + let cases: [(i64, Vec<(i64, u32)>); 6] = [ + (6, vec![(2, 1), (3, 1)]), + (12, vec![(2, 2), (3, 1)]), + (36, vec![(2, 2), (3, 2)]), + (8, vec![(2, 3)]), + (30, vec![(2, 1), (3, 1), (5, 1)]), + (100, vec![(2, 2), (5, 2)]), + ]; + let shapes = [(3, 3), (4, 6), (6, 4), (5, 5), (1, 4)]; + let mut rng = StdRng::seed_from_u64(0xfeed_face_dead_beef); + for (n, factors) in &cases { + for &(rows, cols) in &shapes { + for _ in 0..4 { + let a = Array2::from_shape_fn((rows, cols), |_| rng.gen_range(0..*n)); + let (_, _, s_storj) = smith_normal_form(&a, *n).expect("storjohann"); + let (_, _, s_crt) = crt_snf(&a, *n, factors); + assert_eq!( + normalized_invariants(&s_storj, *n), + normalized_invariants(&s_crt, *n), + "invariant mismatch N={n} shape={rows}x{cols}" + ); + } + } + } + } +} From fd3934089929881d25f21ab8956fc7e50baadf10 Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Sat, 6 Jun 2026 00:42:34 -0400 Subject: [PATCH 18/22] refactor(crt): remove dev scaffolding and tidy the implementation - Drop the CRT_PROFILE timing instrumentation and phase-by-phase development comments from crt.rs; keep the algorithmic invariants (i64 overflow bounds, deferred normalization). - Replace section-divider comments and TRSM/GEMM shorthand with prose matching the rest of the crate; rustfmt ring.rs/snf.rs. - Seed crt.rs tests with rand's StdRng instead of a hand-rolled LCG. - Replace the dev roadmap (docs/PLAN.md) and the CRT prototype/profiling scripts with a concise design note (docs/crt.md). --- crates/modularsnf/src/crt.rs | 171 ++++++------ crates/modularsnf/src/ring.rs | 5 +- crates/modularsnf/src/snf.rs | 5 +- docs/PLAN.md | 133 --------- docs/crt.md | 69 +++++ scripts/crt_snf_profile.py | 75 ----- scripts/crt_snf_prototype.py | 207 -------------- scripts/crt_snf_verify.py | 94 ------- scripts/profile_snf.py | 502 ---------------------------------- 9 files changed, 157 insertions(+), 1104 deletions(-) delete mode 100644 docs/PLAN.md create mode 100644 docs/crt.md delete mode 100644 scripts/crt_snf_profile.py delete mode 100644 scripts/crt_snf_prototype.py delete mode 100644 scripts/crt_snf_verify.py delete mode 100755 scripts/profile_snf.py diff --git a/crates/modularsnf/src/crt.rs b/crates/modularsnf/src/crt.rs index bd913e2..42f95a3 100644 --- a/crates/modularsnf/src/crt.rs +++ b/crates/modularsnf/src/crt.rs @@ -15,15 +15,20 @@ fn pmod(a: i64, n: i64) -> i64 { posmod_i128(a as i128, n) } -/// Threshold below which q^2 (and a - b*c with a,b,c in [0,q)) fits in i64, -/// so the hot loops avoid i128. q < 2^31 => q^2 < 2^62, comfortably in range. +/// Moduli below this bound use i64 arithmetic in the hot loops; larger moduli +/// widen to i128. For `q < 2^31`, products `q^2` stay under 2^62 and cannot +/// overflow i64. const I64_FAST_MAX: i64 = 1i64 << 31; #[inline] fn mulmod(a: i64, b: i64, n: i64) -> i64 { if n < I64_FAST_MAX { let r = (a * b) % n; - if r < 0 { r + n } else { r } + if r < 0 { + r + n + } else { + r + } } else { posmod_i128(a as i128 * b as i128, n) } @@ -35,7 +40,11 @@ fn mulmod(a: i64, b: i64, n: i64) -> i64 { fn sub_mul_mod(a: i64, b: i64, c: i64, q: i64) -> i64 { if q < I64_FAST_MAX { let r = (a - b * c) % q; - if r < 0 { r + q } else { r } + if r < 0 { + r + q + } else { + r + } } else { posmod_i128(a as i128 - b as i128 * c as i128, q) } @@ -79,10 +88,11 @@ fn pval(mut x: i64, p: i64, cap: u32) -> u32 { v } -/// Block width for the right-looking blocked LU (Phase L, unit stage). +/// Panel width for the right-looking blocked LU in Phase L. const PANEL_B: usize = 48; -/// The i64 GEMM accumulator is safe while `PANEL_B * (q-1)^2 < 2^63`. q < 2^28 -/// gives 48 * 2^56 < 2^62, so use the i64 micro-kernel below it, i128 above. +/// Moduli below this bound use the i64 GEMM micro-kernel; larger moduli use the +/// i128 kernel. The i64 accumulator sums `PANEL_B` products without overflow +/// while `PANEL_B * (q-1)^2 < 2^63`; `q < 2^28` gives `48 * 2^56 < 2^62`. const GEMM_I64_MAX: i64 = 1i64 << 28; /// Minimal p-adic valuation over `mat[k.., k..]` (returns `e+1` if all zero). @@ -105,10 +115,15 @@ fn min_valuation(mat: &Array2, k: usize, n: usize, m: usize, p: i64, e: u32 lev } -/// Apply a deferred panel update to `target` over columns `[c0, c1)`: -/// first TRSM the panel rows `[k, k+pb)` by the unit-lower-triangular `L11`, -/// then GEMM the trailing rows `[k+pb, n)` as `-= L21 @ panel (mod q)`. -/// `l11[t*pb + s]` (s < t) and `l21[i*pb + s]` are the stored multipliers. +/// Apply a factored panel's deferred update to `target`, over columns +/// `[c0, c1)`. This is the standard two-step blocked-LU update, all mod `q`: +/// +/// 1. triangular solve — update the panel rows `[k, k+pb)` against the unit +/// lower-triangular factor `L11`; +/// 2. matrix multiply — update the trailing rows `[k+pb, n)` by `-= L21 @ U12`. +/// +/// `l11` and `l21` store the panel's multipliers in row-major order, indexed +/// `l11[t*pb + s]` (with `s < t`) and `l21[i*pb + s]`. fn apply_block_update( target: &mut Array2, l11: &[i64], @@ -124,7 +139,7 @@ fn apply_block_update( } let n = target.nrows(); let pend = k + pb; - // TRSM: forward substitution against unit-lower-triangular L11. + // Triangular solve: forward substitution against unit-lower-triangular L11. for t in 1..pb { for s in 0..t { let f = l11[t * pb + s]; @@ -136,9 +151,9 @@ fn apply_block_update( } } } - // GEMM trailing update: trailing rows -= L21 @ U12 (mod q). - // Pack U12 (post-TRSM panel rows) into a contiguous (pb x ncols) buffer so - // the inner axpy streams contiguously instead of striding by row length. + // Trailing update: rows -= L21 @ U12 (mod q). Pack U12 (the panel rows, + // post-solve) into a contiguous pb-by-ncols buffer so the inner loop reads + // sequentially instead of striding by the matrix row length. let trail = n - pend; if trail == 0 { return; @@ -163,7 +178,7 @@ fn apply_block_update( if l != 0 { let base = s * ncols; for j in 0..ncols { - acc[j] -= l * rpack[base + j]; // contiguous, autovectorizes + acc[j] -= l * rpack[base + j]; } } } @@ -196,10 +211,11 @@ fn apply_block_update( } /// Blocked right-looking LU over the unit (valuation-0) part of `mat`. -/// Eliminates leading columns whose pivot is a unit (mod p), updating `mat` -/// and `u`; returns the first pivot index it could not place (a column with no -/// unit), which the scalar path then finishes. `v` is untouched (no column -/// swaps here). Pre-condition: `min_valuation(mat, 0, ..) == 0`. +/// +/// Eliminates leading columns that have a unit pivot (mod p), updating `mat` +/// and the left transform `u`. Returns the first column with no unit pivot, +/// where the scalar path takes over. `v` is untouched, as this stage performs +/// no column swaps. Precondition: `min_valuation(mat, 0, ..) == 0`. fn phase_l_blocked_unit( mat: &mut Array2, u: &mut Array2, @@ -214,7 +230,7 @@ fn phase_l_blocked_unit( while k < r { let pend_max = (k + PANEL_B).min(r); - // ---- factor the panel columns [k, pend_max): unit row-pivoting ---- + // Factor the panel columns [k, pend_max) with unit row-pivoting. let mut kk = k; while kk < pend_max { // find a unit (valuation 0) in column kk, rows [kk, n) @@ -238,11 +254,11 @@ fn phase_l_blocked_unit( u.swap([kk, col], [pr, col]); } } - // Divide-by-pivot multipliers (the pivot is a unit). The pivot row - // is NOT normalized here; the diagonal is normalized once at the end - // so that mat and u undergo identical row operations — u's elimination - // is deferred (TRSM/GEMM via the stored L), and a normalize-then- - // eliminate / eliminate-then-normalize mismatch would corrupt it. + // Compute divide-by-pivot multipliers (the pivot is a unit). The + // pivot row is left un-normalized; diagonals are normalized once at + // the end so that `mat` and `u` undergo identical row operations. + // `u`'s elimination is deferred via the stored multipliers `L`, so + // normalizing here would desynchronize the two and corrupt `u`. let uinv = inv_mod(mat[[kk, kk]], q); for i in (kk + 1)..n { if mat[[i, kk]] != 0 { @@ -261,7 +277,8 @@ fn phase_l_blocked_unit( } let pb = pend - k; - // ---- extract L11 (unit lower-tri) and L21 (trailing-row) multipliers ---- + // Extract the L11 (unit lower-triangular) and L21 (trailing-row) + // multipliers. let trail = n - pend; let mut l11 = vec![0i64; pb * pb]; for t in 0..pb { @@ -276,11 +293,11 @@ fn phase_l_blocked_unit( } } - // ---- deferred TRSM + GEMM on mat (cols [pend_max, m)) and u (all cols) ---- + // Apply the deferred update to mat (cols [pend_max, m)) and u (all cols). apply_block_update(mat, &l11, &l21, k, pb, pend_max, m, q); apply_block_update(u, &l11, &l21, k, pb, 0, n, q); - // ---- zero the L storage below the diagonal of the pivot columns ---- + // Zero the stored multipliers below the diagonal of the pivot columns. for t in 0..pb { for i in (k + t + 1)..n { mat[[i, k + t]] = 0; @@ -390,7 +407,7 @@ fn local_snf(a: &Array2, p: i64, e: u32) -> (Array2, Array2, Vec< let mut u = Array2::::eye(n); let mut v = Array2::::eye(m); - // ---- Phase L: column elimination -> U, mat becomes upper-triangular ---- + // Phase L: column elimination produces U and leaves mat upper-triangular. // Blocked LU clears the valuation-0 bulk (trailing update = GEMM); the // scalar path finishes the higher-valuation / column-swap tail. let k0 = if min_valuation(&mat, 0, n, m, p, e) == 0 { @@ -421,7 +438,7 @@ fn local_snf(a: &Array2, p: i64, e: u32) -> (Array2, Array2, Vec< } } - // ---- Phase R: diagonalize the upper-triangular mat -> V ---- + // Phase R: diagonalize the upper-triangular mat, producing V. // Clear super-diagonal entries by column operations, processing pivots from // last to first so that fill created above row k lands in not-yet-processed // rows (and is cleared when those pivots are reached). col k of the upper- @@ -480,13 +497,8 @@ pub fn crt_snf( let qs: Vec = factors.iter().map(|&(p, e)| p.pow(e)).collect(); - // Phase 0 profiling: set CRT_PROFILE=1 to print a phase-time breakdown. - let prof = std::env::var("CRT_PROFILE").is_ok(); - let t_local = std::time::Instant::now(); - let mut locals: Vec<(Array2, Array2, Vec)> = factors.iter().map(|&(p, e)| local_snf(a, p, e)).collect(); - let d_local = t_local.elapsed(); // Global invariant factors d_i = prod_p p^{vals_p[i]} (divides N). let mut d = vec![0i64; r]; @@ -499,7 +511,6 @@ pub fn crt_snf( } // Unit-normalize each prime's V so all primes realize the same d_i. - let t_norm = std::time::Instant::now(); for pi in 0..np { let q = qs[pi]; for i in 0..r { @@ -517,10 +528,7 @@ pub fn crt_snf( } } - let d_norm = t_norm.elapsed(); - // CRT-recombine U (n x n) and V (m x m) entrywise across primes. - let t_recomb = std::time::Instant::now(); let mut u = Array2::::zeros((n, n)); let mut resid = vec![0i64; np]; for ar in 0..n { @@ -541,29 +549,11 @@ pub fn crt_snf( } } - let d_recomb = t_recomb.elapsed(); - let mut s = Array2::::zeros((n, m)); for (i, &di) in d.iter().enumerate() { s[[i, i]] = di; } - if prof { - let total = d_local + d_norm + d_recomb; - eprintln!( - "[CRT_PROFILE] n={n} m={m} N={modulus} np={np} | \ - local_snf={:.3}ms ({:.0}%) normalize={:.3}ms ({:.0}%) \ - recombine={:.3}ms ({:.0}%) | total={:.3}ms", - d_local.as_secs_f64() * 1e3, - 100.0 * d_local.as_secs_f64() / total.as_secs_f64(), - d_norm.as_secs_f64() * 1e3, - 100.0 * d_norm.as_secs_f64() / total.as_secs_f64(), - d_recomb.as_secs_f64() * 1e3, - 100.0 * d_recomb.as_secs_f64() / total.as_secs_f64(), - total.as_secs_f64() * 1e3, - ); - } - (u, v, s) } @@ -571,18 +561,8 @@ pub fn crt_snf( mod tests { use super::*; use ndarray::Array2; - - /// Deterministic LCG so tests need no external rng dependency. - struct Lcg(u64); - impl Lcg { - fn next(&mut self) -> u64 { - self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); - self.0 >> 16 - } - fn below(&mut self, n: i64) -> i64 { - (self.next() % n as u64) as i64 - } - } + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; fn matmul_mod(a: &Array2, b: &Array2, q: i64) -> Array2 { let (n, k) = (a.nrows(), a.ncols()); @@ -638,7 +618,7 @@ mod tests { #[test] fn local_snf_is_valid_smith_form() { let cases = [(2u32, 1u32), (2, 3), (3, 2), (5, 1), (5, 2), (7, 3)]; - let mut rng = Lcg(0x1234_5678_9abc_def0); + let mut rng = StdRng::seed_from_u64(0x1234_5678_9abc_def0); let mut checked = 0; for &(p, e) in &cases { let p = p as i64; @@ -646,7 +626,7 @@ mod tests { for n in [1usize, 2, 5, 9, 16, 23] { for m in [1usize, 3, 5, 9, 16] { for _trial in 0..6 { - let a = Array2::from_shape_fn((n, m), |_| rng.below(q)); + let a = Array2::from_shape_fn((n, m), |_| rng.gen_range(0..q)); assert_valid_snf(&a, p, e); checked += 1; } @@ -666,15 +646,32 @@ mod tests { let prod = matmul_mod(&matmul_mod(&uu, a, q), &vv, q); for i in 0..n { for j in 0..m { - let expect = if i == j && i < r { pmod(p.pow(vals[i]), q) } else { 0 }; - assert_eq!(prod[[i, j]], expect, "U*A*V at ({i},{j}) p={p} e={e} n={n} m={m}"); + let expect = if i == j && i < r { + pmod(p.pow(vals[i]), q) + } else { + 0 + }; + assert_eq!( + prod[[i, j]], + expect, + "U*A*V at ({i},{j}) p={p} e={e} n={n} m={m}" + ); } } for i in 1..r { - assert!(vals[i - 1] <= vals[i], "vals not ascending {vals:?} p={p} e={e}"); + assert!( + vals[i - 1] <= vals[i], + "vals not ascending {vals:?} p={p} e={e}" + ); } - assert!(invertible_mod_p(&uu, p), "U not unimodular p={p} e={e} n={n} m={m}"); - assert!(invertible_mod_p(&vv, p), "V not unimodular p={p} e={e} n={n} m={m}"); + assert!( + invertible_mod_p(&uu, p), + "U not unimodular p={p} e={e} n={n} m={m}" + ); + assert!( + invertible_mod_p(&vv, p), + "V not unimodular p={p} e={e} n={n} m={m}" + ); } /// Exercise the blocked-LU paths: large n (multi-panel), rank-deficiency @@ -682,7 +679,7 @@ mod tests { #[test] fn local_snf_blocked_paths() { let cases = [(2u32, 1u32), (2, 4), (3, 2), (3, 3), (5, 2), (7, 2)]; - let mut rng = Lcg(0xdead_beef_0bad_f00d); + let mut rng = StdRng::seed_from_u64(0xdead_beef_0bad_f00d); let mut checked = 0; for &(p, e) in &cases { let p = p as i64; @@ -690,36 +687,40 @@ mod tests { // Larger square sizes spanning the panel width (PANEL_B = 48). for &n in &[40usize, 49, 64, 97] { // (1) dense random: generically full-rank mod p (single long blocked run) - let a = Array2::from_shape_fn((n, n), |_| rng.below(q)); + let a = Array2::from_shape_fn((n, n), |_| rng.gen_range(0..q)); assert_valid_snf(&a, p, e); checked += 1; // (2) low rank mod p: A = B @ C with inner dim rk << n, so the // blocked unit pass exhausts units early and hands to scalar. for &rk in &[1usize, 3, 7] { - let b = Array2::from_shape_fn((n, rk), |_| rng.below(q)); - let c = Array2::from_shape_fn((rk, n), |_| rng.below(q)); + let b = Array2::from_shape_fn((n, rk), |_| rng.gen_range(0..q)); + let c = Array2::from_shape_fn((rk, n), |_| rng.gen_range(0..q)); let a = matmul_mod(&b, &c, q); assert_valid_snf(&a, p, e); checked += 1; } // (3) p | A everywhere: min valuation >= 1, blocked pass skipped. - let a = Array2::from_shape_fn((n, n), |_| (rng.below(q / p.max(1)) * p) % q); + let a = Array2::from_shape_fn((n, n), |_| (rng.gen_range(0..q / p.max(1)) * p) % q); assert_valid_snf(&a, p, e); checked += 1; // (4) mixed: a unit block plus a p-scaled block (multi-level). let a = Array2::from_shape_fn((n, n), |(i, j)| { - if (i + j) % 3 == 0 { rng.below(q) } else { (rng.below(q) * p) % q } + if (i + j) % 3 == 0 { + rng.gen_range(0..q) + } else { + (rng.gen_range(0..q) * p) % q + } }); assert_valid_snf(&a, p, e); checked += 1; // (5) rectangular, both orientations. - let a = Array2::from_shape_fn((n, n / 2 + 1), |_| rng.below(q)); + let a = Array2::from_shape_fn((n, n / 2 + 1), |_| rng.gen_range(0..q)); assert_valid_snf(&a, p, e); - let a = Array2::from_shape_fn((n / 2 + 1, n), |_| rng.below(q)); + let a = Array2::from_shape_fn((n / 2 + 1, n), |_| rng.gen_range(0..q)); assert_valid_snf(&a, p, e); checked += 2; } diff --git a/crates/modularsnf/src/ring.rs b/crates/modularsnf/src/ring.rs index 89c3711..000bdcf 100644 --- a/crates/modularsnf/src/ring.rs +++ b/crates/modularsnf/src/ring.rs @@ -129,10 +129,7 @@ impl RingZModN { let b_val = posmod(b, self.n); let (g, x, _) = egcd(b_val, self.n); if g == 0 || a_val % g != 0 { - return Err(format!( - "{a} not divisible by {b} in Z/{}", - self.n - )); + return Err(format!("{a} not divisible by {b} in Z/{}", self.n)); } Ok(posmod_i128((x as i128) * ((a_val / g) as i128), self.n / g)) } diff --git a/crates/modularsnf/src/snf.rs b/crates/modularsnf/src/snf.rs index 834d222..9ed8f05 100644 --- a/crates/modularsnf/src/snf.rs +++ b/crates/modularsnf/src/snf.rs @@ -15,10 +15,7 @@ fn posmod(a: i64, n: i64) -> i64 { /// Top-level SNF for a square matrix. /// Returns (U, V, S) such that S = U @ A @ V. -pub fn smith_square( - a: &Array2, - ring: &RingZModN, -) -> (Array2, Array2, Array2) { +pub fn smith_square(a: &Array2, ring: &RingZModN) -> (Array2, Array2, Array2) { let n_mod = ring.n(); let n = a.nrows(); diff --git a/docs/PLAN.md b/docs/PLAN.md deleted file mode 100644 index 3de2678..0000000 --- a/docs/PLAN.md +++ /dev/null @@ -1,133 +0,0 @@ -# PLAN: Optimize the CRT-based Smith Normal Form path - -## Goal - -Turn the CRT SNF path (`crates/modularsnf/src/crt.rs`, `rust_crt_snf`) into a -fully-optimized, large-`n`, multi-core implementation for the **small-modulus -(N ≤ a few hundred), large-matrix (n in the thousands)** regime. - -Baseline established (head-to-head, rust-vs-rust, branch `bench/headtohead`): -the CRT path already beats the recursive Storjohann baseline 4.4–35×. This plan -attacks its remaining bottlenecks: i128 in the hot loop, scalar elimination, -and single-threaded execution. - -## Why this design is correct (research-validated) - -- Over a chain/valuation ring `Z/p^e`, the SNF conditioner can be a permutation - matrix — valuation-pivoted Gaussian elimination, **no gcdex/Bézout** - (Storjohann diss. Prop 9.21; ESA). This is exactly what `local_snf` does. -- Local-then-recombine is independently validated (Wilkening–Yu). -- SNF/Howell over `Z/N` reduces to **matrix multiplication** — so "make it fast" - means "turn the trailing update into GEMM." - -## Library strategy (call-into vs hand-roll) - -| Component | Verdict | Rationale | -|---|---|---| -| Local elimination (Z/p^e) | **Hand-roll (Rust)** | Simple chain-ring pivot; no lib does Smith-with-transforms over Z/p^e; license-clean. | -| Modular GEMM (trailing update) | **Hand-roll thin layer + BLAS/faer** | Accumulate f64/i64, reduce mod p^e; FFLAS does this but is LGPL + field-only. | -| CRT recombination | **Hand-roll (Rust)** | Already done; O(n²·np), not the bottleneck. | -| External SNF lib | **Reference only** | FLINT `nmod_mat_howell_form` = Howell≠Smith, single-limb, **LGPL**. PARI = **GPL → incompatible** with Apache-2.0. | - -Decisive constraint: LGPL relinking is fragile in a static manylinux/abi3 wheel; -GPL is out. Hand-roll; use FFLAS-FFPACK / FLINT as algorithmic references. -Prefer **faer** (MIT/Apache) if a Rust GEMM backend is needed. - -Key arithmetic fact for our regime: exact integer GEMM stays exact while -`λ·(p−1)² < 2^53`. For p^e ≈ 256 that's λ ≈ 1.4e8 accumulations before a -reduction is needed — so reductions are essentially free and BLAS `dgemm` / -`faer` can be the inner kernel. (Use only this single-level bound; the nested -Strassen-Winograd `kWinograd` bound did not survive verification.) - -## Roadmap (phased, each verified against the head-to-head harness) - -### Phase 0 — Characterize *(in progress)* -- Env-gated timing in `crt_snf`: measure `local_snf` total vs recombination vs - unit-normalization at large n (500, 1000, 2000) and representative moduli. -- Confirm the `local_snf` i128 inner loop dominates. Record the split. - -### Phase 1 — Kill i128 in the hot loop (single-core constant win) *(next)* -- Add an i64 fast-path reduction (valid when `p^e < 2^31`, i128 fallback else), - mirroring `optimize-modulus`'s `mul_mod`. -- Apply to `local_snf` row/column clears + pivot normalization (`mulmod`). -- Verify: invariant factors + `U·A·V == S (mod N)` unchanged; measure speedup. - -### Phase 2 — Blocked right-looking LU + modular GEMM (asymptotic win) - -**2a — two-phase split (DONE).** `local_snf` now does Phase L (minimal-valuation -column elimination → U, upper-triangular `mat`) + Phase R (reverse-order back- -elimination → V). Groundwork for blocking; also removed wasted work (old inline -row-clear touched all n rows of each column). ~1.1–1.3× over Phase 1. Gated by a -randomized cargo oracle test (`local_snf_is_valid_smith_form`, 500+ cases: -`U·A·V==diag(p^vals)`, ascending vals, `U,V` unimodular) + head-to-head vs dev. - -**2b — blocked GEMM trailing update (DONE).** Implemented as right-looking -blocked LU over the valuation-0 bulk (panel B=48): in-place panel factorization, -TRSM for the deferred panel rows, GEMM trailing update (i64 accumulate + delayed -reduction) on `mat` and `u`; scalar min-valuation path finishes the higher- -valuation/rank-deficient tail; final diagonal normalization. ~1.3–1.5× over 2a -at n=400, 1.6–2.6× over Phase 1 at n=1000–2000 (cumulative ~3× over original -CRT). Gated by two cargo tests (600+ cases incl. multi-panel, rank-deficiency -mod p, p|A, rectangular) + head-to-head vs dev. -**2c — packed cache-friendly GEMM micro-kernel (DONE).** The 2b GEMM read panel -rows with stride-m across the contraction index (~pb cache lines per output). -Fix: pack U12 into a contiguous (pb×ncols) buffer, then contiguous-axpy each -trailing row (autovectorizes); i128 fallback for large q. ~1.10–1.34× over 2b -across single- and multi-prime moduli. (An f64 ndarray/matrixmultiply GEMM was -tried first but regressed multi-prime — n*n product allocation + thin K=48 — so -the packed i64 kernel was kept.) Cumulative ~3.5–4× over original CRT. -NEXT: a true tiled/register-blocked or BLAS/faer kernel could push large-n -further (n=2000 still ~1.1× since the micro-kernel isn't register-tiled), but -diminishing returns vs Phase 3 (rayon). - -Original design notes (for reference): -Key constraint: -minimal-valuation pivoting needs the *global* min over the trailing block, which -defeats naive column-panel/lazy blocking. The blockable formulation is -**valuation-level staging**: process pivots by p-adic level ℓ = 0..e; within a -level every pivot is a unit mod p, so it degenerates to standard blocked LU. -Implementation plan for Phase L: -- Outer loop: `lev` = current global min valuation in `mat[k.., k..]`; `pl=p^lev`. -- Panel (width B≈48) of valuation-`lev` pivots. Bookkeeping per standard - right-looking blocked LU: - - factor panel columns `[pstart,pend)`: full updates to **panel rows** (so the - `U12` block `mat[pstart..pend, pend..]` is correct) and to **panel columns** - for all rows (so pivot search + the `L21` multipliers `mat[pend.., pstart..pend]` - are correct); defer non-panel trailing columns. - - **GEMM trailing update**: `mat[pend.., pend..] -= L21 @ U12 (mod p^e)`, and - the analogous `u[pend.., :] -= L21 @ u[pstart..pend, :]`. - - column-deflation: a panel column with no valuation-`lev` unit → swap out to - the active-region tail (track in V), revisit at a higher level. -- GEMM kernel: i64 accumulate with delayed reduction (small `p^e` ⇒ block - `λ·(p−1)² < 2^53` is huge ⇒ reduce rarely); i128 acc fallback for large `p^e`. - Later swap inner kernel for `faer` (MIT/Apache) or BLAS dgemm. -- RISK: exact-LU bookkeeping (L21/U12 split, panel-row vs below-panel updates, - deflation) is subtle; keep the scalar 2a `local_snf` as the cargo-test oracle - and only switch the default once the blocked version passes the full suite. -- Apply the same blocking to Phase R's V-update (triangular column reduction → - one TRMM/GEMM) once Phase L lands. - -### Phase 3 — Multi-core (rayon) -- Parallelize the trailing-update GEMM first (real scaling). -- Add across-primes parallelism (cheap, low ceiling: np = 1–3). -- **Re-characterize scaling empirically** — published speedups (e.g. 13.6×/32 - cores, FFLAS Euro-Par 2014) are over prime *fields* Z/pZ, not Z/p^e. - -### Phase 4 — GPU (deferred, literature-gated) -- No reusable GPU kernel for modular LU/SNF over Z/p(^e) was found. Revisit only - after rayon scaling is measured. - -## Open questions / risks -- Field-vs-chain-ring: parallel/GEMM benchmark evidence is Z/pZ; Z/p^e parallel - behavior is unmeasured — expect to measure it ourselves. -- Integer-vs-polynomial: Storjohann/Wilkening-Yu state results over K[x]/(p^k); - Z/p^e transfer is structurally sound extrapolation. -- Evaluate `faer` early in Phase 2 as the license-clean GEMM backend. - -## References -- Storjohann, dissertation §9.6 (Prop 9.21, 9.23): cs.uwaterloo.ca/~astorjoh/diss2up.pdf -- Storjohann, ESA (Howell/Smith reduce to matmul): cs.uwaterloo.ca/~astorjoh/esa.pdf -- Dumas/Gautier/Pernet/Sultan, parallel echelon forms (Euro-Par 2014): arxiv.org/abs/1402.3501 -- Dumas/Giorgi/Pernet, FFLAS-FFPACK (exact GEMM, delayed reduction): hal.science/hal-00018223v2 -- Wilkening–Yu, local Smith form: math.berkeley.edu/~wilken/papers/smith.pdf -- FLINT nmod_mat (Howell form over Z/N): flintlib.org/doc/nmod_mat.html diff --git a/docs/crt.md b/docs/crt.md new file mode 100644 index 0000000..8f4e399 --- /dev/null +++ b/docs/crt.md @@ -0,0 +1,69 @@ +# CRT fast path (experimental) + +An alternative SNF algorithm to the Storjohann band reduction +([`docs/algorithm.md`](algorithm.md)), specialized for the **small-modulus +(N up to a few hundred), large-matrix (n in the thousands)** regime. On that +class of inputs it is several times faster than the band-reduction path; for +general moduli, or when the prime factorization of N is unknown or expensive, +prefer the Storjohann path. + +Implemented in `crates/modularsnf/src/crt.rs` (`crt_snf`) and exposed to Python +as `modularsnf.crt_snf`. + +## Idea + +Factor `N = prod p^e` (the factorization is passed in, so it is amortized +across many calls with the same modulus). The SNF over `Z/NZ` decomposes by the +Chinese Remainder Theorem into independent SNFs over each local ring `Z/p^e`: + +1. **Local SNF.** Over a chain ring `Z/p^e` the conditioner can be a + permutation — valuation-pivoted Gaussian elimination, with no `gcdex`/Bézout + step. This yields the divisibility chain for free and produces the unimodular + transforms `U_p, V_p` natively. +2. **Normalize.** Unit-normalize each prime's `V_p` so every prime realizes the + same global invariant factors `d_i = prod_p p^{v_p(i)}`. +3. **Recombine.** CRT-recombine the per-prime `U_p, V_p` entrywise into + `U, V` over `Z/NZ`. + +The result satisfies `S = U A V (mod N)` with `S` diagonal and `U, V` +unimodular, identical in form to the band-reduction path. + +## Local SNF internals + +`local_snf` runs in two phases: + +- **Phase L** — column elimination producing `U` and leaving `mat` upper + triangular. A blocked right-looking LU clears the valuation-0 bulk: a panel is + factored in place, then the trailing block is updated by TRSM + GEMM (i64 + accumulate with delayed reduction; i128 fallback for large `p^e`). A scalar + minimal-valuation path finishes the higher-valuation and rank-deficient tail. +- **Phase R** — back-elimination of the super-diagonal producing `V`, + processing pivots in reverse so fill lands in not-yet-processed rows. + +The blocking exploits the fact that SNF/Howell over `Z/N` reduces to matrix +multiplication, so the trailing update is a modular GEMM. Exact integer GEMM +stays exact while `lambda * (p-1)^2` fits the accumulator, which for small `p^e` +means reductions are rare. + +## Correctness basis + +- Permutation conditioner over chain rings: Storjohann, dissertation §9.6 + (Prop 9.21, 9.23). +- SNF/Howell reduces to matrix multiplication: Storjohann, ESA. +- Local-then-recombine: Wilkening–Yu, *local Smith form*. + +The scalar `local_snf` is the oracle for the blocked path: a randomized Cargo +test suite (`local_snf_is_valid_smith_form`, `local_snf_blocked_paths`) checks +`U A V == diag(p^vals)`, ascending valuations, and unimodularity of `U, V` +across single/multi-panel, rank-deficient, `p | A`, and rectangular cases. + +## References + +- Storjohann, *Algorithms for Matrix Canonical Forms* (diss., §9.6): + cs.uwaterloo.ca/~astorjoh/diss2up.pdf +- Storjohann, *Howell/Smith reduce to matmul* (ESA): + cs.uwaterloo.ca/~astorjoh/esa.pdf +- Wilkening, Yu, *local Smith form*: + math.berkeley.edu/~wilken/papers/smith.pdf +- Dumas, Giorgi, Pernet, *FFLAS-FFPACK* (exact GEMM, delayed reduction): + hal.science/hal-00018223v2 diff --git a/scripts/crt_snf_profile.py b/scripts/crt_snf_profile.py deleted file mode 100644 index 99a59bb..0000000 --- a/scripts/crt_snf_profile.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Profile CRT-style SNF prototype vs the existing implementation. - -Compares three code paths at increasing n for a composite modulus: - * existing pipeline, Rust backend (the current production fast path) - * existing pipeline, pure-Python backend - * CRT prototype (pure-Python, object dtype) -""" - -from __future__ import annotations - -import time - -import numpy as np -from crt_snf_prototype import crt_snf - -import modularsnf.diagonal as diagonal_mod -import modularsnf.ring as ring_mod -import modularsnf.snf as snf_mod -from modularsnf import smith_normal_form_mod - -# Capture real Rust hooks so we can toggle the backend. -_R = { - "ring": ring_mod._RustRing, - "diag": diagonal_mod._rust_diag, - "merge": diagonal_mod._rust_merge, - "snf": snf_mod._rust_snf, -} - - -def set_backend(rust: bool) -> None: - ring_mod._RustRing = _R["ring"] if rust else None - diagonal_mod._rust_diag = _R["diag"] if rust else None - diagonal_mod._rust_merge = _R["merge"] if rust else None - snf_mod._rust_snf = _R["snf"] if rust else None - - -def timeit(fn, *args, repeat: int = 3) -> float: - best = float("inf") - for _ in range(repeat): - t0 = time.perf_counter() - fn(*args) - best = min(best, time.perf_counter() - t0) - return best - - -def main() -> None: - N = 720 # 2^4 * 3^2 * 5 — composite with zero divisors - sizes = [4, 8, 16, 24, 32, 48, 64] - rng = np.random.default_rng(0) - - print(f"modulus N = {N} (factor = 2^4 * 3^2 * 5)") - print(f"{'n':>4} | {'rust (ms)':>12} | {'python (ms)':>12} | {'crt (ms)':>12}") - print("-" * 52) - - for n in sizes: - A = rng.integers(0, N, size=(n, n)) - Al = A.tolist() - - set_backend(True) - t_rust = timeit(lambda: smith_normal_form_mod(Al, modulus=N)) - - set_backend(False) - t_py = timeit(lambda: smith_normal_form_mod(Al, modulus=N), repeat=1) - - t_crt = timeit(lambda: crt_snf(A, N), repeat=1) - - print( - f"{n:>4} | {t_rust*1e3:>12.2f} | {t_py*1e3:>12.2f} | {t_crt*1e3:>12.2f}" - ) - - set_backend(True) - - -if __name__ == "__main__": - main() diff --git a/scripts/crt_snf_prototype.py b/scripts/crt_snf_prototype.py deleted file mode 100644 index 9a4011e..0000000 --- a/scripts/crt_snf_prototype.py +++ /dev/null @@ -1,207 +0,0 @@ -"""Prototype: CRT-style Smith Normal Form over Z/NZ with transforms. - -Non-recursive. Factor N = prod p^e, compute SNF over each local ring Z/p^e -by valuation-pivoted Gaussian elimination (which yields the divisibility -chain for free and unimodular transforms natively), then CRT-recombine the -per-prime transforms into U, V over Z/NZ. - -This is a correctness/feasibility prototype in pure Python (object dtype to -avoid overflow), validated against the existing implementation as oracle. -""" - -from __future__ import annotations - -from math import gcd - -import numpy as np - - -def factorize(n: int) -> dict[int, int]: - factors: dict[int, int] = {} - d = 2 - while d * d <= n: - while n % d == 0: - factors[d] = factors.get(d, 0) + 1 - n //= d - d += 1 - if n > 1: - factors[n] = factors.get(n, 0) + 1 - return factors - - -def pval(x: int, p: int, cap: int) -> int: - """p-adic valuation of x, capped at cap; x == 0 -> cap.""" - x %= p**cap - if x == 0: - return cap - v = 0 - while v < cap and x % p == 0: - x //= p - v += 1 - return v - - -def egcd(a: int, b: int) -> tuple[int, int, int]: - if b == 0: - return (a, 1, 0) - g, x, y = egcd(b, a % b) - return (g, y, x - (a // b) * y) - - -def inv_mod(a: int, m: int) -> int: - g, x, _ = egcd(a % m, m) - assert g == 1, f"{a} not invertible mod {m}" - return x % m - - -def local_snf(A: np.ndarray, p: int, e: int): - """SNF of A over the local ring Z/p^e via valuation pivoting. - - Returns (U, V, diag_vals) with U @ A @ V == diag(p^vals) mod p^e, - U (n x n) and V (m x m) unimodular, vals ascending (divisibility chain). - """ - q = p**e - M = (A.astype(object)) % q - n, m = M.shape - U = np.eye(n, dtype=object) - V = np.eye(m, dtype=object) - r = min(n, m) - - for k in range(r): - # Find pivot in M[k:, k:] with minimal p-adic valuation. - best = None - bestval = e + 1 - for i in range(k, n): - for j in range(k, m): - if int(M[i, j]) % q != 0: - v = pval(int(M[i, j]), p, e) - if v < bestval: - bestval, best = v, (i, j) - if v == 0: - break - if bestval == 0: - break - if best is None: - break # entire trailing block is zero - - pi, pj = best - if pi != k: - M[[k, pi], :] = M[[pi, k], :] - U[[k, pi], :] = U[[pi, k], :] - if pj != k: - M[:, [k, pj]] = M[:, [pj, k]] - V[:, [k, pj]] = V[:, [pj, k]] - - v = bestval - pv = p**v - # Normalize pivot to exactly p^v by scaling row k by the unit inverse. - pivot = int(M[k, k]) % q - unit = pivot // pv # coprime to p - uinv = inv_mod(unit, q) - M[k, :] = (M[k, :] * uinv) % q - U[k, :] = (U[k, :] * uinv) % q - - # Clear column k below the pivot. - for i in range(k + 1, n): - if int(M[i, k]) % q != 0: - c = (int(M[i, k]) % q) // pv # exact: val(M[i,k]) >= v - M[i, :] = (M[i, :] - c * M[k, :]) % q - U[i, :] = (U[i, :] - c * U[k, :]) % q - - # Clear row k to the right of the pivot. - for j in range(k + 1, m): - if int(M[k, j]) % q != 0: - c = (int(M[k, j]) % q) // pv - M[:, j] = (M[:, j] - c * M[:, k]) % q - V[:, j] = (V[:, j] - c * V[:, k]) % q - - vals = [pval(int(M[i, i]), p, e) for i in range(r)] - return U, V, vals - - -def crt_combine(residues: list[int], moduli: list[int]) -> int: - """Combine residues mod pairwise-coprime moduli into a value mod prod.""" - x = 0 - M = 1 - for r, m in zip(residues, moduli): - # solve x ≡ x (mod M), x ≡ r (mod m) - g, s, _ = egcd(M, m) - assert g == 1 - x = (x + M * ((r - x) * s % m)) % (M * m) - M *= m - return x - - -def crt_snf(A: np.ndarray, N: int): - """CRT-style SNF over Z/NZ. Returns (U, V, S) with S = U @ A @ V mod N.""" - A = A.astype(object) % N - n, m = A.shape - r = min(n, m) - fac = factorize(N) - primes = sorted(fac) - qs = [p ** fac[p] for p in primes] - - per_prime = {} - vals = {} - for p in primes: - e = fac[p] - Up, Vp, vp = local_snf(A, p, e) - per_prime[p] = (Up, Vp) - vals[p] = vp - - # Global invariant factors d_i = prod_p p^{vals_p[i]} (divides N). - d = [] - for i in range(r): - di = 1 - for p in primes: - di *= p ** vals[p][i] - d.append(di % N) - - # Unit-normalize each prime's V so all primes realize the SAME global d_i. - for p in primes: - e = fac[p] - q = p**e - Vp = per_prime[p][1] - for i in range(r): - # w = prod_{p'!=p} p'^{vals_{p'}[i]} (a unit mod q) - w = 1 - for p2 in primes: - if p2 != p: - w *= p2 ** vals[p2][i] - w %= q - Vp[:, i] = (Vp[:, i] * w) % q - - # CRT-recombine U and V entrywise across primes. - U = np.zeros((n, n), dtype=object) - V = np.zeros((m, m), dtype=object) - for a in range(n): - for b in range(n): - U[a, b] = crt_combine( - [int(per_prime[p][0][a, b]) for p in primes], qs - ) - for a in range(m): - for b in range(m): - V[a, b] = crt_combine( - [int(per_prime[p][1][a, b]) for p in primes], qs - ) - - S = np.zeros((n, m), dtype=object) - for i in range(r): - S[i, i] = d[i] % N - return U, V, S - - -def det_mod(M: np.ndarray, N: int) -> int: - """Determinant of integer matrix mod N via fraction-free (Bareiss-ish).""" - M = M.astype(object) % N - n = M.shape[0] - if n == 0: - return 1 % N - # Plain cofactor for small n is fine for tests; use integer det via numpy on python ints. - from sympy import Matrix - - return int(Matrix(M.tolist()).det()) % N - - -def is_unimodular(M: np.ndarray, N: int) -> bool: - return gcd(det_mod(M, N) % N, N) == 1 diff --git a/scripts/crt_snf_verify.py b/scripts/crt_snf_verify.py deleted file mode 100644 index 328f791..0000000 --- a/scripts/crt_snf_verify.py +++ /dev/null @@ -1,94 +0,0 @@ -"""Verify the CRT-style SNF prototype against the existing implementation.""" - -from __future__ import annotations - -import random -from math import gcd - -import numpy as np -from crt_snf_prototype import crt_snf, is_unimodular - -from modularsnf import smith_normal_form_mod - - -def check_one(A: np.ndarray, N: int) -> tuple[bool, str]: - n, m = A.shape - U, V, S = crt_snf(A, N) - - # 1. Self-certifying factorization equation. - prod = (U.astype(object) @ A.astype(object) @ V.astype(object)) % N - if not np.array_equal(prod % N, S % N): - return False, "U@A@V != S mod N" - - # 2. S diagonal. - for i in range(n): - for j in range(m): - if i != j and int(S[i, j]) % N != 0: - return False, "S not diagonal" - - # 3. Divisibility chain d_i | d_{i+1} (as ideals: gcd(d_i,N) | gcd(d_{i+1},N)). - r = min(n, m) - g = [gcd(int(S[i, i]) % N, N) for i in range(r)] - for i in range(r - 1): - if g[i + 1] % g[i] != 0: - return False, f"divisibility chain broken: {g}" - - # 4. Unimodularity of U and V. - if not is_unimodular(U, N): - return False, "U not unimodular" - if not is_unimodular(V, N): - return False, "V not unimodular" - - # 5. Diagonal matches the oracle (SNF diagonal is unique up to gcd-with-N). - So, _, _ = smith_normal_form_mod(A.tolist(), modulus=N) - So = np.array(So, dtype=object) - g_oracle = [gcd(int(So[i, i]) % N, N) for i in range(r)] - if g != g_oracle: - return False, f"diagonal mismatch: mine={g} oracle={g_oracle}" - - return True, "ok" - - -def main() -> None: - random.seed(12345) - np.random.seed(12345) - - moduli = [2, 3, 4, 6, 8, 9, 12, 16, 30, 36, 60, 36, 100, 210, 720, 1024] - sizes = [1, 2, 3, 4, 5, 6] - - total = 0 - failures = 0 - for N in moduli: - for n in sizes: - for _ in range(30): - m = n # square; rectangular handled by same padding upstream - A = np.random.randint(0, N, size=(n, m)) - total += 1 - ok, msg = check_one(A, N) - if not ok: - failures += 1 - if failures <= 10: - print(f"FAIL N={N} shape={A.shape}: {msg}") - print(A) - - # Edge cases. - edge = [ - (np.zeros((3, 3), dtype=int), 12), - (np.eye(3, dtype=int), 12), - (np.array([[2, 4, 0], [6, 8, 3], [0, 3, 9]]), 36), - (np.array([[6]]), 12), - (np.full((4, 4), 6), 36), - (np.array([[4, 0], [0, 9]]), 36), - ] - for A, N in edge: - total += 1 - ok, msg = check_one(A, N) - if not ok: - failures += 1 - print(f"FAIL edge N={N} shape={A.shape}: {msg}\n{A}") - - print(f"\n{total - failures}/{total} passed, {failures} failures") - - -if __name__ == "__main__": - main() diff --git a/scripts/profile_snf.py b/scripts/profile_snf.py deleted file mode 100755 index 57db660..0000000 --- a/scripts/profile_snf.py +++ /dev/null @@ -1,502 +0,0 @@ -#!/usr/bin/env python3 -"""Benchmark and profile the modular SNF pipeline. - -The default mode reports wall-clock time, a rough ``n^3 / time`` throughput, -and peak Python memory for a sequence of square matrix sizes. - -Backend selection is explicit: -- ``--backend rust`` benchmarks the native PyO3 fast path. -- ``--backend python`` benchmarks the pure-Python implementation. -- ``--backend auto`` prefers Rust when available. - -For deeper profiling: -- ``--cprofile`` records a Python profiler snapshot for the first size. -- ``--flamegraph`` runs ``py-spy --native`` on a child process so Rust frames - remain visible. -""" - -from __future__ import annotations - -import argparse -import cProfile -import random -import shutil -import subprocess -import sys -import tempfile -import time -import tracemalloc -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import numpy as np - -REPO_ROOT = Path(__file__).resolve().parent.parent - - -@dataclass -class Modules: - """Loaded modularsnf modules and optional Rust entry points.""" - - ring_mod: Any - diagonal_mod: Any - snf_mod: Any - ring_matrix_cls: Any - ring_cls: Any - smith_normal_form: Any - rust_ring_cls: Any - rust_diag_fn: Any - rust_merge_fn: Any - rust_snf_fn: Any - - -@dataclass -class BenchmarkResult: - """Summary for one benchmarked matrix size.""" - - n: int - elapsed: float - rough_n3_per_sec: float - peak_memory_bytes: int - correct: bool - - -def bootstrap_repo() -> None: - """Ensure the repository root is importable.""" - repo_str = str(REPO_ROOT) - if repo_str not in sys.path: - sys.path.insert(0, repo_str) - - -def load_modules() -> Modules: - """Import modularsnf lazily after bootstrapping the repo path.""" - bootstrap_repo() - - import modularsnf.diagonal as diagonal_mod - import modularsnf.ring as ring_mod - import modularsnf.snf as snf_mod - from modularsnf.matrix import RingMatrix - from modularsnf.ring import RingZModN - from modularsnf.snf import smith_normal_form - - rust_ring_cls: Any = None - rust_diag_fn: Any = None - rust_merge_fn: Any = None - rust_snf_fn: Any = None - - try: - from modularsnf._rust import RustRingZModN as rust_ring_cls - from modularsnf._rust import rust_merge_smith_blocks as rust_merge_fn - from modularsnf._rust import rust_smith_from_diagonal as rust_diag_fn - from modularsnf._rust import rust_smith_normal_form as rust_snf_fn - except ImportError: - pass - - return Modules( - ring_mod=ring_mod, - diagonal_mod=diagonal_mod, - snf_mod=snf_mod, - ring_matrix_cls=RingMatrix, - ring_cls=RingZModN, - smith_normal_form=smith_normal_form, - rust_ring_cls=rust_ring_cls, - rust_diag_fn=rust_diag_fn, - rust_merge_fn=rust_merge_fn, - rust_snf_fn=rust_snf_fn, - ) - - -def rust_available(modules: Modules) -> bool: - """Return whether the native extension is fully importable.""" - return all( - value is not None - for value in ( - modules.rust_ring_cls, - modules.rust_diag_fn, - modules.rust_merge_fn, - modules.rust_snf_fn, - ) - ) - - -def configure_backend(modules: Modules, requested: str) -> str: - """Select and activate the requested backend.""" - has_rust = rust_available(modules) - - if requested == "rust": - if not has_rust: - raise RuntimeError( - "Rust backend requested but modularsnf._rust is unavailable" - ) - backend = "rust" - elif requested == "python": - backend = "python" - elif has_rust: - backend = "rust" - else: - backend = "python" - - if backend == "rust": - modules.ring_mod._RustRing = modules.rust_ring_cls - modules.diagonal_mod._rust_diag = modules.rust_diag_fn - modules.diagonal_mod._rust_merge = modules.rust_merge_fn - modules.snf_mod._rust_snf = modules.rust_snf_fn - else: - modules.ring_mod._RustRing = None - modules.diagonal_mod._rust_diag = None - modules.diagonal_mod._rust_merge = None - modules.snf_mod._rust_snf = None - - return backend - - -def make_random_matrix( - modules: Modules, - n: int, - modulus: int, - seed: int, -) -> Any: - """Create a deterministic random square matrix over Z/modulus.""" - rng = random.Random(seed) - ring = modules.ring_cls(modulus) - data = [[rng.randint(0, modulus - 1) for _ in range(n)] for _ in range(n)] - return modules.ring_matrix_cls.from_rows(ring, data) - - -def check_snf_correctness(A: Any, U: Any, V: Any, S: Any) -> bool: - """Verify transform validity, diagonal form, and divisibility chain.""" - modulus = A.ring.N - diag_len = min(S.nrows, S.ncols) - - if not np.array_equal((U @ A @ V).data, S.data): - return False - - for i in range(S.nrows): - for j in range(S.ncols): - if i != j and S.data[i, j] % modulus != 0: - return False - - for i in range(diag_len - 1): - cur = int(S.data[i, i]) % modulus - nxt = int(S.data[i + 1, i + 1]) % modulus - g = np.gcd(cur, modulus) - if g != 0 and nxt % g != 0: - return False - - return True - - -def run_once( - modules: Modules, - n: int, - modulus: int, - seed: int, - verify: bool, -) -> BenchmarkResult: - """Benchmark one square matrix size.""" - matrix = make_random_matrix(modules, n, modulus, seed) - - tracemalloc.start() - t0 = time.perf_counter() - U, V, S = modules.smith_normal_form(matrix) - elapsed = time.perf_counter() - t0 - _, peak = tracemalloc.get_traced_memory() - tracemalloc.stop() - - correct = check_snf_correctness(matrix, U, V, S) if verify else True - rough_n3_per_sec = (n**3 / elapsed) if elapsed > 0 else 0.0 - return BenchmarkResult( - n=n, - elapsed=elapsed, - rough_n3_per_sec=rough_n3_per_sec, - peak_memory_bytes=peak, - correct=correct, - ) - - -def fmt_time(seconds: float) -> str: - """Format seconds into a compact human-readable string.""" - if seconds < 0.001: - return f"{seconds * 1_000_000:.0f}us" - if seconds < 1.0: - return f"{seconds * 1_000:.1f}ms" - return f"{seconds:.3f}s" - - -def fmt_bytes(num_bytes: int) -> str: - """Format bytes into a compact human-readable string.""" - if num_bytes < 1024: - return f"{num_bytes}B" - if num_bytes < 1024**2: - return f"{num_bytes / 1024:.1f}KB" - return f"{num_bytes / 1024**2:.1f}MB" - - -def fmt_rate(rate: float) -> str: - """Format the rough ``n^3 / time`` throughput metric.""" - if rate >= 1_000_000_000: - return f"{rate / 1_000_000_000:.2f}G" - if rate >= 1_000_000: - return f"{rate / 1_000_000:.2f}M" - if rate >= 1_000: - return f"{rate / 1_000:.2f}K" - return f"{rate:.0f}" - - -def print_result(result: BenchmarkResult) -> None: - """Print the outcome for one matrix size.""" - status = "YES" if result.correct else "NO" - print( - f" n={result.n:<5d} time={fmt_time(result.elapsed):>10s} " - f"rough n^3/s={fmt_rate(result.rough_n3_per_sec):>8s} " - f"peak_mem={fmt_bytes(result.peak_memory_bytes):>8s} " - f"correct={status}" - ) - - -def print_scaling_summary(results: list[BenchmarkResult]) -> None: - """Print a compact scaling table and fitted exponent.""" - if not results: - return - - print(f"\n{'=' * 72}") - print(" SCALING SUMMARY") - print(f"{'=' * 72}") - print( - f" {'n':>6s} {'time':>10s} {'rough n^3/s':>12s} " - f"{'peak_mem':>10s} {'correct':>7s}" - ) - print(f" {'─' * 6} {'─' * 10} {'─' * 12} {'─' * 10} {'─' * 7}") - - for result in results: - print( - f" {result.n:>6d} {fmt_time(result.elapsed):>10s} " - f"{fmt_rate(result.rough_n3_per_sec):>12s} " - f"{fmt_bytes(result.peak_memory_bytes):>10s} " - f"{'YES' if result.correct else 'NO':>7s}" - ) - - if len(results) >= 2: - log_n = np.log([result.n for result in results]) - log_t = np.log([result.elapsed for result in results]) - coeffs = np.polyfit(log_n, log_t, 1) - print(f"\n Empirical scaling exponent: n^{coeffs[0]:.2f}") - - print() - - -def run_cprofile( - modules: Modules, - backend: str, - n: int, - modulus: int, - seed: int, -) -> None: - """Record a Python cProfile snapshot for the first size.""" - print(f"\nRunning cProfile on n={n}, modulus={modulus}...") - if backend == "rust": - print( - "Note: cProfile will mostly see the Python wrapper around native code." - ) - - matrix = make_random_matrix(modules, n, modulus, seed) - profiler = cProfile.Profile() - profiler.enable() - modules.smith_normal_form(matrix) - profiler.disable() - - outfile = f"profile_n{n}_mod{modulus}_{backend}.prof" - profiler.dump_stats(outfile) - print(f"Saved to {outfile}") - print(f"Explore with: python -m pstats {outfile}") - print(f" or: pip install snakeviz && snakeviz {outfile}") - - -def run_flamegraph( - backend: str, - n: int, - modulus: int, - seed: int, -) -> None: - """Generate a native flamegraph with py-spy.""" - py_spy = shutil.which("py-spy") - if py_spy is None: - venv_spy = Path(sys.prefix) / "bin" / "py-spy" - if venv_spy.exists(): - py_spy = str(venv_spy) - if py_spy is None: - raise RuntimeError("py-spy not found. Install with: pip install py-spy") - - child_script = f""" -import random -import sys -import time - -repo_root = {str(REPO_ROOT)!r} -if repo_root not in sys.path: - sys.path.insert(0, repo_root) - -import modularsnf.diagonal as diagonal_mod -import modularsnf.ring as ring_mod -import modularsnf.snf as snf_mod -from modularsnf.matrix import RingMatrix -from modularsnf.ring import RingZModN -from modularsnf.snf import smith_normal_form - -USE_RUST = {backend == "rust"!r} -if not USE_RUST: - ring_mod._RustRing = None - diagonal_mod._rust_diag = None - diagonal_mod._rust_merge = None - snf_mod._rust_snf = None - -rng = random.Random({seed}) -ring = RingZModN({modulus}) -data = [[rng.randint(0, {modulus - 1}) for _ in range({n})] for _ in range({n})] -A = RingMatrix.from_rows(ring, data) - -smith_normal_form(A) - -t0 = time.perf_counter() -smith_normal_form(A) -dt = time.perf_counter() - t0 -iters = max(1, int(5.0 / max(dt, 1e-6))) -print(f"Running {{iters}} iterations ({{dt*1000:.1f}}ms each)...", flush=True) -for _ in range(iters): - smith_normal_form(A) -""" - - with tempfile.NamedTemporaryFile( - mode="w", - suffix=".py", - delete=False, - ) as handle: - handle.write(child_script) - child_path = handle.name - - outfile = f"flamegraph_n{n}_mod{modulus}_{backend}.svg" - cmd = [ - py_spy, - "record", - "--native", - "--output", - outfile, - "--rate", - "197", - "--", - sys.executable, - child_path, - ] - - try: - print(f"\nGenerating flamegraph for n={n}, modulus={modulus}...") - print(f"Using py-spy at: {py_spy}") - subprocess.run(cmd, check=True, text=True) - finally: - Path(child_path).unlink(missing_ok=True) - - outpath = Path(outfile) - if not outpath.exists() or outpath.stat().st_size == 0: - raise RuntimeError("flamegraph was not generated") - print(f"Wrote flamegraph to {outfile}") - - -def parse_args() -> argparse.Namespace: - """Parse command-line arguments.""" - parser = argparse.ArgumentParser( - description="Benchmark and profile the modular SNF pipeline." - ) - parser.add_argument( - "--sizes", - default="10,20,50,100,200", - help="Comma-separated square matrix sizes", - ) - parser.add_argument( - "--modulus", - type=int, - default=12, - help="Ring modulus N (default: 12)", - ) - parser.add_argument( - "--seed", - type=int, - default=42, - help="Random seed (default: 42)", - ) - parser.add_argument( - "--backend", - choices=("auto", "python", "rust"), - default="auto", - help="Execution backend (default: auto)", - ) - parser.add_argument( - "--timeout", - type=float, - default=300.0, - help="Stop after a run exceeds this wall-clock time in seconds", - ) - parser.add_argument( - "--skip-check", - action="store_true", - help="Skip the structural SNF correctness check", - ) - parser.add_argument( - "--cprofile", - action="store_true", - help="Dump a cProfile .prof file for the first size", - ) - parser.add_argument( - "--flamegraph", - action="store_true", - help="Generate a py-spy flamegraph SVG for the first size", - ) - return parser.parse_args() - - -def main() -> None: - """CLI entry point.""" - args = parse_args() - sizes = [ - int(part.strip()) for part in args.sizes.split(",") if part.strip() - ] - - modules = load_modules() - backend = configure_backend(modules, args.backend) - - print( - f"Benchmarking SNF: sizes={sizes}, modulus={args.modulus}, " - f"seed={args.seed}, backend={backend}" - ) - if backend == "rust": - print("Note: detailed Python-level counters are intentionally omitted.") - print( - " Use --flamegraph for native profiling and --backend python" - ) - print(" if you specifically want the pure-Python path.") - - if args.cprofile: - run_cprofile(modules, backend, sizes[0], args.modulus, args.seed) - return - - if args.flamegraph: - run_flamegraph(backend, sizes[0], args.modulus, args.seed) - return - - verify = not args.skip_check - results: list[BenchmarkResult] = [] - for n in sizes: - result = run_once(modules, n, args.modulus, args.seed, verify) - print_result(result) - results.append(result) - if result.elapsed > args.timeout: - print( - f" Stopping after n={n}: exceeded timeout {args.timeout:.1f}s" - ) - break - - print_scaling_summary(results) - - -if __name__ == "__main__": - main() From 72f0aa008044ddbe2b50867834bd869fcfadd28b Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Sat, 6 Jun 2026 00:42:53 -0400 Subject: [PATCH 19/22] refactor: collapse the Python package into a thin Rust binding The pure-Python Storjohann implementation is removed; modularsnf is now a thin wrapper over the modularsnf._rust extension. Validation and marshalling live in _core.py; snf.py and crt.py just call the binding. The public API is a single entrance -- smith_normal_form_mod, crt_snf, SNFResult -- with RingZModN/RingMatrix/diagonal dropped. The PyO3 wrapper now delegates to modularsnf::smith_normal_form (which also propagates the tall-matrix fix to Python) and exposes only rust_smith_normal_form and rust_crt_snf; the dead RustRingZModN, diagonal, and merge exports are removed. --- crates/modularsnf-py/src/lib.rs | 188 +----------- modularsnf/__init__.py | 6 +- modularsnf/_core.py | 48 +++ modularsnf/_rust.pyi | 30 +- modularsnf/band.py | 272 ---------------- modularsnf/crt.py | 87 ++++++ modularsnf/diagonal.py | 376 ----------------------- modularsnf/echelon.py | 109 ------- modularsnf/matrix.py | 348 --------------------- modularsnf/ring.py | 177 ----------- modularsnf/snf.py | 528 ++------------------------------ 11 files changed, 177 insertions(+), 1992 deletions(-) create mode 100644 modularsnf/_core.py delete mode 100644 modularsnf/band.py create mode 100644 modularsnf/crt.py delete mode 100644 modularsnf/diagonal.py delete mode 100644 modularsnf/echelon.py delete mode 100644 modularsnf/matrix.py delete mode 100644 modularsnf/ring.py diff --git a/crates/modularsnf-py/src/lib.rs b/crates/modularsnf-py/src/lib.rs index fdbc485..f281352 100644 --- a/crates/modularsnf-py/src/lib.rs +++ b/crates/modularsnf-py/src/lib.rs @@ -1,172 +1,31 @@ -use modularsnf::ring::RingZModN; -use modularsnf::snf::smith_square; use modularsnf::crt::crt_snf; -use modularsnf::diagonal::{merge_raw, smith_from_diagonal}; -use numpy::ndarray::{s, Array2}; +use numpy::ndarray::Array2; use numpy::{IntoPyArray, PyArray2, PyReadonlyArray2}; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -/// Wrapper around RingZModN for Python. -#[pyclass] -#[pyo3(name = "RustRingZModN")] -struct PyRingZModN { - inner: RingZModN, -} - -#[pymethods] -impl PyRingZModN { - #[new] - fn new(n: i64) -> PyResult { - let inner = RingZModN::new(n).map_err(PyValueError::new_err)?; - Ok(Self { inner }) - } - - #[getter] - fn n_val(&self) -> i64 { - self.inner.n() - } - - fn add(&self, a: i64, b: i64) -> i64 { - self.inner.add(a, b) - } - - fn sub(&self, a: i64, b: i64) -> i64 { - self.inner.sub(a, b) - } - - fn mul(&self, a: i64, b: i64) -> i64 { - self.inner.mul(a, b) - } - - fn is_zero(&self, a: i64) -> bool { - self.inner.is_zero(a) - } - - fn gcd(&self, a: i64, b: i64) -> i64 { - self.inner.gcd(a, b) - } - - fn ass(&self, a: i64) -> i64 { - self.inner.ass(a) - } - - fn ann(&self, a: i64) -> i64 { - self.inner.ann(a) - } - - fn rem(&self, a: i64, b: i64) -> i64 { - self.inner.rem(a, b) - } - - fn quo(&self, a: i64, b: i64) -> PyResult { - self.inner.quo(a, b).map_err(PyValueError::new_err) - } - - fn div(&self, a: i64, b: i64) -> PyResult { - self.inner.div(a, b).map_err(PyValueError::new_err) - } - - fn unit(&self, a: i64) -> i64 { - self.inner.unit(a) - } - - fn gcdex(&self, a: i64, b: i64) -> PyResult<(i64, i64, i64, i64, i64)> { - Ok(self.inner.gcdex(a, b)) - } - - fn stab(&self, a: i64, b: i64, c: i64) -> PyResult { - self.inner.stab(a, b, c).map_err(PyValueError::new_err) - } -} - -/// Full Smith Normal Form: takes an n x m matrix and modulus, -/// returns (U, V, S) as numpy arrays with S = U @ A @ V (mod N). -#[pyfunction] -fn rust_smith_normal_form<'py>( - py: Python<'py>, - data: PyReadonlyArray2<'py, i64>, - modulus: i64, -) -> PyResult<( +type SnfArrays<'py> = ( Bound<'py, PyArray2>, Bound<'py, PyArray2>, Bound<'py, PyArray2>, -)> { - let r = RingZModN::new(modulus).map_err(PyValueError::new_err)?; - let a = data.as_array().to_owned(); - let n = a.nrows(); - let m = a.ncols(); - - if n == 0 || m == 0 { - let u = Array2::::eye(n); - let v = Array2::::eye(m); - return Ok((u.into_pyarray(py), v.into_pyarray(py), a.into_pyarray(py))); - } - - // Pad to square if needed - let sz = n.max(m); - let mut a_pad = Array2::zeros((sz, sz)); - a_pad.slice_mut(s![..n, ..m]).assign(&a); +); - let (u_pad, v_pad, s_pad) = smith_square(&a_pad, &r); - - // Crop back - let u = u_pad.slice(s![..n, ..n]).to_owned(); - let v = v_pad.slice(s![..m, ..m]).to_owned(); - let s_mat = s_pad.slice(s![..n, ..m]).to_owned(); - - Ok(( - u.into_pyarray(py), - v.into_pyarray(py), - s_mat.into_pyarray(py), - )) -} - -/// Compute SNF of a diagonal matrix. +/// Smith Normal Form via Storjohann band reduction. Takes an n x m matrix and +/// modulus; returns (U, V, S) as numpy arrays with S = U @ A @ V (mod N). #[pyfunction] -fn rust_smith_from_diagonal<'py>( - py: Python<'py>, - diag_data: PyReadonlyArray2<'py, i64>, - modulus: i64, -) -> PyResult<( - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, -)> { - let ring = RingZModN::new(modulus).map_err(PyValueError::new_err)?; - let arr = diag_data.as_array().to_owned(); - let (u, v, s) = smith_from_diagonal(&arr, &ring); - - Ok(( - u.into_pyarray(py), - v.into_pyarray(py), - s.into_pyarray(py), - )) -} - -/// Merge two SNF blocks. -#[pyfunction] -fn rust_merge_smith_blocks<'py>( +fn rust_smith_normal_form<'py>( py: Python<'py>, - a_data: PyReadonlyArray2<'py, i64>, - b_data: PyReadonlyArray2<'py, i64>, + data: PyReadonlyArray2<'py, i64>, modulus: i64, -) -> PyResult<( - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, -)> { - let ring = RingZModN::new(modulus).map_err(PyValueError::new_err)?; - let a = a_data.as_array().to_owned(); - let b = b_data.as_array().to_owned(); - - let (u, v, s) = merge_raw(&a, &b, &ring); +) -> PyResult> { + let a = data.as_array().to_owned(); + let (u, v, s) = modularsnf::smith_normal_form(&a, modulus).map_err(PyValueError::new_err)?; Ok((u.into_pyarray(py), v.into_pyarray(py), s.into_pyarray(py))) } -/// CRT-style SNF: takes an n x m matrix, modulus, and the prime +/// CRT fast-path SNF. Takes an n x m matrix, modulus, and the prime /// factorization of the modulus as (p, e) pairs (so factoring is amortized). /// Returns (U, V, S) with S = U @ A @ V (mod N). #[pyfunction] @@ -175,36 +34,23 @@ fn rust_crt_snf<'py>( data: PyReadonlyArray2<'py, i64>, modulus: i64, factors: Vec<(i64, u32)>, -) -> PyResult<( - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, -)> { +) -> PyResult> { let a = data.as_array().to_owned(); - let n = a.nrows(); - let m = a.ncols(); - if n == 0 || m == 0 { - let u = Array2::::eye(n); - let v = Array2::::eye(m); + if a.nrows() == 0 || a.ncols() == 0 { + let u = Array2::::eye(a.nrows()); + let v = Array2::::eye(a.ncols()); return Ok((u.into_pyarray(py), v.into_pyarray(py), a.into_pyarray(py))); } - let (u, v, s_mat) = crt_snf(&a, modulus, &factors); + let (u, v, s) = crt_snf(&a, modulus, &factors); - Ok(( - u.into_pyarray(py), - v.into_pyarray(py), - s_mat.into_pyarray(py), - )) + Ok((u.into_pyarray(py), v.into_pyarray(py), s.into_pyarray(py))) } /// Native Rust acceleration for modularsnf. #[pymodule] fn _rust(m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add_class::()?; - m.add_function(wrap_pyfunction!(rust_smith_from_diagonal, m)?)?; - m.add_function(wrap_pyfunction!(rust_merge_smith_blocks, m)?)?; m.add_function(wrap_pyfunction!(rust_smith_normal_form, m)?)?; m.add_function(wrap_pyfunction!(rust_crt_snf, m)?)?; Ok(()) diff --git a/modularsnf/__init__.py b/modularsnf/__init__.py index ef7488b..5cbf413 100644 --- a/modularsnf/__init__.py +++ b/modularsnf/__init__.py @@ -1,10 +1,8 @@ -from .matrix import RingMatrix -from .ring import RingZModN +from .crt import crt_snf from .snf import SNFResult, smith_normal_form_mod __all__ = [ - "RingMatrix", - "RingZModN", "SNFResult", + "crt_snf", "smith_normal_form_mod", ] diff --git a/modularsnf/_core.py b/modularsnf/_core.py new file mode 100644 index 0000000..7f527f3 --- /dev/null +++ b/modularsnf/_core.py @@ -0,0 +1,48 @@ +"""Input validation and marshalling for the modularsnf Python binding.""" + +from collections.abc import Sequence +from numbers import Integral +from typing import Optional + +import numpy as np + + +def validate_modulus(modulus: object) -> int: + """Coerce *modulus* to an int >= 2 or raise ValueError.""" + if isinstance(modulus, bool) or not isinstance(modulus, Integral): + raise ValueError(f"Modulus must be an integer >= 2, got {modulus!r}") + value = int(modulus) + if value < 2: + raise ValueError(f"Modulus must be an integer >= 2, got {modulus!r}") + return value + + +def to_int64_matrix( + matrix: Sequence[Sequence[int]], modulus: int +) -> Optional[np.ndarray]: + """Validate and marshal *matrix* into a 2-D int64 array reduced mod *modulus*. + + Returns ``None`` for an empty matrix (zero rows). + + Raises: + TypeError: If *matrix* is not a list/tuple of rows. + ValueError: If rows have unequal lengths. + OverflowError: If an entry does not fit in a signed 64-bit integer. + """ + if not isinstance(matrix, (list, tuple)): + raise TypeError("matrix must be a list of lists of integers") + + if len(matrix) == 0: + return None + + ncols = len(matrix[0]) + for i, row in enumerate(matrix): + if len(row) != ncols: + raise ValueError( + f"Ragged matrix: row 0 has {ncols} columns " + f"but row {i} has {len(row)} columns" + ) + + # np.asarray raises OverflowError for entries outside int64. + arr = np.asarray(matrix, dtype=np.int64) + return arr % modulus diff --git a/modularsnf/_rust.pyi b/modularsnf/_rust.pyi index 158db5d..a9916de 100644 --- a/modularsnf/_rust.pyi +++ b/modularsnf/_rust.pyi @@ -1,34 +1,12 @@ -class RustRingZModN: - def __init__(self, n: int) -> None: ... - def add(self, a: int, b: int) -> int: ... - def sub(self, a: int, b: int) -> int: ... - def mul(self, a: int, b: int) -> int: ... - def is_zero(self, a: int) -> bool: ... - def gcd(self, a: int, b: int) -> int: ... - def ass(self, a: int) -> int: ... - def ann(self, a: int) -> int: ... - def rem(self, a: int, b: int) -> int: ... - def quo(self, a: int, b: int) -> int: ... - def div(self, a: int, b: int) -> int: ... - def unit(self, a: int) -> int: ... - def gcdex(self, a: int, b: int) -> tuple[int, int, int, int, int]: ... - def stab(self, a: int, b: int, c: int) -> int: ... - import numpy as np from numpy.typing import NDArray -def rust_smith_from_diagonal( - diag_data: NDArray[np.int64], - modulus: int, -) -> tuple[NDArray[np.int64], NDArray[np.int64], NDArray[np.int64]]: ... - -def rust_merge_smith_blocks( - a_data: NDArray[np.int64], - b_data: NDArray[np.int64], +def rust_smith_normal_form( + data: NDArray[np.int64], modulus: int, ) -> tuple[NDArray[np.int64], NDArray[np.int64], NDArray[np.int64]]: ... - -def rust_smith_normal_form( +def rust_crt_snf( data: NDArray[np.int64], modulus: int, + factors: list[tuple[int, int]], ) -> tuple[NDArray[np.int64], NDArray[np.int64], NDArray[np.int64]]: ... diff --git a/modularsnf/band.py b/modularsnf/band.py deleted file mode 100644 index 478602c..0000000 --- a/modularsnf/band.py +++ /dev/null @@ -1,272 +0,0 @@ -"""Band reduction transforms for Storjohann's block algorithms. - -Provides helpers for triang and shift steps that manipulate banded blocks to -support Smith normal form reduction over modular rings. -""" - -from typing import Tuple - -import numpy as np - -from modularsnf.echelon import lemma_3_1 -from modularsnf.matrix import RingMatrix -from modularsnf.ring import RingZModN - - -def triang(B: RingMatrix, b: int) -> Tuple[RingMatrix, RingMatrix]: - """ - Triangulates the top-right block of an upper-b-banded matrix. - - Args: - B: An n1 x n1 upper-b-banded block with s1 = floor(b/2), s2 = b - 1, - and n1 = s1 + s2. - b: Bandwidth parameter; must satisfy b > 2. - - Returns: - A tuple ``(B_prime, W)`` where ``W`` is the s2 x s2 local right - transform such that ``B_prime = B * diag(I_{s1}, W)`` and the - top-right s1 x s2 block of ``B_prime`` has the triangulated structure - induced by ``lemma_3_1`` on ``B2^T``. - - Raises: - ValueError: If ``b <= 2`` or ``B`` is not an n1 x n1 block with - ``n1 = s1 + s2``. - """ - ring: RingZModN = B.ring - if b <= 2: - raise ValueError("Triang requires b > 2") - - s1 = b // 2 - s2 = b - 1 - n1 = s1 + s2 - - if B.nrows != B.ncols or B.nrows != n1: - raise ValueError( - f"Triang expects an n1 x n1 block with n1 = s1 + s2 = {n1}" - ) - - # Extract the top-right s1 x s2 block. - row0, row1 = 0, s1 - col0, col1 = s1, s1 + s2 - B2 = B.submatrix(row0, row1, col0, col1) # shape (s1, s2) - - # Transpose B2 to work with the s2 x s1 view. - C = B2.transpose() - - # Apply Lemma 3.1 to obtain the left transform for C. - U_left, T_ech, _ = lemma_3_1(C) # U_left: s2 x s2, T_ech: s2 x s1 - - # Use the transpose of the left transform as the local right transform. - W = U_left.transpose() # W: s2 x s2 - - # Form the block-diagonal right transform. - I_s1 = RingMatrix.identity(ring, s1) - V_loc = RingMatrix.block_diag(I_s1, W) # n1 x n1 - - # Multiply by the block-diagonal transform. - B_prime = B @ V_loc - - return B_prime, W - - -def shift(C: RingMatrix, b: int) -> Tuple[RingMatrix, RingMatrix, RingMatrix]: - """ - Applies Storjohann's Lemma 7.4 (Shift) to a local block. - - Args: - C: An n2 x n2 block with ``s2 = b - 1`` and ``n2 = 2 * s2`` that is - conceptually upper b-banded. - b: Bandwidth parameter; must satisfy b > 2. - - Returns: - A tuple ``(C_prime, U_block, V_block)`` where ``U_block`` and - ``V_block`` are the s2 x s2 local left and right transforms embedded as - ``diag(U_block, I_{s2})`` and ``diag(I_{s2}, V_block)`` such that - ``C_prime = diag(U_block, I_{s2}) * C * diag(I_{s2}, V_block)`` is the - shifted block from Lemma 7.4. - - Raises: - ValueError: If ``b <= 2`` or ``C`` is not an n2 x n2 block with - ``n2 = 2 * (b - 1)``. - - This function does not attempt to check or enforce bandedness; that is left - to higher-level tests. - """ - ring: RingZModN = C.ring - if b <= 2: - raise ValueError("Shift requires b > 2") - - s2 = b - 1 - n2 = 2 * s2 - - if C.nrows != C.ncols or C.nrows != n2: - raise ValueError( - f"Shift expects an n2 x n2 block with n2 = 2*(b-1) = {n2}, " - f"got shape {C.shape}" - ) - - # Block partition where each block is s2 x s2. - C1 = C.submatrix(0, s2, 0, s2) # top-left - C2 = C.submatrix(0, s2, s2, 2 * s2) # top-right - # C3 = C.submatrix(s2, 2 * s2, 0, s2) # bottom-left (not used directly) - # C4 = C.submatrix(s2, 2 * s2, s2, 2 * s2) - - # Apply Lemma 3.1 to triangularize C1 on the left. - U1, T1, _ = lemma_3_1(C1) # shapes: U1 s2×s2, T1 s2×s2 - - # Triangularize U1 * C2 on the right via Lemma 3.1 applied to its transpose. - C2_prime = U1 @ C2 # s2×s2 - - C2_prime_T = C2_prime.transpose() # s2×s2 - U2, T2, _ = lemma_3_1(C2_prime_T) # U2 s2×s2, T2 s2×s2 - V_block = U2.transpose() # s2×s2 - - # Assemble the local transforms and apply them to C. - I_s2 = RingMatrix.identity(ring, s2) - - U_full = RingMatrix.block_diag(U1, I_s2) # n2×n2 = diag(U1, I_s2) - V_full = RingMatrix.block_diag(I_s2, V_block) # n2×n2 = diag(I_s2, V_block) - - C_prime = U_full @ C @ V_full - - return C_prime, U1, V_block - - -def band_reduction( - A: RingMatrix, - b: int, - t: int = 0, -) -> Tuple[RingMatrix, RingMatrix, RingMatrix, int]: - """ - Proposition 7.1 (band reduction) with transform tracking. - - Input: - A : n x n upper-b-banded matrix (over a PIR) - b : current upper bandwidth - t : number of trailing zero columns (as in the paper, often 0 in practice) - - Output: - A_reduced : n x n upper-b'-banded matrix, with b' = floor(b/2) + 1 - U_band : n x n left transform - V_band : n x n right transform - b_new : new bandwidth (floor(b/2) + 1) - - Such that: - A_reduced = U_band * A * V_band - """ - ring = A.ring - n = A.nrows - if n != A.ncols: - raise ValueError("band_reduction expects a square matrix.") - - # Already at or below 2-banded: nothing to do. - if b <= 2: - I = RingMatrix.identity(ring, n) - return A.copy(), I, I, b - - # Storjohann parameters - s1 = b // 2 - s2 = b - 1 - n1 = s1 + s2 # size of the "triang" block - n2 = 2 * s2 # size of each "shift" block - - # Padding size in both directions (as in the paper / your original code) - pad = 2 * b - t - N_big = n + pad - - # Build padded matrix B with A in the top-left n x n principal block - B_arr = np.zeros((N_big, N_big), dtype=int) - B_arr[:n, :n] = A.data - B = RingMatrix._from_ndarray(ring, B_arr) - - # Global transforms on the padded system - U_big = RingMatrix.identity(ring, N_big) - V_big = RingMatrix.identity(ring, N_big) - - # Number of outer iterations (as in the paper) - num_i = max(0, (n - t + s1 - 1) // s1) - - for i in range(num_i): - top = i * s1 - if top + n1 > N_big: - break # safety guard - - # --- Step: triang on n1 x n1 principal block starting at (top, top) - B_block = B.submatrix(top, top + n1, top, top + n1) - # triang returns B_block_prime, W with B_block_prime = B_block * diag(I_{s1}, W) - B_block_prime, W = triang(B_block, b) - - # V_loc = diag(I_{s1}, W), so only W at position (top+s1, top+s1) - # differs from identity. Apply as sparse block update. - w_start = top + s1 - B.right_apply_block(W.data, w_start) - V_big.right_apply_block(W.data, w_start) - - # --- Inner "shift" blocks - numer = n - t - (i + 1) * s1 - if numer <= 0: - continue - - num_j = max(0, (numer + s2 - 1) // s2) - - for j in range(num_j): - offset = (i + 1) * s1 + j * s2 - if offset + n2 > N_big: - break # safety guard - - # Take current 2*s2 x 2*s2 block from B - C_block = B.submatrix(offset, offset + n2, offset, offset + n2) - - # shift returns C_prime, U_block, V_block with: - # C_prime = diag(U_block, I_{s2}) * C_block * diag(I_{s2}, V_block) - C_prime, U_block, V_block = shift(C_block, b) - - # U_loc = diag(U_block, I_{s2}): non-identity part is U_block - # at (offset, offset). - # V_loc2 = diag(I_{s2}, V_block): non-identity part is V_block - # at (offset+s2, offset+s2). - # Apply as sparse block updates instead of full N_big matmuls. - B.left_apply_block(U_block.data, offset) - B.right_apply_block(V_block.data, offset + s2) - U_big.left_apply_block(U_block.data, offset) - V_big.right_apply_block(V_block.data, offset + s2) - - # Extract the reduced n x n principal block and the corresponding transforms - A_reduced = B.submatrix(0, n, 0, n) - U_band = U_big.submatrix(0, n, 0, n) - V_band = V_big.submatrix(0, n, 0, n) - - b_new = (b // 2) + 1 - return A_reduced, U_band, V_band, b_new - - -def compute_upper_bandwidth(M: RingMatrix) -> int: - """ - Compute the upper bandwidth 'b' in the sense: - b = max{ j - i | M[i,j] ≠ 0, j ≥ i } + 1, - or 0 if the matrix is identically zero. - """ - nrows, ncols = M.shape - N = M.ring.N - nonzero = M.data % N != 0 - rows, cols = np.nonzero(nonzero) - if len(rows) == 0: - return 0 - offsets = cols - rows - upper_offsets = offsets[offsets >= 0] - if len(upper_offsets) == 0: - return 0 - return int(upper_offsets.max()) + 1 - - -def project_to_upper_bandwidth(M: RingMatrix, b: int) -> RingMatrix: - """ - Zero out entries outside the upper bandwidth 'b': - keep entries with 0 ≤ j - i < b, - zero everything else. - """ - nrows, ncols = M.shape - rows, cols = np.ogrid[:nrows, :ncols] - mask = (cols - rows >= 0) & (cols - rows < b) - arr = np.where(mask, M.data, 0) - return RingMatrix._from_ndarray(M.ring, arr) diff --git a/modularsnf/crt.py b/modularsnf/crt.py new file mode 100644 index 0000000..c1b8105 --- /dev/null +++ b/modularsnf/crt.py @@ -0,0 +1,87 @@ +"""CRT-based Smith Normal Form fast path (experimental). + +An alternative to :func:`modularsnf.smith_normal_form_mod`, specialized for the +small-modulus / large-matrix regime. It factors ``N = prod p^e``, solves the +SNF over each local ring ``Z/p^e`` by valuation-pivoted elimination, and +recombines via the Chinese Remainder Theorem. For general moduli prefer +``smith_normal_form_mod``. See ``docs/crt.md`` for the design. +""" + +from typing import List, Optional, Tuple + +from modularsnf._rust import rust_crt_snf as _rust_crt_snf + +from ._core import to_int64_matrix, validate_modulus +from .snf import SNFResult + + +def factorize(n: int) -> List[Tuple[int, int]]: + """Factor *n* into ``[(p, e), ...]`` by trial division. + + Intended for the small moduli this path targets; not a general-purpose + integer factorizer. + """ + if n < 2: + raise ValueError(f"modulus must be >= 2 to factor, got {n}") + factors: List[Tuple[int, int]] = [] + d = 2 + while d * d <= n: + if n % d == 0: + e = 0 + while n % d == 0: + n //= d + e += 1 + factors.append((d, e)) + d += 1 if d == 2 else 2 + if n > 1: + factors.append((n, 1)) + return factors + + +def _resolve_factors( + modulus: int, factors: Optional[List[Tuple[int, int]]] +) -> List[Tuple[int, int]]: + if factors is None: + return factorize(modulus) + prod = 1 + for p, e in factors: + prod *= p**e + if prod != modulus: + raise ValueError( + f"factors {factors} multiply to {prod}, not modulus {modulus}" + ) + return factors + + +def crt_snf( + matrix: list[list[int]], + modulus: int, + *, + factors: Optional[List[Tuple[int, int]]] = None, +) -> SNFResult: + """Compute the Smith Normal Form over Z/NZ via the CRT fast path. + + Returns ``(S, U, V)`` with ``S = U @ A @ V`` (mod *modulus*) — the same + contract as :func:`modularsnf.smith_normal_form_mod`. + + Args: + matrix: 2-D list of signed 64-bit integers. May be rectangular. + modulus: Signed 64-bit integer *N* >= 2 defining the ring Z/NZ. + factors: Prime factorization of *modulus* as ``[(p, e), ...]``. When + omitted it is computed by trial division; pass it explicitly to + amortize factoring across many calls with the same modulus. + + Raises: + OverflowError: If *modulus* or any matrix entry is outside int64. + ValueError: If *modulus* < 2, *factors* does not multiply to *modulus*, + or rows have unequal lengths. + TypeError: If *matrix* is not a list of lists. + """ + modulus = validate_modulus(modulus) + arr = to_int64_matrix(matrix, modulus) + if arr is None: + return SNFResult(S=[], U=[], V=[]) + + factors = _resolve_factors(modulus, factors) + u, v, s = _rust_crt_snf(arr, modulus, factors) + return SNFResult(S=s.tolist(), U=u.tolist(), V=v.tolist()) diff --git a/modularsnf/diagonal.py b/modularsnf/diagonal.py deleted file mode 100644 index 4504100..0000000 --- a/modularsnf/diagonal.py +++ /dev/null @@ -1,376 +0,0 @@ -"""Helpers for diagonalizing matrices into Smith normal form.""" - -from typing import Tuple - -import numpy as np - -from modularsnf.matrix import RingMatrix -from modularsnf.ring import RingZModN - -try: - from modularsnf._rust import ( - rust_merge_smith_blocks as _rust_merge, - ) - from modularsnf._rust import ( - rust_smith_from_diagonal as _rust_diag, - ) -except ImportError: - _rust_diag = None # type: ignore[assignment] - _rust_merge = None # type: ignore[assignment] - - -def smith_from_diagonal( - D: RingMatrix, -) -> Tuple[RingMatrix, RingMatrix, RingMatrix]: - """Compute SNF transforms for a diagonal square matrix. - - Pads the matrix to a power-of-two dimension, runs the recursive merge - routine, and then crops the result back to the original size. - - Args: - D: Diagonal matrix to convert to Smith form. - - Returns: - A tuple ``(U, V, S)`` where ``S`` is the Smith form of ``D`` and - ``U`` and ``V`` are the unimodular transforms such that ``S = U D V``. - - Raises: - ValueError: If ``D`` is not square. - """ - - if D.nrows != D.ncols: - raise ValueError("smith_from_diagonal expects a square matrix") - - n = D.nrows - if n == 0: - ring = D.ring - I0 = RingMatrix.identity(ring, 0) - return I0, I0, D.copy() - - ring = D.ring - - if _rust_diag is not None: - u_arr, v_arr, s_arr = _rust_diag( - D.data.astype(np.int64), - ring.N, - ) - return ( - RingMatrix._from_ndarray(ring, u_arr[:n, :n].copy()), - RingMatrix._from_ndarray(ring, v_arr[:n, :n].copy()), - RingMatrix._from_ndarray(ring, s_arr[:n, :n].copy()), - ) - - # Python fallback: pad to power of two and run raw merge. - p = max(n, 1) - size = 1 << (p - 1).bit_length() - pad_arr = np.zeros((size, size), dtype=int) - pad_arr[:n, :n] = D.data - - U_arr, V_arr, S_arr = _smith_from_diagonal_raw(pad_arr, ring) - - U = RingMatrix._from_ndarray(ring, U_arr[:n, :n].copy()) - V = RingMatrix._from_ndarray(ring, V_arr[:n, :n].copy()) - S = RingMatrix._from_ndarray(ring, S_arr[:n, :n].copy()) - - return U, V, S - - -def merge_smith_blocks( - A: RingMatrix, B: RingMatrix -) -> Tuple[RingMatrix, RingMatrix, RingMatrix]: - """Merge two SNF blocks of equal size. - - Implements the recursive merge step from Theorem 7.11. ``A`` and ``B`` - must be the same power-of-two dimension and already in Smith form. - - Args: - A: Smith-form matrix to merge. - B: Smith-form matrix to merge. - - Returns: - A tuple ``(U, V, S)`` where ``S`` is the merged Smith form and ``U`` - and ``V`` are the associated unimodular transforms. - - Raises: - ValueError: If the matrices use different rings or have mismatched - sizes. - """ - - ring = A.ring - n = A.nrows - - if B.ring is not ring: - raise ValueError("merge_smith_blocks requires same ring") - if n != A.ncols or n != B.nrows or n != B.ncols: - raise ValueError( - "merge_smith_blocks expects A,B to be same square size, " - f"got {A.shape}, {B.shape}" - ) - - if _rust_merge is not None: - u_arr, v_arr, s_arr = _rust_merge( - A.data.astype(np.int64), - B.data.astype(np.int64), - ring.N, - ) - return ( - RingMatrix._from_ndarray(ring, u_arr), - RingMatrix._from_ndarray(ring, v_arr), - RingMatrix._from_ndarray(ring, s_arr), - ) - - u_arr, v_arr, s_arr = _merge_raw(A.data, B.data, ring) - return ( - RingMatrix._from_ndarray(ring, u_arr), - RingMatrix._from_ndarray(ring, v_arr), - RingMatrix._from_ndarray(ring, s_arr), - ) - - -def _merge_scalars( - a: int, b: int, ring: RingZModN -) -> Tuple[RingMatrix, RingMatrix, RingMatrix]: - """Merge two scalar SNF blocks. - - Args: - a: Upper-left scalar entry. - b: Lower-right scalar entry. - ring: Underlying ring for the scalar operations. - - Returns: - A tuple ``(U, V, S)`` where ``S`` is the Smith form of - ``diag(a, b)`` and ``U`` and ``V`` are the unimodular transforms. - """ - u_arr, v_arr, s_arr = _merge_scalars_raw(a, b, ring) - return ( - RingMatrix._from_ndarray(ring, u_arr), - RingMatrix._from_ndarray(ring, v_arr), - RingMatrix._from_ndarray(ring, s_arr), - ) - - -# --------------------------------------------------------------------------- -# Raw numpy implementations — no RingMatrix construction in hot paths -# --------------------------------------------------------------------------- - -_EYE2 = np.eye(2, dtype=int) -_ZERO2 = np.zeros((2, 2), dtype=int) - - -def _merge_scalars_raw( - a: int, b: int, ring: RingZModN -) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: - """Merge two scalar SNF blocks on raw arrays. - - Returns (U, V, S) as 2x2 numpy arrays reduced mod ring.N. - """ - g, s, t, u, v = ring.gcdex(a, b) - N = ring.N - - if g % N == 0: - return _EYE2.copy(), _EYE2.copy(), _ZERO2.copy() - - tb = (t * b) % N - q_raw = ring.div(tb, g) - q = (-q_raw) % N - - u_arr = np.array([[s, t], [u, v]], dtype=int) % N - v_arr = np.array([[1, q], [1, (1 + q) % N]], dtype=int) - s_arr = (u_arr @ np.array([[a, 0], [0, b]], dtype=int) @ v_arr) % N - return u_arr, v_arr, s_arr - - -def _is_zero_raw(arr: np.ndarray, N: int) -> bool: - """Check if a diagonal SNF array is zero (first entry is zero mod N).""" - if arr.shape[0] == 0: - return True - return arr[0, 0] % N == 0 - - -def _get_rank_raw(arr: np.ndarray, N: int) -> int: - """Count nonzero diagonal entries mod N.""" - n = min(arr.shape[0], arr.shape[1]) - rank = 0 - for i in range(n): - if arr[i, i] % N != 0: - rank += 1 - return rank - - -def _merge_raw( - a_arr: np.ndarray, b_arr: np.ndarray, ring: RingZModN -) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: - """Raw-array version of merge_smith_blocks. - - Args: - a_arr: n x n SNF array. - b_arr: n x n SNF array. - ring: Ring for arithmetic. - - Returns: - (U, V, S) as raw 2n x 2n numpy arrays. - """ - N = ring.N - n = a_arr.shape[0] - - if n == 0: - e = np.zeros((0, 0), dtype=int) - return e, e, e - - if n == 1: - u2, v2, s2 = _merge_scalars_raw( - int(a_arr[0, 0]), int(b_arr[0, 0]), ring - ) - return u2, v2, s2 - - t = n // 2 - NN = 2 * n - - A1 = a_arr[0:t, 0:t].copy() - A2 = a_arr[t:n, t:n].copy() - B1 = b_arr[0:t, 0:t].copy() - B2 = b_arr[t:n, t:n].copy() - - U_total = np.eye(NN, dtype=int) - V_total = np.eye(NN, dtype=int) - - index_map = (0, t, n, n + t) # tuple for speed - - def apply_step(block1, block2, idx1, idx2): - if _is_zero_raw(block1, N) and _is_zero_raw(block2, N): - return block1, block2 - - u_loc, v_loc, s_loc = _merge_raw(block1, block2, ring) - - b1_prime = s_loc[0:t, 0:t].copy() - b2_prime = s_loc[t : 2 * t, t : 2 * t].copy() - - s1 = index_map[idx1] - s2_ = index_map[idx2] - - # Left-apply U_loc block pair to U_total - r1 = U_total[s1 : s1 + t, :].copy() - r2 = U_total[s2_ : s2_ + t, :].copy() - U_total[s1 : s1 + t, :] = ( - u_loc[0:t, 0:t] @ r1 + u_loc[0:t, t : 2 * t] @ r2 - ) % N - U_total[s2_ : s2_ + t, :] = ( - u_loc[t : 2 * t, 0:t] @ r1 + u_loc[t : 2 * t, t : 2 * t] @ r2 - ) % N - - # Right-apply V_loc block pair to V_total - c1 = V_total[:, s1 : s1 + t].copy() - c2 = V_total[:, s2_ : s2_ + t].copy() - V_total[:, s1 : s1 + t] = ( - c1 @ v_loc[0:t, 0:t] + c2 @ v_loc[t : 2 * t, 0:t] - ) % N - V_total[:, s2_ : s2_ + t] = ( - c1 @ v_loc[0:t, t : 2 * t] + c2 @ v_loc[t : 2 * t, t : 2 * t] - ) % N - - return b1_prime, b2_prime - - A1, B1 = apply_step(A1, B1, 0, 2) - A2, B2 = apply_step(A2, B2, 1, 3) - A2, B1 = apply_step(A2, B1, 1, 2) - B1, B2 = apply_step(B1, B2, 2, 3) - - if not _is_zero_raw(B2, N): - r_B1 = _get_rank_raw(B1, N) - - if r_B1 < t: - r_B2 = _get_rank_raw(B2, N) - - P_arr = np.zeros((2 * t, 2 * t), dtype=int) - for i in range(r_B1): - P_arr[i, i] = 1 - for i in range(r_B2): - P_arr[r_B1 + i, t + i] = 1 - cr = r_B1 + r_B2 - for i in range(t - r_B1): - P_arr[cr, r_B1 + i] = 1 - cr += 1 - for i in range(t - r_B2): - P_arr[cr, t + r_B2 + i] = 1 - cr += 1 - - # Build P_glob as identity with P_arr at (n, n) - P_glob = np.eye(NN, dtype=int) - P_glob[n : n + 2 * t, n : n + 2 * t] = P_arr - P_glob_T = P_glob.T.copy() - - U_total = (P_glob @ U_total) % N - V_total = (V_total @ P_glob_T) % N - - # Combine B1, B2 into block diag, permute - B_comb = np.zeros((2 * t, 2 * t), dtype=int) - B_comb[0:t, 0:t] = B1 - B_comb[t : 2 * t, t : 2 * t] = B2 - S_target = (P_arr @ B_comb @ P_arr.T) % N - B1 = S_target[0:t, 0:t].copy() - B2 = S_target[t : 2 * t, t : 2 * t].copy() - - # Assemble S_final as block_diag(A1, A2, B1, B2) - S_final = np.zeros((NN, NN), dtype=int) - S_final[0:t, 0:t] = A1 - S_final[t:n, t:n] = A2 - S_final[n : n + t, n : n + t] = B1 - S_final[n + t : NN, n + t : NN] = B2 - - return U_total, V_total, S_final - - -def _smith_from_diagonal_raw( - diag_arr: np.ndarray, ring: RingZModN -) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: - """Raw-array power-of-two diagonal SNF via bottom-up iterative merge. - - Args: - diag_arr: n x n diagonal array with n a power of two. - ring: Ring for arithmetic. - - Returns: - (U, V, S) as n x n numpy arrays. - """ - N = ring.N - n = diag_arr.shape[0] - - if n <= 1: - return np.eye(n, dtype=int), np.eye(n, dtype=int), diag_arr.copy() - - # Bottom-up: start with n 1x1 blocks, merge pairwise - # Each element is (U, V, S) as raw arrays - blocks = [] - for i in range(n): - s = diag_arr[i : i + 1, i : i + 1].copy() - blocks.append( - (np.ones((1, 1), dtype=int), np.ones((1, 1), dtype=int), s) - ) - - size = 1 - while size < n: - new_blocks = [] - for i in range(0, len(blocks), 2): - U1, V1, A = blocks[i] - U2, V2, B = blocks[i + 1] - - U_merge, V_merge, S = _merge_raw(A, B, ring) - - # U_block = block_diag(U1, U2), then U_total = U_merge @ U_block - bsz = size - U_block = np.zeros((2 * bsz, 2 * bsz), dtype=int) - U_block[0:bsz, 0:bsz] = U1 - U_block[bsz : 2 * bsz, bsz : 2 * bsz] = U2 - U_total = (U_merge @ U_block) % N - - V_block = np.zeros((2 * bsz, 2 * bsz), dtype=int) - V_block[0:bsz, 0:bsz] = V1 - V_block[bsz : 2 * bsz, bsz : 2 * bsz] = V2 - V_total = (V_block @ V_merge) % N - - new_blocks.append((U_total, V_total, S)) - - blocks = new_blocks - size *= 2 - - return blocks[0] diff --git a/modularsnf/echelon.py b/modularsnf/echelon.py deleted file mode 100644 index c6f4762..0000000 --- a/modularsnf/echelon.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Row-echelon utilities for modular matrices. - -This module contains the row-reduction helpers used throughout the Smith normal -form pipeline over the ring ``Z/nZ``. The routines mirror the primitives in -Storjohann's dissertation: - -* ``lemma_3_1`` constructs a unimodular left transform that drives a matrix to - row-echelon form while tracking rank. -* ``index1_reduce_on_columns`` applies a targeted column-reduction step for the - index-1 case in the band-reduction phase. - -Each helper returns the unimodular matrix that witnesses the transformation so -callers can compose them into larger reduction sequences. -""" - -from typing import Tuple - -from .matrix import RingMatrix -from .ring import RingZModN - - -def lemma_3_1(A: RingMatrix) -> Tuple[RingMatrix, RingMatrix, int]: - """Compute a row-echelon form using Storjohann's Lemma 3.1. - - The routine walks across the columns of ``A``. For each column it clears - entries below the current pivot row by applying a 2×2 unimodular transform - derived from the extended GCD of the pivot and the entry to eliminate. The - accumulated unimodular matrix ``U`` records the full left transform such - that ``T = U * A`` is in row-echelon form. - - Args: - A: Input matrix over the modular ring ``Z/nZ``. - - Returns: - Tuple containing the unimodular transform ``U``, the row-echelon matrix - ``T``, and the detected rank. - """ - - ring: RingZModN = A.ring - n_rows, n_cols = A.nrows, A.ncols - - U = RingMatrix.identity(ring, n_rows) - T = A.copy() - - r = 0 # Tracks the pivot row index, which doubles as the current rank. - - for k in range(n_cols): - if r >= n_rows: - break - - # Eliminate entries below row r in column k. - for i in range(r + 1, n_rows): - a = T.data[r, k] - b = T.data[i, k] - - if ring.is_zero(b): - continue - - # Compute the unimodular 2×2 transform from gcdex. - g, s, t, u, v = ring.gcdex(a, b) - - # Apply the paired row operations to both T and U. - T.apply_row_2x2(r, i, s, t, u, v) - U.apply_row_2x2(r, i, s, t, u, v) - - # If the pivot is nonzero, lock this row and move down. - if not ring.is_zero(T.data[r, k]): - r += 1 - - return U, T, r - - -def index1_reduce_on_columns( - A: RingMatrix, k: int -) -> Tuple[RingMatrix, RingMatrix]: - """Perform specialized index-1 reduction on the first ``k`` columns. - - This routine targets the "index-1" case from section 7.3 (step 9) of - Storjohann's algorithm. It zeroes entries above the diagonal in the leading - ``k`` columns by adding multiples of each pivot row to the rows above it. - The unimodular matrix ``U`` captures the cumulative left transform such that - ``T = U * A`` reflects the reduced structure. - - Args: - A: Input matrix over the modular ring ``Z/nZ``. - k: Column count to reduce (starting after column ``0``). - - Returns: - Tuple containing the unimodular transform ``U`` and the reduced matrix - ``T``. - """ - - ring = A.ring - n = A.nrows - U = RingMatrix.identity(ring, n) - T = A.copy() - - for j in range(1, k): - sj = T.data[j, j] - for i in range(j): - x = T.data[i, j] - if ring.is_zero(x): - continue - phi = (-ring.quo(x, sj)) % ring.N - # Vectorized row updates to clear the current entry. - T.data[i, :] = (T.data[i, :] + phi * T.data[j, :]) % ring.N - U.data[i, :] = (U.data[i, :] + phi * U.data[j, :]) % ring.N - - return U, T diff --git a/modularsnf/matrix.py b/modularsnf/matrix.py deleted file mode 100644 index abf2bc9..0000000 --- a/modularsnf/matrix.py +++ /dev/null @@ -1,348 +0,0 @@ -"""Ring-aware matrix utilities for modular Smith normal form routines. - -Defines the ``RingMatrix`` dataclass and helpers that normalize data, manage -block operations, and align shapes for modular arithmetic workflows. -""" - -from dataclasses import InitVar, dataclass, field -from typing import List, Tuple - -import numpy as np - -from modularsnf.ring import INT64_MAX, INT64_MIN, RingZModN, _coerce_int64 - - -def _coerce_object_array_to_int64(data: np.ndarray) -> np.ndarray: - """Convert an object array to ``np.int64`` with explicit range checks.""" - flat = np.empty(data.size, dtype=np.int64) - for idx, value in enumerate(data.flat): - flat[idx] = _coerce_int64("matrix entry", value) - return flat.reshape(data.shape) - - -def _normalize_ndarray(data: np.ndarray) -> np.ndarray: - """Normalize ndarray inputs to a 2-D ``np.int64`` array.""" - arr = np.array(data, copy=True) - if arr.size == 0: - return np.zeros((0, 0), dtype=np.int64) - if arr.ndim == 1: - arr = arr.reshape(1, arr.size) - if arr.ndim != 2: - raise ValueError("Matrix data must be 2-dimensional") - - if arr.dtype == object: - return _coerce_object_array_to_int64(arr) - - if np.issubdtype(arr.dtype, np.unsignedinteger): - max_val = int(arr.max()) - if max_val > INT64_MAX: - raise OverflowError( - "matrix entries must fit in a signed 64-bit integer" - ) - return arr.astype(np.int64, copy=False) - - if np.issubdtype(arr.dtype, np.signedinteger): - min_val = int(arr.min()) - max_val = int(arr.max()) - if min_val < INT64_MIN or max_val > INT64_MAX: - raise OverflowError( - "matrix entries must fit in a signed 64-bit integer" - ) - return arr.astype(np.int64, copy=False) - - raise TypeError( - "Matrix data must contain integers representable as signed 64-bit " - "values" - ) - - -def _normalize_matrix_data( - data: list[list[int]] | np.ndarray, -) -> np.ndarray: - """ - Ensure matrix-like input is a 2D ``np.ndarray`` of ints. - - Allows callers to pass either a list-of-lists or an ndarray. Empty inputs - become an explicit (0, 0) array so downstream shape logic remains simple. - """ - if isinstance(data, np.ndarray): - return _normalize_ndarray(data) - - rows = [list(row) for row in data] - if not rows: - return np.zeros((0, 0), dtype=np.int64) - - ncols = len(rows[0]) - for row in rows: - if len(row) != ncols: - raise ValueError("All rows must have the same length") - - arr = np.empty((len(rows), ncols), dtype=np.int64) - for i, row in enumerate(rows): - for j, value in enumerate(row): - arr[i, j] = _coerce_int64("matrix entry", value) - return arr - - -def _is_within_modulus(data: np.ndarray, N: int) -> bool: - """Return True when ``data`` is provably in ``[0, N)``. - - Avoids expensive copies by only skipping the modulus reduction when values - are demonstrably safe. - """ - if data.size == 0: - return True - if not np.issubdtype(data.dtype, np.integer): - return False - try: - min_val = data.min() - max_val = data.max() - except (OverflowError, TypeError, ValueError): - return False - return 0 <= min_val and max_val < N - - -def _to_list_if_array( - data: list[list[int]] | np.ndarray, -) -> list[list[int]]: - """Return nested Python lists, preserving list inputs for callers that expect them.""" - if isinstance(data, np.ndarray): - return data.tolist() - return data - - -@dataclass -class RingMatrix: - ring: RingZModN - _data_init: InitVar[list[list[int]] | np.ndarray] - data: np.ndarray = field(init=False) - _assume_reduced: bool = field(default=False, init=True, repr=False) - - def __post_init__(self, _data_init: list[list[int]] | np.ndarray) -> None: - self.data = _normalize_matrix_data(_data_init) - N = self.ring.N - - # Fast path: internal callers that explicitly guarantee the array is - # already reduced, or data we can verify is inside [0, N). - if not (self._assume_reduced or _is_within_modulus(self.data, N)): - self.data %= N - - @property - def nrows(self) -> int: - return self.data.shape[0] - - @property - def ncols(self) -> int: - return self.data.shape[1] if self.data.size else 0 - - @property - def shape(self) -> Tuple[int, int]: - return self.nrows, self.ncols - - @classmethod - def from_rows(cls, ring: RingZModN, rows: List[List[int]]) -> "RingMatrix": - return cls(ring, rows) - - @classmethod - def identity(cls, ring: RingZModN, n: int) -> "RingMatrix": - return cls(ring, np.eye(n, dtype=np.int64)) - - @classmethod - def diagonal(cls, ring: RingZModN, diag: List[int]) -> "RingMatrix": - n = len(diag) - rows = np.zeros((n, n), dtype=np.int64) - for i, v in enumerate(diag): - rows[i, i] = v - return cls(ring, rows) - - @classmethod - def _from_ndarray(cls, ring: RingZModN, data: np.ndarray) -> "RingMatrix": - """Internal fast constructor for pre-normalized arrays. - - Callers must ensure ``data`` is already 2-D, reduced modulo ``ring.N``, - and uses an appropriate dtype. No copying or normalization is - performed. - """ - if data.ndim != 2: - raise ValueError("Matrix data must be 2-dimensional") - - obj = cls.__new__(cls) - obj.ring = ring - obj.data = data - return obj - - @classmethod - def block_diag(cls, A: "RingMatrix", B: "RingMatrix") -> "RingMatrix": - """ - Block-diagonal composition: - [ A 0 ] - [ 0 B ] - - A and B may be rectangular, but must share the same ring. - """ - if A.ring is not B.ring: - raise ValueError("Cannot form block diagonal over different rings") - ring = A.ring - a_rows, a_cols = A.shape - b_rows, b_cols = B.shape - - total_rows = a_rows + b_rows - total_cols = a_cols + b_cols - - data = np.zeros((total_rows, total_cols), dtype=np.int64) - - # top-left: A - data[:a_rows, :a_cols] = A.data - - # bottom-right: B - data[a_rows:, a_cols:] = B.data - - return cls._from_ndarray(ring, data) - - def copy(self) -> "RingMatrix": - return RingMatrix._from_ndarray(self.ring, np.copy(self.data)) - - def transpose(self) -> "RingMatrix": - return RingMatrix._from_ndarray(self.ring, np.transpose(self.data)) - - def __matmul__(self, other: "RingMatrix") -> "RingMatrix": - if self.ring is not other.ring: - raise ValueError("Cannot multiply matrices over different rings") - - cA = self.ncols - rB = other.nrows - - if cA != rB: - raise ValueError(f"Dimension mismatch: {cA} != {rB}") - - ring = self.ring - N = ring.N - - A = self.data - B = other.data - - C = (A @ B) % N - - return RingMatrix._from_ndarray(ring, C) - - def pad_to(self, target_rows: int, target_cols: int) -> "RingMatrix": - rows, cols = self.shape - if rows == target_rows and cols == target_cols: - return self.copy() - N = self.ring.N - new = np.zeros((target_rows, target_cols), dtype=np.int64) - new[:rows, :cols] = self.data[:rows, :cols] % N - return RingMatrix._from_ndarray(self.ring, new) - - def pad_to_square_power2(self) -> "RingMatrix": - n = max(self.nrows, self.ncols) - if n <= 0: - return self.copy() - size = 1 << (n - 1).bit_length() - return self.pad_to(size, size) - - def submatrix( - self, row_start: int, row_end: int, col_start: int, col_end: int - ) -> "RingMatrix": - """ - Return a view copy of rows [row_start:row_end) and - cols [col_start:col_end). - """ - rows = self.data[row_start:row_end, col_start:col_end] - return RingMatrix._from_ndarray(self.ring, np.copy(rows)) - - def write_block( - self, row_start: int, col_start: int, block: "RingMatrix" - ) -> None: - """ - Overwrite a sub-block of self starting at (row_start, col_start) - with the contents of 'block'. Rings must match. - """ - if self.ring is not block.ring: - raise ValueError("Cannot write block with different ring") - b_rows, b_cols = block.shape - self.data[ - row_start : row_start + b_rows, col_start : col_start + b_cols - ] = block.data - - def apply_row_2x2( - self, r: int, i: int, s: int, t: int, u: int, v: int - ) -> None: - """In-place: [row_r; row_i] <- [s t; u v] [row_r; row_i].""" - N = self.ring.N - rows = self.data[[r, i], :] - transform = np.array([[s, t], [u, v]], dtype=int) % N - new_rows = (transform @ rows) % N - self.data[r, :] = new_rows[0] - self.data[i, :] = new_rows[1] - - def left_apply_block(self, block: np.ndarray, start: int) -> None: - """In-place ``self <- M @ self`` where ``M = I`` with *block* at ``(start, start)``. - - Only rows ``[start:start+k]`` are modified (``k = block.shape[0]``). - """ - N = self.ring.N - k = block.shape[0] - rows = self.data[start : start + k, :].copy() - self.data[start : start + k, :] = (block @ rows) % N - - def right_apply_block(self, block: np.ndarray, start: int) -> None: - """In-place ``self <- self @ M`` where ``M = I`` with *block* at ``(start, start)``. - - Only columns ``[start:start+k]`` are modified. - """ - N = self.ring.N - k = block.shape[0] - cols = self.data[:, start : start + k].copy() - self.data[:, start : start + k] = (cols @ block) % N - - def left_apply_block_pair( - self, - u00: np.ndarray, - u01: np.ndarray, - u10: np.ndarray, - u11: np.ndarray, - start1: int, - start2: int, - ) -> None: - """In-place ``self <- M @ self`` where ``M = I`` with a 2x2 block arrangement. - - ``M[start1:start1+t, start1:start1+t] = u00``, etc. - Only rows ``[start1:start1+t]`` and ``[start2:start2+t]`` change. - """ - N = self.ring.N - t = u00.shape[0] - r1 = self.data[start1 : start1 + t, :].copy() - r2 = self.data[start2 : start2 + t, :].copy() - self.data[start1 : start1 + t, :] = (u00 @ r1 + u01 @ r2) % N - self.data[start2 : start2 + t, :] = (u10 @ r1 + u11 @ r2) % N - - def right_apply_block_pair( - self, - v00: np.ndarray, - v01: np.ndarray, - v10: np.ndarray, - v11: np.ndarray, - start1: int, - start2: int, - ) -> None: - """In-place ``self <- self @ M`` where ``M = I`` with a 2x2 block arrangement. - - Only columns ``[start1:start1+t]`` and ``[start2:start2+t]`` change. - """ - N = self.ring.N - t = v00.shape[0] - c1 = self.data[:, start1 : start1 + t].copy() - c2 = self.data[:, start2 : start2 + t].copy() - self.data[:, start1 : start1 + t] = (c1 @ v00 + c2 @ v10) % N - self.data[:, start2 : start2 + t] = (c1 @ v01 + c2 @ v11) % N - - def to_sympy(self): - import sympy as sp - - return sp.Matrix(_to_list_if_array(self.data)) - - def pprint(self): - from sympy import pprint - - pprint(self.to_sympy()) diff --git a/modularsnf/ring.py b/modularsnf/ring.py deleted file mode 100644 index 085fe20..0000000 --- a/modularsnf/ring.py +++ /dev/null @@ -1,177 +0,0 @@ -"""Modular arithmetic helpers for Z/NZ used across Smith normal form. - -Provides the ``RingZModN`` abstraction with arithmetic, ideal, and gcd tools -needed by the Smith normal form algorithms implemented in this package. - -When the Rust extension is available, hot-path methods delegate to the native -implementation for speed. The pure-Python implementation remains available as -the reference path and as a fallback when the extension cannot be imported. -""" - -import math -from numbers import Integral - -try: - from modularsnf._rust import RustRingZModN as _RustRing -except ImportError: - _RustRing = None # type: ignore[assignment] - -INT64_MIN = -(1 << 63) -INT64_MAX = (1 << 63) - 1 - - -def _coerce_int64(name: str, value: object) -> int: - """Return *value* as a signed 64-bit integer or raise.""" - if isinstance(value, bool) or not isinstance(value, Integral): - raise TypeError(f"{name} must be an integer") - - int_value = int(value) - if int_value < INT64_MIN or int_value > INT64_MAX: - raise OverflowError(f"{name} must fit in a signed 64-bit integer") - return int_value - - -class RingZModN: - """ - Represents Z/NZ. Implements ALL basic operations defined in - Storjohann's Dissertation Section 1.1. - """ - - N: int - - def __init__(self, N: int) -> None: - int_modulus = _coerce_int64("modulus", N) - if int_modulus < 2: - raise ValueError("Modulus N must be >= 2") - self.N = int_modulus - self._rust = _RustRing(int_modulus) if _RustRing is not None else None - - def add(self, a: int, b: int) -> int: - return (a + b) % self.N - - def sub(self, a: int, b: int) -> int: - return (a - b) % self.N - - def mul(self, a: int, b: int) -> int: - return (a * b) % self.N - - def is_zero(self, a: int) -> bool: - return (a % self.N) == 0 - - def gcd(self, a: int, b: int) -> int: - """Principal generator of the ideal (a, b) in Z/NZ.""" - if self._rust is not None: - return self._rust.gcd(a, b) - return math.gcd(a % self.N, math.gcd(b % self.N, self.N)) - - def rem(self, a: int, b: int) -> int: - """ - Return the prescribed residue of a with respect to Ass(b). - In Z, this is typically a % b, but we must respect the ring structure. - Ref: Def 1.4 consistent with Z. - """ - b_ass = self.ass(b) - if b_ass == 0: - return a % self.N - return (a % self.N) % b_ass - - def quo(self, a: int, b: int) -> int: - """ - Return q such that a - qb = rem(a, b). - """ - r = self.rem(a, b) - diff = self.sub(a, r) - # Find q such that q*b equals diff; diff is a multiple of b by definition of the remainder. - return self.div(diff, b) - - def ann(self, a: int) -> int: - """ - Return generator of the annihilator ideal {x | xa = 0}. - In Z/N, Ann(a) is generated by N / gcd(a, N). - """ - g = math.gcd(a % self.N, self.N) - return (self.N // g) % self.N - - def gcdex(self, a: int, b: int) -> tuple[int, int, int, int, int]: - """ - Returns (g, s, t, u, v) such that: - [s t] [a] = [g] - [u v] [b] [0] - and sv - tu = unit. - """ - if self._rust is not None: - return self._rust.gcdex(a, b) - a_val = a % self.N - b_val = b % self.N - - # Return the specialized form when b is a multiple of a in Z/N. - if a_val != 0 and (b_val % math.gcd(a_val, self.N) == 0): - # Check if exact division exists - try: - q = self.div(b_val, a_val) - # s=1, t=0, u=-q, v=1 -> det = 1 - return (a_val, 1, 0, (self.N - q) % self.N, 1) - except ValueError: - pass # Fall through to standard logic if exact division fails. - - # Standard Euclidean - r0, r1 = a_val, b_val - s0, s1 = 1, 0 - t0, t1 = 0, 1 - - while r1 != 0: - q = r0 // r1 - r0, r1 = r1, r0 - q * r1 - s0, s1 = s1, s0 - q * s1 - t0, t1 = t1, t0 - q * t1 - - g, s, t = r0, s0, t0 - - # Compute u, v such that u*a + v*b = 0 and det is a unit (u = -b/g, v = a/g). - if g == 0: - return 0, 1, 0, 0, 1 - - u = -(b_val // g) - v = a_val // g - - return (g % self.N, s % self.N, t % self.N, u % self.N, v % self.N) - - def div(self, a: int, b: int) -> int: - """Exact division a/b. Raises error if b does not divide a.""" - if self._rust is not None: - return self._rust.div(a, b) - a_val, b_val = a % self.N, b % self.N - g, x, _ = self._egcd(b_val, self.N) - if a_val % g != 0: - raise ValueError(f"{a} not divisible by {b} in Z/{self.N}") - return (x * (a_val // g)) % (self.N // g) - - def stab(self, a: int, b: int, c: int) -> int: - if self._rust is not None: - return self._rust.stab(a, b, c) - a, b, c = a % self.N, b % self.N, c % self.N - target = self.gcd(a, self.gcd(b, c)) - for x in range(self.N): - candidate = (a + x * b) % self.N - current = self.gcd(candidate, c) # = gcd(a + x b, c, N) - if current == target: - return x - raise ValueError(f"Stab failed for a={a}, b={b}, c={c} in Z/{self.N}") - - def ass(self, a: int) -> int: - """Associate: gcd(a, N).""" - return math.gcd(a % self.N, self.N) - - def unit(self, a: int) -> int: - """Return u such that u*a = Ass(a).""" - a_val = a % self.N - g, u, _ = self._egcd(a_val, self.N) - return u % self.N - - def _egcd(self, a: int, b: int) -> tuple[int, int, int]: - r0, r1, s0, s1 = a, b, 1, 0 - while r1: - q = r0 // r1 - r0, r1 = r1, r0 - q * r1 - s0, s1 = s1, s0 - q * s1 - return r0, s0, 0 # t not needed here diff --git a/modularsnf/snf.py b/modularsnf/snf.py index 204ce57..80e933e 100644 --- a/modularsnf/snf.py +++ b/modularsnf/snf.py @@ -1,33 +1,23 @@ -"""Utilities for computing Smith normal forms over modular rings.""" +"""Smith Normal Form over Z/NZ via the native Storjohann band reduction. -from numbers import Integral -from typing import NamedTuple, Tuple +A thin wrapper over ``modularsnf._rust.rust_smith_normal_form``; the reduction +itself lives in the Rust ``modularsnf`` crate. For the experimental CRT fast +path see ``modularsnf.crt``. +""" -import numpy as np +from typing import NamedTuple -from modularsnf.band import band_reduction, compute_upper_bandwidth +from modularsnf._rust import rust_smith_normal_form as _rust_snf -from .diagonal import smith_from_diagonal -from .echelon import index1_reduce_on_columns, lemma_3_1 -from .matrix import RingMatrix -from .ring import RingZModN - -try: - from modularsnf._rust import rust_smith_normal_form as _rust_snf -except ImportError: - _rust_snf = None # type: ignore[assignment] +from ._core import to_int64_matrix, validate_modulus class SNFResult(NamedTuple): """Result of a Smith Normal Form decomposition over Z/NZ. - Attributes: - S: The Smith Normal Form (diagonal matrix) as nested lists. - U: Left unimodular transform as nested lists. - V: Right unimodular transform as nested lists. - - The invariant ``S = U @ A @ V`` holds over Z/NZ, where ``@`` - denotes matrix multiplication modulo N. + ``S`` is the diagonal Smith form, ``U`` and ``V`` the unimodular transforms, + each a plain ``list[list[int]]``. The invariant ``S = U @ A @ V`` holds over + Z/NZ. """ S: list[list[int]] @@ -41,10 +31,9 @@ def smith_normal_form_mod( ) -> SNFResult: """Compute the Smith Normal Form of an integer matrix over Z/NZ. - Returns ``(S, U, V)`` where ``S = U @ A @ V`` (mod *modulus*), - *S* is diagonal with the divisibility chain ``s_i | s_{i+1}``, - and *U*, *V* are unimodular. All three are plain Python - ``list[list[int]]``. + Returns ``(S, U, V)`` where ``S = U @ A @ V`` (mod *modulus*), *S* is + diagonal with the divisibility chain ``s_i | s_{i+1}``, and *U*, *V* are + unimodular. All three are plain Python ``list[list[int]]``. Args: matrix: 2-D list of signed 64-bit integers. May be rectangular. @@ -58,489 +47,10 @@ def smith_normal_form_mod( Examples: >>> S, U, V = smith_normal_form_mod([[2, 4], [6, 8]], modulus=12) """ - if isinstance(modulus, bool) or not isinstance(modulus, Integral): - raise ValueError(f"Modulus must be an integer >= 2, got {modulus!r}") - if int(modulus) < 2: - raise ValueError(f"Modulus must be an integer >= 2, got {modulus!r}") - - if not isinstance(matrix, (list, tuple)): - raise TypeError("matrix must be a list of lists of integers") - - # Empty matrix — return empty lists immediately. - if len(matrix) == 0: + modulus = validate_modulus(modulus) + arr = to_int64_matrix(matrix, modulus) + if arr is None: return SNFResult(S=[], U=[], V=[]) - # Validate that all rows have the same length. - ncols = len(matrix[0]) - for i, row in enumerate(matrix): - if len(row) != ncols: - raise ValueError( - f"Ragged matrix: row 0 has {ncols} columns " - f"but row {i} has {len(row)} columns" - ) - - ring = RingZModN(modulus) - A = RingMatrix(ring, matrix) - - U_rm, V_rm, S_rm = smith_normal_form(A) - - return SNFResult( - S=S_rm.data.tolist(), - U=U_rm.data.tolist(), - V=V_rm.data.tolist(), - ) - - -def smith_normal_form( - A: RingMatrix, -) -> Tuple[RingMatrix, RingMatrix, RingMatrix]: - """Computes the Smith normal form over a principal ideal ring. - - For rectangular matrices, the input is zero-padded to square before the - Smith form is computed, then the resulting transforms are cropped back to - the original shape. - - Args: - A: Matrix over a principal ideal ring. - - Returns: - A tuple ``(U, V, S)`` where ``S = U @ A @ V`` is in Smith normal form. - """ - - ring = A.ring - n = A.nrows - m = A.ncols - - # Handle empty matrices by returning identities of matching dimensions. - if n == 0 or m == 0: - U = RingMatrix.identity(ring, n) - V = RingMatrix.identity(ring, m) - return U, V, A.copy() - - if _rust_snf is not None: - u_arr, v_arr, s_arr = _rust_snf( - A.data.astype(np.int64), - ring.N, - ) - return ( - RingMatrix._from_ndarray(ring, u_arr), - RingMatrix._from_ndarray(ring, v_arr), - RingMatrix._from_ndarray(ring, s_arr), - ) - - # If the matrix is already square, proceed without padding. - if n == m: - return _smith_square(A) - - # Pad rectangular matrices to a square ``s x s`` matrix before reduction. - s = max(n, m) - A_pad = A.pad_to(s, s) - - U_pad, V_pad, S_pad = _smith_square(A_pad) - - # Restrict padded transforms back to the original shape. - U = U_pad.submatrix(0, n, 0, n) - V = V_pad.submatrix(0, m, 0, m) - S = S_pad.submatrix(0, n, 0, m) - - return U, V, S - - -def _smith_square(A: RingMatrix) -> Tuple[RingMatrix, RingMatrix, RingMatrix]: - """Computes Smith normal form for a square matrix over a PIR. - - Args: - A: Square matrix. - - Returns: - ``(U, V, S)`` with ``S = U @ A @ V`` diagonal and containing the - invariant factors of ``A``. - - Raises: - ValueError: If ``A`` is not square. - """ - - ring = A.ring - n = A.nrows - if n != A.ncols: - raise ValueError("Expected a square matrix in _smith_square.") - - if n == 0: - I0 = RingMatrix.identity(ring, 0) - return I0, I0, A.copy() - - # Step 1 (Storjohann 7.3): Reduce to row echelon form ``T = U_ech * A``. - U_ech, T, _rank = lemma_3_1(A) - - # Step 2 (Storjohann 7.3): Reduce bandwidth until the matrix is 2-banded. - B = T.copy() - b = compute_upper_bandwidth(B) - - U_band_total = RingMatrix.identity(ring, n) - V_band_total = RingMatrix.identity(ring, n) - - while b > 2: - B, U_step, V_step, b = band_reduction(B, b) - U_band_total = U_step @ U_band_total - V_band_total = V_band_total @ V_step - - U_snf, V_snf, S = smith_from_upper_2_banded(B) - - U_total = U_snf @ U_band_total @ U_ech - V_total = V_band_total @ V_snf - - return U_total, V_total, S - - -def smith_from_upper_2_banded(A: RingMatrix): - """Runs the Smith form pipeline for a 2-banded square matrix. - - Args: - A: Square matrix with bandwidth at most two. - - Returns: - A tuple ``(U, V, S)`` with ``S = U @ A @ V`` in Smith normal form. - - Raises: - ValueError: If ``A`` is not square. - """ - - ring: RingZModN = A.ring - n = A.nrows - if n != A.ncols: - raise ValueError("Expected square matrix") - if n <= 1: - U = RingMatrix.identity(ring, n) - V = RingMatrix.identity(ring, n) - return U, V, A.copy() - - # Step 1 (Storjohann 7.3): Precondition into spike plus principal blocks. - U1, V1, A1, n1, n2 = _step1_split_with_spike(A) - - # Step 2 (Storjohann 7.3): Run recursive Smith forms on both blocks. - U2, V2, A2 = _step2_recursive_blocks(A1, n1, n2) - - # Step 3 (Storjohann 7.3): Permute the spike to the trailing position. - U3, V3, A3 = _step3_permute(A2, n1, n2) - - # Step 4 (Storjohann 7.3): Diagonalize the principal ``(n-1) x (n-1)``. - U4, V4, A4 = _step4_smith_on_n_minus_1(A3) - - # Steps 5–8 (Storjohann 7.3): Build the gcd chain ``s1..sk``. - U5, V5, A5, k = _step5_to_8_gcd_chain(A4) - - # Step 9 (Storjohann 7.3): Reverse first ``k`` rows and run index reduction. - U6, V6, A6 = _step9_index_reduction(A5, k) - - # Compose transforms: U_total, V_total such that U_total A V_total = A6 - U_total = U6 @ U5 @ U4 @ U3 @ U2 @ U1 - V_total = V1 @ V2 @ V3 @ V4 @ V5 @ V6 - - return U_total, V_total, A6 - - -def _step1_split_with_spike( - A: RingMatrix, -) -> Tuple[RingMatrix, RingMatrix, RingMatrix, int, int]: - """Splits a 2-banded matrix into principal, spike, and trailing blocks. - - Args: - A: Square 2-banded matrix. - - Returns: - ``(U, V, A1, n1, n2)`` where ``A1 = U @ A`` has block sizes ``n1`` and - ``n2`` flanking the spike row/column and ``V`` is identity. - """ - - ring = A.ring - n = A.nrows - n1 = (n - 1) // 2 - n2 = n - 1 - n1 - - U = RingMatrix.identity(ring, n) - T = A.copy() - - # Sweep upward to isolate the spike column above ``n1``. - for k in range(n - 1, n1, -1): - r0, r1 = k - 1, k - col = k - a = T.data[r0, col] - b = T.data[r1, col] - g, s, t, u, v = ring.gcdex(a, b) - T.apply_row_2x2(r0, r1, u, v, s, t) - U.apply_row_2x2(r0, r1, u, v, s, t) - - # Sweep downward to clear the spike column below ``n1``. - for k in range(n1 + 1, n - 1): - r0, r1 = k, k + 1 - col = k - a = T.data[r0, col] - b = T.data[r1, col] - g, s, t, u, v = ring.gcdex(a, b) - T.apply_row_2x2(r0, r1, s, t, u, v) - U.apply_row_2x2(r0, r1, s, t, u, v) - - V = RingMatrix.identity(ring, n) - - return U, V, T, n1, n2 - - -def _step2_recursive_blocks(A: RingMatrix, n1: int, n2: int): - """Recursively Smith-forms the principal and trailing diagonal blocks. - - Args: - A: Matrix produced by the spike split. - n1: Size of the top-left block. - n2: Size of the bottom-right block. - - Returns: - ``(U, V, A_new)`` where the principal and trailing blocks are diagonal - but the spike row and column may be dense. - """ - - ring = A.ring - n = A.nrows - - # Extract the independent blocks. - # Block 1: Top-left (0 to n1). - B1 = A.submatrix(0, n1, 0, n1) - - # Block 2: Bottom-right (n1+1 to n) skipping the spike at index ``n1``. - B2 = A.submatrix(n1 + 1, n, n1 + 1, n) - - # Recursively diagonalize both independent blocks. - U1_loc, V1_loc, S1 = smith_from_upper_2_banded(B1) - U2_loc, V2_loc, S2 = smith_from_upper_2_banded(B2) - - # Embed the local transforms with an identity on the spike position. - U = RingMatrix.identity(ring, n) - V = RingMatrix.identity(ring, n) - - U.write_block(0, 0, U1_loc) - U.write_block(n1 + 1, n1 + 1, U2_loc) - - V.write_block(0, 0, V1_loc) - V.write_block(n1 + 1, n1 + 1, V2_loc) - - # Apply to A; S1 and S2 become diagonal while the spike spreads through the - # transforms. - A_new = U @ A @ V - - return U, V, A_new - - -def _step3_permute( - A: RingMatrix, n1: int, n2: int -) -> Tuple[RingMatrix, RingMatrix, RingMatrix]: - """Permutes the spike to the last row/column. - - Args: - A: Matrix from the recursive block step. - n1: Size of the principal block. - n2: Size of the trailing block. - - Returns: - ``(P, P_inv, A3)`` where ``A3 = P @ A @ P_inv`` and the spike sits in - the last row and column. - - Raises: - ValueError: If ``A`` is not square. - """ - - ring = A.ring - n = A.nrows - if n != A.ncols: - raise ValueError("Expected square matrix") - - # Build permutation matrix P as n x n with P[new_i, old_i] = 1. - perm = np.zeros((n, n), dtype=int) - for old_i in range(n): - if old_i < n1: - new_i = old_i - elif old_i == n1: - new_i = n - 1 - else: - new_i = old_i - 1 - perm[new_i, old_i] = 1 - P = RingMatrix._from_ndarray(ring, perm) - P_inv = P.transpose() # permutation matrix => inverse = transpose - - A3 = P @ A @ P_inv - return P, P_inv, A3 - - -def _step4_smith_on_n_minus_1(A: RingMatrix): - """Diagonalizes the top-left ``(n-1) x (n-1)`` block. - - Args: - A: Matrix with the spike in the last row and column. - - Returns: - ``(U, V, A_new)`` where the principal block is diagonalized. - """ - - ring = A.ring - n = A.nrows - assert n == A.ncols - - # Principal ``(n-1) x (n-1)`` block. - B = A.submatrix(0, n - 1, 0, n - 1) - U_loc, V_loc, S_loc = smith_from_diagonal(B) - - # Embed the local transforms with a fixed spike position. - U = RingMatrix.identity(ring, n) - V = RingMatrix.identity(ring, n) - U.write_block(0, 0, U_loc) - V.write_block(0, 0, V_loc) - - A_new = U @ A @ V - return U, V, A_new - - -def _step5_to_8_gcd_chain( - A: RingMatrix, -) -> Tuple[RingMatrix, RingMatrix, RingMatrix, int]: - """Builds the gcd chain that drives the final Smith reduction. - - Args: - A: Matrix with a diagonal principal block and dense last column. - - Returns: - ``(U, V, A_new, k)`` where ``k`` is the 1-based start index of the - dense block as in Storjohann 7.3. - """ - - ring = A.ring - n = A.nrows - U = RingMatrix.identity(ring, n) - V = RingMatrix.identity(ring, n) - T = A.copy() - - idx_k = n - 1 - for i in range(n - 1): - if ring.is_zero(T.data[i, i]): - idx_k = i - break - - last_col = n - 1 - - # Step 5 (Storjohann 7.3): Eliminate entries below ``idx_k`` in last column. - for t in range(idx_k + 1, n): - pivot_val = T.data[idx_k, last_col] - target_val = T.data[t, last_col] - - if ring.is_zero(target_val): - continue - - # Compute extended gcd coefficients on ``(pivot, target)``. - g, s, x, u, v = ring.gcdex(pivot_val, target_val) - - # Apply the 2x2 transform to rows ``idx_k`` and ``t``. - T.apply_row_2x2(idx_k, t, s, x, u, v) - U.apply_row_2x2(idx_k, t, s, x, u, v) - - # Step 6 (Storjohann 7.3): Swap the dense column into position ``idx_k``. - - if idx_k != last_col: - P_swap = RingMatrix.identity(ring, n) - # Swap columns ``idx_k`` and ``last_col`` in the identity to build the - # right multiplication matrix. - - P_swap.data[idx_k, idx_k] = 0 - P_swap.data[last_col, last_col] = 0 - P_swap.data[idx_k, last_col] = 1 - P_swap.data[last_col, idx_k] = 1 - - T = T @ P_swap - V = V @ P_swap - - # Step 7 (Storjohann 7.3): Perform the stabilize loop (ripple up). - col_target = idx_k - - for i in range(idx_k - 1, -1, -1): - a_ik = T.data[i, col_target] - a_i1k = T.data[i + 1, col_target] - a_ii = T.data[i, i] # The diagonal element s_i - - c = ring.stab(a_ik, a_i1k, a_ii) - - s_next = T.data[i + 1, i + 1] - numerator = ring.mul(c, s_next) - q_raw = ring.quo(numerator, a_ii) - q = (-q_raw) % ring.N - - # Operation 1: Add ``c * Row[i+1]`` to ``Row[i]``. - T.apply_row_2x2(i, i + 1, 1, c, 0, 1) - U.apply_row_2x2(i, i + 1, 1, c, 0, 1) - - # Operation 2: Add ``q * Col[i]`` to ``Col[i+1]`` via right transform. - N = ring.N - T.data[:, i + 1] = (T.data[:, i + 1] + q * T.data[:, i]) % N - V.data[:, i + 1] = (V.data[:, i + 1] + q * V.data[:, i]) % N - - # Step 8 (Storjohann 7.3): Run the gcd reduction loop (ripple down). - - for i in range(idx_k): - pivot = T.data[i, i] - target = T.data[i, col_target] # col_target is idx_k - - if ring.is_zero(target): - continue - - g, s, t, u, v = ring.gcdex(pivot, target) - - N = ring.N - ci = T.data[:, i].copy() - ck = T.data[:, col_target].copy() - T.data[:, i] = (s * ci + t * ck) % N - T.data[:, col_target] = (u * ci + v * ck) % N - - vi = V.data[:, i].copy() - vk = V.data[:, col_target].copy() - V.data[:, i] = (s * vi + t * vk) % N - V.data[:, col_target] = (u * vi + v * vk) % N - - # Return idx_k + 1 (1-based) to match paper's "k" for Step 9 - return U, V, T, (idx_k + 1) - - -def _step9_index_reduction( - A: RingMatrix, k: int -) -> Tuple[RingMatrix, RingMatrix, RingMatrix]: - """Performs the final index-1 reduction on the dense block. - - Args: - A: Matrix returned from the gcd chain step. - k: 1-based index where the dense block starts. - - Returns: - ``(U_final, V_final, A_final)`` where ``A_final`` is the Smith form of - the input matrix. - """ - - ring = A.ring - n = A.nrows - idx_k = k - 1 - - # Build ``P`` that reverses the first ``k`` indices. - P_arr = np.zeros((n, n), dtype=int) - for i in range(n): - if i <= idx_k: - P_arr[i, idx_k - i] = 1 - else: - P_arr[i, i] = 1 - P = RingMatrix._from_ndarray(ring, P_arr) - - # Conjugate by ``P`` to move the dense block to the top-left. - A_perm = P @ A @ P - - # Index-1 reduction on first ``k`` columns (row operations only). - U_red, A_red = index1_reduce_on_columns(A_perm, k) - - # Net effect is ``S = (P U_red P) A``. - U_final = P @ U_red @ P - V_final = RingMatrix.identity(ring, n) - - A_final = U_final @ A # no right transform - - return U_final, V_final, A_final + u, v, s = _rust_snf(arr, modulus) + return SNFResult(S=s.tolist(), U=u.tolist(), V=v.tolist()) From 079d53e4d09075ce073945fea93ca94151fd529d Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Sat, 6 Jun 2026 00:42:59 -0400 Subject: [PATCH 20/22] test: replace Python correctness suite with PyO3 boundary tests Algorithm correctness now lives in the Rust crate. The Python suite (test_binding.py) covers only the binding: input validation, type/shape marshalling, and one end-to-end SymPy oracle on the invariant factors. Drop the old correctness/internal suites and the python/rust backend switch (conftest, helpers, pyproject perf marker, justfile). --- justfile | 5 +- pyproject.toml | 3 +- tests/conftest.py | 121 ------------ tests/helpers.py | 112 ----------- tests/test_binding.py | 149 ++++++++++++++ tests/test_echelon.py | 78 -------- tests/test_matrix.py | 28 --- tests/test_oracle.py | 226 --------------------- tests/test_perf.py | 111 ----------- tests/test_regression.py | 413 --------------------------------------- tests/test_ring.py | 102 ---------- tests/test_snf.py | 163 --------------- 12 files changed, 152 insertions(+), 1359 deletions(-) delete mode 100644 tests/conftest.py delete mode 100644 tests/helpers.py create mode 100644 tests/test_binding.py delete mode 100644 tests/test_echelon.py delete mode 100644 tests/test_matrix.py delete mode 100644 tests/test_oracle.py delete mode 100644 tests/test_perf.py delete mode 100644 tests/test_regression.py delete mode 100644 tests/test_ring.py delete mode 100644 tests/test_snf.py diff --git a/justfile b/justfile index 9d0beb7..0940311 100644 --- a/justfile +++ b/justfile @@ -5,11 +5,10 @@ lint: uv run ruff check . typecheck: - uv run ty check modularsnf/ scripts/profile_snf.py tests/conftest.py + uv run ty check modularsnf/ tests/ test: - uv run pytest tests/ -x --backend python - uv run pytest tests/ -x --backend rust + uv run pytest tests/ -x # Verify Cargo workspace and pyproject.toml versions match check-version: diff --git a/pyproject.toml b/pyproject.toml index 7935b0f..7e83fbf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,11 +22,10 @@ dev = ["maturin>=1.7,<2", "pytest", "pytest-cov", "ruff", "ty", "sympy"] [tool.pytest.ini_options] minversion = "6.0" -addopts = "-ra -q -m 'not perf' --cov=modularsnf --cov-report=term-missing" +addopts = "-ra -q --cov=modularsnf --cov-report=term-missing" testpaths = [ "tests", ] -markers = ["perf: performance sanity checks (run with: uv run pytest -m perf)"] [tool.maturin] manifest-path = "crates/modularsnf-py/Cargo.toml" diff --git a/tests/conftest.py b/tests/conftest.py deleted file mode 100644 index 4cffa04..0000000 --- a/tests/conftest.py +++ /dev/null @@ -1,121 +0,0 @@ -"""Shared fixtures for the modularsnf test suite.""" - -import random -from collections.abc import Iterator -from typing import Any - -import numpy as np -import pytest - -import modularsnf.diagonal as diagonal_mod -import modularsnf.ring as ring_mod -import modularsnf.snf as snf_mod -from modularsnf.matrix import RingMatrix -from modularsnf.ring import RingZModN - -_RustRing: Any = None -_rust_diag: Any = None -_rust_merge: Any = None -_rust_snf: Any = None - -try: - from modularsnf._rust import RustRingZModN as _RustRing - from modularsnf._rust import rust_merge_smith_blocks as _rust_merge - from modularsnf._rust import rust_smith_from_diagonal as _rust_diag - from modularsnf._rust import rust_smith_normal_form as _rust_snf -except ImportError: - pass - - -def pytest_addoption(parser: pytest.Parser) -> None: - """Register custom command-line options.""" - parser.addoption( - "--backend", - action="store", - choices=("auto", "python", "rust"), - default="auto", - help="Execution backend for tests", - ) - - -def pytest_configure(config: pytest.Config) -> None: - """Register custom markers.""" - config.addinivalue_line( - "markers", - "perf: performance sanity checks (deselect with -m 'not perf')", - ) - - -@pytest.fixture(scope="session", autouse=True) -def configured_backend(pytestconfig: pytest.Config) -> Iterator[str]: - """Configure the requested execution backend for the full test session.""" - has_rust = all( - value is not None - for value in (_RustRing, _rust_diag, _rust_merge, _rust_snf) - ) - requested = pytestconfig.getoption("backend") - - if requested == "rust": - if not has_rust: - pytest.exit( - "Rust backend requested for tests but modularsnf._rust is " - "unavailable" - ) - backend = "rust" - elif requested == "python": - backend = "python" - elif has_rust: - backend = "rust" - else: - backend = "python" - - old_ring = getattr(ring_mod, "_RustRing") - old_diag = getattr(diagonal_mod, "_rust_diag") - old_merge = getattr(diagonal_mod, "_rust_merge") - old_snf = getattr(snf_mod, "_rust_snf") - - if backend == "rust": - setattr(ring_mod, "_RustRing", _RustRing) - setattr(diagonal_mod, "_rust_diag", _rust_diag) - setattr(diagonal_mod, "_rust_merge", _rust_merge) - setattr(snf_mod, "_rust_snf", _rust_snf) - else: - setattr(ring_mod, "_RustRing", None) - setattr(diagonal_mod, "_rust_diag", None) - setattr(diagonal_mod, "_rust_merge", None) - setattr(snf_mod, "_rust_snf", None) - - yield backend - - setattr(ring_mod, "_RustRing", old_ring) - setattr(diagonal_mod, "_rust_diag", old_diag) - setattr(diagonal_mod, "_rust_merge", old_merge) - setattr(snf_mod, "_rust_snf", old_snf) - - -@pytest.fixture -def seeded_rng(request: pytest.FixtureRequest) -> int: - """Seed both stdlib random and numpy RNG. - - The seed is extracted from ``request.param`` when used with - indirect parametrization, or defaults to 42. - - Returns the seed value for diagnostic printing. - """ - seed = getattr(request, "param", 42) - random.seed(seed) - np.random.seed(seed) - return seed - - -def make_random_matrix( - ring: RingZModN, - nrows: int, - ncols: int, -) -> RingMatrix: - """Generate a random matrix over the given ring.""" - data = [ - [random.randint(0, ring.N - 1) for _ in range(ncols)] - for _ in range(nrows) - ] - return RingMatrix.from_rows(ring, data) diff --git a/tests/helpers.py b/tests/helpers.py deleted file mode 100644 index 77403f8..0000000 --- a/tests/helpers.py +++ /dev/null @@ -1,112 +0,0 @@ -import math -from itertools import product - -import numpy as np - -from modularsnf.matrix import RingMatrix - - -def det_ring_matrix(M: RingMatrix) -> int: - """ - Naive determinant for small square matrices over RingZModN. - Returns an integer representative (mod N is implied by the ring). - """ - ring = M.ring - data = M.data - n = M.nrows - assert n == M.ncols - - if n == 1: - return int(data[0, 0] % ring.N) - if n == 2: - a, b = data[0] - c, d = data[1] - ad = ring.mul(a, d) - bc = ring.mul(b, c) - return ring.sub(ad, bc) - - det = 0 - for j in range(n): - sub_rows = np.concatenate((data[1:, :j], data[1:, j + 1 :]), axis=1) - subM = RingMatrix(ring, sub_rows) - sub_det = det_ring_matrix(subM) - term = ring.mul(data[0, j], sub_det) - if j % 2 == 0: - det = ring.add(det, term) - else: - det = ring.sub(det, term) - return det - - -def verify_echelon_structure(T: RingMatrix) -> bool: - """ - Check: - - For each non-zero row, the first non-zero column index strictly increases. - - Once a zero row appears, all later rows are zero. - """ - ring = T.ring - nrows, ncols = T.nrows, T.ncols - last_pivot_col = -1 - zero_row_seen = False - - for r in range(nrows): - row = T.data[r] - # find first non-zero entry in this row - pivot_col = -1 - for c in range(ncols): - if not ring.is_zero(row[c]): - pivot_col = c - break - - if pivot_col == -1: - # zero row - zero_row_seen = True - # all subsequent rows must be zero - if any( - not all(ring.is_zero(x) for x in T.data[rr]) - for rr in range(r + 1, nrows) - ): - return False - else: - # non-zero row; we must not have seen a zero row before - if zero_row_seen: - return False - if pivot_col <= last_pivot_col: - return False - last_pivot_col = pivot_col - - return True - - -def row_span(M: RingMatrix) -> set[tuple[int, ...]]: - """ - Compute the full row span of M over Z/NZ by brute force. - Only for small matrices (e.g., r <= 3) in tests. - """ - ring = M.ring - N = ring.N - r, c = M.nrows, M.ncols - rows = M.data - - span = set() - for coeffs in product(range(N), repeat=r): - vec = np.zeros(c, dtype=int) - for i, alpha in enumerate(coeffs): - if alpha == 0: - continue - row = rows[i] - vec = (vec + (alpha * row)) % N - span.add(tuple(int(x % N) for x in vec)) - return span - - -def get_normalized_invariants(M: RingMatrix) -> list[int]: - """ - Extracts diagonal entries and normalizes them to ideal generators. - invariant = gcd(d_i, N). - """ - n = min(M.nrows, M.ncols) - diags = np.diag(M.data[:n, :n]) - invariants = [math.gcd(int(d), M.ring.N) for d in diags] - invariants.sort() - return invariants diff --git a/tests/test_binding.py b/tests/test_binding.py new file mode 100644 index 0000000..06b1f70 --- /dev/null +++ b/tests/test_binding.py @@ -0,0 +1,149 @@ +"""PyO3 boundary tests for the modularsnf Python API. + +Exhaustive algorithm correctness lives in the Rust crate (cargo test). These +tests cover what the Python layer is responsible for — input validation and +type/shape marshalling across the extension boundary — plus one end-to-end +oracle that validates the full stack against SymPy, an independent +implementation. +""" + +import math + +import numpy as np +import pytest +import sympy as sp +from sympy import ZZ +from sympy.matrices.normalforms import smith_normal_form as sympy_snf + +from modularsnf import crt_snf, smith_normal_form_mod +from modularsnf.crt import factorize + +ALGORITHMS = [smith_normal_form_mod, crt_snf] + + +def _check_contract(A, modulus, S, U, V): + """S, U, V are list[list[int]] with the documented shapes and S = U A V.""" + rows = len(A) + cols = len(A[0]) if rows else 0 + assert np.array(U).shape == (rows, rows) + assert np.array(V).shape == (cols, cols) + assert np.array(S).shape == (rows, cols) + # Returned as plain Python ints, not numpy scalars. + assert all(isinstance(x, int) for row in S for x in row) + prod = ( + np.array(U, object) @ np.array(A, object) @ np.array(V, object) + ) % modulus + assert np.array_equal(prod, np.array(S, object) % modulus) + + +@pytest.mark.parametrize("fn", ALGORITHMS) +@pytest.mark.parametrize("rows,cols", [(3, 3), (4, 6), (6, 4), (5, 1), (1, 5)]) +def test_shapes_and_contract(fn, rows, cols): + modulus = 12 + A = [[(i * 7 + j * 3) % modulus for j in range(cols)] for i in range(rows)] + S, U, V = fn(A, modulus) + _check_contract(A, modulus, S, U, V) + + +@pytest.mark.parametrize("fn", ALGORITHMS) +def test_negative_entries_reduced(fn): + """Negative inputs are reduced into [0, N) by the binding.""" + S, U, V = fn([[-1, -5], [-7, 2]], 12) + assert all(0 <= x < 12 for row in S for x in row) + _check_contract([[-1, -5], [-7, 2]], 12, S, U, V) + + +@pytest.mark.parametrize("fn", ALGORITHMS) +def test_tuple_input_accepted(fn): + S, U, V = fn(((2, 4), (6, 8)), 12) + assert np.array(S).shape == (2, 2) + + +@pytest.mark.parametrize("fn", ALGORITHMS) +def test_empty_matrix(fn): + assert fn([], 12) == ([], [], []) + + +@pytest.mark.parametrize("fn", ALGORITHMS) +@pytest.mark.parametrize("modulus", [1, 0, -3]) +def test_modulus_too_small(fn, modulus): + with pytest.raises(ValueError): + fn([[1, 2], [3, 4]], modulus) + + +@pytest.mark.parametrize("fn", ALGORITHMS) +@pytest.mark.parametrize("modulus", [True, 3.5]) +def test_modulus_must_be_int(fn, modulus): + with pytest.raises(ValueError): + fn([[1]], modulus) + + +@pytest.mark.parametrize("fn", ALGORITHMS) +def test_modulus_out_of_int64(fn): + with pytest.raises(OverflowError): + fn([[1, 2], [3, 4]], 1 << 70) + + +@pytest.mark.parametrize("fn", ALGORITHMS) +def test_non_list_input(fn): + with pytest.raises(TypeError): + fn(42, 12) + + +@pytest.mark.parametrize("fn", ALGORITHMS) +def test_ragged_rows(fn): + with pytest.raises(ValueError): + fn([[1, 2, 3], [4, 5]], 12) + + +@pytest.mark.parametrize("fn", ALGORITHMS) +def test_entry_out_of_int64(fn): + with pytest.raises(OverflowError): + fn([[1 << 70]], 12) + + +def test_factorize(): + assert factorize(36) == [(2, 2), (3, 2)] + assert factorize(30) == [(2, 1), (3, 1), (5, 1)] + assert factorize(17) == [(17, 1)] + + +def test_crt_explicit_factors_match_auto(): + A = [[2, 4, 0], [6, 8, 3], [0, 3, 9]] + assert crt_snf(A, 36).S == crt_snf(A, 36, factors=[(2, 2), (3, 2)]).S + + +def test_crt_bad_factors_rejected(): + with pytest.raises(ValueError): + crt_snf([[1]], 36, factors=[(2, 1), (3, 2)]) # = 18, not 36 + + +def test_crt_matches_default_on_small_case(): + """Sanity cross-check at the boundary (full parity is tested in Rust).""" + A = [[2, 4, 0], [6, 8, 3], [0, 3, 9]] + assert crt_snf(A, 36).S == smith_normal_form_mod(A, 36).S + + +def _normalized_invariants(diag, modulus): + """gcd of each diagonal entry with the modulus, ascending.""" + r = min(len(diag), len(diag[0])) if diag else 0 + return sorted(math.gcd(int(diag[i][i]), modulus) for i in range(r)) + + +def _sympy_invariants(A, modulus): + """Invariant factors of the integer SNF, projected into Z/NZ.""" + D = sympy_snf(sp.Matrix(A), domain=ZZ) + r = min(D.rows, D.cols) + return sorted(math.gcd(int(D[i, i]), modulus) for i in range(r)) + + +@pytest.mark.parametrize("fn", ALGORITHMS) +@pytest.mark.parametrize("modulus", [6, 12, 36, 8, 30, 100]) +@pytest.mark.parametrize("rows,cols", [(3, 3), (4, 6), (6, 4), (5, 2), (1, 4)]) +def test_invariants_match_sympy(fn, modulus, rows, cols): + """The full stack agrees with SymPy on the invariant factors over Z/NZ.""" + rng = np.random.default_rng(modulus * 1000 + rows * 10 + cols) + A = rng.integers(0, modulus, size=(rows, cols)).tolist() + assert _normalized_invariants( + fn(A, modulus).S, modulus + ) == _sympy_invariants(A, modulus) diff --git a/tests/test_echelon.py b/tests/test_echelon.py deleted file mode 100644 index 36f0add..0000000 --- a/tests/test_echelon.py +++ /dev/null @@ -1,78 +0,0 @@ -import math -import random - -import numpy as np - -import pytest -from modularsnf.ring import RingZModN -from modularsnf.matrix import RingMatrix -from modularsnf.echelon import lemma_3_1 -from tests.helpers import det_ring_matrix, row_span, verify_echelon_structure - -RINGS_TO_TEST = [2, 5, 9, 12] - -@pytest.mark.parametrize("N", RINGS_TO_TEST) -def test_lemma_3_1_equation_and_row_module(N): - ring = RingZModN(N) - - nrows, ncols = 5, 5 - A_rows = [ - [random.randint(0, N - 1) for _ in range(ncols)] - for _ in range(nrows) - ] - A = RingMatrix.from_rows(ring, A_rows) - - U, T, _ = lemma_3_1(A) - - UA = U @ A - assert np.array_equal(UA.data, T.data) - - span_A = row_span(A) - span_T = row_span(T) - assert span_A == span_T - -@pytest.mark.parametrize("N", RINGS_TO_TEST) -def test_lemma_3_1_unimodular_U(N): - ring = RingZModN(N) - - size = 5 - A_rows = [ - [random.randint(0, N - 1) for _ in range(size)] - for _ in range(size) - ] - A = RingMatrix.from_rows(ring, A_rows) - - U, T, _ = lemma_3_1(A) - - det_U = det_ring_matrix(U) % N - assert math.gcd(det_U, N) == 1, f"det(U)={det_U} not a unit mod {N}" - -@pytest.mark.parametrize("N", RINGS_TO_TEST) -def test_lemma_3_1_echelon_structure(N): - ring = RingZModN(N) - - examples = [] - - examples.append([ - [2, 4, 6], - [1, 2, 3], - [0, 1, 5], - [0, 0, 0], - ]) - - for _ in range(3): - nrows, ncols = 4, 5 - rows = [ - [random.randint(0, N - 1) for _ in range(ncols)] - for _ in range(nrows) - ] - examples.append(rows) - - for rows in examples: - A = RingMatrix.from_rows(ring, rows) - _, T, _ = lemma_3_1(A) - - assert verify_echelon_structure(T), ( - f"T is not in echelon form over Z/{N}: {T.data}" - ) - diff --git a/tests/test_matrix.py b/tests/test_matrix.py deleted file mode 100644 index d6675d0..0000000 --- a/tests/test_matrix.py +++ /dev/null @@ -1,28 +0,0 @@ -import numpy as np -import pytest - -from modularsnf.matrix import RingMatrix, _is_within_modulus -from modularsnf.ring import RingZModN - - -def test_normalize_matrix_rejects_out_of_range_integers(): - """Matrix entries must fit in the signed 64-bit contract.""" - - ring = RingZModN(70) - huge = 10**100 - - with pytest.raises(OverflowError, match="signed 64-bit"): - RingMatrix.from_rows(ring, [[huge, 1], [2, 3]]) - - -def test_is_within_modulus_checks_value_range(): - ring = RingZModN(11) - assert _is_within_modulus(np.array([[0, 5, 10]], dtype=int), ring.N) - assert not _is_within_modulus(np.array([[12, -1]], dtype=int), ring.N) - - -def test_post_init_applies_mod_when_needed(): - ring = RingZModN(5) - matrix = RingMatrix(ring, [[-1, 9]]) - - assert np.array_equal(matrix.data, np.array([[4, 4]])) diff --git a/tests/test_oracle.py b/tests/test_oracle.py deleted file mode 100644 index 5e3d725..0000000 --- a/tests/test_oracle.py +++ /dev/null @@ -1,226 +0,0 @@ -"""Oracle comparison tests: modularsnf vs SymPy integer-domain SNF. - -For each test case we: -1. Build a random integer matrix A over Z/NZ. -2. Compute SNF via the modular pipeline -> (U, V, S). -3. Verify four structural properties: - a. S = U @ A @ V (transform validity) - b. S is diagonal (diagonal structure) - c. divisibility chain (d_i | d_{i+1} as ideals) - d. U, V unimodular (det is unit mod N) -4. Compute integer-domain SNF via SymPy -> project to Z/N. -5. Compare invariant factors. -""" - -import math -import random - -import numpy as np -import pytest - -sympy = pytest.importorskip("sympy") -from sympy import ZZ # noqa: E402 -from sympy.matrices.normalforms import ( # noqa: E402 - smith_normal_form as sympy_snf, -) - -from modularsnf.matrix import RingMatrix -from modularsnf.ring import RingZModN -from modularsnf.snf import smith_normal_form as compute_snf -from modularsnf.snf import smith_normal_form_mod -from tests.helpers import det_ring_matrix, get_normalized_invariants - - -def _make_random_matrix( - ring: RingZModN, - nrows: int, - ncols: int, -) -> RingMatrix: - data = [ - [random.randint(0, ring.N - 1) for _ in range(ncols)] - for _ in range(nrows) - ] - return RingMatrix.from_rows(ring, data) - - -# --- Parameter matrices --- - -_ORACLE_PARAMS = [ - pytest.param( - N, - shape, - seed, - id=f"N={N}-{shape[0]}x{shape[1]}-seed={seed}", - ) - for N in [2, 3, 5, 7, 4, 8, 9, 16, 6, 12, 30, 64, 97] - for shape in [(4, 4), (6, 6), (5, 8), (8, 5)] - for seed in [42, 137, 2025] -] - -_PUBLIC_API_PARAMS = [ - pytest.param( - N, - shape, - seed, - id=f"api-N={N}-{shape[0]}x{shape[1]}-seed={seed}", - ) - for N in [6, 12, 8, 97] - for shape in [(4, 4), (5, 3)] - for seed in [42, 2025] -] - -_ZERO_DIVISOR_PARAMS = [ - pytest.param( - N, - seed, - id=f"zerodiv-N={N}-seed={seed}", - ) - for N in [4, 8, 16, 32, 64] - for seed in [42, 137, 2025, 7919, 31337] -] - - -class TestOracleInternal: - """Oracle tests against the internal compute_snf API.""" - - @pytest.mark.parametrize("N, shape, seed", _ORACLE_PARAMS) - def test_structural_properties( - self, - N: int, - shape: tuple[int, int], - seed: int, - ) -> None: - random.seed(seed) - ring = RingZModN(N) - nrows, ncols = shape - A = _make_random_matrix(ring, nrows, ncols) - - U, V, S = compute_snf(A) - - # Property 1: S = U @ A @ V - assert np.array_equal( - (U @ A @ V).data, - S.data, - ), "Transform mismatch" - - # Property 2: S is diagonal - for r in range(S.nrows): - for c in range(S.ncols): - if r != c: - assert ring.is_zero(S.data[r, c]), ( - f"Off-diagonal ({r},{c}) = {S.data[r, c]}" - ) - - # Property 3: Divisibility chain - invariants = get_normalized_invariants(S) - for i in range(len(invariants) - 1): - d_curr, d_next = invariants[i], invariants[i + 1] - if d_curr == 0: - assert d_next == 0 - else: - assert d_next % d_curr == 0, ( - f"Chain break: {d_curr} !| {d_next}" - ) - - # Property 4: U, V unimodular (square inputs only; - # rectangular crops lose unimodularity). - if nrows == ncols: - det_U = det_ring_matrix(U) % N - assert math.gcd(det_U, N) == 1 - det_V = det_ring_matrix(V) % N - assert math.gcd(det_V, N) == 1 - - @pytest.mark.parametrize("N, shape, seed", _ORACLE_PARAMS) - def test_invariants_match_sympy( - self, - N: int, - shape: tuple[int, int], - seed: int, - ) -> None: - random.seed(seed) - ring = RingZModN(N) - nrows, ncols = shape - A = _make_random_matrix(ring, nrows, ncols) - - # SymPy oracle: integer-domain SNF -> project to Z/N - A_sym = A.to_sympy() - S_sym = sympy_snf(A_sym, domain=ZZ) - S_arr = np.array(S_sym.tolist(), dtype=int) - sym_diags: list[int] = S_arr.diagonal().tolist() - expected = sorted(math.gcd(d, N) for d in sym_diags) - - # Our pipeline - _, _, S = compute_snf(A) - actual = get_normalized_invariants(S) - - assert actual == expected, ( - f"Invariant mismatch for N={N}, shape={shape}, " - f"seed={seed}\n" - f" Expected (SymPy): {expected}\n" - f" Actual (ours): {actual}" - ) - - -class TestOraclePublicAPI: - """Oracle tests against the public smith_normal_form_mod API.""" - - @pytest.mark.parametrize("N, shape, seed", _PUBLIC_API_PARAMS) - def test_public_api_structural( - self, - N: int, - shape: tuple[int, int], - seed: int, - ) -> None: - random.seed(seed) - nrows, ncols = shape - matrix = [ - [random.randint(0, N - 1) for _ in range(ncols)] - for _ in range(nrows) - ] - - result = smith_normal_form_mod(matrix, modulus=N) - - # Verify return type - assert isinstance(result.S, list) - assert isinstance(result.U, list) - assert isinstance(result.V, list) - - # Reconstruct as RingMatrix for property checks - ring = RingZModN(N) - S = RingMatrix.from_rows(ring, result.S) - U = RingMatrix.from_rows(ring, result.U) - V = RingMatrix.from_rows(ring, result.V) - A = RingMatrix.from_rows(ring, matrix) - - # S = U @ A @ V - assert np.array_equal((U @ A @ V).data, S.data) - - # S is diagonal - for r in range(S.nrows): - for c in range(S.ncols): - if r != c: - assert ring.is_zero(S.data[r, c]) - - -class TestHighZeroDivisor: - """Stress tests for moduli with many zero divisors.""" - - @pytest.mark.parametrize("N, seed", _ZERO_DIVISOR_PARAMS) - def test_power_of_2_moduli( - self, - N: int, - seed: int, - ) -> None: - random.seed(seed) - ring = RingZModN(N) - A = _make_random_matrix(ring, 6, 6) - - U, V, S = compute_snf(A) - - assert np.array_equal((U @ A @ V).data, S.data) - - invariants = get_normalized_invariants(S) - for i in range(len(invariants) - 1): - d_curr, d_next = invariants[i], invariants[i + 1] - if d_curr != 0: - assert d_next % d_curr == 0 diff --git a/tests/test_perf.py b/tests/test_perf.py deleted file mode 100644 index b862cb2..0000000 --- a/tests/test_perf.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Performance sanity checks for the modular SNF pipeline. - -These tests verify that runtime does not regress catastrophically. -They use generous wall-clock bounds and are marked ``perf`` so they -are excluded from the default test run. - -Run with: uv run pytest -m perf -""" - -import random -import time - -import numpy as np -import pytest - -from modularsnf.matrix import RingMatrix -from modularsnf.ring import RingZModN -from modularsnf.snf import smith_normal_form as compute_snf - - -@pytest.mark.perf -class TestPerformanceSanity: - """Wall-clock sanity checks for representative sizes.""" - - # Bounds are set generously (5x measured baseline) to - # account for CI variability and slow runners. - CASES = [ - pytest.param(12, 10, 5.0, id="10x10-mod12"), - pytest.param(100, 15, 15.0, id="15x15-mod100"), - pytest.param(8, 20, 30.0, id="20x20-mod8"), - pytest.param(30, 20, 30.0, id="20x20-mod30"), - ] - - @pytest.mark.parametrize("N, size, max_seconds", CASES) - def test_runtime_bound( - self, - N: int, - size: int, - max_seconds: float, - ) -> None: - random.seed(42) - ring = RingZModN(N) - data = [ - [random.randint(0, N - 1) for _ in range(size)] for _ in range(size) - ] - A = RingMatrix.from_rows(ring, data) - - t0 = time.perf_counter() - U, V, S = compute_snf(A) - elapsed = time.perf_counter() - t0 - - # Verify correctness so timing doesn't mask a bug - assert np.array_equal((U @ A @ V).data, S.data) - - assert elapsed < max_seconds, ( - f"{size}x{size} mod {N} took {elapsed:.2f}s " - f"(limit {max_seconds:.1f}s)" - ) - - def test_rectangular_perf(self) -> None: - """Rectangular matrices should not be dramatically slower.""" - random.seed(42) - N = 12 - ring = RingZModN(N) - nrows, ncols = 15, 8 - data = [ - [random.randint(0, N - 1) for _ in range(ncols)] - for _ in range(nrows) - ] - A = RingMatrix.from_rows(ring, data) - - t0 = time.perf_counter() - U, V, S = compute_snf(A) - elapsed = time.perf_counter() - t0 - - assert np.array_equal((U @ A @ V).data, S.data) - - assert elapsed < 15.0, f"15x8 mod 12 took {elapsed:.2f}s (limit 15.0s)" - - def test_scaling_ratio(self) -> None: - """Check that doubling size does not cause >16x slowdown. - - If the algorithm is O(n^3), doubling n should give ~8x. - We allow up to 32x to account for constant factors and - CI noise, but catch O(n^4) or worse regressions. - """ - N = 12 - ring = RingZModN(N) - - def time_snf(size: int) -> float: - random.seed(42) - data = [ - [random.randint(0, N - 1) for _ in range(size)] - for _ in range(size) - ] - A = RingMatrix.from_rows(ring, data) - t0 = time.perf_counter() - compute_snf(A) - return time.perf_counter() - t0 - - t_small = time_snf(8) - t_large = time_snf(16) - - # Guard against division by near-zero - if t_small < 0.001: - return # Too fast to measure meaningfully - - ratio = t_large / t_small - assert ratio < 32.0, ( - f"Scaling ratio 8->16: {ratio:.1f}x (expected <32x for O(n^3))" - ) diff --git a/tests/test_regression.py b/tests/test_regression.py deleted file mode 100644 index faf86e9..0000000 --- a/tests/test_regression.py +++ /dev/null @@ -1,413 +0,0 @@ -"""Curated regression corpus for SNF edge cases. - -Each case targets a historically problematic matrix class: -- repeated factors in the modulus -- rank-deficient blocks -- dense random with adversarial structure -- near-diagonal with adversarial superdiagonal entries -- rectangular shapes that exercise padding/cropping -- matrices that trigger the merge_smith_blocks permutation path -""" - -import math -import random - -import numpy as np -import pytest - -from modularsnf.matrix import RingMatrix -from modularsnf.ring import RingZModN -from modularsnf.snf import smith_normal_form as compute_snf -from modularsnf.snf import smith_normal_form_mod -from tests.helpers import det_ring_matrix, get_normalized_invariants - - -def _assert_valid_snf( - A: RingMatrix, - U: RingMatrix, - V: RingMatrix, - S: RingMatrix, - *, - expected_invariants: list[int] | None = None, -) -> None: - """Assert all four SNF structural properties. - - Optionally compares invariant factors against known values. - """ - ring = A.ring - N = ring.N - - # 1. Transform validity - assert np.array_equal( - (U @ A @ V).data, - S.data, - ), "S != U @ A @ V" - - # 2. Diagonal structure - for r in range(S.nrows): - for c in range(S.ncols): - if r != c: - assert ring.is_zero(S.data[r][c]), ( - f"Off-diagonal ({r},{c}) = {S.data[r][c]}" - ) - - # 3. Divisibility chain - inv = get_normalized_invariants(S) - for i in range(len(inv) - 1): - if inv[i] != 0: - assert inv[i + 1] % inv[i] == 0, ( - f"Chain break: {inv[i]} !| {inv[i + 1]}" - ) - - # 4. Unimodularity (square inputs only) - if A.nrows == A.ncols: - if U.nrows == U.ncols and U.nrows > 0: - assert math.gcd(det_ring_matrix(U) % N, N) == 1 - if V.nrows == V.ncols and V.nrows > 0: - assert math.gcd(det_ring_matrix(V) % N, N) == 1 - - # 5. Optional exact invariant check - if expected_invariants is not None: - assert inv == sorted(expected_invariants), ( - f"Invariants {inv} != expected {sorted(expected_invariants)}" - ) - - -class TestRepeatedFactors: - """Matrices whose entries share factors with the modulus.""" - - CASES = [ - pytest.param( - [[2, 4], [6, 8]], - 12, - [2, 4], - id="2x2-all-even-mod12", - ), - pytest.param( - [[3, 6, 9], [6, 3, 6], [9, 6, 3]], - 9, - [3, 3, 3], - id="3x3-multiples-of-3-mod9", - ), - pytest.param( - [[4, 0, 0], [0, 4, 0], [0, 0, 4]], - 8, - [4, 4, 4], - id="3x3-scalar-4-mod8", - ), - pytest.param( - [[2, 0], [0, 4]], - 8, - [2, 4], - id="2x2-diag-powers-of-2-mod8", - ), - pytest.param( - [[6, 0, 0], [0, 10, 0], [0, 0, 15]], - 30, - [1, 30, 30], - id="3x3-diag-mod30-coprime-factors", - ), - ] - - @pytest.mark.parametrize("matrix, N, expected_inv", CASES) - def test_repeated_factors( - self, - matrix: list[list[int]], - N: int, - expected_inv: list[int], - ) -> None: - ring = RingZModN(N) - A = RingMatrix.from_rows(ring, matrix) - U, V, S = compute_snf(A) - _assert_valid_snf( - A, - U, - V, - S, - expected_invariants=expected_inv, - ) - - -class TestRankDeficient: - """Rank-deficient matrices targeting merge_smith_blocks paths.""" - - CASES = [ - pytest.param( - [[0, 0, 0], [0, 0, 0], [0, 0, 1]], - 12, - [1, 12, 12], - id="3x3-rank1-mod12", - ), - pytest.param( - [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 6, 0], [0, 0, 0, 4]], - 12, - [2, 12, 12, 12], - id="4x4-rank2-sparse-mod12", - ), - pytest.param( - [[2, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 6]], - 12, - [2, 6, 12, 12], - id="4x4-two-nonzero-diag-mod12", - ), - pytest.param( - [[0, 0], [0, 0]], - 7, - [7, 7], - id="2x2-zero-mod-prime", - ), - pytest.param( - [[1, 2, 3], [2, 4, 6], [3, 6, 9]], - 12, - None, - id="3x3-rank1-dense-mod12", - ), - ] - - @pytest.mark.parametrize("matrix, N, expected_inv", CASES) - def test_rank_deficient( - self, - matrix: list[list[int]], - N: int, - expected_inv: list[int] | None, - ) -> None: - ring = RingZModN(N) - A = RingMatrix.from_rows(ring, matrix) - U, V, S = compute_snf(A) - _assert_valid_snf( - A, - U, - V, - S, - expected_invariants=expected_inv, - ) - - -class TestDenseRandom: - """Seeded dense random matrices exercising the full pipeline.""" - - @pytest.mark.parametrize( - "N, size, seed", - [ - pytest.param(12, 6, 42, id="dense-6x6-mod12-seed42"), - pytest.param(8, 7, 137, id="dense-7x7-mod8-seed137"), - pytest.param(100, 5, 2025, id="dense-5x5-mod100-seed2025"), - pytest.param(64, 8, 7919, id="dense-8x8-mod64-seed7919"), - pytest.param(30, 6, 31337, id="dense-6x6-mod30-seed31337"), - ], - ) - def test_dense_random( - self, - N: int, - size: int, - seed: int, - ) -> None: - random.seed(seed) - ring = RingZModN(N) - data = [ - [random.randint(0, N - 1) for _ in range(size)] for _ in range(size) - ] - A = RingMatrix.from_rows(ring, data) - U, V, S = compute_snf(A) - _assert_valid_snf(A, U, V, S) - - -class TestAdversarialSuperdiagonal: - """Diagonal + adversarial superdiagonal entries.""" - - CASES = [ - pytest.param( - [[1, 11, 0], [0, 1, 0], [0, 0, 1]], - 12, - [1, 1, 1], - id="3x3-identity-plus-superdiag-mod12", - ), - pytest.param( - [[2, 1, 0, 0], [0, 4, 1, 0], [0, 0, 6, 1], [0, 0, 0, 8]], - 12, - None, - id="4x4-staircase-mod12", - ), - pytest.param( - [[3, 0, 0, 5], [0, 3, 0, 0], [0, 0, 3, 0], [0, 0, 0, 3]], - 9, - None, - id="4x4-scalar3-plus-corner-mod9", - ), - pytest.param( - [[4, 2, 0], [0, 4, 2], [0, 0, 4]], - 8, - None, - id="3x3-upper-triangular-mod8", - ), - pytest.param( - [ - [1, 0, 0, 0, 7], - [0, 1, 0, 0, 0], - [0, 0, 1, 0, 0], - [0, 0, 0, 1, 0], - [0, 0, 0, 0, 1], - ], - 12, - [1, 1, 1, 1, 1], - id="5x5-identity-with-far-corner-mod12", - ), - ] - - @pytest.mark.parametrize("matrix, N, expected_inv", CASES) - def test_adversarial_superdiagonal( - self, - matrix: list[list[int]], - N: int, - expected_inv: list[int] | None, - ) -> None: - ring = RingZModN(N) - A = RingMatrix.from_rows(ring, matrix) - U, V, S = compute_snf(A) - _assert_valid_snf( - A, - U, - V, - S, - expected_invariants=expected_inv, - ) - - -class TestRectangularRegression: - """Rectangular matrices exercising padding/cropping paths.""" - - CASES = [ - pytest.param( - [[1, 2, 3, 4, 5]], - 12, - id="1x5-single-row-mod12", - ), - pytest.param( - [[1], [2], [3], [4], [5]], - 12, - id="5x1-single-col-mod12", - ), - pytest.param( - [[2, 4, 6], [8, 10, 0]], - 12, - id="2x3-fat-mod12", - ), - pytest.param( - [[3, 6], [9, 3], [6, 9]], - 9, - id="3x2-tall-mod9", - ), - ] - - @pytest.mark.parametrize("matrix, N", CASES) - def test_rectangular( - self, - matrix: list[list[int]], - N: int, - ) -> None: - ring = RingZModN(N) - A = RingMatrix.from_rows(ring, matrix) - U, V, S = compute_snf(A) - _assert_valid_snf(A, U, V, S) - - -class TestErrorHandling: - """Exercise snf.py input validation paths.""" - - def test_modulus_too_small(self) -> None: - with pytest.raises(ValueError, match="Modulus must be"): - smith_normal_form_mod([[1]], modulus=1) - - def test_modulus_zero(self) -> None: - with pytest.raises(ValueError, match="Modulus must be"): - smith_normal_form_mod([[1]], modulus=0) - - def test_modulus_negative(self) -> None: - with pytest.raises(ValueError, match="Modulus must be"): - smith_normal_form_mod([[1]], modulus=-5) - - def test_modulus_out_of_int64_range(self) -> None: - with pytest.raises(OverflowError, match="signed 64-bit"): - smith_normal_form_mod([[1]], modulus=10**30) - - def test_non_list_input(self) -> None: - with pytest.raises(TypeError, match="list of lists"): - smith_normal_form_mod("not a matrix", modulus=5) # type: ignore[arg-type] - - def test_ragged_rows(self) -> None: - with pytest.raises(ValueError, match="Ragged matrix"): - smith_normal_form_mod([[1, 2], [3]], modulus=5) - - def test_empty_matrix(self) -> None: - result = smith_normal_form_mod([], modulus=5) - assert result.S == [] - assert result.U == [] - assert result.V == [] - - def test_matrix_entry_out_of_int64_range(self) -> None: - with pytest.raises(OverflowError, match="signed 64-bit"): - smith_normal_form_mod([[10**30]], modulus=5) - - -class TestCoverageTargeted: - """Tests targeting specific uncovered code paths.""" - - def test_merge_smith_blocks_permutation_path(self) -> None: - """Target diagonal.py lines 221-249. - - diag(0, 6, 0, 4) mod 12 forces merge of blocks with - partial ranks, triggering the permutation path in - merge_smith_blocks. - """ - ring = RingZModN(12) - A = RingMatrix.diagonal(ring, [0, 6, 0, 4]) - U, V, S = compute_snf(A) - _assert_valid_snf( - A, - U, - V, - S, - expected_invariants=[2, 12, 12, 12], - ) - - def test_merge_smith_blocks_larger_rank_deficient( - self, - ) -> None: - """Larger rank-deficient case for merge path.""" - ring = RingZModN(8) - A = RingMatrix.diagonal(ring, [0, 0, 2, 0, 4, 0]) - U, V, S = compute_snf(A) - _assert_valid_snf(A, U, V, S) - - def test_project_to_upper_bandwidth(self) -> None: - """Target band.py project_to_upper_bandwidth (dead code).""" - from modularsnf.band import project_to_upper_bandwidth - - ring = RingZModN(12) - data = [ - [1, 2, 3, 4], - [5, 6, 7, 8], - [9, 10, 11, 0], - [1, 2, 3, 4], - ] - M = RingMatrix.from_rows(ring, data) - P = project_to_upper_bandwidth(M, 2) - - # Only entries with 0 <= j - i < 2 should survive - for i in range(4): - for j in range(4): - if 0 <= j - i < 2: - assert P.data[i][j] == M.data[i][j] - else: - assert P.data[i][j] == 0 - - def test_band_reduction_already_2banded(self) -> None: - """band_reduction with b <= 2 is a no-op.""" - from modularsnf.band import band_reduction - - ring = RingZModN(12) - A = RingMatrix.diagonal(ring, [1, 2, 3]) - A_red, U, V, b_new = band_reduction(A, 2) - - assert b_new == 2 - assert np.array_equal(A_red.data, A.data) diff --git a/tests/test_ring.py b/tests/test_ring.py deleted file mode 100644 index 4369ceb..0000000 --- a/tests/test_ring.py +++ /dev/null @@ -1,102 +0,0 @@ -import math - -import pytest - -from modularsnf.ring import RingZModN - -RINGS_TO_TEST = [2, 5, 9, 12] - - -@pytest.mark.parametrize("N", RINGS_TO_TEST) -class TestRing: - @pytest.fixture - def ring(self, N): - return RingZModN(N) - - def test_additive_properties(self, ring, N): - """Verify (a + b) mod N for ALL pairs.""" - for a in range(N): - for b in range(N): - res = ring.add(a, b) - assert res == (a + b) % N - assert ring.sub(res, b) == a - - def test_multiplicative_properties(self, ring, N): - """Verify (a * b) mod N for ALL pairs.""" - for a in range(N): - for b in range(N): - assert ring.mul(a, b) == (a * b) % N - - def test_gcdex_properties(self, ring, N): - """ - For ALL pairs (a,b), verify Gcdex output satisfies: - 1. [s t; u v][a; b] = [g; 0] - 2. det(M) is a unit - """ - for a in range(N): - for b in range(N): - g, s, t, u, v = ring.gcdex(a, b) - - # Check 1: Matrix action - row1 = ring.add(ring.mul(s, a), ring.mul(t, b)) - row2 = ring.add(ring.mul(u, a), ring.mul(v, b)) - assert row1 == g - assert row2 == 0 - - # Check 2: Unimodular - det = ring.sub(ring.mul(s, v), ring.mul(t, u)) - assert ring.gcd(det, N) == 1 - - def test_annihilator(self, ring, N): - """Verify Ann(a) * a = 0 for ALL a.""" - for a in range(N): - x = ring.ann(a) - assert ring.mul(a, x) == 0 - - def test_division_exactness(self, ring, N): - """ - For ALL pairs (a,b), if b|a, verify div(a,b) works. - """ - for b in range(N): - for k in range(N): - a = (b * k) % N - # If b|a is mathematically true in the ring - # (represented by being in the ideal generated by b) - # Note: Our div checks if a is in ideal(b), i.e., gcd(b,N)|a - if a % math.gcd(b, N) == 0: - q = ring.div(a, b) - # Verify q*b = a - assert ring.mul(q, b) == a - - def test_stabilizer(self, ring, N): - """ - For ALL triplets (a,b,c), verify Stab returns x such that - gcd(a+xb, c) = gcd(a, b, c). - """ - # This is O(N^3), safe for N=12 (1728 iterations) - for a in range(N): - for b in range(N): - for c in range(N): - x = ring.stab(a, b, c) - lhs = ring.gcd(a + x * b, c) - rhs = ring.gcd(a, ring.gcd(b, c)) - - # Note: gcd(a,b,c) in ring is gcd(a, gcd(b, gcd(c, N))) - # But since ring.gcd includes N implicitly, we compare outputs - assert lhs == rhs - - def test_rem_quo_consistency(self, ring, N): - """Verify a = q*b + r for ALL pairs.""" - for a in range(N): - for b in range(N): - r = ring.rem(a, b) - q = ring.quo(a, b) - - # Check consistency: a = q*b + r - rhs = ring.add(ring.mul(q, b), r) - assert a == rhs - - -def test_modulus_must_fit_int64() -> None: - with pytest.raises(OverflowError, match="signed 64-bit"): - RingZModN(10**30) diff --git a/tests/test_snf.py b/tests/test_snf.py deleted file mode 100644 index aea92f7..0000000 --- a/tests/test_snf.py +++ /dev/null @@ -1,163 +0,0 @@ -import math -import random - -import numpy as np -import pytest - -from modularsnf.matrix import RingMatrix -from modularsnf.ring import RingZModN -from modularsnf.snf import smith_normal_form as compute_snf -from tests.helpers import det_ring_matrix, get_normalized_invariants - -# --- Test Configuration --- -RINGS_TO_TEST = [2, 4, 6, 9, 12, 100] -DIMENSIONS = [ - (4, 4), # Square - (4, 6), # Fat - (6, 4), # Tall - (5, 1), # Vector column -] - - -class TestSmithNormalForm: - @pytest.fixture - def random_matrix(self, request): - """Fixture to generate random matrices based on params.""" - N, rows, cols = request.param - ring = RingZModN(N) - data = [ - [random.randint(0, N - 1) for _ in range(cols)] for _ in range(rows) - ] - return RingMatrix.from_rows(ring, data) - - @pytest.mark.parametrize( - "random_matrix", - [(N, r, c) for N in [6, 12] for r, c in DIMENSIONS], - indirect=True, - ) - def test_snf_transform_validity(self, random_matrix): - """ - Verify U * A * V == S. - """ - A = random_matrix - U, V, S = compute_snf(A) - - assert U.shape == (A.nrows, A.nrows) - assert V.shape == (A.ncols, A.ncols) - assert S.shape == A.shape - - LHS = U @ A @ V - assert np.array_equal(LHS.data, S.data), ( - "Transform mismatch: U @ A @ V != S" - ) - - @pytest.mark.parametrize( - "random_matrix", - [(N, r, c) for N in [6, 12] for r, c in DIMENSIONS], - indirect=True, - ) - def test_snf_diagonal_structure(self, random_matrix): - """ - Verify S_ij == 0 for all i != j. - """ - A = random_matrix - _, _, S = compute_snf(A) - ring = S.ring - - for r in range(S.nrows): - for c in range(S.ncols): - if r != c: - assert ring.is_zero(S.data[r][c]), ( - f"Non-zero off-diagonal at ({r},{c}): {S.data[r][c]}" - ) - - @pytest.mark.parametrize( - "random_matrix", [(N, 5, 5) for N in [12, 36, 100]], indirect=True - ) - def test_snf_divisibility_chain(self, random_matrix): - """ - Verify d_i | d_{i+1} in the sense of principal ideals. - In Z/N, this means gcd(d_i, N) | gcd(d_{i+1}, N). - """ - A = random_matrix - _, _, S = compute_snf(A) - - invariants = get_normalized_invariants(S) - - for i in range(len(invariants) - 1): - d_curr = invariants[i] - d_next = invariants[i + 1] - - if d_curr == 0: - assert d_next == 0, ( - f"Divisibility break: 0 does not divide {d_next}" - ) - else: - assert d_next % d_curr == 0, ( - f"Divisibility break at index {i}: {d_curr} does not divide {d_next}" - ) - - @pytest.mark.parametrize( - "random_matrix", [(N, 4, 4) for N in [6, 9]], indirect=True - ) - def test_transforms_are_unimodular(self, random_matrix): - """ - Verify det(U) and det(V) are units in Z/N. - i.e., gcd(det, N) == 1. - """ - A = random_matrix - N = A.ring.N - U, V, _ = compute_snf(A) - - det_U = det_ring_matrix(U) - det_V = det_ring_matrix(V) - - assert math.gcd(det_U, N) == 1, ( - f"U is not unimodular. det(U)={det_U} (mod {N})" - ) - assert math.gcd(det_V, N) == 1, ( - f"V is not unimodular. det(V)={det_V} (mod {N})" - ) - - # Oracle comparison tests (formerly test_invariants_against_sympy) - # have moved to tests/test_oracle.py with broader coverage. - - def test_zero_matrix(self): - """Test strict zero matrix.""" - ring = RingZModN(12) - A = RingMatrix.from_rows(ring, [[0, 0], [0, 0]]) - U, V, S = compute_snf(A) - - assert np.array_equal(S.data, np.zeros((2, 2), dtype=int)) - assert np.array_equal((U @ A @ V).data, S.data) - - def test_identity_matrix(self): - """Test already diagonal matrix (identity).""" - ring = RingZModN(12) - A = RingMatrix.identity(ring, 3) - U, V, S = compute_snf(A) - - for i in range(S.nrows): - for j in range(S.ncols): - if i != j: - assert ring.is_zero(S.data[i][j]), ( - f"Non-zero off-diagonal at ({i},{j}): {S.data[i][j]}" - ) - - invariants = get_normalized_invariants(S) - assert invariants == [1, 1, 1] - - def test_already_diagonal_unsorted(self): - """ - Test a diagonal matrix that violates divisibility chain. - diag(2, 1) mod 4 -> should become diag(1, 2). - """ - ring = RingZModN(4) - A = RingMatrix.diagonal(ring, [2, 1]) - - _, _, S = compute_snf(A) - - invariants = get_normalized_invariants(S) - assert invariants == [1, 2] - assert S.data[0][0] == 1 or S.data[0][0] == 3 - assert S.data[1][1] == 2 From d7b662a2d2f0565aec83507889732f949b97249b Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Sat, 6 Jun 2026 00:43:05 -0400 Subject: [PATCH 21/22] docs: position CRT as experimental fast path; reflect thin wrapper Storjohann band reduction stays the documented default; add a CRT fast-path section pointing at docs/crt.md. Update README/AGENTS to describe the Rust-crate-plus-thin-binding layout and that correctness is validated in the Rust suite. --- AGENTS.md | 38 ++++++++++++++++++++------------------ README.md | 34 +++++++++++++++++++++------------- docs/algorithm.md | 4 ++++ 3 files changed, 45 insertions(+), 31 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index fd012fe..90268e9 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -34,25 +34,27 @@ Use `modularsnf.ring` primitives: `gcdex`, `div`, `quo`, `stab`, `ann`. ## Testing -* Structural assertions (`S = U A V`, diagonal shape, divisibility chain, - unimodularity) over hardcoded array equality. -* Validate against SymPy integer-domain projection where feasible. -* Run `just check` (or individually: `just lint`, `just typecheck`, `just test`) before finalizing. +* Algorithm correctness lives in the Rust crate (`cargo test`): structural + assertions (`S = U A V`, diagonal shape, divisibility chain, unimodularity) + over random inputs, plus a Storjohann-vs-CRT cross-check. +* The Python suite (`tests/`) covers only the PyO3 boundary: input validation, + type/shape marshalling, and the public-API contract on small cases. +* Run `cargo test` and `just check` (`just lint`, `just typecheck`, `just test`) + before finalizing. ## File Structure -* `modularsnf/ring.py` — ring primitives (Gcdex, Stab, Div, Ann). +The SNF algorithms live in the Rust workspace; the Python package is a thin +wrapper over the `modularsnf._rust` extension. + +* `crates/modularsnf/` — Rust lib crate: ring primitives, echelon/band + reduction, diagonalization (`snf.rs`/`band.rs`/`echelon.rs`/`diagonal.rs`), + and the CRT fast path (`crt.rs`). +* `crates/modularsnf-py/` — PyO3 bindings exposed as `modularsnf._rust`. +* `modularsnf/ring.py` — `RingZModN`, forwarding ring primitives to Rust. * `modularsnf/matrix.py` — `RingMatrix` data structure. -* `modularsnf/echelon.py` — triangularization (Lemma 3.1). -* `modularsnf/band.py` — band reduction (Lemmas 7.3, 7.4, Prop 7.1). -* `modularsnf/diagonal.py` — diagonalization (Prop 7.7, Theorem 7.11). -* `modularsnf/snf.py` — master pipeline (Lemma 7.14) + public - `smith_normal_form_mod` API. -* `docs/algorithm.md` — mathematical foundation and algorithm description. -* `docs/PLAN.md` — current ExecPlan. - -## ExecPlans - -For work spanning multiple modules or changing public APIs, write a plan in -`docs/PLAN.md` before implementing. Users may say "use an ExecPlan" as -shorthand. +* `modularsnf/diagonal.py` — diagonalization wrappers. +* `modularsnf/snf.py` — public `smith_normal_form_mod` API (default path). +* `modularsnf/crt.py` — public `crt_snf` API (experimental fast path). +* `docs/algorithm.md` — default algorithm (Storjohann band reduction). +* `docs/crt.md` — CRT fast path design and correctness basis. diff --git a/README.md b/README.md index c34b21c..ea4a712 100644 --- a/README.md +++ b/README.md @@ -1,10 +1,10 @@ # Smith Normal form of Integer matrices mod N (Storjohann) -This is a Python module that follows the deterministic algorithms presented in Arne Storjohann's PhD Dissertation *Algorithms for Matrix Canonical Forms* (ETH No. 13922, 2000). +A Rust implementation (with a thin Python binding) of the deterministic algorithms presented in Arne Storjohann's PhD Dissertation *Algorithms for Matrix Canonical Forms* (ETH No. 13922, 2000). -It implements the Lemmas and subsequent subroutines that are necessary for calculating the SNF without exponential intermediate values. +It implements the Lemmas and subsequent subroutines that are necessary for calculating the SNF without exponential intermediate values. The algorithms live in the Rust `modularsnf` crate; the Python package is a thin wrapper over the native extension. -It validates against SymPy using a known equivalence between calculating the Smith Normal form of an integer matrix, and then taking mod N, compared to solving it natively in the ring. +Correctness is validated in the Rust test suite: random inputs are checked against the Smith form contract ($S = UAV$, divisibility chain, unimodular transforms), and the default Storjohann path is cross-checked against the independent CRT path. ## Quick Start @@ -28,22 +28,30 @@ A matrix is **unimodular** over $\mathbb{Z}/N\mathbb{Z}$ when $\gcd(\det(M),\, N) = 1$, the modular analogue of $|\det(U)| = 1$ over $\mathbb{Z}$. -### Lower-level API +For details on the default algorithm (band reduction, diagonalization, +Storjohann's lemmas), see [docs/algorithm.md](docs/algorithm.md). -For direct access to `RingMatrix` objects: +## Alternative: CRT fast path (experimental) + +`smith_normal_form_mod` uses Storjohann's band reduction, which works for any +modulus `N >= 2`. For the **small-modulus, large-matrix** regime there is a +second, experimental algorithm that is several times faster: a CRT-based path +that factors `N = prod p^e`, solves the SNF over each local ring `Z/p^e` by +valuation-pivoted elimination, and recombines via the Chinese Remainder Theorem. ```python -from modularsnf.ring import RingZModN -from modularsnf.matrix import RingMatrix -from modularsnf.snf import smith_normal_form +from modularsnf import crt_snf -ring = RingZModN(12) -A = RingMatrix(ring, [[2, 4], [6, 8]]) -U, V, S = smith_normal_form(A) # note: (U, V, S) order +S, U, V = crt_snf([[2, 4, 0], + [6, 8, 3], + [0, 3, 9]], modulus=36) +# Same (S, U, V) contract as smith_normal_form_mod; S = U @ A @ V (mod 36). ``` -For details on the algorithm (band reduction, diagonalization, Storjohann's -lemmas), see [docs/algorithm.md](docs/algorithm.md). +It requires the prime factorization of `N`; pass it explicitly as +`factors=[(p, e), ...]` to amortize factoring across many calls with the same +modulus. Prefer `smith_normal_form_mod` for general moduli. See +[docs/crt.md](docs/crt.md) for the design and correctness basis. ## Development Workflow (modern Python) diff --git a/docs/algorithm.md b/docs/algorithm.md index 7d4eede..89256a0 100644 --- a/docs/algorithm.md +++ b/docs/algorithm.md @@ -1,5 +1,9 @@ # Mathematical Foundation +This document describes the **default** SNF algorithm (Storjohann band +reduction), used by `smith_normal_form_mod`. For the alternative CRT-based fast +path for small moduli, see [crt.md](crt.md). + The algorithm has two main phases: 1. **Band Reduction:** Transforming an arbitrary matrix into an upper bi-diagonal (2-banded) matrix. 2. **Diagonalization:** Transforming the bi-diagonal matrix into the canonical Smith Normal Form. From 8b517e3601aa445c96f8e1a13932c275f5a7fd26 Mon Sep 17 00:00:00 2001 From: Steven Nguyen <03stevennguyen@gmail.com> Date: Sat, 6 Jun 2026 00:57:57 -0400 Subject: [PATCH 22/22] build: bump version to 0.5.0 --- Cargo.lock | 4 ++-- Cargo.toml | 2 +- pyproject.toml | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e9fcd04..aeb64f1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -67,7 +67,7 @@ dependencies = [ [[package]] name = "modularsnf" -version = "0.4.0" +version = "0.5.0" dependencies = [ "ndarray", "rand", @@ -75,7 +75,7 @@ dependencies = [ [[package]] name = "modularsnf-py" -version = "0.4.0" +version = "0.5.0" dependencies = [ "modularsnf", "ndarray", diff --git a/Cargo.toml b/Cargo.toml index ca5b9fc..88e6e78 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,7 @@ members = ["crates/modularsnf", "crates/modularsnf-py"] resolver = "2" [workspace.package] -version = "0.4.0" +version = "0.5.0" [profile.release] debug = true diff --git a/pyproject.toml b/pyproject.toml index 7e83fbf..e44dcd5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "modularsnf" -version = "0.4.0" +version = "0.5.0" description = "A Python module for Smith Normal Form calculations over Z/N rings" readme = "README.md" license = "Apache-2.0"