use super::least_squares_diagnostics::{
LeastSquaresCheck, LeastSquaresDiagnostics, LeastSquaresObservation, Trace,
};
#[cfg(test)]
mod accounting;
mod step;
use crate::core::constraint::evaluation::BoxProblem;
use crate::core::convergence::{
NativeConvergenceDiagnostics, NativeConvergenceTest,
};
use crate::core::inner::InitialState;
use crate::core::least_squares::evaluation::{BoundedEvaluation, NllsStep};
use crate::core::math::dense_svd::{DenseSvd, norm};
use crate::core::math::{MatrixIndex, Scalar, VectorIndex, VectorLen};
use crate::core::problem::{Jacobian, Problem, Residual};
use crate::core::solver::Solver;
use crate::core::state::{PointState, State};
use crate::core::termination::Termination;
use crate::{LossFunction, RobustLeastSquares, ScaleRowsInPlace};
use step::{Model, interior, number, update_radius};
#[derive(Clone, Debug)]
pub struct TrustRegionReflective<F: Scalar = f64> {
gradient_tolerance: Option<F>,
initial_radius: Option<F>,
rank_tolerance: Option<F>,
max_inner_attempts: usize,
max_subproblem_iterations: usize,
work: Option<Work<F>>,
native_convergence: Vec<NativeConvergenceTest>,
trace: Option<Trace<F>>,
}
impl<F: Scalar> NativeConvergenceDiagnostics for TrustRegionReflective<F> {
fn native_convergence_tests(&self) -> &[NativeConvergenceTest] {
&self.native_convergence
}
}
impl<F: Scalar> LeastSquaresDiagnostics<F> for TrustRegionReflective<F> {
fn least_squares_observations(&self) -> &[LeastSquaresObservation<F>] {
self.trace.as_ref().map_or(&[], |trace| &trace.observations)
}
}
impl<F: Scalar> Default for TrustRegionReflective<F> {
fn default() -> Self {
Self::new()
}
}
impl<F: Scalar> TrustRegionReflective<F> {
pub fn new() -> Self {
Self {
gradient_tolerance: Some(number(1e-8)),
initial_radius: None,
rank_tolerance: None,
max_inner_attempts: 50,
max_subproblem_iterations: 50,
work: None,
native_convergence: Vec::new(),
trace: None,
}
}
pub fn with_trial_diagnostics(mut self, enabled: bool) -> Self {
self.trace = enabled.then(Trace::default);
self
}
pub fn with_absolute_scaled_gradient_tolerance(
mut self,
value: impl Into<Option<F>>,
) -> Self {
self.gradient_tolerance =
crate::core::convergence::optional_tolerance(value);
self
}
pub fn with_initial_radius(mut self, value: impl Into<Option<F>>) -> Self {
let value = value.into();
assert!(
value.is_none_or(|r| r.is_finite() && r > F::zero()),
"initial radius must be finite and positive"
);
self.initial_radius = value;
self
}
pub fn with_rank_tolerance(mut self, value: impl Into<Option<F>>) -> Self {
let value = value.into();
assert!(
value.is_none_or(|r| r.is_finite()
&& r >= F::zero()
&& r < F::one()),
"rank tolerance must be finite and in [0, 1)"
);
self.rank_tolerance = value;
self
}
pub fn with_max_inner_attempts(mut self, value: usize) -> Self {
assert!(value > 0, "inner attempt limit must be positive");
self.max_inner_attempts = value;
self
}
pub fn with_max_subproblem_iterations(mut self, value: usize) -> Self {
assert!(value > 0, "subproblem iteration limit must be positive");
self.max_subproblem_iterations = value;
self
}
}
impl<V: Clone, F: Scalar> InitialState<V> for TrustRegionReflective<F> {
type State = PointState<V, F>;
fn seed(&self, x: &V) -> Self::State {
PointState::new(x.clone())
}
}
#[derive(Clone, Debug)]
struct Work<F> {
free: Vec<usize>,
lower: Vec<F>,
upper: Vec<F>,
residual: Vec<F>,
jacobian: Vec<F>,
gradient: Vec<F>,
optimality: F,
radius: F,
equal_cost_rejection: bool,
failed: bool,
}
fn vector<V: VectorIndex<F> + VectorLen, F: Scalar>(v: &V) -> Vec<F> {
(0..v.vec_len()).map(|i| v.get_scalar(i)).collect()
}
fn jacobian<M: MatrixIndex<F>, F: Scalar>(
j: &M,
rows: usize,
cols: usize,
free: &[usize],
) -> Option<Vec<F>> {
assert_eq!(
j.matrix_rows(),
rows,
"Jacobian row count must equal residual count"
);
assert_eq!(
j.matrix_cols(),
cols,
"Jacobian column count must equal parameter count"
);
if (0..rows).any(|i| (0..cols).any(|k| !j.matrix_entry(i, k).is_finite())) {
return None;
}
Some(
(0..rows)
.flat_map(|i| free.iter().map(move |&k| j.matrix_entry(i, k)))
.collect(),
)
}
fn cost<F: Scalar>(r: &[F]) -> F {
let length = norm(r);
number::<F>(0.5) * length * length
}
impl<F: Scalar> Work<F> {
fn model<V: VectorIndex<F>>(&mut self, x: &V) -> Option<Model<F>> {
let n = self.free.len();
let mut model = Model {
j: self.jacobian.clone(),
g: vec![F::zero(); n],
c: vec![F::zero(); n],
d: vec![F::one(); n],
};
self.optimality = F::zero();
for i in 0..n {
let g = self.gradient[i];
let x = x.get_scalar(self.free[i]);
let v = if g < F::zero() && self.upper[i].is_finite() {
self.upper[i] - x
} else if g > F::zero() && self.lower[i].is_finite() {
x - self.lower[i]
} else {
F::one()
};
model.c[i] = if (g < F::zero() && self.upper[i].is_finite())
|| (g > F::zero() && self.lower[i].is_finite())
{
g.abs()
} else {
F::zero()
};
model.d[i] = v.sqrt();
model.g[i] = model.d[i] * g;
let optimality = v * g.abs();
if !v.is_finite() || v <= F::zero() || !optimality.is_finite() {
return None;
}
self.optimality = self.optimality.max(optimality);
for row in model.j.chunks_mut(n) {
row[i] = row[i] * model.d[i];
}
}
model
.j
.iter()
.chain(&model.g)
.all(|x| x.is_finite())
.then_some(model)
}
fn update_gradient(&mut self) -> bool {
let n = self.free.len();
self.gradient = (0..n)
.map(|i| {
self.residual
.iter()
.enumerate()
.map(|(row, &r)| self.jacobian[row * n + i] * r)
.sum()
})
.collect();
self.gradient.iter().all(|x| x.is_finite())
}
}
impl<P, V, F> Solver<P, PointState<V, F>> for TrustRegionReflective<F>
where
F: Scalar,
P: Residual<Param = V, Output = V> + Jacobian + BoxProblem<Param = V>,
P::Jacobian: MatrixIndex<F>,
V: Clone + VectorLen + VectorIndex<F>,
{
type Error = crate::ConstraintSolveError<<P as Residual>::Error>;
fn reset_convergence(&mut self) {
self.native_convergence.clear();
if let Some(trace) = &mut self.trace {
trace.reset();
}
}
fn check_convergence(
&mut self,
_problem: &Problem<P>,
state: &PointState<V, F>,
) -> Option<Termination<F>> {
self.check_native_convergence(state.cost())
}
fn init(
&mut self,
problem: &mut Problem<P>,
state: PointState<V, F>,
) -> Result<PointState<V, F>, Self::Error> {
crate::core::constraint::evaluation::validate_coordinate_box_problem(
problem.inner(),
state.param().vec_len(),
false,
)
.map_err(crate::core::constraint::evaluation::invalid)?;
self.init_evaluated(problem, state)
.map_err(crate::ConstraintSolveError::Evaluation)
}
fn next_iter(
&mut self,
problem: &mut Problem<P>,
state: PointState<V, F>,
) -> Result<crate::SolverStep<PointState<V, F>>, Self::Error> {
self.next_evaluated(problem, state)
.map_err(crate::ConstraintSolveError::Evaluation)
}
fn terminate(&self, state: &PointState<V, F>) -> Option<Termination<F>> {
self.terminate_evaluated(state)
}
}
impl<P, L, V, F> Solver<RobustLeastSquares<P, L, F>, PointState<V, F>>
for TrustRegionReflective<F>
where
F: Scalar,
P: crate::ProblemConstraints<
Constraints = crate::Constraints<crate::BoxBounds<V>>,
>,
P: Residual<Param = V, Output = V> + Jacobian + BoxProblem<Param = V>,
P::Jacobian: MatrixIndex<F>,
V: Clone + VectorLen + VectorIndex<F>,
L: LossFunction<F>,
P::Jacobian: ScaleRowsInPlace<F>,
{
type Error = crate::ConstraintSolveError<<P as Residual>::Error>;
fn reset_convergence(&mut self) {
self.native_convergence.clear();
if let Some(trace) = &mut self.trace {
trace.reset();
}
}
fn check_convergence(
&mut self,
_problem: &Problem<RobustLeastSquares<P, L, F>>,
state: &PointState<V, F>,
) -> Option<Termination<F>> {
self.check_native_convergence(state.cost())
}
fn init(
&mut self,
problem: &mut Problem<RobustLeastSquares<P, L, F>>,
state: PointState<V, F>,
) -> Result<PointState<V, F>, Self::Error> {
crate::core::constraint::evaluation::validate_coordinate_box_problem(
problem.inner(),
state.param().vec_len(),
false,
)
.map_err(crate::core::constraint::evaluation::invalid)?;
self.init_evaluated(problem, state)
.map_err(crate::ConstraintSolveError::Evaluation)
}
fn next_iter(
&mut self,
problem: &mut Problem<RobustLeastSquares<P, L, F>>,
state: PointState<V, F>,
) -> Result<crate::SolverStep<PointState<V, F>>, Self::Error> {
self.next_evaluated(problem, state)
.map_err(crate::ConstraintSolveError::Evaluation)
}
fn terminate(&self, state: &PointState<V, F>) -> Option<Termination<F>> {
self.terminate_evaluated(state)
}
}
impl<F: Scalar> TrustRegionReflective<F> {
fn init_evaluated<V, M, E>(
&mut self,
problem: &mut E,
state: PointState<V, F>,
) -> Result<PointState<V, F>, E::Error>
where
V: Clone + VectorLen + VectorIndex<F>,
M: MatrixIndex<F>,
E: BoundedEvaluation<V, M, F>,
{
self.native_convergence.clear();
if let Some(trace) = &mut self.trace {
trace.reset();
}
self.work = None;
let mut state = state;
state.reset();
let mut param = state.param().clone();
let n = param.vec_len();
assert!(n > 0, "TRF requires at least one parameter");
let (lo, hi) = (problem.lower(), problem.upper());
assert_eq!(lo.vec_len(), n, "lower bound shape mismatch");
assert_eq!(hi.vec_len(), n, "upper bound shape mismatch");
let mut work = Work {
free: Vec::new(),
lower: Vec::new(),
upper: Vec::new(),
residual: Vec::new(),
jacobian: Vec::new(),
gradient: Vec::new(),
optimality: F::infinity(),
radius: F::one(),
equal_cost_rejection: false,
failed: false,
};
for i in 0..n {
let (x, l, u) =
(param.get_scalar(i), lo.get_scalar(i), hi.get_scalar(i));
assert!(x.is_finite(), "initial parameters must be finite");
assert!(
!l.is_nan()
&& !u.is_nan()
&& l <= u
&& l < F::infinity()
&& u > F::neg_infinity(),
"invalid box bounds"
);
if let Some(x) = interior(x, l, u, true) {
param.set_scalar(i, x);
} else {
work.failed = true;
}
if l < u {
work.free.push(i);
work.lower.push(l);
work.upper.push(u);
}
}
let mut value = F::infinity();
if !work.failed {
if work.free.is_empty() {
let r = problem.residual(¶m)?;
work.residual = vector(&r);
value = problem.cost(&r, |_| cost(&work.residual));
work.optimality = F::zero();
} else {
let (r, j) = problem.residual_and_jacobian(¶m)?;
work.residual = vector(&r);
value = problem.cost(&r, |_| cost(&work.residual));
if let Some((model_r, j)) = problem.model(&r, j) {
if let Some(model_r) = model_r {
work.residual = vector(&model_r);
}
if let Some(j) =
jacobian(&j, work.residual.len(), n, &work.free)
{
work.jacobian = j;
} else {
work.failed = true;
}
} else {
work.failed = true;
}
}
assert!(
!work.residual.is_empty(),
"TRF requires at least one residual"
);
let f = value;
value = if f.is_finite() { f } else { F::infinity() };
work.failed |=
!f.is_finite() || work.residual.iter().any(|x| !x.is_finite());
if !work.failed && !work.free.is_empty() {
work.failed = !work.update_gradient();
if !work.failed {
if let Some(model) = work.model(¶m) {
let scaled: Vec<F> = work
.free
.iter()
.zip(&model.d)
.map(|(&i, &d)| param.get_scalar(i) / d)
.collect();
let auto = norm(&scaled);
work.radius = self.initial_radius.unwrap_or(
if auto > F::zero() { auto } else { F::one() },
);
work.failed = !work.radius.is_finite();
} else {
work.failed = true;
}
}
}
}
self.work = Some(work);
state.replace(param, value);
Ok(state)
}
fn next_evaluated<V, M, E>(
&mut self,
problem: &mut E,
mut state: PointState<V, F>,
) -> NllsStep<V, F, E::Error>
where
V: Clone + VectorLen + VectorIndex<F>,
M: MatrixIndex<F>,
E: BoundedEvaluation<V, M, F>,
{
self.native_convergence.clear();
let work = self.work.as_mut().expect("TRF must be initialized");
let failed = || {
Some(Termination::numerical_failure(
"TRF could not construct or solve a finite feasible trust-region model.",
))
};
if work.failed {
if let Some(trace) = &mut self.trace {
let index = trace.event(
state.cost(),
"construction",
"invalid_initial_model",
);
trace.observations[index].gradient = work.gradient.clone();
trace.observations[index].radius_before = Some(work.radius);
}
return Ok(crate::SolverStep::from((state, failed())));
}
let Some(model) = work.model(state.param()) else {
if let Some(trace) = &mut self.trace {
let index = trace.event(
state.cost(),
"construction",
"invalid_scaling",
);
trace.observations[index].gradient = work.gradient.clone();
trace.observations[index].radius_before = Some(work.radius);
}
work.failed = true;
return Ok(crate::SolverStep::from((state, failed())));
};
let n = work.free.len();
let m = work.residual.len();
let mut augmented = model.j.clone();
augmented.resize((m + n) * n, F::zero());
for i in 0..n {
augmented[(m + i) * n + i] = model.c[i].sqrt();
}
let Some(svd) = DenseSvd::factor(m + n, n, augmented) else {
if let Some(trace) = &mut self.trace {
let index = trace.event(
state.cost(),
"construction",
"factorization_failed",
);
let row = &mut trace.observations[index];
row.gradient = model.g.clone();
row.curvature = model.c.clone();
row.coordinate_scale = model.d.clone();
row.free = work.free.clone();
row.radius_before = Some(work.radius);
}
work.failed = true;
return Ok(crate::SolverStep::from((state, failed())));
};
let mut rhs = work.residual.clone();
rhs.resize(m + n, F::zero());
let x: Vec<F> = work
.free
.iter()
.map(|&i| state.param().get_scalar(i))
.collect();
let theta = number::<F>(0.995).max(F::one() - work.optimality);
let old_cost = state.cost();
for _ in 0..self.max_inner_attempts {
let result = svd.trust_step(
&rhs,
work.radius,
self.rank_tolerance,
self.max_subproblem_iterations,
);
self.trace.as_mut().map(|trace| {
let mut row = LeastSquaresObservation::model(old_cost);
row.event = "subproblem";
row.outcome =
Some(if result.is_some() { "solved" } else { "failed" });
row.model_solves = 1;
row.radius_before = Some(work.radius);
row.radius_after = Some(work.radius);
row.gradient = model.g.clone();
row.curvature = model.c.clone();
row.coordinate_scale = model.d.clone();
row.free = work.free.clone();
if let Some(p) = &result {
row.step = p.clone();
row.step_norm = Some(norm(p));
}
trace.push(row)
});
let Some(p) = result else {
break;
};
let (mut h, _) = model.select(
&x,
&work.lower,
&work.upper,
&p,
work.radius,
theta,
);
let mut trial = state.param().clone();
let mut feasible = true;
let mut roundoff = F::zero();
for i in 0..n {
if let Some(value) = interior(
x[i] + model.d[i] * h[i],
work.lower[i],
work.upper[i],
false,
) {
trial.set_scalar(work.free[i], value);
h[i] = (value - x[i]) / model.d[i];
roundoff = roundoff.hypot(
(F::epsilon() * x[i].abs().max(value.abs()))
/ model.d[i],
);
} else {
feasible = false;
break;
}
}
let length = norm(&h);
let predicted = -model.value(&h);
let selection_record = self.trace.as_mut().map(|trace| {
let mut row = LeastSquaresObservation::model(old_cost);
row.event = "selection";
row.outcome = Some("selected");
row.radius_before = Some(work.radius);
row.radius_after = Some(work.radius);
row.step = h.clone();
row.step_norm = Some(length);
row.predicted_reduction = Some(predicted);
row.gradient = model.g.clone();
row.curvature = model.c.clone();
row.coordinate_scale = model.d.clone();
row.free = work.free.clone();
row.arithmetic = vec![
("coordinate_roundoff", roundoff),
("feasible", if feasible { F::one() } else { F::zero() }),
];
trace.push(row)
});
if feasible
&& work.equal_cost_rejection
&& work
.free
.iter()
.zip(&x)
.all(|(&i, &x)| trial.get_scalar(i) == x)
{
if let Some(index) = selection_record {
self.trace.as_mut().unwrap().observations[index].outcome =
Some("unchanged_after_rejection");
}
return Ok(crate::SolverStep::from((
state,
Some(Termination::numerical_stall(
"A contracted TRF trial changes no parameter after a finite equal-cost rejection.",
)),
)));
}
if !feasible
|| !length.is_finite()
|| length == F::zero()
|| !predicted.is_finite()
|| predicted <= F::zero()
{
if let Some(index) = selection_record {
self.trace.as_mut().unwrap().observations[index].outcome =
Some("invalid_trial_model");
}
break;
}
if length - work.radius
> number::<F>(64.0) * F::epsilon() * work.radius
+ number::<F>(4.0) * roundoff
{
if let Some(index) = selection_record {
self.trace.as_mut().unwrap().observations[index].outcome =
Some("outside_radius");
}
break;
}
let trial_record = self.trace.as_mut().map(|trace| {
let mut row = LeastSquaresObservation::model(old_cost);
row.event = "trial";
row.trial = true;
row.step = h.clone();
row.gradient = model.g.clone();
row.curvature = model.c.clone();
row.coordinate_scale = model.d.clone();
row.free = work.free.clone();
row.radius_before = Some(work.radius);
row.predicted_reduction = Some(predicted);
row.step_norm = Some(length);
row.model_solves = 1;
trace.push(row)
});
let raw_r = match problem.residual(&trial) {
Ok(r) => r,
Err(error) => {
if let Some(index) = trial_record {
self.trace.as_mut().unwrap().observations[index]
.outcome = Some("callback_failed");
}
return Err(error);
}
};
let r = vector(&raw_r);
assert_eq!(r.len(), m, "residual shape changed during solve");
let new_cost = problem.cost(&raw_r, |_| cost(&r));
if let Some(index) = trial_record {
self.trace.as_mut().unwrap().observations[index].trial_cost =
Some(new_cost);
}
if !new_cost.is_finite() || r.iter().any(|x| !x.is_finite()) {
work.equal_cost_rejection = false;
work.radius = number::<F>(0.25) * length;
if let Some(index) = trial_record {
let row =
&mut self.trace.as_mut().unwrap().observations[index];
row.accepted = Some(false);
row.outcome = Some("nonfinite_residual");
row.radius_after = Some(work.radius);
}
continue;
}
let actual = old_cost - new_cost;
work.equal_cost_rejection = actual == F::zero();
let ratio = actual / predicted;
work.radius =
update_radius(work.radius, length, ratio).min(F::max_value());
if let Some(index) = trial_record {
let row = &mut self.trace.as_mut().unwrap().observations[index];
row.actual_reduction = Some(actual);
row.gain_ratio = Some(ratio);
row.accepted = Some(actual > F::zero());
row.outcome = Some(if actual > F::zero() {
"accepted"
} else {
"rejected"
});
row.radius_after = Some(work.radius);
}
if actual > F::zero() {
let j = match problem.jacobian(&trial) {
Ok(j) => j,
Err(error) => {
if let Some(index) = trial_record {
self.trace.as_mut().unwrap().observations[index]
.outcome = Some("accepted_jacobian_failed");
}
return Err(error);
}
};
let Some((model_r, j)) = problem.model(&raw_r, j) else {
if let Some(trace) = &mut self.trace {
trace.event(
old_cost,
"construction",
"invalid_accepted_model",
);
}
break;
};
let Some(j) =
jacobian(&j, m, state.param().vec_len(), &work.free)
else {
if let Some(trace) = &mut self.trace {
trace.event(
old_cost,
"construction",
"invalid_accepted_jacobian",
);
}
break;
};
work.residual = model_r.map_or(r, |r| vector(&r));
work.jacobian = j;
state.replace(trial, new_cost);
work.failed = !work.update_gradient()
|| work.model(state.param()).is_none();
if work.failed {
if let Some(trace) = &mut self.trace {
trace.event(
state.cost(),
"construction",
"invalid_retained_model",
);
}
}
return Ok(crate::SolverStep::from((
state,
if work.failed { failed() } else { None },
)));
}
}
if let Some(trace) = &mut self.trace {
let index =
trace.event(old_cost, "damping_search", "no_step_published");
trace.observations[index].radius_after = Some(work.radius);
}
work.failed = true;
Ok(crate::SolverStep::from((state, failed())))
}
fn terminate_evaluated<V>(
&self,
_: &PointState<V, F>,
) -> Option<Termination<F>> {
self.convergence_test()
.map(|test| self.convergence_report(test))
}
fn convergence_report(
&self,
test: NativeConvergenceTest,
) -> Termination<F> {
match test {
NativeConvergenceTest::NoFreeParameters => {
Termination::converged(crate::ConvergenceCriterion {
test: crate::ConvergenceTest::NoFreeParameters,
evidence: crate::ConvergenceEvidence::NoFreeParameters,
})
}
NativeConvergenceTest::AbsoluteScaledGradient => {
let tolerance = self.gradient_tolerance.unwrap();
Termination::upper_bound(
crate::ConvergenceTest::AbsoluteScaledGradient,
self.work.as_ref().unwrap().optimality,
tolerance,
tolerance,
None,
)
}
_ => unreachable!("TRF only produces its own native predicates"),
}
}
fn convergence_test(&self) -> Option<NativeConvergenceTest> {
let work = self.work.as_ref().filter(|work| !work.failed)?;
if work.free.is_empty() {
Some(NativeConvergenceTest::NoFreeParameters)
} else if self
.gradient_tolerance
.is_some_and(|t| work.optimality <= t)
{
Some(NativeConvergenceTest::AbsoluteScaledGradient)
} else {
None
}
}
fn check_native_convergence(
&mut self,
base_cost: F,
) -> Option<Termination<F>> {
self.native_convergence.clear();
if let Some(trace) = &mut self.trace {
trace.clear();
if let Some(work) = &self.work {
let mut row = LeastSquaresObservation::model(base_cost);
row.gradient = work.gradient.clone();
row.free = work.free.clone();
row.radius_before = Some(work.radius);
if work.free.is_empty() && !work.failed {
row.checks.push(LeastSquaresCheck {
name: "no_free_parameters",
tolerance: None,
passed: Some(true),
evidence: Some(
crate::ConvergenceEvidence::NoFreeParameters,
),
});
} else {
row.checks.push(LeastSquaresCheck {
name: "scaled_gradient",
tolerance: self.gradient_tolerance,
passed: self
.gradient_tolerance
.map(|t| !work.failed && work.optimality <= t),
evidence: self.gradient_tolerance.map(|tolerance| {
crate::ConvergenceEvidence::UpperBound {
value: work.optimality,
bound: tolerance,
tolerance,
reference: None,
}
}),
});
}
trace.push(row);
}
}
let test = self.convergence_test()?;
self.native_convergence.push(test);
Some(self.convergence_report(test))
}
}