use super::*;
pub use gam_problem::PenaltyCoordinate;
#[derive(Clone, Debug)]
pub struct PenaltyLogdetDerivs {
pub value: f64,
pub first: Array1<f64>,
pub second: Option<Array2<f64>>,
}
#[derive(Clone, Debug)]
pub struct PenaltySubspaceTrace {
pub u_s: Array2<f64>,
pub h_proj_inverse: Array2<f64>,
}
impl PenaltySubspaceTrace {
pub fn trace_projected_logdet(&self, a: &Array2<f64>) -> f64 {
gam_terms::construction::trace_penalty_covariance_in_orthogonal_basis(
a,
&self.u_s,
&self.h_proj_inverse,
)
}
pub fn reduce(&self, a: &Array2<f64>) -> Array2<f64> {
let u_s_t_a = gam_linalg::faer_ndarray::fast_atb(&self.u_s, a);
gam_linalg::faer_ndarray::fast_ab(&u_s_t_a, &self.u_s)
}
pub fn trace_projected_logdet_reduced(&self, r_mat: &Array2<f64>) -> f64 {
gam_terms::construction::trace_reduced_penalty_covariance(r_mat, &self.h_proj_inverse)
}
pub fn trace_projected_logdet_cross_reduced(&self, ra: &Array2<f64>, rb: &Array2<f64>) -> f64 {
let left = self.h_proj_inverse.dot(ra);
let right = self.h_proj_inverse.dot(rb);
trace_matrix_product(&left, &right)
}
pub fn reduce_operator<O>(&self, a: &O) -> Array2<f64>
where
O: HyperOperator + ?Sized,
{
let au = a.mul_mat(&self.u_s);
gam_linalg::faer_ndarray::fast_atb(&self.u_s, &au)
}
pub fn trace_operator<O>(&self, a: &O) -> f64
where
O: HyperOperator + ?Sized,
{
self.trace_projected_logdet_reduced(&self.reduce_operator(a))
}
pub fn xt_projected_kernel_x_diagonal(&self, x: &DesignMatrix) -> Array1<f64> {
let n = x.nrows();
let p = x.ncols();
let r = self.u_s.ncols();
assert_eq!(self.u_s.nrows(), p);
assert_eq!(self.h_proj_inverse.nrows(), r);
assert_eq!(self.h_proj_inverse.ncols(), r);
let block = {
const TARGET_CHUNK_FLOATS: usize = 1 << 16;
(TARGET_CHUNK_FLOATS / p.max(1)).clamp(1, n.max(1))
};
let mut h = Array1::<f64>::zeros(n);
let mut start = 0usize;
while start < n {
let end = (start + block).min(n);
let rows = x.try_row_chunk(start..end).unwrap_or_else(|err| {
reml_contract_panic(format!(
"xt_projected_kernel_x_diagonal: row chunk failed: {err}"
))
});
let z_chunk = gam_linalg::faer_ndarray::fast_ab(&rows, &self.u_s);
for (i, row_z) in z_chunk.outer_iter().enumerate() {
let mut acc = 0.0;
for (z_a, h_row) in row_z
.iter()
.copied()
.zip(self.h_proj_inverse.rows().into_iter())
{
let mut inner = 0.0;
for (h_value, z_b) in h_row.iter().copied().zip(row_z.iter().copied()) {
inner += h_value * z_b;
}
acc += z_a * inner;
}
h[start + i] = acc;
}
start = end;
}
h
}
pub fn bilinear_pseudo_inverse(&self, a: &Array1<f64>, b: &Array1<f64>) -> f64 {
let proj_a = gam_linalg::faer_ndarray::fast_atv(&self.u_s, a);
let proj_b = gam_linalg::faer_ndarray::fast_atv(&self.u_s, b);
let h_proj_inv_b = self.h_proj_inverse.dot(&proj_b);
proj_a.dot(&h_proj_inv_b)
}
pub fn project_onto_subspace(&self, a: &Array1<f64>) -> Array1<f64> {
let proj_a = gam_linalg::faer_ndarray::fast_atv(&self.u_s, a);
gam_linalg::faer_ndarray::fast_av(&self.u_s, &proj_a)
}
pub fn apply_pseudo_inverse(&self, a: &Array1<f64>) -> Array1<f64> {
self.sensitivity().apply(a)
}
pub fn sensitivity(&self) -> crate::sensitivity::FitSensitivity<'_> {
crate::sensitivity::FitSensitivity::from_projected(&self.u_s, &self.h_proj_inverse)
}
pub fn with_active_constraints<'a>(
&'a self,
a_act: ndarray::ArrayView2<'a, f64>,
) -> ConstrainedSubspaceKernel<'a> {
let k_active = a_act.nrows();
if k_active == 0 {
return ConstrainedSubspaceKernel {
kernel: self,
z: Array2::zeros((0, self.u_s.nrows())),
a_act,
m_inv: Array2::zeros((0, 0)),
k_active: 0,
};
}
let p = self.u_s.nrows();
let mut z = Array2::<f64>::zeros((p, k_active));
for j in 0..k_active {
let a_row = a_act.row(j).to_owned();
let k_s_a_row = self.apply_pseudo_inverse(&a_row);
z.column_mut(j).assign(&k_s_a_row);
}
let mut m = a_act.dot(&z);
for i in 0..k_active {
for j in 0..i {
let avg = 0.5 * (m[[i, j]] + m[[j, i]]);
m[[i, j]] = avg;
m[[j, i]] = avg;
}
}
let (evals, evecs) = m
.eigh(faer::Side::Lower)
.unwrap_or_else(|_| (Array1::<f64>::zeros(k_active), Array2::<f64>::eye(k_active)));
let sigma_max = evals.iter().copied().fold(0.0_f64, f64::max).max(0.0);
let tol = f64::EPSILON * (k_active as f64) * sigma_max.max(1.0);
let mut m_inv = Array2::<f64>::zeros((k_active, k_active));
let mut dropped = 0usize;
for q in 0..k_active {
if evals[q] > tol {
let inv_sigma = 1.0 / evals[q];
for i in 0..k_active {
for j in 0..k_active {
m_inv[[i, j]] += inv_sigma * evecs[[i, q]] * evecs[[j, q]];
}
}
} else {
dropped += 1;
}
}
if dropped > 0 {
log::debug!(
"[constrained-subspace kernel] dropped {} of {} active-constraint directions \
(rank-deficient on range(S₊)); pseudo-inverse threshold = {:.3e}",
dropped,
k_active,
tol,
);
}
ConstrainedSubspaceKernel {
kernel: self,
z,
a_act,
m_inv,
k_active,
}
}
}
pub struct ConstrainedSubspaceKernel<'a> {
pub(crate) kernel: &'a PenaltySubspaceTrace,
pub(crate) z: Array2<f64>,
pub(crate) a_act: ndarray::ArrayView2<'a, f64>,
pub(crate) m_inv: Array2<f64>,
pub(crate) k_active: usize,
}
impl<'a> ConstrainedSubspaceKernel<'a> {
pub fn apply_pseudo_inverse(&self, a: &Array1<f64>) -> Array1<f64> {
let v_s = self.kernel.apply_pseudo_inverse(a);
if self.k_active == 0 {
return v_s;
}
let t = self.a_act.dot(&v_s);
let mu = self.m_inv.dot(&t);
let correction = self.z.dot(&mu);
v_s - &correction
}
pub fn has_active_constraints(&self) -> bool {
self.k_active > 0
}
}
pub(crate) const THETA_MODE_RESPONSE_TANGENCY_GATE: f64 = 1e-6;
pub(crate) struct ThetaModeResponseKernel<'s> {
pub(crate) hop: &'s dyn HessianOperator,
pub(crate) constrained: Option<ConstrainedSubspaceKernel<'s>>,
}
impl<'s> ThetaModeResponseKernel<'s> {
pub(crate) fn select(
subspace: Option<&'s PenaltySubspaceTrace>,
active_constraints: Option<&'s ActiveLinearConstraintBlock>,
hop: &'s dyn HessianOperator,
) -> Self {
let constrained = match (subspace, active_constraints) {
(Some(kernel), Some(block)) => {
let ck = kernel.with_active_constraints(block.a.view());
ck.has_active_constraints().then_some(ck)
}
_ => None,
};
Self { hop, constrained }
}
pub(crate) fn respond_one(&self, rhs: &Array1<f64>) -> Array1<f64> {
match self.constrained.as_ref() {
Some(ck) => {
let v = ck.apply_pseudo_inverse(rhs);
self.certify_tangency(ck, &v);
v
}
None => self.hop.solve(rhs),
}
}
pub(crate) fn respond_stack(&self, rhs_stack: &Array2<f64>) -> Array2<f64> {
match self.constrained.as_ref() {
Some(ck) => {
let mut out = Array2::<f64>::zeros(rhs_stack.raw_dim());
for (j, col) in rhs_stack.columns().into_iter().enumerate() {
let v = ck.apply_pseudo_inverse(&col.to_owned());
self.certify_tangency(ck, &v);
out.column_mut(j).assign(&v);
}
out
}
None => self.hop.solve_multi(rhs_stack),
}
}
pub(crate) fn certify_tangency(&self, ck: &ConstrainedSubspaceKernel<'_>, v: &Array1<f64>) {
let residual = ck.a_act.dot(v);
for (row, r) in residual.iter().enumerate() {
let scale: f64 = ck
.a_act
.row(row)
.iter()
.zip(v.iter())
.map(|(a, x)| (a * x).abs())
.sum();
if r.abs() > THETA_MODE_RESPONSE_TANGENCY_GATE * (scale + f64::EPSILON) {
log::warn!(
"[CERTIFICATE warning] atom \"theta_mode_response\": constrained IFT \
mode response left ker(A_act) — active row {row} residual {:.3e} \
exceeds gate {:.1e}·{:.3e}; the lifted kernel K_T and its emission \
have desynced (#931 pass-2 invariant)",
r.abs(),
THETA_MODE_RESPONSE_TANGENCY_GATE,
scale,
);
}
}
}
}
impl ProjectedKktResidual {
pub(crate) fn projected_into_reduced_range(
&self,
kernel: &PenaltySubspaceTrace,
) -> Result<Self, String> {
match self.subspace {
KktResidualSubspace::ReducedRange => Ok(self.clone()),
KktResidualSubspace::ActiveProjected => {
let reduced_residual = kernel.project_onto_subspace(&self.residual);
let dropped_inf = self
.residual
.iter()
.zip(reduced_residual.iter())
.map(|(full, reduced)| (full - reduced).abs())
.fold(0.0_f64, f64::max);
let residual_inf = self
.residual
.iter()
.map(|value| value.abs())
.fold(0.0_f64, f64::max);
const DEFAULT_KKT_RESIDUAL_REL_TOL: f64 = 1e-10;
let tol = self
.residual_tol
.unwrap_or_else(|| DEFAULT_KKT_RESIDUAL_REL_TOL * (1.0 + residual_inf));
let gate = tol;
if dropped_inf > gate {
return Err(format!(
"projected KKT residual contains unresolved mass outside the reduced \
Hessian/penalty range: |r_A - r_R|∞={dropped_inf:.3e} > tol={gate:.3e}; \
range-projected IFT correction is valid only after the null direction is \
explicitly removed/fixed or after the active-projected residual is small"
));
}
let mut reduced = Self::from_reduced_range(reduced_residual);
reduced.residual_tol = self.residual_tol;
reduced.free_rank = self.free_rank;
Ok(reduced)
}
}
}
pub fn residual_tol(&self) -> Option<f64> {
self.residual_tol
}
pub fn free_rank(&self) -> Option<usize> {
self.free_rank
}
}