zsfm_tabfm/ensemble/
nnls.rs

1//! Lawson-Hanson active-set NNLS (`scipy.optimize.nnls`'s algorithm): solves
2//! `min ||Ax - b||^2` subject to `x >= 0`. Self-contained (no external linear-algebra crate) —
3//! columns are few (`n_estimators <= 32`), so dense Gaussian elimination on the small
4//! passive-set normal-equations system is more than adequate.
5
6/// `a` is column-major: `a[j]` is the `j`-th column (length `m`, matching `b`'s length).
7/// Returns the `n = a.len()`-length non-negative solution.
8pub fn nnls(a: &[Vec<f64>], b: &[f64]) -> Vec<f64> {
9    let n = a.len();
10    if n == 0 {
11        return vec![];
12    }
13    let m = b.len();
14    let mut x = vec![0.0f64; n];
15    let mut passive: Vec<usize> = Vec::new();
16    let mut active: Vec<usize> = (0..n).collect();
17
18    const TOL: f64 = 1e-10;
19    let max_iter = 3 * n + 10;
20
21    for _ in 0..max_iter {
22        // w = A^T (b - A x)
23        let residual: Vec<f64> = {
24            let ax = matvec(a, &x, m);
25            (0..m).map(|i| b[i] - ax[i]).collect()
26        };
27        let w: Vec<f64> = active.iter().map(|&j| dot(&a[j], &residual)).collect();
28
29        let Some((best_pos, &best_w)) = w.iter().enumerate().max_by(|a, b| a.1.partial_cmp(b.1).unwrap()) else {
30            break;
31        };
32        if active.is_empty() || best_w <= TOL {
33            break;
34        }
35        let j = active.remove(best_pos);
36        passive.push(j);
37
38        loop {
39            let z_passive = solve_normal_equations(&passive.iter().map(|&j| a[j].clone()).collect::<Vec<_>>(), b);
40            if z_passive.iter().all(|&v| v > TOL) {
41                for (idx, &j) in passive.iter().enumerate() {
42                    x[j] = z_passive[idx];
43                }
44                break;
45            }
46            // Feasibility step: find alpha shrinking x toward z_passive without crossing 0.
47            let mut alpha = f64::INFINITY;
48            for (idx, &j) in passive.iter().enumerate() {
49                if z_passive[idx] <= TOL {
50                    let denom = x[j] - z_passive[idx];
51                    if denom > 0.0 {
52                        alpha = alpha.min(x[j] / denom);
53                    }
54                }
55            }
56            if !alpha.is_finite() {
57                alpha = 0.0;
58            }
59            for (idx, &j) in passive.iter().enumerate() {
60                x[j] += alpha * (z_passive[idx] - x[j]);
61            }
62            // Move near-zero passive indices back to active.
63            let mut still_passive = Vec::new();
64            for &j in &passive {
65                if x[j] <= TOL {
66                    x[j] = 0.0;
67                    active.push(j);
68                } else {
69                    still_passive.push(j);
70                }
71            }
72            passive = still_passive;
73            if passive.is_empty() {
74                break;
75            }
76        }
77    }
78    x
79}
80
81fn matvec(a: &[Vec<f64>], x: &[f64], m: usize) -> Vec<f64> {
82    let mut out = vec![0.0f64; m];
83    for (j, col) in a.iter().enumerate() {
84        if x[j] == 0.0 {
85            continue;
86        }
87        for i in 0..m {
88            out[i] += col[i] * x[j];
89        }
90    }
91    out
92}
93
94fn dot(a: &[f64], b: &[f64]) -> f64 {
95    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
96}
97
98/// Solves `argmin_z ||A_cols z - b||^2` (unconstrained) via the normal equations
99/// `(A^T A) z = A^T b`, Gaussian elimination with partial pivoting.
100fn solve_normal_equations(a_cols: &[Vec<f64>], b: &[f64]) -> Vec<f64> {
101    let k = a_cols.len();
102    if k == 0 {
103        return vec![];
104    }
105    let mut ata = vec![vec![0.0f64; k]; k];
106    let mut atb = vec![0.0f64; k];
107    for i in 0..k {
108        atb[i] = dot(&a_cols[i], b);
109        for j in 0..k {
110            ata[i][j] = dot(&a_cols[i], &a_cols[j]);
111        }
112        ata[i][i] += 1e-10; // ridge for numerical stability on near-singular passive sets
113    }
114    gaussian_solve(&mut ata, &mut atb)
115}
116
117fn gaussian_solve(a: &mut [Vec<f64>], b: &mut [f64]) -> Vec<f64> {
118    let n = b.len();
119    for col in 0..n {
120        // Partial pivot.
121        let mut pivot_row = col;
122        let mut pivot_val = a[col][col].abs();
123        for row in (col + 1)..n {
124            if a[row][col].abs() > pivot_val {
125                pivot_val = a[row][col].abs();
126                pivot_row = row;
127            }
128        }
129        if pivot_val < 1e-14 {
130            continue; // singular column; leave corresponding solution as 0 later
131        }
132        a.swap(col, pivot_row);
133        b.swap(col, pivot_row);
134        let diag = a[col][col];
135        for row in (col + 1)..n {
136            let factor = a[row][col] / diag;
137            if factor == 0.0 {
138                continue;
139            }
140            for k in col..n {
141                a[row][k] -= factor * a[col][k];
142            }
143            b[row] -= factor * b[col];
144        }
145    }
146    let mut x = vec![0.0f64; n];
147    for i in (0..n).rev() {
148        if a[i][i].abs() < 1e-14 {
149            x[i] = 0.0;
150            continue;
151        }
152        let mut sum = b[i];
153        for j in (i + 1)..n {
154            sum -= a[i][j] * x[j];
155        }
156        x[i] = sum / a[i][i];
157    }
158    x
159}
160
161/// The wrapper's ensemble-weight post-processing: normalize NNLS weights to sum to 1 (fallback
162/// to uniform if the sum is 0), then blend with uniform via `nnls_beta`.
163pub fn finalize_weights(raw_weights: &[f64], nnls_beta: f64) -> Vec<f64> {
164    let n = raw_weights.len();
165    let sum: f64 = raw_weights.iter().sum();
166    let normalized: Vec<f64> = if sum > 0.0 {
167        raw_weights.iter().map(|&w| w / sum).collect()
168    } else {
169        vec![1.0 / n as f64; n]
170    };
171    let uniform = 1.0 / n as f64;
172    normalized.iter().map(|&w| nnls_beta * w + (1.0 - nnls_beta) * uniform).collect()
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178
179    #[test]
180    fn test_nnls_exact_solution() {
181        // A = [[1,0],[0,1],[1,1]], b = [1,2,3] -> exact solution x=[1,2]
182        let a = vec![vec![1.0, 0.0, 1.0], vec![0.0, 1.0, 1.0]];
183        let b = vec![1.0, 2.0, 3.0];
184        let x = nnls(&a, &b);
185        assert!((x[0] - 1.0).abs() < 1e-6, "x0={}", x[0]);
186        assert!((x[1] - 2.0).abs() < 1e-6, "x1={}", x[1]);
187    }
188
189    #[test]
190    fn test_nnls_enforces_nonnegativity() {
191        // A single column strongly anti-correlated with b should be driven to 0, not negative.
192        let a = vec![vec![-1.0, -1.0, -1.0]];
193        let b = vec![1.0, 1.0, 1.0];
194        let x = nnls(&a, &b);
195        assert!(x[0] >= 0.0);
196    }
197
198    #[test]
199    fn test_nnls_matches_scipy() {
200        // Generated via: rng = np.random.default_rng(3); A = rng.uniform(-1,1,(10,4));
201        // b = rng.uniform(0,1,10); scipy.optimize.nnls(A, b).
202        let a = vec![
203            vec![-0.8287016657127513, -0.8117427155192016, 0.46915430281842907, -0.1387439591716444, -0.4315976725024171, -0.9970198329823277, 0.7834221408903144, -0.9393079846750576, 0.3210001348557896, -0.4036738186851505],
204            vec![-0.5263789868078006, -0.1337461195270524, -0.7726559601571932, 0.17359714287628147, 0.2970944141596501, 0.9469205495328255, 0.17032587978181613, 0.413930191311247, 0.862927709482709, 0.48351336013866075],
205            vec![0.6025489304127938, -0.04189740371833195, -0.21754361900867591, 0.4756755745843204, 0.39243199334031087, -0.4031975539662487, -0.057380669636337256, -0.2515123330430584, -0.5856176638379975, 0.44432961628423495],
206            vec![0.16432407212873557, -0.6805221707258429, 0.03348036524272735, 0.9125345096721971, -0.41455850197502575, -0.3720279959313264, 0.5465540192976328, -0.8182945729914843, 0.26018039957068595, -0.5625691508623909],
207        ];
208        let b = vec![0.8298868742743123, 0.6576522108732432, 0.6827989078603502, 0.820075750170535, 0.42857290429846195, 0.758705461154919, 0.8784801846662539, 0.1023199219220744, 0.8497683374661538, 0.39392733263233515];
209        let expected = [0.0, 0.4483466391927005, 0.283691444258291, 0.19757999070784782];
210        let x = nnls(&a, &b);
211        for (g, e) in x.iter().zip(expected.iter()) {
212            assert!((g - e).abs() < 1e-4, "got {g} vs expected {e}");
213        }
214    }
215
216    #[test]
217    fn test_finalize_weights_sums_to_one_when_beta_1() {
218        let w = finalize_weights(&[1.0, 3.0], 1.0);
219        assert!((w[0] + w[1] - 1.0).abs() < 1e-9);
220        assert!((w[0] - 0.25).abs() < 1e-9);
221    }
222
223    #[test]
224    fn test_finalize_weights_uniform_when_beta_0() {
225        let w = finalize_weights(&[1.0, 3.0], 0.0);
226        assert!((w[0] - 0.5).abs() < 1e-9);
227        assert!((w[1] - 0.5).abs() < 1e-9);
228    }
229}
zsfm_tabfm/ensemble/nnls.rs

zsfm_tabfm/ensemble/
nnls.rs
