use std::{
collections::{HashMap, HashSet},
fs::File,
io::{Error, Write},
sync::{Arc, Mutex},
};
use faer::{Col, Mat, MatRef, sparse::SparseColMat};
use nalgebra::DVector;
use rayon::prelude::*;
use tracing::warn;
use crate::error::ErrorLogging;
use crate::{
core::CoreResult,
core::{
CoreError, corrector::Corrector, loss_functions::LossFunction,
residual_block::ResidualBlock, variable::Variable,
},
factors::Factor,
linalg::{JacobianMode, LinearSolver, SparseMode, extract_variable_covariances},
};
use apex_manifolds::{ManifoldType, rn, se2, se3, se23, sgal3, sim3, so2, so3};
pub use crate::linearizer::cpu::sparse::SymbolicStructure;
#[derive(Clone)]
pub enum VariableEnum {
Rn(Variable<rn::Rn>),
SE2(Variable<se2::SE2>),
SE3(Variable<se3::SE3>),
SE23(Variable<se23::SE23>),
SGal3(Variable<sgal3::SGal3>),
Sim3(Variable<sim3::Sim3>),
SO2(Variable<so2::SO2>),
SO3(Variable<so3::SO3>),
}
impl VariableEnum {
pub fn manifold_type(&self) -> ManifoldType {
match self {
VariableEnum::Rn(_) => ManifoldType::RN,
VariableEnum::SE2(_) => ManifoldType::SE2,
VariableEnum::SE3(_) => ManifoldType::SE3,
VariableEnum::SE23(_) => ManifoldType::SE23,
VariableEnum::SGal3(_) => ManifoldType::SGal3,
VariableEnum::Sim3(_) => ManifoldType::Sim3,
VariableEnum::SO2(_) => ManifoldType::SO2,
VariableEnum::SO3(_) => ManifoldType::SO3,
}
}
pub fn get_size(&self) -> usize {
match self {
VariableEnum::Rn(var) => var.get_size(),
VariableEnum::SE2(var) => var.get_size(),
VariableEnum::SE3(var) => var.get_size(),
VariableEnum::SE23(var) => var.get_size(),
VariableEnum::SGal3(var) => var.get_size(),
VariableEnum::Sim3(var) => var.get_size(),
VariableEnum::SO2(var) => var.get_size(),
VariableEnum::SO3(var) => var.get_size(),
}
}
pub fn to_vector(&self) -> DVector<f64> {
match self {
VariableEnum::Rn(var) => var.value.clone().into(),
VariableEnum::SE2(var) => var.value.clone().into(),
VariableEnum::SE3(var) => var.value.clone().into(),
VariableEnum::SE23(var) => var.value.clone().into(),
VariableEnum::SGal3(var) => var.value.clone().into(),
VariableEnum::Sim3(var) => var.value.clone().into(),
VariableEnum::SO2(var) => var.value.clone().into(),
VariableEnum::SO3(var) => var.value.clone().into(),
}
}
pub fn apply_tangent_step(&mut self, step_slice: MatRef<f64>) {
match self {
VariableEnum::SE3(var) => {
let mut step_data = DVector::from_fn(6, |i, _| step_slice[(i, 0)]);
for &fixed_idx in &var.fixed_indices {
if fixed_idx < 6 {
step_data[fixed_idx] = 0.0;
}
}
let tangent = se3::SE3Tangent::from(step_data);
let new_value = var.plus(&tangent);
var.set_value(new_value);
}
VariableEnum::SE2(var) => {
let mut step_data = DVector::from_fn(3, |i, _| step_slice[(i, 0)]);
for &fixed_idx in &var.fixed_indices {
if fixed_idx < 3 {
step_data[fixed_idx] = 0.0;
}
}
let tangent = se2::SE2Tangent::from(step_data);
let new_value = var.plus(&tangent);
var.set_value(new_value);
}
VariableEnum::SO3(var) => {
let mut step_data = DVector::from_fn(3, |i, _| step_slice[(i, 0)]);
for &fixed_idx in &var.fixed_indices {
if fixed_idx < 3 {
step_data[fixed_idx] = 0.0;
}
}
let tangent = so3::SO3Tangent::from(step_data);
let new_value = var.plus(&tangent);
var.set_value(new_value);
}
VariableEnum::SO2(var) => {
let mut step_data = DVector::from_fn(1, |i, _| step_slice[(i, 0)]);
for &fixed_idx in &var.fixed_indices {
if fixed_idx < 1 {
step_data[fixed_idx] = 0.0;
}
}
let tangent = so2::SO2Tangent::from(step_data);
let new_value = var.plus(&tangent);
var.set_value(new_value);
}
VariableEnum::SE23(var) => {
let mut step_data = DVector::from_fn(9, |i, _| step_slice[(i, 0)]);
for &fixed_idx in &var.fixed_indices {
if fixed_idx < 9 {
step_data[fixed_idx] = 0.0;
}
}
let tangent = se23::SE23Tangent::from(step_data);
let new_value = var.plus(&tangent);
var.set_value(new_value);
}
VariableEnum::SGal3(var) => {
let mut step_data = DVector::from_fn(10, |i, _| step_slice[(i, 0)]);
for &fixed_idx in &var.fixed_indices {
if fixed_idx < 10 {
step_data[fixed_idx] = 0.0;
}
}
let tangent = sgal3::SGal3Tangent::from(step_data);
let new_value = var.plus(&tangent);
var.set_value(new_value);
}
VariableEnum::Sim3(var) => {
let mut step_data = DVector::from_fn(7, |i, _| step_slice[(i, 0)]);
for &fixed_idx in &var.fixed_indices {
if fixed_idx < 7 {
step_data[fixed_idx] = 0.0;
}
}
let tangent = sim3::Sim3Tangent::from(step_data);
let new_value = var.plus(&tangent);
var.set_value(new_value);
}
VariableEnum::Rn(var) => {
let size = var.get_size();
let mut step_data = DVector::from_fn(size, |i, _| step_slice[(i, 0)]);
for &fixed_idx in &var.fixed_indices {
if fixed_idx < size {
step_data[fixed_idx] = 0.0;
}
}
let tangent = rn::RnTangent::new(step_data);
let new_value = var.plus(&tangent);
var.set_value(new_value);
}
}
}
pub fn get_covariance(&self) -> Option<&Mat<f64>> {
match self {
VariableEnum::Rn(var) => var.get_covariance(),
VariableEnum::SE2(var) => var.get_covariance(),
VariableEnum::SE3(var) => var.get_covariance(),
VariableEnum::SE23(var) => var.get_covariance(),
VariableEnum::SGal3(var) => var.get_covariance(),
VariableEnum::Sim3(var) => var.get_covariance(),
VariableEnum::SO2(var) => var.get_covariance(),
VariableEnum::SO3(var) => var.get_covariance(),
}
}
pub fn set_covariance(&mut self, cov: Mat<f64>) {
match self {
VariableEnum::Rn(var) => var.set_covariance(cov),
VariableEnum::SE2(var) => var.set_covariance(cov),
VariableEnum::SE3(var) => var.set_covariance(cov),
VariableEnum::SE23(var) => var.set_covariance(cov),
VariableEnum::SGal3(var) => var.set_covariance(cov),
VariableEnum::Sim3(var) => var.set_covariance(cov),
VariableEnum::SO2(var) => var.set_covariance(cov),
VariableEnum::SO3(var) => var.set_covariance(cov),
}
}
pub fn clear_covariance(&mut self) {
match self {
VariableEnum::Rn(var) => var.clear_covariance(),
VariableEnum::SE2(var) => var.clear_covariance(),
VariableEnum::SE3(var) => var.clear_covariance(),
VariableEnum::SE23(var) => var.clear_covariance(),
VariableEnum::SGal3(var) => var.clear_covariance(),
VariableEnum::Sim3(var) => var.clear_covariance(),
VariableEnum::SO2(var) => var.clear_covariance(),
VariableEnum::SO3(var) => var.clear_covariance(),
}
}
pub fn get_bounds(&self) -> &HashMap<usize, (f64, f64)> {
match self {
VariableEnum::Rn(var) => &var.bounds,
VariableEnum::SE2(var) => &var.bounds,
VariableEnum::SE3(var) => &var.bounds,
VariableEnum::SE23(var) => &var.bounds,
VariableEnum::SGal3(var) => &var.bounds,
VariableEnum::Sim3(var) => &var.bounds,
VariableEnum::SO2(var) => &var.bounds,
VariableEnum::SO3(var) => &var.bounds,
}
}
pub fn get_fixed_indices(&self) -> &HashSet<usize> {
match self {
VariableEnum::Rn(var) => &var.fixed_indices,
VariableEnum::SE2(var) => &var.fixed_indices,
VariableEnum::SE3(var) => &var.fixed_indices,
VariableEnum::SE23(var) => &var.fixed_indices,
VariableEnum::SGal3(var) => &var.fixed_indices,
VariableEnum::Sim3(var) => &var.fixed_indices,
VariableEnum::SO2(var) => &var.fixed_indices,
VariableEnum::SO3(var) => &var.fixed_indices,
}
}
pub fn set_from_vector(&mut self, vec: &DVector<f64>) {
match self {
VariableEnum::Rn(var) => {
var.set_value(rn::Rn::new(vec.clone()));
}
VariableEnum::SE2(var) => {
let new_se2: se2::SE2 = vec.clone().into();
var.set_value(new_se2);
}
VariableEnum::SE3(var) => {
let new_se3: se3::SE3 = vec.clone().into();
var.set_value(new_se3);
}
VariableEnum::SE23(var) => {
let new_se23: se23::SE23 = vec.clone().into();
var.set_value(new_se23);
}
VariableEnum::SGal3(var) => {
let new_sgal3: sgal3::SGal3 = vec.clone().into();
var.set_value(new_sgal3);
}
VariableEnum::Sim3(var) => {
let new_sim3: sim3::Sim3 = vec.clone().into();
var.set_value(new_sim3);
}
VariableEnum::SO2(var) => {
let new_so2: so2::SO2 = vec.clone().into();
var.set_value(new_so2);
}
VariableEnum::SO3(var) => {
let new_so3: so3::SO3 = vec.clone().into();
var.set_value(new_so3);
}
}
}
}
pub struct Problem {
pub(crate) total_residual_dimension: usize,
pub(crate) jacobian_mode: JacobianMode,
residual_id_count: usize,
residual_blocks: HashMap<usize, ResidualBlock>,
pub(crate) fixed_variable_indexes: HashMap<String, HashSet<usize>>,
pub(crate) variable_bounds: HashMap<String, HashMap<usize, (f64, f64)>>,
}
impl Default for Problem {
fn default() -> Self {
Self::new(JacobianMode::Sparse)
}
}
impl Problem {
pub fn new(jacobian_mode: JacobianMode) -> Self {
Self {
total_residual_dimension: 0,
jacobian_mode,
residual_id_count: 0,
residual_blocks: HashMap::new(),
fixed_variable_indexes: HashMap::new(),
variable_bounds: HashMap::new(),
}
}
pub fn add_residual_block(
&mut self,
variable_key_size_list: &[&str],
factor: Box<dyn Factor + Send>,
loss_func: Option<Box<dyn LossFunction + Send>>,
) -> usize {
let new_residual_dimension = factor.get_dimension();
self.residual_blocks.insert(
self.residual_id_count,
ResidualBlock::new(
self.residual_id_count,
self.total_residual_dimension,
variable_key_size_list,
factor,
loss_func,
),
);
let block_id = self.residual_id_count;
self.residual_id_count += 1;
self.total_residual_dimension += new_residual_dimension;
block_id
}
pub fn remove_residual_block(&mut self, block_id: usize) -> Option<ResidualBlock> {
if let Some(residual_block) = self.residual_blocks.remove(&block_id) {
self.total_residual_dimension -= residual_block.factor.get_dimension();
Some(residual_block)
} else {
None
}
}
pub fn fix_variable(&mut self, var_to_fix: &str, idx: usize) {
if let Some(var_mut) = self.fixed_variable_indexes.get_mut(var_to_fix) {
var_mut.insert(idx);
} else {
self.fixed_variable_indexes
.insert(var_to_fix.to_owned(), HashSet::from([idx]));
}
}
pub fn unfix_variable(&mut self, var_to_unfix: &str) {
self.fixed_variable_indexes.remove(var_to_unfix);
}
pub fn set_variable_bounds(
&mut self,
var_to_bound: &str,
idx: usize,
lower_bound: f64,
upper_bound: f64,
) {
if lower_bound > upper_bound {
warn!("lower bound is larger than upper bound");
} else if let Some(var_mut) = self.variable_bounds.get_mut(var_to_bound) {
var_mut.insert(idx, (lower_bound, upper_bound));
} else {
self.variable_bounds.insert(
var_to_bound.to_owned(),
HashMap::from([(idx, (lower_bound, upper_bound))]),
);
}
}
pub fn remove_variable_bounds(&mut self, var_to_unbound: &str) {
self.variable_bounds.remove(var_to_unbound);
}
pub fn initialize_variables(
&self,
initial_values: &HashMap<String, (ManifoldType, DVector<f64>)>,
) -> HashMap<String, VariableEnum> {
let variables: HashMap<String, VariableEnum> = initial_values
.iter()
.map(|(k, v)| {
let variable_enum = match v.0 {
ManifoldType::SO2 => {
assert_eq!(
v.1.len(),
1,
"Variable '{}': SO2 expects 1 element, got {}",
k,
v.1.len()
);
let mut var = Variable::new(so2::SO2::from(v.1.clone()));
if let Some(indexes) = self.fixed_variable_indexes.get(k) {
var.fixed_indices = indexes.clone();
}
if let Some(bounds) = self.variable_bounds.get(k) {
var.bounds = bounds.clone();
}
VariableEnum::SO2(var)
}
ManifoldType::SO3 => {
assert_eq!(
v.1.len(),
4,
"Variable '{}': SO3 expects 4 elements, got {}",
k,
v.1.len()
);
let mut var = Variable::new(so3::SO3::from(v.1.clone()));
if let Some(indexes) = self.fixed_variable_indexes.get(k) {
var.fixed_indices = indexes.clone();
}
if let Some(bounds) = self.variable_bounds.get(k) {
var.bounds = bounds.clone();
}
VariableEnum::SO3(var)
}
ManifoldType::SE2 => {
assert_eq!(
v.1.len(),
3,
"Variable '{}': SE2 expects 3 elements, got {}",
k,
v.1.len()
);
let mut var = Variable::new(se2::SE2::from(v.1.clone()));
if let Some(indexes) = self.fixed_variable_indexes.get(k) {
var.fixed_indices = indexes.clone();
}
if let Some(bounds) = self.variable_bounds.get(k) {
var.bounds = bounds.clone();
}
VariableEnum::SE2(var)
}
ManifoldType::SE3 => {
assert_eq!(
v.1.len(),
7,
"Variable '{}': SE3 expects 7 elements, got {}",
k,
v.1.len()
);
let mut var = Variable::new(se3::SE3::from(v.1.clone()));
if let Some(indexes) = self.fixed_variable_indexes.get(k) {
var.fixed_indices = indexes.clone();
}
if let Some(bounds) = self.variable_bounds.get(k) {
var.bounds = bounds.clone();
}
VariableEnum::SE3(var)
}
ManifoldType::RN => {
let mut var = Variable::new(rn::Rn::from(v.1.clone()));
if let Some(indexes) = self.fixed_variable_indexes.get(k) {
var.fixed_indices = indexes.clone();
}
if let Some(bounds) = self.variable_bounds.get(k) {
var.bounds = bounds.clone();
}
VariableEnum::Rn(var)
}
ManifoldType::SE23 => {
let mut var = Variable::new(se23::SE23::from(v.1.clone()));
if let Some(indexes) = self.fixed_variable_indexes.get(k) {
var.fixed_indices = indexes.clone();
}
if let Some(bounds) = self.variable_bounds.get(k) {
var.bounds = bounds.clone();
}
VariableEnum::SE23(var)
}
ManifoldType::SGal3 => {
let mut var = Variable::new(sgal3::SGal3::from(v.1.clone()));
if let Some(indexes) = self.fixed_variable_indexes.get(k) {
var.fixed_indices = indexes.clone();
}
if let Some(bounds) = self.variable_bounds.get(k) {
var.bounds = bounds.clone();
}
VariableEnum::SGal3(var)
}
ManifoldType::Sim3 => {
let mut var = Variable::new(sim3::Sim3::from(v.1.clone()));
if let Some(indexes) = self.fixed_variable_indexes.get(k) {
var.fixed_indices = indexes.clone();
}
if let Some(bounds) = self.variable_bounds.get(k) {
var.bounds = bounds.clone();
}
VariableEnum::Sim3(var)
}
};
(k.to_owned(), variable_enum)
})
.collect();
variables
}
pub fn num_residual_blocks(&self) -> usize {
self.residual_blocks.len()
}
pub(crate) fn residual_blocks(&self) -> &HashMap<usize, ResidualBlock> {
&self.residual_blocks
}
pub fn compute_residual_sparse(
&self,
variables: &HashMap<String, VariableEnum>,
) -> CoreResult<Mat<f64>> {
let total_residual = Arc::new(Mutex::new(Col::<f64>::zeros(self.total_residual_dimension)));
let result: Result<Vec<()>, CoreError> = self
.residual_blocks
.par_iter()
.map(|(_, residual_block)| {
self.compute_residual_block(residual_block, variables, &total_residual)
})
.collect();
result?;
let total_residual = Arc::try_unwrap(total_residual)
.map_err(|_| {
CoreError::ParallelComputation(
"Failed to unwrap Arc for total residual".to_string(),
)
.log()
})?
.into_inner()
.map_err(|e| {
CoreError::ParallelComputation(
"Failed to extract mutex inner value for total residual".to_string(),
)
.log_with_source(e)
})?;
let residual_faer = total_residual.as_ref().as_mat().to_owned();
Ok(residual_faer)
}
pub fn compute_residual_and_jacobian_sparse(
&self,
variables: &HashMap<String, VariableEnum>,
variable_index_sparce_matrix: &HashMap<String, usize>,
symbolic_structure: &SymbolicStructure,
) -> CoreResult<(Mat<f64>, SparseColMat<usize, f64>)> {
Ok(crate::linearizer::cpu::sparse::assemble_sparse(
self,
variables,
variable_index_sparce_matrix,
symbolic_structure,
)?)
}
pub fn compute_residual_and_jacobian_dense(
&self,
variables: &HashMap<String, VariableEnum>,
variable_index_map: &HashMap<String, usize>,
total_dof: usize,
) -> CoreResult<(Mat<f64>, Mat<f64>)> {
Ok(crate::linearizer::cpu::dense::assemble_dense(
self,
variables,
variable_index_map,
total_dof,
)?)
}
fn compute_residual_block(
&self,
residual_block: &ResidualBlock,
variables: &HashMap<String, VariableEnum>,
total_residual: &Arc<Mutex<Col<f64>>>,
) -> CoreResult<()> {
let mut param_vectors: Vec<DVector<f64>> = Vec::new();
for var_key in &residual_block.variable_key_list {
if let Some(variable) = variables.get(var_key) {
param_vectors.push(variable.to_vector());
}
}
let (mut res, _) = residual_block.factor.linearize(¶m_vectors, false);
if let Some(loss_func) = &residual_block.loss_func {
let squared_norm = res.dot(&res);
let corrector = Corrector::new(loss_func.as_ref(), squared_norm);
corrector.correct_residuals(&mut res);
}
let mut total_residual = total_residual.lock().map_err(|e| {
CoreError::ParallelComputation("Failed to acquire lock on total residual".to_string())
.log_with_source(e)
})?;
let start_idx = residual_block.residual_row_start_idx;
let dim = residual_block.factor.get_dimension();
let mut total_residual_mut = total_residual.as_mut();
for i in 0..dim {
total_residual_mut[start_idx + i] = res[i];
}
Ok(())
}
pub fn log_residual_to_file(
&self,
residual: &nalgebra::DVector<f64>,
filename: &str,
) -> Result<(), Error> {
let mut file = File::create(filename)?;
writeln!(file, "# Residual vector - {} elements", residual.len())?;
for (i, &value) in residual.iter().enumerate() {
writeln!(file, "{}: {:.12}", i, value)?;
}
Ok(())
}
pub fn log_sparse_jacobian_to_file(
&self,
jacobian: &SparseColMat<usize, f64>,
filename: &str,
) -> Result<(), Error> {
let mut file = File::create(filename)?;
writeln!(
file,
"# Sparse Jacobian matrix - {} x {} ({} non-zeros)",
jacobian.nrows(),
jacobian.ncols(),
jacobian.compute_nnz()
)?;
writeln!(file, "# Matrix saved as dimensions and non-zero count only")?;
writeln!(file, "# For detailed access, convert to dense matrix first")?;
Ok(())
}
pub fn log_variables_to_file(
&self,
variables: &HashMap<String, VariableEnum>,
filename: &str,
) -> Result<(), Error> {
let mut file = File::create(filename)?;
writeln!(file, "# Variables - {} total", variables.len())?;
writeln!(file, "# Format: variable_name: [values...]")?;
let mut sorted_vars: Vec<_> = variables.keys().collect();
sorted_vars.sort();
for var_name in sorted_vars {
let var_vector = variables[var_name].to_vector();
write!(file, "{}: [", var_name)?;
for (i, &value) in var_vector.iter().enumerate() {
write!(file, "{:.12}", value)?;
if i < var_vector.len() - 1 {
write!(file, ", ")?;
}
}
writeln!(file, "]")?;
}
Ok(())
}
pub fn compute_and_set_covariances(
&self,
linear_solver: &mut Box<dyn LinearSolver<SparseMode>>,
variables: &mut HashMap<String, VariableEnum>,
variable_index_map: &HashMap<String, usize>,
) -> Option<HashMap<String, Mat<f64>>> {
linear_solver.compute_covariance_matrix()?;
let full_cov = linear_solver.get_covariance_matrix()?.clone();
let per_var_covariances =
extract_variable_covariances(&full_cov, variables, variable_index_map);
for (var_name, cov) in &per_var_covariances {
if let Some(var) = variables.get_mut(var_name) {
var.set_covariance(cov.clone());
}
}
Some(per_var_covariances)
}
pub fn compute_and_set_covariances_generic<M: crate::linalg::LinearizationMode>(
&self,
linear_solver: &mut dyn crate::linalg::LinearSolver<M>,
variables: &mut HashMap<String, VariableEnum>,
variable_index_map: &HashMap<String, usize>,
) -> Option<HashMap<String, Mat<f64>>> {
linear_solver.compute_covariance_matrix()?;
let full_cov = linear_solver.get_covariance_matrix()?.clone();
let per_var_covariances =
extract_variable_covariances(&full_cov, variables, variable_index_map);
for (var_name, cov) in &per_var_covariances {
if let Some(var) = variables.get_mut(var_name) {
var.set_covariance(cov.clone());
}
}
Some(per_var_covariances)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::loss_functions::HuberLoss;
use crate::factors::{BetweenFactor, PriorFactor};
use apex_manifolds::{ManifoldType, se2::SE2, se3::SE3};
use nalgebra::{Quaternion, Vector3, dvector};
use std::collections::HashMap;
type TestResult = Result<(), Box<dyn std::error::Error>>;
type TestProblemResult = Result<
(
Problem,
HashMap<String, (ManifoldType, nalgebra::DVector<f64>)>,
),
Box<dyn std::error::Error>,
>;
fn create_se2_test_problem() -> TestProblemResult {
let mut problem = Problem::new(JacobianMode::Sparse);
let mut initial_values = HashMap::new();
let poses = vec![
(0.0, 0.0, 0.0), (1.0, 0.0, 0.1), (1.5, 1.0, 0.5), (1.0, 2.0, 1.0), (0.0, 2.5, 1.5), (-1.0, 2.0, 2.0), (-1.5, 1.0, 2.5), (-1.0, 0.0, 3.0), (-0.5, -0.5, -2.8), (0.5, -0.5, -2.3), ];
for (i, (x, y, theta)) in poses.iter().enumerate() {
let var_name = format!("x{}", i);
let se2_data = dvector![*x, *y, *theta];
initial_values.insert(var_name, (ManifoldType::SE2, se2_data));
}
for i in 0..9 {
let from_pose = poses[i];
let to_pose = poses[i + 1];
let dx = to_pose.0 - from_pose.0;
let dy = to_pose.1 - from_pose.1;
let dtheta = to_pose.2 - from_pose.2;
let between_factor = BetweenFactor::new(SE2::from_xy_angle(dx, dy, dtheta));
problem.add_residual_block(
&[&format!("x{}", i), &format!("x{}", i + 1)],
Box::new(between_factor),
Some(Box::new(HuberLoss::new(1.0)?)),
);
}
let dx = poses[0].0 - poses[9].0;
let dy = poses[0].1 - poses[9].1;
let dtheta = poses[0].2 - poses[9].2;
let loop_closure = BetweenFactor::new(SE2::from_xy_angle(dx, dy, dtheta));
problem.add_residual_block(
&["x9", "x0"],
Box::new(loop_closure),
Some(Box::new(HuberLoss::new(1.0)?)),
);
let prior_factor = PriorFactor {
data: dvector![0.0, 0.0, 0.0],
};
problem.add_residual_block(&["x0"], Box::new(prior_factor), None);
Ok((problem, initial_values))
}
fn create_se3_test_problem() -> TestProblemResult {
let mut problem = Problem::new(JacobianMode::Sparse);
let mut initial_values = HashMap::new();
let poses = [
(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), (1.0, 0.0, 0.0, 0.0, 0.0, 0.1, 0.995), (1.0, 1.0, 0.0, 0.0, 0.0, 0.2, 0.98), (0.0, 1.0, 0.0, 0.0, 0.0, 0.3, 0.955), (0.0, 0.0, 1.0, 0.1, 0.0, 0.0, 0.995), (1.0, 0.0, 1.0, 0.1, 0.0, 0.1, 0.99), (1.0, 1.0, 1.0, 0.1, 0.0, 0.2, 0.975), (0.0, 1.0, 1.0, 0.1, 0.0, 0.3, 0.95), ];
for (i, (tx, ty, tz, qx, qy, qz, qw)) in poses.iter().enumerate() {
let var_name = format!("x{}", i);
let se3_data = dvector![*tx, *ty, *tz, *qw, *qx, *qy, *qz];
initial_values.insert(var_name, (ManifoldType::SE3, se3_data));
}
let edges = vec![
(0, 1),
(1, 2),
(2, 3),
(3, 0), (4, 5),
(5, 6),
(6, 7),
(7, 4), (0, 4),
(1, 5),
(2, 6),
(3, 7), ];
for (from_idx, to_idx) in edges {
let from_pose = poses[from_idx];
let to_pose = poses[to_idx];
let relative_se3 = SE3::from_translation_quaternion(
Vector3::new(
to_pose.0 - from_pose.0, to_pose.1 - from_pose.1, to_pose.2 - from_pose.2, ),
Quaternion::new(1.0, 0.0, 0.0, 0.0), );
let between_factor = BetweenFactor::new(relative_se3);
problem.add_residual_block(
&[&format!("x{}", from_idx), &format!("x{}", to_idx)],
Box::new(between_factor),
Some(Box::new(HuberLoss::new(1.0)?)),
);
}
let prior_factor = PriorFactor {
data: dvector![0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
};
problem.add_residual_block(&["x0"], Box::new(prior_factor), None);
Ok((problem, initial_values))
}
#[test]
fn test_problem_construction_se2() -> TestResult {
let (problem, initial_values) = create_se2_test_problem()?;
assert_eq!(problem.num_residual_blocks(), 11); assert_eq!(problem.total_residual_dimension, 33); assert_eq!(initial_values.len(), 10);
Ok(())
}
#[test]
fn test_problem_construction_se3() -> TestResult {
let (problem, initial_values) = create_se3_test_problem()?;
assert_eq!(problem.num_residual_blocks(), 13); assert_eq!(problem.total_residual_dimension, 79); assert_eq!(initial_values.len(), 8);
Ok(())
}
#[test]
fn test_variable_initialization_se2() -> TestResult {
let (problem, initial_values) = create_se2_test_problem()?;
let variables = problem.initialize_variables(&initial_values);
assert_eq!(variables.len(), 10);
for (name, var) in &variables {
assert_eq!(
var.get_size(),
3,
"SE2 variable {} should have size 3",
name
);
}
for (name, var) in &variables {
let vec = var.to_vector();
assert_eq!(
vec.len(),
3,
"SE2 variable {} vector should have length 3",
name
);
}
Ok(())
}
#[test]
fn test_variable_initialization_se3() -> TestResult {
let (problem, initial_values) = create_se3_test_problem()?;
let variables = problem.initialize_variables(&initial_values);
assert_eq!(variables.len(), 8);
for (name, var) in &variables {
assert_eq!(
var.get_size(),
6,
"SE3 variable {} should have size 6 (DOF)",
name
);
}
for (name, var) in &variables {
let vec = var.to_vector();
assert_eq!(
vec.len(),
7,
"SE3 variable {} vector should have length 7",
name
);
}
Ok(())
}
#[test]
fn test_column_mapping_se2() -> TestResult {
let (problem, initial_values) = create_se2_test_problem()?;
let variables = problem.initialize_variables(&initial_values);
let mut variable_index_sparce_matrix = HashMap::new();
let mut col_offset = 0;
let mut sorted_vars: Vec<_> = variables.keys().collect();
sorted_vars.sort();
for var_name in sorted_vars {
variable_index_sparce_matrix.insert(var_name.clone(), col_offset);
col_offset += variables[var_name].get_size();
}
let total_dof: usize = variables.values().map(|v| v.get_size()).sum();
assert_eq!(total_dof, 30); assert_eq!(col_offset, 30);
for (var_name, &col_idx) in &variable_index_sparce_matrix {
assert!(
col_idx < total_dof,
"Column index {} for {} should be < {}",
col_idx,
var_name,
total_dof
);
}
Ok(())
}
#[test]
fn test_symbolic_structure_se2() -> TestResult {
let (problem, initial_values) = create_se2_test_problem()?;
let variables = problem.initialize_variables(&initial_values);
let mut variable_index_sparce_matrix = HashMap::new();
let mut col_offset = 0;
let mut sorted_vars: Vec<_> = variables.keys().collect();
sorted_vars.sort();
for var_name in sorted_vars {
variable_index_sparce_matrix.insert(var_name.clone(), col_offset);
col_offset += variables[var_name].get_size();
}
let symbolic_structure = crate::linearizer::cpu::sparse::build_symbolic_structure(
&problem,
&variables,
&variable_index_sparce_matrix,
col_offset,
)?;
assert_eq!(
symbolic_structure.pattern.nrows(),
problem.total_residual_dimension
);
assert_eq!(symbolic_structure.pattern.ncols(), 30);
Ok(())
}
#[test]
fn test_residual_jacobian_computation_se2() -> TestResult {
let (problem, initial_values) = create_se2_test_problem()?;
let variables = problem.initialize_variables(&initial_values);
let mut variable_index_sparce_matrix = HashMap::new();
let mut col_offset = 0;
let mut sorted_vars: Vec<_> = variables.keys().collect();
sorted_vars.sort();
for var_name in sorted_vars {
variable_index_sparce_matrix.insert(var_name.clone(), col_offset);
col_offset += variables[var_name].get_size();
}
let symbolic_structure = crate::linearizer::cpu::sparse::build_symbolic_structure(
&problem,
&variables,
&variable_index_sparce_matrix,
col_offset,
)?;
let (residual_sparse, jacobian_sparse) = problem.compute_residual_and_jacobian_sparse(
&variables,
&variable_index_sparce_matrix,
&symbolic_structure,
)?;
assert_eq!(residual_sparse.nrows(), problem.total_residual_dimension);
assert_eq!(jacobian_sparse.nrows(), problem.total_residual_dimension);
assert_eq!(jacobian_sparse.ncols(), 30);
Ok(())
}
#[test]
fn test_residual_jacobian_computation_se3() -> TestResult {
let (problem, initial_values) = create_se3_test_problem()?;
let variables = problem.initialize_variables(&initial_values);
let mut variable_index_sparce_matrix = HashMap::new();
let mut col_offset = 0;
let mut sorted_vars: Vec<_> = variables.keys().collect();
sorted_vars.sort();
for var_name in sorted_vars {
variable_index_sparce_matrix.insert(var_name.clone(), col_offset);
col_offset += variables[var_name].get_size();
}
let symbolic_structure = crate::linearizer::cpu::sparse::build_symbolic_structure(
&problem,
&variables,
&variable_index_sparce_matrix,
col_offset,
)?;
let (residual_sparse, jacobian_sparse) = problem.compute_residual_and_jacobian_sparse(
&variables,
&variable_index_sparce_matrix,
&symbolic_structure,
)?;
assert_eq!(residual_sparse.nrows(), problem.total_residual_dimension);
assert_eq!(jacobian_sparse.nrows(), problem.total_residual_dimension);
assert_eq!(jacobian_sparse.ncols(), 48);
Ok(())
}
#[test]
fn test_residual_block_operations() -> TestResult {
let mut problem = Problem::new(JacobianMode::Sparse);
let block_id1 = problem.add_residual_block(
&["x0", "x1"],
Box::new(BetweenFactor::new(SE2::from_xy_angle(1.0, 0.0, 0.1))),
Some(Box::new(HuberLoss::new(1.0)?)),
);
let block_id2 = problem.add_residual_block(
&["x0"],
Box::new(PriorFactor {
data: dvector![0.0, 0.0, 0.0],
}),
None,
);
assert_eq!(problem.num_residual_blocks(), 2);
assert_eq!(problem.total_residual_dimension, 6); assert_eq!(block_id1, 0);
assert_eq!(block_id2, 1);
let removed_block = problem.remove_residual_block(block_id1);
assert!(removed_block.is_some());
assert_eq!(problem.num_residual_blocks(), 1);
assert_eq!(problem.total_residual_dimension, 3);
let non_existent = problem.remove_residual_block(999);
assert!(non_existent.is_none());
Ok(())
}
#[test]
fn test_variable_constraints() -> TestResult {
let mut problem = Problem::new(JacobianMode::Sparse);
problem.fix_variable("x0", 0);
problem.fix_variable("x0", 1);
problem.fix_variable("x1", 2);
assert!(problem.fixed_variable_indexes.contains_key("x0"));
assert!(problem.fixed_variable_indexes.contains_key("x1"));
assert_eq!(problem.fixed_variable_indexes["x0"].len(), 2);
assert_eq!(problem.fixed_variable_indexes["x1"].len(), 1);
problem.unfix_variable("x0");
assert!(!problem.fixed_variable_indexes.contains_key("x0"));
assert!(problem.fixed_variable_indexes.contains_key("x1"));
problem.set_variable_bounds("x2", 0, -1.0, 1.0);
problem.set_variable_bounds("x2", 1, -2.0, 2.0);
problem.set_variable_bounds("x3", 0, 0.0, 5.0);
assert!(problem.variable_bounds.contains_key("x2"));
assert!(problem.variable_bounds.contains_key("x3"));
assert_eq!(problem.variable_bounds["x2"].len(), 2);
assert_eq!(problem.variable_bounds["x3"].len(), 1);
problem.remove_variable_bounds("x2");
assert!(!problem.variable_bounds.contains_key("x2"));
assert!(problem.variable_bounds.contains_key("x3"));
Ok(())
}
#[test]
fn test_problem_default_equals_new_sparse() {
let default_problem = Problem::default();
let new_problem = Problem::new(JacobianMode::Sparse);
assert_eq!(default_problem.jacobian_mode, new_problem.jacobian_mode);
assert_eq!(
default_problem.num_residual_blocks(),
new_problem.num_residual_blocks()
);
}
#[test]
fn test_initialize_so2_variable() -> TestResult {
let problem = Problem::new(JacobianMode::Sparse);
let mut initial = HashMap::new();
initial.insert("angle".to_string(), (ManifoldType::SO2, dvector![0.5]));
let variables = problem.initialize_variables(&initial);
assert_eq!(variables.len(), 1);
let var = variables.get("angle").ok_or("angle not found")?;
assert_eq!(var.manifold_type(), ManifoldType::SO2);
assert_eq!(var.get_size(), 1);
assert_eq!(var.to_vector().len(), 1);
Ok(())
}
#[test]
fn test_initialize_so3_variable() -> TestResult {
let problem = Problem::new(JacobianMode::Sparse);
let mut initial = HashMap::new();
initial.insert(
"rot".to_string(),
(ManifoldType::SO3, dvector![1.0, 0.0, 0.0, 0.0]),
);
let variables = problem.initialize_variables(&initial);
assert_eq!(variables.len(), 1);
let var = variables.get("rot").ok_or("rot not found")?;
assert_eq!(var.manifold_type(), ManifoldType::SO3);
assert_eq!(var.get_size(), 3); Ok(())
}
#[test]
fn test_initialize_rn_variable() -> TestResult {
let problem = Problem::new(JacobianMode::Sparse);
let mut initial = HashMap::new();
initial.insert(
"pt".to_string(),
(ManifoldType::RN, dvector![1.0, 2.0, 3.0]),
);
let variables = problem.initialize_variables(&initial);
let var = variables.get("pt").ok_or("pt not found")?;
assert_eq!(var.manifold_type(), ManifoldType::RN);
assert_eq!(var.get_size(), 3);
let vec = var.to_vector();
assert!((vec[0] - 1.0).abs() < 1e-12);
assert!((vec[1] - 2.0).abs() < 1e-12);
assert!((vec[2] - 3.0).abs() < 1e-12);
Ok(())
}
#[test]
fn test_variable_enum_covariance_lifecycle() -> TestResult {
use faer::Mat;
let problem = Problem::new(JacobianMode::Sparse);
let mut initial = HashMap::new();
initial.insert(
"x0".to_string(),
(ManifoldType::SE2, dvector![0.0, 0.0, 0.0]),
);
let mut variables = problem.initialize_variables(&initial);
let var = variables.get_mut("x0").ok_or("x0 not found")?;
assert!(var.get_covariance().is_none());
let cov = Mat::identity(3, 3);
var.set_covariance(cov);
let retrieved = var.get_covariance().ok_or("covariance not set")?;
assert_eq!(retrieved.nrows(), 3);
assert!((retrieved[(0, 0)] - 1.0).abs() < 1e-12);
var.clear_covariance();
assert!(var.get_covariance().is_none());
Ok(())
}
#[test]
fn test_variable_enum_bounds_propagation() -> TestResult {
let mut problem = Problem::new(JacobianMode::Sparse);
problem.set_variable_bounds("x0", 0, -1.0, 1.0);
problem.set_variable_bounds("x0", 1, -2.0, 2.0);
let mut initial = HashMap::new();
initial.insert(
"x0".to_string(),
(ManifoldType::SE2, dvector![0.0, 0.0, 0.0]),
);
let variables = problem.initialize_variables(&initial);
let var = variables.get("x0").ok_or("x0 not found")?;
let bounds = var.get_bounds();
assert_eq!(bounds.len(), 2);
let (lo, hi) = bounds[&0];
assert!((lo + 1.0).abs() < 1e-12);
assert!((hi - 1.0).abs() < 1e-12);
Ok(())
}
#[test]
fn test_variable_enum_fixed_indices_propagation() -> TestResult {
let mut problem = Problem::new(JacobianMode::Sparse);
problem.fix_variable("x0", 0);
problem.fix_variable("x0", 2);
let mut initial = HashMap::new();
initial.insert(
"x0".to_string(),
(ManifoldType::SE2, dvector![0.0, 0.0, 0.0]),
);
let variables = problem.initialize_variables(&initial);
let var = variables.get("x0").ok_or("x0 not found")?;
let fixed = var.get_fixed_indices();
assert_eq!(fixed.len(), 2);
assert!(fixed.contains(&0));
assert!(fixed.contains(&2));
Ok(())
}
#[test]
fn test_variable_enum_set_from_vector_all_variants() -> TestResult {
let problem = Problem::new(JacobianMode::Sparse);
let mut initial = HashMap::new();
initial.insert(
"se2".to_string(),
(ManifoldType::SE2, dvector![0.0, 0.0, 0.0]),
);
initial.insert(
"se3".to_string(),
(
ManifoldType::SE3,
dvector![0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
),
);
initial.insert("so2".to_string(), (ManifoldType::SO2, dvector![0.0]));
initial.insert(
"so3".to_string(),
(ManifoldType::SO3, dvector![1.0, 0.0, 0.0, 0.0]),
);
initial.insert("rn".to_string(), (ManifoldType::RN, dvector![1.0, 2.0]));
let mut variables = problem.initialize_variables(&initial);
let new_se2 = dvector![1.0, 2.0, 0.5];
variables
.get_mut("se2")
.ok_or("se2 not found")?
.set_from_vector(&new_se2);
let got = variables.get("se2").ok_or("se2 not found")?.to_vector();
assert!((got[0] - 1.0).abs() < 1e-10);
let new_se3 = dvector![1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0];
variables
.get_mut("se3")
.ok_or("se3 not found")?
.set_from_vector(&new_se3);
let got = variables.get("se3").ok_or("se3 not found")?.to_vector();
assert!((got[0] - 1.0).abs() < 1e-10);
let new_so2 = dvector![0.5];
variables
.get_mut("so2")
.ok_or("so2 not found")?
.set_from_vector(&new_so2);
assert_eq!(variables.get("so2").ok_or("so2 not found")?.get_size(), 1);
let new_so3 = dvector![1.0, 0.0, 0.0, 0.0];
variables
.get_mut("so3")
.ok_or("so3 not found")?
.set_from_vector(&new_so3);
assert_eq!(variables.get("so3").ok_or("so3 not found")?.get_size(), 3);
let new_rn = dvector![5.0, 6.0];
variables
.get_mut("rn")
.ok_or("rn not found")?
.set_from_vector(&new_rn);
let got = variables.get("rn").ok_or("rn not found")?.to_vector();
assert!((got[0] - 5.0).abs() < 1e-10);
assert!((got[1] - 6.0).abs() < 1e-10);
Ok(())
}
#[test]
fn test_apply_tangent_step_se2_fixed_index() -> TestResult {
use faer::Mat;
let mut problem = Problem::new(JacobianMode::Sparse);
problem.fix_variable("x0", 0);
let mut initial = HashMap::new();
initial.insert(
"x0".to_string(),
(ManifoldType::SE2, dvector![0.0, 0.0, 0.0]),
);
let mut variables = problem.initialize_variables(&initial);
let mut step = Mat::zeros(3, 1);
step[(0, 0)] = 1.0;
step[(1, 0)] = 2.0;
step[(2, 0)] = 3.0;
let var = variables.get_mut("x0").ok_or("x0 not found")?;
var.apply_tangent_step(step.as_ref());
assert_eq!(var.get_size(), 3);
Ok(())
}
#[test]
fn test_apply_tangent_step_se3_fixed_index() -> TestResult {
use faer::Mat;
let mut problem = Problem::new(JacobianMode::Sparse);
problem.fix_variable("p", 0);
let mut initial = HashMap::new();
initial.insert(
"p".to_string(),
(
ManifoldType::SE3,
dvector![0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
),
);
let mut variables = problem.initialize_variables(&initial);
let mut step = Mat::zeros(6, 1);
for i in 0..6 {
step[(i, 0)] = (i + 1) as f64;
}
let var = variables.get_mut("p").ok_or("p not found")?;
var.apply_tangent_step(step.as_ref());
assert_eq!(var.get_size(), 6);
Ok(())
}
#[test]
fn test_apply_tangent_step_remaining_manifolds() -> TestResult {
use faer::Mat;
let problem = Problem::new(JacobianMode::Sparse);
let mut initial = HashMap::new();
initial.insert("so2".to_string(), (ManifoldType::SO2, dvector![0.0]));
initial.insert(
"so3".to_string(),
(ManifoldType::SO3, dvector![1.0, 0.0, 0.0, 0.0]),
);
initial.insert("rn".to_string(), (ManifoldType::RN, dvector![0.0, 0.0]));
let mut variables = problem.initialize_variables(&initial);
let mut step_so2 = Mat::zeros(1, 1);
step_so2[(0, 0)] = 0.1;
variables
.get_mut("so2")
.ok_or("so2 not found")?
.apply_tangent_step(step_so2.as_ref());
let mut step_so3 = Mat::zeros(3, 1);
step_so3[(0, 0)] = 0.01;
variables
.get_mut("so3")
.ok_or("so3 not found")?
.apply_tangent_step(step_so3.as_ref());
let mut step_rn = Mat::zeros(2, 1);
step_rn[(0, 0)] = 1.0;
step_rn[(1, 0)] = 2.0;
variables
.get_mut("rn")
.ok_or("rn not found")?
.apply_tangent_step(step_rn.as_ref());
let vec = variables.get("rn").ok_or("rn not found")?.to_vector();
assert!((vec[0] - 1.0).abs() < 1e-10);
assert!((vec[1] - 2.0).abs() < 1e-10);
Ok(())
}
#[test]
fn test_compute_residual_sparse_smoke() -> TestResult {
let (problem, initial_values) = create_se2_test_problem()?;
let variables = problem.initialize_variables(&initial_values);
let residual = problem.compute_residual_sparse(&variables)?;
let norm_sq: f64 = (0..residual.nrows())
.map(|i| residual[(i, 0)].powi(2))
.sum();
assert!(norm_sq >= 0.0);
assert_eq!(residual.nrows(), problem.total_residual_dimension);
Ok(())
}
#[test]
fn test_log_residual_to_file() -> TestResult {
let problem = Problem::new(JacobianMode::Sparse);
let residual = nalgebra::dvector![1.0, 2.0, 3.0];
let path = std::env::temp_dir().join("apex_test_residual.txt");
let path_str = path.to_str().ok_or("temp path is not valid UTF-8")?;
problem.log_residual_to_file(&residual, path_str)?;
assert!(path.exists());
Ok(())
}
#[test]
fn test_log_variables_to_file() -> TestResult {
let problem = Problem::new(JacobianMode::Sparse);
let mut initial = HashMap::new();
initial.insert(
"x0".to_string(),
(ManifoldType::SE2, dvector![1.0, 2.0, 0.3]),
);
let variables = problem.initialize_variables(&initial);
let path = std::env::temp_dir().join("apex_test_variables.txt");
let path_str = path.to_str().ok_or("temp path is not valid UTF-8")?;
problem.log_variables_to_file(&variables, path_str)?;
assert!(path.exists());
Ok(())
}
#[test]
fn test_log_sparse_jacobian_to_file() -> TestResult {
use faer::sparse::SparseColMat;
let problem = Problem::new(JacobianMode::Sparse);
let triplets = vec![faer::sparse::Triplet::new(0usize, 0usize, 1.0f64)];
let jacobian =
SparseColMat::try_new_from_triplets(1, 1, &triplets).map_err(|e| format!("{e:?}"))?;
let path = std::env::temp_dir().join("apex_test_jacobian.txt");
let path_str = path.to_str().ok_or("temp path is not valid UTF-8")?;
problem.log_sparse_jacobian_to_file(&jacobian, path_str)?;
assert!(path.exists());
Ok(())
}
#[test]
fn test_set_variable_bounds_invalid_order() {
let mut problem = Problem::new(JacobianMode::Sparse);
problem.set_variable_bounds("x0", 0, 5.0, 1.0);
assert!(!problem.variable_bounds.contains_key("x0"));
}
#[test]
fn test_variable_enum_manifold_type() {
let problem = Problem::new(JacobianMode::Sparse);
let mut initial = HashMap::new();
initial.insert(
"se2".to_string(),
(ManifoldType::SE2, dvector![0.0, 0.0, 0.0]),
);
initial.insert(
"se3".to_string(),
(
ManifoldType::SE3,
dvector![0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
),
);
initial.insert("so2".to_string(), (ManifoldType::SO2, dvector![0.0]));
initial.insert(
"so3".to_string(),
(ManifoldType::SO3, dvector![1.0, 0.0, 0.0, 0.0]),
);
initial.insert("rn".to_string(), (ManifoldType::RN, dvector![0.0]));
let variables = problem.initialize_variables(&initial);
assert_eq!(variables["se2"].manifold_type(), ManifoldType::SE2);
assert_eq!(variables["se3"].manifold_type(), ManifoldType::SE3);
assert_eq!(variables["so2"].manifold_type(), ManifoldType::SO2);
assert_eq!(variables["so3"].manifold_type(), ManifoldType::SO3);
assert_eq!(variables["rn"].manifold_type(), ManifoldType::RN);
}
}