Skip to content

Commit

Permalink
linear solvers
Browse files Browse the repository at this point in the history
  • Loading branch information
strasdat committed Jan 7, 2025
1 parent 0bfee58 commit 31e0b93
Show file tree
Hide file tree
Showing 22 changed files with 799 additions and 813 deletions.
1 change: 1 addition & 0 deletions crates/sophus_opt/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ nalgebra.workspace = true
ndarray.workspace = true
rand.workspace = true
rayon = "1.10"
snafu.workspace = true

[features]
simd = [
Expand Down
210 changes: 4 additions & 206 deletions crates/sophus_opt/src/block.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use nalgebra::Const;
use sophus_autodiff::linalg::MatF64;
use sophus_autodiff::linalg::VecF64;
/// Block Hessian matrix
pub mod block_hessian;
/// Block gradient vector
pub mod block_gradient;

/// Range of a block
#[derive(Clone, Debug, Copy, Default)]
Expand All @@ -10,206 +11,3 @@ pub struct BlockRange {
/// Dimension of the block
pub dim: usize,
}

/// Block vector
#[derive(Debug, Clone)]
pub struct BlockVector<const NUM: usize, const NUM_ARGS: usize> {
/// vector storage
pub vec: nalgebra::SVector<f64, NUM>,
/// ranges, one for each block
pub ranges: [BlockRange; NUM_ARGS],
}

impl<const NUM: usize, const NUM_ARGS: usize> BlockVector<NUM, NUM_ARGS> {
/// create a new block vector
pub fn new(dims: &[usize; NUM_ARGS]) -> Self {
debug_assert!(!dims.is_empty());

let num_blocks = dims.len();

let mut ranges = [BlockRange::default(); NUM_ARGS];

let mut num_rows: usize = 0;

for i in 0..num_blocks {
let dim = dims[i];
ranges[i] = BlockRange {
index: num_rows as i64,
dim,
};
num_rows += dim;
}
Self {
vec: nalgebra::SVector::zeros(),
ranges,
}
}

/// Number of blocks
pub fn num_blocks(&self) -> usize {
self.ranges.len()
}

/// set the ith block
pub fn set_block<const R: usize>(&mut self, ith: usize, v: VecF64<R>) {
debug_assert!(ith < self.num_blocks());
debug_assert_eq!(R, self.ranges[ith].dim);
self.mut_block::<R>(ith).copy_from(&v);
}

/// get the ith block
pub fn block(
&self,
ith: usize,
) -> nalgebra::Matrix<
f64,
nalgebra::Dyn,
nalgebra::Const<1>,
nalgebra::ViewStorage<
'_,
f64,
nalgebra::Dyn,
nalgebra::Const<1>,
nalgebra::Const<1>,
Const<NUM>,
>,
> {
let idx = self.ranges[ith].index as usize;
self.vec.rows(idx, self.ranges[ith].dim)
}

/// mutable reference to the ith block
pub fn mut_block<const ROWS: usize>(
&mut self,
ith: usize,
) -> nalgebra::Matrix<
f64,
nalgebra::Const<ROWS>,
nalgebra::Const<1>,
nalgebra::ViewStorageMut<
'_,
f64,
nalgebra::Const<ROWS>,
nalgebra::Const<1>,
nalgebra::Const<1>,
Const<NUM>,
>,
> {
let idx = self.ranges[ith].index as usize;
self.vec.fixed_rows_mut::<ROWS>(idx)
}
}

/// Block matrix
#[derive(Clone, Debug)]
pub struct BlockMatrix<const NUM: usize, const NUM_ARGS: usize> {
/// matrix storage
pub mat: nalgebra::SMatrix<f64, NUM, NUM>,
/// ranges, one for each block
pub ranges: [BlockRange; NUM_ARGS],
}

impl<const NUM: usize, const NUM_ARGS: usize> BlockMatrix<NUM, NUM_ARGS> {
/// create a new block matrix
pub fn new(dims: &[usize]) -> Self {
debug_assert!(!dims.is_empty());

let num_blocks = dims.len();

let mut ranges = [BlockRange::default(); NUM_ARGS];
let mut num_rows: usize = 0;

for i in 0..num_blocks {
let dim = dims[i];
ranges[i] = BlockRange {
index: num_rows as i64,
dim,
};
num_rows += dim;
}
Self {
mat: nalgebra::SMatrix::zeros(),
ranges,
}
}

/// Number of blocks
pub fn num_blocks(&self) -> usize {
self.ranges.len()
}

/// set block (i, j)
pub fn set_block<const R: usize, const C: usize>(
&mut self,
ith: usize,
jth: usize,
m: MatF64<R, C>,
) {
debug_assert!(ith < self.num_blocks());
debug_assert!(jth < self.num_blocks());
debug_assert_eq!(R, self.ranges[ith].dim);
debug_assert_eq!(C, self.ranges[jth].dim);

if ith == jth {
debug_assert_eq!(R, C);

let mut mut_block = self.mut_block::<R, C>(ith, jth);
mut_block.copy_from(&m);
} else {
debug_assert!(ith < jth);
{
self.mut_block::<R, C>(ith, jth).copy_from(&m);
}
{
self.mut_block::<C, R>(jth, ith).copy_from(&m.transpose());
}
}
}

/// get block (i, j)
pub fn block(
&self,
ith: usize,
jth: usize,
) -> nalgebra::Matrix<
f64,
nalgebra::Dyn,
nalgebra::Dyn,
nalgebra::ViewStorage<
'_,
f64,
nalgebra::Dyn,
nalgebra::Dyn,
nalgebra::Const<1>,
nalgebra::Const<NUM>,
>,
> {
let idx_i = self.ranges[ith].index as usize;
let idx_j = self.ranges[jth].index as usize;
self.mat
.view((idx_i, idx_j), (self.ranges[ith].dim, self.ranges[jth].dim))
}

/// mutable reference to block (i, j)
pub fn mut_block<const R: usize, const C: usize>(
&mut self,
ith: usize,
jth: usize,
) -> nalgebra::Matrix<
f64,
nalgebra::Const<R>,
nalgebra::Const<C>,
nalgebra::ViewStorageMut<
'_,
f64,
nalgebra::Const<R>,
nalgebra::Const<C>,
nalgebra::Const<1>,
nalgebra::Const<NUM>,
>,
> {
let idx_i = self.ranges[ith].index as usize;
let idx_j = self.ranges[jth].index as usize;
self.mat.fixed_view_mut::<R, C>(idx_i, idx_j)
}
}
92 changes: 92 additions & 0 deletions crates/sophus_opt/src/block/block_gradient.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
use super::BlockRange;
use nalgebra::Const;
use sophus_autodiff::linalg::VecF64;

/// Block gradient vector
#[derive(Debug, Clone)]
pub struct BlockGradient<const NUM: usize, const NUM_ARGS: usize> {
/// vector storage
pub vec: nalgebra::SVector<f64, NUM>,
/// ranges, one for each block
pub ranges: [BlockRange; NUM_ARGS],
}

impl<const NUM: usize, const NUM_ARGS: usize> BlockGradient<NUM, NUM_ARGS> {
/// create a new block vector
pub fn new(dims: &[usize; NUM_ARGS]) -> Self {
debug_assert!(!dims.is_empty());

let num_blocks = dims.len();

let mut ranges = [BlockRange::default(); NUM_ARGS];

let mut num_rows: usize = 0;

for i in 0..num_blocks {
let dim = dims[i];
ranges[i] = BlockRange {
index: num_rows as i64,
dim,
};
num_rows += dim;
}
Self {
vec: nalgebra::SVector::zeros(),
ranges,
}
}

/// Number of blocks
pub fn num_blocks(&self) -> usize {
self.ranges.len()
}

/// set the ith block
pub fn set_block<const R: usize>(&mut self, ith: usize, v: VecF64<R>) {
debug_assert!(ith < self.num_blocks());
debug_assert_eq!(R, self.ranges[ith].dim);
self.mut_block::<R>(ith).copy_from(&v);
}

/// get the ith block
pub fn block(
&self,
ith: usize,
) -> nalgebra::Matrix<
f64,
nalgebra::Dyn,
nalgebra::Const<1>,
nalgebra::ViewStorage<
'_,
f64,
nalgebra::Dyn,
nalgebra::Const<1>,
nalgebra::Const<1>,
Const<NUM>,
>,
> {
let idx = self.ranges[ith].index as usize;
self.vec.rows(idx, self.ranges[ith].dim)
}

/// mutable reference to the ith block
pub fn mut_block<const ROWS: usize>(
&mut self,
ith: usize,
) -> nalgebra::Matrix<
f64,
nalgebra::Const<ROWS>,
nalgebra::Const<1>,
nalgebra::ViewStorageMut<
'_,
f64,
nalgebra::Const<ROWS>,
nalgebra::Const<1>,
nalgebra::Const<1>,
Const<NUM>,
>,
> {
let idx = self.ranges[ith].index as usize;
self.vec.fixed_rows_mut::<ROWS>(idx)
}
}
Loading

0 comments on commit 31e0b93

Please sign in to comment.