use super::*;
pub(crate) const OPERATOR_TRUST_RESTART_RADIUS_FLOOR: f64 = 1.0e-6;
#[derive(Clone, Debug)]
pub(crate) struct OuterConfig {
pub(crate) tolerance: f64,
pub(crate) rel_cost_tolerance: Option<f64>,
pub(crate) max_iter: usize,
pub(crate) bounds: Option<(Array1<f64>, Array1<f64>)>,
pub(crate) seed_config: crate::seeding::SeedConfig,
pub(crate) rho_bound: f64,
pub(crate) heuristic_lambdas: Option<Vec<f64>>,
pub(crate) initial_rho: Option<Array1<f64>>,
pub(crate) fallback_policy: FallbackPolicy,
pub(crate) screening_cap: Option<Arc<AtomicUsize>>,
pub(crate) screen_initial_rho: bool,
pub(crate) outer_inner_cap: Option<InnerProgressFeedback>,
pub(crate) operator_initial_trust_radius: Option<f64>,
pub(crate) arc_initial_regularization: Option<f64>,
pub(crate) objective_scale: Option<f64>,
pub(crate) bfgs_step_cap: Option<f64>,
pub(crate) bfgs_step_cap_psi: Option<f64>,
pub(crate) cache_session: Option<Arc<CacheSession>>,
pub(crate) cache_mirror_sessions: Vec<Arc<CacheSession>>,
pub(crate) rho_uncertainty_problem_size: crate::rho_uncertainty::RhoUncertaintyProblemSize,
pub(crate) warm_start_cache_hit: bool,
pub(crate) warm_start_outer_hessian: Option<Array2<f64>>,
pub(crate) rho_canonical_keys: Option<Vec<u64>>,
}
impl Default for OuterConfig {
fn default() -> Self {
Self {
tolerance: 1e-5,
rel_cost_tolerance: None,
max_iter: 200,
bounds: None,
seed_config: crate::seeding::SeedConfig::default(),
rho_bound: 30.0,
heuristic_lambdas: None,
initial_rho: None,
fallback_policy: FallbackPolicy::Automatic,
screening_cap: None,
screen_initial_rho: false,
outer_inner_cap: None,
operator_initial_trust_radius: None,
arc_initial_regularization: None,
objective_scale: None,
bfgs_step_cap: None,
bfgs_step_cap_psi: None,
cache_session: None,
cache_mirror_sessions: Vec::new(),
rho_uncertainty_problem_size:
crate::rho_uncertainty::RhoUncertaintyProblemSize::default(),
warm_start_cache_hit: false,
warm_start_outer_hessian: None,
rho_canonical_keys: None,
}
}
}
pub struct OuterProblem {
n_params: usize,
gradient: Derivative,
hessian: DeclaredHessianForm,
prefer_gradient_only: bool,
disable_fixed_point: bool,
psi_dim: usize,
barrier_config: Option<BarrierConfig>,
tolerance: f64,
rel_cost_tolerance: Option<f64>,
max_iter: usize,
bounds: Option<(Array1<f64>, Array1<f64>)>,
rho_bound: f64,
seed_config: crate::seeding::SeedConfig,
heuristic_lambdas: Option<Vec<f64>>,
initial_rho: Option<Array1<f64>>,
fallback_policy: FallbackPolicy,
screening_cap: Option<Arc<AtomicUsize>>,
screen_initial_rho: bool,
outer_inner_cap: Option<InnerProgressFeedback>,
operator_initial_trust_radius: Option<f64>,
arc_initial_regularization: Option<f64>,
objective_scale: Option<f64>,
bfgs_step_cap: Option<f64>,
bfgs_step_cap_psi: Option<f64>,
cache_session: Option<Arc<CacheSession>>,
cache_mirror_sessions: Vec<Arc<CacheSession>>,
rho_uncertainty_problem_size: crate::rho_uncertainty::RhoUncertaintyProblemSize,
continuation_prewarm: bool,
rho_canonical_keys: Option<Vec<u64>>,
}
impl OuterProblem {
pub fn new(n_params: usize) -> Self {
Self {
n_params,
gradient: Derivative::Unavailable,
hessian: DeclaredHessianForm::Unavailable,
prefer_gradient_only: false,
disable_fixed_point: false,
psi_dim: 0,
barrier_config: None,
tolerance: 1e-5,
rel_cost_tolerance: None,
max_iter: 200,
bounds: None,
rho_bound: 30.0,
seed_config: crate::seeding::SeedConfig::default(),
heuristic_lambdas: None,
initial_rho: None,
fallback_policy: FallbackPolicy::Automatic,
screening_cap: None,
screen_initial_rho: false,
outer_inner_cap: None,
operator_initial_trust_radius: None,
arc_initial_regularization: None,
objective_scale: None,
bfgs_step_cap: None,
bfgs_step_cap_psi: None,
cache_session: None,
cache_mirror_sessions: Vec::new(),
rho_uncertainty_problem_size:
crate::rho_uncertainty::RhoUncertaintyProblemSize::default(),
continuation_prewarm: true,
rho_canonical_keys: None,
}
}
pub fn with_rho_canonical_keys(mut self, keys: Option<Vec<u64>>) -> Self {
self.rho_canonical_keys = keys;
self
}
pub fn with_gradient(mut self, d: Derivative) -> Self {
self.gradient = d;
self
}
pub fn with_hessian(mut self, form: DeclaredHessianForm) -> Self {
self.hessian = form;
self
}
pub fn with_prefer_gradient_only(mut self, prefer_gradient_only: bool) -> Self {
self.prefer_gradient_only = prefer_gradient_only;
self
}
pub fn with_disable_fixed_point(mut self, disable: bool) -> Self {
self.disable_fixed_point = disable;
self
}
pub fn with_psi_dim(mut self, dim: usize) -> Self {
self.psi_dim = dim;
self
}
pub fn with_barrier(mut self, cfg: Option<BarrierConfig>) -> Self {
self.barrier_config = cfg;
self
}
pub fn with_tolerance(mut self, tol: f64) -> Self {
self.tolerance = tol;
self
}
pub fn with_max_iter(mut self, n: usize) -> Self {
self.max_iter = n;
self
}
pub fn with_bounds(mut self, lo: Array1<f64>, hi: Array1<f64>) -> Self {
self.bounds = Some((lo, hi));
self
}
pub fn with_rho_bound(mut self, b: f64) -> Self {
self.rho_bound = b;
self
}
pub fn with_seed_config(mut self, sc: crate::seeding::SeedConfig) -> Self {
self.seed_config = sc;
self
}
pub fn with_heuristic_lambdas(mut self, h: Vec<f64>) -> Self {
self.heuristic_lambdas = Some(h);
self
}
pub fn with_initial_rho(mut self, rho: Array1<f64>) -> Self {
self.initial_rho = Some(rho);
self
}
pub fn with_continuation_prewarm(mut self, enabled: bool) -> Self {
self.continuation_prewarm = enabled;
self
}
pub fn with_screening_cap(mut self, screening_cap: Arc<AtomicUsize>) -> Self {
self.screening_cap = Some(screening_cap);
self
}
pub fn with_screen_initial_rho(mut self, screen_initial_rho: bool) -> Self {
self.screen_initial_rho = screen_initial_rho;
self
}
pub fn with_outer_inner_cap(mut self, feedback: InnerProgressFeedback) -> Self {
self.outer_inner_cap = Some(feedback);
self
}
pub fn with_operator_initial_trust_radius(mut self, radius: Option<f64>) -> Self {
self.operator_initial_trust_radius = sanitized_operator_trust_restart_radius(radius);
self
}
pub fn with_arc_initial_regularization(mut self, sigma: Option<f64>) -> Self {
self.arc_initial_regularization = sigma.filter(|v| v.is_finite() && *v > 0.0);
self
}
pub fn with_objective_scale(mut self, scale: Option<f64>) -> Self {
self.objective_scale = scale.filter(|v| v.is_finite() && *v > 0.0);
self
}
pub fn with_rel_cost_tolerance(mut self, rel_cost: Option<f64>) -> Self {
self.rel_cost_tolerance = rel_cost.filter(|v| v.is_finite() && *v > 0.0);
self
}
pub fn with_bfgs_step_cap(mut self, cap: Option<f64>) -> Self {
self.bfgs_step_cap = cap.filter(|v| v.is_finite() && *v > 0.0);
self
}
pub fn with_bfgs_step_cap_psi(mut self, cap: Option<f64>) -> Self {
self.bfgs_step_cap_psi = cap.filter(|v| v.is_finite() && *v > 0.0);
self
}
pub fn with_cache_session(mut self, session: Arc<CacheSession>) -> Self {
self.cache_session = Some(session);
self
}
pub fn with_cache_mirror_sessions(mut self, sessions: Vec<Arc<CacheSession>>) -> Self {
self.cache_mirror_sessions = sessions;
self
}
pub fn with_problem_size(mut self, n_obs: usize, p_coefficients: usize) -> Self {
self.rho_uncertainty_problem_size = crate::rho_uncertainty::RhoUncertaintyProblemSize {
n_obs: Some(n_obs),
p_coefficients: Some(p_coefficients),
};
self
}
pub fn with_fallback_policy(mut self, policy: FallbackPolicy) -> Self {
self.fallback_policy = policy;
self
}
fn capability(&self) -> OuterCapability {
OuterCapability {
gradient: self.gradient,
hessian: self.hessian,
prefer_gradient_only: self.prefer_gradient_only,
disable_fixed_point: self.disable_fixed_point,
n_params: self.n_params,
psi_dim: self.psi_dim,
fixed_point_available: false,
barrier_config: self.barrier_config.clone(),
}
}
pub(crate) fn config(&self) -> OuterConfig {
OuterConfig {
tolerance: self.tolerance,
rel_cost_tolerance: self.rel_cost_tolerance,
max_iter: self.max_iter,
bounds: self.bounds.clone(),
seed_config: self.seed_config,
rho_bound: self.rho_bound,
heuristic_lambdas: self.heuristic_lambdas.clone(),
initial_rho: self.initial_rho.clone(),
fallback_policy: self.fallback_policy,
screening_cap: self.screening_cap.clone(),
screen_initial_rho: self.screen_initial_rho,
outer_inner_cap: self.outer_inner_cap.clone(),
operator_initial_trust_radius: self.operator_initial_trust_radius,
arc_initial_regularization: self.arc_initial_regularization,
objective_scale: self.objective_scale,
bfgs_step_cap: self.bfgs_step_cap,
bfgs_step_cap_psi: self.bfgs_step_cap_psi,
cache_session: self.cache_session.clone(),
cache_mirror_sessions: self.cache_mirror_sessions.clone(),
rho_uncertainty_problem_size: self.rho_uncertainty_problem_size,
warm_start_cache_hit: false,
warm_start_outer_hessian: None,
rho_canonical_keys: self.rho_canonical_keys.clone(),
}
}
pub fn build_objective<S, Fc, Fe, Fr, Fefs>(
&self,
state: S,
cost_fn: Fc,
eval_fn: Fe,
reset_fn: Option<Fr>,
efs_fn: Option<Fefs>,
) -> ClosureObjective<S, Fc, Fe, Fr, Fefs>
where
Fc: FnMut(&mut S, &Array1<f64>) -> Result<f64, EstimationError>,
Fe: FnMut(&mut S, &Array1<f64>) -> Result<OuterEval, EstimationError>,
Fr: FnMut(&mut S),
Fefs: FnMut(&mut S, &Array1<f64>) -> Result<EfsEval, EstimationError>,
{
let mut cap = self.capability();
cap.fixed_point_available = efs_fn.is_some();
ClosureObjective {
state,
cap,
cost_fn,
eval_fn,
eval_order_fn: None,
reset_fn,
efs_fn,
screening_proxy_fn: None::<fn(&mut S, &Array1<f64>) -> Result<f64, EstimationError>>,
seed_fn: None::<fn(&mut S, &Array1<f64>) -> Result<SeedOutcome, EstimationError>>,
continuation_prewarm: self.continuation_prewarm,
}
}
pub fn build_objective_with_eval_order<S, Fc, Fe, Feo, Fr, Fefs>(
&self,
state: S,
cost_fn: Fc,
eval_fn: Fe,
eval_order_fn: Feo,
reset_fn: Option<Fr>,
efs_fn: Option<Fefs>,
) -> ClosureObjective<S, Fc, Fe, Fr, Fefs, Feo>
where
Fc: FnMut(&mut S, &Array1<f64>) -> Result<f64, EstimationError>,
Fe: FnMut(&mut S, &Array1<f64>) -> Result<OuterEval, EstimationError>,
Feo: FnMut(&mut S, &Array1<f64>, OuterEvalOrder) -> Result<OuterEval, EstimationError>,
Fr: FnMut(&mut S),
Fefs: FnMut(&mut S, &Array1<f64>) -> Result<EfsEval, EstimationError>,
{
let mut cap = self.capability();
cap.fixed_point_available = efs_fn.is_some();
ClosureObjective {
state,
cap,
cost_fn,
eval_fn,
eval_order_fn: Some(eval_order_fn),
reset_fn,
efs_fn,
screening_proxy_fn: None::<fn(&mut S, &Array1<f64>) -> Result<f64, EstimationError>>,
seed_fn: None::<fn(&mut S, &Array1<f64>) -> Result<SeedOutcome, EstimationError>>,
continuation_prewarm: self.continuation_prewarm,
}
}
pub fn build_objective_with_screening_proxy<S, Fc, Fe, Feo, Fr, Fefs, Fsp>(
&self,
state: S,
cost_fn: Fc,
eval_fn: Fe,
eval_order_fn: Feo,
reset_fn: Option<Fr>,
efs_fn: Option<Fefs>,
screening_proxy_fn: Fsp,
) -> ClosureObjective<S, Fc, Fe, Fr, Fefs, Feo, Fsp>
where
Fc: FnMut(&mut S, &Array1<f64>) -> Result<f64, EstimationError>,
Fe: FnMut(&mut S, &Array1<f64>) -> Result<OuterEval, EstimationError>,
Feo: FnMut(&mut S, &Array1<f64>, OuterEvalOrder) -> Result<OuterEval, EstimationError>,
Fr: FnMut(&mut S),
Fefs: FnMut(&mut S, &Array1<f64>) -> Result<EfsEval, EstimationError>,
Fsp: FnMut(&mut S, &Array1<f64>) -> Result<f64, EstimationError>,
{
let mut cap = self.capability();
cap.fixed_point_available = efs_fn.is_some();
ClosureObjective {
state,
cap,
cost_fn,
eval_fn,
eval_order_fn: Some(eval_order_fn),
reset_fn,
efs_fn,
screening_proxy_fn: Some(screening_proxy_fn),
seed_fn: None::<fn(&mut S, &Array1<f64>) -> Result<SeedOutcome, EstimationError>>,
continuation_prewarm: self.continuation_prewarm,
}
}
pub fn run(
&self,
obj: &mut dyn OuterObjective,
context: &str,
) -> Result<OuterResult, EstimationError> {
let mut config = self.config();
let Some(session) = config.cache_session.clone() else {
return run_outer(obj, &config, context);
};
let key_hex = session.key().to_hex();
let short_key = &key_hex[..8.min(key_hex.len())];
let mut had_hit = false;
let mut cached_inner_beta: Option<Array1<f64>> = None;
if let Some(loaded) = session.try_load_with_source() {
match classify_cache_entry_for_outer(&loaded, self.n_params) {
CacheSeedDecision::ExactFinal {
rho,
beta: _beta_final,
final_value,
iterations,
prior_obj_display,
} => {
let cap = primary_capability_for_config(obj.capability(), &config, context);
let plan_used = plan(&cap);
log::info!(
"[CACHE] final-hit key={}.. context={} rho_dim={} prior_obj={:.6e} iter={} action=skip-outer-validation",
short_key,
context,
rho.len(),
prior_obj_display,
iterations,
);
let mut result =
OuterResult::new(rho, final_value, iterations, true, plan_used);
result.rho_uncertainty_diagnostic = Some(compute_rho_uncertainty_diagnostic(
obj,
&config,
context,
&mut result,
));
return Ok(result);
}
CacheSeedDecision::Seed {
rho,
beta,
hessian,
prior_obj_display,
iteration,
} => {
let beta_len = beta.len();
let beta_arr = if beta.is_empty() {
None
} else {
Some(Array1::from_vec(beta))
};
config.warm_start_outer_hessian = hessian.and_then(|(dim, flat)| {
if dim == self.n_params && flat.len() == dim * dim {
Array2::from_shape_vec((dim, dim), flat).ok()
} else {
None
}
});
if config
.initial_rho
.as_ref()
.is_none_or(|initial| initial != rho)
{
log::info!(
"[CACHE] hit key={}.. context={} rho_dim={} beta_dim={} prior_obj={:.6e} iter={}",
short_key,
context,
rho.len(),
beta_len,
prior_obj_display,
iteration,
);
config.initial_rho = Some(rho);
config.screen_initial_rho = false;
had_hit = true;
} else {
log::info!(
"[CACHE] hit key={}.. context={} rho_dim={} beta_dim={} already-aligned prior_obj={:.6e}",
short_key,
context,
rho.len(),
beta_len,
prior_obj_display,
);
had_hit = true;
}
cached_inner_beta = beta_arr;
}
CacheSeedDecision::Discard {
reason: "payload-shape-mismatch",
..
} => {
log::info!(
"[CACHE] skip key={}.. context={} reason=payload-shape-mismatch n_params={}",
short_key,
context,
self.n_params,
);
}
CacheSeedDecision::Discard {
reason,
prior_obj_display,
all_rho_finite,
} => {
log::info!(
"[CACHE] skip key={}.. context={} reason={} prior_obj={:.6e} all_rho_finite={}",
short_key,
context,
reason,
prior_obj_display,
all_rho_finite.unwrap_or(false),
);
}
}
} else {
log::info!(
"[CACHE] miss key={}.. context={} reason=fresh-fingerprint n_params={}",
short_key,
context,
self.n_params,
);
}
config.warm_start_cache_hit = had_hit;
let mut checkpointing = CheckpointingObjective::new(
obj,
Arc::clone(&session),
config.cache_mirror_sessions.clone(),
);
if let Some(beta) = cached_inner_beta.as_ref() {
match checkpointing.seed_inner_state(beta) {
Ok(SeedOutcome::Installed) => log::info!(
"[CACHE] beta-warm key={}.. context={} beta_dim={} action=installed",
short_key,
context,
beta.len(),
),
Ok(SeedOutcome::NoSlot) => log::warn!(
"[CACHE] beta-warm key={}.. context={} beta_dim={} action=skip \
reason=objective_has_no_inner_beta_slot",
short_key,
context,
beta.len(),
),
Ok(SeedOutcome::Incompatible) => log::info!(
"[CACHE] beta-warm key={}.. context={} beta_dim={} action=rho-only \
reason=seed_beta_length_incompatible_with_inner_blocks",
short_key,
context,
beta.len(),
),
Err(err) => log::warn!(
"[CACHE] beta-warm key={}.. context={} beta_dim={} action=skip err={}",
short_key,
context,
beta.len(),
err,
),
}
}
let result = run_outer(&mut checkpointing, &config, context);
let final_beta = checkpointing.last_inner_beta();
if let Ok(result) = result.as_ref()
&& result.final_value.is_finite()
&& result.converged
&& let Some(bytes) = encode_iterate(
&result.rho,
final_beta.as_ref(),
result.final_hessian.as_ref(),
result.final_value,
result.iterations as u64,
)
{
let saved = session.finalize(
&bytes,
Some(result.final_value),
Some(result.iterations as u64),
);
if saved {
log::info!(
"[CACHE] save key={}.. context={} final_obj={:.6e} iter={} resumed={}",
short_key,
context,
result.final_value,
result.iterations,
had_hit,
);
}
for mirror in &config.cache_mirror_sessions {
let mirror_saved = mirror.finalize(
&bytes,
Some(result.final_value),
Some(result.iterations as u64),
);
if mirror_saved {
let mirror_hex = mirror.key().to_hex();
log::info!(
"[CACHE] save key={}.. context={} mirror final_obj={:.6e} iter={}",
&mirror_hex[..8.min(mirror_hex.len())],
context,
result.final_value,
result.iterations,
);
}
}
}
result
}
}
#[derive(Clone, Debug)]
pub struct OuterResult {
pub rho: Array1<f64>,
pub final_value: f64,
pub iterations: usize,
pub final_grad_norm: Option<f64>,
pub final_gradient: Option<Array1<f64>>,
pub final_hessian: Option<Array2<f64>>,
pub converged: bool,
pub plan_used: OuterPlan,
pub operator_trust_radius: Option<f64>,
pub operator_stop_reason: Option<OperatorTrustRegionStopReason>,
pub criterion_certificate: Option<CriterionCertificate>,
pub rho_uncertainty_diagnostic: Option<crate::rho_uncertainty::RhoUncertaintyDiagnostic>,
}
impl OuterResult {
pub fn new(
rho: Array1<f64>,
final_value: f64,
iterations: usize,
converged: bool,
plan_used: OuterPlan,
) -> Self {
Self {
rho,
final_value,
iterations,
final_grad_norm: None,
final_gradient: None,
final_hessian: None,
converged,
plan_used,
operator_trust_radius: None,
operator_stop_reason: None,
criterion_certificate: None,
rho_uncertainty_diagnostic: None,
}
}
pub fn final_grad_norm_report(&self) -> String {
match self.final_grad_norm {
Some(g) => format!("{g:.3e}"),
None => "n/a".to_string(),
}
}
}
pub(crate) fn certificate_audit_direction(theta: &Array1<f64>, context: &str) -> Array1<f64> {
let mut seed: u64 = 0xcbf2_9ce4_8422_2325;
let mut fnv = |byte: u8| {
seed ^= u64::from(byte);
seed = seed.wrapping_mul(0x0000_0100_0000_01b3);
};
for byte in context.bytes() {
fnv(byte);
}
for &x in theta.iter() {
for byte in x.to_bits().to_le_bytes() {
fnv(byte);
}
}
let mut state = seed;
let mut next_unit = move || {
state = state.wrapping_add(0x9e37_79b9_7f4a_7c15);
let mut z = state;
z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
z ^= z >> 31;
((z >> 11) as f64 + 0.5) / (1u64 << 53) as f64
};
let mut direction = Array1::<f64>::zeros(theta.len());
let mut i = 0;
while i < direction.len() {
let (u1, u2) = (next_unit(), next_unit());
let radius = (-2.0 * u1.ln()).sqrt();
let angle = 2.0 * std::f64::consts::PI * u2;
direction[i] = radius * angle.cos();
if i + 1 < direction.len() {
direction[i + 1] = radius * angle.sin();
}
i += 2;
}
let norm = direction.dot(&direction).sqrt();
if norm.is_finite() && norm > f64::EPSILON {
direction.mapv_inplace(|v| v / norm);
direction
} else {
let mut fallback = Array1::<f64>::zeros(theta.len());
fallback[0] = 1.0;
fallback
}
}
pub(crate) fn certificate_hessian_is_pd(hessian: &Array2<f64>) -> Option<bool> {
let n = hessian.nrows();
if n == 0 || hessian.ncols() != n || hessian.iter().any(|v| !v.is_finite()) {
return None;
}
let mut chol = hessian.clone();
for j in 0..n {
for k in 0..j {
let l_jk = chol[[j, k]];
for i in j..n {
chol[[i, j]] -= chol[[i, k]] * l_jk;
}
}
let pivot = chol[[j, j]];
if !(pivot > 0.0) || !pivot.is_finite() {
return Some(false);
}
let inv_sqrt = 1.0 / pivot.sqrt();
for i in j..n {
chol[[i, j]] *= inv_sqrt;
}
}
Some(true)
}
pub(crate) fn certificate_railed_lambdas(
rho: &Array1<f64>,
rho_dim: usize,
config: &OuterConfig,
) -> Vec<usize> {
(0..rho_dim.min(rho.len()))
.filter(|&k| {
let (lo, hi) = match config.bounds.as_ref() {
Some((lo, hi)) if k < lo.len() && k < hi.len() => (lo[k], hi[k]),
Some(_) => return false,
None => (-config.rho_bound, config.rho_bound),
};
(rho[k] - lo).abs() <= CERTIFICATE_RAIL_MARGIN
|| (hi - rho[k]).abs() <= CERTIFICATE_RAIL_MARGIN
})
.collect()
}
pub(crate) fn audit_first_order_optimality(
obj: &mut dyn OuterObjective,
config: &OuterConfig,
context: &str,
result: &OuterResult,
) -> Option<CriterionCertificate> {
let gradient = result.final_gradient.as_ref()?;
if gradient.is_empty()
|| gradient.len() != result.rho.len()
|| gradient.iter().any(|g| !g.is_finite())
|| result.rho.iter().any(|r| !r.is_finite())
{
return None;
}
let theta = &result.rho;
let rho_dim = obj.capability().theta_layout().rho_dim();
let railed = certificate_railed_lambdas(theta, rho_dim, config);
let full_direction = certificate_audit_direction(theta, context);
let direction = if railed.is_empty() {
full_direction
} else {
let mut projected = full_direction;
for &k in &railed {
if k < projected.len() {
projected[k] = 0.0;
}
}
let norm = projected.dot(&projected).sqrt();
if norm.is_finite() && norm > f64::EPSILON {
projected.mapv_inplace(|v| v / norm);
projected
} else {
log::debug!(
"[CERTIFICATE] {context}: every audited coordinate railed \
(railed={railed:?}); no free direction to audit, certificate skipped"
);
return None;
}
};
let theta_scale = theta.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
let step = f64::EPSILON.cbrt() * (1.0 + theta_scale);
let mut probe = |scale: f64| -> Option<f64> {
let point = theta + &(scale * &direction);
match obj.eval_cost(&point) {
Ok(value) if value.is_finite() => Some(value),
Ok(value) => {
log::debug!(
"[CERTIFICATE] {context}: audit probe at θ̂{scale:+.3e}·v returned \
non-finite criterion value {value}; certificate skipped"
);
None
}
Err(err) => {
log::debug!(
"[CERTIFICATE] {context}: audit probe at θ̂{scale:+.3e}·v failed ({err}); \
certificate skipped"
);
None
}
}
};
let f_plus_h = probe(step)?;
let f_minus_h = probe(-step)?;
let f_plus_2h = probe(2.0 * step)?;
let f_minus_2h = probe(-2.0 * step)?;
let d_h = (f_plus_h - f_minus_h) / (2.0 * step);
let d_2h = (f_plus_2h - f_minus_2h) / (4.0 * step);
let fd_directional = (4.0 * d_h - d_2h) / 3.0; let value_scale = f_plus_h
.abs()
.max(f_minus_h.abs())
.max(f_plus_2h.abs())
.max(f_minus_2h.abs());
let roundoff = f64::EPSILON * (1.0 + value_scale) / step;
let fd_error = (d_h - d_2h).abs().max(roundoff);
let analytic_directional = gradient.dot(&direction);
let grad_norm = gradient.dot(gradient).sqrt();
let agreement_z = (analytic_directional - fd_directional).abs() / fd_error;
let certificate = CriterionCertificate {
grad_norm,
analytic_directional,
fd_directional, fd_error, agreement_z,
fd_step: step, hessian_pd: result
.final_hessian
.as_ref()
.and_then(certificate_hessian_is_pd),
lambdas_railed: railed,
};
if certificate.is_clean() {
log::info!("[CERTIFICATE] {context}: {}", certificate.summary());
} else {
log::warn!(
"[CERTIFICATE warning] {context}: optimality self-audit flagged the returned \
optimum — {}",
certificate.summary(),
);
}
Some(certificate)
}
pub(crate) fn compute_rho_uncertainty_diagnostic(
obj: &mut dyn OuterObjective,
config: &OuterConfig,
context: &str,
result: &mut OuterResult,
) -> crate::rho_uncertainty::RhoUncertaintyDiagnostic {
let cap = obj.capability();
let layout = cap.theta_layout();
let rho_dim = layout.rho_dim();
let gate = crate::rho_uncertainty::RhoUncertaintyCostGate {
sample_count: 32,
problem_size: config.rho_uncertainty_problem_size,
};
if let Err(reason) = crate::rho_uncertainty::cost_gate_allows(rho_dim, gate) {
return crate::rho_uncertainty::RhoUncertaintyDiagnostic::skipped(reason, 0);
}
if result.rho.len() != layout.n_params {
return crate::rho_uncertainty::RhoUncertaintyDiagnostic::skipped(
format!(
"final outer point length {} does not match objective dimension {}",
result.rho.len(),
layout.n_params
),
0,
);
}
let final_eval = match obj.eval_with_order(&result.rho, OuterEvalOrder::ValueGradientHessian) {
Ok(eval) => eval,
Err(err) => {
return crate::rho_uncertainty::RhoUncertaintyDiagnostic::skipped(
format!("final exact Hessian evaluation failed: {err}"),
1,
);
}
};
let hessian = match final_eval.hessian.materialize_dense() {
Ok(Some(hessian)) => hessian,
Ok(None) => {
return crate::rho_uncertainty::RhoUncertaintyDiagnostic::skipped(
"exact outer Hessian unavailable at fitted rho",
1,
);
}
Err(message) => {
return crate::rho_uncertainty::RhoUncertaintyDiagnostic::skipped(
format!("exact outer Hessian materialization failed: {message}"),
1,
);
}
};
if hessian.nrows() != layout.n_params || hessian.ncols() != layout.n_params {
return crate::rho_uncertainty::RhoUncertaintyDiagnostic::skipped(
format!(
"exact outer Hessian shape {}x{} does not match objective dimension {}",
hessian.nrows(),
hessian.ncols(),
layout.n_params
),
1,
);
}
if result.final_hessian.is_none() && hessian.iter().all(|v| v.is_finite()) {
result.final_hessian = Some(hessian.clone());
}
let mut hessian_rho = Array2::<f64>::zeros((rho_dim, rho_dim));
for row in 0..rho_dim {
for col in 0..rho_dim {
hessian_rho[[row, col]] = hessian[[row, col]];
}
}
let rho_hat = result.rho.slice(ndarray::s![..rho_dim]).to_owned();
let theta_hat = result.rho.clone();
let cost_hat = final_eval.cost;
let final_beta_hint = final_eval.inner_beta_hint.clone();
let diagnostic = {
let mut served_hat_cost = false;
let mut criterion = |rho: &Array1<f64>| -> Option<f64> {
let is_hat = rho.len() == rho_hat.len()
&& rho
.iter()
.zip(rho_hat.iter())
.all(|(&left, &right)| left.to_bits() == right.to_bits());
if is_hat && !served_hat_cost {
served_hat_cost = true;
return Some(cost_hat);
}
let mut theta = theta_hat.clone();
for idx in 0..rho_dim {
theta[idx] = rho[idx];
}
if let Some(beta) = final_beta_hint.as_ref()
&& obj.seed_inner_state(beta).is_err()
{
return None;
}
obj.eval_cost(&theta).ok()
};
crate::rho_uncertainty::rho_uncertainty_diagnostic(
&rho_hat,
&hessian_rho,
gate,
&mut criterion,
)
};
if let Some(beta) = final_beta_hint.as_ref()
&& let Err(err) = obj.seed_inner_state(beta)
{
log::debug!(
"[RHO uncertainty] {context}: final inner-state restore skipped after diagnostic ({err})"
);
}
match &diagnostic.status {
crate::rho_uncertainty::RhoUncertaintyStatus::NoEvidenceOfHeavyTails => {
log::info!(
"[RHO uncertainty] {context}: no heavy-tail evidence at sampled rho proposals k_hat={:.3} evals={}",
diagnostic.k_hat.unwrap_or(f64::NAN),
diagnostic.n_evaluations,
);
}
crate::rho_uncertainty::RhoUncertaintyStatus::HeavyTailsDetected { k_hat } => {
log::warn!(
"[RHO uncertainty] {context}: heavy rho-importance tail detected k_hat={:.3} evals={}",
k_hat,
diagnostic.n_evaluations,
);
}
crate::rho_uncertainty::RhoUncertaintyStatus::Skipped { reason } => {
log::info!("[RHO uncertainty] {context}: skipped ({reason})");
}
}
diagnostic
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OperatorTrustRegionStopReason {
Converged,
RejectFloor,
IterationBudget,
CostStallFlatValley,
RoutingMismatch,
}
pub(crate) fn run_outer(
obj: &mut dyn OuterObjective,
config: &OuterConfig,
context: &str,
) -> Result<OuterResult, EstimationError> {
if let Some(keys) = config.rho_canonical_keys.as_ref()
&& let Some(perm) = canonical_permutation(keys)
{
let canonical_config = canonicalize_outer_config(config, &perm);
let mut canonical_obj = CanonicalizedObjective::new(obj, perm.clone());
let result = run_outer(&mut canonical_obj, &canonical_config, context)?;
return Ok(outer_result_to_native(result, &perm));
}
let mut result = run_outer_uncertified(obj, config, context)?;
if config.max_iter <= 1 {
return Ok(result);
}
result.criterion_certificate = audit_first_order_optimality(obj, config, context, &result);
result.rho_uncertainty_diagnostic = Some(compute_rho_uncertainty_diagnostic(
obj,
config,
context,
&mut result,
));
Ok(result)
}
fn canonicalize_outer_config(config: &OuterConfig, perm: &[usize]) -> OuterConfig {
let permute_vec = |v: &[f64]| -> Vec<f64> {
if v.len() == perm.len() {
perm.iter().map(|&i| v[i]).collect()
} else {
v.to_vec()
}
};
let permute_arr = |a: &Array1<f64>| -> Array1<f64> {
if a.len() == perm.len() {
Array1::from_iter(perm.iter().map(|&i| a[i]))
} else {
a.clone()
}
};
let mut canonical = config.clone();
canonical.rho_canonical_keys = None;
if let Some(initial) = config.initial_rho.as_ref() {
canonical.initial_rho = Some(permute_arr(initial));
}
if let Some(h) = config.heuristic_lambdas.as_ref() {
canonical.heuristic_lambdas = Some(permute_vec(h));
}
if let Some((lower, upper)) = config.bounds.as_ref() {
canonical.bounds = Some((permute_arr(lower), permute_arr(upper)));
}
if let Some(h) = config.warm_start_outer_hessian.as_ref()
&& h.nrows() == perm.len()
&& h.ncols() == perm.len()
{
let mut hc = Array2::<f64>::zeros((perm.len(), perm.len()));
for (a, &ia) in perm.iter().enumerate() {
for (b, &ib) in perm.iter().enumerate() {
hc[[a, b]] = h[[ia, ib]];
}
}
canonical.warm_start_outer_hessian = Some(hc);
}
canonical
}
pub(crate) fn run_outer_uncertified(
obj: &mut dyn OuterObjective,
config: &OuterConfig,
context: &str,
) -> Result<OuterResult, EstimationError> {
let cap = primary_capability_for_config(obj.capability(), config, context);
cap.validate_layout(context)?;
if let Some(initial_rho) = config.initial_rho.as_ref() {
cap.theta_layout()
.validate_point_len(initial_rho, "initial outer seed")
.map_err(|err| match err {
ObjectiveEvalError::Recoverable { message }
| ObjectiveEvalError::Fatal { message } => {
EstimationError::RemlOptimizationFailed(format!("{context}: {message}"))
}
})?;
}
crate::solver::estimate::reml::outer_eval::clear_outer_ift_residual_energy_for_fit();
if let Some(result) = run_per_atom_efs_if_frontier(obj, config, context)? {
return Ok(result);
}
if cap.n_params == 0 {
let cost = obj.eval_cost(&Array1::zeros(0))?;
let the_plan = plan(&cap);
return Ok(outer_result_with_gradient_norm(
Array1::zeros(0),
cost,
0,
Some(0.0),
true,
the_plan,
));
}
let fallback_attempts = match config.fallback_policy {
FallbackPolicy::Automatic => automatic_fallback_attempts(&cap),
FallbackPolicy::Disabled => Vec::new(),
};
let mut attempts: Vec<OuterCapability> = Vec::with_capacity(1 + fallback_attempts.len());
attempts.push(cap.clone());
for degraded in fallback_attempts {
attempts.push(degraded);
}
let mut last_error: Option<EstimationError> = None;
for (attempt_idx, attempt_cap) in attempts.iter().enumerate() {
let the_plan = plan(attempt_cap);
if attempt_idx > 0 {
log::debug!("[OUTER] {context}: primary plan failed; falling back to {the_plan}");
}
log_plan(context, attempt_cap, &the_plan);
obj.reset();
let mut arc_retries_left: u32 = if matches!(the_plan.solver, Solver::Arc) {
2
} else {
0
};
let mut retry_config: Option<OuterConfig> = None;
let mut prev_attempt_grad_norm: Option<f64> = None;
let outcome = loop {
let active_config_owned: OuterConfig =
retry_config.clone().unwrap_or_else(|| config.clone());
let active_config: &OuterConfig = &active_config_owned;
match run_outer_with_plan(obj, active_config, context, attempt_cap, &the_plan) {
Ok(result) => {
if result.converged
|| arc_retries_left == 0
|| matches!(
result.operator_stop_reason,
Some(OperatorTrustRegionStopReason::RejectFloor)
)
{
break Ok(result);
}
let Some(cur_grad_norm) = result.final_grad_norm else {
log::info!(
"[OUTER] {context}: ARC attempt exhausted budget at \
iter={} cost={:.6e} without a final gradient norm; \
falling through to degraded plan",
result.iterations,
result.final_value,
);
break Ok(result);
};
if let Some(prev_g) = prev_attempt_grad_norm {
let progressed = cur_grad_norm.is_finite()
&& prev_g.is_finite()
&& cur_grad_norm < 0.5 * prev_g;
if !progressed {
log::info!(
"[OUTER] {context}: ARC retry stalled at \
iter={} cost={:.6e} |g|={:.6e} (prev |g|={:.6e}); \
deterministic replay suspected, falling through \
to degraded plan",
result.iterations,
result.final_value,
cur_grad_norm,
prev_g,
);
break Ok(result);
}
}
let next_trust_radius =
sanitized_operator_trust_restart_radius(result.operator_trust_radius);
log::info!(
"[OUTER] {context}: ARC attempt exhausted budget at \
iter={} cost={:.6e} |g|={:.6e}; resuming from last \
rho + trust_radius={:?}, inner-PIRLS uncapped \
(objective caches wiped; operator-TR Cauchy/Newton \
state is not resumable)",
result.iterations,
result.final_value,
cur_grad_norm,
next_trust_radius,
);
let cap_feedback = active_config.outer_inner_cap.clone();
let mut next = active_config.clone();
prev_attempt_grad_norm = Some(cur_grad_norm);
next.initial_rho = Some(result.rho.clone());
next.operator_initial_trust_radius = next_trust_radius;
retry_config = Some(next);
arc_retries_left -= 1;
obj.reset();
if let Some(feedback) = cap_feedback.as_ref() {
feedback.cap.store(0, Ordering::Relaxed);
}
}
Err(e) => break Err(e),
}
};
match outcome {
Ok(result) => {
if result.converged || attempt_idx + 1 == attempts.len() {
if !result.converged {
log::warn!(
"[OUTER warning] {context}: final outer attempt returned without convergence \
(plan={the_plan}, iterations={}, final_value={:.6e}, |g|={})",
result.iterations,
result.final_value,
result.final_grad_norm_report(),
);
}
return Ok(result);
}
let message = format!(
"{context}: attempt {} (plan={the_plan}) exhausted without convergence",
attempt_idx + 1
);
log::debug!("[OUTER] {message}; trying degraded fallback plan");
last_error = Some(EstimationError::RemlOptimizationFailed(message));
}
Err(e) => {
log::debug!(
"[OUTER] {context}: attempt {} (plan={the_plan}) failed: {e}",
attempt_idx + 1
);
last_error = Some(e);
}
}
}
Err(last_error.unwrap_or_else(|| {
EstimationError::RemlOptimizationFailed(format!("all plan attempts exhausted ({context})"))
}))
}
pub fn is_per_atom_efs_frontier(cap: &OuterCapability) -> bool {
crate::solver::estimate::reml::per_atom_efs::per_atom_efs_eligible(cap)
}
pub(crate) fn run_per_atom_efs_if_frontier(
obj: &mut dyn OuterObjective,
config: &OuterConfig,
context: &str,
) -> Result<Option<OuterResult>, EstimationError> {
let cap = primary_capability_for_config(obj.capability(), config, context);
cap.validate_layout(context)?;
if !is_per_atom_efs_frontier(&cap) {
return Ok(None);
}
let the_plan = plan(&cap);
let rho_dim = cap.theta_layout().rho_dim();
let (lower, upper) = outer_bounds_template(config, cap.n_params);
let seed = match config.initial_rho.as_ref() {
Some(initial) if initial.len() == cap.n_params => initial.clone(),
_ => {
let generated = crate::seeding::generate_rho_candidates(
cap.n_params,
config.heuristic_lambdas.as_deref(),
&config.seed_config,
);
match generated.into_iter().next() {
Some(first) => first,
None => Array1::<f64>::zeros(cap.n_params),
}
}
};
log::info!(
"[OUTER] {context}: frontier ρ-scaling (rho_dim={rho_dim}) → per-atom decoupled EFS primary"
);
let pa_cfg = crate::solver::estimate::reml::per_atom_efs::PerAtomEfsConfig::new(
config.tolerance,
config.max_iter,
lower,
upper,
);
let topology =
crate::solver::estimate::reml::per_atom_efs::SharedBorderTopology::disjoint(rho_dim);
obj.reset();
let result = crate::solver::estimate::reml::per_atom_efs::run_per_atom_efs(
obj, &seed, &pa_cfg, &topology,
)?;
Ok(Some(result.into_outer_result(the_plan)))
}
pub(crate) fn outer_bounds(lo: &Array1<f64>, hi: &Array1<f64>) -> Result<Bounds, EstimationError> {
Bounds::new(lo.clone(), hi.clone(), 1e-6).map_err(|err| {
EstimationError::InvalidInput(format!("outer rho bounds are invalid: {err}"))
})
}
pub(crate) fn outer_bounds_template(config: &OuterConfig, n: usize) -> (Array1<f64>, Array1<f64>) {
config.bounds.clone().unwrap_or_else(|| {
(
Array1::<f64>::from_elem(n, -config.rho_bound),
Array1::<f64>::from_elem(n, config.rho_bound),
)
})
}
pub(crate) fn outer_tolerance(value: f64) -> Result<Tolerance, EstimationError> {
Tolerance::new(value)
.map_err(|err| EstimationError::InvalidInput(format!("outer tolerance is invalid: {err}")))
}
pub(crate) fn outer_gradient_tolerance(config: &OuterConfig) -> GradientTolerance {
let abs = config
.objective_scale
.map(|scale| config.tolerance.max(scale * 1.0e-9))
.unwrap_or(config.tolerance);
GradientTolerance {
abs,
rel_initial_grad: None,
rel_cost: Some(config.rel_cost_tolerance.unwrap_or(config.tolerance)),
projected: true,
}
}
pub(crate) fn outer_max_iterations(value: usize) -> Result<MaxIterations, EstimationError> {
MaxIterations::new(value)
.map_err(|err| EstimationError::InvalidInput(format!("outer max_iter is invalid: {err}")))
}
pub(crate) fn sanitized_operator_trust_restart_radius(radius: Option<f64>) -> Option<f64> {
radius
.filter(|value| value.is_finite() && *value > 0.0)
.map(|value| value.max(OPERATOR_TRUST_RESTART_RADIUS_FLOOR))
}
pub(crate) fn bfgs_axis_step_caps(
config: &OuterConfig,
layout: OuterThetaLayout,
) -> Option<Array1<f64>> {
if config.bfgs_step_cap.is_none() && config.bfgs_step_cap_psi.is_none() {
return None;
}
let mut caps = Array1::from_elem(layout.n_params, f64::INFINITY);
if let Some(cap) = config.bfgs_step_cap {
for i in 0..layout.rho_dim() {
caps[i] = cap;
}
}
if let Some(cap) = config.bfgs_step_cap_psi {
for i in layout.rho_dim()..layout.n_params {
caps[i] = cap;
}
}
Some(caps)
}
pub(crate) enum FixedPointOuterRunError {
SeedRejected(EstimationError),
ImmediateFallback(EstimationError),
Failed(EstimationError),
}
pub(crate) fn run_fixed_point_outer_solver(
obj: &mut dyn OuterObjective,
layout: OuterThetaLayout,
barrier_config: Option<BarrierConfig>,
config: &OuterConfig,
context: &str,
seed: &Array1<f64>,
the_plan: OuterPlan,
label: &str,
failure_prefix: &str,
) -> Result<OuterResult, FixedPointOuterRunError> {
let mut objective = OuterFixedPointBridge {
obj,
layout,
barrier_config,
fixed_point_tolerance: config.tolerance,
consecutive_psi_zero_iters: 0,
};
match objective.eval_step(seed) {
Ok(_) => {}
Err(err) => {
let err = match err {
ObjectiveEvalError::Recoverable { message }
| ObjectiveEvalError::Fatal { message } => {
EstimationError::RemlOptimizationFailed(message)
}
};
if requests_immediate_first_order_fallback(&err.to_string()) {
return Err(FixedPointOuterRunError::ImmediateFallback(err));
}
return Err(FixedPointOuterRunError::SeedRejected(err));
}
};
let (lo, hi) = outer_bounds_template(config, layout.n_params);
let bounds = outer_bounds(&lo, &hi).map_err(FixedPointOuterRunError::Failed)?;
let tol = outer_tolerance(config.tolerance).map_err(FixedPointOuterRunError::Failed)?;
let max_iter =
outer_max_iterations(config.max_iter).map_err(FixedPointOuterRunError::Failed)?;
let mut optimizer = FixedPoint::new(seed.clone(), objective)
.with_bounds(bounds)
.with_tolerance(tol)
.with_max_iterations(max_iter);
match optimizer.run() {
Ok(sol) => Ok(solution_into_outer_result(sol, true, the_plan)),
Err(FixedPointError::MaxIterationsReached { last_solution }) => {
log::warn!(
"[OUTER warning] {context}: {label} hit max_iter={} at final_value={:.6e} step_norm={:.3e}",
config.max_iter,
last_solution.final_value,
last_solution.final_gradient_norm.unwrap_or(f64::NAN),
);
Ok(solution_into_outer_result(*last_solution, false, the_plan))
}
Err(e) => Err(FixedPointOuterRunError::Failed(
EstimationError::RemlOptimizationFailed(format!("{failure_prefix}: {e:?}")),
)),
}
}