use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use crate::behavioral_fidelity::report::BehavioralFidelityReport;
use super::knob::{CalibrationKnob, KnobValue};
use super::objective::CalibrationObjective;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum RollbackPolicy {
#[default]
Revert,
Keep,
HalveDamping,
}
#[derive(Debug, Clone)]
pub struct CalibrationConfig {
pub max_iterations: usize,
pub seeds_per_iteration: usize,
pub patience: usize,
pub min_improvement: f64,
pub damping: f64,
pub rollback: RollbackPolicy,
}
impl Default for CalibrationConfig {
fn default() -> Self {
Self {
max_iterations: 20,
seeds_per_iteration: 3,
patience: 3,
min_improvement: 1.0,
damping: 0.5,
rollback: RollbackPolicy::default(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StepReport {
pub iter: usize,
pub loss_before_mean: f64,
pub loss_before_std: f64,
pub proposed_patch: Option<ProposedPatch>,
pub loss_after_mean: Option<f64>,
pub loss_after_std: Option<f64>,
pub knob_values: BTreeMap<String, KnobValue>,
pub outcome: StepOutcome,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum StepOutcome {
Improved,
AcceptedNoNoiseFloorBeat,
Reverted,
ProposerExhausted,
TargetMet,
PatienceExhausted,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProposedPatch {
pub knob_index: usize,
pub proposed_value: KnobValue,
pub rationale: String,
}
pub trait Proposer {
fn propose(
&mut self,
knobs: &[CalibrationKnob],
current_loss: (f64, f64),
history: &[StepReport],
) -> Option<ProposedPatch>;
}
pub trait Evaluator {
fn evaluate(
&self,
knobs: &[CalibrationKnob],
seed: u64,
) -> Result<BehavioralFidelityReport, EvaluatorError>;
}
#[derive(Debug)]
pub struct EvaluatorError(pub String);
impl std::fmt::Display for EvaluatorError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "evaluator error: {}", self.0)
}
}
impl std::error::Error for EvaluatorError {}
pub struct CalibrationLoop {
pub objective: CalibrationObjective,
pub knobs: Vec<CalibrationKnob>,
pub config: CalibrationConfig,
pub history: Vec<StepReport>,
pub best_loss: Option<(f64, f64)>,
pub best_knob_values: BTreeMap<String, KnobValue>,
effective_damping: f64,
steps_since_improvement: usize,
}
impl CalibrationLoop {
pub fn new(
objective: CalibrationObjective,
knobs: Vec<CalibrationKnob>,
config: CalibrationConfig,
) -> Self {
let damping = config.damping;
Self {
objective,
knobs,
config,
history: Vec::new(),
best_loss: None,
best_knob_values: BTreeMap::new(),
effective_damping: damping,
steps_since_improvement: 0,
}
}
pub fn step<E: Evaluator, P: Proposer>(
&mut self,
evaluator: &E,
proposer: &mut P,
) -> Result<&StepReport, EvaluatorError> {
let iter = self.history.len();
let (mean_before, std_before) = self.measure_loss(evaluator)?;
if let Some(target) = self.objective.target {
if mean_before <= target {
return Ok(self.record(StepReport {
iter,
loss_before_mean: mean_before,
loss_before_std: std_before,
proposed_patch: None,
loss_after_mean: None,
loss_after_std: None,
knob_values: self.snapshot_knobs(),
outcome: StepOutcome::TargetMet,
}));
}
}
if self.best_loss.map(|(m, _)| mean_before < m).unwrap_or(true) {
self.best_loss = Some((mean_before, std_before));
self.best_knob_values = self.snapshot_knobs();
self.steps_since_improvement = 0;
}
let raw_patch = proposer.propose(&self.knobs, (mean_before, std_before), &self.history);
let Some(patch) = raw_patch else {
return Ok(self.record(StepReport {
iter,
loss_before_mean: mean_before,
loss_before_std: std_before,
proposed_patch: None,
loss_after_mean: None,
loss_after_std: None,
knob_values: self.snapshot_knobs(),
outcome: StepOutcome::ProposerExhausted,
}));
};
let damped_value = damp_value(
self.knobs[patch.knob_index].current,
patch.proposed_value,
self.effective_damping,
);
let pre_value = self.knobs[patch.knob_index].current;
let _clip_result = self.knobs[patch.knob_index].apply(damped_value);
let (mean_after, std_after) = self.measure_loss(evaluator)?;
let outcome = self.decide_outcome(
mean_before,
std_before,
mean_after,
patch.knob_index,
pre_value,
);
match outcome {
StepOutcome::Improved => {
self.steps_since_improvement = 0;
}
_ => {
self.steps_since_improvement += 1;
}
}
Ok(self.record(StepReport {
iter,
loss_before_mean: mean_before,
loss_before_std: std_before,
proposed_patch: Some(ProposedPatch {
knob_index: patch.knob_index,
proposed_value: damped_value,
rationale: patch.rationale,
}),
loss_after_mean: Some(mean_after),
loss_after_std: Some(std_after),
knob_values: self.snapshot_knobs(),
outcome,
}))
}
pub fn run<E: Evaluator, P: Proposer>(
&mut self,
evaluator: &E,
proposer: &mut P,
) -> Result<&[StepReport], EvaluatorError> {
for _ in 0..self.config.max_iterations {
let outcome = self.step(evaluator, proposer)?.outcome;
if matches!(
outcome,
StepOutcome::TargetMet
| StepOutcome::ProposerExhausted
| StepOutcome::PatienceExhausted
) {
break;
}
if self.steps_since_improvement >= self.config.patience {
let last = self.history.last().expect("history non-empty after step");
let mut term = last.clone();
term.iter = self.history.len();
term.outcome = StepOutcome::PatienceExhausted;
term.proposed_patch = None;
term.loss_after_mean = None;
term.loss_after_std = None;
self.history.push(term);
break;
}
}
Ok(&self.history)
}
fn measure_loss<E: Evaluator>(&self, evaluator: &E) -> Result<(f64, f64), EvaluatorError> {
let mut reports = Vec::with_capacity(self.config.seeds_per_iteration);
for seed in 0..self.config.seeds_per_iteration as u64 {
reports.push(evaluator.evaluate(&self.knobs, seed)?);
}
self.objective.aggregate(&reports).ok_or_else(|| {
EvaluatorError("objective returned None from non-empty report set".into())
})
}
fn snapshot_knobs(&self) -> BTreeMap<String, KnobValue> {
self.knobs
.iter()
.map(|k| (k.path.clone(), k.current))
.collect()
}
fn decide_outcome(
&mut self,
mean_before: f64,
std_before: f64,
mean_after: f64,
knob_idx: usize,
pre_value: KnobValue,
) -> StepOutcome {
let beat_noise_floor = std_before > 0.0
&& (mean_before - mean_after) > self.config.min_improvement * std_before;
if mean_after < mean_before && beat_noise_floor {
self.best_loss = Some((mean_after, 0.0));
self.best_knob_values = self.snapshot_knobs();
return StepOutcome::Improved;
}
if mean_after >= mean_before {
match self.config.rollback {
RollbackPolicy::Revert => {
self.knobs[knob_idx].current = pre_value;
StepOutcome::Reverted
}
RollbackPolicy::Keep => StepOutcome::AcceptedNoNoiseFloorBeat,
RollbackPolicy::HalveDamping => {
self.effective_damping *= 0.5;
self.knobs[knob_idx].current = pre_value;
StepOutcome::Reverted
}
}
} else {
StepOutcome::AcceptedNoNoiseFloorBeat
}
}
fn record(&mut self, report: StepReport) -> &StepReport {
self.history.push(report);
self.history.last().expect("just pushed")
}
}
fn damp_value(current: KnobValue, proposed: KnobValue, damping: f64) -> KnobValue {
let cur = current.as_f64();
let prop = proposed.as_f64();
let damped = cur + (prop - cur) * damping;
match proposed {
KnobValue::F64(_) => KnobValue::F64(damped),
KnobValue::Usize(_) => KnobValue::Usize(damped.round().max(0.0) as usize),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::behavioral_fidelity::report::{
BaselineValues, BehavioralFidelityReport, CorpusSummary, EntityMetrics, GateResult,
PerMetric,
};
use chrono::Utc;
fn empty_per_metric() -> PerMetric {
PerMetric {
raw: 0.0,
baseline: 0.0,
dr: 0.0,
is_degenerate_baseline: false,
is_volume_bounded: false,
}
}
fn empty_em() -> EntityMetrics {
EntityMetrics {
entity_column: "t".into(),
p1_ietd: empty_per_metric(),
p1_autocorr: empty_per_metric(),
p2_active_lifetime: empty_per_metric(),
p2_burst_len_by_threshold: BTreeMap::new(),
p2_je_line_burst: empty_per_metric(),
p3_fanout_by_attr: BTreeMap::new(),
p3_clustering: empty_per_metric(),
p3_triangle_log_ratio: empty_per_metric(),
p4_rule_results: vec![],
p4_mean_gap: empty_per_metric(),
}
}
fn make_report(composite: f64) -> BehavioralFidelityReport {
BehavioralFidelityReport {
profile: "t".into(),
generator_id: "t".into(),
generator_version: "v5.x".into(),
seed: 0,
generated_at: Utc::now(),
reference_corpus: CorpusSummary {
path: "/dev/null".into(),
n_rows: 0,
n_entities_primary: 0,
n_entities_secondary: 0,
period_start: None,
period_end: None,
},
synthetic: CorpusSummary {
path: "/dev/null".into(),
n_rows: 0,
n_entities_primary: 0,
n_entities_secondary: 0,
period_start: None,
period_end: None,
},
noise_floor: BaselineValues {
p1_ietd_w1_days: 0.0,
p1_autocorr_gap: 0.0,
p2_active_lifetime_w1: 0.0,
p2_burst_len_by_threshold: BTreeMap::new(),
p2_je_line_burst_w1: 0.0,
p3_fanout_by_attr: BTreeMap::new(),
p3_clustering_gap: 0.0,
p3_triangle_log_ratio: 0.0,
p4_mean_gap: 0.0,
},
per_entity: {
let mut m = BTreeMap::new();
m.insert("t".into(), empty_em());
m
},
composite_bf_score: composite,
composite_bf_median: composite,
n_metrics_aggregated: 1,
n_metrics_excluded_degenerate: 0,
composite_bf_volume_corrected: composite,
n_metrics_excluded_volume: 0,
intraday_structural: None,
gates: GateResult {
fail_if_dr_above: 100.0,
fail_if_composite_above: 100.0,
passed: true,
failures: vec![],
},
}
}
struct LinearMockEvaluator {
optimum: f64,
noise: f64,
}
impl Evaluator for LinearMockEvaluator {
fn evaluate(
&self,
knobs: &[CalibrationKnob],
seed: u64,
) -> Result<BehavioralFidelityReport, EvaluatorError> {
let v = knobs[0].current.as_f64();
let noise = self.noise * (seed as f64 - 1.0); let composite = (v - self.optimum).abs() + noise;
Ok(make_report(composite))
}
}
struct StepTowardProposer {
target: f64,
}
impl Proposer for StepTowardProposer {
fn propose(
&mut self,
knobs: &[CalibrationKnob],
_current_loss: (f64, f64),
_history: &[StepReport],
) -> Option<ProposedPatch> {
let cur = knobs[0].current.as_f64();
if (cur - self.target).abs() < 1e-9 {
return None;
}
let direction = (self.target - cur).signum();
let step = direction * knobs[0].max_step;
Some(ProposedPatch {
knob_index: 0,
proposed_value: KnobValue::F64(cur + step),
rationale: format!("step toward {target}", target = self.target),
})
}
}
#[test]
fn step_reduces_loss_when_moving_toward_optimum() {
let knobs = vec![CalibrationKnob::new_f64("test.rate", 0.10, 0.0, 1.0, 0.05)];
let mut loop_ = CalibrationLoop::new(
CalibrationObjective::bf_composite(),
knobs,
CalibrationConfig {
seeds_per_iteration: 1,
max_iterations: 1,
min_improvement: 0.0, damping: 1.0, ..CalibrationConfig::default()
},
);
let eval = LinearMockEvaluator {
optimum: 0.02,
noise: 0.0,
};
let mut prop = StepTowardProposer { target: 0.02 };
let report = loop_.step(&eval, &mut prop).expect("step ok").clone();
assert!(report.loss_before_mean > 0.0);
let after = report.loss_after_mean.unwrap();
assert!(
after < report.loss_before_mean,
"step should reduce loss (before={}, after={})",
report.loss_before_mean,
after
);
}
#[test]
fn run_converges_to_optimum_within_max_iter() {
let knobs = vec![CalibrationKnob::new_f64("test.rate", 0.10, 0.0, 1.0, 0.02)];
let mut loop_ = CalibrationLoop::new(
CalibrationObjective::bf_composite().with_target(0.001),
knobs,
CalibrationConfig {
seeds_per_iteration: 1,
max_iterations: 10,
min_improvement: 0.0,
damping: 1.0,
patience: 20, ..CalibrationConfig::default()
},
);
let eval = LinearMockEvaluator {
optimum: 0.02,
noise: 0.0,
};
let mut prop = StepTowardProposer { target: 0.02 };
let history = loop_.run(&eval, &mut prop).unwrap().to_vec();
assert!(!history.is_empty());
let final_value = loop_.knobs[0].current.as_f64();
assert!(
(final_value - 0.02).abs() < 1e-6,
"final knob value should converge to optimum: got {final_value}"
);
assert!(
history
.iter()
.any(|s| matches!(s.outcome, StepOutcome::TargetMet)),
"convergence target should have been met before max_iter"
);
}
#[test]
fn proposer_exhaustion_stops_the_loop() {
let knobs = vec![CalibrationKnob::new_f64("test.rate", 0.02, 0.0, 1.0, 0.02)];
let mut loop_ = CalibrationLoop::new(
CalibrationObjective::bf_composite(),
knobs,
CalibrationConfig {
seeds_per_iteration: 1,
max_iterations: 10,
..CalibrationConfig::default()
},
);
let eval = LinearMockEvaluator {
optimum: 0.02,
noise: 0.0,
};
let mut prop = StepTowardProposer { target: 0.02 };
let history = loop_.run(&eval, &mut prop).unwrap().to_vec();
assert_eq!(history.len(), 1);
assert!(matches!(history[0].outcome, StepOutcome::ProposerExhausted));
}
#[test]
fn rollback_revert_restores_pre_step_value() {
let knobs = vec![CalibrationKnob::new_f64("test.rate", 0.02, 0.0, 1.0, 0.05)];
let mut loop_ = CalibrationLoop::new(
CalibrationObjective::bf_composite(),
knobs,
CalibrationConfig {
seeds_per_iteration: 1,
max_iterations: 3,
min_improvement: 0.0,
damping: 1.0,
rollback: RollbackPolicy::Revert,
patience: 20,
},
);
let eval = LinearMockEvaluator {
optimum: 0.02,
noise: 0.0,
};
let mut prop = StepTowardProposer { target: 0.5 };
loop_.run(&eval, &mut prop).unwrap();
let final_value = loop_.knobs[0].current.as_f64();
assert!(
(final_value - 0.02).abs() < 1e-9,
"Revert policy should restore the starting value; got {final_value}"
);
assert!(
loop_
.history
.iter()
.any(|s| matches!(s.outcome, StepOutcome::Reverted)),
"at least one step should have been Reverted"
);
}
#[test]
fn multi_seed_aggregate_produces_std() {
let knobs = vec![CalibrationKnob::new_f64("test.rate", 0.10, 0.0, 1.0, 0.05)];
let loop_ = CalibrationLoop::new(
CalibrationObjective::bf_composite(),
knobs,
CalibrationConfig {
seeds_per_iteration: 3,
..CalibrationConfig::default()
},
);
let eval = LinearMockEvaluator {
optimum: 0.02,
noise: 0.01,
};
let (mean, std) = loop_.measure_loss(&eval).unwrap();
assert!((mean - 0.08).abs() < 1e-9, "expected mean 0.08, got {mean}");
assert!(
(std - (2.0_f64 / 30000.0).sqrt()).abs() < 1e-9,
"expected std ≈ 0.00816, got {std}"
);
}
}