gam_solve/reml/reml_outer_engine/penalty_coordinate.rs
1use super::*;
2pub use gam_problem::PenaltyCoordinate;
3
4/// Exact pseudo-logdeterminant log|S|₊ and its derivatives with respect to ρ.
5///
6/// # Exact pseudo-logdet on the positive eigenspace
7///
8/// For S(ρ) = Σ exp(ρ_k) S_k with S_k ⪰ 0, the nullspace
9/// N(S) = ∩_k N(S_k) is structurally fixed (independent of ρ).
10/// No eigenvalue of S crosses zero during optimization, so the
11/// pseudo-logdet L = Σ_{σ_i > 0} log σ_i is C∞ in ρ.
12///
13/// ## Computation
14///
15/// Eigendecompose S, identify positive eigenvalues σ_i > ε (where ε is a
16/// relative threshold for numerical zero detection), then:
17///
18/// L(S) = Σ_{positive} log σ_i
19/// ∂_k L = tr(S⁺ A_k) where A_k = λ_k S_k
20/// ∂²_kl L = δ_{kl} ∂_k L − tr(S⁺ A_l S⁺ A_k)
21///
22/// S⁺ is the Moore-Penrose pseudoinverse restricted to the positive
23/// eigenspace. These are the exact derivatives of L — no δ-regularization,
24/// no nullity metadata, no chain-rule inconsistencies.
25#[derive(Clone, Debug)]
26pub struct PenaltyLogdetDerivs {
27 /// L(S) = log|S|₊ — the exact pseudo-logdeterminant on the positive eigenspace.
28 ///
29 /// L(S) = Σ_{σ_i > ε} log σ_i, where ε is a relative threshold that
30 /// identifies the structural nullspace directly from the eigenspectrum.
31 pub value: f64,
32 /// ∂/∂ρₖ L(S) — first derivatives (one per smoothing parameter).
33 ///
34 /// ∂_k L = tr(S⁺ Aₖ) where Aₖ = λₖ Sₖ and S⁺ is the pseudoinverse
35 /// restricted to the positive eigenspace.
36 pub first: Array1<f64>,
37 /// ∂²/(∂ρₖ∂ρₗ) L(S) — second derivatives (for outer Hessian).
38 ///
39 /// ∂²_kl L = δ_{kl} ∂_k L − λₖ λₗ tr(S⁺ Sₖ S⁺ Sₗ).
40 pub second: Option<Array2<f64>>,
41}
42
43/// Unified representation of a single smoothing-parameter penalty coordinate.
44///
45
46// PenaltyLogdetEigenspace, build_penalty_logdet_eigenspace,
47// scaled_penalty_logdet_nullspace_leakage, and frobenius_inner_same_shape
48// have been replaced by the canonical PenaltyPseudologdet in
49// super::super::penalty_logdet. All callers now use that module directly.
50
51/// Reduced trace kernel `K = U · M · Uᵀ` for pseudo-logdet REML/LAML
52/// criteria: an orthonormal column basis `u_s` (p × r) plus the r × r
53/// symmetric reduced kernel `h_proj_inverse`, with `tr(K · A)` evaluated as
54/// `tr(M · Uᵀ A U)` so contractions run on the r-dimensional subspace.
55///
56/// Two producers install it, with different (documented) exactness domains:
57///
58/// 1. **Intrinsic spectral form (#901, the GLM dense paths in runtime.rs —
59/// `intrinsic_hessian_pseudo_logdet_parts`):** `u_s = U_H`, the kept
60/// eigenvectors of the penalized Hessian `H_pen`, and `h_proj_inverse =
61/// diag(1/σ_a)`. Then `K = H_pen⁺` exactly, and `tr(K · Ḣ)` is the exact
62/// first derivative of the cost's `log|H_pen|₊` along **every** drift
63/// direction — penalty-supported or not, moving-subspace ψ drifts
64/// included — because on a constant-rank stratum first-order eigenvector
65/// motion cancels out of the pseudo-logdet derivative. This object can be
66/// traced against the GLM IFT correction `D_β H[v] = X' diag(c ⊙ X v) X`
67/// (which leaks onto `null(S)` via the intercept column) without error.
68///
69/// 2. **Range(Sλ) Schur block (#752, `joint_penalty_subspace_trace_parts`
70/// in custom_family.rs):** `u_s` spans `range(Sλ)` and `h_proj_inverse =
71/// U_Sᵀ (H+Sλ)⁺ U_S`. For penalty-supported `A` (`A = ∂Sλ/∂ρ`), the
72/// identity `U_S U_Sᵀ A U_S U_Sᵀ = A` gives `tr(K · A) = tr((H+Sλ)⁺ A) =
73/// d log|H+Sλ|₊/dρ` — exact for the ρ family. It is **not** exact for
74/// drifts with `null(Sλ)` support (GLM cubic corrections, ψ basis
75/// drifts); paths that carry such drifts must install form 1.
76///
77/// Historically this struct carried a third reading — `(U_Sᵀ H U_S)⁻¹`, the
78/// plain projected inverse paired with the projected cost `log|U_Sᵀ H U_S|₊`.
79/// That object is WRONG as a REML determinant term: splitting `H` over
80/// `range(S) ⊕ ker(S)` as `[[A,B],[Bᵀ,C]]`, the projected logdet is
81/// `log det A`, dropping the θ-dependent Schur curvature
82/// `log det(C − BᵀA⁻¹B)` of the likelihood-identified, penalty-null block
83/// (sign-flipped ρ-gradients, ~1e5 ψ blow-ups vs FD — #901). No producer
84/// builds it anymore.
85#[derive(Clone, Debug)]
86pub struct PenaltySubspaceTrace {
87 pub u_s: Array2<f64>,
88 pub h_proj_inverse: Array2<f64>,
89}
90
91impl PenaltySubspaceTrace {
92 /// Compute `tr(K · A)` where `K = U_S · h_proj_inverse · U_Sᵀ` — the
93 /// pseudo-logdet trace kernel (see the struct doc for the two producer
94 /// forms and their exactness domains).
95 ///
96 /// Uses the identity `tr(K · A) = tr(h_proj_inverse · U_Sᵀ A U_S)` so the
97 /// reduction runs on the r × r subspace rather than materializing K.
98 pub fn trace_projected_logdet(&self, a: &Array2<f64>) -> f64 {
99 gam_terms::construction::trace_penalty_covariance_in_orthogonal_basis(
100 a,
101 &self.u_s,
102 &self.h_proj_inverse,
103 )
104 }
105
106 /// Reduce a p × p matrix `A` to its r × r projection `U_Sᵀ · A · U_S`.
107 ///
108 /// Exposed so callers that need the same reduced matrix for both the
109 /// single-trace `tr(K · A)` and the cross-trace `tr(K · A · K · B)`
110 /// can avoid repeating the p × p · p × r matmuls. Routes through
111 /// faer's parallel SIMD GEMM (`fast_atb` / `fast_ab`) so the p-large
112 /// contraction axis amortizes across all cores.
113 pub fn reduce(&self, a: &Array2<f64>) -> Array2<f64> {
114 let u_s_t_a = gam_linalg::faer_ndarray::fast_atb(&self.u_s, a);
115 gam_linalg::faer_ndarray::fast_ab(&u_s_t_a, &self.u_s)
116 }
117
118 /// Compute `tr(H_proj⁻¹ · R)` given an already-reduced `R = U_Sᵀ A U_S`.
119 pub fn trace_projected_logdet_reduced(&self, r_mat: &Array2<f64>) -> f64 {
120 gam_terms::construction::trace_reduced_penalty_covariance(r_mat, &self.h_proj_inverse)
121 }
122
123 /// Cross-trace given pre-reduced blocks `R_A = U_Sᵀ A U_S`, `R_B = U_Sᵀ B U_S`.
124 pub fn trace_projected_logdet_cross_reduced(&self, ra: &Array2<f64>, rb: &Array2<f64>) -> f64 {
125 // left = H_proj⁻¹ · R_A ; right = H_proj⁻¹ · R_B ; tr(left · right).
126 let left = self.h_proj_inverse.dot(ra);
127 let right = self.h_proj_inverse.dot(rb);
128 dense::trace_product(&left, &right)
129 }
130
131 /// Reduce a `HyperOperator` `A` to its `r × r` projection
132 /// `U_Sᵀ · A · U_S` without materializing the dense `p × p` block.
133 /// Uses `A.mul_mat(U_S)` so an Hv-only operator is probed in `r` matvecs
134 /// (each `O(work_of_A)`), then a single `r × p × r` reduction routed
135 /// through faer's parallel SIMD GEMM (`fast_atb`).
136 pub fn reduce_operator<O>(&self, a: &O) -> Array2<f64>
137 where
138 O: HyperOperator + ?Sized,
139 {
140 let au = a.mul_mat(&self.u_s);
141 gam_linalg::faer_ndarray::fast_atb(&self.u_s, &au)
142 }
143
144 /// `tr(K · A)` for `A` exposed only as a `HyperOperator`. Mirrors
145 /// [`Self::trace_projected_logdet`] without forcing dense materialization
146 /// of `A`.
147 pub fn trace_operator<O>(&self, a: &O) -> f64
148 where
149 O: HyperOperator + ?Sized,
150 {
151 self.trace_projected_logdet_reduced(&self.reduce_operator(a))
152 }
153
154 /// Projected leverage `h^{G,proj}_i = Xᵢᵀ · K · Xᵢ` for every row of `x`.
155 ///
156 /// Computed in bulk as `Z = X · U_S` (`n × r`) then
157 /// `h^{G,proj}_i = (Z H_proj⁻¹ Zᵀ)_{ii} = Σ_{a,b} Z_{ia} (H_proj⁻¹)_{ab} Z_{ib}`,
158 /// total cost `O(n · p · r + n · r²)` — strictly cheaper than `n` calls
159 /// to [`Self::apply`] because the `n × p · p × r` GEMM streams the
160 /// `p`-axis once. Streams `X` through `try_row_chunk` so operator-backed
161 /// (Lazy) designs at large scale never densify the full `(n × p)` block.
162 pub fn xt_projected_kernel_x_diagonal(&self, x: &DesignMatrix) -> Array1<f64> {
163 let n = x.nrows();
164 let p = x.ncols();
165 let r = self.u_s.ncols();
166 assert_eq!(self.u_s.nrows(), p);
167 assert_eq!(self.h_proj_inverse.nrows(), r);
168 assert_eq!(self.h_proj_inverse.ncols(), r);
169
170 let block = {
171 const TARGET_CHUNK_FLOATS: usize = 1 << 16;
172 (TARGET_CHUNK_FLOATS / p.max(1)).clamp(1, n.max(1))
173 };
174
175 let mut h = Array1::<f64>::zeros(n);
176 let mut start = 0usize;
177 while start < n {
178 let end = (start + block).min(n);
179 let rows = x.try_row_chunk(start..end).unwrap_or_else(|err| {
180 // SAFETY: `start..end` is constructed from
181 // `0..n = 0..x.nrows()` with `end = (start+block).min(n)`,
182 // so it is always a valid sub-range of `x`. Failure means
183 // the operator broke its row-chunk contract.
184 // SAFETY: row range built from 0..x.nrows(); failure means operator broke its contract.
185 reml_contract_panic(format!(
186 "xt_projected_kernel_x_diagonal: row chunk failed: {err}"
187 ))
188 });
189 // Z_chunk = rows · U_S ((end-start) × r).
190 let z_chunk = gam_linalg::faer_ndarray::fast_ab(&rows, &self.u_s);
191 // h_i = Σ_{a,b} Z_{ia} (H_proj⁻¹)_{ab} Z_{ib}.
192 for (i, row_z) in z_chunk.outer_iter().enumerate() {
193 let mut acc = 0.0;
194 for (z_a, h_row) in row_z
195 .iter()
196 .copied()
197 .zip(self.h_proj_inverse.rows().into_iter())
198 {
199 let mut inner = 0.0;
200 for (h_value, z_b) in h_row.iter().copied().zip(row_z.iter().copied()) {
201 inner += h_value * z_b;
202 }
203 acc += z_a * inner;
204 }
205 h[start + i] = acc;
206 }
207 start = end;
208 }
209 h
210 }
211
212 /// Projected bilinear pseudo-inverse `aᵀ · K⁺ · b` where
213 /// `K⁺ = U_S · H_proj⁻¹ · U_Sᵀ`.
214 ///
215 /// Used by the rank-deficient LAML IFT correction path: when `b ∈
216 /// col(S_k) ⊂ range(S_+)`, applying the projected pseudo-inverse
217 /// instead of the full `H⁻¹` strips spurious null-space noise from
218 /// `a` (≈ the outer-stationarity residual `r`) before the inverse,
219 /// without biasing the numerator. Costs `O(p·r + r²)` versus the
220 /// `O(p²·r)` full solve.
221 pub fn bilinear_pseudo_inverse(&self, a: &Array1<f64>, b: &Array1<f64>) -> f64 {
222 let proj_a = gam_linalg::faer_ndarray::fast_atv(&self.u_s, a);
223 let proj_b = gam_linalg::faer_ndarray::fast_atv(&self.u_s, b);
224 let h_proj_inv_b = self.h_proj_inverse.dot(&proj_b);
225 proj_a.dot(&h_proj_inv_b)
226 }
227
228 /// Euclidean projection onto the retained penalty/Hessian range used by
229 /// this projected kernel: `P_S a = U_S U_Sᵀ a`.
230 pub fn project_onto_subspace(&self, a: &Array1<f64>) -> Array1<f64> {
231 let proj_a = gam_linalg::faer_ndarray::fast_atv(&self.u_s, a);
232 gam_linalg::faer_ndarray::fast_av(&self.u_s, &proj_a)
233 }
234
235 /// Apply the projected pseudo-inverse `K = U_S · H_proj⁻¹ · U_Sᵀ` to a
236 /// vector `a`, returning the minimum-norm solution `v = K · a` of the
237 /// system `H v = a` restricted to `range(S₊)`.
238 ///
239 /// This is the correct stand-in for `H⁻¹ · a` in all per-coordinate
240 /// outer-gradient/Hessian formulas when the rank-deficient LAML fix is
241 /// active (`penalty_subspace_trace = Some`). The full `H⁻¹ · a` solve
242 /// amplifies any component of `a` outside `range(H_free)` by
243 /// `1/σ_min(H_active_normal)` — which on large-scale survival
244 /// marginal-slope is ~10¹² and propagates into outer gradients of
245 /// magnitude 10¹⁴, suppressed by the envelope tripwire downstream and
246 /// killing every seed before the fit can take a step. Dropping the
247 /// `ker(M)` component of `a` is exactly what the cost `½ log|M|₊` demands —
248 /// that subspace is the cost's gauge, so the drop is bit-invariant for the
249 /// correction regardless of the residual's gauge-direction magnitude;
250 /// `ProjectedKktResidual::projected_into_reduced_range` certifies only that
251 /// the dropped component is genuinely invisible to this kernel (no retained-
252 /// range leakage / orthonormal `U_S`) before the IFT correction uses it. The
253 /// returned gradient lives on the identifiable manifold, matching the
254 /// projected `½ log|M|₊` term.
255 ///
256 /// Costs `O(p·r + r²)` for the two `U_S`-contractions plus the `r × r`
257 /// solve — strictly cheaper than the `O(p²)` full `hop.solve_multi`
258 /// when `r ≪ p`, and bounded regardless of `σ_min(H)`.
259 pub fn apply_pseudo_inverse(&self, a: &Array1<f64>) -> Array1<f64> {
260 // The one sensitivity operator (#935): the projected inverse action
261 // `U_S · H_proj⁻¹ · U_Sᵀ · a` has a single spelling, shared with every
262 // other consumer of `FittedInverse::Projected`.
263 self.sensitivity().apply(a)
264 }
265
266 /// View this projected trace kernel as the unified [`FitSensitivity`]
267 /// (#935) over the rank-deficient LAML convention `K = U_S · H_proj⁻¹ ·
268 /// U_Sᵀ`. The trace machinery stays here; the *inverse action* is the
269 /// shared operator, so no site can disagree about what `H⁻¹` means.
270 pub fn sensitivity(&self) -> crate::sensitivity::FitSensitivity<'_> {
271 crate::sensitivity::FitSensitivity::from_projected(&self.u_s, &self.h_proj_inverse)
272 }
273
274 /// Build the **constrained pseudo-inverse kernel**
275 /// `K_T = K_S − K_S Aᵀ (A K_S Aᵀ)⁻¹ A K_S`
276 /// from this penalty-projected kernel `K_S` and the *active* row block
277 /// `A_act` of the joint linear inequality constraint matrix.
278 ///
279 /// `K_T` is the **Moore-Penrose pseudo-inverse of `H` restricted to
280 /// `T = range(S₊) ∩ ker(A_act)`** — the smooth manifold the inner
281 /// solver actually moves on at a constrained-stationary point. It is
282 /// exactly the kernel that solves the per-coordinate saddle-point
283 /// IFT system
284 ///
285 /// ```text
286 /// [ H Aᵀ_act ] [ ∂β/∂ρ_k ] [ −a_k ]
287 /// [ A_act 0 ] [ ∂λ/∂ρ_k ] = [ 0 ]
288 /// ```
289 ///
290 /// with `∂β/∂ρ_k = −K_T · a_k`. Using `K_T` for the per-coordinate
291 /// mode response `v_k` makes the outer gradient the *exact* derivative
292 /// of the projected Laplace cost `log|U_Tᵀ H U_T|`, where `U_T` is an
293 /// orthonormal basis of `T` — the marginal-likelihood determinant the
294 /// inner is actually drawing on.
295 ///
296 /// Returns a [`ConstrainedSubspaceKernel`] handle that caches the
297 /// small `k_active × k_active` Schur complement so subsequent
298 /// `apply_pseudo_inverse` calls for different RHS reuse it. When the
299 /// active set is empty the handle degrades to a pass-through over
300 /// `self` (no extra work).
301 ///
302 /// Total precompute cost: `k_active` calls to
303 /// [`Self::apply_pseudo_inverse`] (one per active row) plus a
304 /// `k_active × k_active` Cholesky/QR. Per-vector `apply` cost: one
305 /// `K_S` apply + one `k_active × p` matvec + one small triangular
306 /// solve + one `p × k_active` matvec.
307 pub fn with_active_constraints<'a>(
308 &'a self,
309 a_act: ndarray::ArrayView2<'a, f64>,
310 ) -> ConstrainedSubspaceKernel<'a> {
311 let k_active = a_act.nrows();
312 if k_active == 0 {
313 return ConstrainedSubspaceKernel {
314 kernel: self,
315 z: Array2::zeros((0, self.u_s.nrows())),
316 a_act,
317 m_inv: Array2::zeros((0, 0)),
318 k_active: 0,
319 };
320 }
321 // Z = K_S · Aᵀ_act, shape (p × k_active).
322 let p = self.u_s.nrows();
323 let mut z = Array2::<f64>::zeros((p, k_active));
324 for j in 0..k_active {
325 let a_row = a_act.row(j).to_owned();
326 let k_s_a_row = self.apply_pseudo_inverse(&a_row);
327 z.column_mut(j).assign(&k_s_a_row);
328 }
329 // M = A_act · Z (shape k_active × k_active, symmetric PSD on
330 // range(K_S) ∩ image(A_actᵀ); on a rank-deficient overlap we
331 // add a tiny diagonal regulariser so the inversion remains
332 // bounded — same noise-floor strategy as elsewhere in this
333 // module).
334 let mut m = a_act.dot(&z);
335 // Symmetrise (numerical noise from the matmul leaves small skew).
336 for i in 0..k_active {
337 for j in 0..i {
338 let avg = 0.5 * (m[[i, j]] + m[[j, i]]);
339 m[[i, j]] = avg;
340 m[[j, i]] = avg;
341 }
342 }
343 // Eigendecomposition-based Moore-Penrose pseudo-inverse with a
344 // relative spectral cutoff. This is the principled treatment of
345 // rank deficiency in `A_act` when restricted to `range(S₊)`:
346 // some active constraint rows may be linearly dependent after
347 // projection (e.g. several monotonicity rows pinning the same
348 // flat region all reduce to the same row in `range(S₊)`).
349 // A plain `M⁻¹` then amplifies near-null directions; the
350 // pseudo-inverse drops them at a relative threshold
351 // `tol = eps · k_active · σ_max(M)`, which is the standard
352 // NumPy/LAPACK convention and exactly what Codex flagged as
353 // necessary in the math review.
354 let (evals, evecs) = m
355 .eigh(faer::Side::Lower)
356 .unwrap_or_else(|_| (Array1::<f64>::zeros(k_active), Array2::<f64>::eye(k_active)));
357 let sigma_max = evals.iter().copied().fold(0.0_f64, f64::max).max(0.0);
358 let tol = f64::EPSILON * (k_active as f64) * sigma_max.max(1.0);
359 let mut m_inv = Array2::<f64>::zeros((k_active, k_active));
360 let mut dropped = 0usize;
361 for q in 0..k_active {
362 if evals[q] > tol {
363 let inv_sigma = 1.0 / evals[q];
364 // Outer product u_q u_qᵀ scaled by 1/σ_q.
365 for i in 0..k_active {
366 for j in 0..k_active {
367 m_inv[[i, j]] += inv_sigma * evecs[[i, q]] * evecs[[j, q]];
368 }
369 }
370 } else {
371 dropped += 1;
372 }
373 }
374 if dropped > 0 {
375 log::debug!(
376 "[constrained-subspace kernel] dropped {} of {} active-constraint directions \
377 (rank-deficient on range(S₊)); pseudo-inverse threshold = {:.3e}",
378 dropped,
379 k_active,
380 tol,
381 );
382 }
383 ConstrainedSubspaceKernel {
384 kernel: self,
385 z,
386 a_act,
387 m_inv,
388 k_active,
389 }
390 }
391}
392
393/// Per-evaluation handle that combines a penalty-projected
394/// [`PenaltySubspaceTrace`] with an active inequality-constraint block,
395/// producing the constraint-aware pseudo-inverse
396/// `K_T = K_S − K_S Aᵀ (A K_S Aᵀ)⁻¹ A K_S`. See
397/// [`PenaltySubspaceTrace::with_active_constraints`] for the math.
398///
399/// Caches the small `k_active × k_active` Schur inverse so subsequent
400/// per-coordinate `apply` calls only do `O(p · k_active)` work each.
401pub struct ConstrainedSubspaceKernel<'a> {
402 pub(crate) kernel: &'a PenaltySubspaceTrace,
403 /// `Z = K_S · Aᵀ_act`, shape `(p × k_active)`.
404 pub(crate) z: Array2<f64>,
405 /// Active-row block of the joint constraint matrix.
406 pub(crate) a_act: ndarray::ArrayView2<'a, f64>,
407 /// `(A_act · K_S · Aᵀ_act)⁻¹`, shape `(k_active × k_active)`.
408 pub(crate) m_inv: Array2<f64>,
409 pub(crate) k_active: usize,
410}
411
412impl<'a> ConstrainedSubspaceKernel<'a> {
413 /// Apply `K_T = K_S − K_S Aᵀ (A K_S Aᵀ)⁻¹ A K_S` to `a`. The result
414 /// lies in `range(S₊) ∩ ker(A_act)` — the smooth manifold the inner
415 /// solver actually moves on at a constrained-stationary point.
416 pub fn apply_pseudo_inverse(&self, a: &Array1<f64>) -> Array1<f64> {
417 let v_s = self.kernel.apply_pseudo_inverse(a);
418 if self.k_active == 0 {
419 return v_s;
420 }
421 // mu = M_inv · (A_act · v_s)
422 let t = self.a_act.dot(&v_s);
423 let mu = self.m_inv.dot(&t);
424 // v = v_s - Z · mu
425 let correction = self.z.dot(&mu);
426 v_s - &correction
427 }
428
429 /// Whether any active constraints contribute (when false this kernel
430 /// is identical to the bare [`PenaltySubspaceTrace::apply_pseudo_inverse`]).
431 pub fn has_active_constraints(&self) -> bool {
432 self.k_active > 0
433 }
434}
435
436/// Tangency self-audit gate for the constrained mode-response arm: the
437/// emitted `v = K_T · rhs` must lie in `ker(A_act)` by construction, so
438/// `|A_act · v|` is compared against this fraction of the cancellation
439/// scale `|A_act| · |v|` (per active row). Generous enough that legitimate
440/// rank-deficient active sets (whose dropped Schur directions leave
441/// ε-level residue, see [`PenaltySubspaceTrace::with_active_constraints`])
442/// never trip it; the historical failure mode it guards (the d6b17a7f
443/// `1/σ_min ≈ 10¹²` null-space amplification) exceeds it by six orders.
444pub(crate) const THETA_MODE_RESPONSE_TANGENCY_GATE: f64 = 1e-6;
445
446/// #931 migration pass 2 — the ThetaDirection shared-drift pass: the ONE
447/// per-evaluation selection of the IFT mode-response kernel behind every
448/// `dβ̂/dθ = −K · ∂g/∂θ` solve in the outer gradient/Hessian assembly.
449///
450/// Before this object existed, four sites (the gradient solve stack in
451/// `reml_laml_evaluate`, the ρ- and ext-coordinate standalone fallbacks in
452/// `compute_outer_hessian`, and the standalone fallback in
453/// `build_outer_hessian_operator`) each re-implemented the same selection
454/// rule by hand, with comments warning each other to "mirror the
455/// selection exactly, otherwise the operator-form Hessian and dense
456/// materialization disagree on every entry". A hand-copied convention every
457/// caller must remember is precisely the objective↔gradient desync surface
458/// (#748/#752/#901 class) the criterion-as-atoms architecture (#931)
459/// removes. Now the rule is DECIDED in exactly one constructor and every
460/// consumer is a contraction of the same kernel object — the gradient and
461/// both Hessian representations structurally cannot pick different
462/// inverses for the same evaluation point:
463///
464/// * Active inequality constraints recorded on the inner solution → the
465/// lifted constrained kernel
466/// `K_T = K_S − K_S Aᵀ (A K_S Aᵀ)⁻¹ A K_S`. The inner SCOP solver
467/// clamps β̂(θ) onto `T = range(S₊) ∩ ker(A_act)`, so the true IFT
468/// derivative lives in T and the lifted kernel gives the minimum-norm
469/// solution there; the full solve would amplify any RHS component
470/// outside `range(H_free)` by `1/σ_min(H_active_normal)` — ~10¹² on
471/// large-scale survival marginal-slope (commit d6b17a7f).
472/// * Otherwise → the FULL Hessian solve `v = H⁻¹ · rhs`, even when the
473/// LAML cost surface uses the projected logdet `½ log|U_Sᵀ H U_S|`:
474/// the inner solver converges β̂ ∈ R^p in the unconstrained full
475/// space, so the IFT identity demands the full inverse, and the
476/// penalty-subspace projection acts on the TRACE contraction side
477/// only. Routing through bare `K_S` here would discard the
478/// `null(S₊)` component of dβ̂/dθ — the near-separable ψ-gradient
479/// blow-up pinned by `duchon_probit_per_row_dnu_dpsi_fd_vs_analytic`.
480///
481/// The two emission shapes (`respond_one` per-vector, `respond_stack`
482/// batched) exist because the call sites have different RHS layouts and
483/// their solve shapes must stay bit-identical to the pre-port assembly
484/// (per-column GEMV vs blocked GEMM sum in different orders) — NOT because
485/// a site may choose a different kernel. Both shapes dispatch on the same
486/// stored decision.
487///
488/// This is the `Sensitivity`-operator half of the `ThetaDirection`
489/// calculus sketched in `atoms.rs`: the direction's `β̇` channel is a
490/// contraction of this kernel, so atoms borrowing the shared drift can no
491/// longer see a different chain rule than their neighbors.
492pub(crate) struct ThetaModeResponseKernel<'s> {
493 pub(crate) hop: &'s dyn HessianFactorization,
494 /// `Some` exactly when the selection rule chose the lifted constrained
495 /// kernel. Built once per evaluation point (one Schur-complement
496 /// factorization), shared by every gradient/Hessian consumer — the
497 /// pre-port code rebuilt it per consumer site.
498 pub(crate) constrained: Option<ConstrainedSubspaceKernel<'s>>,
499}
500
501impl<'s> ThetaModeResponseKernel<'s> {
502 /// The ONE place the mode-response kernel selection rule lives.
503 pub(crate) fn select(
504 subspace: Option<&'s PenaltySubspaceTrace>,
505 active_constraints: Option<&'s ActiveLinearConstraintBlock>,
506 hop: &'s dyn HessianFactorization,
507 ) -> Self {
508 let constrained = match (subspace, active_constraints) {
509 (Some(kernel), Some(block)) => {
510 let ck = kernel.with_active_constraints(block.a.view());
511 ck.has_active_constraints().then_some(ck)
512 }
513 _ => None,
514 };
515 Self { hop, constrained }
516 }
517
518 /// Mode response for one right-hand side: `K_T · rhs` under active
519 /// constraints, `H⁻¹ · rhs` (single-RHS `solve`) otherwise. Used by the
520 /// per-coordinate fallbacks whose pre-port assembly solved one vector at
521 /// a time — the single-RHS shape is preserved bit-identically.
522 pub(crate) fn respond_one(&self, rhs: &Array1<f64>) -> Array1<f64> {
523 match self.constrained.as_ref() {
524 Some(ck) => {
525 let v = ck.apply_pseudo_inverse(rhs);
526 self.certify_tangency(ck, &v);
527 v
528 }
529 None => self.hop.solve(rhs),
530 }
531 }
532
533 /// Mode responses for a column-stacked RHS block: per-column `K_T`
534 /// applies under active constraints (the lifted kernel has no blocked
535 /// form), one batched `solve_multi` otherwise (BLAS-3 / GPU batched
536 /// route) — exactly the shapes the stacked call sites used pre-port.
537 /// Zero RHS columns (box-masked ρ coordinates) emit exact zeros through
538 /// either arm, since both kernels are linear.
539 pub(crate) fn respond_stack(&self, rhs_stack: &Array2<f64>) -> Array2<f64> {
540 match self.constrained.as_ref() {
541 Some(ck) => {
542 let mut out = Array2::<f64>::zeros(rhs_stack.raw_dim());
543 for (j, col) in rhs_stack.columns().into_iter().enumerate() {
544 let v = ck.apply_pseudo_inverse(&col.to_owned());
545 self.certify_tangency(ck, &v);
546 out.column_mut(j).assign(&v);
547 }
548 out
549 }
550 None => self.hop.solve_multi(rhs_stack),
551 }
552 }
553
554 /// Per-atom certify body (#934 FD-self-audit pattern, applied as an
555 /// exact structural invariant): every constrained emission must lie in
556 /// `ker(A_act)` — `A_act · v = 0` is the defining property of `K_T`'s
557 /// range, so a violation can only mean the kernel object and the
558 /// emission desynced. Checked on every constrained response (cost
559 /// `O(k_active · p)`, negligible next to the apply itself) against the
560 /// row-wise cancellation scale `|A_act| · |v|`; a violation does not
561 /// fail the fit — it names the atom loudly in the `[CERTIFICATE]`
562 /// stream, exactly like the outer-optimum criterion audit. The
563 /// unconstrained arm carries no separate certify: its coherence with
564 /// the criterion VALUE is audited end-to-end by the #934
565 /// `OuterCriterionCertificate` at every returned optimum.
566 pub(crate) fn certify_tangency(&self, ck: &ConstrainedSubspaceKernel<'_>, v: &Array1<f64>) {
567 let residual = ck.a_act.dot(v);
568 for (row, r) in residual.iter().enumerate() {
569 let scale: f64 = ck
570 .a_act
571 .row(row)
572 .iter()
573 .zip(v.iter())
574 .map(|(a, x)| (a * x).abs())
575 .sum();
576 if r.abs() > THETA_MODE_RESPONSE_TANGENCY_GATE * (scale + f64::EPSILON) {
577 log::warn!(
578 "[CERTIFICATE warning] atom \"theta_mode_response\": constrained IFT \
579 mode response left ker(A_act) — active row {row} residual {:.3e} \
580 exceeds gate {:.1e}·{:.3e}; the lifted kernel K_T and its emission \
581 have desynced (#931 pass-2 invariant)",
582 r.abs(),
583 THETA_MODE_RESPONSE_TANGENCY_GATE,
584 scale,
585 );
586 }
587 }
588 }
589}
590
591impl ProjectedKktResidual {
592 pub(crate) fn projected_into_reduced_range(
593 &self,
594 kernel: &PenaltySubspaceTrace,
595 ) -> Result<Self, String> {
596 match self.subspace {
597 KktResidualSubspace::ReducedRange => Ok(self.clone()),
598 KktResidualSubspace::ActiveProjected => {
599 let reduced_residual = kernel.project_onto_subspace(&self.residual);
600 let dropped = &self.residual - &reduced_residual;
601 // Validity invariant for reducing `r_A → r_R = U_S U_Sᵀ r_A`
602 // before the projected IFT correction `−½ rᵀ K r`, `K = U_S
603 // H_proj⁻¹ U_Sᵀ`.
604 //
605 // Since #901 the kernel's `U_S` is the ORTHONORMAL eigenbasis of
606 // the full LAML Hessian `M = H + Sλ (+ H_Φ)` over its identifiable
607 // range (`|σ| > positive_eigenvalue_threshold`), the SAME kept-set
608 // that defines the cost's pseudo-logdet `½ log|M|₊`
609 // (`joint_penalty_subspace_trace_parts`). The dropped component
610 // `r_A − r_R = (I − U_S U_Sᵀ) r_A` therefore lives in `ker(M)` —
611 // the genuine gauge of the cost surface — and is ANNIHILATED by the
612 // kernel: `K (r_A − r_R) = U_S H_proj⁻¹ (U_Sᵀ(I − U_S U_Sᵀ) r_A) = 0`
613 // because `U_Sᵀ(I − U_S U_Sᵀ) = 0` for orthonormal `U_S`. Dropping
614 // it leaves the correction `−½ rᵀ K r` (and every `tr(K·)` trace)
615 // bit-identical, so a gauge-direction residual of ANY magnitude is
616 // valid to reduce.
617 //
618 // The magnitude the inner solver leaves in `ker(M)` is NOT bounded
619 // by the KKT stationarity tolerance: those are the weakly-/un-
620 // identified directions where `M`'s curvature is below threshold, so
621 // a Newton step to resolve their residual is `r/σ` — far outside the
622 // trust region — and the inner correctly leaves them unmoved
623 // (gam#1040 gauge/plateau exit), the outer projected pseudo-inverse
624 // correctly discards them, and the cost `½ log|M|₊` never included
625 // them. Comparing that gauge mass to `residual_tol` (the previous
626 // gate) is a category error inherited from the pre-#901 `range(Sλ)`
627 // kernel, whose dropped subspace `null(Sλ)` DID contain likelihood-
628 // identified curvature; it rejected every oversmoothed marginal-slope
629 // fit as a false "unresolved mass" contract violation.
630 //
631 // The one property that IS required — and that this reduction must
632 // certify — is that the dropped component is genuinely invisible to
633 // the kernel, i.e. carries no RETAINED-range mass. For an orthonormal
634 // `U_S` the leakage `U_Sᵀ(r_A − r_R)` is identically zero; a nonzero
635 // value can only mean `U_S` is not orthonormal / the kernel and
636 // residual desynced (the actual failure class the guard exists to
637 // catch), in which case reducing WOULD alter the correction.
638 let retained_leak = gam_linalg::faer_ndarray::fast_atv(&kernel.u_s, &dropped);
639 let leak_inf = retained_leak
640 .iter()
641 .map(|value| value.abs())
642 .fold(0.0_f64, f64::max);
643 let residual_inf = self
644 .residual
645 .iter()
646 .map(|value| value.abs())
647 .fold(0.0_f64, f64::max);
648 // Orthonormal-projection round-off floor: `U_Sᵀ(I − U_S U_Sᵀ) r`
649 // accumulates only `O(r_cols · ε · ‖r‖)` numerical noise, so a
650 // relative floor of `1e-8·(1 + ‖r‖∞)` clears every well-formed
651 // kernel by orders of magnitude while a genuinely non-orthonormal
652 // `U_S` leaks `O(‖r‖)` and trips it.
653 let leak_floor = 1e-8 * (1.0 + residual_inf);
654 if leak_inf > leak_floor {
655 return Err(format!(
656 "projected KKT residual reduction leaks retained-range mass: \
657 |U_Sᵀ(r_A − r_R)|∞={leak_inf:.3e} > floor={leak_floor:.3e}; the \
658 reduced-range kernel U_S is not orthonormal or is desynced from the \
659 residual, so projecting r_A onto range(U_S) would change the IFT \
660 correction −½ rᵀ U_S H_proj⁻¹ U_Sᵀ r instead of leaving it invariant"
661 ));
662 }
663 let mut reduced = Self::from_reduced_range(reduced_residual);
664 reduced.residual_tol = self.residual_tol;
665 reduced.free_rank = self.free_rank;
666 Ok(reduced)
667 }
668 }
669 }
670
671 /// The KKT-stationarity tolerance the inner solver applied at the
672 /// producing iterate. Returns `None` when the residual was built
673 /// from a legacy site that hasn't been threaded yet; downstream
674 /// consumers should substitute `f64::NAN` in that case.
675 pub fn residual_tol(&self) -> Option<f64> {
676 self.residual_tol
677 }
678
679 /// Dimensionality of the free subspace: `total_p - active_set_size`
680 /// at the producing iterate. `None` from legacy construction sites.
681 pub fn free_rank(&self) -> Option<usize> {
682 self.free_rank
683 }
684}