diff --git a/crates/track/src/camtrack/bundle.rs b/crates/track/src/camtrack/bundle.rs index b8421f144..a4412dbce 100644 --- a/crates/track/src/camtrack/bundle.rs +++ b/crates/track/src/camtrack/bundle.rs @@ -11,7 +11,7 @@ use rayon::prelude::*; use super::geometry::Pose; -use super::linalg::{M3, V3, add, cholesky_solve, dot, inv3, mm, mtv, mv, rodrigues, scale, skew, sub}; +use super::linalg::{M3, V3, add, cholesky_solve, conjugate_gradient, dot, inv3, mm, mtv, mv, rodrigues, scale, skew, sub}; use super::{Distortion, SolveError}; /// A camera in the adjustment. @@ -287,6 +287,101 @@ 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, +} + +fn camera_preconditioner(u: &[f64], l: &Layout) -> Option> { + 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], l: &Layout) -> Option> { + 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> = + 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(); + } + } + }, + ) +} + /// Smooth per-camera log focal lengths in the normal equations. fn add_focal_prior(cams: &[BaCam], o: &BaOpts, l: &Layout, s_mat: &mut [f64], bc: &mut [f64]) { let n = l.n; @@ -310,10 +405,56 @@ fn add_focal_prior(cams: &[BaCam], o: &BaOpts, l: &Layout, s_mat: &mut [f64], bc } } +/// Eliminate the points into the camera system explicitly (`S -= W V^-1 W^T`, `rhs -= W V^-1 bp`). +fn reduce_points(systems: &[PtSys], vinv: &[Option], n: usize, s: &mut [f64], rhs: &mut [f64]) { + let npts = systems.len(); + let reduce = |range: std::ops::Range, 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 = 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, Vec)> = (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, s, rhs); + } +} + /// Run Levenberg–Marquardt. `cams` and `pts` are updated in place. /// Invalid inputs or workspaces exceeding 256 MiB return an error before mutation. pub fn adjust(cams: &mut [BaCam], pts: &mut [V3], obs: &[BaObs], o: &BaOpts) -> Result { validate(cams, pts, obs, o)?; + adjust_impl(cams, pts, obs, o, true) +} + +fn adjust_impl(cams: &mut [BaCam], pts: &mut [V3], obs: &[BaObs], o: &BaOpts, allow_iterative: bool) -> Result { let l = layout(cams.len(), o); let matrix_len = matrix_len(l.n, pts.len(), obs.len())?; let npts = pts.len(); @@ -341,11 +482,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, bc: &mut Vec| -> PtSys { let mut ps = PtSys { v: [[0.0; 3]; 3], bp: [0.0; 3], w: vec![] }; let mut jc = CameraJacobian::new(); @@ -445,6 +581,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::() / 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 { @@ -472,50 +612,26 @@ pub fn adjust(cams: &mut [BaCam], pts: &mut [V3], obs: &[BaObs], o: &BaOpts) -> }) .collect() }; - if !o.fix_points { - let reduce = |range: std::ops::Range, 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 = 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, Vec)> = (0..chunks) - .into_par_iter() - .map(|t| { - let (mut ds, mut dr) = (vec![0.0; matrix_len], 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); + 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 { + reduce_points(&systems, &vinv, n, &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 = cams.to_vec(); for (i, c) in nc.iter_mut().enumerate() { @@ -593,6 +709,112 @@ pub fn adjust(cams: &mut [BaCam], pts: &mut [V3], obs: &[BaObs], o: &BaOpts) -> mod tests { use super::*; use crate::camtrack::geometry::Rng; + use crate::camtrack::linalg::{I3, log_rot, mt, norm}; + + #[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 = (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).unwrap(); + let direct = adjust_impl(&mut direct_cams, &mut direct_pts, &obs, &opts, false).unwrap(); + 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::()); + u[0] = -1.0; + assert!(iterative_schur(&u, &rhs, &systems, &vinv, &l).is_none()); + } fn test_camera() -> BaCam { BaCam { pose: Pose::IDENTITY, f: 900.0, dist: Distortion::NONE } diff --git a/crates/track/src/camtrack/linalg.rs b/crates/track/src/camtrack/linalg.rs index 692eef34e..71e88e109 100644 --- a/crates/track/src/camtrack/linalg.rs +++ b/crates/track/src/camtrack/linalg.rs @@ -398,6 +398,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> { + 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 { @@ -586,6 +640,30 @@ mod tests { assert!(cholesky_solve(&mut [], 0, &mut [])); } + #[test] + fn conjugate_gradient_checks_residual_and_bad_curvature() { + let n = 64; + let rhs: Vec = (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];