Avoid quadratic Schur assembly for long tracks

Solve expensive camera systems with matrix-free Schur products and a block-preconditioned iterative solver. Check the true linear residual and fall back to the direct solve when convergence fails. Keep short systems on the existing path and compare reprojections, camera parameters and outlier decisions against direct solves.

Fixes #2328
This commit is contained in:
vangel76 committed 2026-10-24 23:59:36 +02:00
1 parent 991904888a
commit b813294fc8
2 files changed
+336 -43

No files matched your search

+258 -43
View File
@@ -174,8 +174,107 @@ fn cam_jac(l: &Layout, ci: usize, c: &BaCam, xc: V3, a: &[[f64; 3]; 2], p: [f64;
}
}
struct PtSys {
v: M3,
bp: V3,
w: Vec<(usize, V3)>,
}
struct PreconditionerBlock {
start: usize,
dim: usize,
inverse: Vec<f64>,
}
fn camera_preconditioner(u: &[f64], l: &Layout) -> Option<Vec<PreconditionerBlock>> {
let mut ranges = Vec::new();
let pose_dim = if l.rot_only { 3 } else { 6 };
for (&pose, &focal) in l.start.iter().zip(&l.fidx) {
if pose != usize::MAX {
ranges.push((pose, pose_dim + usize::from(focal != usize::MAX)));
} else if focal != usize::MAX {
ranges.push((focal, 1));
}
}
if let Some(start) = l.shared.or(l.dist) {
ranges.push((start, l.n - start));
}
ranges
.into_iter()
.map(|(start, dim)| {
let mut matrix = Vec::with_capacity(dim * dim);
for i in 0..dim {
for j in 0..dim {
matrix.push(*u.get((start + i) * l.n + start + j)?);
}
}
let mut inverse = vec![0.0; dim * dim];
for col in 0..dim {
let mut a = matrix.clone();
let mut b = vec![0.0; dim];
*b.get_mut(col)? = 1.0;
if !cholesky_solve(&mut a, dim, &mut b) {
return None;
}
for (row, value) in b.into_iter().enumerate() {
inverse[row * dim + col] = value;
}
}
Some(PreconditionerBlock { start, dim, inverse })
})
.collect()
}
/// Apply the point-eliminated system without materialising W V^-1 W^T.
fn iterative_schur(u: &[f64], rhs: &[f64], systems: &[PtSys], vinv: &[Option<M3>], l: &Layout) -> Option<Vec<f64>> {
if l.n == 0 || l.n.checked_mul(l.n) != Some(u.len()) || rhs.len() != l.n || systems.len() != vinv.len() {
return None;
}
let blocks = camera_preconditioner(u, l)?;
// U contains camera blocks, global parameters and the focal prior. Store its nonzeros
// once; all subsequent products are linear in these and the point-camera connections.
let rows: Vec<Vec<(usize, f64)>> =
u.chunks_exact(l.n).map(|row| row.iter().enumerate().filter(|(_, v)| **v != 0.0).map(|(i, v)| (i, *v)).collect()).collect();
let mut reduced_rhs = rhs.to_vec();
for (ps, vi) in systems.iter().zip(vinv) {
let Some(vi) = vi else { continue };
let vb = mv(vi, ps.bp);
for (i, w) in &ps.w {
*reduced_rhs.get_mut(*i)? -= dot(*w, vb);
}
}
conjugate_gradient(
&reduced_rhs,
|x, y| {
for (i, row) in rows.iter().enumerate() {
y[i] = row.iter().map(|(j, v)| v * x[*j]).sum();
}
for (ps, vi) in systems.iter().zip(vinv) {
let Some(vi) = vi else { continue };
let q = ps.w.iter().fold([0.0; 3], |q, (i, w)| add(q, scale(*w, x[*i])));
let q = mv(vi, q);
for (i, w) in &ps.w {
y[*i] -= dot(*w, q);
}
}
},
|r, z| {
z.fill(0.0);
for block in &blocks {
for i in 0..block.dim {
z[block.start + i] = (0..block.dim).map(|j| block.inverse[i * block.dim + j] * r[block.start + j]).sum();
}
}
},
)
}
/// Run Levenberg–Marquardt. `cams` and `pts` are updated in place.
pub fn adjust(cams: &mut [BaCam], pts: &mut [V3], obs: &[BaObs], o: &BaOpts) -> BaReport {
adjust_impl(cams, pts, obs, o, true)
}
fn adjust_impl(cams: &mut [BaCam], pts: &mut [V3], obs: &[BaObs], o: &BaOpts, allow_iterative: bool) -> BaReport {
let l = layout(cams.len(), o);
let npts = pts.len();
// Observations by point.
@@ -199,11 +298,6 @@ pub fn adjust(cams: &mut [BaCam], pts: &mut [V3], obs: &[BaObs], o: &BaOpts) ->
iters += 1;
// ---- Normal equations (Gauss–Newton with IRLS weights).
// Per point: V (3×3), bp (3), W entries (cam-side index → 3).
struct PtSys {
v: [[f64; 3]; 3],
bp: V3,
w: Vec<(usize, V3)>,
}
let build_point = |j: usize, s_acc: &mut Vec<f64>, bc: &mut Vec<f64>| -> PtSys {
let mut ps = PtSys { v: [[0.0; 3]; 3], bp: [0.0; 3], w: vec![] };
let mut jc: Vec<(usize, [f64; 2])> = Vec::with_capacity(8);
@@ -319,6 +413,10 @@ pub fn adjust(cams: &mut [BaCam], pts: &mut [V3], obs: &[BaObs], o: &BaOpts) ->
let base_s = s_mat.clone();
let base_bc = bc.clone();
let mean_diag = ((0..n).map(|i| base_s[i * n + i]).sum::<f64>() / n.max(1) as f64).max(1e-12);
let schur_pairs = systems.iter().fold(0usize, |sum, ps| sum.saturating_add(ps.w.len().saturating_mul(ps.w.len().saturating_add(1)) / 2));
// Reserve the iterative solve for expensive long-track reductions; small systems
// retain the direct path and its rounding. Failed solves use the direct path too.
let mut use_iterative = allow_iterative && !o.fix_points && n >= 120 && schur_pairs > 4_000_000;
// ---- Try steps with increasing damping.
let mut accepted = false;
for _try in 0..12 {
@@ -346,50 +444,62 @@ pub fn adjust(cams: &mut [BaCam], pts: &mut [V3], obs: &[BaObs], o: &BaOpts) ->
})
.collect()
};
if !o.fix_points {
let reduce = |range: std::ops::Range<usize>, s: &mut [f64], rhs: &mut [f64]| {
for j in range {
let (Some(vi), ps) = (&vinv[j], &systems[j]) else { continue };
let vb = mv(vi, ps.bp);
let wv: Vec<V3> = ps.w.iter().map(|(_, w)| mv(vi, *w)).collect();
for (a, (ia, wa)) in ps.w.iter().enumerate() {
rhs[*ia] -= dot(*wa, vb);
for (ib, wb) in &ps.w[a..] {
let d = dot(wv[a], *wb);
s[ia * n + ib] -= d;
if ia != ib {
s[ib * n + ia] -= d;
let iterative_step = if use_iterative { iterative_schur(&s, &rhs, &systems, &vinv, &l) } else { None };
if dbg && use_iterative {
eprintln!("ba iterative n={n} pairs={schur_pairs} solved={}", iterative_step.is_some());
}
if use_iterative && iterative_step.is_none() {
use_iterative = false;
}
let dc = if let Some(dc) = iterative_step {
dc
} else {
if !o.fix_points {
let reduce = |range: std::ops::Range<usize>, s: &mut [f64], rhs: &mut [f64]| {
for j in range {
let (Some(vi), ps) = (&vinv[j], &systems[j]) else { continue };
let vb = mv(vi, ps.bp);
let wv: Vec<V3> = ps.w.iter().map(|(_, w)| mv(vi, *w)).collect();
for (a, (ia, wa)) in ps.w.iter().enumerate() {
rhs[*ia] -= dot(*wa, vb);
for (ib, wb) in &ps.w[a..] {
let d = dot(wv[a], *wb);
s[ia * n + ib] -= d;
if ia != ib {
s[ib * n + ia] -= d;
}
}
}
}
}
};
let chunks = rayon::current_num_threads().clamp(1, 8);
if n <= 1200 && npts >= 64 && chunks > 1 {
let per = npts.div_ceil(chunks);
let parts: Vec<(Vec<f64>, Vec<f64>)> = (0..chunks)
.into_par_iter()
.map(|t| {
let (mut ds, mut dr) = (vec![0.0; n * n], vec![0.0; n]);
reduce(t * per..((t + 1) * per).min(npts), &mut ds, &mut dr);
(ds, dr)
})
.collect();
for (ds, dr) in parts {
s.par_iter_mut().zip(ds.par_iter()).for_each(|(x, y)| *x += y);
for (x, y) in rhs.iter_mut().zip(dr) {
*x += y;
};
let chunks = rayon::current_num_threads().clamp(1, 8);
if n <= 1200 && npts >= 64 && chunks > 1 {
let per = npts.div_ceil(chunks);
let parts: Vec<(Vec<f64>, Vec<f64>)> = (0..chunks)
.into_par_iter()
.map(|t| {
let (mut ds, mut dr) = (vec![0.0; n * n], vec![0.0; n]);
reduce(t * per..((t + 1) * per).min(npts), &mut ds, &mut dr);
(ds, dr)
})
.collect();
for (ds, dr) in parts {
s.par_iter_mut().zip(ds.par_iter()).for_each(|(x, y)| *x += y);
for (x, y) in rhs.iter_mut().zip(dr) {
*x += y;
}
}
} else {
reduce(0..npts, &mut s, &mut rhs);
}
} else {
reduce(0..npts, &mut s, &mut rhs);
}
}
let mut dc = rhs;
if n > 0 && !cholesky_solve(&mut s, n, &mut dc) {
lambda *= 10.0;
continue;
}
let mut dc = rhs;
if n > 0 && !cholesky_solve(&mut s, n, &mut dc) {
lambda *= 10.0;
continue;
}
dc
};
// Apply.
let mut nc: Vec<BaCam> = cams.to_vec();
for (i, c) in nc.iter_mut().enumerate() {
@@ -464,6 +574,111 @@ mod tests {
use super::*;
use crate::camtrack::geometry::Rng;
#[test]
fn long_track_adjustment_matches_direct_reprojections_and_outliers() {
rayon::ThreadPoolBuilder::new().num_threads(1).build().unwrap().install(|| {
for (focal, fix_centers, distortion) in
[(FocalMode::Fixed, false, false), (FocalMode::Shared, false, true), (FocalMode::PerCamera(30.0), false, true), (FocalMode::Shared, true, true)]
{
let mut rng = Rng(71);
let mut noise = || (rng.next_u64() % 100_000) as f64 / 100_000.0 - 0.5;
let truth: Vec<V3> = (0..320).map(|_| [noise() * 8.0, noise() * 5.0, 10.0 + noise() * 5.0]).collect();
let true_cams: Vec<_> = (0..60)
.map(|i| BaCam {
pose: Pose {
r: rodrigues([0.001 * i as f64, -0.003 * i as f64, 0.0]),
c: if fix_centers { [0.0; 3] } else { [0.03 * i as f64, 0.01 * (i as f64).sin(), 0.0] },
},
f: if matches!(focal, FocalMode::PerCamera(_)) { 900.0 * (1.0 + 0.05 * (i as f64 * 0.07).sin()) } else { 900.0 },
dist: if distortion { Distortion { k1: 0.015, k2: -0.005, radius: 1100.0 } } else { Distortion::NONE },
})
.collect();
let mut obs = Vec::new();
for (ci, c) in true_cams.iter().enumerate() {
for (pi, p) in truth.iter().enumerate() {
let mut uv = residual(c, *p, [0.0; 2]).unwrap();
uv[0] += noise() * 0.05;
uv[1] += noise() * 0.05;
if (ci * truth.len() + pi).is_multiple_of(37) {
uv[0] += 25.0;
uv[1] -= 30.0;
}
obs.push(BaObs { cam: ci as u32, pt: pi as u32, uv });
}
}
let mut cams = true_cams.clone();
for c in cams.iter_mut().skip(2) {
c.pose.r = mm(&rodrigues([0.001, -0.001, 0.0005]), &c.pose.r);
if !fix_centers {
c.pose.c = add(c.pose.c, [0.002, -0.003, 0.001]);
}
}
let mut pts: Vec<_> = truth.iter().map(|p| add(*p, [0.01, -0.01, 0.005])).collect();
let (mut direct_cams, mut direct_pts) = (cams.clone(), pts.clone());
let opts = BaOpts { focal, fix_centers, distortion, fixed: (0..60).map(|i| i < 2).collect(), max_iter: 10, ..Default::default() };
let report = adjust(&mut cams, &mut pts, &obs, &opts);
let direct = adjust_impl(&mut direct_cams, &mut direct_pts, &obs, &opts, false);
assert!((report.cost - direct.cost).abs() < 1e-5 * direct.cost.max(1.0));
for ob in &obs {
let a = residual(&cams[ob.cam as usize], pts[ob.pt as usize], ob.uv).unwrap();
let b = residual(&direct_cams[ob.cam as usize], direct_pts[ob.pt as usize], ob.uv).unwrap();
assert!((a[0] - b[0]).hypot(a[1] - b[1]) < 0.001);
assert_eq!(a[0].hypot(a[1]) <= 4.0, b[0].hypot(b[1]) <= 4.0);
}
for (a, b) in cams.iter().zip(direct_cams) {
assert!(norm(sub(a.pose.c, b.pose.c)) < 0.001);
assert!(norm(log_rot(&mm(&a.pose.r, &mt(&b.pose.r)))) < 1e-4);
assert!((a.f - b.f).abs() < 1e-4 * b.f);
assert!((a.dist.k1 - b.dist.k1).abs() < 1e-4 && (a.dist.k2 - b.dist.k2).abs() < 1e-4);
}
}
});
}
#[test]
fn matrix_free_schur_matches_explicit_system_and_rejects_bad_preconditioner() {
let l = layout(20, &BaOpts::default());
let n = l.n;
let mut u = vec![0.0; n * n];
for i in 0..n {
u[i * n + i] = 10.0;
}
let systems: Vec<_> = (0..80)
.map(|j| PtSys {
v: I3,
bp: [0.01, -0.02, 0.03],
w: (0..n).map(|i| (i, [((j * 13 + i * 7) % 19) as f64 * 0.001 - 0.009, (i % 7) as f64 * 0.001, (j % 11) as f64 * 0.001])).collect(),
})
.collect();
let vinv = vec![Some(I3); systems.len()];
let rhs: Vec<_> = (0..n).map(|i| (i % 7) as f64 - 3.0).collect();
let x = iterative_schur(&u, &rhs, &systems, &vinv, &l).unwrap();
let mut explicit = u.clone();
let mut reduced_rhs = rhs.clone();
for ps in &systems {
for (i, a) in &ps.w {
reduced_rhs[*i] -= dot(*a, ps.bp);
for (j, b) in &ps.w {
explicit[i * n + j] -= dot(*a, *b);
}
}
}
let mut direct = reduced_rhs.clone();
let mut factor = explicit.clone();
assert!(cholesky_solve(&mut factor, n, &mut direct));
for (a, b) in x.iter().zip(direct) {
assert!((a - b).abs() < 1e-9);
}
let mut error2 = 0.0;
for (i, row) in explicit.chunks_exact(n).enumerate() {
let ax: f64 = row.iter().zip(&x).map(|(a, x)| a * x).sum();
error2 += (ax - reduced_rhs[i]).powi(2);
}
assert!(error2 < 1e-18 * reduced_rhs.iter().map(|v| v * v).sum::<f64>());
u[0] = -1.0;
assert!(iterative_schur(&u, &rhs, &systems, &vinv, &l).is_none());
}
#[test]
fn bundle_adjustment_converges_from_noise() {
let mut rng = Rng(99);
+78
View File
@@ -353,6 +353,60 @@ pub fn cholesky_solve(a: &mut [f64], n: usize, b: &mut [f64]) -> bool {
b.iter().all(|v| v.is_finite())
}
/// Preconditioned conjugate gradients with an independently recomputed residual check.
/// Failure to converge (or non-finite/non-positive curvature) requests the direct fallback.
pub(crate) fn conjugate_gradient(rhs: &[f64], apply: impl Fn(&[f64], &mut [f64]), precondition: impl Fn(&[f64], &mut [f64])) -> Option<Vec<f64>> {
let n = rhs.len();
let norm2 = dot_n(rhs, rhs);
if !norm2.is_finite() {
return None;
}
if rhs.iter().all(|v| *v == 0.0) {
return Some(vec![0.0; n]);
}
let tolerance2 = norm2 * 1e-22;
if tolerance2 == 0.0 {
return None;
}
let mut x = vec![0.0; n];
let mut r = rhs.to_vec();
let mut z = vec![0.0; n];
let mut ap = vec![0.0; n];
precondition(&r, &mut z);
let mut p = z.clone();
let mut rz = dot_n(&r, &z);
for _ in 0..n.saturating_mul(2).min(512) {
if !rz.is_finite() || rz <= 0.0 {
return None;
}
apply(&p, &mut ap);
let curvature = dot_n(&p, &ap);
if !curvature.is_finite() || curvature <= 0.0 {
return None;
}
let alpha = rz / curvature;
for i in 0..n {
x[i] += alpha * p[i];
r[i] -= alpha * ap[i];
}
if dot_n(&r, &r) <= tolerance2 {
apply(&x, &mut ap);
for i in 0..n {
ap[i] = rhs[i] - ap[i];
}
return (x.iter().all(|v| v.is_finite()) && dot_n(&ap, &ap) <= 100.0 * tolerance2).then_some(x);
}
precondition(&r, &mut z);
let next_rz = dot_n(&r, &z);
let beta = next_rz / rz;
for i in 0..n {
p[i] = z[i] + beta * p[i];
}
rz = next_rz;
}
None
}
/// A similarity transform `x ↦ s R x + t`.
#[derive(Clone, Copy, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct Similarity {
@@ -428,6 +482,30 @@ pub fn from_euler_xyz(e: V3) -> M3 {
mod tests {
use super::*;
#[test]
fn conjugate_gradient_checks_residual_and_bad_curvature() {
let n = 64;
let rhs: Vec<f64> = (0..n).map(|i| (i % 11) as f64 - 5.0).collect();
let apply = |x: &[f64], y: &mut [f64]| {
for i in 0..n {
y[i] = 4.0 * x[i] + if i > 0 { x[i - 1] } else { 0.0 } + if i + 1 < n { x[i + 1] } else { 0.0 };
}
};
let x = conjugate_gradient(&rhs, apply, |r, z| {
for (r, z) in r.iter().zip(z) {
*z = *r / 4.0;
}
})
.unwrap();
let mut product = vec![0.0; n];
apply(&x, &mut product);
let residual: f64 = rhs.iter().zip(product).map(|(b, ax)| (b - ax).powi(2)).sum();
assert!(residual < 1e-18 * dot_n(&rhs, &rhs));
assert!(conjugate_gradient(&[1.0], |x, y| y[0] = -x[0], |r, z| z.copy_from_slice(r)).is_none());
assert!(conjugate_gradient(&[f64::NAN], |_, _| {}, |_, _| {}).is_none());
assert_eq!(conjugate_gradient(&[0.0], |_, _| {}, |_, _| {}).unwrap(), vec![0.0]);
}
#[test]
fn svd_reconstructs() {
let a = [3.0, 1.0, 2.0, -1.0, 4.0, 0.5, 2.0, 2.0, 7.0, 0.1, -3.0, 1.0];