use ndarray::{Array1, ArrayView1};
use serde::{Deserialize, Serialize};
use std::ops::{Deref, DerefMut};
pub use gam_linalg::RidgePolicy;
pub use gam_spec::*;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum RidgeMatrixForm {
ScaledIdentity,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InvalidStabilization {
reason: String,
}
impl InvalidStabilization {
fn new(reason: impl Into<String>) -> Self {
Self {
reason: reason.into(),
}
}
pub fn reason(&self) -> &str {
&self.reason
}
}
impl std::fmt::Display for InvalidStabilization {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "invalid stabilization metadata: {}", self.reason)
}
}
impl std::error::Error for InvalidStabilization {}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
struct RidgePassportWire {
delta: f64,
matrix_form: RidgeMatrixForm,
policy: RidgePolicy,
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "RidgePassportWire", into = "RidgePassportWire")]
pub struct RidgePassport {
delta: f64,
matrix_form: RidgeMatrixForm,
policy: RidgePolicy,
}
impl RidgePassport {
pub fn scaled_identity(delta: f64, policy: RidgePolicy) -> Result<Self, InvalidStabilization> {
if !(delta.is_finite() && delta >= 0.0) {
return Err(InvalidStabilization::new(format!(
"ridge delta must be finite and non-negative, got {delta:?}"
)));
}
Ok(Self {
delta: if delta == 0.0 { 0.0 } else { delta },
matrix_form: RidgeMatrixForm::ScaledIdentity,
policy,
})
}
pub const fn zero(policy: RidgePolicy) -> Self {
Self {
delta: 0.0,
matrix_form: RidgeMatrixForm::ScaledIdentity,
policy,
}
}
#[inline]
pub const fn delta(self) -> f64 {
self.delta
}
#[inline]
pub const fn matrix_form(self) -> RidgeMatrixForm {
self.matrix_form
}
#[inline]
pub const fn policy(self) -> RidgePolicy {
self.policy
}
#[inline]
pub const fn penalty_logdet_ridge(self) -> f64 {
if self.policy.accounts_for_objective() {
self.delta
} else {
0.0
}
}
}
impl TryFrom<RidgePassportWire> for RidgePassport {
type Error = InvalidStabilization;
fn try_from(wire: RidgePassportWire) -> Result<Self, Self::Error> {
let mut passport = Self::scaled_identity(wire.delta, wire.policy)?;
passport.matrix_form = wire.matrix_form;
Ok(passport)
}
}
impl From<RidgePassport> for RidgePassportWire {
fn from(passport: RidgePassport) -> Self {
Self {
delta: passport.delta,
matrix_form: passport.matrix_form,
policy: passport.policy,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(try_from = "InertiaWire", into = "InertiaWire")]
pub struct Inertia {
positive: usize,
zero: usize,
negative: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
struct InertiaWire {
positive: usize,
zero: usize,
negative: usize,
}
impl Inertia {
pub fn new(
positive: usize,
zero: usize,
negative: usize,
) -> Result<Self, InvalidStabilization> {
let total = positive
.checked_add(zero)
.and_then(|value| value.checked_add(negative))
.ok_or_else(|| InvalidStabilization::new("inertia count sum overflows usize"))?;
if total == 0 {
return Err(InvalidStabilization::new(
"inertia must describe a non-empty matrix",
));
}
Ok(Self {
positive,
zero,
negative,
})
}
pub const fn positive(self) -> usize {
self.positive
}
pub const fn zero(self) -> usize {
self.zero
}
pub const fn negative(self) -> usize {
self.negative
}
pub fn total(self) -> usize {
self.positive + self.zero + self.negative
}
}
impl TryFrom<InertiaWire> for Inertia {
type Error = InvalidStabilization;
fn try_from(wire: InertiaWire) -> Result<Self, Self::Error> {
Self::new(wire.positive, wire.zero, wire.negative)
}
}
impl From<Inertia> for InertiaWire {
fn from(inertia: Inertia) -> Self {
Self {
positive: inertia.positive,
zero: inertia.zero,
negative: inertia.negative,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub enum StabilizationRule {
FixedConstant,
InertiaTarget { spd_floor: f64 },
Heuristic,
UserSpecified,
BackoffEscalation { attempts: usize },
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub enum StabilizationKind {
None,
SolverDampingOnly,
NumericalPerturbation,
ApproximationOnly,
ObjectiveStabilization,
ExplicitPrior,
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
struct StabilizationLedgerWire {
kind: StabilizationKind,
delta: f64,
matrix_form: RidgeMatrixForm,
chosen_by: StabilizationRule,
objective_policy: Option<RidgePolicy>,
backward_error_bound: Option<f64>,
inertia_before: Option<Inertia>,
inertia_after: Option<Inertia>,
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "StabilizationLedgerWire", into = "StabilizationLedgerWire")]
pub struct StabilizationLedger {
kind: StabilizationKind,
delta: f64,
matrix_form: RidgeMatrixForm,
chosen_by: StabilizationRule,
objective_policy: Option<RidgePolicy>,
backward_error_bound: Option<f64>,
inertia_before: Option<Inertia>,
inertia_after: Option<Inertia>,
}
impl StabilizationLedger {
pub const fn none() -> Self {
Self {
kind: StabilizationKind::None,
delta: 0.0,
matrix_form: RidgeMatrixForm::ScaledIdentity,
chosen_by: StabilizationRule::FixedConstant,
objective_policy: None,
backward_error_bound: None,
inertia_before: None,
inertia_after: None,
}
}
fn try_new(
kind: StabilizationKind,
delta: f64,
chosen_by: StabilizationRule,
backward_error_bound: Option<f64>,
) -> Result<Self, InvalidStabilization> {
if matches!(kind, StabilizationKind::None) {
return Err(InvalidStabilization::new(
"None stabilization must be constructed with StabilizationLedger::none",
));
}
if !(delta.is_finite() && delta >= 0.0) {
return Err(InvalidStabilization::new(format!(
"stabilization delta must be finite and non-negative, got {delta:?}"
)));
}
Self::validate_rule(chosen_by)?;
if let Some(bound) = backward_error_bound
&& !(bound.is_finite() && bound >= 0.0)
{
return Err(InvalidStabilization::new(format!(
"backward-error bound must be finite and non-negative, got {bound:?}"
)));
}
if !matches!(kind, StabilizationKind::NumericalPerturbation)
&& backward_error_bound.is_some()
{
return Err(InvalidStabilization::new(
"only a numerical perturbation may carry a backward-error bound",
));
}
Ok(Self {
kind,
delta: if delta == 0.0 { 0.0 } else { delta },
matrix_form: RidgeMatrixForm::ScaledIdentity,
chosen_by,
objective_policy: None,
backward_error_bound,
inertia_before: None,
inertia_after: None,
})
}
fn validate_rule(rule: StabilizationRule) -> Result<(), InvalidStabilization> {
match rule {
StabilizationRule::InertiaTarget { spd_floor }
if !(spd_floor.is_finite() && spd_floor > 0.0) =>
{
Err(InvalidStabilization::new(format!(
"inertia-target SPD floor must be finite and strictly positive, got {spd_floor:?}"
)))
}
StabilizationRule::BackoffEscalation { attempts } if attempts == 0 => Err(
InvalidStabilization::new("backoff escalation must record at least one attempt"),
),
_ => Ok(()),
}
}
pub fn with_inertia(
mut self,
before: Option<Inertia>,
after: Option<Inertia>,
) -> Result<Self, InvalidStabilization> {
if before.is_some() != after.is_some() {
return Err(InvalidStabilization::new(
"inertia diagnostics must record both the pre- and post-stabilization matrix",
));
}
if let (Some(before), Some(after)) = (before, after)
&& before.total() != after.total()
{
return Err(InvalidStabilization::new(format!(
"inertia dimensions disagree: before={}, after={}",
before.total(),
after.total()
)));
}
if matches!(self.chosen_by, StabilizationRule::InertiaTarget { .. }) {
let Some(after) = after else {
return Err(InvalidStabilization::new(
"inertia-target stabilization must record post-stabilization inertia",
));
};
if after.zero() != 0 || after.negative() != 0 {
return Err(InvalidStabilization::new(format!(
"inertia-target stabilization did not certify SPD curvature: zero={}, negative={}",
after.zero(),
after.negative()
)));
}
}
self.inertia_before = before;
self.inertia_after = after;
Ok(self)
}
pub const fn kind(self) -> StabilizationKind {
self.kind
}
pub const fn delta(self) -> f64 {
self.delta
}
pub const fn matrix_form(self) -> RidgeMatrixForm {
self.matrix_form
}
pub const fn chosen_by(self) -> StabilizationRule {
self.chosen_by
}
pub const fn objective_policy(self) -> Option<RidgePolicy> {
self.objective_policy
}
pub const fn backward_error_bound(self) -> Option<f64> {
self.backward_error_bound
}
pub const fn inertia_before(self) -> Option<Inertia> {
self.inertia_before
}
pub const fn inertia_after(self) -> Option<Inertia> {
self.inertia_after
}
}
impl TryFrom<StabilizationLedgerWire> for StabilizationLedger {
type Error = InvalidStabilization;
fn try_from(wire: StabilizationLedgerWire) -> Result<Self, Self::Error> {
if matches!(wire.kind, StabilizationKind::None) {
if wire.delta != 0.0
|| wire.chosen_by != StabilizationRule::FixedConstant
|| wire.objective_policy.is_some()
|| wire.backward_error_bound.is_some()
|| wire.inertia_before.is_some()
|| wire.inertia_after.is_some()
{
return Err(InvalidStabilization::new(
"None stabilization must have zero delta and no diagnostic payload",
));
}
return Ok(Self::none());
}
let mut ledger = Self::try_new(
wire.kind,
wire.delta,
wire.chosen_by,
wire.backward_error_bound,
)?;
if matches!(wire.kind, StabilizationKind::ExplicitPrior)
&& wire.chosen_by != StabilizationRule::UserSpecified
{
return Err(InvalidStabilization::new(
"an explicit prior must be recorded as user specified",
));
}
match (wire.kind, wire.objective_policy) {
(
StabilizationKind::ExplicitPrior | StabilizationKind::ObjectiveStabilization,
Some(policy),
) if policy.accounts_for_objective() => {
ledger.objective_policy = Some(policy);
}
(StabilizationKind::ExplicitPrior | StabilizationKind::ObjectiveStabilization, _) => {
return Err(InvalidStabilization::new(
"objective-accounted stabilization must preserve its ridge policy",
));
}
(_, Some(_)) => {
return Err(InvalidStabilization::new(
"only objective-accounted stabilization may carry objective ridge provenance",
));
}
(_, None) => {}
}
ledger.matrix_form = wire.matrix_form;
ledger.with_inertia(wire.inertia_before, wire.inertia_after)
}
}
impl From<StabilizationLedger> for StabilizationLedgerWire {
fn from(ledger: StabilizationLedger) -> Self {
Self {
kind: ledger.kind,
delta: ledger.delta,
matrix_form: ledger.matrix_form,
chosen_by: ledger.chosen_by,
objective_policy: ledger.objective_policy,
backward_error_bound: ledger.backward_error_bound,
inertia_before: ledger.inertia_before,
inertia_after: ledger.inertia_after,
}
}
}
macro_rules! array1_f64_newtype {
($name:ident) => {
#[repr(transparent)]
#[derive(Clone, Debug, PartialEq)]
pub struct $name(pub Array1<f64>);
impl $name {
#[inline]
pub fn new(values: Array1<f64>) -> Self {
Self(values)
}
#[inline]
pub fn zeros(len: usize) -> Self {
Self(Array1::zeros(len))
}
}
impl Deref for $name {
type Target = Array1<f64>;
#[inline]
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl DerefMut for $name {
#[inline]
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
impl AsRef<Array1<f64>> for $name {
#[inline]
fn as_ref(&self) -> &Array1<f64> {
&self.0
}
}
impl From<Array1<f64>> for $name {
#[inline]
fn from(values: Array1<f64>) -> Self {
Self(values)
}
}
impl From<$name> for Array1<f64> {
#[inline]
fn from(values: $name) -> Self {
values.0
}
}
};
}
array1_f64_newtype!(Coefficients);
array1_f64_newtype!(LinearPredictor);
#[repr(transparent)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct SmoothTermIdx(usize);
impl SmoothTermIdx {
#[inline]
pub const fn new(idx: usize) -> Self {
Self(idx)
}
#[inline]
pub const fn placeholder() -> Self {
Self(usize::MAX)
}
#[inline]
pub const fn get(self) -> usize {
self.0
}
}
impl std::fmt::Display for SmoothTermIdx {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[repr(transparent)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct PenaltyIdx(usize);
impl PenaltyIdx {
#[inline]
pub const fn new(idx: usize) -> Self {
Self(idx)
}
#[inline]
pub const fn get(self) -> usize {
self.0
}
}
impl std::fmt::Display for PenaltyIdx {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[repr(transparent)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct BasisIdx(usize);
impl BasisIdx {
#[inline]
pub const fn new(idx: usize) -> Self {
Self(idx)
}
#[inline]
pub const fn get(self) -> usize {
self.0
}
}
impl std::fmt::Display for BasisIdx {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[repr(transparent)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct ColIdx(usize);
impl ColIdx {
#[inline]
pub const fn new(idx: usize) -> Self {
Self(idx)
}
#[inline]
pub const fn get(self) -> usize {
self.0
}
}
impl std::fmt::Display for ColIdx {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[repr(transparent)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct RowIdx(usize);
impl RowIdx {
#[inline]
pub const fn new(idx: usize) -> Self {
Self(idx)
}
#[inline]
pub const fn get(self) -> usize {
self.0
}
}
impl std::fmt::Display for RowIdx {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[repr(transparent)]
#[derive(Clone, Copy, Debug)]
pub struct LogSmoothingParamsView<'a>(ArrayView1<'a, f64>);
impl<'a> LogSmoothingParamsView<'a> {
pub fn new(values: ArrayView1<'a, f64>) -> Result<Self, crate::IndexedLogStrengthDomainError> {
crate::validate_log_strengths(values.iter().copied())?;
Ok(Self(values))
}
pub fn exact_exp(&self) -> Array1<f64> {
self.0.mapv(f64::exp)
}
}
impl<'a> Deref for LogSmoothingParamsView<'a> {
type Target = ArrayView1<'a, f64>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
#[cfg(test)]
mod newtype_tests {
use super::*;
use ndarray::array;
#[test]
fn smooth_term_idx_ordering() {
let a = SmoothTermIdx::new(1);
let b = SmoothTermIdx::new(2);
assert!(a < b);
assert_eq!(a, SmoothTermIdx::new(1));
}
#[test]
fn coefficients_zeros_and_deref() {
let c = Coefficients::zeros(3);
assert_eq!(c.len(), 3);
assert!(c.iter().all(|&v| v == 0.0));
}
#[test]
fn coefficients_from_array1() {
let arr = array![1.0, 2.0, 3.0];
let c = Coefficients::from(arr.clone());
assert_eq!(*c, arr);
}
#[test]
fn log_smoothing_params_view_is_validated_and_exponentiates_exactly() {
let arr = array![crate::LOG_STRENGTH_MIN, 0.0, crate::LOG_STRENGTH_MAX];
let rho = LogSmoothingParamsView::new(arr.view()).expect("closed domain");
for (actual, expected) in rho.exact_exp().iter().zip(arr.iter()) {
assert_eq!(actual.to_bits(), expected.exp().to_bits());
}
let invalid = array![0.0, crate::LOG_STRENGTH_MAX + 1.0];
let error = LogSmoothingParamsView::new(invalid.view()).unwrap_err();
assert_eq!(error.coordinate, 1);
assert_eq!(error.value, crate::LOG_STRENGTH_MAX + 1.0);
}
#[test]
fn linear_predictor_zeros_and_deref() {
let lp = LinearPredictor::zeros(4);
assert_eq!(lp.len(), 4);
assert!(lp.iter().all(|&v| v == 0.0));
}
}
#[cfg(test)]
mod ridge_policy_tests {
use super::{RidgePassport, RidgePolicy, StabilizationLedger};
use serde_json::json;
#[test]
fn serde_cannot_bypass_passport_validation() {
let negative = json!({
"delta": -1.0,
"matrix_form": "ScaledIdentity",
"policy": "SolverOnly"
});
assert!(serde_json::from_value::<RidgePassport>(negative).is_err());
let passport = RidgePassport::scaled_identity(
2.5e-7,
RidgePolicy::exact_full_objective(),
)
.expect("valid ridge");
let roundtrip: RidgePassport =
serde_json::from_value(serde_json::to_value(passport).expect("serialize passport"))
.expect("deserialize validated passport");
assert_eq!(roundtrip, passport);
}
#[test]
fn serde_cannot_bypass_ledger_semantics() {
let invalid_none = json!({
"kind": "None",
"delta": 1.0,
"matrix_form": "ScaledIdentity",
"chosen_by": "FixedConstant",
"objective_policy": null,
"backward_error_bound": null,
"inertia_before": null,
"inertia_after": null
});
assert!(serde_json::from_value::<StabilizationLedger>(invalid_none).is_err());
let invalid_prior_rule = json!({
"kind": "ExplicitPrior",
"delta": 1.0,
"matrix_form": "ScaledIdentity",
"chosen_by": "Heuristic",
"objective_policy": "ExactFullObjective",
"backward_error_bound": null,
"inertia_before": null,
"inertia_after": null
});
assert!(serde_json::from_value::<StabilizationLedger>(invalid_prior_rule).is_err());
let bound_on_wrong_kind = json!({
"kind": "ApproximationOnly",
"delta": 1.0,
"matrix_form": "ScaledIdentity",
"chosen_by": "FixedConstant",
"objective_policy": null,
"backward_error_bound": 1.0e-10,
"inertia_before": null,
"inertia_after": null
});
assert!(serde_json::from_value::<StabilizationLedger>(bound_on_wrong_kind).is_err());
}
}