apex-solver 1.5.0

//! Prior factors: unary constraints anchoring a variable to a known value.
//!
//! Two flavors, chosen by the variable's manifold:
//!
//! - [`PriorFactor<T>`] — **tangent-space anchor** for Lie groups
//!   (`r = Log(T_prior⁻¹ ∘ X)`). Correct for SE(2)/SE(3)/Sim(3)/…: no
//!   quaternion double-cover ambiguity, tangent-dimension Jacobian with the
//!   full rotation–translation coupling.
//! - [`EuclideanPriorFactor`] — plain parameter difference (`r = x − x_prior`)
//!   for [`Rn`](apex_manifolds::rn::Rn) variables **only**; registration
//!   rejects any other manifold.

use super::BetweenFactor;
use crate::core::variable::ManifoldVariable;
use crate::factors::Factor;
use crate::factors::common::validate::expect_block_sizes;
use apex_manifolds::LieGroup;
use faer::prelude::ReborrowMut;
use nalgebra::DVector;

/// Tangent-space prior (unary anchor) on a Lie-group variable.
///
/// # Mathematical Formulation
///
/// ```text
/// r = Log(T_prior⁻¹ ∘ X) ∈ ℝ^dof,   J = J_log(·) · J_between|_X
/// ```
///
/// The residual vanishes iff `X == T_prior`. Because the error lives in the
/// tangent space:
///
/// - no quaternion double-cover ambiguity (a prior near a sign flip is fine),
/// - the Jacobian is `dof × dof` — full rotation–translation coupling, no
///   dropped 7th column,
/// - the angle difference wraps correctly (SE(2) priors at θ_prior = 3.14
///   against θ = −3.14 give a residual of ≈ 0.003, not ≈ 6.28).
///
/// The chain rule reuses the [`BetweenFactor`](super::BetweenFactor)
/// machinery with a constant origin.
///
/// # Use Cases
///
/// - **Anchoring**: fix the first pose in SLAM to prevent drift
/// - **Loop-closure priors**: incorporate a known global pose
/// - **Regularization**: keep a pose near a trusted value
///
/// # Restrictions
///
/// Fixed-DOF groups (SO(2), SO(3), SE(2), SE(3), SE₂(3), Sim(3), SGal(3)).
/// For `Rn` variables use [`EuclideanPriorFactor`] — the tangent space and the
/// parameter space coincide there, and `Rn`'s dimension is dynamic.
///
/// # Examples
///
/// ```
/// # use apex_solver::factors::pose::PriorFactor;
/// # use apex_solver::manifold::se3::SE3;
/// # use nalgebra::Vector3;
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// let prior = SE3::from_translation_quaternion(
///     nalgebra::Vector3::new(0.0, 0.0, 0.0),
///     nalgebra::Quaternion::new(1.0, 0.0, 0.0, 0.0),
/// );
/// let factor = PriorFactor::new(prior);
/// # Ok(())
/// # }
/// ```
#[derive(Debug, Clone)]
pub struct PriorFactor<T>
where
    T: LieGroup + Clone + Send + Sync,
{
    /// The anchored value: the residual is zero when the variable equals it.
    prior: T,
    /// The identity-measurement between constraint this factor delegates to,
    /// built once so `linearize` allocates nothing.
    between: BetweenFactor<T>,
}

impl<T> PriorFactor<T>
where
    T: LieGroup + Clone + Send + Sync,
{
    /// Anchor a variable to `prior`.
    pub fn new(prior: T) -> Self {
        // `x.between(&x)` is `x⁻¹ ∘ x` — the group identity, and for
        // dimension-carrying groups such as `Rn` it is the identity of the
        // *right* dimension, which an associated `identity()` could not be.
        let identity = prior.between(&prior, None, None);
        Self {
            prior,
            between: BetweenFactor::new(identity),
        }
    }
}

impl<T> Factor for PriorFactor<T>
where
    T: LieGroup + Clone + Send + Sync,
{
    fn linearize(
        &self,
        params: &[&[f64]],
        residual: &mut [f64],
        jacobian: Option<faer::mat::MatMut<'_, f64>>,
    ) {
        // r = Log(prior⁻¹ ∘ X) expressed through the BetweenFactor chain with
        // an identity measurement — identical to anchoring with
        // `BetweenFactor::new(T::identity())` between (prior, X). Delegating
        // keeps this factor's Jacobian bit-for-bit consistent with the
        // production between machinery (same chain, same conventions).
        //
        // Note the residual sign follows the between convention (the overall
        // sign of a least-squares residual is immaterial: the zero and the
        // optimum are the same).
        //
        // `linearize_wrt_second` writes only the block this factor needs, so
        // nothing is allocated here — `linearize` runs once per factor per
        // iteration and the trait documents it as allocation-free.
        self.between.linearize_wrt_second(
            self.prior.as_param_slice(),
            params[0],
            residual,
            jacobian,
        );
    }

    fn residual_dim(&self) -> usize {
        self.prior.tangent_dim()
    }

    fn jacobian_shape(&self) -> (usize, usize) {
        let dof = self.prior.tangent_dim();
        (dof, dof)
    }

    fn validate_variables(&self, variables: &[&dyn ManifoldVariable]) -> Result<(), String> {
        expect_block_sizes(
            variables,
            &[self.prior.as_param_slice().len()],
            "PriorFactor expects one variable of the prior's own group",
        )
    }
}

/// Euclidean prior on an [`Rn`](apex_manifolds::rn::Rn) variable.
///
/// # Mathematical Formulation
///
/// ```text
/// r = x − x_prior,   J = I
/// ```
///
/// In the parameter space — which for `Rn` *is* the tangent space, so this is
/// exact there. Registration rejects non-`Rn` manifolds: on SE(3) this
/// residual is meaningless (quaternion double cover, dropped rotation–
/// translation coupling, no angle wrap) — use [`PriorFactor<T>`] instead.
///
/// Combined with a diagonal [`noise model`](crate::core::noise::NoiseModel)
/// this is the weighted prior of the `weighted_prior_constrained_ls` branch
/// (`WeightedPriorFactor ≡ EuclideanPriorFactor + diagonal sqrt-info`), which
/// is why no separate weighted factor type exists.
///
/// # Use Cases
///
/// - Anchoring landmarks, velocities, IMU biases (`Rn` states)
/// - Priors whose components carry different units, weighted per-axis by the
///   noise model
#[derive(Debug, Clone)]
pub struct EuclideanPriorFactor {
    /// The prior value (measurement or known value)
    pub data: DVector<f64>,
}

impl EuclideanPriorFactor {
    /// Anchor an `Rn` variable to `data`.
    pub fn new(data: DVector<f64>) -> Self {
        Self { data }
    }
}

impl Factor for EuclideanPriorFactor {
    fn linearize(
        &self,
        params: &[&[f64]],
        residual: &mut [f64],
        jacobian: Option<faer::mat::MatMut<'_, f64>>,
    ) {
        let n = self.data.len();
        for i in 0..n {
            residual[i] = params[0][i] - self.data[i];
        }
        if let Some(mut jac) = jacobian {
            for i in 0..n {
                *jac.rb_mut().get_mut(i, i) = 1.0;
            }
        }
    }

    fn residual_dim(&self) -> usize {
        self.data.len()
    }

    fn jacobian_shape(&self) -> (usize, usize) {
        let n = self.data.len();
        (n, n)
    }

    /// Restrict to `Rn` variables: on anything else the ambient difference is
    /// not the tangent-space error (see the type docs).
    fn validate_variables(&self, variables: &[&dyn ManifoldVariable]) -> Result<(), String> {
        for variable in variables {
            if variable.manifold_type_name() != "Rn" {
                return Err(format!(
                    "EuclideanPriorFactor operates on ambient parameters and is restricted to Rn \
                     variables, but was registered on a {} variable. Use PriorFactor<T> for \
                     Lie-group anchoring.",
                    variable.manifold_type_name()
                ));
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::factors::pose::BetweenFactor;
    use apex_manifolds::se2::{SE2, SE2Tangent};
    use apex_manifolds::se3::{SE3, SE3Tangent};
    use nalgebra::{Quaternion, Vector3};

    type TestResult = Result<(), Box<dyn std::error::Error>>;

    /// Central-difference Jacobian of an SE(3) prior's residual w.r.t. the
    /// variable's tangent coordinates.
    fn fd_jacobian_se3(factor: &PriorFactor<SE3>, x: &SE3) -> Vec<Vec<f64>> {
        let dof = 6;
        let eps = 1e-6;
        let mut fd = vec![vec![0.0; dof]; factor.residual_dim()];
        for k in 0..dof {
            let mut rho = Vector3::zeros();
            let mut theta = Vector3::zeros();
            if k < 3 {
                rho[k] = eps;
            } else {
                theta[k - 3] = eps;
            }
            let step_p = SE3Tangent::new(rho, theta);
            let step_m = SE3Tangent::new(-rho, -theta);
            let x_p = x.plus(&step_p, None, None);
            let x_m = x.plus(&step_m, None, None);
            let mut rp = vec![0.0; factor.residual_dim()];
            let mut rm = vec![0.0; factor.residual_dim()];
            factor.linearize(&[x_p.as_param_slice()], &mut rp, None);
            factor.linearize(&[x_m.as_param_slice()], &mut rm, None);
            for i in 0..factor.residual_dim() {
                fd[i][k] = (rp[i] - rm[i]) / (2.0 * eps);
            }
        }
        fd
    }

    /// Central-difference Jacobian for an SE(2) prior.
    fn fd_jacobian_se2(factor: &PriorFactor<SE2>, x: &SE2) -> Vec<Vec<f64>> {
        let dof = 3;
        let eps = 1e-6;
        let mut fd = vec![vec![0.0; dof]; factor.residual_dim()];
        for k in 0..dof {
            let mut c = [0.0f64; 3];
            c[k] = eps;
            let step_p = SE2Tangent::new(c[0], c[1], c[2]);
            let step_m = SE2Tangent::new(-c[0], -c[1], -c[2]);
            let x_p = x.plus(&step_p, None, None);
            let x_m = x.plus(&step_m, None, None);
            let mut rp = vec![0.0; factor.residual_dim()];
            let mut rm = vec![0.0; factor.residual_dim()];
            factor.linearize(&[x_p.as_param_slice()], &mut rp, None);
            factor.linearize(&[x_m.as_param_slice()], &mut rm, None);
            for i in 0..factor.residual_dim() {
                fd[i][k] = (rp[i] - rm[i]) / (2.0 * eps);
            }
        }
        fd
    }

    /// FD-vs-analytic check for the SE(3) prior Jacobian.
    ///
    /// Was ignored (ISSUE-0001): `SE3Tangent::q_block_jacobian_matrix`'s last
    /// term computed `(c - 3.0) * e` instead of `0.5 * (c - 3.0 * e)` and
    /// multiplied only `T·P·T²`, omitting the distinct `T²·P·T` term. Fixed;
    /// this now agrees with central differences like the SE(2) variant below.
    #[test]
    fn prior_factor_se3_jacobian_matches_central_difference() -> TestResult {
        let prior = SE3::from_translation_quaternion(
            Vector3::new(0.2, -0.4, 0.7),
            Quaternion::new(0.9, 0.1, -0.2, 0.3),
        );
        let x = SE3::from_translation_quaternion(
            Vector3::new(1.0, 0.5, -0.3),
            Quaternion::new(0.8, 0.2, 0.1, -0.4),
        );
        let factor = PriorFactor::new(prior.clone());

        let (rows, cols) = factor.jacobian_shape();
        let mut jac_buf = vec![0.0; rows * cols];
        let jac_mut = faer::mat::MatMut::from_column_major_slice_mut(&mut jac_buf, rows, cols);
        let mut res_buf = vec![0.0; rows];
        factor.linearize(&[x.as_param_slice()], &mut res_buf, Some(jac_mut));

        let fd = fd_jacobian_se3(&factor, &x);
        for i in 0..rows {
            for k in 0..cols {
                let analytic = jac_buf[k * rows + i]; // column-major
                assert!(
                    (analytic - fd[i][k]).abs() < 1e-5,
                    "SE3 prior J[{i}][{k}]: analytic {analytic} vs FD {}",
                    fd[i][k]
                );
            }
        }
        Ok(())
    }

    #[test]
    fn prior_factor_se2_jacobian_matches_central_difference() -> TestResult {
        let prior = SE2::from_xy_angle(0.7, -0.2, 1.2);
        let x = SE2::from_xy_angle(-0.3, 0.9, -2.5);
        let factor = PriorFactor::new(prior);

        let (rows, cols) = factor.jacobian_shape();
        let mut jac_buf = vec![0.0; rows * cols];
        let jac_mut = faer::mat::MatMut::from_column_major_slice_mut(&mut jac_buf, rows, cols);
        let mut res_buf = vec![0.0; rows];
        factor.linearize(&[x.as_param_slice()], &mut res_buf, Some(jac_mut));

        let fd = fd_jacobian_se2(&factor, &x);
        for i in 0..rows {
            for k in 0..cols {
                let analytic = jac_buf[k * rows + i];
                assert!(
                    (analytic - fd[i][k]).abs() < 1e-5,
                    "SE2 prior J[{i}][{k}]: analytic {analytic} vs FD {}",
                    fd[i][k]
                );
            }
        }
        Ok(())
    }

    /// The prior must equal the equivalent between construction:
    /// `Log(prior⁻¹∘X)` == `Log(between(prior, X) ∘ I)`.
    #[test]
    fn prior_factor_matches_between_with_identity_measurement() -> TestResult {
        let prior = SE3::from_translation_quaternion(
            Vector3::new(0.2, -0.4, 0.7),
            Quaternion::new(0.9, 0.1, -0.2, 0.3),
        );
        let x = SE3::from_translation_quaternion(
            Vector3::new(1.0, 0.5, -0.3),
            Quaternion::new(0.8, 0.2, 0.1, -0.4),
        );

        let mut rp = vec![0.0; 6];
        PriorFactor::<SE3>::new(prior.clone()).linearize(&[x.as_param_slice()], &mut rp, None);

        let mut rb = vec![0.0; 6];
        let between = BetweenFactor::new(SE3::identity());
        between.linearize(&[prior.as_param_slice(), x.as_param_slice()], &mut rb, None);

        for i in 0..6 {
            assert!(
                (rp[i] - rb[i]).abs() < 1e-12,
                "row {i}: {} vs {}",
                rp[i],
                rb[i]
            );
        }
        Ok(())
    }

    /// SE(2) angle wrap: a prior at +3.1 rad against a state at −3.1 rad must
    /// give a small residual through the wrap (the old ambient factor gave
    /// ≈ 6.2).
    #[test]
    fn prior_factor_se2_wraps_angle() -> TestResult {
        let prior = SE2::from_xy_angle(0.0, 0.0, 3.1);
        let x = SE2::from_xy_angle(0.0, 0.0, -3.1);
        let mut r = vec![0.0; 3];
        PriorFactor::<SE2>::new(prior).linearize(&[x.as_param_slice()], &mut r, None);
        assert!(
            r[2].abs() < 0.1,
            "angle residual must wrap: got {} (expected ≈ -0.083)",
            r[2]
        );
        Ok(())
    }

    /// EuclideanPriorFactor must refuse non-Rn variables at registration.
    #[test]
    fn euclidean_prior_rejects_non_rn_variables() -> TestResult {
        use crate::core::variable::{ManifoldVariable, Variable};
        use apex_manifolds::se3::SE3;
        use slotmap::SlotMap;

        let mut variables: SlotMap<crate::core::VarKey, Box<dyn ManifoldVariable>> =
            SlotMap::with_key();
        let key = variables.insert(Box::new(Variable::new(SE3::identity())));

        let factor = EuclideanPriorFactor::new(DVector::from_element(7, 0.0));
        let err = factor
            .validate_variables(&[variables.get(key).ok_or("key missing")?.as_ref()])
            .err()
            .ok_or("expected rejection")?;
        assert!(err.contains("Rn"), "rejection must name Rn: {err}");
        Ok(())
    }
}