use ndarray::{Array1, Array2, ArrayView1};
use crate::encode::EncodeAtlas;
use crate::manifold::{SaeManifoldAtom, SaeManifoldTerm};
use gam_problem::{FisherFactorKind, MetricProvenance, RowMetric};
const STEER_VALIDITY_STEPS: usize = 64;
const VALIDITY_DIVERGENCE_FRACTION: f64 = 0.1;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FisherDoseKind {
Unavailable,
ExactFull,
CertifiedPsdLowerBound,
UncertifiedApproximation,
}
impl FisherDoseKind {
pub const fn as_str(self) -> &'static str {
match self {
Self::Unavailable => "unavailable",
Self::ExactFull => "exact_full",
Self::CertifiedPsdLowerBound => "certified_psd_lower_bound",
Self::UncertifiedApproximation => "uncertified_approximation",
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct SteerPlan {
pub atom: usize,
pub atom_name: String,
pub t_from: Vec<f64>,
pub t_to: Vec<f64>,
pub amplitude: f64,
pub metric_row: usize,
pub delta: Array1<f64>,
pub predicted_nats: Option<f64>,
pub predicted_nats_kind: FisherDoseKind,
pub fisher_mass_captured: Option<f64>,
pub fisher_mass_residual: Option<f64>,
pub fisher_mass_residual_fraction: Option<f64>,
pub validity_radius: Option<f64>,
pub off_manifold_norm: f64,
pub metric_provenance: MetricProvenance,
}
#[derive(Clone, Debug)]
pub struct CoordinateSetResult {
pub edited: Array1<f64>,
pub t_from_certified: Array1<f64>,
pub encode_certificate: crate::encode::RowCertificate,
pub steer: SteerPlan,
}
pub fn set_coordinate(
model: &SaeManifoldTerm,
metric: &RowMetric,
atlas: &EncodeAtlas,
x: ArrayView1<'_, f64>,
atom_k: usize,
metric_row: usize,
amplitude: f64,
t_to: &[f64],
) -> Result<CoordinateSetResult, String> {
let atom = model.atoms.get(atom_k).ok_or_else(|| {
format!(
"set_coordinate: atom index {atom_k} out of range (term has {} atoms)",
model.k_atoms()
)
})?;
if x.len() != atom.output_dim() {
return Err(format!(
"set_coordinate: input row has length {} but atom {atom_k} output_dim is {}",
x.len(),
atom.output_dim()
));
}
let (t_from, cert) = atlas.certified_encode_row(atom, atom_k, x, amplitude)?;
let steer = steer_delta(
model,
metric,
atom_k,
metric_row,
amplitude,
t_from.as_slice().unwrap_or(&[]),
t_to,
)?;
let mut edited = x.to_owned();
if edited.len() != steer.delta.len() {
return Err(format!(
"set_coordinate: steering delta length {} does not match row length {}",
steer.delta.len(),
edited.len()
));
}
for i in 0..edited.len() {
edited[i] += steer.delta[i];
}
Ok(CoordinateSetResult {
edited,
t_from_certified: t_from,
encode_certificate: cert,
steer,
})
}
#[derive(Clone, Debug)]
pub struct InterchangeResult {
pub edited_target: Array1<f64>,
pub donor_t: Array1<f64>,
pub target_t_before: Array1<f64>,
pub target_t_after: Array1<f64>,
pub predicted_nats: Option<f64>,
pub off_manifold_norm: f64,
pub validity_radius: Option<f64>,
pub landing_error: f64,
pub set_result: CoordinateSetResult,
}
pub fn interchange(
model: &SaeManifoldTerm,
metric: &RowMetric,
atlas: &EncodeAtlas,
x_target: ArrayView1<'_, f64>,
target_amplitude: f64,
x_source: ArrayView1<'_, f64>,
source_amplitude: f64,
atom_k: usize,
target_metric_row: usize,
) -> Result<InterchangeResult, String> {
let atom = model.atoms.get(atom_k).ok_or_else(|| {
format!(
"interchange: atom index {atom_k} out of range (term has {} atoms)",
model.k_atoms()
)
})?;
let (donor_t, _donor_cert) =
atlas.certified_encode_row(atom, atom_k, x_source, source_amplitude)?;
let set = set_coordinate(
model,
metric,
atlas,
x_target,
atom_k,
target_metric_row,
target_amplitude,
donor_t.as_slice().unwrap_or(&[]),
)?;
let (target_t_after, _after_cert) =
atlas.certified_encode_row(atom, atom_k, set.edited.view(), target_amplitude)?;
let periods = model.assignment.coords[atom_k].effective_axis_periods();
let landing_error = coordinate_l2_distance(
donor_t.as_slice().unwrap_or(&[]),
target_t_after.as_slice().unwrap_or(&[]),
&periods,
)?;
Ok(InterchangeResult {
edited_target: set.edited.clone(),
donor_t,
target_t_before: set.t_from_certified.clone(),
target_t_after,
predicted_nats: set.steer.predicted_nats,
off_manifold_norm: set.steer.off_manifold_norm,
validity_radius: set.steer.validity_radius,
landing_error,
set_result: set,
})
}
fn shortest_coordinate_delta(
from: &[f64],
to: &[f64],
periods: &[Option<f64>],
) -> Result<Vec<f64>, String> {
if from.len() != to.len() || from.len() != periods.len() {
return Err(format!(
"coordinate displacement length mismatch: from={}, to={}, periods={}",
from.len(),
to.len(),
periods.len()
));
}
let mut delta = Vec::with_capacity(from.len());
for axis in 0..from.len() {
let mut d = to[axis] - from[axis];
if let Some(period) = periods[axis] {
if !(period.is_finite() && period > 0.0) {
return Err(format!(
"coordinate axis {axis} has invalid period {period}"
));
}
d -= period * (d / period).round();
}
delta.push(d);
}
Ok(delta)
}
fn coordinate_l2_distance(a: &[f64], b: &[f64], periods: &[Option<f64>]) -> Result<f64, String> {
Ok(shortest_coordinate_delta(a, b, periods)?
.iter()
.map(|d| d * d)
.sum::<f64>()
.sqrt())
}
fn path_coordinate(
from: &[f64],
delta: &[f64],
periods: &[Option<f64>],
fraction: f64,
) -> Vec<f64> {
from.iter()
.zip(delta.iter())
.zip(periods.iter())
.map(|((&start, &step), &period)| {
let value = start + fraction * step;
period.map_or(value, |p| value.rem_euclid(p))
})
.collect()
}
pub fn steer_delta(
model: &SaeManifoldTerm,
metric: &RowMetric,
atom_k: usize,
metric_row: usize,
amplitude: f64,
t_from: &[f64],
t_to: &[f64],
) -> Result<SteerPlan, String> {
if !(amplitude.is_finite() && amplitude > 0.0) {
return Err(format!(
"steer_delta: amplitude must be finite and positive, got {amplitude}"
));
}
let k = model.k_atoms();
if atom_k >= k {
return Err(format!(
"steer_delta: atom index {atom_k} out of range (term has {k} atoms)"
));
}
let atom = &model.atoms[atom_k];
let d = atom.latent_dim();
let p = atom.output_dim();
if t_from.len() != d || t_to.len() != d {
return Err(format!(
"steer_delta: t_from/t_to must have length latent_dim={d}; got {} and {}",
t_from.len(),
t_to.len()
));
}
atom.basis_evaluator.as_ref().ok_or_else(|| {
format!(
"steer_delta: atom {atom_k} ('{}') has no installed basis evaluator; \
arbitrary-t decoder evaluation requires one",
atom.name
)
})?;
let periods = model.assignment.coords[atom_k].effective_axis_periods();
let coordinate_delta = shortest_coordinate_delta(t_from, t_to, &periods)?;
let n = model.n_obs();
if metric.n_rows() != n || metric.p_out() != p {
return Err(format!(
"steer_delta: metric shape ({}, {}) must equal fitted term shape ({n}, {p})",
metric.n_rows(),
metric.p_out()
));
}
if metric_row >= n {
return Err(format!(
"steer_delta: metric_row={metric_row} out of range for {n} fitted rows"
));
}
let tier0_scale = model.tier0_scale();
let g_from = decode_at(atom, t_from, tier0_scale)?;
let g_to = decode_at(atom, t_to, tier0_scale)?;
let mut delta = Array1::<f64>::zeros(p);
for i in 0..p {
delta[i] = amplitude * (g_to[i] - g_from[i]);
}
let provenance = metric.provenance();
let behavior_available = metric_carries_behavior(provenance);
let fisher_mass_captured = behavior_available.then(|| metric.row_traces()[metric_row]);
let fisher_mass_residual = behavior_available
.then(|| metric.truncation_mass_residual(metric_row))
.flatten();
let fisher_mass_residual_fraction = behavior_available
.then(|| metric.truncation_mass_residual_fraction(metric_row))
.flatten();
let predicted_nats_kind = if !behavior_available {
FisherDoseKind::Unavailable
} else {
match metric.fisher_factor_kind() {
Some(FisherFactorKind::ExactFull) => FisherDoseKind::ExactFull,
Some(FisherFactorKind::CertifiedPsdLowerBound) => {
FisherDoseKind::CertifiedPsdLowerBound
}
Some(FisherFactorKind::UncertifiedApproximation) => {
FisherDoseKind::UncertifiedApproximation
}
None => {
return Err(format!(
"steer_delta: behavioral metric provenance {provenance:?} has no explicit Fisher factor status"
));
}
}
};
let mut t_mid = vec![0.0_f64; d];
for a in 0..d {
t_mid[a] = t_from[a] + 0.5 * coordinate_delta[a];
if let Some(period) = periods[a] {
t_mid[a] = t_mid[a].rem_euclid(period);
}
}
let tangents = decode_tangents_at(atom, &t_mid, tier0_scale)?;
let off_manifold_norm = off_manifold_residual_norm(&tangents, delta.view());
let (predicted_nats, validity_radius) = if !behavior_available {
(None, None)
} else {
let ctx = SteerContext {
atom,
scale: tier0_scale,
metric,
row: metric_row,
p,
d,
amplitude,
coordinate_delta: &coordinate_delta,
periods: &periods,
};
let dose = 0.5 * metric.fisher_mass(metric_row, delta.view());
let radius = validity_radius(&ctx, t_from)?;
(Some(dose), Some(radius))
};
Ok(SteerPlan {
atom: atom_k,
atom_name: atom.name.clone(),
t_from: t_from.to_vec(),
t_to: t_to.to_vec(),
amplitude,
metric_row,
delta,
predicted_nats,
predicted_nats_kind,
fisher_mass_captured,
fisher_mass_residual,
fisher_mass_residual_fraction,
validity_radius,
off_manifold_norm,
metric_provenance: provenance,
})
}
pub fn predicted_response(
model: &SaeManifoldTerm,
atom_k: usize,
t_at: &[f64],
delta: ArrayView1<'_, f64>,
) -> Result<Array1<f64>, String> {
let k = model.k_atoms();
if atom_k >= k {
return Err(format!(
"predicted_response: atom index {atom_k} out of range (term has {k} atoms)"
));
}
let atom = &model.atoms[atom_k];
let d = atom.latent_dim();
let p = atom.output_dim();
if t_at.len() != d {
return Err(format!(
"predicted_response: t_at must have length latent_dim={d}; got {}",
t_at.len()
));
}
if delta.len() != p {
return Err(format!(
"predicted_response: delta must have length output_dim={p}; got {}",
delta.len()
));
}
atom.basis_evaluator.as_ref().ok_or_else(|| {
format!(
"predicted_response: atom {atom_k} ('{}') has no installed basis evaluator",
atom.name
)
})?;
let tangents = decode_tangents_at(atom, t_at, model.tier0_scale())?;
Ok(project_onto_tangent_span(&tangents, delta))
}
#[derive(Clone, Debug, PartialEq)]
pub struct AppliedDoseObservation {
pub effective_delta: Array1<f64>,
pub exact_directional_nats: f64,
pub measured_nats: f64,
pub certified_attainable_upper_nats: Option<f64>,
}
pub type AppliedDoseProbe<'a> =
dyn FnMut(&SteerPlan) -> Result<AppliedDoseObservation, String> + 'a;
#[derive(Clone, Copy, Debug)]
pub struct TargetDoseConfig {
pub tol_rel: f64,
pub max_iter: usize,
pub readout_tol_rel: f64,
}
impl Default for TargetDoseConfig {
fn default() -> Self {
Self {
tol_rel: 1.0e-2,
max_iter: 12,
readout_tol_rel: 1.0e-1,
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct TargetDoseRequest<'a> {
pub atom_k: usize,
pub metric_row: usize,
pub t_from: &'a [f64],
pub t_to: &'a [f64],
pub target_nats: f64,
pub config: TargetDoseConfig,
}
#[derive(Clone, Debug)]
pub struct TargetDosePlan {
pub target_nats: f64,
pub seed_amplitude: f64,
pub steer: SteerPlan,
pub applied_probe: Option<AppliedDoseObservation>,
pub iterations: usize,
pub readout_kl_radius: Option<f64>,
pub certified_attainable_upper_nats: Option<f64>,
}
#[derive(Clone, Debug, PartialEq)]
pub enum TargetDoseError {
InvalidRequest(String),
Steering(String),
Probe(String),
FactorNeedsAppliedDoseProbe {
kind: FisherDoseKind,
},
UnreachableTarget {
target_nats: f64,
certified_attainable_upper_nats: f64,
},
UnbracketedTarget {
target_nats: f64,
max_probed_amplitude: f64,
max_measured_nats: f64,
probes: usize,
},
BracketResolutionExhausted {
target_nats: f64,
lower_amplitude: f64,
lower_nats: f64,
upper_amplitude: f64,
upper_nats: f64,
probes: usize,
},
}
impl std::fmt::Display for TargetDoseError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::InvalidRequest(message) | Self::Steering(message) | Self::Probe(message) => {
f.write_str(message)
}
Self::FactorNeedsAppliedDoseProbe { kind } => write!(
f,
"steer_to_target_nats: factor kind {} cannot solve a full-KL target without an applied-dose probe",
kind.as_str()
),
Self::UnreachableTarget {
target_nats,
certified_attainable_upper_nats,
} => write!(
f,
"steer_to_target_nats: target {target_nats} nats is outside the certified \
attainable envelope, whose global upper bound is \
{certified_attainable_upper_nats} nats"
),
Self::UnbracketedTarget {
target_nats,
max_probed_amplitude,
max_measured_nats,
probes,
} => write!(
f,
"steer_to_target_nats: could not bracket target {target_nats} nats after \
{probes} probes through amplitude {max_probed_amplitude}; the largest \
observed dose was {max_measured_nats} nats and no global attainable \
envelope certified the target unreachable"
),
Self::BracketResolutionExhausted {
target_nats,
lower_amplitude,
lower_nats,
upper_amplitude,
upper_nats,
probes,
} => write!(
f,
"steer_to_target_nats: exhausted {probes} probes before resolving target \
{target_nats} nats inside measured bracket \
({lower_amplitude}, {lower_nats})..({upper_amplitude}, {upper_nats})"
),
}
}
}
impl std::error::Error for TargetDoseError {}
fn probe_applied_dose(
probe: &mut AppliedDoseProbe<'_>,
plan: &SteerPlan,
) -> Result<AppliedDoseObservation, TargetDoseError> {
let observation = probe(plan).map_err(TargetDoseError::Probe)?;
if observation.effective_delta.len() != plan.delta.len() {
return Err(TargetDoseError::Probe(format!(
"steer_to_target_nats: probe effective_delta length {} does not match plan delta length {}",
observation.effective_delta.len(),
plan.delta.len()
)));
}
if !observation
.effective_delta
.iter()
.all(|value| value.is_finite())
{
return Err(TargetDoseError::Probe(
"steer_to_target_nats: probe effective_delta must be finite".to_string(),
));
}
if !(observation.exact_directional_nats.is_finite()
&& observation.exact_directional_nats >= 0.0)
{
return Err(TargetDoseError::Probe(format!(
"steer_to_target_nats: probe exact_directional_nats must be finite and non-negative; got {}",
observation.exact_directional_nats
)));
}
if !(observation.measured_nats.is_finite() && observation.measured_nats >= 0.0) {
return Err(TargetDoseError::Probe(format!(
"steer_to_target_nats: probe measured_nats must be finite and non-negative; got {}",
observation.measured_nats
)));
}
if let Some(upper) = observation.certified_attainable_upper_nats {
if !(upper.is_finite() && upper >= 0.0) {
return Err(TargetDoseError::Probe(format!(
"steer_to_target_nats: probe certified_attainable_upper_nats must be finite \
and non-negative when present; got {upper}"
)));
}
if observation.measured_nats > upper {
return Err(TargetDoseError::Probe(format!(
"steer_to_target_nats: measured dose {} exceeds the probe's certified \
global attainable upper bound {upper}",
observation.measured_nats
)));
}
}
Ok(observation)
}
fn merge_attainable_envelope(
observation: &AppliedDoseObservation,
max_measured_nats: f64,
envelope: &mut Option<f64>,
) -> Result<(), TargetDoseError> {
if let Some(upper) = observation.certified_attainable_upper_nats {
*envelope = Some(envelope.map_or(upper, |current| current.min(upper)));
}
if let Some(upper) = *envelope
&& max_measured_nats > upper
{
return Err(TargetDoseError::Probe(format!(
"steer_to_target_nats: observed dose {max_measured_nats} exceeds an earlier \
certified global attainable upper bound {upper}"
)));
}
Ok(())
}
fn record_readout_probe(
amplitude: f64,
measured: f64,
predicted: f64,
tolerance: f64,
first_failure: &mut Option<f64>,
radius: &mut Option<f64>,
) {
let agrees = predicted > 0.0 && (measured - predicted).abs() / predicted <= tolerance;
if agrees {
if (*first_failure).is_none_or(|failed| amplitude < failed) {
*radius = Some((*radius).map_or(amplitude, |current| current.max(amplitude)));
}
} else {
*first_failure = Some((*first_failure).map_or(amplitude, |failed| failed.min(amplitude)));
if (*radius).is_some_and(|current| current >= amplitude) {
*radius = None;
}
}
}
pub fn steer_to_target_nats(
model: &SaeManifoldTerm,
metric: &RowMetric,
request: TargetDoseRequest<'_>,
probe: Option<&mut AppliedDoseProbe<'_>>,
) -> Result<TargetDosePlan, TargetDoseError> {
let TargetDoseRequest {
atom_k,
metric_row,
t_from,
t_to,
target_nats,
config,
} = request;
if !(target_nats.is_finite() && target_nats > 0.0) {
return Err(TargetDoseError::InvalidRequest(format!(
"steer_to_target_nats: target_nats must be finite and positive, got {target_nats}"
)));
}
if !(config.tol_rel.is_finite() && (0.0..1.0).contains(&config.tol_rel))
|| config.max_iter == 0
|| !(config.readout_tol_rel.is_finite() && (0.0..1.0).contains(&config.readout_tol_rel))
{
return Err(TargetDoseError::InvalidRequest(format!(
"steer_to_target_nats: config must have finite 0<=tol_rel<1, max_iter>0, \
finite 0<=readout_tol_rel<1; \
got {config:?}"
)));
}
let unit = steer_delta(model, metric, atom_k, metric_row, 1.0, t_from, t_to)
.map_err(TargetDoseError::Steering)?;
let unit_nats = unit.predicted_nats.ok_or_else(|| {
TargetDoseError::InvalidRequest(format!(
"steer_to_target_nats: atom {atom_k} has no behavioral (nats) metric \
(provenance {:?}); a target-nats dose is undefined",
unit.metric_provenance
))
})?;
if !(unit_nats.is_finite() && unit_nats > 0.0) {
return Err(TargetDoseError::InvalidRequest(format!(
"steer_to_target_nats: unit-amplitude dose must be finite and positive; got \
{unit_nats} in metric row {metric_row}"
)));
}
if probe.is_none() && unit.predicted_nats_kind != FisherDoseKind::ExactFull {
return Err(TargetDoseError::FactorNeedsAppliedDoseProbe {
kind: unit.predicted_nats_kind,
});
}
let seed_amplitude = (target_nats / unit_nats).sqrt();
if !(seed_amplitude.is_finite() && seed_amplitude > 0.0) {
return Err(TargetDoseError::InvalidRequest(format!(
"steer_to_target_nats: target {target_nats} nats and unit dose {unit_nats} \
imply an unrepresentable amplitude {seed_amplitude}"
)));
}
let plan_at = |amplitude: f64| {
steer_delta(model, metric, atom_k, metric_row, amplitude, t_from, t_to)
.map_err(TargetDoseError::Steering)
};
let finish = |steer: SteerPlan,
applied_probe: Option<AppliedDoseObservation>,
iterations: usize,
readout_kl_radius: Option<f64>,
certified_attainable_upper_nats: Option<f64>|
-> Result<TargetDosePlan, TargetDoseError> {
Ok(TargetDosePlan {
target_nats,
seed_amplitude,
steer,
applied_probe,
iterations,
readout_kl_radius,
certified_attainable_upper_nats,
})
};
let probe = match probe {
Some(probe) => probe,
None => return finish(plan_at(seed_amplitude)?, None, 0, None, None),
};
let mut first_readout_failure: Option<f64> = None;
let mut readout_kl_radius: Option<f64> = None;
let mut probes = 0usize;
let mut lo_a = 0.0_f64;
let mut lo_kl = 0.0_f64;
let mut hi_a = seed_amplitude;
let hi_plan = plan_at(hi_a)?;
let hi_probe = probe_applied_dose(probe, &hi_plan)?;
let mut hi_kl = hi_probe.measured_nats;
probes += 1;
let mut max_probed_amplitude = hi_a;
let mut max_measured_nats = hi_kl;
let mut certified_attainable_upper_nats = None;
merge_attainable_envelope(
&hi_probe,
max_measured_nats,
&mut certified_attainable_upper_nats,
)?;
record_readout_probe(
hi_a,
hi_kl,
hi_probe.exact_directional_nats,
config.readout_tol_rel,
&mut first_readout_failure,
&mut readout_kl_radius,
);
if (hi_kl - target_nats).abs() / target_nats <= config.tol_rel {
return finish(
hi_plan,
Some(hi_probe),
probes,
readout_kl_radius,
certified_attainable_upper_nats,
);
}
let accepted_lower_nats = target_nats * (1.0 - config.tol_rel);
if let Some(upper) = certified_attainable_upper_nats
&& upper < accepted_lower_nats
{
return Err(TargetDoseError::UnreachableTarget {
target_nats,
certified_attainable_upper_nats: upper,
});
}
while hi_kl < target_nats {
if probes >= config.max_iter {
return Err(TargetDoseError::UnbracketedTarget {
target_nats,
max_probed_amplitude,
max_measured_nats,
probes,
});
}
let next_a = hi_a * 2.0;
if !(next_a.is_finite() && next_a > hi_a) {
return Err(TargetDoseError::UnbracketedTarget {
target_nats,
max_probed_amplitude,
max_measured_nats,
probes,
});
}
let next_plan = plan_at(next_a)?;
let next_probe = probe_applied_dose(probe, &next_plan)?;
let next_kl = next_probe.measured_nats;
probes += 1;
max_probed_amplitude = next_a;
max_measured_nats = max_measured_nats.max(next_kl);
merge_attainable_envelope(
&next_probe,
max_measured_nats,
&mut certified_attainable_upper_nats,
)?;
record_readout_probe(
next_a,
next_kl,
next_probe.exact_directional_nats,
config.readout_tol_rel,
&mut first_readout_failure,
&mut readout_kl_radius,
);
if (next_kl - target_nats).abs() / target_nats <= config.tol_rel {
return finish(
next_plan,
Some(next_probe),
probes,
readout_kl_radius,
certified_attainable_upper_nats,
);
}
if let Some(upper) = certified_attainable_upper_nats
&& upper < accepted_lower_nats
{
return Err(TargetDoseError::UnreachableTarget {
target_nats,
certified_attainable_upper_nats: upper,
});
}
lo_a = hi_a;
lo_kl = hi_kl;
hi_a = next_a;
hi_kl = next_kl;
}
while probes < config.max_iter {
let denominator = hi_kl - lo_kl;
let secant = hi_a - (hi_kl - target_nats) * (hi_a - lo_a) / denominator;
let candidate = if secant.is_finite() && secant > lo_a && secant < hi_a {
secant
} else {
0.5 * (lo_a + hi_a)
};
let candidate_plan = plan_at(candidate)?;
let candidate_probe = probe_applied_dose(probe, &candidate_plan)?;
let measured = candidate_probe.measured_nats;
probes += 1;
max_measured_nats = max_measured_nats.max(measured);
merge_attainable_envelope(
&candidate_probe,
max_measured_nats,
&mut certified_attainable_upper_nats,
)?;
record_readout_probe(
candidate,
measured,
candidate_probe.exact_directional_nats,
config.readout_tol_rel,
&mut first_readout_failure,
&mut readout_kl_radius,
);
if (measured - target_nats).abs() / target_nats <= config.tol_rel {
return finish(
candidate_plan,
Some(candidate_probe),
probes,
readout_kl_radius,
certified_attainable_upper_nats,
);
}
if measured < target_nats {
lo_a = candidate;
lo_kl = measured;
} else {
hi_a = candidate;
hi_kl = measured;
}
}
Err(TargetDoseError::BracketResolutionExhausted {
target_nats,
lower_amplitude: lo_a,
lower_nats: lo_kl,
upper_amplitude: hi_a,
upper_nats: hi_kl,
probes,
})
}
fn metric_carries_behavior(p: MetricProvenance) -> bool {
match p {
MetricProvenance::Euclidean | MetricProvenance::WhitenedStructured { .. } => false,
MetricProvenance::OutputFisher { .. }
| MetricProvenance::OutputFisherDownstream { .. }
| MetricProvenance::BehavioralFisher { .. } => true,
}
}
fn decode_at(
atom: &SaeManifoldAtom,
t: &[f64],
scale: Option<&Array1<f64>>,
) -> Result<Array1<f64>, String> {
let d = t.len();
let coords = Array2::from_shape_vec((1, d), t.to_vec())
.map_err(|e| format!("steer_delta::decode_at: coord shape: {e}"))?;
let mut out = atom.decode_at_coords(coords.view())?.row(0).to_owned();
if let Some(scale) = scale {
if scale.len() != out.len() {
return Err(format!(
"steer_delta::decode_at: tier0 scale length {} != output_dim {}",
scale.len(),
out.len()
));
}
for (v, &s) in out.iter_mut().zip(scale.iter()) {
*v *= s;
}
}
Ok(out)
}
fn decode_tangents_at(
atom: &SaeManifoldAtom,
t: &[f64],
scale: Option<&Array1<f64>>,
) -> Result<Array2<f64>, String> {
let evaluator = atom.basis_evaluator.as_ref().ok_or_else(|| {
"steer_delta::decode_tangents_at: atom has no installed basis evaluator".to_string()
})?;
let p = atom.output_dim();
let d = atom.latent_dim();
let coords = Array2::from_shape_vec((1, d), t.to_vec())
.map_err(|e| format!("steer_delta::decode_tangents_at: coord shape: {e}"))?;
let jet = if atom.homotopy_eta == 1.0 {
evaluator.evaluate(coords.view())?.1
} else {
evaluator
.evaluate_phi_eta(coords.view(), atom.homotopy_eta)?
.jet
};
let decoder = atom.decoder_coefficients();
let m = decoder.nrows();
if jet.dim() != (1, m, d) {
return Err(format!(
"steer_delta::decode_tangents_at: evaluator jet {:?} != (1, {m}, {d})",
jet.dim()
));
}
let mut tang = Array2::<f64>::zeros((p, d));
for axis in 0..d {
for basis_col in 0..m {
let dphi = jet[[0, basis_col, axis]];
if dphi == 0.0 {
continue;
}
for out_col in 0..p {
tang[[out_col, axis]] += dphi * decoder[[basis_col, out_col]];
}
}
}
if let Some(scale) = scale {
if scale.len() != p {
return Err(format!(
"steer_delta::decode_tangents_at: tier0 scale length {} != output_dim {p}",
scale.len()
));
}
for (out_col, &s) in scale.iter().enumerate() {
tang.row_mut(out_col).mapv_inplace(|v| v * s);
}
}
Ok(tang)
}
fn project_onto_tangent_span(tangents: &Array2<f64>, delta: ArrayView1<'_, f64>) -> Array1<f64> {
let p = tangents.nrows();
let d = tangents.ncols();
if d == 0 {
return Array1::<f64>::zeros(p);
}
let mut gram = Array2::<f64>::zeros((d, d));
let mut rhs = Array1::<f64>::zeros(d);
for a in 0..d {
let mut r = 0.0_f64;
for i in 0..p {
r += tangents[[i, a]] * delta[i];
}
rhs[a] = r;
for b in a..d {
let mut acc = 0.0_f64;
for i in 0..p {
acc += tangents[[i, a]] * tangents[[i, b]];
}
gram[[a, b]] = acc;
gram[[b, a]] = acc;
}
}
let trace: f64 = (0..d).map(|a| gram[[a, a]]).sum();
let jitter = if trace > 0.0 { 1e-12 * trace } else { 1e-12 };
for a in 0..d {
gram[[a, a]] += jitter;
}
let coeffs = solve_spd_small(&gram, &rhs);
let mut proj = Array1::<f64>::zeros(p);
for i in 0..p {
for a in 0..d {
proj[i] += tangents[[i, a]] * coeffs[a];
}
}
proj
}
fn off_manifold_residual_norm(tangents: &Array2<f64>, delta: ArrayView1<'_, f64>) -> f64 {
let proj = project_onto_tangent_span(tangents, delta);
let mut res_sq = 0.0_f64;
for i in 0..delta.len() {
let r = delta[i] - proj[i];
res_sq += r * r;
}
res_sq.max(0.0).sqrt()
}
fn solve_spd_small(gram: &Array2<f64>, rhs: &Array1<f64>) -> Array1<f64> {
let d = gram.nrows();
let mut l = Array2::<f64>::zeros((d, d));
for i in 0..d {
for j in 0..=i {
let mut sum = gram[[i, j]];
for k in 0..j {
sum -= l[[i, k]] * l[[j, k]];
}
if i == j {
if sum <= 0.0 {
return Array1::<f64>::zeros(d);
}
l[[i, j]] = sum.sqrt();
} else {
l[[i, j]] = sum / l[[j, j]];
}
}
}
let mut y = Array1::<f64>::zeros(d);
for i in 0..d {
let mut sum = rhs[i];
for k in 0..i {
sum -= l[[i, k]] * y[k];
}
y[i] = sum / l[[i, i]];
}
let mut x = Array1::<f64>::zeros(d);
for i in (0..d).rev() {
let mut sum = y[i];
for k in (i + 1)..d {
sum -= l[[k, i]] * x[k];
}
x[i] = sum / l[[i, i]];
}
x
}
struct SteerContext<'a> {
atom: &'a SaeManifoldAtom,
scale: Option<&'a Array1<f64>>,
metric: &'a RowMetric,
row: usize,
p: usize,
d: usize,
amplitude: f64,
coordinate_delta: &'a [f64],
periods: &'a [Option<f64>],
}
fn validity_radius(ctx: &SteerContext<'_>, t_from: &[f64]) -> Result<f64, String> {
let d = ctx.d;
let p = ctx.p;
let full_len: f64 = ctx
.coordinate_delta
.iter()
.map(|d| d * d)
.sum::<f64>()
.sqrt();
if full_len == 0.0 {
return Ok(0.0);
}
let dt = ctx.coordinate_delta;
let amp = ctx.amplitude;
let tang0 = decode_tangents_at(ctx.atom, t_from, ctx.scale)?;
let mut v0 = Array1::<f64>::zeros(p);
for i in 0..p {
let mut acc = 0.0_f64;
for a in 0..d {
acc += tang0[[i, a]] * dt[a];
}
v0[i] = acc;
}
let lin_coeff = 0.5 * amp * amp * ctx.metric.fisher_mass(ctx.row, v0.view());
if !(lin_coeff > 0.0) {
return Ok(full_len);
}
let g_from = decode_at(ctx.atom, t_from, ctx.scale)?;
let steps = STEER_VALIDITY_STEPS;
for s in 0..steps {
let tau = (s as f64 + 1.0) / steps as f64;
let t_mid = path_coordinate(t_from, dt, ctx.periods, tau);
let g_tau = decode_at(ctx.atom, &t_mid, ctx.scale)?;
let mut chord = Array1::<f64>::zeros(p);
for i in 0..p {
chord[i] = amp * (g_tau[i] - g_from[i]);
}
let chord_kl = 0.5 * ctx.metric.fisher_mass(ctx.row, chord.view());
let lin_kl = tau * tau * lin_coeff;
let rel = (chord_kl - lin_kl).abs() / lin_kl;
if rel > VALIDITY_DIVERGENCE_FRACTION {
return Ok(tau * full_len);
}
}
Ok(full_len)
}
#[derive(Clone, Debug, PartialEq, serde::Serialize)]
pub struct CollateralPoint {
pub dose: f64,
pub on_target_effect: f64,
pub collateral: f64,
pub cross_feature: f64,
}
#[derive(Clone, Debug, PartialEq, serde::Serialize)]
pub struct CollateralArm {
pub points: Vec<CollateralPoint>,
pub efficiency: f64,
}
#[derive(Clone, Debug, PartialEq, serde::Serialize)]
pub struct CollateralCurve {
pub atom: usize,
pub axis: usize,
pub others: Vec<usize>,
pub manifold: CollateralArm,
pub flat: CollateralArm,
pub manifold_is_cleaner: bool,
}
fn frame_landed_norm(frame: &Array2<f64>, delta: ArrayView1<'_, f64>) -> f64 {
let proj = project_onto_tangent_span(frame, delta);
proj.iter().map(|&x| x * x).sum::<f64>().sqrt()
}
pub fn collateral_curve(
model: &SaeManifoldTerm,
atom_k: usize,
axis: usize,
others: &[usize],
doses: &[f64],
) -> Result<CollateralCurve, String> {
let k = model.k_atoms();
if atom_k >= k {
return Err(format!(
"collateral_curve: atom index {atom_k} out of range (term has {k} atoms)"
));
}
let d_k = model.atoms[atom_k].latent_dim();
if axis >= d_k {
return Err(format!(
"collateral_curve: axis {axis} out of range for atom {atom_k} latent_dim {d_k}"
));
}
if doses.is_empty() {
return Err("collateral_curve: doses must be non-empty".to_string());
}
for &j in others {
if j >= k {
return Err(format!(
"collateral_curve: other atom index {j} out of range (term has {k} atoms)"
));
}
}
let n = model.n_obs();
let p = model.output_dim();
let rows: Vec<usize> = (0..n).collect();
let frame_at = |atom_idx: usize| -> Result<Vec<Array2<f64>>, String> {
let coords = model.assignment.coords[atom_idx].as_matrix();
let mut frames = Vec::with_capacity(n);
for row in 0..n {
let t: Vec<f64> = coords.row(row).to_vec();
frames.push(decode_tangents_at(
&model.atoms[atom_idx],
&t,
model.tier0_scale(),
)?);
}
Ok(frames)
};
let target_frames = frame_at(atom_k)?;
let mut other_frames: Vec<Vec<Array2<f64>>> = Vec::with_capacity(others.len());
for &j in others {
other_frames.push(frame_at(j)?);
}
let mut gram = Array2::<f64>::zeros((p, p));
for frame in &target_frames {
for i in 0..p {
let gi = frame[[i, axis]];
if gi == 0.0 {
continue;
}
for j in 0..p {
gram[[i, j]] += gi * frame[[j, axis]];
}
}
}
let mut w = Array1::<f64>::from_elem(p, 1.0 / (p as f64).sqrt());
for _ in 0..128 {
let mut next = Array1::<f64>::zeros(p);
for i in 0..p {
let mut acc = 0.0_f64;
for j in 0..p {
acc += gram[[i, j]] * w[j];
}
next[i] = acc;
}
let norm = next.iter().map(|&x| x * x).sum::<f64>().sqrt();
if !(norm > 0.0) {
return Err(format!(
"collateral_curve: atom {atom_k} has a vanishing tangent field along axis {axis}; \
no fixed direction to define the flat control"
));
}
next.mapv_inplace(|x| x / norm);
w = next;
}
let decompose = |field: &Array2<f64>| -> CollateralPoint {
let mut eff_sq = 0.0_f64;
let mut col_sq = 0.0_f64;
let mut cross_sq = 0.0_f64;
for row in 0..n {
let delta = field.row(row);
let on_target = project_onto_tangent_span(&target_frames[row], delta);
let mut e = 0.0_f64;
let mut c = 0.0_f64;
for i in 0..p {
e += on_target[i] * on_target[i];
let residual = delta[i] - on_target[i];
c += residual * residual;
}
eff_sq += e;
col_sq += c;
let mut cross = 0.0_f64;
for frames in &other_frames {
let l = frame_landed_norm(&frames[row], delta);
cross += l * l;
}
cross_sq += cross;
}
let denom = n.max(1) as f64;
CollateralPoint {
dose: 0.0,
on_target_effect: (eff_sq / denom).sqrt(),
collateral: (col_sq / denom).sqrt(),
cross_feature: (cross_sq / denom).sqrt(),
}
};
let mut manifold_pts = Vec::with_capacity(doses.len());
let mut flat_pts = Vec::with_capacity(doses.len());
for &dose in doses {
let mut step = Array1::<f64>::zeros(d_k);
step[axis] = dose;
let on_field = model.steer_rows(atom_k, &rows, step.view())?;
let mut flat_field = Array2::<f64>::zeros((n, p));
for row in 0..n {
let norm = on_field.row(row).iter().map(|&x| x * x).sum::<f64>().sqrt();
for i in 0..p {
flat_field[[row, i]] = norm * w[i];
}
}
let mut m = decompose(&on_field);
m.dose = dose;
manifold_pts.push(m);
let mut f = decompose(&flat_field);
f.dose = dose;
flat_pts.push(f);
}
let efficiency = |pts: &[CollateralPoint]| -> f64 {
let eff_sq: f64 = pts
.iter()
.map(|q| q.on_target_effect * q.on_target_effect)
.sum();
let col_sq: f64 = pts.iter().map(|q| q.collateral * q.collateral).sum();
if eff_sq > 0.0 {
(col_sq / eff_sq).sqrt()
} else {
f64::NAN
}
};
let manifold = CollateralArm {
efficiency: efficiency(&manifold_pts),
points: manifold_pts,
};
let flat = CollateralArm {
efficiency: efficiency(&flat_pts),
points: flat_pts,
};
let manifold_is_cleaner = manifold.efficiency.is_finite()
&& flat.efficiency.is_finite()
&& manifold.efficiency < flat.efficiency;
Ok(CollateralCurve {
atom: atom_k,
axis,
others: others.to_vec(),
manifold,
flat,
manifold_is_cleaner,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn periodic_steering_uses_shortest_path_across_seam() {
let periods = [Some(1.0)];
let delta = shortest_coordinate_delta(&[0.99], &[0.01], &periods).unwrap();
assert!((delta[0] - 0.02).abs() < 1e-12);
let midpoint = path_coordinate(&[0.99], &delta, &periods, 0.5);
assert!(midpoint[0].abs() < 1e-12 || (midpoint[0] - 1.0).abs() < 1e-12);
let distance = coordinate_l2_distance(&[0.99], &[0.01], &periods).unwrap();
assert!((distance - 0.02).abs() < 1e-12);
}
}