pub fn build_term_collection_designs_joint(
data: ArrayView2<'_, f64>,
specs: &[TermCollectionSpec],
) -> Result<Vec<TermCollectionDesign>, BasisError> {
for spec in specs {
validate_term_collection_finite_inputs(data, spec)?;
}
let smooth_blocks = specs
.iter()
.map(|spec| spec.smooth_terms.clone())
.collect::<Vec<_>>();
let planned_blocks = plan_joint_spatial_centers_for_term_blocks(data, &smooth_blocks)?;
let mut out = Vec::with_capacity(specs.len());
for (spec, planned_terms) in specs.iter().zip(planned_blocks.into_iter()) {
let mut planned_spec = spec.clone();
planned_spec.smooth_terms = planned_terms;
out.push(build_term_collection_design_inner(data, &planned_spec)?);
}
Ok(out)
}
pub fn build_term_collection_designs_and_freeze_joint(
data: ArrayView2<'_, f64>,
specs: &[TermCollectionSpec],
) -> Result<(Vec<TermCollectionDesign>, Vec<TermCollectionSpec>), EstimationError> {
let designs = build_term_collection_designs_joint(data, specs)?;
let mut resolved_specs = Vec::with_capacity(specs.len());
for (spec, design) in specs.iter().zip(designs.iter()) {
resolved_specs.push(freeze_term_collection_from_design(spec, design)?);
}
Ok((designs, resolved_specs))
}
pub fn fit_term_collection_forspec(
data: ArrayView2<'_, f64>,
y: ArrayView1<'_, f64>,
weights: ArrayView1<'_, f64>,
offset: ArrayView1<'_, f64>,
spec: &TermCollectionSpec,
family: LikelihoodSpec,
options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError> {
fit_term_collection_forspecwith_heuristic_lambdas(
data, y, weights, offset, spec, None, family, options,
)
}
pub fn fit_term_collection_with_coefficient_groups(
data: ArrayView2<'_, f64>,
y: ArrayView1<'_, f64>,
weights: ArrayView1<'_, f64>,
offset: ArrayView1<'_, f64>,
spec: &TermCollectionSpec,
groups: &[CoefficientGroupSpec],
family: LikelihoodSpec,
options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError> {
if groups.is_empty() {
return fit_term_collection_forspec(data, y, weights, offset, spec, family, options);
}
let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
let base_fit_opts = adaptive_fit_options_base(options, &design);
let realized = design
.realize_coefficient_groups(groups, &base_fit_opts.rho_prior)
.map_err(EstimationError::BasisError)?;
let effective_offset = design
.compose_offset(offset, "coefficient-group fit")
.map_err(EstimationError::BasisError)?;
let mut grouped_options = base_fit_opts.clone();
grouped_options.rho_prior = realized.rho_prior;
let fitted = FittedTermCollection {
fit: gam_solve::estimate::fit_gam_with_penalty_specs(
design.design.clone(),
y,
weights,
effective_offset.view(),
realized.penalty_specs,
realized.nullspace_dims,
family.clone(),
&grouped_options,
)?,
design,
adaptive_diagnostics: None,
};
enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
Ok(fitted)
}
pub fn fit_term_collection_with_penalty_block_gamma_prior_callback<F>(
data: ArrayView2<'_, f64>,
y: ArrayView1<'_, f64>,
weights: ArrayView1<'_, f64>,
offset: ArrayView1<'_, f64>,
spec: &TermCollectionSpec,
callback: F,
family: LikelihoodSpec,
options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError>
where
F: FnMut(&PenaltyBlockGammaPriorMetadata<'_>) -> Option<(f64, f64)>,
{
let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
let effective_offset = design
.compose_offset(offset, "penalty-prior callback fit")
.map_err(EstimationError::BasisError)?;
let mut fit_opts = adaptive_fit_options_base(options, &design);
fit_opts.rho_prior = realize_penalty_block_gamma_priors(&design, callback)
.map_err(EstimationError::BasisError)?;
let fitted = FittedTermCollection {
fit: fit_gamwith_heuristic_lambdas(
design.design.clone(),
y,
weights,
effective_offset.view(),
&design.penalties,
None,
family.clone(),
&fit_opts,
)?,
design,
adaptive_diagnostics: None,
};
enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
Ok(fitted)
}
pub fn fit_term_collection_with_penalty_block_gamma_priors(
data: ArrayView2<'_, f64>,
y: ArrayView1<'_, f64>,
weights: ArrayView1<'_, f64>,
offset: ArrayView1<'_, f64>,
spec: &TermCollectionSpec,
priors: &[(String, f64, f64)],
family: LikelihoodSpec,
options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError> {
let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
let effective_offset = design
.compose_offset(offset, "penalty-prior fit")
.map_err(EstimationError::BasisError)?;
let mut fit_opts = adaptive_fit_options_base(options, &design);
fit_opts.rho_prior = realize_keyed_penalty_block_gamma_priors(&design, priors)
.map_err(EstimationError::BasisError)?;
let fitted = FittedTermCollection {
fit: fit_gamwith_heuristic_lambdas(
design.design.clone(),
y,
weights,
effective_offset.view(),
&design.penalties,
None,
family.clone(),
&fit_opts,
)?,
design,
adaptive_diagnostics: None,
};
enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
Ok(fitted)
}
pub fn fit_term_collection_with_coefficient_groups_and_penalty_block_gamma_priors(
data: ArrayView2<'_, f64>,
y: ArrayView1<'_, f64>,
weights: ArrayView1<'_, f64>,
offset: ArrayView1<'_, f64>,
spec: &TermCollectionSpec,
groups: &[CoefficientGroupSpec],
priors: &[(String, f64, f64)],
family: LikelihoodSpec,
options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError> {
if groups.is_empty() {
return fit_term_collection_with_penalty_block_gamma_priors(
data, y, weights, offset, spec, priors, family, options,
);
}
if priors.is_empty() {
return fit_term_collection_with_coefficient_groups(
data, y, weights, offset, spec, groups, family, options,
);
}
let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
let base_fit_opts = adaptive_fit_options_base(options, &design);
let base_rho_prior = realize_keyed_penalty_block_gamma_priors(&design, priors)
.map_err(EstimationError::BasisError)?;
let realized = design
.realize_coefficient_groups(groups, &base_rho_prior)
.map_err(EstimationError::BasisError)?;
let effective_offset = design
.compose_offset(offset, "coefficient-group and penalty-prior fit")
.map_err(EstimationError::BasisError)?;
let mut grouped_options = base_fit_opts.clone();
grouped_options.rho_prior = realized.rho_prior;
let fitted = FittedTermCollection {
fit: gam_solve::estimate::fit_gam_with_penalty_specs(
design.design.clone(),
y,
weights,
effective_offset.view(),
realized.penalty_specs,
realized.nullspace_dims,
family.clone(),
&grouped_options,
)?,
design,
adaptive_diagnostics: None,
};
enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
Ok(fitted)
}
fn fit_term_collection_forspecwith_heuristic_lambdas(
data: ArrayView2<'_, f64>,
y: ArrayView1<'_, f64>,
weights: ArrayView1<'_, f64>,
offset: ArrayView1<'_, f64>,
spec: &TermCollectionSpec,
heuristic_lambdas: Option<&[f64]>,
family: LikelihoodSpec,
options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError> {
let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
let resolved_spec;
let design_spec = if adaptive_opts.enabled {
resolved_spec = ensure_matern_adaptive_center_resolution(spec, data.nrows());
&resolved_spec
} else {
spec
};
let base_design =
build_term_collection_design_with_policy(data, design_spec, &options.resource_policy)?;
fit_term_collection_on_realized_design(
y,
weights,
offset,
design_spec,
&base_design,
heuristic_lambdas,
family,
options,
)
}
fn ensure_matern_adaptive_center_resolution(
spec: &TermCollectionSpec,
n_rows: usize,
) -> TermCollectionSpec {
let mut out = spec.clone();
for term in &mut out.smooth_terms {
let gam_terms::smooth::SmoothBasisSpec::Matern {
feature_cols,
spec: matern,
..
} = &mut term.basis
else {
continue;
};
if let gam_terms::basis::CenterStrategy::FarthestPoint { num_centers } =
&mut matern.center_strategy
{
let min_centers = (4 * feature_cols.len()).min(n_rows).max(*num_centers);
*num_centers = min_centers;
}
}
out
}
fn has_bounded_linear_terms(spec: &TermCollectionSpec) -> bool {
spec.linear_terms.iter().any(|term| {
matches!(
term.coefficient_geometry,
LinearCoefficientGeometry::Bounded { .. }
)
})
}
fn fit_term_collection_on_realized_design(
y: ArrayView1<'_, f64>,
weights: ArrayView1<'_, f64>,
offset: ArrayView1<'_, f64>,
spec: &TermCollectionSpec,
design: &TermCollectionDesign,
heuristic_lambdas: Option<&[f64]>,
family: LikelihoodSpec,
options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError> {
let effective_offset = design
.compose_offset(offset, "term-collection fit")
.map_err(EstimationError::BasisError)?;
let offset = effective_offset.view();
if has_bounded_linear_terms(spec) {
return fit_bounded_term_collection_with_design(
y,
weights,
offset,
spec,
design,
heuristic_lambdas,
family,
options,
);
}
let mut base_fit_opts = adaptive_fit_options_base(options, design);
base_fit_opts.rho_prior = relax_smoothing_rho_prior(options, design, y, weights);
let fitted = FittedTermCollection {
fit: fit_gamwith_heuristic_lambdas(
design.design.clone(),
y,
weights,
offset,
&design.penalties,
heuristic_lambdas,
family.clone(),
&base_fit_opts,
)?,
design: design.clone(),
adaptive_diagnostics: None,
};
enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
if !adaptive_opts.enabled {
return Ok(fitted);
}
let runtime_caches = extract_spatial_operator_runtime_caches(spec, &fitted.design)?;
if runtime_caches.is_empty() {
return Ok(fitted);
}
fit_term_collectionwith_exact_spatial_adaptive_regularization(
fitted,
y,
weights,
offset,
family,
options,
&runtime_caches,
)
}
#[derive(Clone)]
struct SpatialOperatorRuntimeCache {
termname: String,
feature_cols: Vec<usize>,
coeff_global_range: Range<usize>,
mass_penalty_global_idx: usize,
tension_penalty_global_idx: usize,
stiffness_penalty_global_idx: usize,
d0: Array2<f64>,
d1: Array2<f64>,
d2: Array2<f64>,
collocation_points: Array2<f64>,
dimension: usize,
}
#[derive(Clone)]
struct SpatialAdaptiveWeights {
inv_magweight: Array1<f64>,
invgradweight: Array1<f64>,
inv_lapweight: Array1<f64>,
}
#[derive(Clone)]
struct CharbonnierScalarBlockState {
signal: Array1<f64>,
radius: Array1<f64>,
epsilon: f64,
}
impl CharbonnierScalarBlockState {
fn from_signal(signal: Array1<f64>, epsilon: f64) -> Self {
let eps = epsilon.max(1e-12);
let radius = signal.mapv(|t| (t * t + eps * eps).sqrt());
Self {
signal,
radius,
epsilon: eps,
}
}
fn absolute_signal(&self) -> Array1<f64> {
self.signal.mapv(f64::abs)
}
fn penalty_value(&self) -> f64 {
self.radius.iter().map(|r| r - self.epsilon).sum::<f64>()
}
fn betagradient_coeff(&self) -> Array1<f64> {
Array1::from_iter(
self.signal
.iter()
.zip(self.radius.iter())
.map(|(t, r)| t / r),
)
}
fn betahessian_diag(&self) -> Array1<f64> {
let eps2 = self.epsilon * self.epsilon;
self.radius.mapv(|r| eps2 / r.powi(3))
}
fn log_epsilon_gradient_terms(&self) -> Array1<f64> {
let epsilon = self.epsilon;
let eps2 = epsilon * epsilon;
self.radius.mapv(|r| eps2 / r - epsilon)
}
fn log_epsilon_betagradient_coeff(&self) -> Array1<f64> {
let eps2 = self.epsilon * self.epsilon;
Array1::from_iter(
self.signal
.iter()
.zip(self.radius.iter())
.map(|(t, r)| -eps2 * t / r.powi(3)),
)
}
fn log_epsilon_hessian_terms(&self) -> Array1<f64> {
let epsilon = self.epsilon;
let eps2 = epsilon * epsilon;
let eps4 = eps2 * eps2;
self.radius
.mapv(|r| 2.0 * eps2 / r - eps4 / r.powi(3) - epsilon)
}
fn surrogateweights_posterior_snr(
&self,
variance: &Array1<f64>,
weight_floor: f64,
weight_ceiling: f64,
) -> (Array1<f64>, Array1<f64>) {
let eps2 = self.epsilon * self.epsilon;
let weight = Array1::from_iter(self.signal.iter().zip(variance.iter()).map(|(&t, &v)| {
let credible2 = (t * t - v.max(0.0)).max(0.0);
let r = (credible2 + eps2).sqrt();
(1.0 / r).clamp(weight_floor, weight_ceiling)
}));
let invweight = weight.mapv(|u| 1.0 / u);
(weight, invweight)
}
fn directionalhessian_diag(&self, direction_signal: &Array1<f64>) -> Array1<f64> {
let eps2 = self.epsilon * self.epsilon;
Array1::from_iter(
self.signal
.iter()
.zip(direction_signal.iter())
.zip(self.radius.iter())
.map(|((t, q), r)| -3.0 * eps2 * t * q / r.powi(5)),
)
}
fn second_directionalhessian_diag(
&self,
direction1_signal: &Array1<f64>,
direction2_signal: &Array1<f64>,
) -> Array1<f64> {
let eps2 = self.epsilon * self.epsilon;
Array1::from_iter(
self.signal
.iter()
.zip(direction1_signal.iter())
.zip(direction2_signal.iter())
.zip(self.radius.iter())
.map(|(((t, q1), q2), r)| {
let r2 = r * r;
let psi4 = -3.0 * eps2 / r.powi(5) + 15.0 * eps2 * t * t / (r.powi(5) * r2);
psi4 * q1 * q2
}),
)
}
fn log_epsilon_betahessian_diag(&self) -> Array1<f64> {
let eps2 = self.epsilon * self.epsilon;
let eps4 = eps2 * eps2;
Array1::from_iter(
self.signal
.iter()
.zip(self.radius.iter())
.map(|(_, r)| 2.0 * eps2 / r.powi(3) - 3.0 * eps4 / r.powi(5)),
)
}
fn log_epsilon_beta_mixed_second_coeff(&self) -> Array1<f64> {
let eps2 = self.epsilon * self.epsilon;
Array1::from_iter(
self.signal
.iter()
.zip(self.radius.iter())
.map(|(t, r)| eps2 * t * (eps2 - 2.0 * t * t) / r.powi(5)),
)
}
fn log_epsilon_betahessian_second_diag(&self) -> Array1<f64> {
let eps2 = self.epsilon * self.epsilon;
let eps4 = eps2 * eps2;
let eps6 = eps4 * eps2;
Array1::from_iter(
self.radius.iter().map(|r| {
4.0 * eps2 / r.powi(3) - 18.0 * eps4 / r.powi(5) + 15.0 * eps6 / r.powi(7)
}),
)
}
fn log_epsilon_betahessian_directional_diag(
&self,
direction_signal: &Array1<f64>,
) -> Array1<f64> {
let eps2 = self.epsilon * self.epsilon;
let eps4 = eps2 * eps2;
Array1::from_iter(
self.signal
.iter()
.zip(direction_signal.iter())
.zip(self.radius.iter())
.map(|((t, q), r)| (-6.0 * eps2 * t / r.powi(5) + 15.0 * eps4 * t / r.powi(7)) * q),
)
}
}
#[derive(Clone)]
struct CharbonnierGroupedBlockState {
norm: Array1<f64>,
radius: Array1<f64>,
signal_blocks: Array2<f64>,
epsilon: f64,
}
impl CharbonnierGroupedBlockState {
fn from_signal_blocks(signal_blocks: Array2<f64>, epsilon: f64) -> Self {
let eps = epsilon.max(1e-12);
let norm = Array1::from_iter(
signal_blocks
.rows()
.into_iter()
.map(|row| row.iter().map(|v| v * v).sum::<f64>().sqrt()),
);
let radius = norm.mapv(|g| (g * g + eps * eps).sqrt());
Self {
norm,
radius,
signal_blocks,
epsilon: eps,
}
}
fn penalty_value(&self) -> f64 {
self.radius.iter().map(|r| r - self.epsilon).sum::<f64>()
}
fn norm_signal(&self) -> Array1<f64> {
self.norm.clone()
}
fn betagradient_blocks(&self) -> Array2<f64> {
let mut out = self.signal_blocks.clone();
for (k, mut row) in out.rows_mut().into_iter().enumerate() {
let scale = 1.0 / self.radius[k];
row.mapv_inplace(|v| v * scale);
}
out
}
fn betahessian_blocks(&self) -> Vec<Array2<f64>> {
let mut out = Vec::with_capacity(self.signal_blocks.nrows());
for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
let dim = row.len();
let mut block = Array2::<f64>::eye(dim);
block.mapv_inplace(|v| v / self.radius[k]);
for i in 0..dim {
for j in 0..dim {
block[[i, j]] -= row[i] * row[j] / self.radius[k].powi(3);
}
}
out.push(block);
}
out
}
fn log_epsilon_gradient_terms(&self) -> Array1<f64> {
let epsilon = self.epsilon;
let eps2 = epsilon * epsilon;
self.radius.mapv(|r| eps2 / r - epsilon)
}
fn log_epsilon_betagradient_blocks(&self) -> Array2<f64> {
let mut out = self.signal_blocks.clone();
let eps2 = self.epsilon * self.epsilon;
for (k, mut row) in out.rows_mut().into_iter().enumerate() {
let scale = -eps2 / self.radius[k].powi(3);
row.mapv_inplace(|v| v * scale);
}
out
}
fn log_epsilon_hessian_terms(&self) -> Array1<f64> {
let epsilon = self.epsilon;
let eps2 = epsilon * epsilon;
let eps4 = eps2 * eps2;
self.radius
.mapv(|r| 2.0 * eps2 / r - eps4 / r.powi(3) - epsilon)
}
fn surrogateweights_posterior_snr(
&self,
variance: &Array1<f64>,
weight_floor: f64,
weight_ceiling: f64,
) -> (Array1<f64>, Array1<f64>) {
let eps2 = self.epsilon * self.epsilon;
let weight = Array1::from_iter(self.norm.iter().zip(variance.iter()).map(|(&g, &v)| {
let credible2 = (g * g - v.max(0.0)).max(0.0);
let r = (credible2 + eps2).sqrt();
(1.0 / r).clamp(weight_floor, weight_ceiling)
}));
let invweight = weight.mapv(|u| 1.0 / u);
(weight, invweight)
}
fn directionalhessian_blocks(&self, direction_blocks: &Array2<f64>) -> Vec<Array2<f64>> {
let mut out = Vec::with_capacity(self.signal_blocks.nrows());
for (k, (v, q)) in self
.signal_blocks
.rows()
.into_iter()
.zip(direction_blocks.rows().into_iter())
.enumerate()
{
let dim = v.len();
let dot = v.iter().zip(q.iter()).map(|(a, b)| a * b).sum::<f64>();
let r3 = self.radius[k].powi(3);
let r5 = self.radius[k].powi(5);
let mut block = Array2::<f64>::eye(dim);
block.mapv_inplace(|x| -dot * x / r3);
for i in 0..dim {
for j in 0..dim {
block[[i, j]] -= (q[i] * v[j] + v[i] * q[j]) / r3;
block[[i, j]] += 3.0 * dot * v[i] * v[j] / r5;
}
}
out.push(block);
}
out
}
fn second_directionalhessian_blocks(
&self,
direction1_blocks: &Array2<f64>,
direction2_blocks: &Array2<f64>,
) -> Vec<Array2<f64>> {
let mut out = Vec::with_capacity(self.signal_blocks.nrows());
for ((k, v), (a, b)) in self.signal_blocks.rows().into_iter().enumerate().zip(
direction1_blocks
.rows()
.into_iter()
.zip(direction2_blocks.rows().into_iter()),
) {
let dim = v.len();
let dot = |x: ndarray::ArrayView1<'_, f64>, y: ndarray::ArrayView1<'_, f64>| {
x.iter().zip(y.iter()).map(|(p, q)| p * q).sum::<f64>()
};
let sa = dot(v, a);
let sb = dot(v, b);
let ab = dot(a, b);
let r = self.radius[k];
let r3 = r.powi(3);
let r5 = r.powi(5);
let r7 = r5 * r * r;
let diag = -ab / r3 + 3.0 * sa * sb / r5;
let mut block = Array2::<f64>::eye(dim);
block.mapv_inplace(|x| diag * x);
for i in 0..dim {
for j in 0..dim {
block[[i, j]] -= (a[i] * b[j] + b[i] * a[j]) / r3;
block[[i, j]] += 3.0 * sb * (a[i] * v[j] + v[i] * a[j]) / r5;
block[[i, j]] += 3.0 * ab * v[i] * v[j] / r5;
block[[i, j]] += 3.0 * sa * (b[i] * v[j] + v[i] * b[j]) / r5;
block[[i, j]] -= 15.0 * sa * sb * v[i] * v[j] / r7;
}
}
out.push(block);
}
out
}
fn log_epsilon_betahessian_blocks(&self) -> Vec<Array2<f64>> {
let mut out = Vec::with_capacity(self.signal_blocks.nrows());
for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
let dim = row.len();
let r3 = self.radius[k].powi(3);
let r5 = self.radius[k].powi(5);
let mut block = Array2::<f64>::eye(dim);
let eps2 = self.epsilon * self.epsilon;
block.mapv_inplace(|v| -eps2 * v / r3);
for i in 0..dim {
for j in 0..dim {
block[[i, j]] += 3.0 * eps2 * row[i] * row[j] / r5;
}
}
out.push(block);
}
out
}
fn log_epsilon_beta_mixed_second_blocks(&self) -> Array2<f64> {
let mut out = self.signal_blocks.clone();
let eps2 = self.epsilon * self.epsilon;
for (k, mut row) in out.rows_mut().into_iter().enumerate() {
let norm2 = self.norm[k] * self.norm[k];
let scale = eps2 * (eps2 - 2.0 * norm2) / self.radius[k].powi(5);
row.mapv_inplace(|v| v * scale);
}
out
}
fn log_epsilon_betahessian_second_blocks(&self) -> Vec<Array2<f64>> {
let mut out = Vec::with_capacity(self.signal_blocks.nrows());
let eps2 = self.epsilon * self.epsilon;
for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
let dim = row.len();
let norm2 = self.norm[k] * self.norm[k];
let r5 = self.radius[k].powi(5);
let r7 = self.radius[k].powi(7);
let mut block = Array2::<f64>::eye(dim);
block.mapv_inplace(|v| eps2 * (eps2 - 2.0 * norm2) * v / r5);
for i in 0..dim {
for j in 0..dim {
block[[i, j]] += 3.0 * eps2 * (2.0 * norm2 - 3.0 * eps2) * row[i] * row[j] / r7;
}
}
out.push(block);
}
out
}
fn log_epsilon_betahessian_directional_blocks(
&self,
direction_blocks: &Array2<f64>,
) -> Vec<Array2<f64>> {
let mut out = Vec::with_capacity(self.signal_blocks.nrows());
let eps2 = self.epsilon * self.epsilon;
for (k, (v, q)) in self
.signal_blocks
.rows()
.into_iter()
.zip(direction_blocks.rows().into_iter())
.enumerate()
{
let dim = v.len();
let dot = v.iter().zip(q.iter()).map(|(a, b)| a * b).sum::<f64>();
let r5 = self.radius[k].powi(5);
let r7 = self.radius[k].powi(7);
let mut block = Array2::<f64>::eye(dim);
block.mapv_inplace(|x| 3.0 * eps2 * dot * x / r5);
for i in 0..dim {
for j in 0..dim {
block[[i, j]] += 3.0 * eps2 * (q[i] * v[j] + v[i] * q[j]) / r5;
block[[i, j]] -= 15.0 * eps2 * dot * v[i] * v[j] / r7;
}
}
out.push(block);
}
out
}
}
fn scalar_operatorgradient(operator: &Array2<f64>, coeff: &Array1<f64>) -> Array1<f64> {
operator.t().dot(coeff)
}
fn scalar_operatorhessian(operator: &Array2<f64>, diag: &Array1<f64>) -> Array2<f64> {
let mut weighted = operator.clone();
for (k, &w) in diag.iter().enumerate() {
weighted.row_mut(k).mapv_inplace(|v| v * w);
}
let gram = operator.t().dot(&weighted);
(&gram + &gram.t().to_owned()) * 0.5
}
fn grouped_operatorgradient(
d1: &Array2<f64>,
dimension: usize,
blocks: &Array2<f64>,
) -> Result<Array1<f64>, EstimationError> {
if blocks.ncols() != dimension {
crate::bail_invalid_estim!(
"grouped gradient block dimension mismatch: got {}, expected {dimension}",
blocks.ncols()
);
}
if d1.nrows() != blocks.nrows() * dimension {
crate::bail_invalid_estim!(
"grouped gradient row mismatch: D1 has {} rows, blocks imply {}",
d1.nrows(),
blocks.nrows() * dimension
);
}
let mut out = Array1::<f64>::zeros(d1.ncols());
for k in 0..blocks.nrows() {
let gk = d1
.slice(s![k * dimension..(k + 1) * dimension, ..])
.to_owned();
out += &gk.t().dot(&blocks.row(k));
}
Ok(out)
}
fn grouped_operatorhessian(
d1: &Array2<f64>,
dimension: usize,
blocks: &[Array2<f64>],
) -> Result<Array2<f64>, EstimationError> {
if d1.nrows() != blocks.len() * dimension {
crate::bail_invalid_estim!(
"grouped Hessian row mismatch: D1 has {} rows, blocks imply {}",
d1.nrows(),
blocks.len() * dimension
);
}
let p = d1.ncols();
let mut out = Array2::<f64>::zeros((p, p));
for (k, block) in blocks.iter().enumerate() {
if block.nrows() != dimension || block.ncols() != dimension {
crate::bail_invalid_estim!(
"grouped Hessian block {k} has shape {}x{}, expected {}x{}",
block.nrows(),
block.ncols(),
dimension,
dimension
);
}
let gk = d1
.slice(s![k * dimension..(k + 1) * dimension, ..])
.to_owned();
out += &gk.t().dot(&block.dot(&gk));
}
Ok((&out + &out.t().to_owned()) * 0.5)
}
#[derive(Clone)]
struct SpatialPenaltyExactState {
magnitude: CharbonnierScalarBlockState,
gradient: CharbonnierGroupedBlockState,
curvature: CharbonnierGroupedBlockState,
}
fn collocationgradient_blocks(
gradrows: &Array1<f64>,
dimension: usize,
) -> Result<Array2<f64>, EstimationError> {
if dimension == 0 || !gradrows.len().is_multiple_of(dimension) {
crate::bail_invalid_estim!(
"invalid collocation gradient layout: rows={}, dimension={dimension}",
gradrows.len()
);
}
let p = gradrows.len() / dimension;
let mut out = Array2::<f64>::zeros((p, dimension));
for k in 0..p {
for axis in 0..dimension {
out[[k, axis]] = gradrows[k * dimension + axis];
}
}
Ok(out)
}
fn collocationhessian_blocks(
hessianrows: &Array1<f64>,
dimension: usize,
) -> Result<Array2<f64>, EstimationError> {
let block_dim = dimension.checked_mul(dimension).ok_or_else(|| {
EstimationError::InvalidInput("invalid collocation Hessian dimension overflow".to_string())
})?;
if block_dim == 0 || !hessianrows.len().is_multiple_of(block_dim) {
crate::bail_invalid_estim!(
"invalid collocation Hessian layout: rows={}, dimension={dimension}",
hessianrows.len()
);
}
let p = hessianrows.len() / block_dim;
let mut out = Array2::<f64>::zeros((p, block_dim));
for k in 0..p {
for idx in 0..block_dim {
out[[k, idx]] = hessianrows[k * block_dim + idx];
}
}
Ok(out)
}
impl SpatialPenaltyExactState {
fn from_beta_local(
beta_local: ArrayView1<'_, f64>,
cache: &SpatialOperatorRuntimeCache,
epsilons: [f64; 3],
) -> Result<Self, EstimationError> {
let gradientrows = cache.d1.dot(&beta_local);
let hessianrows = cache.d2.dot(&beta_local);
Ok(Self {
magnitude: CharbonnierScalarBlockState::from_signal(
cache.d0.dot(&beta_local),
epsilons[0],
),
gradient: CharbonnierGroupedBlockState::from_signal_blocks(
collocationgradient_blocks(&gradientrows, cache.dimension)?,
epsilons[1],
),
curvature: CharbonnierGroupedBlockState::from_signal_blocks(
collocationhessian_blocks(&hessianrows, cache.dimension)?,
epsilons[2],
),
})
}
fn absolute_collocation_magnitudes(&self) -> (Array1<f64>, Array1<f64>, Array1<f64>) {
(
self.magnitude.absolute_signal(),
self.gradient.norm_signal(),
self.curvature.norm_signal(),
)
}
}
fn robust_epsilon_from_samples(values: &[f64], min_epsilon_cfg: f64) -> f64 {
if values.is_empty() {
return min_epsilon_cfg.max(1e-12);
}
let mut clean = values
.iter()
.copied()
.filter(|v| v.is_finite() && *v >= 0.0)
.collect::<Vec<_>>();
if clean.is_empty() {
return min_epsilon_cfg.max(1e-12);
}
clean.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let n = clean.len();
let median = quantile_from_sorted(&clean, 0.5);
let q75 = quantile_from_sorted(&clean, 0.75);
let q95 = quantile_from_sorted(&clean, 0.95);
let mut abs_dev = clean
.iter()
.map(|v| (v - median).abs())
.filter(|v| v.is_finite())
.collect::<Vec<_>>();
abs_dev.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let mad = 1.4826 * quantile_from_sorted(&abs_dev, 0.5);
let mut scale = median.max(mad).max(q75);
let delta = (f64::EPSILON.sqrt() * q95.max(1.0))
.max(min_epsilon_cfg)
.max(1e-12);
let s_min = min_epsilon_cfg.max(1e-12);
if scale <= delta {
let rms = (clean.iter().map(|v| v * v).sum::<f64>() / n as f64).sqrt();
scale = q95.max(rms);
}
if scale <= delta {
scale = s_min;
}
let kappa = 1.0_f64;
(kappa * scale).max(s_min)
}
fn extract_spatial_operator_runtime_caches(
spec: &TermCollectionSpec,
design: &TermCollectionDesign,
) -> Result<Vec<SpatialOperatorRuntimeCache>, EstimationError> {
let smooth_start = design
.design
.ncols()
.saturating_sub(design.smooth.total_smooth_cols());
let mut out = Vec::<SpatialOperatorRuntimeCache>::new();
for (term_idx, (termspec, term_fit)) in spec
.smooth_terms
.iter()
.zip(design.smooth.terms.iter())
.enumerate()
{
let Some(global_range) = design
.smooth_term_penalty_range(term_idx)
.map_err(EstimationError::InvalidInput)?
else {
continue;
};
let global_base_idx = global_range.start;
let mut mass_local_idx = None;
let mut tension_local_idx = None;
let mut stiffness_local_idx = None;
let mut mass_norm = None;
let mut tension_norm = None;
let mut stiffness_norm = None;
for (active_local_idx, penalty) in term_fit.active_penalties.iter().enumerate() {
match penalty.info.source {
PenaltySource::OperatorMass => {
mass_local_idx = Some(active_local_idx);
mass_norm = Some(penalty.info.normalization_scale);
}
PenaltySource::OperatorTension => {
tension_local_idx = Some(active_local_idx);
tension_norm = Some(penalty.info.normalization_scale);
}
PenaltySource::OperatorStiffness => {
stiffness_local_idx = Some(active_local_idx);
stiffness_norm = Some(penalty.info.normalization_scale);
}
_ => {}
}
}
let (
Some(mass_local),
Some(tension_local),
Some(stiffness_local),
Some(mass_scale),
Some(tension_scale),
Some(stiffness_scale),
) = (
mass_local_idx,
tension_local_idx,
stiffness_local_idx,
mass_norm,
tension_norm,
stiffness_norm,
)
else {
continue;
};
let mass_global_idx = global_base_idx + mass_local;
let tension_global_idx = global_base_idx + tension_local;
let stiffness_global_idx = global_base_idx + stiffness_local;
let (feature_cols, mut d0, mut d1, mut d2, collocation_points, dim, center_mass_rows) =
match (&termspec.basis, &term_fit.metadata) {
(
SmoothBasisSpec::Matern { feature_cols, .. },
BasisMetadata::Matern {
centers,
length_scale,
nu,
include_intercept,
identifiability_transform,
aniso_log_scales,
input_scale,
..
},
) => {
let collocation_length_scale =
input_scale.to_standardized_units(*length_scale);
let ops = build_matern_collocation_operator_matrices(
centers.view(),
None,
collocation_length_scale,
*nu,
*include_intercept,
identifiability_transform.as_ref().map(|z| z.view()),
aniso_log_scales.as_deref(),
)?;
(
feature_cols.clone(),
ops.d0,
ops.d1,
ops.d2,
ops.collocation_points,
centers.ncols(),
false,
)
}
(
SmoothBasisSpec::Duchon { feature_cols, .. },
BasisMetadata::Duchon {
centers,
length_scale,
power,
nullspace_order,
identifiability_transform,
input_scale,
aniso_log_scales,
operator_collocation_points: Some(collocation_points),
radial_reparam,
..
},
) => {
let collocation_length_scale = (*length_scale)
.map(|length| input_scale.to_standardized_units(length));
let ops =
gam_terms::basis::build_duchon_collocation_operator_matriceswithworkspace(
centers.view(),
collocation_points.view(),
None,
collocation_length_scale,
*power,
*nullspace_order,
aniso_log_scales.as_deref(),
identifiability_transform.as_ref().map(|z| z.view()),
2,
radial_reparam.as_ref().map(|v| v.view()),
&mut BasisWorkspace::default(),
)?;
(
feature_cols.clone(),
ops.d0,
ops.d1,
ops.d2,
ops.collocation_points,
centers.ncols(),
true,
)
}
_ => continue,
};
if center_mass_rows && d0.nrows() > 0 && d0.ncols() > 0 {
let means = d0.sum_axis(Axis(0)).mapv(|v| v / d0.nrows() as f64);
for mut row in d0.rows_mut() {
row -= &means;
}
}
let mass_scale = mass_scale.max(1e-12).sqrt();
let tension_scale = tension_scale.max(1e-12).sqrt();
let stiffness_scale = stiffness_scale.max(1e-12).sqrt();
d0.mapv_inplace(|v| v / mass_scale);
d1.mapv_inplace(|v| v / tension_scale);
d2.mapv_inplace(|v| v / stiffness_scale);
let coeff_global_range =
(smooth_start + term_fit.coeff_range.start)..(smooth_start + term_fit.coeff_range.end);
if d0.ncols() != coeff_global_range.len()
|| d1.ncols() != coeff_global_range.len()
|| d2.ncols() != coeff_global_range.len()
{
crate::bail_invalid_estim!(
"spatial operator dimension mismatch for term '{}': D0 cols={}, D1 cols={}, D2 cols={}, coeffs={}",
term_fit.name,
d0.ncols(),
d1.ncols(),
d2.ncols(),
coeff_global_range.len()
);
}
out.push(SpatialOperatorRuntimeCache {
termname: term_fit.name.clone(),
feature_cols,
coeff_global_range,
mass_penalty_global_idx: mass_global_idx,
tension_penalty_global_idx: tension_global_idx,
stiffness_penalty_global_idx: stiffness_global_idx,
d0,
d1,
d2,
collocation_points,
dimension: dim,
});
}
Ok(out)
}
fn scalar_operator_response_variance(
operator: &Array2<f64>,
cov_local: &Array2<f64>,
) -> Array1<f64> {
Array1::from_iter(operator.rows().into_iter().map(|row| {
let s = cov_local.dot(&row);
row.dot(&s).max(0.0)
}))
}
fn grouped_operator_response_variance(
operator: &Array2<f64>,
block_dim: usize,
cov_local: &Array2<f64>,
) -> Result<Array1<f64>, EstimationError> {
if block_dim == 0 || !operator.nrows().is_multiple_of(block_dim) {
crate::bail_invalid_estim!(
"grouped variance row layout invalid: rows={}, block_dim={block_dim}",
operator.nrows()
);
}
let p = operator.nrows() / block_dim;
let mut out = Array1::<f64>::zeros(p);
for k in 0..p {
let mut acc = 0.0;
for axis in 0..block_dim {
let row = operator.row(k * block_dim + axis);
let s = cov_local.dot(&row);
acc += row.dot(&s);
}
out[k] = acc.max(0.0);
}
Ok(out)
}
fn compute_spatial_adaptiveweights_for_beta(
beta: &Array1<f64>,
caches: &[SpatialOperatorRuntimeCache],
epsilon_0: f64,
epsilon_g: f64,
epsilon_c: f64,
weight_floor: f64,
weight_ceiling: f64,
beta_covariance: Option<&Array2<f64>>,
) -> Result<Vec<SpatialAdaptiveWeights>, EstimationError> {
caches
.iter()
.map(|cache| {
let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
let exact = SpatialPenaltyExactState::from_beta_local(
beta_local,
cache,
[epsilon_0, epsilon_g, epsilon_c],
)?;
let cov_local = beta_covariance.map(|cov| {
cov.slice(s![
cache.coeff_global_range.clone(),
cache.coeff_global_range.clone()
])
.to_owned()
});
let dim = cache.dimension;
let (var_0, var_g, var_c) = match cov_local.as_ref() {
Some(cov) => (
scalar_operator_response_variance(&cache.d0, cov),
grouped_operator_response_variance(&cache.d1, dim, cov)?,
grouped_operator_response_variance(&cache.d2, dim * dim, cov)?,
),
None => (
Array1::<f64>::zeros(exact.magnitude.signal.len()),
Array1::<f64>::zeros(exact.gradient.norm.len()),
Array1::<f64>::zeros(exact.curvature.norm.len()),
),
};
let (_, inv_0) = exact.magnitude.surrogateweights_posterior_snr(
&var_0,
weight_floor,
weight_ceiling,
);
let (_, inv_g) =
exact
.gradient
.surrogateweights_posterior_snr(&var_g, weight_floor, weight_ceiling);
let (_, inv_c) = exact.curvature.surrogateweights_posterior_snr(
&var_c,
weight_floor,
weight_ceiling,
);
Ok(SpatialAdaptiveWeights {
inv_magweight: inv_0,
invgradweight: inv_g,
inv_lapweight: inv_c,
})
})
.collect()
}
fn compute_initial_epsilons(
beta: &Array1<f64>,
caches: &[SpatialOperatorRuntimeCache],
min_epsilon: f64,
) -> Result<(f64, f64, f64), EstimationError> {
let mut fvals = Vec::<f64>::new();
let mut gvals = Vec::<f64>::new();
let mut cvals = Vec::<f64>::new();
for cache in caches {
let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
let exact = SpatialPenaltyExactState::from_beta_local(
beta_local,
cache,
[min_epsilon, min_epsilon, min_epsilon],
)?;
let (f, g, c) = exact.absolute_collocation_magnitudes();
fvals.extend(f.iter().copied());
gvals.extend(g.iter().copied());
cvals.extend(c.iter().copied());
}
let eps_0 = robust_epsilon_from_samples(&fvals, min_epsilon);
let eps_g = robust_epsilon_from_samples(&gvals, min_epsilon);
let eps_c = robust_epsilon_from_samples(&cvals, min_epsilon);
Ok((eps_0, eps_g, eps_c))
}
fn exact_spatial_adaptive_penalty_index_set(
caches: &[SpatialOperatorRuntimeCache],
) -> BTreeSet<usize> {
let mut out = BTreeSet::new();
for cache in caches {
out.insert(cache.mass_penalty_global_idx);
out.insert(cache.tension_penalty_global_idx);
out.insert(cache.stiffness_penalty_global_idx);
}
out
}
fn checked_fit_log_lambdas(
lambdas: &Array1<f64>,
context: &str,
) -> Result<Array1<f64>, EstimationError> {
let values = lambdas
.iter()
.copied()
.enumerate()
.map(|(coordinate, lambda)| {
gam_problem::checked_log_strength(lambda).map_err(|error| {
EstimationError::InvalidInput(format!(
"{context} lambda coordinate {coordinate} is outside the canonical physical-strength domain: {error}"
))
})
})
.collect::<Result<Vec<_>, _>>()?;
Ok(Array1::from_vec(values))
}
fn build_spatial_adaptive_hyperspecs(cache_count: usize) -> Vec<SpatialAdaptiveHyperSpec> {
let mut out = Vec::with_capacity(cache_count * 3 + 3);
for cache_index in 0..cache_count {
out.push(SpatialAdaptiveHyperSpec {
cache_index,
kind: SpatialAdaptiveHyperKind::LogLambdaMagnitude,
});
out.push(SpatialAdaptiveHyperSpec {
cache_index,
kind: SpatialAdaptiveHyperKind::LogLambdaGradient,
});
out.push(SpatialAdaptiveHyperSpec {
cache_index,
kind: SpatialAdaptiveHyperKind::LogLambdaCurvature,
});
}
out.push(SpatialAdaptiveHyperSpec {
cache_index: 0,
kind: SpatialAdaptiveHyperKind::LogEpsilonMagnitude,
});
out.push(SpatialAdaptiveHyperSpec {
cache_index: 0,
kind: SpatialAdaptiveHyperKind::LogEpsilonGradient,
});
out.push(SpatialAdaptiveHyperSpec {
cache_index: 0,
kind: SpatialAdaptiveHyperKind::LogEpsilonCurvature,
});
out
}
fn penalty_matrixwith_local_block(
total_dim: usize,
coeff_range: Range<usize>,
local: &Array2<f64>,
) -> Array2<f64> {
let mut out = Array2::<f64>::zeros((total_dim, total_dim));
out.slice_mut(s![coeff_range.clone(), coeff_range])
.assign(local);
out
}
fn fit_term_collectionwith_exact_spatial_adaptive_regularization(
baseline: FittedTermCollection,
y: ArrayView1<'_, f64>,
weights: ArrayView1<'_, f64>,
offset: ArrayView1<'_, f64>,
family: LikelihoodSpec,
options: &FitOptions,
runtime_caches: &[SpatialOperatorRuntimeCache],
) -> Result<FittedTermCollection, EstimationError> {
let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
let adaptive_penalty_indices = exact_spatial_adaptive_penalty_index_set(runtime_caches);
let p_total = baseline.design.design.ncols();
if baseline.fit.lambdas.len() != baseline.design.penalties.len() {
crate::bail_invalid_estim!(
"exact spatial adaptive fit received {} baseline lambdas for {} penalties",
baseline.fit.lambdas.len(),
baseline.design.penalties.len(),
);
}
let baseline_log_lambdas =
checked_fit_log_lambdas(&baseline.fit.lambdas, "exact spatial adaptive baseline")?;
for (cache_idx, cache) in runtime_caches.iter().enumerate() {
for (operator, penalty_idx) in [
("mass", cache.mass_penalty_global_idx),
("tension", cache.tension_penalty_global_idx),
("stiffness", cache.stiffness_penalty_global_idx),
] {
if penalty_idx >= baseline.fit.lambdas.len() {
crate::bail_invalid_estim!(
"exact spatial adaptive cache {cache_idx} {operator} penalty index {penalty_idx} is out of bounds for {} baseline lambdas",
baseline.fit.lambdas.len(),
);
}
}
}
struct RetainedPenaltySetup {
global_idx: usize,
global_penalty: Array2<f64>,
nullspace_dim: usize,
log_lambda: f64,
}
use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
let retained_setups = baseline
.design
.penalties
.par_iter()
.enumerate()
.map(|(idx, bp)| {
if adaptive_penalty_indices.contains(&idx) {
return None;
}
Some(RetainedPenaltySetup {
global_idx: idx,
global_penalty: bp.to_global(p_total),
nullspace_dim: baseline
.design
.nullspace_dims
.get(idx)
.copied()
.unwrap_or(0),
log_lambda: baseline_log_lambdas[idx],
})
})
.collect::<Vec<_>>();
let retained_count = retained_setups
.iter()
.filter(|setup| setup.is_some())
.count();
let mut retained_penalties = Vec::<Array2<f64>>::with_capacity(retained_count);
let mut retained_nullspace_dims = Vec::<usize>::with_capacity(retained_count);
let mut retained_log_lambdas = Vec::<f64>::with_capacity(retained_count);
let mut retained_global_indices = Vec::<usize>::with_capacity(retained_count);
for setup in retained_setups.into_iter().flatten() {
retained_penalties.push(setup.global_penalty);
retained_nullspace_dims.push(setup.nullspace_dim);
retained_log_lambdas.push(setup.log_lambda);
retained_global_indices.push(setup.global_idx);
}
let (eps_0_init, eps_g_init, eps_c_init) = compute_initial_epsilons(
&baseline.fit.beta,
runtime_caches,
adaptive_opts.min_epsilon,
)?;
let mut initial_theta =
Array1::<f64>::zeros(retained_penalties.len() + runtime_caches.len() * 3 + 3);
for (idx, value) in retained_log_lambdas.iter().enumerate() {
initial_theta[idx] = *value;
}
let adaptive_log_lambda_components = runtime_caches
.par_iter()
.map(|cache| {
[
baseline_log_lambdas[cache.mass_penalty_global_idx],
baseline_log_lambdas[cache.tension_penalty_global_idx],
baseline_log_lambdas[cache.stiffness_penalty_global_idx],
]
})
.collect::<Vec<_>>();
let mut at = retained_penalties.len();
for logs in &adaptive_log_lambda_components {
initial_theta[at] = logs[0];
initial_theta[at + 1] = logs[1];
initial_theta[at + 2] = logs[2];
at += 3;
}
let minimum_log_epsilon = gam_problem::checked_log_strength(adaptive_opts.min_epsilon)
.map_err(|error| {
EstimationError::InvalidInput(format!(
"adaptive minimum epsilon is outside the canonical positive-strength domain: {error}"
))
})?;
for (slot, epsilon) in [eps_0_init, eps_g_init, eps_c_init].into_iter().enumerate() {
initial_theta[at + slot] =
gam_problem::checked_log_strength(epsilon.max(adaptive_opts.min_epsilon)).map_err(
|error| {
EstimationError::InvalidInput(format!(
"adaptive initial epsilon coordinate {slot} is outside the canonical positive-strength domain: {error}"
))
},
)?;
}
let hyperspecs = build_spatial_adaptive_hyperspecs(runtime_caches.len());
let zero_psi_op: std::sync::Arc<dyn gam_custom_family::CustomFamilyPsiDerivativeOperator> =
std::sync::Arc::new(gam_custom_family::ZeroPsiDerivativeOperator::new(
baseline.design.design.nrows(),
baseline.design.design.ncols(),
));
let derivative_blocks = vec![
hyperspecs
.par_iter()
.map(|_| CustomFamilyBlockPsiDerivative {
penalty_index: None,
x_psi: Array2::<f64>::zeros((0, 0)),
s_psi: Array2::<f64>::zeros((0, 0)),
s_psi_components: None,
s_psi_penalty_components: None,
x_psi_psi: None,
s_psi_psi: None,
s_psi_psi_components: None,
s_psi_psi_penalty_components: None,
implicit_operator: Some(std::sync::Arc::clone(&zero_psi_op)),
implicit_axis: 0,
implicit_group_id: None,
})
.collect::<Vec<_>>(),
];
let mixture_link_state = options
.mixture_link
.clone()
.as_ref()
.map(state_fromspec)
.transpose()
.map_err(EstimationError::InvalidInput)?;
let sas_link_state = options
.sas_link
.map(|spec| {
if family.is_binomial_beta_logistic() {
state_from_beta_logisticspec(spec)
} else {
state_from_sasspec(spec)
}
})
.transpose()
.map_err(EstimationError::InvalidInput)?;
let latent_cloglog_state = options.latent_cloglog;
let shared_y = Arc::new(y.to_owned());
let sharedweights = Arc::new(weights.to_owned());
let shared_design = baseline
.design
.design
.try_to_dense_arc("spatial adaptive exact hyperfit design")
.map_err(EstimationError::InvalidInput)?;
let shared_offset = Arc::new(offset.to_owned());
let shared_runtime_caches = Arc::new(runtime_caches.to_vec());
let shared_hyperspecs = Arc::new(hyperspecs.clone());
let zero_quadratic = ValidatedFixedQuadraticHessian::zero(
baseline.design.design.ncols(),
)
.map_err(EstimationError::InvalidInput)?;
let base_family = SpatialAdaptiveExactFamily {
family: family.clone(),
latent_cloglog_state,
mixture_link_state: mixture_link_state.clone(),
sas_link_state,
y: shared_y.clone(),
weights: sharedweights.clone(),
design: shared_design.clone(),
offset: shared_offset.clone(),
linear_constraints: baseline.design.linear_constraints.clone(),
runtime_caches: shared_runtime_caches.clone(),
adaptive_params: Vec::new(),
fixed_quadratic_hessian: zero_quadratic.clone(),
hyperspecs: shared_hyperspecs.clone(),
exact_eval_cache: Arc::new(Mutex::new(None)),
};
let rho_dim = retained_penalties.len();
let operator_slots_end = rho_dim + runtime_caches.len() * 3;
const UNIFIED_LOG_WINDOW: f64 = 6.0;
const RETAINED_LAMBDA_LOG_LOWER_FLOOR: f64 = -30.0;
const RETAINED_LAMBDA_LOG_UPPER_CAP: f64 = 30.0;
const OPERATOR_LAMBDA_LOG_LOWER_FLOOR: f64 = -10.0;
const OPERATOR_LAMBDA_LOG_UPPER_CAP: f64 = 30.0;
let epsilon_floor_log = minimum_log_epsilon;
let anchored_bound = |idx: usize, sign: f64| -> f64 {
let raw = initial_theta[idx] + sign * UNIFIED_LOG_WINDOW;
if idx < rho_dim {
raw.clamp(
RETAINED_LAMBDA_LOG_LOWER_FLOOR,
RETAINED_LAMBDA_LOG_UPPER_CAP,
)
} else if idx < operator_slots_end {
raw.clamp(
OPERATOR_LAMBDA_LOG_LOWER_FLOOR,
OPERATOR_LAMBDA_LOG_UPPER_CAP,
)
} else {
raw.clamp(epsilon_floor_log, gam_problem::LOG_STRENGTH_MAX)
}
};
let eps_lower =
Array1::from_iter((0..initial_theta.len()).map(|idx| anchored_bound(idx, -1.0)));
let eps_upper = Array1::from_iter((0..initial_theta.len()).map(|idx| anchored_bound(idx, 1.0)));
let blockspec = ParameterBlockSpec {
name: "eta".to_string(),
design: baseline.design.design.clone(),
offset: offset.to_owned(),
penalties: retained_penalties
.iter()
.cloned()
.map(PenaltyMatrix::Dense)
.collect(),
nullspace_dims: retained_nullspace_dims.clone(),
initial_log_lambdas: Array1::from_vec(retained_log_lambdas.clone()),
initial_beta: Some(baseline.fit.beta.clone()),
gauge_priority: 100,
jacobian_callback: None,
stacked_design: None,
stacked_offset: None,
};
let screening_cap = Arc::new(AtomicUsize::new(0));
let outer_opts = BlockwiseFitOptions {
inner_max_cycles: options.max_iter,
inner_tol: options.tol,
outer_max_iter: options.max_iter,
outer_tol: options.tol,
compute_covariance: false,
screening_max_inner_iterations: Some(Arc::clone(&screening_cap)),
..BlockwiseFitOptions::default()
};
use gam_problem::{DeclaredHessianForm, Derivative, HessianValue, OuterEval};
use gam_solve::rho_optimizer::OuterProblem;
struct SpatialAdaptiveOuterState {
warm_cache: Option<CustomFamilyWarmStart>,
terminal_mode: Option<(Array1<f64>, f64, CustomFamilyOwnedMode)>,
last_eval: Option<(
Array1<f64>,
f64,
Array1<f64>,
HessianValue,
CustomFamilyWarmStart,
)>,
}
struct DecodedSpatialAdaptiveTheta {
rho: Array1<f64>,
retained_lambdas: Array1<f64>,
adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
epsilon: [f64; 3],
}
let n_theta = initial_theta.len();
let theta_bounds = Some((eps_lower.clone(), eps_upper.clone()));
let decode_theta =
|theta: &Array1<f64>| -> Result<DecodedSpatialAdaptiveTheta, EstimationError> {
let physical = gam_problem::checked_exp_log_strengths(theta.iter().copied()).map_err(
|error| {
EstimationError::InvalidInput(format!(
"exact spatial adaptive outer coordinate is outside the canonical log-strength domain: {error}"
))
},
)?;
let rho = theta.slice(s![..rho_dim]).to_owned();
let retained_lambdas = Array1::from_vec(physical[..rho_dim].to_vec());
let adaptive_lambda_start = rho_dim;
let adaptive_lambda_end = adaptive_lambda_start + runtime_caches.len() * 3;
let eps = [
physical[adaptive_lambda_end],
physical[adaptive_lambda_end + 1],
physical[adaptive_lambda_end + 2],
];
let adaptive_params = runtime_caches
.iter()
.enumerate()
.map(|(cache_idx, _)| SpatialAdaptiveTermHyperParams {
lambda: [
physical[adaptive_lambda_start + cache_idx * 3],
physical[adaptive_lambda_start + cache_idx * 3 + 1],
physical[adaptive_lambda_start + cache_idx * 3 + 2],
],
epsilon: eps,
})
.collect::<Vec<_>>();
Ok(DecodedSpatialAdaptiveTheta {
rho,
retained_lambdas,
adaptive_params,
epsilon: eps,
})
};
let clamp_theta = |theta: &Array1<f64>| -> Array1<f64> {
Array1::from_shape_fn(theta.len(), |i| theta[i].clamp(eps_lower[i], eps_upper[i]))
};
let realize_hyper_layout = |theta: &Array1<f64>| {
gam_custom_family::CustomFamilyHyperLayout::new(
derivative_blocks.clone(),
Vec::new(),
theta.slice(s![rho_dim..]).to_owned(),
)
.map_err(EstimationError::InvalidInput)
};
let analytic_outer_hessian_available =
gam_custom_family::joint_exact_analytic_outer_hessian_available()
&& base_family
.exact_outer_derivative_order(std::slice::from_ref(&blockspec), &outer_opts)
.has_hessian()
&& gam_custom_family::exact_newton_outer_geometry_supports_second_order_solver(
&base_family,
);
let problem = OuterProblem::new(n_theta)
.with_gradient(Derivative::Analytic)
.with_hessian(if analytic_outer_hessian_available {
DeclaredHessianForm::Either
} else {
DeclaredHessianForm::Unavailable
})
.with_prefer_gradient_only(true)
.with_fallback_policy(gam_solve::rho_optimizer::FallbackPolicy::Disabled)
.with_psi_dim(n_theta.saturating_sub(rho_dim))
.with_tolerance(options.tol)
.with_rel_cost_tolerance(Some(options.tol))
.with_max_iter(options.max_iter)
.with_seed_config(gam_problem::SeedConfig::default())
.with_screening_cap(Arc::clone(&screening_cap))
.with_initial_rho(initial_theta.clone());
let problem = if let Some((lo, hi)) = theta_bounds {
problem.with_bounds(lo, hi)
} else {
problem
};
let eval_outer = |st: &mut SpatialAdaptiveOuterState,
theta: &Array1<f64>,
order: gam_solve::rho_optimizer::OuterEvalOrder|
-> Result<OuterEval, EstimationError> {
let decoded = decode_theta(theta)?;
if let Some((cached_theta, cached_cost, cached_grad, cached_hess, cached_warm)) =
&st.last_eval
&& cached_theta.len() == theta.len()
&& cached_theta
.iter()
.zip(theta.iter())
.all(|(&a, &b)| a.to_bits() == b.to_bits())
&& st
.terminal_mode
.as_ref()
.is_some_and(|(mode_theta, mode_objective, _)| {
mode_theta.len() == theta.len()
&& mode_theta
.iter()
.zip(theta.iter())
.all(|(&a, &b)| a.to_bits() == b.to_bits())
&& mode_objective.to_bits() == cached_cost.to_bits()
})
&& (!matches!(
order,
gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
) || analytic_outer_hessian_available)
{
st.warm_cache = Some(cached_warm.clone());
return Ok(OuterEval {
cost: *cached_cost,
gradient: cached_grad.clone(),
hessian: if matches!(
order,
gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
) && analytic_outer_hessian_available
{
cached_hess.clone()
} else {
HessianValue::Unavailable
},
inner_beta_hint: None,
});
}
let family_eval =
base_family.with_adaptive_params(decoded.adaptive_params, zero_quadratic.clone());
let hyper_layout = realize_hyper_layout(theta)?;
let need_hessian = matches!(
order,
gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
) && analytic_outer_hessian_available;
let owned = evaluate_custom_family_joint_hyper_owned(
&family_eval,
std::slice::from_ref(&blockspec),
&outer_opts,
&decoded.rho,
&hyper_layout,
st.warm_cache.as_ref(),
if need_hessian {
gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueGradientHessian
} else {
gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueAndGradient
},
)
.map_err(|e| {
EstimationError::RemlOptimizationFailed(format!("spatial adaptive eval failed: {e}"))
})?;
if !owned.result.inner_converged {
st.warm_cache = Some(owned.result.warm_start.clone());
return Err(EstimationError::RemlOptimizationFailed(
"exact spatial adaptive inner solve did not converge".to_string(),
));
}
if !owned.result.objective.is_finite()
|| owned.result.gradient.iter().any(|v| !v.is_finite())
{
return Err(EstimationError::RemlOptimizationFailed(
"exact spatial adaptive objective returned non-finite values".to_string(),
));
}
let hessian_result = if need_hessian {
if !owned.result.outer_hessian.is_analytic() {
return Err(EstimationError::RemlOptimizationFailed(
"exact spatial adaptive objective did not return an exact outer Hessian"
.to_string(),
));
}
match owned.result.outer_hessian.dim() {
Some(dim) if dim == theta.len() => {}
Some(dim) => {
return Err(EstimationError::RemlOptimizationFailed(format!(
"exact spatial adaptive outer Hessian dimension mismatch: got {dim}, expected {}",
theta.len(),
)));
}
None => {
return Err(EstimationError::RemlOptimizationFailed(
"exact spatial adaptive objective did not report an outer Hessian dimension"
.to_string(),
));
}
}
st.last_eval = Some((
theta.to_owned(),
owned.result.objective,
owned.result.gradient.clone(),
owned.result.outer_hessian.clone(),
owned.result.warm_start.clone(),
));
owned.result.outer_hessian
} else {
HessianValue::Unavailable
};
let objective = owned.result.objective;
let gradient = owned.result.gradient;
st.warm_cache = Some(owned.result.warm_start);
st.terminal_mode = Some((theta.to_owned(), objective, owned.mode));
Ok(OuterEval {
cost: objective,
gradient,
hessian: hessian_result,
inner_beta_hint: None,
})
};
let mut obj = problem.build_objective_with_screening_proxy(
SpatialAdaptiveOuterState {
warm_cache: None,
terminal_mode: None,
last_eval: None,
},
|st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
let theta = clamp_theta(theta);
let DecodedSpatialAdaptiveTheta {
rho,
adaptive_params,
..
} = decode_theta(&theta)?;
let family_eval =
base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
let hyper_layout = realize_hyper_layout(&theta)?;
let owned = evaluate_custom_family_joint_hyper_owned(
&family_eval,
std::slice::from_ref(&blockspec),
&outer_opts,
&rho,
&hyper_layout,
st.warm_cache.as_ref(),
gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
)
.map_err(|e| {
EstimationError::RemlOptimizationFailed(format!(
"spatial adaptive cost eval failed: {e}"
))
})?;
if !owned.result.inner_converged {
st.warm_cache = Some(owned.result.warm_start);
return Err(EstimationError::RemlOptimizationFailed(
"exact spatial adaptive cost inner solve did not converge".to_string(),
));
}
let objective = owned.result.objective;
st.warm_cache = Some(owned.result.warm_start);
st.terminal_mode = Some((theta, objective, owned.mode));
Ok(objective)
},
|st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
eval_outer(
st,
theta,
if analytic_outer_hessian_available {
gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
} else {
gam_solve::rho_optimizer::OuterEvalOrder::ValueAndGradient
},
)
},
|st: &mut SpatialAdaptiveOuterState,
theta: &Array1<f64>,
order: gam_solve::rho_optimizer::OuterEvalOrder| { eval_outer(st, theta, order) },
Some(|st: &mut SpatialAdaptiveOuterState| {
st.warm_cache = None;
st.terminal_mode = None;
st.last_eval = None;
}),
Some(|st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
let theta = clamp_theta(theta);
let DecodedSpatialAdaptiveTheta {
rho,
adaptive_params,
..
} = decode_theta(&theta)?;
let family_eval =
base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
let hyper_layout = realize_hyper_layout(&theta)?;
let owned = evaluate_custom_family_joint_hyper_efs_owned(
&family_eval,
std::slice::from_ref(&blockspec),
&outer_opts,
&rho,
&hyper_layout,
st.warm_cache.as_ref(),
)
.map_err(|e| {
EstimationError::RemlOptimizationFailed(format!(
"spatial adaptive EFS eval failed: {e}"
))
})?;
if !owned.result.inner_converged {
st.warm_cache = Some(owned.result.warm_start);
return Err(EstimationError::RemlOptimizationFailed(
"exact spatial adaptive EFS inner solve did not converge".to_string(),
));
}
let objective = owned.result.efs_eval.cost;
st.warm_cache = Some(owned.result.warm_start);
st.terminal_mode = Some((theta, objective, owned.mode));
Ok(owned.result.efs_eval)
}),
|st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
let theta = clamp_theta(theta);
let DecodedSpatialAdaptiveTheta {
rho,
adaptive_params,
..
} = decode_theta(&theta)?;
let family_eval =
base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
let hyper_layout = realize_hyper_layout(&theta)?;
let owned = evaluate_custom_family_joint_hyper_owned(
&family_eval,
std::slice::from_ref(&blockspec),
&outer_opts,
&rho,
&hyper_layout,
st.warm_cache.as_ref(),
gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
)
.map_err(|e| {
EstimationError::RemlOptimizationFailed(format!(
"spatial adaptive screening eval failed: {e}"
))
})?;
st.warm_cache = Some(owned.result.warm_start);
Ok(owned.result.objective)
},
);
let certified_outer = problem
.run_certified(&mut obj, "exact spatial adaptive regularization")
.map_err(|e| {
EstimationError::InvalidInput(format!(
"exact spatial adaptive outer optimization failed: {e}"
))
})?;
let outer_iterations = certified_outer.iterations();
let outer_grad_norm = certified_outer.final_grad_norm();
let theta_star = certified_outer.rho().clone();
let (mode_theta, mode_objective, terminal_mode) =
obj.state.terminal_mode.take().ok_or_else(|| {
EstimationError::InvalidInput(
"exact spatial adaptive optimization certified without retaining its terminal coefficient mode"
.to_string(),
)
})?;
if mode_theta.len() != theta_star.len()
|| mode_theta
.iter()
.zip(theta_star.iter())
.any(|(mode, certified)| mode.to_bits() != certified.to_bits())
{
return Err(EstimationError::InvalidInput(
"exact spatial adaptive terminal coefficient mode does not bitwise match the certified hyperparameter vector"
.to_string(),
));
}
if mode_objective.to_bits() != certified_outer.final_value().to_bits() {
return Err(EstimationError::InvalidInput(format!(
"exact spatial adaptive terminal coefficient mode objective does not bitwise match the certified objective: mode={mode_objective:.17e}, certified={:.17e}",
certified_outer.final_value(),
)));
}
let DecodedSpatialAdaptiveTheta {
rho: _,
retained_lambdas,
adaptive_params,
epsilon: eps_star,
} = decode_theta(&theta_star)?;
let mut fixed_total = Array2::<f64>::zeros((
baseline.design.design.ncols(),
baseline.design.design.ncols(),
));
for (idx, penalty) in retained_penalties.iter().enumerate() {
fixed_total.scaled_add(retained_lambdas[idx], penalty);
}
let certified_final_family = base_family.with_adaptive_params(
adaptive_params.clone(),
zero_quadratic.clone(),
);
let fixed_total = ValidatedFixedQuadraticHessian::try_from_dense(
fixed_total,
baseline.design.design.ncols(),
)
.map_err(|error| {
EstimationError::InvalidInput(format!(
"optimized spatial adaptive fixed quadratic Hessian is invalid: {error}"
))
})?;
let final_family =
base_family.with_adaptive_params(adaptive_params.clone(), fixed_total.clone());
let final_blockspec = ParameterBlockSpec {
name: "eta".to_string(),
design: baseline.design.design.clone(),
offset: offset.to_owned(),
penalties: retained_penalties
.iter()
.cloned()
.map(PenaltyMatrix::Dense)
.collect(),
nullspace_dims: retained_nullspace_dims.clone(),
initial_log_lambdas: theta_star.slice(s![..rho_dim]).to_owned(),
initial_beta: Some(baseline.fit.beta.clone()),
gauge_priority: 100,
jacobian_callback: None,
stacked_design: None,
stacked_offset: None,
};
let final_fit = fit_custom_family_fixed_log_lambdas_from_owned_mode(
&certified_final_family,
&[final_blockspec],
&BlockwiseFitOptions {
inner_max_cycles: options.max_iter,
inner_tol: options.tol,
outer_max_iter: 1,
outer_tol: options.tol,
compute_covariance: true,
..BlockwiseFitOptions::default()
},
terminal_mode,
&theta_star,
&certified_outer,
)
.map_err(EstimationError::CustomFamily)?;
let beta = final_fit.block_states[0].beta.clone();
let final_eval = final_family
.exact_evaluation(&beta)
.map_err(EstimationError::InvalidInput)?;
let penalized_hessian = final_eval
.totalobjectivehessian(&final_family.design)
.map_err(EstimationError::InvalidInput)?;
let beta_covariance = final_fit.covariance_conditional.clone();
let beta_standard_errors = beta_covariance
.as_ref()
.map(|cov| Array1::from_iter((0..cov.nrows()).map(|i| cov[[i, i]].max(0.0).sqrt())));
let mut full_lambdas = baseline.fit.lambdas.clone();
for (idx, &global_idx) in retained_global_indices.iter().enumerate() {
full_lambdas[global_idx] = retained_lambdas[idx];
}
for (cache_idx, cache) in runtime_caches.iter().enumerate() {
full_lambdas[cache.mass_penalty_global_idx] = adaptive_params[cache_idx].lambda[0];
full_lambdas[cache.tension_penalty_global_idx] = adaptive_params[cache_idx].lambda[1];
full_lambdas[cache.stiffness_penalty_global_idx] = adaptive_params[cache_idx].lambda[2];
}
let deviance = -2.0 * final_eval.obs.log_likelihood;
let mut local_penalty_blocks =
Vec::<PenaltySpec>::with_capacity(baseline.design.penalties.len());
for (global_idx, bp) in baseline.design.penalties.iter().enumerate() {
if adaptive_penalty_indices.contains(&global_idx) {
let cache = runtime_caches
.iter()
.find(|cache| {
cache.mass_penalty_global_idx == global_idx
|| cache.tension_penalty_global_idx == global_idx
|| cache.stiffness_penalty_global_idx == global_idx
})
.ok_or_else(|| {
EstimationError::InvalidInput(format!(
"missing runtime cache for adaptive penalty index {global_idx}"
))
})?;
let cache_idx = runtime_caches
.iter()
.position(|c| {
c.mass_penalty_global_idx == global_idx
|| c.tension_penalty_global_idx == global_idx
|| c.stiffness_penalty_global_idx == global_idx
})
.ok_or_else(|| {
EstimationError::InvalidInput(format!(
"missing adaptive cache position for penalty index {global_idx}"
))
})?;
let state = &final_eval.adaptive_states[cache_idx];
let local = if cache.mass_penalty_global_idx == global_idx {
scalar_operatorhessian(&cache.d0, &state.magnitude.betahessian_diag())
.mapv(|v| adaptive_params[cache_idx].lambda[0] * v)
} else if cache.tension_penalty_global_idx == global_idx {
grouped_operatorhessian(
&cache.d1,
cache.dimension,
&state.gradient.betahessian_blocks(),
)?
.mapv(|v| adaptive_params[cache_idx].lambda[1] * v)
} else {
grouped_operatorhessian(
&cache.d2,
cache.dimension * cache.dimension,
&state.curvature.betahessian_blocks(),
)?
.mapv(|v| adaptive_params[cache_idx].lambda[2] * v)
};
local_penalty_blocks.push(PenaltySpec::Dense(penalty_matrixwith_local_block(
baseline.design.design.ncols(),
cache.coeff_global_range.clone(),
&local,
)));
} else {
local_penalty_blocks.push(PenaltySpec::Dense(
bp.to_global(p_total).mapv(|v| v * full_lambdas[global_idx]),
));
}
}
let (edf_by_block, penalty_block_trace, edf_total) = if let Some(cov) = beta_covariance.as_ref()
{
exact_bounded_edf(
&local_penalty_blocks,
&Array1::from_elem(local_penalty_blocks.len(), 1.0),
cov,
)?
} else {
(
vec![0.0; local_penalty_blocks.len()],
vec![0.0; local_penalty_blocks.len()],
0.0,
)
};
let stable_penalty_term = 2.0 * final_eval.adaptive_penalty_value
+ beta.dot(&fixed_total.as_dense().dot(&beta));
let standard_deviation = if family.is_gaussian_identity() {
let denom = (y.len() as f64 - edf_total).max(1.0);
(deviance / denom).sqrt()
} else {
1.0
};
let maps = compute_spatial_adaptiveweights_for_beta(
&beta,
runtime_caches,
eps_star[0],
eps_star[1],
eps_star[2],
adaptive_opts.weight_floor,
adaptive_opts.weight_ceiling,
beta_covariance.as_ref(),
)?
.into_iter()
.zip(runtime_caches.iter())
.map(|(w, cache)| AdaptiveSpatialMap {
termname: cache.termname.clone(),
feature_cols: cache.feature_cols.clone(),
collocation_points: cache.collocation_points.clone(),
inv_magweight: w.inv_magweight,
invgradweight: w.invgradweight,
inv_lapweight: w.inv_lapweight,
})
.collect::<Vec<_>>();
let fitted_link = if family.is_latent_cloglog() {
FittedLinkState::LatentCLogLog {
state: latent_cloglog_state
.expect("BinomialLatentCLogLog requires an explicit latent-cloglog state"),
}
} else if family.is_binomial_mixture() {
mixture_link_state
.clone()
.map(|state| FittedLinkState::Mixture {
state,
covariance: None,
})
.unwrap_or(FittedLinkState::Standard(None))
} else if family.is_binomial_sas() {
sas_link_state
.map(|state| FittedLinkState::Sas {
state,
covariance: None,
})
.unwrap_or(FittedLinkState::Standard(None))
} else if family.is_binomial_beta_logistic() {
sas_link_state
.map(|state| FittedLinkState::BetaLogistic {
state,
covariance: None,
})
.unwrap_or(FittedLinkState::Standard(None))
} else {
FittedLinkState::Standard(None)
};
let max_abs_eta = final_eval
.obs
.eta
.iter()
.fold(0.0_f64, |acc, &v| acc.max(v.abs()));
let fitted = FittedTermCollection {
fit: {
let log_lambdas =
checked_fit_log_lambdas(&full_lambdas, "final exact spatial adaptive fit")?;
let working = gam_solve::estimate::WorkingGeometry {
weights: final_eval.obs.fisherweight.clone(),
response: exact_standard_working_response(&final_eval.obs)?,
};
let inf = FitInference {
edf_by_block,
penalty_block_trace,
edf_total,
smoothing_correction: None,
smoothing_correction_method: None,
smoothing_correction_first_order: None,
smoothing_correction_method_first_order: None,
penalized_hessian: penalized_hessian.clone().into(),
reparam_qs: None,
dispersion: gam_solve::estimate::Dispersion::UNIT,
beta_covariance: beta_covariance
.clone()
.map(gam_problem::dispersion_cov::PhiScaledCovariance::from),
beta_standard_errors,
beta_covariance_corrected: None,
beta_standard_errors_corrected: None,
beta_covariance_frequentist: None,
coefficient_influence: None,
weighted_gram: None,
bias_correction_beta: None,
bias_correction_jacobian: None,
};
let geometry = Some(gam_solve::estimate::FitGeometry {
coefficient_gauge: gam_problem::gauge::Gauge::identity(&[beta.len()]),
penalized_hessian: penalized_hessian.into(),
constrained_posterior: None,
working: Some(working),
});
let covariance_conditional = beta_covariance;
let convergence = final_fit.convergence_evidence();
let pirls_status_val = convergence.inner_status();
let certified_outer_present = convergence.outer_certificate().is_some();
UnifiedFitResult::try_from_parts(UnifiedFitResultParts {
blocks: vec![gam_solve::estimate::FittedBlock {
beta: beta.clone(),
role: gam_problem::BlockRole::Mean,
edf: edf_total,
lambdas: full_lambdas.clone(),
}],
log_lambdas,
lambdas: full_lambdas,
likelihood_scale: family.default_scale_metadata(),
likelihood_family: Some(family),
log_likelihood_normalization: gam_spec::LogLikelihoodNormalization::UserProvided,
log_likelihood: final_eval.obs.log_likelihood,
deviance,
reml_score: final_fit.penalized_objective,
stable_penalty_term,
penalized_objective: final_fit.penalized_objective,
used_device: false,
outer_iterations,
outer_converged: certified_outer_present,
outer_gradient_norm: outer_grad_norm,
standard_deviation,
covariance_conditional,
covariance_corrected: None,
inference: Some(inf),
fitted_link,
geometry,
block_states: Vec::new(),
pirls_status: pirls_status_val,
max_abs_eta,
constraint_kkt: None,
artifacts: gam_solve::estimate::FitArtifacts {
pirls: None,
criterion_certificate: final_fit.artifacts.criterion_certificate.clone(),
..Default::default()
},
inner_cycles: 0,
})?
},
design: baseline.design,
adaptive_diagnostics: Some(AdaptiveRegularizationDiagnostics {
epsilon_0: eps_star[0],
epsilon_g: eps_star[1],
epsilon_c: eps_star[2],
epsilon_outer_iterations: outer_iterations,
mm_iterations: 0,
converged: true,
maps,
}),
};
enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
Ok(fitted)
}
fn relax_smoothing_rho_prior(
options: &FitOptions,
design: &TermCollectionDesign,
y: ArrayView1<'_, f64>,
weights: ArrayView1<'_, f64>,
) -> gam_spec::RhoPrior {
use gam_terms::basis::BasisMetadata;
let base = &options.rho_prior;
if !base.is_unset() {
return base.clone();
}
let has_link_aux = options.sas_link.is_some()
|| options.optimize_sas
|| options.mixture_link.is_some()
|| options.optimize_mixture;
let has_moving_kappa = design.smooth.terms.iter().any(|t| {
if let BasisMetadata::Duchon {
length_scale,
aniso_log_scales,
..
} = &t.metadata
{
return length_scale.is_some() || aniso_log_scales.is_some();
}
matches!(
t.metadata,
BasisMetadata::Matern { .. }
| BasisMetadata::Sphere { .. }
| BasisMetadata::SphereHarmonics { .. }
| BasisMetadata::ConstantCurvature { .. }
| BasisMetadata::MeasureJet { .. }
)
});
let length_safe = !has_link_aux && !has_moving_kappa;
if !length_safe {
return base.clone();
}
let coords = &design.penaltyinfo;
if coords.is_empty() {
return base.clone();
}
let n_obs = design.design.nrows();
let p_total = design.design.ncols();
let underdetermined = n_obs < 2 * p_total;
let relaxable_terms: std::collections::HashSet<&str> = design
.smooth
.terms
.iter()
.filter(|t| {
(matches!(
t.metadata,
BasisMetadata::BSpline1D { .. }
| BasisMetadata::ThinPlate { .. }
| BasisMetadata::TensorBSpline { .. }
)
|| matches!(
t.metadata,
BasisMetadata::Duchon {
length_scale: None,
aniso_log_scales: None,
..
}
))
&& matches!(t.shape, gam_terms::smooth::ShapeConstraint::None)
})
.map(|t| t.name.as_str())
.collect();
let any_relaxed = coords.iter().any(|info| {
info.termname
.as_deref()
.is_some_and(|name| relaxable_terms.contains(name))
});
if !any_relaxed {
return base.clone();
}
let relaxed_prior = if underdetermined {
gam_spec::RhoPrior::Normal {
mean: 0.0,
sd: RELAX_UNDERDETERMINED_RHO_SD,
}
} else {
gam_spec::RhoPrior::Flat
};
let nullspace_select_prior = gam_spec::RhoPrior::PenalizedComplexity {
upper: NULLSPACE_SELECT_PC_UPPER,
tail_prob: NULLSPACE_SELECT_PC_TAIL_PROB,
};
let nullspace_degeneracy_prior = gam_spec::RhoPrior::Normal {
mean: 0.0,
sd: NULLSPACE_WELLDET_DEGENERACY_RHO_SD,
};
let per_coord = coords
.iter()
.enumerate()
.map(|(coord_idx, info)| {
let relax = info
.termname
.as_deref()
.is_some_and(|name| relaxable_terms.contains(name));
if !relax {
return base.clone();
}
let is_nullspace = matches!(info.penalty.source, PenaltySource::DoublePenaltyNullspace);
if is_nullspace {
if underdetermined
&& !nullspace_directions_are_supported(design, coord_idx, y, weights)
{
nullspace_select_prior.clone()
} else {
nullspace_degeneracy_prior.clone()
}
} else {
relaxed_prior.clone()
}
})
.collect::<Vec<_>>();
gam_spec::RhoPrior::Independent(per_coord)
}
const NULLSPACE_SUPPORT_FRACTION_THRESHOLD: f64 = 0.5;
fn nullspace_directions_are_supported(
design: &TermCollectionDesign,
penalty_idx: usize,
y: ArrayView1<'_, f64>,
weights: ArrayView1<'_, f64>,
) -> bool {
use gam_linalg::faer_ndarray::FaerEigh;
let Some(pen) = design.penalties.get(penalty_idx) else {
return false;
};
let col_range = pen.col_range.clone();
if col_range.is_empty() {
return false;
}
let x = design.design.to_dense();
let n = x.nrows();
if n == 0 || y.len() != n || weights.len() != n || x.ncols() < col_range.end {
return false;
}
let mut resp = y.to_owned();
if design.affine_offset.len() == n {
resp -= &design.affine_offset;
}
let Ok((evals, evecs)) = pen.local.eigh(faer::Side::Lower) else {
return false;
};
let max_eig = evals.iter().cloned().fold(0.0_f64, |m, v| m.max(v));
if !(max_eig > 0.0) {
return false;
}
let tol = 1.0e-9 * max_eig;
let pos_cols: Vec<usize> = (0..evals.len()).filter(|&j| evals[j] > tol).collect();
if pos_cols.is_empty() {
return false;
}
let xblock = x.slice(s![.., col_range.clone()]);
let mut z = Array2::<f64>::zeros((n, pos_cols.len()));
for (out_j, &j) in pos_cols.iter().enumerate() {
let v = evecs.column(j);
if v.len() != xblock.ncols() {
return false;
}
z.column_mut(out_j).assign(&xblock.dot(&v));
}
let mut control: Vec<usize> = design.intercept_range.clone().collect();
for (_, r) in &design.linear_ranges {
control.extend(r.clone());
}
control.retain(|&c| c < x.ncols());
let mut cmat = Array2::<f64>::zeros((n, control.len().max(1)));
if control.is_empty() {
cmat.column_mut(0).fill(1.0);
} else {
for (out_j, &c) in control.iter().enumerate() {
cmat.column_mut(out_j).assign(&x.column(c));
}
}
let rss_c = weighted_regression_rss(cmat.view(), resp.view(), weights);
let Some(rss_c) = rss_c else { return false };
let base_scale = weighted_total_ss(resp.view(), weights);
if !(rss_c > 1.0e-12 * base_scale.max(f64::MIN_POSITIVE)) {
return false;
}
let mut cz = Array2::<f64>::zeros((n, cmat.ncols() + z.ncols()));
cz.slice_mut(s![.., ..cmat.ncols()]).assign(&cmat);
cz.slice_mut(s![.., cmat.ncols()..]).assign(&z);
let Some(rss_cz) = weighted_regression_rss(cz.view(), resp.view(), weights) else {
return false;
};
let support = (rss_c - rss_cz) / rss_c;
support.is_finite() && support > NULLSPACE_SUPPORT_FRACTION_THRESHOLD
}
fn weighted_total_ss(y: ArrayView1<'_, f64>, w: ArrayView1<'_, f64>) -> f64 {
let mut sw = 0.0;
let mut swy = 0.0;
for (&yi, &wi) in y.iter().zip(w.iter()) {
if wi > 0.0 && yi.is_finite() {
sw += wi;
swy += wi * yi;
}
}
if sw <= 0.0 {
return 0.0;
}
let mean = swy / sw;
let mut ss = 0.0;
for (&yi, &wi) in y.iter().zip(w.iter()) {
if wi > 0.0 && yi.is_finite() {
ss += wi * (yi - mean) * (yi - mean);
}
}
ss
}
fn weighted_regression_rss(
d: ArrayView2<'_, f64>,
y: ArrayView1<'_, f64>,
w: ArrayView1<'_, f64>,
) -> Option<f64> {
use gam_linalg::faer_ndarray::FaerCholesky;
let m = d.ncols();
if m == 0 {
return Some(weighted_total_ss(y, w));
}
let mut gram = Array2::<f64>::zeros((m, m));
let mut rhs = Array1::<f64>::zeros(m);
for row in 0..d.nrows() {
let wi = w[row];
if !(wi > 0.0) || !y[row].is_finite() {
continue;
}
let dr = d.row(row);
for a in 0..m {
let wda = wi * dr[a];
rhs[a] += wda * y[row];
for b in a..m {
gram[[a, b]] += wda * dr[b];
}
}
}
for a in 0..m {
for b in (a + 1)..m {
gram[[b, a]] = gram[[a, b]];
}
}
let trace = (0..m).map(|i| gram[[i, i]]).sum::<f64>();
if !(trace > 0.0) {
return Some(weighted_total_ss(y, w));
}
let ridge = 1.0e-10 * trace / (m as f64);
for i in 0..m {
gram[[i, i]] += ridge;
}
let chol = gram.cholesky(faer::Side::Lower).ok()?;
let beta = chol.solvevec(&rhs);
let mut rss = 0.0;
for row in 0..d.nrows() {
let wi = w[row];
if !(wi > 0.0) || !y[row].is_finite() {
continue;
}
let fitted = d.row(row).dot(&beta);
let resid = y[row] - fitted;
rss += wi * resid * resid;
}
Some(rss)
}
const RELAX_UNDERDETERMINED_RHO_SD: f64 = 15.0;
const NULLSPACE_SELECT_PC_UPPER: f64 = 0.05;
const NULLSPACE_SELECT_PC_TAIL_PROB: f64 = 0.01;
fn adaptive_fit_options_base(options: &FitOptions, design: &TermCollectionDesign) -> FitOptions {
FitOptions {
resource_policy: options.resource_policy.clone(),
latent_cloglog: options.latent_cloglog,
mixture_link: options.mixture_link.clone(),
optimize_mixture: options.optimize_mixture,
sas_link: options.sas_link,
optimize_sas: options.optimize_sas,
compute_inference: options.compute_inference,
skip_rho_posterior_inference: options.skip_rho_posterior_inference,
max_iter: options.max_iter,
tol: options.tol,
nullspace_dims: design.nullspace_dims.clone(),
linear_constraints: design.linear_constraints.clone(),
firth_bias_reduction: options.firth_bias_reduction,
adaptive_regularization: None,
penalty_shrinkage_floor: options.penalty_shrinkage_floor,
rho_prior: options.rho_prior.clone(),
kronecker_penalty_system: design.kronecker_penalty_system(),
kronecker_factored: design
.smooth
.terms
.iter()
.find_map(|t| t.kronecker_factored.clone()),
persist_warm_start_disk: options.persist_warm_start_disk,
}
}
fn superseded_fit_options(options: &FitOptions) -> FitOptions {
let mut fit_options = options.clone();
fit_options.skip_rho_posterior_inference = true;
fit_options
}
#[derive(Clone)]
struct BoundedLinearTermMeta {
col_idx: usize,
min: f64,
max: f64,
prior: BoundedCoefficientPriorSpec,
}
struct BoundedEffectiveJacobian {
design: Array2<f64>,
bounded_terms: Vec<BoundedLinearTermMeta>,
}
impl BlockEffectiveJacobian for BoundedEffectiveJacobian {
fn effective_jacobian_rows(
&self,
state: &FamilyLinearizationState<'_>,
rows: std::ops::Range<usize>,
) -> Result<Array2<f64>, String> {
let p = self.design.ncols();
let n = self.design.nrows();
let rows = rows.start.min(n)..rows.end.min(n);
if !state.beta.is_empty() {
if state.beta.len() != p {
return Err(format!(
"BoundedEffectiveJacobian::effective_jacobian_at: beta length {} != design \
ncols {p}",
state.beta.len(),
));
}
if state.beta.iter().any(|v| !v.is_finite()) {
return Err(
"BoundedEffectiveJacobian::effective_jacobian_at: beta contains a non-finite value"
.to_string(),
);
}
}
let mut jac = self
.design
.slice(ndarray::s![rows.start..rows.end, ..])
.to_owned();
for term in &self.bounded_terms {
if term.col_idx >= p {
return Err(format!(
"BoundedEffectiveJacobian::effective_jacobian_at: bounded column {} is outside {p} columns",
term.col_idx
));
}
let theta = if state.beta.is_empty() {
0.0
} else {
state.beta[term.col_idx]
};
let (_, _, db_dtheta, _, _) = bounded_latent_derivatives(theta, term.min, term.max);
if !(db_dtheta.is_finite() && db_dtheta > 0.0) {
return Err(format!(
"BoundedEffectiveJacobian::effective_jacobian_at: bounded column {} has unrepresentable derivative {db_dtheta} at theta={theta}",
term.col_idx
));
}
jac.column_mut(term.col_idx).mapv_inplace(|v| v * db_dtheta);
}
Ok(jac)
}
}
#[derive(Clone)]
struct BoundedLinearFamily {
likelihood: gam_spec::GlmLikelihoodSpec,
latent_cloglog_state: Option<LatentCLogLogState>,
mixture_link_state: Option<MixtureLinkState>,
sas_link_state: Option<SasLinkState>,
y: Array1<f64>,
weights: Array1<f64>,
design: Array2<f64>,
designzeroed: Array2<f64>,
offset: Array1<f64>,
bounded_terms: Vec<BoundedLinearTermMeta>,
}
#[derive(Clone, Debug)]
struct StandardFamilyObservationState {
eta: Array1<f64>,
score: Array1<f64>,
fisherweight: Array1<f64>,
neghessian_eta: Array1<f64>,
neghessian_eta_derivative: Array1<f64>,
log_likelihood: f64,
}
fn bounded_latent_to_user(theta: f64, min: f64, max: f64) -> (f64, f64, f64) {
let jet = logit_inverse_link_jet5(theta);
let z = jet.mu;
let width = max - min;
let beta = min + width * z;
let db_dtheta = width * jet.d1;
(beta, z, db_dtheta)
}
fn bounded_user_to_latent(beta: f64, min: f64, max: f64) -> f64 {
(beta - min).ln() - (max - beta).ln()
}
#[derive(Debug, Clone, Copy)]
pub struct BoundedSampleColumn {
pub col_idx: usize,
pub min: f64,
pub max: f64,
}
pub fn sample_bounded_latent_posterior_internal(
beta_user: &Array1<f64>,
user_hessian: &Array2<f64>,
bounded_columns: &[BoundedSampleColumn],
n_draws: usize,
sqrt_cov_scale: f64,
base_seed: u64,
) -> Result<Array2<f64>, EstimationError> {
let p = beta_user.len();
if user_hessian.nrows() != p || user_hessian.ncols() != p {
crate::bail_invalid_estim!(
"bounded posterior sampling dimension mismatch: mode has {p} entries, user Hessian is {}x{}",
user_hessian.nrows(),
user_hessian.ncols()
);
}
if beta_user.iter().any(|value| !value.is_finite()) {
crate::bail_invalid_estim!("bounded posterior sampling requires a finite mode");
}
if user_hessian.iter().any(|value| !value.is_finite()) {
crate::bail_invalid_estim!("bounded posterior sampling requires a finite Hessian");
}
if !(sqrt_cov_scale.is_finite() && sqrt_cov_scale >= 0.0) {
crate::bail_invalid_estim!(
"bounded posterior sampling covariance scale must be finite and non-negative, got {sqrt_cov_scale}"
);
}
let mut theta_mode = beta_user.clone();
let mut jac_diag = Array1::<f64>::ones(p);
for bc in bounded_columns {
if bc.col_idx >= p {
crate::bail_invalid_estim!(
"bounded posterior sampling: bounded column index {} out of range for {p} coefficients",
bc.col_idx
);
}
if !(bc.min.is_finite()
&& bc.max.is_finite()
&& (bc.max - bc.min).is_finite()
&& bc.min < beta_user[bc.col_idx]
&& beta_user[bc.col_idx] < bc.max)
{
crate::bail_invalid_estim!(
"bounded posterior sampling column {} requires finite bounds with a finite width and a mode strictly inside ({}, {}); got {}",
bc.col_idx,
bc.min,
bc.max,
beta_user[bc.col_idx]
);
}
let theta_i = bounded_user_to_latent(beta_user[bc.col_idx], bc.min, bc.max);
let (_, _, db_dtheta) = bounded_latent_to_user(theta_i, bc.min, bc.max);
if !(theta_i.is_finite() && db_dtheta.is_finite() && db_dtheta > 0.0) {
crate::bail_invalid_estim!(
"bounded posterior sampling column {} has unrepresentable latent geometry: theta={theta_i}, d_beta/d_theta={db_dtheta}",
bc.col_idx
);
}
theta_mode[bc.col_idx] = theta_i;
jac_diag[bc.col_idx] = db_dtheta;
}
let mut h_latent = user_hessian.clone();
for i in 0..p {
let ji = jac_diag[i];
if ji != 1.0 {
h_latent.row_mut(i).mapv_inplace(|v| v * ji);
h_latent.column_mut(i).mapv_inplace(|v| v * ji);
}
}
use gam_linalg::faer_ndarray::FaerCholesky as _;
use rand::SeedableRng as _;
let chol = h_latent.cholesky(faer::Side::Lower).map_err(|err| {
EstimationError::InvalidInput(format!(
"bounded posterior sampling: Cholesky of the latent penalized Hessian failed: {err:?}"
))
})?;
let l = chol.lower_triangular();
let mut draws = Array2::<f64>::zeros((n_draws, p));
let mut eps = Array1::<f64>::zeros(p);
let mut delta = Array1::<f64>::zeros(p);
let mut rng = rand::rngs::StdRng::seed_from_u64(base_seed);
for k in 0..n_draws {
for e in eps.iter_mut() {
*e = standard_normal_draw(&mut rng);
}
solve_lower_transpose_into(&l, &eps, &mut delta)?;
for i in 0..p {
draws[(k, i)] = theta_mode[i] + sqrt_cov_scale * delta[i];
}
for bc in bounded_columns {
let (beta_draw, _, _) = bounded_latent_to_user(draws[(k, bc.col_idx)], bc.min, bc.max);
draws[(k, bc.col_idx)] = beta_draw;
}
}
Ok(draws)
}
#[inline]
fn standard_normal_draw<R: rand::Rng + ?Sized>(rng: &mut R) -> f64 {
use rand::RngExt as _;
let u1 = loop {
let candidate = rng.random::<f64>();
if candidate > 0.0 {
break candidate;
}
};
let u2 = rng.random::<f64>();
(-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
}
fn solve_lower_transpose_into(
l: &Array2<f64>,
b: &Array1<f64>,
out: &mut Array1<f64>,
) -> Result<(), EstimationError> {
let p = l.nrows();
if l.ncols() != p || b.len() != p || out.len() != p {
crate::bail_invalid_estim!(
"bounded triangular solve dimension mismatch: L={}x{}, b={}, out={}",
l.nrows(),
l.ncols(),
b.len(),
out.len()
);
}
for i in (0..p).rev() {
let mut acc = b[i];
for j in (i + 1)..p {
acc -= l[(j, i)] * out[j];
}
let diag = l[(i, i)];
if !(diag.is_finite() && diag > 0.0 && acc.is_finite()) {
crate::bail_invalid_estim!(
"bounded triangular solve has invalid row {i}: diagonal={diag}, residual={acc}"
);
}
let value = acc / diag;
if !value.is_finite() {
crate::bail_invalid_estim!(
"bounded triangular solve produced a non-finite value at row {i}: {acc}/{diag}"
);
}
out[i] = value;
}
Ok(())
}
fn bounded_latent_derivatives(theta: f64, min: f64, max: f64) -> (f64, f64, f64, f64, f64) {
let jet = logit_inverse_link_jet5(theta);
let z = jet.mu;
let width = max - min;
let beta = min + width * z;
let db_dtheta = width * jet.d1;
let d2b_dtheta2 = width * jet.d2;
let d3b_dtheta3 = width * jet.d3;
(beta, z, db_dtheta, d2b_dtheta2, d3b_dtheta3)
}
fn bounded_prior_terms(
theta: f64,
prior: &BoundedCoefficientPriorSpec,
) -> Result<(f64, f64, f64, f64), String> {
if !theta.is_finite() {
return Err(format!(
"bounded coefficient prior requires a finite latent coordinate, got {theta}"
));
}
let (a, b) = match prior {
BoundedCoefficientPriorSpec::None => return Ok((0.0, 0.0, 0.0, 0.0)),
BoundedCoefficientPriorSpec::Uniform => (1.0, 1.0),
BoundedCoefficientPriorSpec::Beta { a, b } => (*a, *b),
};
if !(a.is_finite() && a > 0.0 && b.is_finite() && b > 0.0) {
return Err(format!(
"bounded coefficient Beta prior requires finite positive shapes, got ({a}, {b})"
));
}
let jet = logit_inverse_link_jet5(theta);
let z = jet.mu;
let logp = -a * gam_linalg::utils::stable_softplus(-theta)
- b * gam_linalg::utils::stable_softplus(theta);
let grad = a - (a + b) * z;
let neghess = (a + b) * jet.d1;
let neghess_derivative = (a + b) * jet.d2;
let terms = (logp, grad, neghess, neghess_derivative);
if [terms.0, terms.1, terms.2, terms.3]
.iter()
.any(|value| !value.is_finite())
{
return Err(format!(
"bounded coefficient prior geometry is not representable at theta={theta}: {terms:?}"
));
}
Ok(terms)
}
#[derive(Clone, Copy)]
struct ExactStandardObservationRow {
mu: f64,
score: f64,
fisherweight: f64,
neghessian_eta: f64,
neghessian_eta_derivative: f64,
log_likelihood: f64,
}
impl ExactStandardObservationRow {
#[inline]
fn zero_weight(mu: f64) -> Self {
Self {
mu,
score: 0.0,
fisherweight: 0.0,
neghessian_eta: 0.0,
neghessian_eta_derivative: 0.0,
log_likelihood: 0.0,
}
}
}
#[inline]
fn bounded_row_error(row: usize, quantity: &'static str, eta: f64, value: f64) -> EstimationError {
EstimationError::PirlsRowGeometryUnrepresentable {
row,
quantity,
eta,
value,
}
}
#[inline]
fn certify_bounded_row(
row: usize,
eta: f64,
state: ExactStandardObservationRow,
) -> Result<ExactStandardObservationRow, EstimationError> {
for (quantity, value) in [
("bounded-family mean", state.mu),
("bounded-family score", state.score),
("bounded-family Fisher weight", state.fisherweight),
("bounded-family observed Hessian", state.neghessian_eta),
(
"bounded-family observed Hessian derivative",
state.neghessian_eta_derivative,
),
("bounded-family log likelihood", state.log_likelihood),
] {
if !value.is_finite() {
return Err(bounded_row_error(row, quantity, eta, value));
}
}
if state.fisherweight < 0.0 {
return Err(bounded_row_error(
row,
"bounded-family Fisher weight",
eta,
state.fisherweight,
));
}
Ok(state)
}
#[inline]
fn weighted_positive_from_log(weight: f64, log_value: f64) -> f64 {
if weight == 0.0 {
return 0.0;
}
(weight.ln() + log_value).exp()
}
#[inline]
fn weighted_product3(a: f64, b: f64, c: f64) -> f64 {
crate::gamlss::scaled_signed_product3(a, b, c)
}
#[inline]
fn convex_combination(y: f64, left: f64, right: f64) -> f64 {
if y == 0.0 {
right
} else if y == 1.0 {
left
} else {
y.mul_add(left, (1.0 - y) * right)
}
}
#[derive(Clone, Copy)]
struct BernoulliNaturalJet {
mu: f64,
log_mu: [f64; 4],
log_one_minus_mu: [f64; 4],
log_fisher: f64,
}
#[inline]
fn probit_natural_jet(eta: f64) -> BernoulliNaturalJet {
let left = gam_math::probability::normal_logcdf_derivatives(eta);
let right_at_neg_eta = gam_math::probability::normal_logcdf_derivatives(-eta);
let log_pdf = if eta.abs() <= f64::MAX.sqrt() {
-0.5 * eta * eta - 0.5 * (2.0 * std::f64::consts::PI).ln()
} else {
f64::NEG_INFINITY
};
BernoulliNaturalJet {
mu: left[0].exp(),
log_mu: [left[0], left[1], left[2], left[3]],
log_one_minus_mu: [
right_at_neg_eta[0],
-right_at_neg_eta[1],
right_at_neg_eta[2],
-right_at_neg_eta[3],
],
log_fisher: 2.0 * log_pdf - left[0] - right_at_neg_eta[0],
}
}
#[inline]
fn cloglog_natural_jet(eta: f64) -> BernoulliNaturalJet {
let x = eta.exp();
if x == f64::INFINITY {
return BernoulliNaturalJet {
mu: 1.0,
log_mu: [0.0; 4],
log_one_minus_mu: [f64::NEG_INFINITY; 4],
log_fisher: f64::NEG_INFINITY,
};
}
if x == 0.0 {
return BernoulliNaturalJet {
mu: 0.0,
log_mu: [eta, 1.0, 0.0, 0.0],
log_one_minus_mu: [0.0; 4],
log_fisher: eta,
};
}
let mu = -(-x).exp_m1();
let log_mu = if x < 0.5 {
eta + (mu / x).ln()
} else {
mu.ln()
};
let h = if x < 1.0 {
x / x.exp_m1()
} else {
let exp_neg_x = (-x).exp();
x * exp_neg_x / (1.0 - exp_neg_x)
};
let a = 1.0 - x - h;
let d2_log_mu = h * a;
let d3_log_mu = h * (a * a - x - h * a);
BernoulliNaturalJet {
mu,
log_mu: [log_mu, h, d2_log_mu, d3_log_mu],
log_one_minus_mu: [-x, -x, -x, -x],
log_fisher: 2.0 * eta - x - log_mu,
}
}
#[inline]
fn loglog_natural_jet(eta: f64) -> BernoulliNaturalJet {
let mirrored = cloglog_natural_jet(-eta);
BernoulliNaturalJet {
mu: mirrored.log_one_minus_mu[0].exp(),
log_mu: [
mirrored.log_one_minus_mu[0],
-mirrored.log_one_minus_mu[1],
mirrored.log_one_minus_mu[2],
-mirrored.log_one_minus_mu[3],
],
log_one_minus_mu: [
mirrored.log_mu[0],
-mirrored.log_mu[1],
mirrored.log_mu[2],
-mirrored.log_mu[3],
],
log_fisher: mirrored.log_fisher,
}
}
#[inline]
fn cauchit_natural_jet(eta: f64) -> BernoulliNaturalJet {
let (mu, one_minus_mu) = if eta > 0.0 {
let q = (eta.recip()).atan() / std::f64::consts::PI;
(1.0 - q, q)
} else if eta < 0.0 {
let p = (-eta.recip()).atan() / std::f64::consts::PI;
(p, 1.0 - p)
} else {
(0.5, 0.5)
};
let abs_eta = eta.abs();
let log_one_plus_eta_sq = if abs_eta <= f64::MAX.sqrt() {
(eta * eta).ln_1p()
} else {
2.0 * abs_eta.ln() + eta.recip().powi(2).ln_1p()
};
let log_d1 = -std::f64::consts::PI.ln() - log_one_plus_eta_sq;
let ratio = if abs_eta <= 1.0 {
eta / (1.0 + eta * eta)
} else {
1.0 / (eta + eta.recip())
};
let d2_over_d1 = -2.0 * ratio;
let inv_one_plus_sq = if abs_eta <= 1.0 {
1.0 / (1.0 + eta * eta)
} else {
let inv = eta.recip();
inv * inv / (1.0 + inv * inv)
};
let d3_over_d1 = inv_one_plus_sq * (6.0 * (eta * ratio) - 2.0 * inv_one_plus_sq);
let d1_over_mu = (log_d1 - mu.ln()).exp();
let d1_over_q = (log_d1 - one_minus_mu.ln()).exp();
let left_d2_ratio = d2_over_d1 * d1_over_mu;
let right_d2_ratio = d2_over_d1 * d1_over_q;
BernoulliNaturalJet {
mu,
log_mu: [
mu.ln(),
d1_over_mu,
left_d2_ratio - d1_over_mu * d1_over_mu,
d3_over_d1 * d1_over_mu - 3.0 * d1_over_mu * left_d2_ratio + 2.0 * d1_over_mu.powi(3),
],
log_one_minus_mu: [
one_minus_mu.ln(),
-d1_over_q,
-right_d2_ratio - d1_over_q * d1_over_q,
-d3_over_d1 * d1_over_q - 3.0 * d1_over_q * right_d2_ratio - 2.0 * d1_over_q.powi(3),
],
log_fisher: 2.0 * log_d1 - mu.ln() - one_minus_mu.ln(),
}
}
#[inline]
fn generic_bernoulli_natural_jet(
row: usize,
eta: f64,
link: &InverseLink,
) -> Result<BernoulliNaturalJet, EstimationError> {
let jet = inverse_link_jet_for_inverse_link(link, eta)?;
if !(jet.mu.is_finite()
&& jet.mu > 0.0
&& jet.mu < 1.0
&& jet.d1.is_finite()
&& jet.d1 > 0.0
&& jet.d2.is_finite()
&& jet.d3.is_finite())
{
return Err(bounded_row_error(
row,
"bounded-family inverse-link jet",
eta,
jet.mu,
));
}
let mu = jet.mu;
let q = 1.0 - mu;
let r1 = jet.d1 / mu;
let r2 = jet.d2 / mu;
let r3 = jet.d3 / mu;
let s1 = jet.d1 / q;
let s2 = jet.d2 / q;
let s3 = jet.d3 / q;
Ok(BernoulliNaturalJet {
mu,
log_mu: [
mu.ln(),
r1,
r2 - r1 * r1,
r3 - 3.0 * r1 * r2 + 2.0 * r1.powi(3),
],
log_one_minus_mu: [
(-mu).ln_1p(),
-s1,
-s2 - s1 * s1,
-s3 - 3.0 * s1 * s2 - 2.0 * s1.powi(3),
],
log_fisher: 2.0 * jet.d1.ln() - mu.ln() - q.ln(),
})
}
fn resolved_bounded_binomial_link(
family: &LikelihoodSpec,
latent_cloglog_state: Option<&LatentCLogLogState>,
mixture_link_state: Option<&MixtureLinkState>,
sas_link_state: Option<&SasLinkState>,
) -> InverseLink {
match &family.link {
InverseLink::LatentCLogLog(_) => latent_cloglog_state
.copied()
.map(InverseLink::LatentCLogLog)
.unwrap_or_else(|| family.link.clone()),
InverseLink::Mixture(_) => mixture_link_state
.cloned()
.map(InverseLink::Mixture)
.unwrap_or_else(|| family.link.clone()),
InverseLink::Sas(_) => sas_link_state
.copied()
.map(InverseLink::Sas)
.unwrap_or_else(|| family.link.clone()),
InverseLink::BetaLogistic(_) => sas_link_state
.copied()
.map(InverseLink::BetaLogistic)
.unwrap_or_else(|| family.link.clone()),
InverseLink::Standard(_) => family.link.clone(),
}
}
fn binomial_natural_jet(
row: usize,
eta: f64,
link: &InverseLink,
) -> Result<BernoulliNaturalJet, EstimationError> {
match link {
InverseLink::Standard(StandardLink::Probit) => Ok(probit_natural_jet(eta)),
InverseLink::Standard(StandardLink::CLogLog) => Ok(cloglog_natural_jet(eta)),
InverseLink::Standard(StandardLink::LogLog) => Ok(loglog_natural_jet(eta)),
InverseLink::Standard(StandardLink::Cauchit) => Ok(cauchit_natural_jet(eta)),
_ => generic_bernoulli_natural_jet(row, eta, link),
}
}
fn exact_logit_observation_row(
row: usize,
y: f64,
weight: f64,
eta: f64,
) -> Result<ExactStandardObservationRow, EstimationError> {
let tail = (-eta.abs()).exp();
let (mu, one_minus_mu) = if eta >= 0.0 {
let q = tail / (1.0 + tail);
(1.0 - q, q)
} else {
let p = tail / (1.0 + tail);
(p, 1.0 - p)
};
if weight == 0.0 {
return Ok(ExactStandardObservationRow::zero_weight(mu));
}
let log_fisher =
-gam_linalg::utils::stable_softplus(eta) - gam_linalg::utils::stable_softplus(-eta);
let fisherweight = weighted_positive_from_log(weight, log_fisher);
if !(fisherweight.is_finite() && fisherweight > 0.0) {
return Err(bounded_row_error(
row,
"bounded logit Fisher weight",
eta,
fisherweight,
));
}
let residual = if eta >= 0.0 {
if y == 1.0 {
one_minus_mu
} else {
(y - 1.0) + one_minus_mu
}
} else {
y - mu
};
let log_likelihood_unit = if eta >= 0.0 {
-(1.0 - y) * eta - gam_linalg::utils::stable_softplus(-eta)
} else {
y * eta - gam_linalg::utils::stable_softplus(eta)
};
certify_bounded_row(
row,
eta,
ExactStandardObservationRow {
mu,
score: weight * residual,
fisherweight,
neghessian_eta: fisherweight,
neghessian_eta_derivative: fisherweight * (one_minus_mu - mu),
log_likelihood: weight * log_likelihood_unit,
},
)
}
fn exact_noncanonical_binomial_observation_row(
row: usize,
y: f64,
weight: f64,
eta: f64,
link: &InverseLink,
) -> Result<ExactStandardObservationRow, EstimationError> {
let jet = binomial_natural_jet(row, eta, link)?;
if weight == 0.0 {
return Ok(ExactStandardObservationRow::zero_weight(jet.mu));
}
let fisherweight = weighted_positive_from_log(weight, jet.log_fisher);
if !(fisherweight.is_finite() && fisherweight > 0.0) {
return Err(bounded_row_error(
row,
"bounded binomial Fisher weight",
eta,
fisherweight,
));
}
let log_likelihood = weight * convex_combination(y, jet.log_mu[0], jet.log_one_minus_mu[0]);
let score = weight * convex_combination(y, jet.log_mu[1], jet.log_one_minus_mu[1]);
let neghessian_eta = -weight * convex_combination(y, jet.log_mu[2], jet.log_one_minus_mu[2]);
let neghessian_eta_derivative =
-weight * convex_combination(y, jet.log_mu[3], jet.log_one_minus_mu[3]);
certify_bounded_row(
row,
eta,
ExactStandardObservationRow {
mu: jet.mu,
score,
fisherweight,
neghessian_eta,
neghessian_eta_derivative,
log_likelihood,
},
)
}
#[inline]
fn eta_exprel(rate: f64, eta: f64) -> f64 {
(rate * eta).exp_m1() / rate
}
fn validate_bounded_observation_inputs(
likelihood: &gam_spec::GlmLikelihoodSpec,
y: &Array1<f64>,
weights: &Array1<f64>,
eta: &Array1<f64>,
) -> Result<gam_spec::ResolvedLikelihoodScale, EstimationError> {
let family = &likelihood.spec;
if weights.len() != y.len() || eta.len() != y.len() {
crate::bail_invalid_estim!(
"bounded family observation size mismatch: y={}, weights={}, eta={}",
y.len(),
weights.len(),
eta.len()
);
}
if !LikelihoodSpec::is_legal_cell(&family.response, &family.link) {
crate::bail_invalid_estim!(
"bounded family received illegal likelihood cell response={} link={}",
family.response.name(),
family.link.link_function().name()
);
}
let resolved_scale = likelihood
.resolved_scale()
.map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
match &family.response {
ResponseFamily::Tweedie { p } if !(p.is_finite() && *p > 1.0 && *p < 2.0) => {
crate::bail_invalid_estim!(
"bounded Tweedie power must be finite and strictly inside (1, 2), got {p}"
);
}
ResponseFamily::NegativeBinomial { theta, .. } if !(theta.is_finite() && *theta > 0.0) => {
crate::bail_invalid_estim!(
"bounded negative-binomial theta must be finite and positive, got {theta}"
);
}
_ => {}
}
for (i, &wi) in weights.iter().enumerate() {
if !(wi.is_finite() && wi >= 0.0) {
return Err(EstimationError::InvalidInput(format!(
"bounded-family row {} has invalid prior weight {wi:?}; expected finite weight >= 0",
i + 1
)));
}
}
for i in 0..y.len() {
let wi = weights[i];
if wi == 0.0 {
continue;
}
if !eta[i].is_finite() {
return Err(bounded_row_error(i, "linear predictor", eta[i], eta[i]));
}
if !y[i].is_finite() {
return Err(bounded_row_error(
i,
"bounded-family response",
eta[i],
y[i],
));
}
let yi = y[i];
let valid = match &family.response {
ResponseFamily::Gaussian => yi.is_finite(),
ResponseFamily::Binomial => yi.is_finite() && (0.0..=1.0).contains(&yi),
ResponseFamily::Poisson | ResponseFamily::NegativeBinomial { .. } => {
yi.is_finite() && yi >= 0.0 && (yi - yi.round()).abs() <= 1e-9
}
ResponseFamily::Tweedie { .. } => yi.is_finite() && yi >= 0.0,
ResponseFamily::Gamma => yi.is_finite() && yi > 0.0,
ResponseFamily::Beta { .. } | ResponseFamily::RoystonParmar => false,
};
if !valid {
return Err(bounded_row_error(i, "bounded-family response", eta[i], yi));
}
}
Ok(resolved_scale)
}
fn exact_standard_observation_row(
likelihood: &gam_spec::GlmLikelihoodSpec,
resolved_scale: gam_spec::ResolvedLikelihoodScale,
binomial_link: &InverseLink,
row: usize,
y: f64,
weight: f64,
eta: f64,
) -> Result<ExactStandardObservationRow, EstimationError> {
if weight == 0.0 {
return Ok(ExactStandardObservationRow::zero_weight(0.0));
}
let family = &likelihood.spec;
match &family.response {
ResponseFamily::Gaussian => {
let scaled_weight = match resolved_scale {
gam_spec::ResolvedLikelihoodScale::ProfiledGaussian => weight,
gam_spec::ResolvedLikelihoodScale::FixedGaussian { phi } => {
crate::gamlss::scaled_positive_product_quotient(weight, 1.0, 1.0, phi.value())
}
_ => {
crate::bail_invalid_estim!(
"bounded Gaussian received a non-Gaussian resolved scale"
);
}
};
if !(scaled_weight.is_finite() && scaled_weight > 0.0) {
return Err(bounded_row_error(
row,
"bounded Gaussian dispersion-scaled weight",
eta,
scaled_weight,
));
}
let residual = y - eta;
let loss = if residual == 0.0 {
0.0
} else {
crate::gamlss::scaled_positive_product_quotient(
scaled_weight,
residual.abs(),
residual.abs(),
2.0,
)
};
certify_bounded_row(
row,
eta,
ExactStandardObservationRow {
mu: eta,
score: scaled_weight * residual,
fisherweight: scaled_weight,
neghessian_eta: scaled_weight,
neghessian_eta_derivative: 0.0,
log_likelihood: -loss,
},
)
}
ResponseFamily::Binomial
if matches!(binomial_link, InverseLink::Standard(StandardLink::Logit)) =>
{
exact_logit_observation_row(row, y, weight, eta)
}
ResponseFamily::Binomial => {
exact_noncanonical_binomial_observation_row(row, y, weight, eta, binomial_link)
}
ResponseFamily::Poisson => {
let mu = inverse_link_jet_for_inverse_link(&family.link, eta)?.mu;
let fisherweight = weight * mu;
let score = weight * (y - mu);
let raw_log_likelihood = y.mul_add(eta, -mu);
let log_likelihood = if raw_log_likelihood.is_finite() {
weight * raw_log_likelihood
} else {
weighted_product3(weight, y, eta) - weight * mu
};
if !(fisherweight.is_finite() && fisherweight > 0.0) {
return Err(bounded_row_error(
row,
"bounded Poisson Fisher weight",
eta,
fisherweight,
));
}
certify_bounded_row(
row,
eta,
ExactStandardObservationRow {
mu,
score,
fisherweight,
neghessian_eta: fisherweight,
neghessian_eta_derivative: fisherweight,
log_likelihood,
},
)
}
ResponseFamily::Gamma => {
let mu = inverse_link_jet_for_inverse_link(&family.link, eta)?.mu;
let shape = resolved_scale
.gamma_shape()
.map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
let weighted_shape = weight * shape;
if !(weighted_shape.is_finite() && weighted_shape > 0.0) {
return Err(bounded_row_error(
row,
"bounded Gamma shape-scaled weight",
eta,
weighted_shape,
));
}
let weighted_ratio =
crate::gamlss::scaled_positive_product_quotient(weight, y, shape, mu);
if !(weighted_ratio.is_finite() && weighted_ratio > 0.0) {
return Err(bounded_row_error(
row,
"bounded Gamma observed Hessian",
eta,
weighted_ratio,
));
}
certify_bounded_row(
row,
eta,
ExactStandardObservationRow {
mu,
score: weighted_ratio - weighted_shape,
fisherweight: weighted_shape,
neghessian_eta: weighted_ratio,
neghessian_eta_derivative: -weighted_ratio,
log_likelihood: -weighted_ratio - weighted_shape * eta,
},
)
}
ResponseFamily::Tweedie { p } => {
let p = *p;
let mu = inverse_link_jet_for_inverse_link(&family.link, eta)?.mu;
let phi = resolved_scale
.tweedie_phi()
.map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
let weight = crate::gamlss::scaled_positive_product_quotient(weight, 1.0, 1.0, phi);
if !(weight.is_finite() && weight > 0.0) {
return Err(bounded_row_error(
row,
"bounded Tweedie dispersion-scaled weight",
eta,
weight,
));
}
let a = ((1.0 - p) * eta).exp();
let b = ((2.0 - p) * eta).exp();
let score_unit = y.mul_add(a, -b);
let score = if score_unit.is_finite() {
weight * score_unit
} else {
weighted_product3(weight, y, a) - weight * b
};
let fisherweight = weight * b;
let observed_unit = (p - 1.0) * y * a + (2.0 - p) * b;
let neghessian_eta = if observed_unit.is_finite() {
weight * observed_unit
} else {
weighted_product3(weight * (p - 1.0), y, a) + weight * (2.0 - p) * b
};
let observed_derivative_unit = -(p - 1.0).powi(2) * y * a + (2.0 - p).powi(2) * b;
let neghessian_eta_derivative = if observed_derivative_unit.is_finite() {
weight * observed_derivative_unit
} else {
-weighted_product3(weight * (p - 1.0).powi(2), y, a)
+ weight * (2.0 - p).powi(2) * b
};
let q_left = eta_exprel(1.0 - p, eta);
let q_right = eta_exprel(2.0 - p, eta);
let q = y.mul_add(q_left, -q_right);
let log_likelihood = if q.is_finite() {
weight * q
} else {
weighted_product3(weight, y, q_left) - weight * q_right
};
if !(fisherweight.is_finite() && fisherweight > 0.0) {
return Err(bounded_row_error(
row,
"bounded Tweedie Fisher weight",
eta,
fisherweight,
));
}
certify_bounded_row(
row,
eta,
ExactStandardObservationRow {
mu,
score,
fisherweight,
neghessian_eta,
neghessian_eta_derivative,
log_likelihood,
},
)
}
ResponseFamily::NegativeBinomial { .. } => {
let theta = resolved_scale
.negative_binomial_theta()
.map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
let mu = inverse_link_jet_for_inverse_link(&family.link, eta)?.mu;
let log_theta = theta.ln();
let delta = eta - log_theta;
let log_q = -gam_linalg::utils::stable_softplus(-delta);
let log_r = -gam_linalg::utils::stable_softplus(delta);
let q = log_q.exp();
let r = log_r.exp();
let y_r = if y == 0.0 {
0.0
} else {
(y.ln() + log_r).exp()
};
let theta_q = (log_theta + log_q).exp();
let score = weight * (y_r - theta_q);
let fisherweight = weighted_positive_from_log(weight, log_theta + log_q);
let log_qr = log_q + log_r;
let observed_y = if y == 0.0 {
0.0
} else {
weighted_positive_from_log(weight, y.ln() + log_qr)
};
let observed_theta = weighted_positive_from_log(weight, log_theta + log_qr);
let neghessian_eta = observed_y + observed_theta;
let neghessian_eta_derivative = neghessian_eta * (r - q);
let softplus_tail = if delta >= 0.0 {
gam_linalg::utils::stable_softplus(-delta)
} else {
gam_linalg::utils::stable_softplus(delta)
};
let log_likelihood = if delta >= 0.0 {
-weighted_product3(weight, theta, delta)
- weighted_product3(weight, y, softplus_tail)
- weighted_product3(weight, theta, softplus_tail)
} else {
weighted_product3(weight, y, delta)
- weighted_product3(weight, y, softplus_tail)
- weighted_product3(weight, theta, softplus_tail)
};
if !(fisherweight.is_finite() && fisherweight > 0.0) {
return Err(bounded_row_error(
row,
"bounded negative-binomial Fisher weight",
eta,
fisherweight,
));
}
certify_bounded_row(
row,
eta,
ExactStandardObservationRow {
mu,
score,
fisherweight,
neghessian_eta,
neghessian_eta_derivative,
log_likelihood,
},
)
}
ResponseFamily::Beta { .. } => {
crate::bail_invalid_estim!("bounded linear terms are not supported for BetaLogit fits");
}
ResponseFamily::RoystonParmar => {
crate::bail_invalid_estim!(
"bounded linear terms are not supported for survival model fits"
);
}
}
}
fn evaluate_resolved_standard_family_observations(
likelihood: &gam_spec::GlmLikelihoodSpec,
latent_cloglog_state: Option<&LatentCLogLogState>,
mixture_link_state: Option<&MixtureLinkState>,
sas_link_state: Option<&SasLinkState>,
y: &Array1<f64>,
weights: &Array1<f64>,
eta: &Array1<f64>,
) -> Result<StandardFamilyObservationState, EstimationError> {
let n = y.len();
let resolved_scale = validate_bounded_observation_inputs(likelihood, y, weights, eta)?;
let family = &likelihood.spec;
let binomial_link = resolved_bounded_binomial_link(
&family,
latent_cloglog_state,
mixture_link_state,
sas_link_state,
);
let mut score = Array1::<f64>::zeros(n);
let mut fisherweight = Array1::<f64>::zeros(n);
let mut neghessian_eta = Array1::<f64>::zeros(n);
let mut neghessian_eta_derivative = Array1::<f64>::zeros(n);
let mut log_likelihood = 0.0;
let mut log_likelihood_compensation = 0.0;
for i in 0..n {
let row = exact_standard_observation_row(
likelihood,
resolved_scale,
&binomial_link,
i,
y[i],
weights[i],
eta[i],
)?;
score[i] = row.score;
fisherweight[i] = row.fisherweight;
neghessian_eta[i] = row.neghessian_eta;
neghessian_eta_derivative[i] = row.neghessian_eta_derivative;
let adjusted = row.log_likelihood - log_likelihood_compensation;
let updated = log_likelihood + adjusted;
log_likelihood_compensation = (updated - log_likelihood) - adjusted;
log_likelihood = updated;
if !log_likelihood.is_finite() {
return Err(bounded_row_error(
i,
"bounded-family cumulative log likelihood",
eta[i],
log_likelihood,
));
}
}
Ok(StandardFamilyObservationState {
eta: eta.clone(),
score,
fisherweight,
neghessian_eta,
neghessian_eta_derivative,
log_likelihood,
})
}
fn evaluate_standard_familyobservations(
family: LikelihoodSpec,
latent_cloglog_state: Option<&LatentCLogLogState>,
mixture_link_state: Option<&MixtureLinkState>,
sas_link_state: Option<&SasLinkState>,
y: &Array1<f64>,
weights: &Array1<f64>,
eta: &Array1<f64>,
) -> Result<StandardFamilyObservationState, EstimationError> {
let likelihood = gam_spec::GlmLikelihoodSpec::canonical(family);
evaluate_resolved_standard_family_observations(
&likelihood,
latent_cloglog_state,
mixture_link_state,
sas_link_state,
y,
weights,
eta,
)
}
fn exact_standard_working_response(
state: &StandardFamilyObservationState,
) -> Result<Array1<f64>, EstimationError> {
let mut out = state.eta.clone();
for i in 0..out.len() {
let weight = state.fisherweight[i];
let score = state.score[i];
if weight == 0.0 {
if score != 0.0 {
return Err(bounded_row_error(
i,
"zero-Fisher row with nonzero score",
state.eta[i],
score,
));
}
continue;
}
let increment = score / weight;
let value = out[i] + increment;
if !increment.is_finite() || !value.is_finite() {
return Err(bounded_row_error(
i,
"bounded-family working response",
state.eta[i],
value,
));
}
out[i] = value;
}
Ok(out)
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum SpatialAdaptiveHyperKind {
LogLambdaMagnitude,
LogLambdaGradient,
LogLambdaCurvature,
LogEpsilonMagnitude,
LogEpsilonGradient,
LogEpsilonCurvature,
}
impl SpatialAdaptiveHyperKind {
fn component_index(self) -> usize {
match self {
SpatialAdaptiveHyperKind::LogLambdaMagnitude
| SpatialAdaptiveHyperKind::LogEpsilonMagnitude => 0,
SpatialAdaptiveHyperKind::LogLambdaGradient
| SpatialAdaptiveHyperKind::LogEpsilonGradient => 1,
SpatialAdaptiveHyperKind::LogLambdaCurvature
| SpatialAdaptiveHyperKind::LogEpsilonCurvature => 2,
}
}
fn is_log_lambda(self) -> bool {
matches!(
self,
SpatialAdaptiveHyperKind::LogLambdaMagnitude
| SpatialAdaptiveHyperKind::LogLambdaGradient
| SpatialAdaptiveHyperKind::LogLambdaCurvature
)
}
fn is_log_epsilon(self) -> bool {
matches!(
self,
SpatialAdaptiveHyperKind::LogEpsilonMagnitude
| SpatialAdaptiveHyperKind::LogEpsilonGradient
| SpatialAdaptiveHyperKind::LogEpsilonCurvature
)
}
}
#[derive(Clone, Copy, Debug)]
struct SpatialAdaptiveHyperSpec {
cache_index: usize,
kind: SpatialAdaptiveHyperKind,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum SpatialAdaptiveExplicitSecondOrderKind {
StructuralZero,
LocalAlphaAlpha,
LocalAlphaEta,
SharedEtaEta,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum AdaptiveComponent {
Magnitude,
Gradient,
Curvature,
}
impl AdaptiveComponent {
fn from_index(index: usize) -> Result<Self, String> {
match index {
0 => Ok(AdaptiveComponent::Magnitude),
1 => Ok(AdaptiveComponent::Gradient),
2 => Ok(AdaptiveComponent::Curvature),
other => Err(SmoothError::invalid_index(format!(
"invalid adaptive component index {}",
other
))
.into()),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum HyperDerivativeKind {
Rho,
LogEpsilonFirst,
LogEpsilonSecond,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum HyperDriftKind {
Rho,
LogEpsilon,
}
impl SpatialAdaptiveHyperSpec {
fn component_index(self) -> usize {
self.kind.component_index()
}
fn explicit_second_order_kind(self, other: Self) -> SpatialAdaptiveExplicitSecondOrderKind {
if self.component_index() != other.component_index() {
return SpatialAdaptiveExplicitSecondOrderKind::StructuralZero;
}
match (
self.kind.is_log_lambda(),
other.kind.is_log_lambda(),
self.kind.is_log_epsilon(),
other.kind.is_log_epsilon(),
) {
(true, true, false, false) if self.cache_index == other.cache_index => {
SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaAlpha
}
(true, false, false, true) | (false, true, true, false) => {
SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaEta
}
(false, false, true, true) => SpatialAdaptiveExplicitSecondOrderKind::SharedEtaEta,
_ => SpatialAdaptiveExplicitSecondOrderKind::StructuralZero,
}
}
}
#[derive(Clone, Debug)]
struct SpatialAdaptiveTermHyperParams {
lambda: [f64; 3],
epsilon: [f64; 3],
}
#[derive(Clone, Debug)]
struct ValidatedFixedQuadraticHessian {
dense: Arc<Array2<f64>>,
}
impl ValidatedFixedQuadraticHessian {
fn try_from_dense(dense: Array2<f64>, coefficient_dim: usize) -> Result<Self, String> {
gam_linalg::utils::validate_finite_symmetric_matrix(
&dense,
"spatial adaptive fixed quadratic Hessian",
)
.map_err(|error| error.to_string())?;
PenaltyMatrix::Dense(dense.clone())
.validate(coefficient_dim)
.map_err(|error| {
format!(
"spatial adaptive fixed quadratic Hessian failed quadratic-form validation: {error}"
)
})?;
Ok(Self {
dense: Arc::new(dense),
})
}
fn zero(coefficient_dim: usize) -> Result<Self, String> {
Self::try_from_dense(
Array2::<f64>::zeros((coefficient_dim, coefficient_dim)),
coefficient_dim,
)
}
fn as_dense(&self) -> &Array2<f64> {
self.dense.as_ref()
}
fn quadratic_terms(&self, beta: &Array1<f64>) -> Result<(f64, Array1<f64>), String> {
if beta.len() != self.dense.ncols() {
return Err(format!(
"spatial adaptive fixed quadratic beta length {} does not match validated Hessian dimension {}",
beta.len(),
self.dense.ncols()
));
}
let gradient = self.dense.dot(beta);
let value = 0.5 * beta.dot(&gradient);
Ok((value, gradient))
}
}
#[derive(Clone)]
struct SpatialAdaptiveExactEvaluation {
obs: StandardFamilyObservationState,
adaptive_states: Vec<SpatialPenaltyExactState>,
adaptive_penalty_value: f64,
adaptive_penaltygradient: Array1<f64>,
adaptive_penaltyhessian: Array2<f64>,
fixed_quadraticvalue: f64,
fixed_quadraticgradient: Array1<f64>,
fixed_quadratic_hessian: ValidatedFixedQuadraticHessian,
}
#[derive(Clone)]
struct CachedSpatialAdaptiveExactEvaluation {
beta: Array1<f64>,
eval: Arc<SpatialAdaptiveExactEvaluation>,
}
impl SpatialAdaptiveExactEvaluation {
fn total_penalty_value(&self) -> f64 {
self.adaptive_penalty_value + self.fixed_quadraticvalue
}
fn total_penaltygradient(&self) -> Array1<f64> {
&self.adaptive_penaltygradient + &self.fixed_quadraticgradient
}
fn total_penaltyhessian(&self) -> Array2<f64> {
&self.adaptive_penaltyhessian + self.fixed_quadratic_hessian.as_dense()
}
fn totalobjectivehessian(&self, design: &Array2<f64>) -> Result<Array2<f64>, String> {
let mut out = xt_diag_x_dense(design.view(), self.obs.neghessian_eta.view())?;
out += &self.total_penaltyhessian();
Ok(out)
}
}
#[derive(Clone)]
struct SpatialAdaptiveExactFamily {
family: LikelihoodSpec,
latent_cloglog_state: Option<LatentCLogLogState>,
mixture_link_state: Option<MixtureLinkState>,
sas_link_state: Option<SasLinkState>,
y: Arc<Array1<f64>>,
weights: Arc<Array1<f64>>,
design: Arc<Array2<f64>>,
offset: Arc<Array1<f64>>,
linear_constraints: Option<LinearInequalityConstraints>,
runtime_caches: Arc<Vec<SpatialOperatorRuntimeCache>>,
adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
fixed_quadratic_hessian: ValidatedFixedQuadraticHessian,
hyperspecs: Arc<Vec<SpatialAdaptiveHyperSpec>>,
exact_eval_cache: Arc<Mutex<Option<CachedSpatialAdaptiveExactEvaluation>>>,
}
impl SpatialAdaptiveExactFamily {
fn with_adaptive_params(
&self,
adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
fixed_quadratic_hessian: ValidatedFixedQuadraticHessian,
) -> Self {
Self {
family: self.family.clone(),
latent_cloglog_state: self.latent_cloglog_state,
mixture_link_state: self.mixture_link_state.clone(),
sas_link_state: self.sas_link_state,
y: self.y.clone(),
weights: self.weights.clone(),
design: self.design.clone(),
offset: self.offset.clone(),
linear_constraints: self.linear_constraints.clone(),
runtime_caches: self.runtime_caches.clone(),
adaptive_params,
fixed_quadratic_hessian,
hyperspecs: self.hyperspecs.clone(),
exact_eval_cache: Arc::new(Mutex::new(None)),
}
}
fn total_eta(&self, beta: &Array1<f64>) -> Array1<f64> {
gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), beta) + self.offset.as_ref()
}
fn fixed_quadratic_terms(
&self,
beta: &Array1<f64>,
) -> Result<(f64, Array1<f64>), String> {
self.fixed_quadratic_hessian.quadratic_terms(beta)
}
fn adaptive_penalty_value_only(&self, beta: &Array1<f64>) -> Result<f64, String> {
let mut penalty_value = 0.0;
for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
format!(
"missing adaptive parameter block for cache {}",
cache.termname
)
})?;
let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
let state =
SpatialPenaltyExactState::from_beta_local(beta_local, cache, params.epsilon)
.map_err(|e| e.to_string())?;
penalty_value += params.lambda[0] * state.magnitude.penalty_value();
penalty_value += params.lambda[1] * state.gradient.penalty_value();
penalty_value += params.lambda[2] * state.curvature.penalty_value();
}
Ok(penalty_value)
}
fn zero_hyper_parts(&self) -> (Array1<f64>, Array2<f64>) {
let total_dim = self.design.ncols();
(
Array1::<f64>::zeros(total_dim),
Array2::<f64>::zeros((total_dim, total_dim)),
)
}
fn embed_local_hyper_parts(
&self,
coeff_range: &Range<usize>,
local_grad: &Array1<f64>,
local_hess: &Array2<f64>,
) -> (Array1<f64>, Array2<f64>) {
let (mut beta_mixed, mut betahessian) = self.zero_hyper_parts();
beta_mixed
.slice_mut(s![coeff_range.clone()])
.assign(local_grad);
betahessian
.slice_mut(s![coeff_range.clone(), coeff_range.clone()])
.assign(local_hess);
(beta_mixed, betahessian)
}
fn embed_local_hyper_hessian(
&self,
coeff_range: &Range<usize>,
local_hess: &Array2<f64>,
) -> Array2<f64> {
let total_dim = self.design.ncols();
let mut out = Array2::<f64>::zeros((total_dim, total_dim));
out.slice_mut(s![coeff_range.clone(), coeff_range.clone()])
.assign(local_hess);
out
}
fn adaptive_block_eval(
&self,
eval: &SpatialAdaptiveExactEvaluation,
cache_idx: usize,
component: AdaptiveComponent,
derivative: HyperDerivativeKind,
) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
let cache = self
.runtime_caches
.get(cache_idx)
.ok_or_else(|| format!("adaptive cache index {} out of bounds", cache_idx))?;
let params = self
.adaptive_params
.get(cache_idx)
.ok_or_else(|| format!("adaptive hyperparameter block {} out of bounds", cache_idx))?;
let state = eval
.adaptive_states
.get(cache_idx)
.ok_or_else(|| format!("adaptive exact state index {} out of bounds", cache_idx))?;
let (objective_local, beta_mixed_local, betahessian_local) = match component {
AdaptiveComponent::Magnitude => {
let lambda = params.lambda[0];
let mag = &state.magnitude;
let (objective, gradient_coeff, hessian_diag) = match derivative {
HyperDerivativeKind::Rho => (
mag.penalty_value(),
mag.betagradient_coeff(),
mag.betahessian_diag(),
),
HyperDerivativeKind::LogEpsilonFirst => (
mag.log_epsilon_gradient_terms().sum(),
mag.log_epsilon_betagradient_coeff(),
mag.log_epsilon_betahessian_diag(),
),
HyperDerivativeKind::LogEpsilonSecond => (
mag.log_epsilon_hessian_terms().sum(),
mag.log_epsilon_beta_mixed_second_coeff(),
mag.log_epsilon_betahessian_second_diag(),
),
};
(
lambda * objective,
lambda * scalar_operatorgradient(&cache.d0, &gradient_coeff),
lambda * scalar_operatorhessian(&cache.d0, &hessian_diag),
)
}
AdaptiveComponent::Gradient => {
let lambda = params.lambda[1];
let grad = &state.gradient;
let (objective, gradient_blocks, hessian_blocks) = match derivative {
HyperDerivativeKind::Rho => (
grad.penalty_value(),
grad.betagradient_blocks(),
grad.betahessian_blocks(),
),
HyperDerivativeKind::LogEpsilonFirst => (
grad.log_epsilon_gradient_terms().sum(),
grad.log_epsilon_betagradient_blocks(),
grad.log_epsilon_betahessian_blocks(),
),
HyperDerivativeKind::LogEpsilonSecond => (
grad.log_epsilon_hessian_terms().sum(),
grad.log_epsilon_beta_mixed_second_blocks(),
grad.log_epsilon_betahessian_second_blocks(),
),
};
(
lambda * objective,
lambda
* grouped_operatorgradient(&cache.d1, cache.dimension, &gradient_blocks)
.map_err(|e| e.to_string())?,
lambda
* grouped_operatorhessian(&cache.d1, cache.dimension, &hessian_blocks)
.map_err(|e| e.to_string())?,
)
}
AdaptiveComponent::Curvature => {
let lambda = params.lambda[2];
let group = cache.dimension * cache.dimension;
let curv = &state.curvature;
let (objective, gradient_blocks, hessian_blocks) = match derivative {
HyperDerivativeKind::Rho => (
curv.penalty_value(),
curv.betagradient_blocks(),
curv.betahessian_blocks(),
),
HyperDerivativeKind::LogEpsilonFirst => (
curv.log_epsilon_gradient_terms().sum(),
curv.log_epsilon_betagradient_blocks(),
curv.log_epsilon_betahessian_blocks(),
),
HyperDerivativeKind::LogEpsilonSecond => (
curv.log_epsilon_hessian_terms().sum(),
curv.log_epsilon_beta_mixed_second_blocks(),
curv.log_epsilon_betahessian_second_blocks(),
),
};
(
lambda * objective,
lambda
* grouped_operatorgradient(&cache.d2, group, &gradient_blocks)
.map_err(|e| e.to_string())?,
lambda
* grouped_operatorhessian(&cache.d2, group, &hessian_blocks)
.map_err(|e| e.to_string())?,
)
}
};
let (beta_mixed, betahessian) = self.embed_local_hyper_parts(
&cache.coeff_global_range,
&beta_mixed_local,
&betahessian_local,
);
Ok((objective_local, beta_mixed, betahessian))
}
fn adaptive_shared_log_epsilon_parts(
&self,
eval: &SpatialAdaptiveExactEvaluation,
component: usize,
) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
self.adaptive_shared_block_eval(eval, component, HyperDerivativeKind::LogEpsilonFirst)
}
fn adaptive_shared_log_epsilon_second_parts(
&self,
eval: &SpatialAdaptiveExactEvaluation,
component: usize,
) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
self.adaptive_shared_block_eval(eval, component, HyperDerivativeKind::LogEpsilonSecond)
}
fn adaptive_shared_block_eval(
&self,
eval: &SpatialAdaptiveExactEvaluation,
component: usize,
derivative: HyperDerivativeKind,
) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
let component = AdaptiveComponent::from_index(component)?;
let (mut score, mut hessian) = self.zero_hyper_parts();
let mut objective = 0.0;
for cache_idx in 0..self.runtime_caches.len() {
let (local_objective, local_score, local_hessian) =
self.adaptive_block_eval(eval, cache_idx, component, derivative)?;
objective += local_objective;
score += &local_score;
hessian += &local_hessian;
}
Ok((objective, score, hessian))
}
fn adaptive_shared_log_epsilon_drift(
&self,
eval: &SpatialAdaptiveExactEvaluation,
component: usize,
direction: &Array1<f64>,
) -> Result<Array2<f64>, String> {
let component = AdaptiveComponent::from_index(component)?;
let total_dim = self.design.ncols();
let mut total = Array2::<f64>::zeros((total_dim, total_dim));
for cache_idx in 0..self.runtime_caches.len() {
total += &self.adaptive_block_drift_eval(
eval,
cache_idx,
component,
HyperDriftKind::LogEpsilon,
direction,
)?;
}
Ok(total)
}
fn adaptive_explicit_second_order_parts(
&self,
eval: &SpatialAdaptiveExactEvaluation,
left: SpatialAdaptiveHyperSpec,
right: SpatialAdaptiveHyperSpec,
) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
match left.explicit_second_order_kind(right) {
SpatialAdaptiveExplicitSecondOrderKind::StructuralZero => {
let (score, hessian) = self.zero_hyper_parts();
Ok((0.0, score, hessian))
}
SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaAlpha => self.adaptive_block_eval(
eval,
left.cache_index,
AdaptiveComponent::from_index(left.component_index())?,
HyperDerivativeKind::Rho,
),
SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaEta => {
let local_alpha = if left.kind.is_log_lambda() {
left
} else {
right
};
self.adaptive_block_eval(
eval,
local_alpha.cache_index,
AdaptiveComponent::from_index(local_alpha.component_index())?,
HyperDerivativeKind::LogEpsilonFirst,
)
}
SpatialAdaptiveExplicitSecondOrderKind::SharedEtaEta => {
self.adaptive_shared_log_epsilon_second_parts(eval, left.component_index())
}
}
}
fn adaptive_block_drift_eval(
&self,
eval: &SpatialAdaptiveExactEvaluation,
cache_idx: usize,
component: AdaptiveComponent,
drift: HyperDriftKind,
direction: &Array1<f64>,
) -> Result<Array2<f64>, String> {
let cache = self
.runtime_caches
.get(cache_idx)
.ok_or_else(|| format!("adaptive cache index {} out of bounds", cache_idx))?;
let params = self
.adaptive_params
.get(cache_idx)
.ok_or_else(|| format!("adaptive hyperparameter block {} out of bounds", cache_idx))?;
let state = eval
.adaptive_states
.get(cache_idx)
.ok_or_else(|| format!("adaptive exact state index {} out of bounds", cache_idx))?;
let direction_local = direction.slice(s![cache.coeff_global_range.clone()]);
let local_hessian = match component {
AdaptiveComponent::Magnitude => {
let d0_u = cache.d0.dot(&direction_local);
let mag = &state.magnitude;
let diag = match drift {
HyperDriftKind::Rho => mag.directionalhessian_diag(&d0_u),
HyperDriftKind::LogEpsilon => {
mag.log_epsilon_betahessian_directional_diag(&d0_u)
}
};
params.lambda[0] * scalar_operatorhessian(&cache.d0, &diag)
}
AdaptiveComponent::Gradient => {
let d1_u = cache.d1.dot(&direction_local);
let direction_blocks = collocationgradient_blocks(&d1_u, cache.dimension)
.map_err(|e| e.to_string())?;
let grad = &state.gradient;
let blocks = match drift {
HyperDriftKind::Rho => grad.directionalhessian_blocks(&direction_blocks),
HyperDriftKind::LogEpsilon => {
grad.log_epsilon_betahessian_directional_blocks(&direction_blocks)
}
};
params.lambda[1]
* grouped_operatorhessian(&cache.d1, cache.dimension, &blocks)
.map_err(|e| e.to_string())?
}
AdaptiveComponent::Curvature => {
let group = cache.dimension * cache.dimension;
let d2_u = cache.d2.dot(&direction_local);
let direction_blocks =
collocationhessian_blocks(&d2_u, cache.dimension).map_err(|e| e.to_string())?;
let curv = &state.curvature;
let blocks = match drift {
HyperDriftKind::Rho => curv.directionalhessian_blocks(&direction_blocks),
HyperDriftKind::LogEpsilon => {
curv.log_epsilon_betahessian_directional_blocks(&direction_blocks)
}
};
params.lambda[2]
* grouped_operatorhessian(&cache.d2, group, &blocks)
.map_err(|e| e.to_string())?
}
};
Ok(self.embed_local_hyper_hessian(&cache.coeff_global_range, &local_hessian))
}
fn adaptive_hyper_parts(
&self,
eval: &SpatialAdaptiveExactEvaluation,
hyper: SpatialAdaptiveHyperSpec,
) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
match hyper.kind {
SpatialAdaptiveHyperKind::LogLambdaMagnitude
| SpatialAdaptiveHyperKind::LogLambdaGradient
| SpatialAdaptiveHyperKind::LogLambdaCurvature => self.adaptive_block_eval(
eval,
hyper.cache_index,
AdaptiveComponent::from_index(hyper.component_index())?,
HyperDerivativeKind::Rho,
),
SpatialAdaptiveHyperKind::LogEpsilonMagnitude
| SpatialAdaptiveHyperKind::LogEpsilonGradient
| SpatialAdaptiveHyperKind::LogEpsilonCurvature => {
self.adaptive_shared_log_epsilon_parts(eval, hyper.component_index())
}
}
}
fn exact_evaluation_uncached(
&self,
beta: &Array1<f64>,
) -> Result<SpatialAdaptiveExactEvaluation, String> {
let eta = self.total_eta(beta);
let obs = evaluate_standard_familyobservations(
self.family.clone(),
self.latent_cloglog_state.as_ref(),
self.mixture_link_state.as_ref(),
self.sas_link_state.as_ref(),
&self.y,
&self.weights,
&eta,
)
.map_err(|e| e.to_string())?;
let p = beta.len();
let mut penalty_value = 0.0;
let mut penaltygradient = Array1::<f64>::zeros(p);
let mut penaltyhessian = Array2::<f64>::zeros((p, p));
let mut adaptive_states = Vec::with_capacity(self.runtime_caches.len());
for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
format!(
"missing adaptive parameter block for cache {}",
cache.termname
)
})?;
let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
let state =
SpatialPenaltyExactState::from_beta_local(beta_local, cache, params.epsilon)
.map_err(|e| e.to_string())?;
let g0 = scalar_operatorgradient(&cache.d0, &state.magnitude.betagradient_coeff());
let gg = grouped_operatorgradient(
&cache.d1,
cache.dimension,
&state.gradient.betagradient_blocks(),
)
.map_err(|e| e.to_string())?;
let gc = grouped_operatorgradient(
&cache.d2,
cache.dimension * cache.dimension,
&state.curvature.betagradient_blocks(),
)
.map_err(|e| e.to_string())?;
let h0 = scalar_operatorhessian(&cache.d0, &state.magnitude.betahessian_diag());
let hg = grouped_operatorhessian(
&cache.d1,
cache.dimension,
&state.gradient.betahessian_blocks(),
)
.map_err(|e| e.to_string())?;
let hc = grouped_operatorhessian(
&cache.d2,
cache.dimension * cache.dimension,
&state.curvature.betahessian_blocks(),
)
.map_err(|e| e.to_string())?;
let lambda0 = params.lambda[0];
let lambdag = params.lambda[1];
let lambdac = params.lambda[2];
penalty_value += lambda0 * state.magnitude.penalty_value();
penalty_value += lambdag * state.gradient.penalty_value();
penalty_value += lambdac * state.curvature.penalty_value();
let range = cache.coeff_global_range.clone();
{
let mut grad_local = penaltygradient.slice_mut(s![range.clone()]);
grad_local += &(g0.mapv(|v| lambda0 * v));
grad_local += &(gg.mapv(|v| lambdag * v));
grad_local += &(gc.mapv(|v| lambdac * v));
}
{
let mut h_local = penaltyhessian.slice_mut(s![range.clone(), range]);
h_local += &h0.mapv(|v| lambda0 * v);
h_local += &hg.mapv(|v| lambdag * v);
h_local += &hc.mapv(|v| lambdac * v);
}
adaptive_states.push(state);
}
let (fixed_quadraticvalue, fixed_quadraticgradient) =
self.fixed_quadratic_terms(beta)?;
Ok(SpatialAdaptiveExactEvaluation {
obs,
adaptive_states,
adaptive_penalty_value: penalty_value,
adaptive_penaltygradient: penaltygradient,
adaptive_penaltyhessian: penaltyhessian,
fixed_quadraticvalue,
fixed_quadraticgradient,
fixed_quadratic_hessian: self.fixed_quadratic_hessian.clone(),
})
}
fn exact_evaluation(
&self,
beta: &Array1<f64>,
) -> Result<Arc<SpatialAdaptiveExactEvaluation>, String> {
{
let cache = self
.exact_eval_cache
.lock()
.map_err(|_| "spatial adaptive exact-evaluation cache lock poisoned".to_string())?;
if let Some(cached) = cache.as_ref()
&& cached.beta.len() == beta.len()
&& cached
.beta
.iter()
.zip(beta.iter())
.all(|(&left, &right)| left == right)
{
return Ok(Arc::clone(&cached.eval));
}
}
let eval = Arc::new(self.exact_evaluation_uncached(beta)?);
let mut cache = self
.exact_eval_cache
.lock()
.map_err(|_| "spatial adaptive exact-evaluation cache lock poisoned".to_string())?;
*cache = Some(CachedSpatialAdaptiveExactEvaluation {
beta: beta.clone(),
eval: Arc::clone(&eval),
});
Ok(eval)
}
fn exacthessian_directional_derivative_from_evaluation(
&self,
beta: &Array1<f64>,
eval: &SpatialAdaptiveExactEvaluation,
direction: &Array1<f64>,
) -> Result<Array2<f64>, String> {
assert_eq!(
beta.len(),
direction.len(),
"beta/direction length mismatch",
);
let d_eta = gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), direction);
let mut total = xt_diag_x_dense(
self.design.view(),
(&eval.obs.neghessian_eta_derivative * &d_eta).view(),
)?;
for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
format!(
"missing adaptive parameter block for cache {}",
cache.termname
)
})?;
let state = eval
.adaptive_states
.get(cache_idx)
.ok_or_else(|| format!("missing adaptive state for cache {}", cache.termname))?;
let direction_local = direction.slice(s![cache.coeff_global_range.clone()]);
let d0_u = cache.d0.dot(&direction_local);
let d1_u = cache.d1.dot(&direction_local);
let d2_u = cache.d2.dot(&direction_local);
let h0 =
scalar_operatorhessian(&cache.d0, &state.magnitude.directionalhessian_diag(&d0_u))
.mapv(|v| params.lambda[0] * v);
let hg = grouped_operatorhessian(
&cache.d1,
cache.dimension,
&state.gradient.directionalhessian_blocks(
&collocationgradient_blocks(&d1_u, cache.dimension)
.map_err(|e| e.to_string())?,
),
)
.map_err(|e| e.to_string())?
.mapv(|v| params.lambda[1] * v);
let hc = grouped_operatorhessian(
&cache.d2,
cache.dimension * cache.dimension,
&state.curvature.directionalhessian_blocks(
&collocationhessian_blocks(&d2_u, cache.dimension)
.map_err(|e| e.to_string())?,
),
)
.map_err(|e| e.to_string())?
.mapv(|v| params.lambda[2] * v);
let range = cache.coeff_global_range.clone();
let mut local = total.slice_mut(s![range.clone(), range]);
local += &h0;
local += &hg;
local += &hc;
}
Ok(total)
}
fn exacthessian_second_directional_derivative_from_evaluation(
&self,
eval: &SpatialAdaptiveExactEvaluation,
direction_u: &Array1<f64>,
direction_v: &Array1<f64>,
) -> Result<Option<Array2<f64>>, String> {
let p = self.design.ncols();
if eval.obs.neghessian_eta_derivative.iter().any(|&w| w != 0.0) {
return Ok(None);
}
let mut total = Array2::<f64>::zeros((p, p));
for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
format!(
"missing adaptive parameter block for cache {}",
cache.termname
)
})?;
let state = eval
.adaptive_states
.get(cache_idx)
.ok_or_else(|| format!("missing adaptive state for cache {}", cache.termname))?;
let u_local = direction_u.slice(s![cache.coeff_global_range.clone()]);
let v_local = direction_v.slice(s![cache.coeff_global_range.clone()]);
let q0_u = cache.d0.dot(&u_local);
let q0_v = cache.d0.dot(&v_local);
let h0 = scalar_operatorhessian(
&cache.d0,
&state.magnitude.second_directionalhessian_diag(&q0_u, &q0_v),
)
.mapv(|x| params.lambda[0] * x);
let a1 = collocationgradient_blocks(&cache.d1.dot(&u_local), cache.dimension)
.map_err(|e| e.to_string())?;
let b1 = collocationgradient_blocks(&cache.d1.dot(&v_local), cache.dimension)
.map_err(|e| e.to_string())?;
let hg = grouped_operatorhessian(
&cache.d1,
cache.dimension,
&state.gradient.second_directionalhessian_blocks(&a1, &b1),
)
.map_err(|e| e.to_string())?
.mapv(|x| params.lambda[1] * x);
let a2 = collocationhessian_blocks(&cache.d2.dot(&u_local), cache.dimension)
.map_err(|e| e.to_string())?;
let b2 = collocationhessian_blocks(&cache.d2.dot(&v_local), cache.dimension)
.map_err(|e| e.to_string())?;
let hc = grouped_operatorhessian(
&cache.d2,
cache.dimension * cache.dimension,
&state.curvature.second_directionalhessian_blocks(&a2, &b2),
)
.map_err(|e| e.to_string())?
.mapv(|x| params.lambda[2] * x);
let range = cache.coeff_global_range.clone();
let mut local = total.slice_mut(s![range.clone(), range]);
local += &h0;
local += &hg;
local += &hc;
}
Ok(Some(total))
}
}
impl CustomFamily for SpatialAdaptiveExactFamily {
fn joint_jeffreys_term_required(&self) -> bool {
true
}
fn joint_jeffreys_information_with_specs(
&self,
block_states: &[ParameterBlockState],
specs: &[ParameterBlockSpec],
) -> Result<Option<Array2<f64>>, String> {
let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
if spec.design.ncols() != beta.len() {
return Err(SmoothError::dimension_mismatch(format!(
"spatial adaptive Jeffreys information: spec design has {} columns, beta has {}",
spec.design.ncols(),
beta.len()
))
.into());
}
let eval = self.exact_evaluation(beta)?;
Ok(Some(xt_diag_x_dense(
self.design.view(),
eval.obs.neghessian_eta.view(),
)?))
}
fn joint_jeffreys_information_directional_derivative_with_specs(
&self,
block_states: &[ParameterBlockState],
specs: &[ParameterBlockSpec],
d_beta_flat: &Array1<f64>,
) -> Result<Option<Array2<f64>>, String> {
let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
if spec.design.ncols() != d_beta_flat.len() {
return Err(SmoothError::dimension_mismatch(format!(
"spatial adaptive Jeffreys directional derivative: spec design has {} columns, direction has {}",
spec.design.ncols(),
d_beta_flat.len()
))
.into());
}
let eval = self.exact_evaluation(beta)?;
let d_eta = gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), d_beta_flat);
Ok(Some(xt_diag_x_dense(
self.design.view(),
(&eval.obs.neghessian_eta_derivative * &d_eta).view(),
)?))
}
fn joint_jeffreys_information_second_directional_derivative_with_specs(
&self,
block_states: &[ParameterBlockState],
specs: &[ParameterBlockSpec],
d_beta_u_flat: &Array1<f64>,
d_betav_flat: &Array1<f64>,
) -> Result<Option<Array2<f64>>, String> {
let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
if spec.design.ncols() != beta.len()
|| d_beta_u_flat.len() != beta.len()
|| d_betav_flat.len() != beta.len()
{
return Err(SmoothError::dimension_mismatch(format!(
"spatial adaptive Jeffreys second-direction length mismatch: spec cols={}, dirs=({}, {}), expected {}",
spec.design.ncols(),
d_beta_u_flat.len(),
d_betav_flat.len(),
beta.len()
))
.into());
}
let eval = self.exact_evaluation(beta)?;
if eval.obs.neghessian_eta_derivative.iter().any(|&w| w != 0.0) {
return Ok(None);
}
Ok(Some(Array2::<f64>::zeros((beta.len(), beta.len()))))
}
fn joint_jeffreys_information_matches_observed_hessian(&self) -> bool {
false
}
fn joint_jeffreys_information_depends_on_psi(&self) -> bool {
false
}
fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
let eval = self.exact_evaluation(beta)?;
let mut gradient = fast_atv(&self.design, &eval.obs.score);
gradient -= &eval.total_penaltygradient();
let mut hessian = xt_diag_x_dense(self.design.view(), eval.obs.neghessian_eta.view())?;
hessian += &eval.total_penaltyhessian();
Ok(FamilyEvaluation {
log_likelihood: eval.obs.log_likelihood - eval.total_penalty_value(),
blockworking_sets: vec![BlockWorkingSet::ExactNewton {
gradient,
hessian: SymmetricMatrix::Dense(hessian),
}],
})
}
fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
let state = expect_single_block_state(block_states, "spatial adaptive exact family")?;
let beta = &state.beta;
let obs = evaluate_standard_familyobservations(
self.family.clone(),
self.latent_cloglog_state.as_ref(),
self.mixture_link_state.as_ref(),
self.sas_link_state.as_ref(),
&self.y,
&self.weights,
&state.eta,
)
.map_err(|e| e.to_string())?;
let adaptive_penalty = self.adaptive_penalty_value_only(beta)?;
let (fixed_quadratic, _) = self.fixed_quadratic_terms(beta)?;
Ok(obs.log_likelihood - adaptive_penalty - fixed_quadratic)
}
fn exact_newton_outerobjective(&self) -> ExactNewtonOuterObjective {
ExactNewtonOuterObjective::StrictPseudoLaplace
}
fn exact_newton_joint_hessian(
&self,
block_states: &[ParameterBlockState],
) -> Result<Option<Array2<f64>>, String> {
let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
let eval = self.exact_evaluation(beta)?;
Ok(Some(eval.totalobjectivehessian(&self.design)?))
}
fn exact_newton_hessian_directional_derivative(
&self,
block_states: &[ParameterBlockState],
block_idx: usize,
d_beta: &Array1<f64>,
) -> Result<Option<Array2<f64>>, String> {
expect_block_idx_zero(block_idx, "spatial adaptive exact family", "")?;
self.exact_newton_joint_hessian_directional_derivative(block_states, d_beta)
}
fn exact_newton_joint_hessian_directional_derivative(
&self,
block_states: &[ParameterBlockState],
d_beta_flat: &Array1<f64>,
) -> Result<Option<Array2<f64>>, String> {
let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
if d_beta_flat.len() != beta.len() {
return Err(SmoothError::dimension_mismatch(format!(
"spatial adaptive exact family direction length mismatch: got {}, expected {}",
d_beta_flat.len(),
beta.len()
))
.into());
}
let eval = self.exact_evaluation(beta)?;
Ok(Some(
self.exacthessian_directional_derivative_from_evaluation(beta, &eval, d_beta_flat)?,
))
}
fn exact_newton_joint_hessiansecond_directional_derivative(
&self,
block_states: &[ParameterBlockState],
d_beta_u_flat: &Array1<f64>,
d_betav_flat: &Array1<f64>,
) -> Result<Option<Array2<f64>>, String> {
let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
if d_beta_u_flat.len() != beta.len() || d_betav_flat.len() != beta.len() {
return Err(SmoothError::dimension_mismatch(format!(
"spatial adaptive exact family second-direction length mismatch: got ({}, {}), expected {}",
d_beta_u_flat.len(),
d_betav_flat.len(),
beta.len()
))
.into());
}
let eval = self.exact_evaluation(beta)?;
self.exacthessian_second_directional_derivative_from_evaluation(
&eval,
d_beta_u_flat,
d_betav_flat,
)
}
fn block_linear_constraints(
&self,
block_states: &[ParameterBlockState],
block_idx: usize,
block_spec: &ParameterBlockSpec,
) -> Result<Option<ConstraintSet>, String> {
assert!(!block_states.is_empty(), "block_states must be non-empty");
assert!(
!block_spec.name.is_empty(),
"block spec name must be non-empty",
);
expect_block_idx_zero(block_idx, "spatial adaptive exact family", "")?;
Ok(self.linear_constraints.clone().map(ConstraintSet::Dense))
}
fn exact_newton_joint_psi_terms(
&self,
block_states: &[ParameterBlockState],
specs: &[ParameterBlockSpec],
hyper_layout: &crate::custom_family::CustomFamilyHyperLayout,
psi_index: usize,
) -> Result<Option<ExactNewtonJointPsiTerms>, String> {
if hyper_layout.family_axis_count() != 0 {
return Err(
"spatial adaptive exact family does not declare family-owned hyper axes"
.to_string(),
);
}
let derivative_blocks = hyper_layout.design_derivative_blocks();
if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
return Err(SmoothError::dimension_mismatch(format!(
"spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
block_states.len(),
specs.len(),
derivative_blocks.len()
))
.into());
}
derivative_blocks[0]
.get(psi_index)
.ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
let hyper = self
.hyperspecs
.get(psi_index)
.ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
let beta = &block_states[0].beta;
let eval = self.exact_evaluation(beta)?;
let (direct, beta_mixed, betahessian_explicit) =
self.adaptive_hyper_parts(&eval, *hyper)?;
Ok(Some(ExactNewtonJointPsiTerms {
objective_psi: direct,
score_psi: beta_mixed,
hessian_psi: betahessian_explicit,
hessian_psi_operator: None,
}))
}
fn exact_newton_joint_psisecond_order_terms(
&self,
block_states: &[ParameterBlockState],
specs: &[ParameterBlockSpec],
hyper_layout: &crate::custom_family::CustomFamilyHyperLayout,
psi_i: usize,
psi_j: usize,
) -> Result<Option<gam_problem::ExactNewtonJointPsiSecondOrderTerms>, String> {
if hyper_layout.family_axis_count() != 0 {
return Err(
"spatial adaptive exact family does not declare family-owned hyper axes"
.to_string(),
);
}
let derivative_blocks = hyper_layout.design_derivative_blocks();
if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
return Err(SmoothError::dimension_mismatch(format!(
"spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
block_states.len(),
specs.len(),
derivative_blocks.len()
))
.into());
}
derivative_blocks[0]
.get(psi_i)
.ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_i))?;
derivative_blocks[0]
.get(psi_j)
.ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_j))?;
let hyper_i = self
.hyperspecs
.get(psi_i)
.ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_i))?;
let hyper_j = self
.hyperspecs
.get(psi_j)
.ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_j))?;
let beta = &block_states[0].beta;
let eval = self.exact_evaluation(beta)?;
let (objective_psi_psi, score_psi_psi, hessian_psi_psi) =
self.adaptive_explicit_second_order_parts(&eval, *hyper_i, *hyper_j)?;
Ok(Some(gam_problem::ExactNewtonJointPsiSecondOrderTerms {
objective_psi_psi,
score_psi_psi,
hessian_psi_psi,
hessian_psi_psi_operator: None,
}))
}
fn exact_newton_joint_psihessian_directional_derivative(
&self,
block_states: &[ParameterBlockState],
specs: &[ParameterBlockSpec],
hyper_layout: &crate::custom_family::CustomFamilyHyperLayout,
psi_index: usize,
direction: &Array1<f64>,
) -> Result<Option<Array2<f64>>, String> {
if hyper_layout.family_axis_count() != 0 {
return Err(
"spatial adaptive exact family does not declare family-owned hyper axes"
.to_string(),
);
}
let derivative_blocks = hyper_layout.design_derivative_blocks();
if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
return Err(SmoothError::dimension_mismatch(format!(
"spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
block_states.len(),
specs.len(),
derivative_blocks.len()
))
.into());
}
let beta = &block_states[0].beta;
if direction.len() != beta.len() {
return Err(SmoothError::dimension_mismatch(format!(
"spatial adaptive exact family direction length mismatch: got {}, expected {}",
direction.len(),
beta.len()
))
.into());
}
derivative_blocks[0]
.get(psi_index)
.ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
let hyper = self
.hyperspecs
.get(psi_index)
.ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
let eval = self.exact_evaluation(beta)?;
let drift = match hyper.kind {
SpatialAdaptiveHyperKind::LogLambdaMagnitude
| SpatialAdaptiveHyperKind::LogLambdaGradient
| SpatialAdaptiveHyperKind::LogLambdaCurvature => self.adaptive_block_drift_eval(
&eval,
hyper.cache_index,
AdaptiveComponent::from_index(hyper.kind.component_index())?,
HyperDriftKind::Rho,
direction,
)?,
SpatialAdaptiveHyperKind::LogEpsilonMagnitude
| SpatialAdaptiveHyperKind::LogEpsilonGradient
| SpatialAdaptiveHyperKind::LogEpsilonCurvature => self
.adaptive_shared_log_epsilon_drift(
&eval,
hyper.kind.component_index(),
direction,
)?,
};
Ok(Some(drift))
}
}
fn expect_single_block_state<'a>(
block_states: &'a [ParameterBlockState],
family_name: &str,
) -> Result<&'a ParameterBlockState, String> {
crate::block_layout::block_count::validate_block_count::<SmoothError>(
family_name,
1,
block_states.len(),
)?;
Ok(&block_states[0])
}
fn expect_single_blockspec<'a>(
specs: &'a [ParameterBlockSpec],
family_name: &str,
) -> Result<&'a ParameterBlockSpec, String> {
crate::block_layout::block_count::validate_block_count::<SmoothError>(
family_name,
1,
specs.len(),
)?;
Ok(&specs[0])
}
fn expect_block_idx_zero(block_idx: usize, family_name: &str, context: &str) -> Result<(), String> {
if block_idx != 0 {
return Err(SmoothError::invalid_index(format!(
"{family_name} expects block_idx 0{context}, got {block_idx}"
))
.into());
}
Ok::<(), _>(())
}
impl BoundedLinearFamily {
fn bounded_term_derivative_data(
&self,
latent_beta: &Array1<f64>,
) -> Result<
(
Array1<f64>,
Array1<f64>,
Array1<f64>,
Array1<f64>,
Array1<f64>,
),
String,
> {
let p = latent_beta.len();
if p != self.design.ncols() || latent_beta.iter().any(|value| !value.is_finite()) {
return Err(format!(
"bounded coefficient geometry requires {} finite latent coefficients, got {}",
self.design.ncols(),
p
));
}
let mut beta_user = latent_beta.clone();
let mut jac_diag = Array1::<f64>::ones(p);
let mut second_diag = Array1::<f64>::zeros(p);
let mut third_diag = Array1::<f64>::zeros(p);
let mut priorthird = Array1::<f64>::zeros(p);
for term in &self.bounded_terms {
let width = term.max - term.min;
if term.col_idx >= p
|| !term.min.is_finite()
|| !term.max.is_finite()
|| !(width.is_finite() && width > 0.0)
{
return Err(format!(
"bounded coefficient geometry has invalid column/bounds: col={}, p={p}, bounds=({}, {})",
term.col_idx, term.min, term.max
));
}
let (beta, _, db_dtheta, d2b_dtheta2, d3b_dtheta3) =
bounded_latent_derivatives(latent_beta[term.col_idx], term.min, term.max);
if [beta, db_dtheta, d2b_dtheta2, d3b_dtheta3]
.iter()
.any(|value| !value.is_finite())
{
return Err(format!(
"bounded coefficient transform is not representable at column {} and theta={}",
term.col_idx, latent_beta[term.col_idx]
));
}
beta_user[term.col_idx] = beta;
jac_diag[term.col_idx] = db_dtheta;
second_diag[term.col_idx] = d2b_dtheta2;
third_diag[term.col_idx] = d3b_dtheta3;
let (_, _, _, prior_neghess_derivative) =
bounded_prior_terms(latent_beta[term.col_idx], &term.prior)?;
priorthird[term.col_idx] = prior_neghess_derivative;
}
Ok((beta_user, jac_diag, second_diag, third_diag, priorthird))
}
fn user_beta_and_jacobian(
&self,
latent_beta: &Array1<f64>,
) -> Result<(Array1<f64>, Array1<f64>), String> {
let (beta_user, jac_diag, _, _, _) = self.bounded_term_derivative_data(latent_beta)?;
Ok((beta_user, jac_diag))
}
fn nonlinear_offset_from_latent(
&self,
latent_beta: &Array1<f64>,
) -> Result<Array1<f64>, String> {
self.bounded_term_derivative_data(latent_beta)?;
let mut offset = self.offset.clone();
for term in &self.bounded_terms {
let (beta, _, _) =
bounded_latent_to_user(latent_beta[term.col_idx], term.min, term.max);
offset.scaled_add(beta, &self.design.column(term.col_idx));
}
if offset.iter().any(|value| !value.is_finite()) {
return Err("bounded nonlinear offset is not representable".to_string());
}
Ok(offset)
}
fn effective_design_for_latent(&self, jac_diag: &Array1<f64>) -> Array2<f64> {
let mut x_eff = self.design.clone();
for term in &self.bounded_terms {
x_eff
.column_mut(term.col_idx)
.mapv_inplace(|v| v * jac_diag[term.col_idx]);
}
x_eff
}
fn exacthessian_andgradient(
&self,
latent_beta: &Array1<f64>,
) -> Result<
(
StandardFamilyObservationState,
Array2<f64>,
Array1<f64>,
f64,
Array1<f64>,
Array1<f64>,
Array1<f64>,
),
String,
> {
let (_, jac_diag, second_diag, third_diag, priorthird) =
self.bounded_term_derivative_data(latent_beta)?;
let x_eff = self.effective_design_for_latent(&jac_diag);
let eta =
self.designzeroed.dot(latent_beta) + self.nonlinear_offset_from_latent(latent_beta)?;
let obs = evaluate_resolved_standard_family_observations(
&self.likelihood,
self.latent_cloglog_state.as_ref(),
self.mixture_link_state.as_ref(),
self.sas_link_state.as_ref(),
&self.y,
&self.weights,
&eta,
)
.map_err(|e| e.to_string())?;
let mut priorgrad = Array1::<f64>::zeros(latent_beta.len());
let mut prior_neghess = Array2::<f64>::zeros((latent_beta.len(), latent_beta.len()));
let mut prior_loglik = 0.0;
for term in &self.bounded_terms {
let (logp, grad, neghess, _) =
bounded_prior_terms(latent_beta[term.col_idx], &term.prior)?;
prior_loglik += logp;
priorgrad[term.col_idx] += grad;
prior_neghess[[term.col_idx, term.col_idx]] += neghess;
}
let mut hessian = xt_diag_x_dense(x_eff.view(), obs.neghessian_eta.view())?;
let mut gradient = fast_atv(&x_eff, &obs.score);
for term in &self.bounded_terms {
let score_beta = self.design.column(term.col_idx).dot(&obs.score);
hessian[[term.col_idx, term.col_idx]] -= score_beta * second_diag[term.col_idx];
}
hessian += &prior_neghess;
gradient += &priorgrad;
Ok((
obs,
hessian,
gradient,
prior_loglik,
second_diag,
third_diag,
priorthird,
))
}
fn evaluation_from_latent(
&self,
latent_beta: &Array1<f64>,
) -> Result<
(
StandardFamilyObservationState,
Array2<f64>,
Array1<f64>,
f64,
),
String,
> {
let (obs, hessian, gradient, prior_loglik, _, _, _) =
self.exacthessian_andgradient(latent_beta)?;
Ok((obs, hessian, gradient, prior_loglik))
}
}
impl CustomFamily for BoundedLinearFamily {
fn joint_jeffreys_term_required(&self) -> bool {
true
}
fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
let (obs, hessian, gradient, prior_loglik) = self.evaluation_from_latent(latent_beta)?;
Ok(FamilyEvaluation {
log_likelihood: obs.log_likelihood + prior_loglik,
blockworking_sets: vec![BlockWorkingSet::ExactNewton {
gradient,
hessian: SymmetricMatrix::Dense(hessian),
}],
})
}
fn exact_newton_joint_hessian(
&self,
block_states: &[ParameterBlockState],
) -> Result<Option<Array2<f64>>, String> {
let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
let (_, hessian, _, _) = self.evaluation_from_latent(latent_beta)?;
Ok(Some(hessian))
}
fn exact_newton_hessian_directional_derivative(
&self,
block_states: &[ParameterBlockState],
block_idx: usize,
d_beta: &Array1<f64>,
) -> Result<Option<Array2<f64>>, String> {
expect_block_idx_zero(block_idx, "bounded linear family", "")?;
self.exact_newton_joint_hessian_directional_derivative(block_states, d_beta)
}
fn exact_newton_joint_hessian_directional_derivative(
&self,
block_states: &[ParameterBlockState],
d_beta_flat: &Array1<f64>,
) -> Result<Option<Array2<f64>>, String> {
let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
if d_beta_flat.len() != latent_beta.len() {
return Err(SmoothError::dimension_mismatch(format!(
"bounded linear family directional derivative length mismatch: got {}, expected {}",
d_beta_flat.len(),
latent_beta.len()
))
.into());
}
let (obs, _, _, _, second_diag, third_diag, priorthird) =
self.exacthessian_andgradient(latent_beta)?;
let (_, jac_diag, _, _, _) = self.bounded_term_derivative_data(latent_beta)?;
let x_eff = self.effective_design_for_latent(&jac_diag);
let deta = x_eff.dot(d_beta_flat);
let d_neghess_eta = &obs.neghessian_eta_derivative * &deta;
let mut dx_eff = Array2::<f64>::zeros(x_eff.raw_dim());
for term in &self.bounded_terms {
let scale = second_diag[term.col_idx] * d_beta_flat[term.col_idx];
if scale != 0.0 {
let mut col = dx_eff.column_mut(term.col_idx);
col.assign(&self.design.column(term.col_idx));
col.mapv_inplace(|v| v * scale);
}
}
let mut dhessian = xt_diag_x_dense(x_eff.view(), d_neghess_eta.view())?;
let mut wxdx = Array2::<f64>::zeros((x_eff.ncols(), x_eff.ncols()));
for i in 0..x_eff.nrows() {
let wi = obs.neghessian_eta[i];
if wi == 0.0 {
continue;
}
for a in 0..x_eff.ncols() {
let xa = x_eff[[i, a]];
for b in 0..x_eff.ncols() {
wxdx[[a, b]] += wi * (dx_eff[[i, a]] * x_eff[[i, b]] + xa * dx_eff[[i, b]]);
}
}
}
dhessian += &wxdx;
let d_score = -&obs.neghessian_eta * &deta;
for term in &self.bounded_terms {
let score_beta = self.design.column(term.col_idx).dot(&obs.score);
let d_score_beta = self.design.column(term.col_idx).dot(&d_score);
dhessian[[term.col_idx, term.col_idx]] -= d_score_beta * second_diag[term.col_idx]
+ score_beta * third_diag[term.col_idx] * d_beta_flat[term.col_idx];
dhessian[[term.col_idx, term.col_idx]] +=
priorthird[term.col_idx] * d_beta_flat[term.col_idx];
}
Ok(Some(dhessian))
}
fn block_geometry(
&self,
block_states: &[ParameterBlockState],
spec: &ParameterBlockSpec,
) -> Result<(DesignMatrix, Array1<f64>), String> {
if block_states.is_empty() {
return Ok((
DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
self.designzeroed.clone(),
)),
self.offset.clone(),
));
}
let offset = self.nonlinear_offset_from_latent(
&expect_single_block_state(block_states, "bounded linear family")?.beta,
)?;
let x = if spec.design.ncols() == self.designzeroed.ncols() {
self.designzeroed.clone()
} else {
return Err(SmoothError::dimension_mismatch(
"bounded linear family design column mismatch",
)
.into());
};
Ok((
DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(x)),
offset,
))
}
fn block_geometry_is_dynamic(&self) -> bool {
true
}
fn block_geometry_directional_derivative(
&self,
block_states: &[ParameterBlockState],
block_idx: usize,
spec: &ParameterBlockSpec,
d_beta: &Array1<f64>,
) -> Result<Option<BlockGeometryDirectionalDerivative>, String> {
expect_block_idx_zero(
block_idx,
"bounded linear family",
" for geometry derivative",
)?;
expect_single_block_state(block_states, "bounded linear family")?;
if d_beta.len() != spec.design.ncols() {
return Err(SmoothError::dimension_mismatch(format!(
"bounded linear family geometry derivative direction mismatch: got {}, expected {}",
d_beta.len(),
spec.design.ncols()
))
.into());
}
let (_, jac_diag, _, _, _) = self.bounded_term_derivative_data(&block_states[0].beta)?;
let mut d_offset = Array1::<f64>::zeros(self.offset.len());
let has_drift = self
.bounded_terms
.iter()
.any(|term| jac_diag[term.col_idx] != 0.0 && d_beta[term.col_idx] != 0.0);
if !has_drift {
return Ok(Some(BlockGeometryDirectionalDerivative {
d_design: None,
d_offset,
}));
}
for term in &self.bounded_terms {
let col = term.col_idx;
let drift = jac_diag[col] * d_beta[col];
if drift != 0.0 {
d_offset.scaled_add(drift, &self.design.column(col));
}
}
Ok(Some(BlockGeometryDirectionalDerivative {
d_design: None,
d_offset,
}))
}
}
#[inline]
fn dense_diag_gram_chunkrows(p: usize) -> usize {
const MIN_ROWS: usize = 512;
const MAX_ROWS: usize = 2048;
const TARGET_BYTES: usize = 2 * 1024 * 1024;
let bytes_per_row = p.max(1) * std::mem::size_of::<f64>();
(TARGET_BYTES / bytes_per_row).clamp(MIN_ROWS, MAX_ROWS)
}
fn xt_diag_x_dense(x: ArrayView2<'_, f64>, w: ArrayView1<'_, f64>) -> Result<Array2<f64>, String> {
if x.nrows() != w.len() {
return Err(SmoothError::dimension_mismatch("xt_diag_x_dense row mismatch").into());
}
let (n, p) = x.dim();
if n == 0 || p == 0 {
return Ok(Array2::<f64>::zeros((p, p)));
}
const STREAMING_BYTES_THRESHOLD: usize = 8 * 1024 * 1024;
let dense_work_bytes = n
.checked_mul(p)
.and_then(|cells| cells.checked_mul(std::mem::size_of::<f64>()))
.unwrap_or(usize::MAX);
if dense_work_bytes <= STREAMING_BYTES_THRESHOLD {
let mut weighted = x.to_owned();
ndarray::Zip::from(weighted.rows_mut())
.and(w)
.par_for_each(|mut row, wi| row *= *wi);
return Ok(fast_atb(&x, &weighted));
}
let chunkrows = dense_diag_gram_chunkrows(p).min(n);
let mut weighted_chunk = Array2::<f64>::zeros((chunkrows, p));
let mut out = Array2::<f64>::zeros((p, p));
for row_start in (0..n).step_by(chunkrows) {
let rows = (n - row_start).min(chunkrows);
let x_chunk = x.slice(s![row_start..row_start + rows, ..]);
{
let mut chunk = weighted_chunk.slice_mut(s![0..rows, ..]);
for local_row in 0..rows {
let scale = w[row_start + local_row];
if scale == 0.0 {
chunk.row_mut(local_row).fill(0.0);
continue;
}
for col in 0..p {
chunk[[local_row, col]] = x_chunk[[local_row, col]] * scale;
}
}
}
out += &fast_atb(&x_chunk, &weighted_chunk.slice(s![0..rows, ..]));
}
Ok(out)
}
fn trace_of_dense_product(a: &Array2<f64>, b: &Array2<f64>) -> Result<f64, String> {
if a.nrows() != a.ncols() || b.nrows() != b.ncols() || a.nrows() != b.nrows() {
return Err(
SmoothError::dimension_mismatch("trace_of_dense_product dimension mismatch").into(),
);
}
if a.iter().chain(b.iter()).any(|value| !value.is_finite()) {
return Err("trace_of_dense_product requires finite matrices".to_string());
}
let mut trace = gam_linalg::utils::KahanSum::default();
for i in 0..a.nrows() {
for j in 0..a.ncols() {
let term = a[[i, j]] * b[[j, i]];
if !term.is_finite() {
return Err(format!(
"trace_of_dense_product term ({i}, {j}) is not representable"
));
}
trace.add(term);
}
}
let trace = trace.sum();
if !trace.is_finite() {
return Err("trace_of_dense_product sum is not representable".to_string());
}
Ok(trace)
}
fn certify_bounded_edf_interval(
value: f64,
lower: f64,
upper: f64,
dimension: usize,
label: &str,
) -> Result<f64, EstimationError> {
if !(value.is_finite() && lower.is_finite() && upper.is_finite() && lower <= upper) {
crate::bail_invalid_estim!(
"{label} has invalid EDF interval/value: value={value}, interval=[{lower}, {upper}]"
);
}
let scale = 1.0_f64.max(value.abs()).max(lower.abs()).max(upper.abs());
let allowed = 256.0 * f64::EPSILON * (dimension.max(1) as f64).powi(2) * scale;
if value < lower {
if lower - value <= allowed {
return Ok(lower);
}
} else if value > upper {
if value - upper <= allowed {
return Ok(upper);
}
} else {
return Ok(value);
}
crate::bail_invalid_estim!(
"{label}={value} lies outside [{lower}, {upper}] by more than the dense-trace backward-error allowance {allowed}"
)
}
fn exact_bounded_edf(
penalties: &[PenaltySpec],
lambdas: &Array1<f64>,
latent_cov: &Array2<f64>,
) -> Result<(Vec<f64>, Vec<f64>, f64), EstimationError> {
if penalties.len() != lambdas.len() {
crate::bail_invalid_estim!(
"bounded EDF penalty/lambda mismatch: {} penalties vs {} lambdas",
penalties.len(),
lambdas.len()
);
}
if latent_cov.nrows() != latent_cov.ncols() {
crate::bail_invalid_estim!("bounded EDF covariance must be square");
}
let p = latent_cov.nrows();
let mut s_lambda = Array2::<f64>::zeros((p, p));
let mut edf_by_block = Vec::with_capacity(penalties.len());
let mut penalty_block_trace = Vec::with_capacity(penalties.len());
let mut trace_sum = gam_linalg::utils::KahanSum::default();
for (k, ps) in penalties.iter().enumerate() {
let lambda_k = lambdas[k];
if !(lambda_k.is_finite() && lambda_k >= 0.0) {
crate::bail_invalid_estim!(
"bounded EDF smoothing strength at block {k} must be finite and non-negative, got {lambda_k}"
);
}
match ps {
PenaltySpec::Block {
local, col_range, ..
} => {
s_lambda
.slice_mut(ndarray::s![col_range.clone(), col_range.clone()])
.scaled_add(lambda_k, local);
let penalty_rank =
local
.nrows()
.saturating_sub(estimate_penalty_nullity(local).map_err(|e| {
EstimationError::InvalidInput(format!("bounded EDF rank failed: {e}"))
})?);
let cov_block = latent_cov.slice(ndarray::s![col_range.clone(), col_range.clone()]);
let trace_k = lambda_k
* trace_of_dense_product(&cov_block.to_owned(), local)
.map_err(EstimationError::InvalidInput)?;
trace_sum.add(trace_k);
penalty_block_trace.push(trace_k);
let p_k = penalty_rank as f64;
edf_by_block.push(certify_bounded_edf_interval(
p_k - trace_k,
0.0,
p_k,
p,
&format!("bounded EDF block {k}"),
)?);
}
PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
s_lambda.scaled_add(lambda_k, m);
let penalty_rank = p.saturating_sub(estimate_penalty_nullity(m).map_err(|e| {
EstimationError::InvalidInput(format!("bounded EDF rank failed: {e}"))
})?);
let trace_k = lambda_k
* trace_of_dense_product(latent_cov, m)
.map_err(EstimationError::InvalidInput)?;
trace_sum.add(trace_k);
penalty_block_trace.push(trace_k);
let p_k = penalty_rank as f64;
edf_by_block.push(certify_bounded_edf_interval(
p_k - trace_k,
0.0,
p_k,
p,
&format!("bounded EDF block {k}"),
)?);
}
}
}
let nullity_total = estimate_penalty_nullity(&s_lambda)
.map_err(|e| EstimationError::InvalidInput(format!("bounded EDF nullity failed: {e}")))?
as f64;
let trace_sum = trace_sum.sum();
let edf_total = certify_bounded_edf_interval(
p as f64 - trace_sum,
nullity_total,
p as f64,
p,
"bounded total EDF",
)?;
Ok((edf_by_block, penalty_block_trace, edf_total))
}
fn certified_bounded_posterior_covariance(
precision: &Array2<f64>,
label: &'static str,
) -> Result<Array2<f64>, EstimationError> {
gam_linalg::utils::certified_spd_inverse(precision, label)
.map(gam_linalg::utils::CertifiedSpdInverse::into_inverse)
.map_err(|error| {
EstimationError::InvalidInput(format!(
"bounded posterior covariance requires an exact SPD precision: {error}"
))
})
}
fn transform_bounded_latent_precision_to_user_internal(
latent_precision: &Array2<f64>,
jac_diag: &Array1<f64>,
) -> Result<Array2<f64>, EstimationError> {
let p = latent_precision.nrows();
if latent_precision.ncols() != p || jac_diag.len() != p {
crate::bail_invalid_estim!(
"bounded precision transform dimension mismatch: precision is {}x{}, jacobian has {} entries",
latent_precision.nrows(),
latent_precision.ncols(),
jac_diag.len()
);
}
let mut out = latent_precision.clone();
for i in 0..p {
let scale = jac_diag[i];
if !scale.is_finite() || scale <= 0.0 {
crate::bail_invalid_estim!(
"bounded precision transform requires a positive finite coefficient jacobian; column {i} has {scale}"
);
}
if scale != 1.0 {
out.row_mut(i).mapv_inplace(|v| v / scale);
out.column_mut(i).mapv_inplace(|v| v / scale);
}
}
Ok(out)
}
fn fit_bounded_term_collection_with_design(
y: ArrayView1<'_, f64>,
weights: ArrayView1<'_, f64>,
offset: ArrayView1<'_, f64>,
spec: &TermCollectionSpec,
design: &TermCollectionDesign,
heuristic_lambdas: Option<&[f64]>,
family: LikelihoodSpec,
options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError> {
let conditioning_cols: Vec<usize> = spec
.linear_terms
.iter()
.enumerate()
.filter_map(|(j, linear)| {
(!linear.double_penalty).then_some(design.intercept_range.end + j)
})
.collect();
let conditioning = LinearFitConditioning::from_columns(design, &conditioning_cols);
let dense_design = design.design.to_dense_cow();
let fit_design = conditioning.apply_to_design(&dense_design);
let fit_penalties = conditioning
.transform_blockwise_penalties_to_internal(&design.penalties, design.design.ncols());
if design.linear_constraints.is_some() {
crate::bail_invalid_estim!(
"bounded() terms are not yet compatible with explicit linear constraints"
);
}
let mut bounded_terms = Vec::<BoundedLinearTermMeta>::new();
for (j, term) in spec.linear_terms.iter().enumerate() {
if term.double_penalty
&& matches!(
term.coefficient_geometry,
LinearCoefficientGeometry::Bounded { .. }
)
{
crate::bail_invalid_estim!(
"bounded linear term '{}' cannot also use double_penalty",
term.name
);
}
if let LinearCoefficientGeometry::Bounded { min, max, prior } =
term.coefficient_geometry.clone()
{
let col_idx = design.intercept_range.end + j;
let (min_internal, max_internal) = conditioning.internal_bounds_for(col_idx, min, max);
bounded_terms.push(BoundedLinearTermMeta {
col_idx,
min: min_internal,
max: max_internal,
prior,
});
}
}
if bounded_terms.is_empty() {
crate::bail_invalid_estim!("internal bounded fit path called with no bounded terms");
}
let mut designzeroed = fit_design.clone();
let mut initial_beta = Array1::<f64>::zeros(fit_design.ncols());
for term in &bounded_terms {
designzeroed.column_mut(term.col_idx).fill(0.0);
initial_beta[term.col_idx] = 0.0;
}
let initial_log_lambdas = heuristic_lambdas
.map(|vals| Array1::from_vec(vals.to_vec()))
.unwrap_or_else(|| Array1::zeros(fit_penalties.len()));
if initial_log_lambdas.len() != fit_penalties.len() {
crate::bail_invalid_estim!(
"heuristic lambda length mismatch for bounded model: got {}, expected {}",
initial_log_lambdas.len(),
fit_penalties.len()
);
}
let glm_likelihood = gam_spec::GlmLikelihoodSpec::canonical(family);
let resolved_likelihood_scale = glm_likelihood
.resolved_scale()
.map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
let is_beta_logistic = glm_likelihood.spec.is_binomial_beta_logistic();
let family_adapter = BoundedLinearFamily {
likelihood: glm_likelihood.clone(),
latent_cloglog_state: options.latent_cloglog,
mixture_link_state: options
.mixture_link
.clone()
.as_ref()
.map(state_fromspec)
.transpose()
.map_err(EstimationError::InvalidInput)?,
sas_link_state: options
.sas_link
.map(|spec| {
if is_beta_logistic {
state_from_beta_logisticspec(spec)
} else {
state_from_sasspec(spec)
}
})
.transpose()
.map_err(EstimationError::InvalidInput)?,
y: y.to_owned(),
weights: weights.to_owned(),
design: fit_design.clone(),
designzeroed: designzeroed.clone(),
offset: offset.to_owned(),
bounded_terms: bounded_terms.clone(),
};
let blockspec = ParameterBlockSpec {
name: "eta".to_string(),
design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(designzeroed)),
offset: offset.to_owned(),
penalties: fit_penalties
.iter()
.map(|ps| match ps {
PenaltySpec::Block {
local, col_range, ..
} => PenaltyMatrix::Blockwise {
local: local.clone(),
col_range: col_range.clone(),
total_dim: design.design.ncols(),
},
PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
PenaltyMatrix::Dense(m.clone())
}
})
.collect(),
nullspace_dims: design.nullspace_dims.clone(),
initial_log_lambdas,
initial_beta: Some(initial_beta),
gauge_priority: 100,
jacobian_callback: Some(Arc::new(BoundedEffectiveJacobian {
design: fit_design.clone(),
bounded_terms: bounded_terms.clone(),
})),
stacked_design: None,
stacked_offset: None,
};
let fit = fit_custom_family(
&family_adapter,
&[blockspec],
&BlockwiseFitOptions {
inner_max_cycles: options.max_iter,
inner_tol: options.tol,
outer_max_iter: options.max_iter,
outer_tol: options.tol,
compute_covariance: false,
..BlockwiseFitOptions::default()
},
)
.map_err(EstimationError::CustomFamily)?;
let latent_beta = fit.block_states[0].beta.clone();
let (beta_user_internal, jac_diag) = family_adapter
.user_beta_and_jacobian(&latent_beta)
.map_err(EstimationError::InvalidInput)?;
let beta_user = conditioning.backtransform_beta(&beta_user_internal);
let (eta_state, h_data, _, _) = family_adapter
.evaluation_from_latent(&latent_beta)
.map_err(EstimationError::InvalidInput)?;
let p_fit = fit_design.ncols();
let mut s_lambda_internal = Array2::<f64>::zeros((p_fit, p_fit));
for (k, penalty) in fit_penalties.iter().enumerate() {
match penalty {
PenaltySpec::Block {
local, col_range, ..
} => {
s_lambda_internal
.slice_mut(ndarray::s![col_range.clone(), col_range.clone()])
.scaled_add(fit.lambdas[k], local);
}
PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
s_lambda_internal.scaled_add(fit.lambdas[k], m);
}
}
}
let mut latent_precision = h_data.clone();
latent_precision += &s_lambda_internal;
let user_precision_internal =
transform_bounded_latent_precision_to_user_internal(&latent_precision, &jac_diag)?;
let penalized_hessian =
conditioning.transform_penalized_hessian_to_original(&user_precision_internal);
let beta_covariance_unscaled = if options.compute_inference {
Some(certified_bounded_posterior_covariance(
&penalized_hessian,
"bounded user-scale posterior precision",
)?)
} else {
None
};
let latent_cov = if options.compute_inference {
Some(certified_bounded_posterior_covariance(
&latent_precision,
"bounded latent posterior precision",
)?)
} else {
None
};
let s_lambda_original = weighted_blockwise_penalty_sum(
&design.penalties,
fit.lambdas.as_slice().unwrap(),
design.design.ncols(),
);
let penalty_term = beta_user.dot(&s_lambda_original.dot(&beta_user));
let deviance = -2.0 * eta_state.log_likelihood;
let (edf_by_block, penalty_block_trace, edf_total) = if let Some(cov) = latent_cov.as_ref() {
exact_bounded_edf(&fit_penalties, &fit.lambdas, cov)?
} else {
(
vec![0.0; fit_penalties.len()],
vec![0.0; fit_penalties.len()],
0.0,
)
};
let profiled_gaussian_standard_deviation = if matches!(
resolved_likelihood_scale,
gam_spec::ResolvedLikelihoodScale::ProfiledGaussian
) {
let residual_dof = if options.compute_inference {
y.len() as f64 - edf_total
} else {
y.len() as f64
};
if !(residual_dof.is_finite() && residual_dof > 0.0) {
return Err(EstimationError::InvalidInput(format!(
"bounded Gaussian residual degrees of freedom must be finite and positive, got n={} minus edf={edf_total} = {residual_dof}",
y.len()
)));
}
if !(deviance.is_finite() && deviance >= 0.0) {
return Err(EstimationError::InvalidInput(format!(
"bounded Gaussian deviance must be finite and non-negative, got {deviance}"
)));
}
let variance = deviance / residual_dof;
if !variance.is_finite() {
return Err(EstimationError::InvalidInput(format!(
"bounded Gaussian residual variance is not representable: {deviance}/{residual_dof}"
)));
}
Some(variance.sqrt())
} else {
None
};
let dispersion = gam_solve::estimate::dispersion_from_likelihood(
&glm_likelihood,
profiled_gaussian_standard_deviation,
)?;
let standard_deviation = dispersion.phi().sqrt();
let cov_scale = glm_likelihood
.coefficient_covariance_scale(dispersion.phi())
.map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
let beta_covariance = beta_covariance_unscaled.map(|mut cov| {
if cov_scale != 1.0 {
cov.mapv_inplace(|v| v * cov_scale);
}
cov
});
if let Some(covariance) = beta_covariance.as_ref()
&& covariance.iter().any(|value| !value.is_finite())
{
return Err(EstimationError::InvalidInput(
"bounded coefficient covariance scaling produced a non-finite value".to_string(),
));
}
let beta_standard_errors = beta_covariance
.as_ref()
.map(gam_problem::se_from_covariance)
.transpose()
.map_err(|err| {
EstimationError::InvalidInput(format!(
"bounded coefficient covariance cannot produce standard errors: {err}"
))
})?;
let working_response = exact_standard_working_response(&eta_state)?;
let geometry = Some(gam_solve::estimate::FitGeometry {
coefficient_gauge: gam_problem::gauge::Gauge::identity(&[beta_user.len()]),
penalized_hessian: penalized_hessian.clone().into(),
constrained_posterior: None,
working: Some(gam_solve::estimate::WorkingGeometry {
weights: eta_state.fisherweight.clone(),
response: working_response,
}),
});
let max_abs_eta = eta_state
.eta
.iter()
.fold(0.0_f64, |acc, &v| acc.max(v.abs()));
Ok(FittedTermCollection {
fit: {
let log_lambdas =
checked_fit_log_lambdas(&fit.lambdas, "final fitted term collection")?;
let inf = FitInference {
edf_by_block,
penalty_block_trace,
edf_total,
smoothing_correction: None,
smoothing_correction_method: None,
smoothing_correction_first_order: None,
smoothing_correction_method_first_order: None,
penalized_hessian: penalized_hessian.clone().into(),
reparam_qs: None,
dispersion,
beta_covariance: beta_covariance
.clone()
.map(gam_problem::dispersion_cov::PhiScaledCovariance::from),
beta_standard_errors,
beta_covariance_corrected: None,
beta_standard_errors_corrected: None,
beta_covariance_frequentist: None,
coefficient_influence: None,
weighted_gram: None,
bias_correction_beta: None,
bias_correction_jacobian: None,
};
let covariance_conditional = beta_covariance;
let pirls_status_val = gam_solve::pirls::PirlsStatus::Converged;
UnifiedFitResult::try_from_parts(UnifiedFitResultParts {
blocks: vec![gam_solve::estimate::FittedBlock {
beta: beta_user.clone(),
role: gam_problem::BlockRole::Mean,
edf: edf_total,
lambdas: fit.lambdas.clone(),
}],
log_lambdas,
lambdas: fit.lambdas,
likelihood_scale: glm_likelihood.scale,
likelihood_family: Some(glm_likelihood.spec),
log_likelihood_normalization: gam_spec::LogLikelihoodNormalization::UserProvided,
log_likelihood: eta_state.log_likelihood,
deviance,
reml_score: fit.penalized_objective,
stable_penalty_term: penalty_term,
penalized_objective: fit.penalized_objective,
used_device: false,
outer_iterations: fit.outer_iterations,
outer_converged: true,
outer_gradient_norm: fit.outer_gradient_norm,
standard_deviation,
covariance_conditional,
covariance_corrected: None,
inference: Some(inf),
fitted_link: gam_solve::estimate::FittedLinkState::Standard(None),
geometry,
block_states: Vec::new(),
pirls_status: pirls_status_val,
max_abs_eta,
constraint_kkt: None,
artifacts: gam_solve::estimate::FitArtifacts {
pirls: None,
..Default::default()
},
inner_cycles: 0,
})?
},
design: design.clone(),
adaptive_diagnostics: None,
})
}
fn enforce_term_constraint_feasibility(
design: &TermCollectionDesign,
fit: &UnifiedFitResult,
) -> Result<(), EstimationError> {
const CONSTRAINT_FEASIBILITY_RAW_TOL: f64 = 1e-7;
let tol = CONSTRAINT_FEASIBILITY_RAW_TOL;
let smooth_start = design
.design
.ncols()
.saturating_sub(design.smooth.total_smooth_cols());
let mut violations: Vec<String> = Vec::new();
for term in &design.smooth.terms {
let gr = (smooth_start + term.coeff_range.start)..(smooth_start + term.coeff_range.end);
let beta_local = fit.beta.slice(s![gr.clone()]).to_owned();
if let Some(lb) = term.lower_bounds_local.as_ref() {
let mut worst = 0.0_f64;
let mut worst_idx = 0usize;
for i in 0..lb.len().min(beta_local.len()) {
if lb[i].is_finite() {
let viol = (lb[i] - beta_local[i]).max(0.0);
if viol > worst {
worst = viol;
worst_idx = i;
}
}
}
if worst > tol {
violations.push(format!(
"term='{}' kind=lower-bound maxviolation={:.3e} coeff_index={}",
term.name, worst, worst_idx
));
}
}
if let Some(lin) = term.linear_constraints_local.as_ref() {
let mut worst = 0.0_f64;
let mut worstrow = 0usize;
for i in 0..lin.a.nrows() {
let norm = lin.a.row(i).dot(&lin.a.row(i)).sqrt();
let inv = if norm > 0.0 { 1.0 / norm } else { 0.0 };
let s = (lin.a.row(i).dot(&beta_local) - lin.b[i]) * inv;
let viol = (-s).max(0.0);
if viol > worst {
worst = viol;
worstrow = i;
}
}
if worst > tol {
violations.push(format!(
"term='{}' kind=linear-inequality maxviolation={:.3e} row={}",
term.name, worst, worstrow
));
}
}
}
if !violations.is_empty() {
let mut msg = format!(
"constraint violation after fit ({} violating term constraints): {}",
violations.len(),
violations.join(" | ")
);
if let Some(kkt) = fit.constraint_kkt.as_ref() {
msg.push_str(&format!(
"; KKT[primal={:.3e}, dual={:.3e}, comp={:.3e}, stat={:.3e}]",
kkt.primal_feasibility, kkt.dual_feasibility, kkt.complementarity, kkt.stationarity
));
}
return Err(EstimationError::ParameterConstraintViolation(msg));
}
Ok(())
}
fn stratified_spatial_subsample(
data: ArrayView2<'_, f64>,
spec: &TermCollectionSpec,
target_size: usize,
) -> Vec<usize> {
use rand::SeedableRng;
use rand::rngs::StdRng;
use rand::seq::SliceRandom;
let n = data.nrows();
if n <= target_size {
return (0..n).collect();
}
let spatial_cols: Option<Vec<usize>> =
spec.smooth_terms.iter().find_map(|term| match &term.basis {
SmoothBasisSpec::ThinPlate { feature_cols, .. }
| SmoothBasisSpec::Matern { feature_cols, .. }
| SmoothBasisSpec::Duchon { feature_cols, .. } => {
if !feature_cols.is_empty() {
Some(feature_cols.clone())
} else {
None
}
}
_ => None,
});
let cols = match spatial_cols {
Some(c) if !c.is_empty() => c,
_ => {
let mut rng = StdRng::seed_from_u64(spatial_subsample_seed(data, &[], target_size));
let mut indices: Vec<usize> = (0..n).collect();
indices.shuffle(&mut rng);
indices.truncate(target_size);
indices.sort_unstable();
return indices;
}
};
let mut rng = StdRng::seed_from_u64(spatial_subsample_seed(data, &cols, target_size));
let d = cols.len();
let mut mins = vec![f64::INFINITY; d];
let mut maxs = vec![f64::NEG_INFINITY; d];
for i in 0..n {
for (ax, &col) in cols.iter().enumerate() {
let v = data[[i, col]];
if v < mins[ax] {
mins[ax] = v;
}
if v > maxs[ax] {
maxs[ax] = v;
}
}
}
const TARGET_POINTS_PER_CELL: usize = 5;
let total_cells_target = (target_size / TARGET_POINTS_PER_CELL).max(1);
let cells_per_axis = ((total_cells_target as f64).powf(1.0 / d as f64)).ceil() as usize;
let cells_per_axis = cells_per_axis.max(1);
let mut cell_members: std::collections::HashMap<Vec<usize>, Vec<usize>> =
std::collections::HashMap::new();
for i in 0..n {
let mut cell_key = Vec::with_capacity(d);
for (ax, &col) in cols.iter().enumerate() {
let range = maxs[ax] - mins[ax];
let cell = if range <= 0.0 {
0
} else {
let frac = (data[[i, col]] - mins[ax]) / range;
(frac * cells_per_axis as f64).floor() as usize
};
cell_key.push(cell.min(cells_per_axis - 1));
}
cell_members.entry(cell_key).or_default().push(i);
}
let mut selected: Vec<usize> = Vec::with_capacity(target_size);
let mut remaining_budget = target_size;
let mut remaining_population = n;
let mut cells: Vec<(Vec<usize>, Vec<usize>)> = cell_members.into_iter().collect();
cells.sort_by(|a, b| a.0.cmp(&b.0));
for (_, members) in &mut cells {
if remaining_budget == 0 {
break;
}
let alloc = ((members.len() as f64 / remaining_population as f64) * remaining_budget as f64)
.round() as usize;
let alloc = alloc.max(1).min(members.len()).min(remaining_budget);
members.shuffle(&mut rng);
selected.extend_from_slice(&members[..alloc]);
remaining_budget = remaining_budget.saturating_sub(alloc);
remaining_population = remaining_population.saturating_sub(members.len());
}
if selected.len() > target_size {
selected.shuffle(&mut rng);
selected.truncate(target_size);
}
selected.sort_unstable();
selected
}
fn spatial_subsample_seed(
data: ArrayView2<'_, f64>,
spatial_cols: &[usize],
target_size: usize,
) -> u64 {
let mut state = 0x5350_4154_4941_4C53_u64;
spatial_seed_mix(&mut state, data.nrows() as u64);
spatial_seed_mix(&mut state, data.ncols() as u64);
spatial_seed_mix(&mut state, target_size as u64);
spatial_seed_mix(&mut state, spatial_cols.len() as u64);
for &col in spatial_cols {
spatial_seed_mix(&mut state, col as u64);
}
if data.nrows() > 0 {
let mid = data.nrows() / 2;
let last = data.nrows() - 1;
for &row in &[0usize, mid, last] {
for &col in spatial_cols {
let value = data[[row, col]];
spatial_seed_mix(&mut state, value.to_bits());
}
}
}
state
}
#[inline]
fn spatial_seed_mix(state: &mut u64, value: u64) {
let mut s = value.wrapping_add(*state);
let z = gam_linalg::utils::splitmix64(&mut s);
*state ^= z;
*state = (*state).rotate_left(27).wrapping_mul(0x3C79_AC49_2BA7_B653);
}
fn sampled_rows(data: ArrayView2<'_, f64>, indices: &[usize]) -> Array2<f64> {
let mut sampled = Array2::<f64>::zeros((indices.len(), data.ncols()));
for (new_row, &orig_row) in indices.iter().enumerate() {
sampled.row_mut(new_row).assign(&data.row(orig_row));
}
sampled
}
fn spatial_term_user_centers(term: &SmoothTermSpec) -> Option<ArrayView2<'_, f64>> {
match spatial_term_center_strategy(term) {
Some(CenterStrategy::UserProvided(centers)) => Some(centers.view()),
_ => None,
}
}
fn finite_centered_axis_contrasts(values: &[f64], expected_dim: usize) -> Option<Vec<f64>> {
if values.len() != expected_dim || expected_dim <= 1 {
return None;
}
if values.iter().any(|value| !value.is_finite()) {
return None;
}
Some(center_aniso_log_scales(values))
}
fn blended_pilot_axis_contrasts(
pilot_data: ArrayView2<'_, f64>,
term: &SmoothTermSpec,
centers: ArrayView2<'_, f64>,
) -> Result<Option<Vec<f64>>, BasisError> {
let d = centers.ncols();
if d <= 1 {
return Ok(None);
}
let center_eta = initial_aniso_contrasts(centers);
let standardized_data = standardized_spatial_term_data(pilot_data, term)?;
let data_eta = finite_centered_axis_contrasts(
&initial_aniso_contrasts(standardized_data.view()),
d,
);
let Some(center_eta) = finite_centered_axis_contrasts(¢er_eta, d) else {
return Ok(None);
};
let blended = match data_eta {
Some(data_eta) => center_eta
.iter()
.zip(data_eta.iter())
.map(|(&from_centers, &from_data)| 0.5 * (from_centers + from_data))
.collect::<Vec<_>>(),
None => center_eta,
};
Ok(finite_centered_axis_contrasts(&blended, d))
}
fn apply_pilot_spatial_psi_reseed(
pilot_data: ArrayView2<'_, f64>,
spec: &TermCollectionSpec,
spatial_terms: &[usize],
kappa_options: &SpatialLengthScaleOptimizationOptions,
) -> Result<TermCollectionSpec, EstimationError> {
let dims_per_term = spatial_dims_per_term(spec, spatial_terms);
let use_aniso = has_aniso_terms(spec, spatial_terms);
let log_kappa0 = if use_aniso {
SpatialLogKappaCoords::from_length_scales_aniso(spec, spatial_terms, kappa_options)
} else {
SpatialLogKappaCoords::from_length_scales(spec, spatial_terms, kappa_options)
};
let log_kappa0 = log_kappa0
.reseed_from_data(pilot_data, spec, spatial_terms, kappa_options)
.map_err(EstimationError::BasisError)?;
let log_kappa_lower = if use_aniso {
SpatialLogKappaCoords::lower_bounds_aniso_from_data(
pilot_data,
spec,
spatial_terms,
&dims_per_term,
kappa_options,
)
} else {
SpatialLogKappaCoords::lower_bounds_from_data(
pilot_data,
spec,
spatial_terms,
kappa_options,
)
}
.map_err(EstimationError::BasisError)?;
let log_kappa_upper = if use_aniso {
SpatialLogKappaCoords::upper_bounds_aniso_from_data(
pilot_data,
spec,
spatial_terms,
&dims_per_term,
kappa_options,
)
} else {
SpatialLogKappaCoords::upper_bounds_from_data(
pilot_data,
spec,
spatial_terms,
kappa_options,
)
}
.map_err(EstimationError::BasisError)?;
log_kappa0
.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper)
.apply_tospec(spec, spatial_terms)
}
pub(crate) fn apply_spatial_anisotropy_pilot_initializer(
data: ArrayView2<'_, f64>,
spec: &mut TermCollectionSpec,
spatial_terms: &[usize],
target_size: usize,
kappa_options: &SpatialLengthScaleOptimizationOptions,
) -> Result<usize, EstimationError> {
if target_size == 0 || data.nrows() <= target_size.saturating_mul(2) || spatial_terms.is_empty()
{
return Ok(0);
}
if !has_aniso_terms(spec, spatial_terms) {
return Ok(0);
}
let indices = stratified_spatial_subsample(data, spec, target_size);
let pilot_data = sampled_rows(data, &indices);
let mut working = spec.clone();
let mut updated_terms = 0usize;
const GEOMETRY_UPDATES: usize = 2;
for pass in 0..GEOMETRY_UPDATES {
let planned_terms = plan_joint_spatial_centers_for_term_blocks(
pilot_data.view(),
&[working.smooth_terms.clone()],
)
.and_then(|mut blocks| {
blocks.pop().ok_or_else(|| {
BasisError::InvalidInput(
"pilot geometry initializer produced no smooth-term block".to_string(),
)
})
})
.map_err(EstimationError::BasisError)?;
for &term_idx in spatial_terms {
let Some(current_eta) = get_spatial_aniso_log_scales(&working, term_idx) else {
continue;
};
let Some(d) = get_spatial_feature_dim(&working, term_idx) else {
continue;
};
if d <= 1 || current_eta.len() != d {
continue;
}
let Some(planned_term) = planned_terms.get(term_idx) else {
continue;
};
let Some(centers) = spatial_term_user_centers(planned_term) else {
continue;
};
let Some(eta) = blended_pilot_axis_contrasts(
pilot_data.view(),
planned_term,
centers,
)
.map_err(EstimationError::BasisError)?
else {
continue;
};
set_spatial_aniso_log_scales(&mut working, term_idx, eta)?;
updated_terms += usize::from(pass == 0);
}
working = apply_pilot_spatial_psi_reseed(
pilot_data.view(),
&working,
spatial_terms,
kappa_options,
)?;
}
if updated_terms > 0 {
log::info!(
"[spatial-kappa] initialized anisotropy from {}-row pilot geometry for {} spatial term(s); proceeding to full-data optimization",
indices.len(),
updated_terms
);
*spec = working;
}
Ok(updated_terms)
}
pub(crate) fn spatial_length_scale_term_indices(spec: &TermCollectionSpec) -> Vec<usize> {
spec.smooth_terms
.iter()
.enumerate()
.filter_map(|(idx, _)| spatial_term_supports_hyper_optimization(spec, idx).then_some(idx))
.collect()
}
fn fit_score(fit: &UnifiedFitResult) -> f64 {
if fit.reml_score.is_finite() {
return fit.reml_score;
}
let score = 0.5 * fit.deviance + 0.5 * fit.stable_penalty_term;
if score.is_finite() {
score
} else {
f64::INFINITY
}
}
fn is_recoverable_trial_point_error(err: &EstimationError) -> bool {
matches!(err, EstimationError::BasisError(_))
|| err.is_inner_solve_retreat()
|| is_recoverable_fit_inference_finiteness_error(err)
}
fn is_recoverable_fit_inference_finiteness_error(err: &EstimationError) -> bool {
let EstimationError::InvalidInput(message) = err else {
return false;
};
message.contains("must be finite")
&& [
"fit_result.beta_covariance_frequentist",
"fit_result.coefficient_influence",
"fit_result.weighted_gram",
]
.iter()
.any(|field| message.contains(field))
}
#[cfg(test)]
mod spatial_trial_recovery_tests {
use super::*;
#[test]
fn nonfinite_frequentist_covariance_is_recoverable_trial_point() {
let err = EstimationError::InvalidInput(
"fit_result.beta_covariance_frequentist[0] must be finite, got NaN".to_string(),
);
assert!(
is_recoverable_trial_point_error(&err),
"singular trial-point curvature should make spatial κ retreat, not abort"
);
}
#[test]
fn arbitrary_invalid_input_remains_fatal_trial_point_error() {
let err = EstimationError::InvalidInput("outer rho bounds are invalid".to_string());
assert!(
!is_recoverable_trial_point_error(&err),
"the spatial κ recovery gate must not mask unrelated invalid inputs"
);
}
}
fn require_successful_spatial_optimization_result<T>(
initial_score: f64,
result: Result<Option<(T, f64)>, EstimationError>,
) -> Result<T, EstimationError> {
match result {
Ok(Some((value, exact_score))) => {
const SCORE_DRIFT_ABS_TOL: f64 = 1e-6;
const SCORE_DRIFT_REL_TOL: f64 = 1e-8;
let tol = SCORE_DRIFT_ABS_TOL.max(initial_score.abs() * SCORE_DRIFT_REL_TOL);
if exact_score <= initial_score + tol {
Ok(value)
} else {
Err(EstimationError::RemlOptimizationFailed(format!(
"spatial kappa optimization made REML score worse ({initial_score:.6e} -> {exact_score:.6e})"
)))
}
}
Ok(None) => Err(EstimationError::RemlOptimizationFailed(
"spatial kappa optimization is unavailable for one or more eligible spatial terms"
.to_string(),
)),
Err(err) => Err(EstimationError::RemlOptimizationFailed(format!(
"spatial kappa optimization failed: {err}"
))),
}
}
fn external_opts_for_design(
family: &LikelihoodSpec,
design: &TermCollectionDesign,
options: &FitOptions,
) -> ExternalOptimOptions {
ExternalOptimOptions {
family: family.clone(),
latent_cloglog: options.latent_cloglog,
mixture_link: options.mixture_link.clone(),
optimize_mixture: options.optimize_mixture,
sas_link: options.sas_link,
optimize_sas: options.optimize_sas,
compute_inference: options.compute_inference,
skip_rho_posterior_inference: options.skip_rho_posterior_inference,
max_iter: options.max_iter,
tol: options.tol,
nullspace_dims: design.nullspace_dims.clone(),
linear_constraints: design.linear_constraints.clone(),
firth_bias_reduction: Some(options.firth_bias_reduction),
penalty_shrinkage_floor: options.penalty_shrinkage_floor,
rho_prior: options.rho_prior.clone(),
kronecker_penalty_system: design.kronecker_penalty_system(),
kronecker_factored: design
.smooth
.terms
.iter()
.find_map(|t| t.kronecker_factored.clone()),
persist_warm_start_disk: options.persist_warm_start_disk,
}
}
fn evaluate_joint_reml_outer_eval_at_theta(
evaluator: &mut gam_solve::estimate::ExternalJointHyperEvaluator<'_>,
design: &TermCollectionDesign,
theta: &Array1<f64>,
rho_dim: usize,
hyper_dirs: Vec<gam_solve::estimate::reml::DirectionalHyperParam>,
warm_start_beta: Option<ArrayView1<'_, f64>>,
order: gam_solve::rho_optimizer::OuterEvalOrder,
design_revision: Option<u64>,
) -> Result<(f64, Array1<f64>, gam_problem::HessianValue), EstimationError> {
evaluator.evaluate_with_order(
&design.design,
&design.penalties,
&design.nullspace_dims,
design.linear_constraints.clone(),
theta,
rho_dim,
hyper_dirs,
warm_start_beta,
"evaluate_joint_reml_outer_eval_at_theta",
order,
design_revision,
)
}
fn evaluate_joint_reml_efs_at_theta(
evaluator: &mut gam_solve::estimate::ExternalJointHyperEvaluator<'_>,
design: &TermCollectionDesign,
theta: &Array1<f64>,
rho_dim: usize,
hyper_dirs: Vec<gam_solve::estimate::reml::DirectionalHyperParam>,
warm_start_beta: Option<ArrayView1<'_, f64>>,
design_revision: Option<u64>,
) -> Result<gam_problem::EfsEval, EstimationError> {
evaluator.evaluate_efs(
&design.design,
&design.penalties,
&design.nullspace_dims,
design.linear_constraints.clone(),
theta,
rho_dim,
hyper_dirs,
warm_start_beta,
"evaluate_joint_reml_efs_at_theta",
design_revision,
)
}
fn exact_joint_spatial_outer_hessian_available(
family: &LikelihoodSpec,
design: &TermCollectionDesign,
) -> bool {
let family_supported = match &family.response {
ResponseFamily::Gaussian
| ResponseFamily::Binomial
| ResponseFamily::Poisson
| ResponseFamily::Tweedie { .. }
| ResponseFamily::NegativeBinomial { .. }
| ResponseFamily::Beta { .. }
| ResponseFamily::Gamma
| ResponseFamily::RoystonParmar => true,
};
family_supported && design.design.ncols() > 0
}
fn try_build_spatial_term_log_kappa_derivativeinfo(
data: ArrayView2<'_, f64>,
resolvedspec: &TermCollectionSpec,
design: &TermCollectionDesign,
term_idx: usize,
) -> Result<Option<SpatialPsiDerivative>, EstimationError> {
let Some((
global_range,
total_p,
x_psi_local,
s_psi_local_check,
x_psi_psi_local,
s_psi_psi_local,
s_psi_components_local,
s_psi_psi_components_local,
implicit_operator,
)) = try_build_spatial_term_log_kappa_derivative(data, resolvedspec, design, term_idx)?
else {
return Ok(None);
};
let Some(penalty_range) = design
.smooth_term_penalty_range(term_idx)
.map_err(EstimationError::InvalidInput)?
else {
return Ok(None);
};
let penalty_start = penalty_range.start;
if s_psi_components_local.is_empty() || s_psi_psi_components_local.is_empty() {
return Ok(None);
}
if s_psi_components_local.len() != s_psi_psi_components_local.len() {
return Ok(None);
}
let penalty_indices = (0..s_psi_components_local.len())
.map(|j| penalty_start + j)
.collect::<Vec<_>>();
let penalty_index = penalty_indices[0];
if s_psi_local_check.nrows() == 0 || s_psi_psi_local.nrows() == 0 {
return Ok(None);
}
Ok(Some(SpatialPsiDerivative {
penalty_index,
penalty_indices,
global_range,
total_p,
x_psi_local,
s_psi_components_local,
x_psi_psi_local,
s_psi_psi_components_local,
aniso_group_id: None,
aniso_cross_designs: None,
aniso_cross_penalty_provider: None,
implicit_operator,
implicit_axis: 0,
}))
}
pub(crate) fn try_build_spatial_log_kappa_derivativeinfo_list(
data: ArrayView2<'_, f64>,
resolvedspec: &TermCollectionSpec,
design: &TermCollectionDesign,
spatial_terms: &[usize],
) -> Result<Option<Vec<SpatialPsiDerivative>>, EstimationError> {
let mut out = Vec::new();
let mut aniso_gid = 0usize;
for &term_idx in spatial_terms {
if spatial_term_uses_per_axis_psi(resolvedspec, term_idx) {
if let Some(entries) = try_build_spatial_term_log_kappa_aniso_derivativeinfos(
data,
resolvedspec,
design,
term_idx,
aniso_gid,
)? {
aniso_gid += 1;
out.extend(entries);
continue;
} else {
return Ok(None);
}
}
let Some(info) =
try_build_spatial_term_log_kappa_derivativeinfo(data, resolvedspec, design, term_idx)?
else {
return Ok(None);
};
out.push(info);
}
Ok(Some(out))
}
fn try_build_spatial_term_log_kappa_aniso_derivativeinfos(
data: ArrayView2<'_, f64>,
resolvedspec: &TermCollectionSpec,
design: &TermCollectionDesign,
term_idx: usize,
aniso_group_id: usize,
) -> Result<Option<Vec<SpatialPsiDerivative>>, EstimationError> {
let Some(smooth_term) = design.smooth.terms.get(term_idx) else {
return Ok(None);
};
let Some(termspec) = resolvedspec.smooth_terms.get(term_idx) else {
return Ok(None);
};
let mut aniso_result = match &termspec.basis {
SmoothBasisSpec::Sphere { .. } => return Ok(None),
SmoothBasisSpec::Matern {
feature_cols,
spec,
input_scale,
} => {
let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
let mut spec_operator = spec.clone();
if let Some(scale) = input_scale {
scale.standardize(&mut x);
let length_scale = spec.length_scale.resolved().ok_or_else(|| {
EstimationError::InvalidInput(
"anisotropic Matérn Auto length_scale reached derivative construction \
unresolved"
.to_string(),
)
})?;
spec_operator
.length_scale
.set_resolved(scale.to_standardized_units(length_scale));
}
spec_operator.double_penalty = false;
build_matern_basis_log_kappa_aniso_derivatives(x.view(), &spec_operator)
.map_err(EstimationError::from)?
}
SmoothBasisSpec::MeasureJet {
feature_cols,
spec,
input_scale,
} => {
let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
if let Some(scale) = input_scale {
scale.standardize(&mut x);
}
build_measure_jet_basis_psi_derivatives(x.view(), spec)
.map_err(EstimationError::from)?
}
_ => return Ok(None),
};
let d = if let Some(ref op) = aniso_result.implicit_operator {
op.n_axes()
} else if !aniso_result.design_first.is_empty() {
aniso_result.design_first.len()
} else {
0
};
if d == 0 {
return Ok(None);
}
let Some(penalty_range) = design
.smooth_term_penalty_range(term_idx)
.map_err(EstimationError::InvalidInput)?
else {
return Ok(None);
};
let penalty_start = penalty_range.start;
let p_total = design.design.ncols();
let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
let global_range = (smooth_start + smooth_term.coeff_range.start)
..(smooth_start + smooth_term.coeff_range.end);
let num_penalties = aniso_result.penalties_first[0].len();
let penalty_indices: Vec<usize> = (0..num_penalties).map(|j| penalty_start + j).collect();
let penalties_cross_provider = aniso_result.penalties_cross_provider.clone();
let use_implicit_design = aniso_result.design_first.is_empty();
let implicit_op_arc = aniso_result
.implicit_operator
.as_ref()
.map(|op| std::sync::Arc::new(op.clone()));
let mut entries = Vec::with_capacity(d);
for a in 0..d {
let (x_psi_local, x_psi_psi_local) = if use_implicit_design {
(Array2::<f64>::zeros((0, 0)), Array2::<f64>::zeros((0, 0)))
} else {
let x_first = std::mem::take(&mut aniso_result.design_first[a]);
let x_second = std::mem::take(&mut aniso_result.design_second_diag[a]);
if x_first.ncols() != smooth_term.coeff_range.len() {
return Ok(None);
}
(x_first, x_second)
};
let s_psi_components = std::mem::take(&mut aniso_result.penalties_first[a]);
let s_psi_psi_components = std::mem::take(&mut aniso_result.penalties_second_diag[a]);
let cross_designs = if implicit_op_arc.is_some() {
let mut cd = Vec::with_capacity(d - 1);
for b in 0..d {
if b == a {
continue;
}
cd.push((b, Array2::<f64>::zeros((0, 0))));
}
cd
} else if !aniso_result.design_second_cross.is_empty() {
let mut cd = Vec::new();
for (cross_idx, &(pa, pb)) in aniso_result.design_second_cross_pairs.iter().enumerate()
{
if pa == a {
cd.push((pb, aniso_result.design_second_cross[cross_idx].clone()));
} else if pb == a {
cd.push((pa, aniso_result.design_second_cross[cross_idx].clone()));
}
}
cd
} else {
Vec::new()
};
let cross_penalty_provider = if d > 1 {
let penalties_cross_provider = penalties_cross_provider.clone();
Some(std::sync::Arc::new(
move |b_axis: usize| -> Result<Vec<Array2<f64>>, EstimationError> {
if b_axis == a {
return Ok(Vec::new());
}
let (axis_lo, axis_hi) = if a < b_axis { (a, b_axis) } else { (b_axis, a) };
if let Some(provider) = penalties_cross_provider.as_ref() {
provider
.evaluate(axis_lo, axis_hi)
.map_err(EstimationError::from)
} else {
Ok(Vec::new())
}
},
)
as std::sync::Arc<
dyn Fn(usize) -> Result<Vec<Array2<f64>>, EstimationError>
+ Send
+ Sync
+ 'static,
>)
} else {
None
};
entries.push(SpatialPsiDerivative {
penalty_index: penalty_indices[0],
penalty_indices: penalty_indices.clone(),
global_range: global_range.clone(),
total_p: p_total,
x_psi_local,
s_psi_components_local: s_psi_components,
x_psi_psi_local,
s_psi_psi_components_local: s_psi_psi_components,
aniso_group_id: Some(aniso_group_id),
aniso_cross_designs: if cross_designs.is_empty() {
None
} else {
Some(cross_designs)
},
aniso_cross_penalty_provider: cross_penalty_provider,
implicit_operator: implicit_op_arc.clone(),
implicit_axis: a,
});
}
Ok(Some(entries))
}
#[cfg(test)]
mod glm_eta_observation_fd_tests {
use super::*;
use ndarray::array;
fn one_obs_weight(
spec: &LikelihoodSpec,
y: f64,
weight: f64,
eta: f64,
) -> StandardFamilyObservationState {
let yv = Array1::from_vec(vec![y]);
let wv = Array1::from_vec(vec![weight]);
let ev = Array1::from_vec(vec![eta]);
evaluate_standard_familyobservations(spec.clone(), None, None, None, &yv, &wv, &ev)
.expect("standard family observation state assembles")
}
fn one_obs(spec: &LikelihoodSpec, y: f64, eta: f64) -> StandardFamilyObservationState {
one_obs_weight(spec, y, 1.0, eta)
}
fn one_obs_resolved(
likelihood: &gam_spec::GlmLikelihoodSpec,
y: f64,
weight: f64,
eta: f64,
) -> StandardFamilyObservationState {
evaluate_resolved_standard_family_observations(
likelihood,
None,
None,
None,
&array![y],
&array![weight],
&array![eta],
)
.expect("resolved standard family observation state assembles")
}
#[test]
fn bounded_gamma_and_tweedie_use_the_resolved_likelihood_scale() {
let gamma_unit = gam_spec::GlmLikelihoodSpec {
spec: LikelihoodSpec::gamma_log(),
scale: gam_spec::LikelihoodScaleMetadata::FixedGammaShape { shape: 1.0 },
};
let gamma_scaled = gam_spec::GlmLikelihoodSpec {
spec: LikelihoodSpec::gamma_log(),
scale: gam_spec::LikelihoodScaleMetadata::FixedGammaShape { shape: 8.0 },
};
let unit = one_obs_resolved(&gamma_unit, 2.3, 0.7, 0.2);
let scaled = one_obs_resolved(&gamma_scaled, 2.3, 0.7, 0.2);
for (label, actual, base) in [
("Gamma score", scaled.score[0], unit.score[0]),
(
"Gamma Fisher weight",
scaled.fisherweight[0],
unit.fisherweight[0],
),
(
"Gamma observed Hessian",
scaled.neghessian_eta[0],
unit.neghessian_eta[0],
),
(
"Gamma Hessian derivative",
scaled.neghessian_eta_derivative[0],
unit.neghessian_eta_derivative[0],
),
(
"Gamma log likelihood",
scaled.log_likelihood,
unit.log_likelihood,
),
] {
let expected = 8.0 * base;
assert!(
(actual - expected).abs() <= 32.0 * f64::EPSILON * expected.abs().max(1.0),
"{label} scale mismatch: actual={actual}, expected={expected}"
);
}
let tweedie_unit = gam_spec::GlmLikelihoodSpec {
spec: LikelihoodSpec::tweedie_log(1.5),
scale: gam_spec::LikelihoodScaleMetadata::FixedDispersion { phi: 1.0 },
};
let tweedie_scaled = gam_spec::GlmLikelihoodSpec {
spec: LikelihoodSpec::tweedie_log(1.5),
scale: gam_spec::LikelihoodScaleMetadata::FixedDispersion { phi: 0.25 },
};
let unit = one_obs_resolved(&tweedie_unit, 1.7, 0.8, -0.1);
let scaled = one_obs_resolved(&tweedie_scaled, 1.7, 0.8, -0.1);
for (actual, base) in [
(scaled.score[0], unit.score[0]),
(scaled.fisherweight[0], unit.fisherweight[0]),
(scaled.neghessian_eta[0], unit.neghessian_eta[0]),
(
scaled.neghessian_eta_derivative[0],
unit.neghessian_eta_derivative[0],
),
(scaled.log_likelihood, unit.log_likelihood),
] {
let expected = 4.0 * base;
assert!((actual - expected).abs() <= 32.0 * f64::EPSILON * expected.abs().max(1.0));
}
}
#[test]
fn bounded_zero_rows_are_dormant_and_weight_preflight_is_atomic() {
let likelihood = gam_spec::GlmLikelihoodSpec::canonical(LikelihoodSpec::poisson_log());
let dormant = evaluate_resolved_standard_family_observations(
&likelihood,
None,
None,
None,
&array![f64::NAN, 2.0],
&array![0.0, 1.0],
&array![f64::NAN, 0.2],
)
.expect("zero-weight response and predictor are dormant");
assert_eq!(dormant.score[0], 0.0);
assert_eq!(dormant.fisherweight[0], 0.0);
let error = evaluate_resolved_standard_family_observations(
&likelihood,
None,
None,
None,
&array![f64::NAN, 2.0],
&array![1.0, f64::NAN],
&array![f64::NAN, 0.2],
)
.expect_err("later invalid weight must refuse before row evaluation");
assert!(
error.to_string().contains("row 2 has invalid prior weight"),
"unexpected atomic preflight error: {error}"
);
}
fn check_fd(label: &str, spec: &LikelihoodSpec, y: f64, eta: f64) {
let h = 1e-5;
let s0 = one_obs(spec, y, eta);
let sp = one_obs(spec, y, eta + h);
let sm = one_obs(spec, y, eta - h);
let score_fd = (sp.log_likelihood - sm.log_likelihood) / (2.0 * h);
let score = s0.score[0];
assert!(
(score - score_fd).abs() <= 1e-4 * (1.0 + score.abs()),
"{label}: score {score} vs FD {score_fd}"
);
let neghess_fd = -(sp.score[0] - sm.score[0]) / (2.0 * h);
let neghess = s0.neghessian_eta[0];
assert!(
(neghess - neghess_fd).abs() <= 1e-3 * (1.0 + neghess.abs()),
"{label}: neghessian_eta {neghess} vs FD {neghess_fd}"
);
let nhd_fd = (sp.neghessian_eta[0] - sm.neghessian_eta[0]) / (2.0 * h);
let nhd = s0.neghessian_eta_derivative[0];
assert!(
(nhd - nhd_fd).abs() <= 1e-2 * (1.0 + nhd.abs()),
"{label}: neghessian_eta_derivative {nhd} vs FD {nhd_fd}"
);
}
#[test]
fn poisson_gamma_nb_tweedie_arms_match_finite_differences_1615_1616() {
let log = InverseLink::Standard(StandardLink::Log);
let poisson = LikelihoodSpec {
response: ResponseFamily::Poisson,
link: log.clone(),
};
check_fd("poisson y=3", &poisson, 3.0, 0.4);
check_fd("poisson y=0", &poisson, 0.0, -0.2);
let gamma = LikelihoodSpec {
response: ResponseFamily::Gamma,
link: log.clone(),
};
check_fd("gamma y=2.5", &gamma, 2.5, 0.3);
check_fd("gamma y=0.7", &gamma, 0.7, -0.1);
let nb = LikelihoodSpec {
response: ResponseFamily::NegativeBinomial {
theta: 1.5,
theta_fixed: true,
},
link: log.clone(),
};
check_fd("negbin y=4", &nb, 4.0, 0.5);
check_fd("negbin y=0", &nb, 0.0, -0.3);
let tweedie = LikelihoodSpec {
response: ResponseFamily::Tweedie { p: 1.5 },
link: log.clone(),
};
check_fd("tweedie y=2", &tweedie, 2.0, 0.25);
check_fd("tweedie y=0.5", &tweedie, 0.5, -0.15);
}
#[test]
fn binomial_natural_coordinate_towers_match_finite_differences() {
for (label, family, eta) in [
("logit", LikelihoodSpec::binomial_logit(), 0.7),
("probit", LikelihoodSpec::binomial_probit(), -1.1),
("cloglog", LikelihoodSpec::binomial_cloglog(), 0.4),
(
"loglog",
LikelihoodSpec::try_new(
ResponseFamily::Binomial,
InverseLink::Standard(StandardLink::LogLog),
)
.unwrap(),
-0.35,
),
(
"cauchit",
LikelihoodSpec::try_new(
ResponseFamily::Binomial,
InverseLink::Standard(StandardLink::Cauchit),
)
.unwrap(),
1.25,
),
] {
check_fd(label, &family, 0.37, eta);
}
}
#[test]
fn logit_observation_geometry_carries_the_prior_weight_everywhere() {
let eta = 1.75;
let y = 0.3;
let weight = 7.25;
let state = one_obs_weight(&LikelihoodSpec::binomial_logit(), y, weight, eta);
let jet = logit_inverse_link_jet5(eta);
for (got, expected) in [
(state.fisherweight[0], weight * jet.d1),
(state.neghessian_eta[0], weight * jet.d1),
(state.neghessian_eta_derivative[0], weight * jet.d2),
(state.score[0], weight * (y - jet.mu)),
] {
assert!((got - expected).abs() <= 4.0 * f64::EPSILON * (1.0 + expected.abs()));
}
}
#[test]
fn tiny_positive_and_zero_weights_are_not_projected() {
let tiny = 1e-200;
let logit = one_obs_weight(&LikelihoodSpec::binomial_logit(), 0.4, tiny, 0.0);
assert!((logit.fisherweight[0] / tiny - 0.25).abs() <= 2.0 * f64::EPSILON);
assert!(logit.fisherweight[0] < 1e-190);
let zero = one_obs_weight(&LikelihoodSpec::gaussian_identity(), 3.0, 0.0, -2.0);
assert_eq!(zero.score[0], 0.0);
assert_eq!(zero.fisherweight[0], 0.0);
assert_eq!(zero.neghessian_eta[0], 0.0);
assert_eq!(zero.neghessian_eta_derivative[0], 0.0);
assert_eq!(zero.log_likelihood, 0.0);
assert_eq!(exact_standard_working_response(&zero).unwrap()[0], -2.0);
}
#[test]
fn log_link_tails_balance_tiny_weights_before_certification() {
let poisson = one_obs_weight(&LikelihoodSpec::poisson_log(), 0.0, 1e-300, 700.0);
assert!(poisson.fisherweight[0].is_finite() && poisson.fisherweight[0] > 1.0);
assert!(poisson.score[0].is_finite());
assert!(poisson.log_likelihood.is_finite());
let gamma = one_obs_weight(&LikelihoodSpec::gamma_log(), 1.0, 1e-300, -700.0);
assert!(gamma.neghessian_eta[0].is_finite() && gamma.neghessian_eta[0] > 1.0);
assert!(gamma.score[0].is_finite());
assert!(gamma.log_likelihood.is_finite());
}
#[test]
fn invalid_weights_and_nonfinite_inputs_are_refused_in_row_order() {
let family = LikelihoodSpec::gaussian_identity();
let y = array![1.0, 2.0];
let eta = array![0.0, 0.0];
for weights in [array![-1.0, 1.0], array![f64::NAN, 1.0]] {
let err = evaluate_standard_familyobservations(
family.clone(),
None,
None,
None,
&y,
&weights,
&eta,
)
.expect_err("invalid prior weight must be refused");
assert!(err.to_string().contains("row 0"), "{err}");
}
let err = evaluate_standard_familyobservations(
family,
None,
None,
None,
&array![f64::NAN],
&array![0.0],
&array![0.0],
)
.expect_err("a non-finite response may not hide behind zero weight");
assert!(err.to_string().contains("row 0"), "{err}");
}
#[test]
fn unrepresentable_cloglog_curvature_is_refused_without_a_floor() {
let err = evaluate_standard_familyobservations(
LikelihoodSpec::binomial_cloglog(),
None,
None,
None,
&array![1.0],
&array![1.0],
&array![18.0],
)
.expect_err("mathematically sub-f64 Fisher information must be refused");
assert!(err.to_string().contains("Fisher weight"), "{err}");
}
#[test]
fn bounded_covariance_requires_a_certified_strict_spd_precision() {
let covariance = certified_bounded_posterior_covariance(
&array![[4.0, 1.0], [1.0, 3.0]],
"bounded covariance regression",
)
.expect("strict SPD precision");
assert!((covariance[[0, 0]] - 3.0 / 11.0).abs() < 1e-14);
assert!((covariance[[0, 1]] + 1.0 / 11.0).abs() < 1e-14);
assert!((covariance[[1, 1]] - 4.0 / 11.0).abs() < 1e-14);
for invalid in [
array![[1.0, 1.0], [1.0, 1.0]],
array![[1.0, 2.0], [2.0, 1.0]],
] {
assert!(
certified_bounded_posterior_covariance(
&invalid,
"invalid bounded covariance regression"
)
.is_err(),
"singular/indefinite precision must not become a pseudo-covariance"
);
}
}
}