use crate::math::Fix128;
#[must_use]
pub fn residual_norm_inf(residual: &[Fix128]) -> Fix128 {
let mut worst = Fix128::ZERO;
for &component in residual {
let magnitude = component.abs();
if magnitude > worst {
worst = magnitude;
}
}
worst
}
pub const L2_TERM_FLOOR: Fix128 = Fix128::from_raw(0, 1 << 32);
pub fn residual_norm_l2_checked(residual: &[Fix128]) -> Result<Fix128, CoupledIterationError> {
let mut sum = Fix128::ZERO;
for (index, &component) in residual.iter().enumerate() {
let magnitude = component.abs();
if !magnitude.is_negative() && magnitude < L2_TERM_FLOOR {
continue;
}
let Some(square) = magnitude.checked_mul(magnitude) else {
let at = u32::try_from(index).unwrap_or(u32::MAX);
return Err(CoupledIterationError::ArithmeticWrapped { sweeps: at });
};
let next = sum + square;
if next < sum {
let at = u32::try_from(index).unwrap_or(u32::MAX);
return Err(CoupledIterationError::ArithmeticWrapped { sweeps: at });
}
sum = next;
}
Ok(sum.sqrt())
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct EquilibrationScale {
exponent: u32,
}
impl EquilibrationScale {
pub const MAX_EXPONENT: u32 = 62;
pub const IDENTITY: Self = Self { exponent: 0 };
pub fn covering(magnitude: Fix128) -> Result<Self, ConfigFault> {
if magnitude <= Fix128::ZERO {
return Ok(Self::IDENTITY);
}
let mut exponent = 0u32;
let mut bound = Fix128::ONE;
while bound < magnitude {
exponent += 1;
if exponent > Self::MAX_EXPONENT {
return Err(ConfigFault::ScaleOutOfRange);
}
bound = Fix128::from_int(1i64 << exponent);
}
Ok(Self { exponent })
}
#[must_use]
pub const fn exponent(self) -> u32 {
self.exponent
}
#[must_use]
pub fn factor(self) -> Fix128 {
Fix128::from_int(1i64 << self.exponent)
}
#[must_use]
pub const fn scale_down(self, value: Fix128) -> Fix128 {
value.shr_bits(self.exponent)
}
#[must_use]
pub fn scale_up(self, value: Fix128) -> Fix128 {
value * self.factor()
}
#[must_use]
pub fn round_trip_bound(self) -> Fix128 {
if self.exponent == 0 {
return Fix128::ZERO;
}
Fix128::from_raw(0, (1u64 << self.exponent) - 1)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum ConfigFault {
StagnationFractionOutOfRange,
ZeroSweepBudget,
NonPositiveTolerance,
DivergenceRatioNotAboveOne,
InsufficientRatioSamples,
ScaleOutOfRange,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum CoupledIterationError {
Diverging {
observed_ratio: Fix128,
sweeps: u32,
},
ArithmeticWrapped {
sweeps: u32,
},
NotConverged {
sweeps: u32,
residual: Fix128,
},
Stagnated {
sweeps: u32,
best_residual: Fix128,
first_residual: Fix128,
sweeps_since_improvement: u32,
},
InvalidConfig {
fault: ConfigFault,
},
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SubIterationConfig {
pub max_sweeps: u32,
pub relative_tolerance: Fix128,
pub stagnation_fraction: Fix128,
pub divergence_ratio: Fix128,
pub ratio_samples: u32,
}
impl SubIterationConfig {
pub fn new(
max_sweeps: u32,
relative_tolerance: Fix128,
stagnation_fraction: Fix128,
divergence_ratio: Fix128,
ratio_samples: u32,
) -> Result<Self, ConfigFault> {
let config = Self {
max_sweeps,
relative_tolerance,
stagnation_fraction,
divergence_ratio,
ratio_samples,
};
config.validate()?;
Ok(config)
}
pub fn validate(&self) -> Result<(), ConfigFault> {
if self.max_sweeps == 0 {
return Err(ConfigFault::ZeroSweepBudget);
}
if self.relative_tolerance <= Fix128::ZERO {
return Err(ConfigFault::NonPositiveTolerance);
}
if self.stagnation_fraction <= Fix128::ZERO || self.stagnation_fraction >= Fix128::ONE {
return Err(ConfigFault::StagnationFractionOutOfRange);
}
if self.divergence_ratio <= Fix128::ONE {
return Err(ConfigFault::DivergenceRatioNotAboveOne);
}
if self.ratio_samples < 2 {
return Err(ConfigFault::InsufficientRatioSamples);
}
Ok(())
}
}
impl Default for SubIterationConfig {
fn default() -> Self {
Self {
max_sweeps: 64,
relative_tolerance: Fix128::from_ratio(1, 1024),
stagnation_fraction: Fix128::from_ratio(1, 2),
divergence_ratio: Fix128::ONE + Fix128::from_ratio(1, 16),
ratio_samples: 4,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MonitorVerdict {
Continue,
Converged,
Stop(CoupledIterationError),
}
#[derive(Clone, Copy, Debug)]
pub struct ContractionMonitor {
config: SubIterationConfig,
sweeps: u32,
first: Fix128,
previous: Fix128,
best: Fix128,
best_sweep: u32,
growth_streak: u32,
observed_ratio: Option<Fix128>,
}
impl ContractionMonitor {
pub fn new(config: SubIterationConfig) -> Result<Self, CoupledIterationError> {
config
.validate()
.map_err(|fault| CoupledIterationError::InvalidConfig { fault })?;
Ok(Self {
config,
sweeps: 0,
first: Fix128::ZERO,
previous: Fix128::ZERO,
best: Fix128::ZERO,
best_sweep: 0,
growth_streak: 0,
observed_ratio: None,
})
}
#[must_use]
pub const fn sweeps(&self) -> u32 {
self.sweeps
}
#[must_use]
pub const fn observed_ratio(&self) -> Option<Fix128> {
self.observed_ratio
}
#[must_use]
pub const fn best_residual(&self) -> Fix128 {
self.best
}
pub fn observe(&mut self, residual: Fix128) -> MonitorVerdict {
let previous = self.previous;
let previous_streak = self.growth_streak;
self.sweeps += 1;
self.previous = residual;
if self.sweeps == 1 {
self.first = residual;
self.best = residual;
self.best_sweep = 1;
return if residual.is_zero() {
MonitorVerdict::Converged
} else {
MonitorVerdict::Continue
};
}
if residual > previous {
self.growth_streak = previous_streak + 1;
} else {
self.growth_streak = 0;
if previous_streak >= 2 {
return MonitorVerdict::Stop(CoupledIterationError::ArithmeticWrapped {
sweeps: self.sweeps,
});
}
}
if self.observed_ratio.is_none() && self.sweeps >= self.config.ratio_samples {
let ratio = if previous.is_zero() {
Fix128::ZERO
} else {
residual / previous
};
self.observed_ratio = Some(ratio);
if ratio > self.config.divergence_ratio {
return MonitorVerdict::Stop(CoupledIterationError::Diverging {
observed_ratio: ratio,
sweeps: self.sweeps,
});
}
}
if residual < self.best {
self.best = residual;
self.best_sweep = self.sweeps;
}
if residual <= self.first * self.config.relative_tolerance {
return MonitorVerdict::Converged;
}
if self.growth_streak == 0 && self.sweeps > self.config.ratio_samples {
let denominator = Fix128::ONE - self.config.stagnation_fraction;
let stop_at = Fix128::from_int(i64::from(self.best_sweep)) / denominator;
if Fix128::from_int(i64::from(self.sweeps)) > stop_at {
return MonitorVerdict::Stop(CoupledIterationError::Stagnated {
sweeps: self.sweeps,
best_residual: self.best,
first_residual: self.first,
sweeps_since_improvement: self.sweeps - self.best_sweep,
});
}
}
if self.sweeps >= self.config.max_sweeps {
return MonitorVerdict::Stop(CoupledIterationError::NotConverged {
sweeps: self.sweeps,
residual,
});
}
MonitorVerdict::Continue
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SubIterationReport {
pub sweeps: u32,
pub residual: Fix128,
pub observed_ratio: Option<Fix128>,
}
pub fn run_sub_iteration<F>(
config: SubIterationConfig,
mut sweep: F,
) -> Result<SubIterationReport, CoupledIterationError>
where
F: FnMut(u32) -> Fix128,
{
let mut monitor = ContractionMonitor::new(config)?;
loop {
let index = monitor.sweeps();
let residual = sweep(index);
match monitor.observe(residual) {
MonitorVerdict::Continue => {}
MonitorVerdict::Converged => {
return Ok(SubIterationReport {
sweeps: monitor.sweeps(),
residual,
observed_ratio: monitor.observed_ratio(),
});
}
MonitorVerdict::Stop(error) => return Err(error),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn added_mass_sweep(a: Fix128, mass_ratio: Fix128) -> Fix128 {
-Fix128::ONE - mass_ratio * a
}
fn fixed_point(mass_ratio: Fix128) -> Fix128 {
-Fix128::ONE / (Fix128::ONE + mass_ratio)
}
fn raw(v: Fix128) -> i128 {
(i128::from(v.hi) << 64) | i128::from(v.lo)
}
const RATIO_ULP_BOUND: u128 = 256;
#[track_caller]
fn assert_ratio_matches_closed_form(observed: Option<Fix128>, closed: Fix128) {
let observed = observed.expect("the ratio window closed, so a ratio exists");
let deviation = raw(observed).abs_diff(raw(closed));
assert!(
deviation <= RATIO_ULP_BOUND,
"observed ratio {observed:?} deviates from the closed form {closed:?} \
by {deviation} ulp, above the measured bound {RATIO_ULP_BOUND}"
);
}
#[test]
fn inf_norm_picks_the_largest_magnitude() {
let r = [
Fix128::from_ratio(1, 4),
Fix128::from_ratio(-3, 4),
Fix128::from_ratio(1, 2),
];
assert_eq!(residual_norm_inf(&r), Fix128::from_ratio(3, 4));
assert_eq!(residual_norm_inf(&[]), Fix128::ZERO);
}
#[test]
fn l2_floor_is_the_square_root_of_the_truncation_floor() {
assert_eq!(L2_TERM_FLOOR * L2_TERM_FLOOR, Fix128::from_raw(0, 1));
let below = Fix128::from_raw(0, (1u64 << 32) - 1);
assert!(below < L2_TERM_FLOOR);
assert!((below * below).is_zero());
}
#[test]
fn l2_accepts_components_that_legitimately_square_to_zero() {
let tiny = Fix128::from_raw(0, 1);
assert!(tiny < L2_TERM_FLOOR);
assert_eq!(residual_norm_l2_checked(&[tiny]), Ok(Fix128::ZERO));
}
#[test]
fn l2_reports_a_wrap_instead_of_a_plausible_zero() {
let huge = Fix128::from_int(4_000_000_000);
assert!(huge.abs() >= L2_TERM_FLOOR);
assert!(
(huge * huge).is_negative(),
"this case exercises the sum path"
);
assert_eq!(
residual_norm_l2_checked(&[huge]),
Err(CoupledIterationError::ArithmeticWrapped { sweeps: 0 })
);
}
#[test]
fn l2_reports_a_term_whose_square_wrapped_to_exactly_zero() {
for exponent in [32u32, 33, 40, 62] {
let magnitude = Fix128::from_int(1i64 << exponent);
let square = magnitude * magnitude;
assert!(
square.is_zero() && !square.is_negative(),
"2^{exponent} squared should wrap to exactly zero, got {square:?}"
);
assert!(magnitude >= L2_TERM_FLOOR);
assert_eq!(
residual_norm_l2_checked(&[magnitude]),
Err(CoupledIterationError::ArithmeticWrapped { sweeps: 0 }),
"a silently dropped term at 2^{exponent} must be refused"
);
}
}
#[test]
fn an_equilibration_factor_is_the_next_power_of_two_at_or_above_the_magnitude() {
for (magnitude, expected_exponent) in [
(Fix128::from_ratio(1, 4), 0u32),
(Fix128::ONE, 0),
(Fix128::from_int(2), 1),
(Fix128::from_int(3), 2),
(Fix128::from_int(4), 2),
(Fix128::from_int(5), 3),
(Fix128::from_int(1024), 10),
] {
let scale = EquilibrationScale::covering(magnitude).expect("within range");
assert_eq!(
scale.exponent(),
expected_exponent,
"magnitude {magnitude:?} should need 2^{expected_exponent}"
);
assert!(
scale.factor() >= magnitude,
"the factor must cover the magnitude"
);
}
assert_eq!(
EquilibrationScale::covering(Fix128::ZERO),
Ok(EquilibrationScale::IDENTITY)
);
assert_eq!(
EquilibrationScale::covering(Fix128::from_int(-8)),
Ok(EquilibrationScale::IDENTITY)
);
}
#[test]
fn equilibration_refuses_a_factor_it_cannot_represent() {
assert_eq!(
EquilibrationScale::covering(Fix128::from_int(1i64 << 62))
.expect("2^62 is the boundary and must be accepted")
.exponent(),
EquilibrationScale::MAX_EXPONENT
);
assert_eq!(
EquilibrationScale::covering(Fix128::from_int(i64::MAX)),
Err(ConfigFault::ScaleOutOfRange)
);
}
#[test]
fn scaling_up_then_down_is_exact() {
for exponent in [0u32, 1, 8, 20] {
let scale = EquilibrationScale::covering(Fix128::from_int(1i64 << exponent))
.expect("within range");
for raw_value in [1i64, 3, 1_000, -7] {
let value = Fix128::from_ratio(raw_value, 64);
assert_eq!(
scale.scale_down(scale.scale_up(value)),
value,
"up-then-down must be exact at 2^{exponent} for {value:?}"
);
}
}
}
#[test]
fn scaling_down_then_up_loses_only_the_shifted_bits() {
for exponent in [1u32, 4, 12, 30] {
let scale = EquilibrationScale::covering(Fix128::from_int(1i64 << exponent))
.expect("within range");
let bound = scale.round_trip_bound();
assert!(bound > Fix128::ZERO, "a real shift must have a real bound");
let value = Fix128::from_raw(3, u64::MAX);
let round_tripped = scale.scale_up(scale.scale_down(value));
let loss = (value - round_tripped).abs();
assert!(
loss <= bound,
"2^{exponent}: round-trip loss {loss:?} exceeded the stated bound {bound:?}"
);
assert!(
loss > Fix128::ZERO,
"2^{exponent}: this value has low bits set, so some loss must show; \
if none does, the bound is describing something that never happens"
);
}
let identity = EquilibrationScale::IDENTITY;
assert_eq!(identity.round_trip_bound(), Fix128::ZERO);
let value = Fix128::from_raw(3, u64::MAX);
assert_eq!(identity.scale_up(identity.scale_down(value)), value);
}
#[test]
fn equilibration_lifts_a_block_above_the_product_floor() {
let elastic = Fix128::from_int(3_500_000);
let scale = EquilibrationScale::covering(elastic).expect("within range");
let thermal_raw = Fix128::from_ratio(13, 100_000);
let equilibrated_elastic = scale.scale_down(elastic);
assert!(
equilibrated_elastic <= Fix128::ONE && equilibrated_elastic > Fix128::from_ratio(1, 2),
"the dominant block should land in (1/2, 1], got {equilibrated_elastic:?}"
);
assert!(
equilibrated_elastic >= L2_TERM_FLOOR,
"the equilibrated block must stay above the product floor"
);
assert!(thermal_raw < elastic);
}
#[test]
fn config_rejects_a_stagnation_fraction_that_can_never_fire() {
for fraction in [Fix128::ONE, Fix128::ZERO, Fix128::from_int(2)] {
let config = SubIterationConfig {
stagnation_fraction: fraction,
..SubIterationConfig::default()
};
assert_eq!(
config.validate(),
Err(ConfigFault::StagnationFractionOutOfRange)
);
}
}
#[test]
fn config_rejects_an_unreachable_zero_tolerance() {
let config = SubIterationConfig {
relative_tolerance: Fix128::ZERO,
..SubIterationConfig::default()
};
assert_eq!(config.validate(), Err(ConfigFault::NonPositiveTolerance));
}
#[test]
fn config_rejects_a_divergence_threshold_at_the_marginal_ratio() {
let config = SubIterationConfig {
divergence_ratio: Fix128::ONE,
..SubIterationConfig::default()
};
assert_eq!(
config.validate(),
Err(ConfigFault::DivergenceRatioNotAboveOne)
);
}
#[test]
fn config_rejects_an_empty_budget_and_too_few_ratio_samples() {
let config = SubIterationConfig {
max_sweeps: 0,
..SubIterationConfig::default()
};
assert_eq!(config.validate(), Err(ConfigFault::ZeroSweepBudget));
let config = SubIterationConfig {
ratio_samples: 1,
..SubIterationConfig::default()
};
assert_eq!(
config.validate(),
Err(ConfigFault::InsufficientRatioSamples)
);
}
#[test]
fn new_refuses_what_validate_refuses() {
assert_eq!(
SubIterationConfig::new(
0,
Fix128::from_ratio(1, 1024),
Fix128::from_ratio(1, 2),
Fix128::from_int(2),
4,
),
Err(ConfigFault::ZeroSweepBudget)
);
assert!(SubIterationConfig::new(
64,
Fix128::from_ratio(1, 1024),
Fix128::from_ratio(1, 2),
Fix128::from_int(2),
4,
)
.is_ok());
}
#[test]
fn default_config_is_pinned_and_valid() {
let config = SubIterationConfig::default();
assert_eq!(config.validate(), Ok(()));
assert_eq!(config.max_sweeps, 64);
assert_eq!(config.relative_tolerance, Fix128::from_ratio(1, 1024));
assert_eq!(config.stagnation_fraction, Fix128::from_ratio(1, 2));
assert_eq!(
config.divergence_ratio,
Fix128::ONE + Fix128::from_ratio(1, 16)
);
assert_eq!(config.ratio_samples, 4);
}
#[test]
fn a_contracting_splitting_converges_and_reports_the_mass_ratio() {
let mass_ratio = Fix128::from_ratio(1, 2);
let target = fixed_point(mass_ratio);
let mut a = Fix128::ZERO;
let report = run_sub_iteration(SubIterationConfig::default(), |_| {
a = added_mass_sweep(a, mass_ratio);
(a - target).abs()
})
.expect("m_f/m = 1/2 is a contraction");
assert_ratio_matches_closed_form(report.observed_ratio, mass_ratio);
}
#[test]
fn a_divergent_splitting_is_reported_as_divergent_not_as_a_wrap() {
let mass_ratio = Fix128::from_int(2);
let target = fixed_point(mass_ratio);
let mut a = Fix128::ZERO;
let error = run_sub_iteration(SubIterationConfig::default(), |_| {
a = added_mass_sweep(a, mass_ratio);
(a - target).abs()
})
.expect_err("m_f/m = 2 is not a contraction");
match error {
CoupledIterationError::Diverging { observed_ratio, .. } => {
assert_ratio_matches_closed_form(Some(observed_ratio), mass_ratio);
}
other => panic!("expected Diverging, got {other:?}"),
}
}
#[test]
fn the_wrap_guard_fires_on_growth_followed_by_a_decrease() {
let config = SubIterationConfig {
divergence_ratio: Fix128::from_int(1000),
..SubIterationConfig::default()
};
let mut monitor = ContractionMonitor::new(config).expect("valid config");
for value in [1i64, 2, 4, 8] {
assert_eq!(
monitor.observe(Fix128::from_int(value)),
MonitorVerdict::Continue
);
}
assert_eq!(
monitor.observe(Fix128::from_int(3)),
MonitorVerdict::Stop(CoupledIterationError::ArithmeticWrapped { sweeps: 5 })
);
}
#[test]
fn the_wrap_guard_fires_at_exactly_two_growths() {
let config = SubIterationConfig {
divergence_ratio: Fix128::from_int(1000),
..SubIterationConfig::default()
};
let mut monitor = ContractionMonitor::new(config).expect("valid config");
for value in [1i64, 2, 4] {
assert_eq!(
monitor.observe(Fix128::from_int(value)),
MonitorVerdict::Continue
);
}
assert_eq!(
monitor.observe(Fix128::from_int(3)),
MonitorVerdict::Stop(CoupledIterationError::ArithmeticWrapped { sweeps: 4 })
);
}
#[test]
fn the_wrap_guard_stays_silent_on_a_decreasing_sequence() {
let mut monitor =
ContractionMonitor::new(SubIterationConfig::default()).expect("valid config");
let mut residual = Fix128::ONE;
for _ in 0..8 {
let verdict = monitor.observe(residual);
assert!(
!matches!(
verdict,
MonitorVerdict::Stop(CoupledIterationError::ArithmeticWrapped { .. })
),
"wrap guard fired on a monotonically decreasing sequence"
);
residual = residual / Fix128::from_int(2);
}
}
#[test]
fn the_wrap_guard_tolerates_a_single_overshoot() {
let mut monitor =
ContractionMonitor::new(SubIterationConfig::default()).expect("valid config");
assert_eq!(monitor.observe(Fix128::ONE), MonitorVerdict::Continue);
assert_eq!(
monitor.observe(Fix128::from_ratio(3, 2)),
MonitorVerdict::Continue
);
assert!(!matches!(
monitor.observe(Fix128::from_ratio(1, 2)),
MonitorVerdict::Stop(CoupledIterationError::ArithmeticWrapped { .. })
));
}
#[test]
fn stagnation_reports_enough_to_tell_the_mechanisms_apart() {
let mut monitor =
ContractionMonitor::new(SubIterationConfig::default()).expect("valid config");
assert_eq!(monitor.observe(Fix128::ONE), MonitorVerdict::Continue);
let plateau = Fix128::from_ratio(9, 10);
let mut verdict = MonitorVerdict::Continue;
for _ in 0..8 {
verdict = monitor.observe(plateau);
if matches!(verdict, MonitorVerdict::Stop(_)) {
break;
}
}
match verdict {
MonitorVerdict::Stop(CoupledIterationError::Stagnated {
best_residual,
first_residual,
sweeps_since_improvement,
..
}) => {
assert_eq!(best_residual, plateau);
assert_eq!(first_residual, Fix128::ONE);
assert!(
sweeps_since_improvement > 0,
"a stagnation report must say how long it has been flat"
);
assert!(
best_residual > first_residual / Fix128::from_int(2),
"this scene barely improved, which is what marks it as a short \
window rather than a reached floor"
);
}
other => panic!("expected Stagnated, got {other:?}"),
}
}
#[test]
fn an_invalid_config_is_reported_before_any_sweep_runs() {
let config = SubIterationConfig {
max_sweeps: 0,
..SubIterationConfig::default()
};
let mut sweeps_run = 0u32;
let error = run_sub_iteration(config, |_| {
sweeps_run += 1;
Fix128::ONE
})
.expect_err("zero budget is invalid");
assert_eq!(
error,
CoupledIterationError::InvalidConfig {
fault: ConfigFault::ZeroSweepBudget
}
);
assert_eq!(sweeps_run, 0, "no sweep may run under an invalid config");
}
}