├── .github └── workflows │ └── rust.yml ├── .gitignore ├── Cargo.toml ├── LICENSE ├── README.md ├── docs ├── m3500_rs.png ├── parking_garage.png └── sphere2500_rs.png ├── examples ├── m3500_benchmark.rs ├── parking-garage.rs ├── python │ ├── m3500.py │ ├── small_problem.py │ └── try_import.py ├── small_problem.rs └── sphere2500.rs ├── pyproject.toml ├── scripts ├── compile_aarch64.sh ├── compile_local.sh └── fmt.sh ├── src ├── corrector.rs ├── factors.rs ├── helper.rs ├── lib.rs ├── linear │ ├── mod.rs │ ├── sparse.rs │ ├── sparse_cholesky.rs │ └── sparse_qr.rs ├── loss_functions.rs ├── manifold │ ├── mod.rs │ ├── se3.rs │ └── so3.rs ├── optimizer │ ├── common.rs │ ├── gauss_newton_optimizer.rs │ ├── levenberg_marquardt_optimizer.rs │ └── mod.rs ├── parameter_block.rs ├── problem.rs ├── python │ ├── mod.rs │ ├── py_factors.rs │ ├── py_linear.rs │ ├── py_loss_functions.rs │ ├── py_optimizer.rs │ └── py_problem.rs └── residual_block.rs ├── tests ├── data │ ├── input_M3500_g2o.g2o │ ├── parking-garage.g2o │ └── sphere2500.g2o ├── test_factors.rs ├── test_loss.rs ├── test_manifold.rs ├── test_problem.rs └── test_tiny_solver.rs └── tiny_solver ├── __init__.py ├── factors └── __init__.pyi ├── loss_functions └── __init__.pyi ├── py.typed └── tiny_solver.pyi /.github/workflows/rust.yml: -------------------------------------------------------------------------------- 1 | name: Rust 2 | 3 | on: 4 | push: 5 | branches: [ "master" ] 6 | pull_request: 7 | branches: [ "master" ] 8 | 9 | env: 10 | CARGO_TERM_COLOR: always 11 | 12 | jobs: 13 | build: 14 | 15 | runs-on: ubuntu-latest 16 | 17 | steps: 18 | - uses: actions/checkout@v4 19 | - name: Dependencies 20 | run: sudo apt install pkg-config libfreetype6-dev libfontconfig1-dev 21 | - name: Format 22 | run: cargo fmt --check --verbose 23 | - name: Linting 24 | run: cargo clippy 25 | - name: Build 26 | run: cargo build --verbose 27 | - name: Run tests 28 | run: cargo test --verbose 29 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | *.pyc 3 | __pycache__ 4 | *.so 5 | .DS_Store 6 | .vscode/* 7 | .venv 8 | .pytest_cache 9 | Cargo.lock 10 | m3500_rs.png 11 | .vscode/ 12 | venv/ -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "tiny-solver" 3 | version = "0.17.1" 4 | edition = "2021" 5 | authors = ["Powei Lin , Hossam R. "] 6 | readme = "README.md" 7 | license = "MIT OR Apache-2.0" 8 | description = "Factor graph solver" 9 | homepage = "https://github.com/powei-lin/tiny-solver-rs" 10 | repository = "https://github.com/powei-lin/tiny-solver-rs" 11 | keywords = ["factor-graph", "ceres-solver", "minisam"] 12 | categories = ["data-structures", "science", "mathematics"] 13 | exclude = ["/.github/*", "*.ipynb", "./scripts/*", "examples/*", "tests/"] 14 | 15 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 16 | 17 | [dependencies] 18 | faer = "0.22.4" 19 | faer-ext = { version = "0.6.0", features = ["nalgebra"] } 20 | log = "0.4.27" 21 | nalgebra = "0.33.2" 22 | num-dual = "0.11.1" 23 | num-traits = "0.2.19" 24 | numpy = { version = "0.23.0", features = ["nalgebra"], optional = true } 25 | pyo3 = { version = "0.23.3", features = ["abi3", "abi3-py38"], optional = true } 26 | # pyo3-log = { version = "0.9.0", optional = true } 27 | rayon = "1.9.0" 28 | simba = "0.9.0" 29 | 30 | [[example]] 31 | name = "m3500_benchmark" 32 | path = "examples/m3500_benchmark.rs" 33 | 34 | [[example]] 35 | name = "sphere2500" 36 | 37 | [[example]] 38 | name = "parking-garage" 39 | 40 | [features] 41 | python = ["num-dual/python", "numpy", "pyo3"] 42 | 43 | [dev-dependencies] 44 | env_logger = "0.11.8" 45 | itertools = "0.14.0" 46 | nalgebra = { version = "0.33.2", features = ["rand"] } 47 | plotters = "0.3.6" 48 | rand = "0.9.1" 49 | 50 | [profile.dev.package.faer] 51 | opt-level = 3 52 | 53 | [lib] 54 | name = "tiny_solver" 55 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Powei Lin 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # tiny-solver-rs 2 | [![crate](https://img.shields.io/crates/v/tiny-solver.svg)](https://crates.io/crates/tiny-solver) 3 | [![crate downloads](https://img.shields.io/crates/d/tiny-solver)](https://crates.io/crates/tiny-solver) 4 | 5 | Inspired by [ceres-solver](https://github.com/ceres-solver/ceres-solver), [tiny-solver](https://github.com/keir/tinysolver), and [minisam](https://github.com/dongjing3309/minisam). 6 | 7 | This is a general optimizer written in Rust, including bindings for Python. If you're familiar with ceres-solver or factor-graph optimizers, you'll find it very easy to use. 8 | 9 | #### Other great rust optimizers 10 | * [fact.rs](https://github.com/rpl-cmu/fact.rs) 11 | * [sophus-rs](https://github.com/sophus-vision/sophus-rs) 12 | 13 | ## Installation 14 | ### rust 15 | ```sh 16 | cargo add tiny-solver 17 | ``` 18 | 19 | ## Current Features 20 | 21 | - [x] Automatic Derivatives using [num-dual](https://github.com/itt-ustutt/num-dual) 22 | - [x] Sparse QR, Sparse Cholesky using [faer](https://github.com/sarah-quinones/faer-rs) 23 | - [x] GaussNewtonOptimizer 24 | - [x] LevenbergMarquardtOptimizer 25 | - [x] Multithreading jacobian 26 | - [x] loss functions (Huber, CauchyLoss, ArctanLoss) 27 | - [x] Parameter on manifold (SO3, SE3) 28 | 29 | #### TODO 30 | - [ ] information matrix 31 | 32 | ## Benchmark 33 | On m3 macbook air 34 | | dataset | tiny-solver | gtsam | minisam | 35 | |---------|-------------|---------|----------| 36 | | m3500 | 128.3ms | 130.7ms | 123.6 ms | 37 | 38 | It's not extremely optimized, but it's easy to install and use. 39 | 40 | ## Usage 41 | Rust 42 | ```rust 43 | // define your own Cost/Factor struct 44 | // impl residual function 45 | // and the jacobian will be auto generated 46 | use nalgebra as na; 47 | struct CustomFactor {} 48 | impl tiny_solver::factors::Factor for CustomFactor { 49 | fn residual_func(&self, params: &[nalgebra::DVector]) -> nalgebra::DVector { 50 | let x = ¶ms[0][0]; 51 | let y = ¶ms[1][0]; 52 | let z = ¶ms[1][1]; 53 | 54 | na::dvector![ 55 | x.clone() 56 | + y.clone() * T::from_f64(2.0).unwrap() 57 | + z.clone() * T::from_f64(4.0).unwrap(), 58 | y.clone() * z.clone() 59 | ] 60 | } 61 | } 62 | 63 | fn main() { 64 | // init logger, `export RUST_LOG=trace` to see more log 65 | env_logger::init(); 66 | 67 | // init problem (factor graph) 68 | let mut problem = tiny_solver::Problem::new(); 69 | 70 | // add residual blocks (factors) 71 | // add residual x needs to be close to 3.0 72 | problem.add_residual_block( 73 | 1, 74 | &["x"], 75 | Box::new(tiny_solver::factors::PriorFactor { 76 | v: na::dvector![3.0], 77 | }), 78 | None, 79 | ); 80 | // add custom residual for x and yz 81 | problem.add_residual_block(2, &["x", "yz"], Box::new(CustomFactor), None); 82 | 83 | // the initial values for x is 0.7 and yz is [-30.2, 123.4] 84 | let initial_values = HashMap::>::from([ 85 | ("x".to_string(), na::dvector![0.7]), 86 | ("yz".to_string(), na::dvector![-30.2, 123.4]), 87 | ]); 88 | 89 | // initialize optimizer 90 | let optimizer = tiny_solver::GaussNewtonOptimizer::new(); 91 | 92 | // optimize 93 | let result = optimizer.optimize(&problem, &initial_values, None); 94 | 95 | // result 96 | for (k, v) in result { 97 | println!("{}: {}", k, v); 98 | } 99 | } 100 | ``` 101 | You can find other examples in 102 | * https://github.com/powei-lin/camera-intrinsic-calibration-rs/blob/main/src/util.rs 103 | * https://github.com/powei-lin/b-spline/blob/main/src/so3bspline.rs 104 | 105 | Python (Currently not maintaining) 106 | ```py 107 | import numpy as np 108 | from tiny_solver import Problem, GaussNewtonOptimizer 109 | from tiny_solver.factors import PriorFactor, PyFactor 110 | 111 | # define custom cost function in python 112 | # the trade off is the jacobian for the problem cannot be done in parallel 113 | # because of gil 114 | def cost(x: np.ndarray, yz: np.ndarray) -> np.ndarray: 115 | r0 = x[0] + 2 * yz[0] + 4 * yz[1] 116 | r1 = yz[0] * yz[0] 117 | return np.array([r0, r1]) 118 | 119 | 120 | def main(): 121 | 122 | # initialize problem (factor graph) 123 | problem = Problem() 124 | 125 | # factor defined in python 126 | custom_factor = PyFactor(cost) 127 | problem.add_residual_block( 128 | 2, 129 | [ 130 | ("x", 1), 131 | ("yz", 2), 132 | ], 133 | custom_factor, 134 | None, 135 | ) 136 | 137 | # prior factor import from rust 138 | prior_factor = PriorFactor(np.array([3.0])) 139 | problem.add_residual_block(1, [("x", 1)], prior_factor, None) 140 | 141 | # initial values 142 | init_values = {"x": np.array([0.7]), "yz": np.array([-30.2, 123.4])} 143 | 144 | # optimizer 145 | optimizer = GaussNewtonOptimizer() 146 | result_values = optimizer.optimize(problem, init_values) 147 | 148 | # result 149 | for k, v in result_values.items(): 150 | print(f"{k}: {v}") 151 | 152 | 153 | if __name__ == "__main__": 154 | main() 155 | ``` 156 | 157 | ## Example 158 | ### Basic example 159 | ```sh 160 | cargo run -r --example small_problem 161 | ``` 162 | ### M3500 dataset 163 | m3500 dataset rust result. 164 | 165 | ```sh 166 | git clone https://github.com/powei-lin/tiny-solver-rs.git 167 | cd tiny-solver-rs 168 | 169 | # run rust version 170 | cargo run -r --example m3500_benchmark 171 | 172 | # run python version 173 | pip install tiny-solver matplotlib 174 | python3 examples/python/m3500.py 175 | ``` 176 | ### Sphere 2500 dataset 177 | ``` 178 | cargo run -r --example sphere2500 179 | ``` 180 | sp2500 dataset rust result. 181 | 182 | ### Parking garage dataset 183 | ``` 184 | cargo run -r --example parking-garage 185 | ``` 186 | parking dataset rust result. 187 | -------------------------------------------------------------------------------- /docs/m3500_rs.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/powei-lin/tiny-solver-rs/094690931a475fcd17a6d5fe9968daf7ece00cf0/docs/m3500_rs.png -------------------------------------------------------------------------------- /docs/parking_garage.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/powei-lin/tiny-solver-rs/094690931a475fcd17a6d5fe9968daf7ece00cf0/docs/parking_garage.png -------------------------------------------------------------------------------- /docs/sphere2500_rs.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/powei-lin/tiny-solver-rs/094690931a475fcd17a6d5fe9968daf7ece00cf0/docs/sphere2500_rs.png -------------------------------------------------------------------------------- /examples/m3500_benchmark.rs: -------------------------------------------------------------------------------- 1 | use std::time::Instant; 2 | 3 | use plotters::prelude::*; 4 | 5 | use tiny_solver::{helper::read_g2o, optimizer::Optimizer, GaussNewtonOptimizer}; 6 | 7 | fn main() { 8 | // init logger 9 | env_logger::init(); 10 | 11 | let (problem, init_values) = read_g2o("tests/data/input_M3500_g2o.g2o"); 12 | let init_points: Vec<(f64, f64)> = init_values.values().map(|v| (v[1], v[2])).collect(); 13 | let root_drawing_area = BitMapBackend::new("m3500_rs.png", (1024, 1024)).into_drawing_area(); 14 | 15 | root_drawing_area.fill(&WHITE).unwrap(); 16 | 17 | let mut scatter_ctx = ChartBuilder::on(&root_drawing_area) 18 | .x_label_area_size(40) 19 | .y_label_area_size(40) 20 | .build_cartesian_2d(-50f64..50f64, -80f64..20f64) 21 | .unwrap(); 22 | scatter_ctx 23 | .configure_mesh() 24 | .disable_x_mesh() 25 | .disable_y_mesh() 26 | .draw() 27 | .unwrap(); 28 | scatter_ctx 29 | .draw_series( 30 | init_points 31 | .iter() 32 | .map(|(x, y)| Circle::new((*x, *y), 2, GREEN.filled())), 33 | ) 34 | .unwrap(); 35 | let gn = GaussNewtonOptimizer::new(); 36 | let start = Instant::now(); 37 | let result = gn.optimize(&problem, &init_values, None); 38 | let duration = start.elapsed(); 39 | println!("Time elapsed in total is: {:?}", duration); 40 | let result_points: Vec<(f64, f64)> = result.unwrap().values().map(|v| (v[1], v[2])).collect(); 41 | scatter_ctx 42 | .draw_series( 43 | result_points 44 | .iter() 45 | .map(|(x, y)| Circle::new((*x, *y), 2, BLUE.filled())), 46 | ) 47 | .unwrap(); 48 | } 49 | -------------------------------------------------------------------------------- /examples/parking-garage.rs: -------------------------------------------------------------------------------- 1 | use std::time::Instant; 2 | 3 | use plotters::prelude::*; 4 | 5 | use tiny_solver::helper::read_g2o; 6 | use tiny_solver::{optimizer::Optimizer, GaussNewtonOptimizer}; 7 | 8 | fn main() { 9 | // init logger 10 | env_logger::init(); 11 | 12 | let (problem, init_values) = read_g2o("tests/data/parking-garage.g2o"); 13 | let init_points: Vec<(f64, f64)> = init_values.values().map(|v| (v[4], v[5])).collect(); 14 | let root_drawing_area = 15 | BitMapBackend::new("parking_garage.png", (1024, 1024)).into_drawing_area(); 16 | 17 | root_drawing_area.fill(&WHITE).unwrap(); 18 | 19 | let mut scatter_ctx = ChartBuilder::on(&root_drawing_area) 20 | .x_label_area_size(40) 21 | .y_label_area_size(40) 22 | .build_cartesian_2d(-220f64..80f64, -20f64..280f64) 23 | .unwrap(); 24 | scatter_ctx 25 | .configure_mesh() 26 | .disable_x_mesh() 27 | .disable_y_mesh() 28 | .draw() 29 | .unwrap(); 30 | scatter_ctx 31 | .draw_series( 32 | init_points 33 | .iter() 34 | .map(|(x, y)| Circle::new((*x, *y), 2, GREEN.filled())), 35 | ) 36 | .unwrap(); 37 | let gn = GaussNewtonOptimizer::new(); 38 | let start = Instant::now(); 39 | let result = gn.optimize(&problem, &init_values, None); 40 | let duration = start.elapsed(); 41 | println!("Time elapsed in total is: {:?}", duration); 42 | let result_points: Vec<(f64, f64)> = result.unwrap().values().map(|v| (v[4], v[5])).collect(); 43 | scatter_ctx 44 | .draw_series( 45 | result_points 46 | .iter() 47 | .map(|(x, y)| Circle::new((*x, *y), 2, BLUE.filled())), 48 | ) 49 | .unwrap(); 50 | } 51 | -------------------------------------------------------------------------------- /examples/python/m3500.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from time import perf_counter 3 | from typing import Dict, Tuple 4 | 5 | import matplotlib.pyplot as plt 6 | import numpy as np 7 | 8 | from tiny_solver import GaussNewtonOptimizer, Problem, OptimizerOptions 9 | from tiny_solver.factors import PriorFactor, BetweenFactorSE2 10 | from tiny_solver.loss_functions import HuberLoss 11 | 12 | 13 | def load_g2o(file_path: str) -> Tuple[Problem, Dict[str, np.ndarray]]: 14 | init_values = {} 15 | factor_graph = Problem() 16 | with open(file_path) as ifile: 17 | for line in ifile.readlines(): 18 | items = line[:-1].split(" ") 19 | if items[0] == "EDGE_SE2": 20 | point_id0 = f"x{int(items[1])}" 21 | point_id1 = f"x{int(items[2])}" 22 | items_float = [float(i) for i in items[3:]] 23 | dx = items_float[0] 24 | dy = items_float[1] 25 | dtheta = items_float[2] 26 | dpose = np.array([dtheta, dx, dy]) 27 | 28 | # if point_id0 == "x10" and point_id1 == "x11": 29 | # print(dpose) 30 | i11, i12, i13, i22, i23, i33 = items_float[3:] 31 | matrix_i = np.array([[i11, i12, i13], [i12, i22, i23], [i13, i23, i33]]) 32 | factor = BetweenFactorSE2(dx, dy, dtheta) 33 | factor_graph.add_residual_block(3, [(point_id0, 3), (point_id1, 3)], factor, HuberLoss()) 34 | elif items[0] == "VERTEX_SE2": 35 | point_id = f"x{int(items[1])}" 36 | x = float(items[2]) 37 | y = float(items[3]) 38 | theta = float(items[4]) 39 | 40 | init_values[point_id] = np.array([theta, x, y], dtype=np.float64) 41 | else: 42 | print(items) 43 | break 44 | return factor_graph, init_values 45 | 46 | 47 | def show_pose(init_values, color): 48 | data_x = [x[1] for x in init_values.values()] 49 | data_y = [x[2] for x in init_values.values()] 50 | plt.scatter(data_x, data_y, s=1, c=color) 51 | 52 | 53 | def main(): 54 | FORMAT = "%(levelname)s %(name)s %(asctime)-15s %(filename)s:%(lineno)d %(message)s" 55 | logging.basicConfig(format=FORMAT) 56 | logging.getLogger().setLevel(logging.INFO) 57 | file_path = "tests/data/input_M3500_g2o.g2o" 58 | factor_graph, init_values = load_g2o(file_path) 59 | 60 | prior_factor = PriorFactor(np.zeros(3)) 61 | factor_graph.add_residual_block(3, [("x0", 3)], prior_factor, None) 62 | solver = GaussNewtonOptimizer() 63 | # gn = LevenbergMarquardtOptimizer() 64 | draw = True 65 | if draw: 66 | plt.figure(figsize=(8, 8)) 67 | show_pose(init_values, "red") 68 | 69 | optimizer_option = OptimizerOptions(max_iteration=50) 70 | start_time = perf_counter() 71 | init_values = solver.optimize(factor_graph, init_values, optimizer_option) 72 | end_time = perf_counter() 73 | print(f"{solver.__class__.__name__} takes {end_time-start_time:.3f} sec") 74 | if draw: 75 | show_pose(init_values, "blue") 76 | ax = plt.gca() 77 | ax.set_xlim((-50, 50)) 78 | ax.set_ylim((-80, 20)) 79 | plt.tight_layout() 80 | plt.savefig("m3500_py.png") 81 | print("end") 82 | 83 | 84 | if __name__ == "__main__": 85 | main() 86 | -------------------------------------------------------------------------------- /examples/python/small_problem.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from tiny_solver import Problem, GaussNewtonOptimizer 3 | from tiny_solver.factors import PriorFactor, PyFactor 4 | 5 | # define custom cost function in python 6 | def cost(x: np.ndarray, yz: np.ndarray) -> np.ndarray: 7 | r0 = x[0] + 2 * yz[0] + 4 * yz[1] 8 | r1 = yz[0] * yz[0] 9 | return np.array([r0, r1]) 10 | 11 | 12 | def main(): 13 | 14 | # initialize problem (factor graph) 15 | problem = Problem() 16 | 17 | # factor defined in python 18 | custom_factor = PyFactor(cost) 19 | problem.add_residual_block( 20 | 2, 21 | [ 22 | ("x", 1), 23 | ("yz", 2), 24 | ], 25 | custom_factor, 26 | None, 27 | ) 28 | 29 | # prior factor import from rust 30 | prior_factor = PriorFactor(np.array([3.0])) 31 | problem.add_residual_block(1, [("x", 1)], prior_factor, None) 32 | 33 | # initial values 34 | init_values = {"x": np.array([0.7]), "yz": np.array([-30.2, 123.4])} 35 | 36 | # optimizer 37 | optimizer = GaussNewtonOptimizer() 38 | result_values = optimizer.optimize(problem, init_values) 39 | 40 | # result 41 | for k, v in result_values.items(): 42 | print(f"{k}: {v}") 43 | 44 | 45 | if __name__ == "__main__": 46 | main() 47 | -------------------------------------------------------------------------------- /examples/python/try_import.py: -------------------------------------------------------------------------------- 1 | import tiny_solver 2 | from tiny_solver import GaussNewtonOptimizer, Problem, LinearSolver, OptimizerOptions 3 | from tiny_solver.factors import PriorFactor, BetweenFactorSE2, PyFactor 4 | from tiny_solver.loss_functions import HuberLoss 5 | import numpy as np 6 | 7 | 8 | def f(x: np.ndarray, y: np.ndarray): 9 | # print("py ", x*x) 10 | return np.array([2 * x[0], x[1] * x[1] * x[1], y[1] * 4.0]) 11 | 12 | 13 | def fa(): 14 | print("fa") 15 | return 123 16 | 17 | 18 | def main(): 19 | print(f"{tiny_solver.__version__=}") 20 | print(dir(tiny_solver)) 21 | 22 | # print(dir(LinearSolver.SparseCholesky)) 23 | # opt_option = OptimizerOptions(linear_solver_type=LinearSolver.SparseQR, max_iteration=12, verbosity_level=1) 24 | # print(opt_option) 25 | # loss = HuberLoss(1.0) 26 | # print(loss) 27 | # a = np.array([1.0, 2.0]) 28 | # # j = first_derivative_test(f, a) 29 | # # print(j) 30 | # a = PyFactor(f) 31 | # # print(tiny_solver.sum_as_string(1, 2)) 32 | 33 | # # tiny_solver.mult(np.zeros((1, 2))) 34 | # # a = tiny_solver.Dual64() 35 | # # print(a.first_derivative) 36 | # b = BetweenFactorSE2(1.0, 2.0, 3.0) 37 | # # print("factor module\n", dir(factors)) 38 | # # b = Costf(1.0, 2.0, 3.0) 39 | 40 | # print(type(b)) 41 | # print(dir(b)) 42 | # print(b) 43 | # problem = Problem() 44 | # print(dir(problem)) 45 | # problem.num = 200 46 | # print(problem.num) 47 | # d = PriorFactor(np.array([1.0, 2.0, 3.0])) 48 | # problem.add_residual_block(1, [("aa", 1)], d) 49 | # problem.add_residual_block(1, [("aaa", 1)], b) 50 | # # c.add_residual_block(1, [("aaa", 1)]) 51 | # # c.add_residual_block(1, [("aa", 1)]) 52 | # # d = tiny_solver.BetweenFactor() 53 | # # d.ttt() 54 | # # tiny_solver.te(d) 55 | # optimizer = GaussNewtonOptimizer() 56 | # optimizer.optimize(problem, {"aa": np.array([123, 2, 3, 4], dtype=np.float64)}) 57 | 58 | 59 | if __name__ == "__main__": 60 | main() 61 | -------------------------------------------------------------------------------- /examples/small_problem.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | 3 | use tiny_solver::{self, na, Optimizer}; 4 | 5 | struct CustomFactor; 6 | // define your own residual function and the jacobian will be auto generated 7 | impl tiny_solver::factors::Factor for CustomFactor { 8 | fn residual_func(&self, params: &[nalgebra::DVector]) -> nalgebra::DVector { 9 | let x = ¶ms[0][0]; 10 | let y = ¶ms[1][0]; 11 | let z = ¶ms[1][1]; 12 | 13 | na::dvector![ 14 | x.clone() 15 | + y.clone() * T::from_f64(2.0).unwrap() 16 | + z.clone() * T::from_f64(4.0).unwrap(), 17 | y.clone() * z.clone() 18 | ] 19 | } 20 | } 21 | 22 | fn main() { 23 | // init logger, `export RUST_LOG=trace` to see more log 24 | env_logger::init(); 25 | 26 | // init problem (factor graph) 27 | let mut problem = tiny_solver::Problem::new(); 28 | 29 | // add residual blocks (factors) 30 | // add residual x needs to be close to 3.0 31 | problem.add_residual_block( 32 | 1, 33 | &["x"], 34 | Box::new(tiny_solver::factors::PriorFactor { 35 | v: na::dvector![3.0], 36 | }), 37 | None, 38 | ); 39 | // add custom residual for x and yz 40 | problem.add_residual_block(2, &["x", "yz"], Box::new(CustomFactor), None); 41 | 42 | // the initial values for x is 0.7 and yz is [-30.2, 123.4] 43 | let initial_values = HashMap::>::from([ 44 | ("x".to_string(), na::dvector![0.7]), 45 | ("yz".to_string(), na::dvector![-30.2, 123.4]), 46 | ]); 47 | 48 | // initialize optimizer 49 | let optimizer = tiny_solver::GaussNewtonOptimizer {}; 50 | 51 | // optimize 52 | let result = optimizer.optimize(&problem, &initial_values, None); 53 | 54 | // result 55 | for (k, v) in result.unwrap() { 56 | println!("{}: {}", k, v); 57 | } 58 | } 59 | -------------------------------------------------------------------------------- /examples/sphere2500.rs: -------------------------------------------------------------------------------- 1 | use std::time::Instant; 2 | 3 | use plotters::prelude::*; 4 | 5 | use tiny_solver::helper::read_g2o; 6 | use tiny_solver::{optimizer::Optimizer, GaussNewtonOptimizer}; 7 | 8 | fn main() { 9 | // init logger 10 | env_logger::init(); 11 | 12 | let (problem, init_values) = read_g2o("tests/data/sphere2500.g2o"); 13 | let init_points: Vec<(f64, f64)> = init_values.values().map(|v| (v[4], v[6])).collect(); 14 | let root_drawing_area = 15 | BitMapBackend::new("sphere2500_rs.png", (1024, 1024)).into_drawing_area(); 16 | 17 | root_drawing_area.fill(&WHITE).unwrap(); 18 | 19 | let mut scatter_ctx = ChartBuilder::on(&root_drawing_area) 20 | .x_label_area_size(40) 21 | .y_label_area_size(40) 22 | .build_cartesian_2d(-60f64..60f64, -110f64..10f64) 23 | .unwrap(); 24 | scatter_ctx 25 | .configure_mesh() 26 | .disable_x_mesh() 27 | .disable_y_mesh() 28 | .draw() 29 | .unwrap(); 30 | scatter_ctx 31 | .draw_series( 32 | init_points 33 | .iter() 34 | .map(|(x, y)| Circle::new((*x, *y), 2, GREEN.filled())), 35 | ) 36 | .unwrap(); 37 | let gn = GaussNewtonOptimizer::new(); 38 | let start = Instant::now(); 39 | let result = gn.optimize(&problem, &init_values, None); 40 | let duration = start.elapsed(); 41 | println!("Time elapsed in total is: {:?}", duration); 42 | let result_points: Vec<(f64, f64)> = result.unwrap().values().map(|v| (v[4], v[6])).collect(); 43 | scatter_ctx 44 | .draw_series( 45 | result_points 46 | .iter() 47 | .map(|(x, y)| Circle::new((*x, *y), 2, BLUE.filled())), 48 | ) 49 | .unwrap(); 50 | } 51 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["maturin>=1.4,<2.0"] 3 | build-backend = "maturin" 4 | 5 | [project] 6 | name = "tiny_solver" 7 | requires-python = ">=3.8" 8 | classifiers = [ 9 | "Programming Language :: Rust", 10 | "Programming Language :: Python :: Implementation :: CPython", 11 | "Programming Language :: Python :: Implementation :: PyPy", 12 | "Programming Language :: Python :: 3.8", 13 | "Programming Language :: Python :: 3.9", 14 | "Programming Language :: Python :: 3.10", 15 | "Programming Language :: Python :: 3.11", 16 | "Programming Language :: Python :: 3.12", 17 | ] 18 | dynamic = ["version"] 19 | dependencies = ["numpy"] 20 | 21 | 22 | [tool.maturin] 23 | features = ["pyo3/extension-module", "python"] 24 | -------------------------------------------------------------------------------- /scripts/compile_aarch64.sh: -------------------------------------------------------------------------------- 1 | maturin build --release --target aarch64-unknown-linux-gnu --zig -------------------------------------------------------------------------------- /scripts/compile_local.sh: -------------------------------------------------------------------------------- 1 | maturin build --release -------------------------------------------------------------------------------- /scripts/fmt.sh: -------------------------------------------------------------------------------- 1 | cargo fmt 2 | black -l 120 tiny_solver 3 | black -l 120 examples/python -------------------------------------------------------------------------------- /src/corrector.rs: -------------------------------------------------------------------------------- 1 | use nalgebra as na; 2 | /// https://github.com/ceres-solver/ceres-solver/blob/master/internal/ceres/corrector.cc 3 | use std::ops::Mul; 4 | 5 | pub(crate) struct Corrector { 6 | sqrt_rho1: f64, 7 | residual_scaling: f64, 8 | alpha_sq_norm: f64, 9 | } 10 | 11 | impl Corrector { 12 | pub fn new(sq_norm: f64, rho: &[f64; 3]) -> Corrector { 13 | let sqrt_rho1 = rho[1].sqrt(); 14 | // is always well formed. 15 | if (sq_norm == 0.0) || (rho[2] <= 0.0) { 16 | let residual_scaling = sqrt_rho1; 17 | let alpha_sq_norm = 0.0; 18 | Corrector { 19 | sqrt_rho1, 20 | residual_scaling, 21 | alpha_sq_norm, 22 | } 23 | } else { 24 | let d = 1.0 + 2.0 * sq_norm * rho[2] / rho[1]; 25 | 26 | let alpha = 1.0 - d.sqrt(); 27 | 28 | let residual_scaling = sqrt_rho1 / (1.0 - alpha); 29 | let alpha_sq_norm = alpha / sq_norm; 30 | Corrector { 31 | sqrt_rho1, 32 | residual_scaling, 33 | alpha_sq_norm, 34 | } 35 | } 36 | } 37 | pub fn correct_jacobian(&self, residuals: &na::DVector, jacobian: &mut na::DMatrix) { 38 | // The common case (rho[2] <= 0). 39 | if self.alpha_sq_norm == 0.0 { 40 | *jacobian = jacobian.clone().mul(self.sqrt_rho1); 41 | return; 42 | } 43 | 44 | // J = sqrt(rho) * (J - alpha^2 r * r' J) 45 | let r_rtj = residuals.clone() * residuals.transpose() * jacobian.clone(); 46 | *jacobian = (jacobian.clone() - r_rtj.mul(self.alpha_sq_norm)).mul(self.sqrt_rho1); 47 | } 48 | pub fn correct_residuals(&self, residuals: &mut na::DVector) { 49 | *residuals = residuals.clone().mul(self.residual_scaling); 50 | } 51 | } 52 | -------------------------------------------------------------------------------- /src/factors.rs: -------------------------------------------------------------------------------- 1 | pub use nalgebra as na; 2 | 3 | use crate::manifold::se3::SE3; 4 | 5 | pub trait Factor: Send + Sync { 6 | fn residual_func(&self, params: &[na::DVector]) -> na::DVector; 7 | } 8 | pub trait FactorImpl: Factor + Factor { 9 | fn residual_func_dual( 10 | &self, 11 | params: &[na::DVector], 12 | ) -> na::DVector { 13 | self.residual_func(params) 14 | } 15 | fn residual_func_f64(&self, params: &[na::DVector]) -> na::DVector { 16 | self.residual_func(params) 17 | } 18 | } 19 | 20 | impl FactorImpl for T 21 | where 22 | T: Factor + Factor, 23 | { 24 | fn residual_func_dual( 25 | &self, 26 | params: &[na::DVector], 27 | ) -> na::DVector { 28 | self.residual_func(params) 29 | } 30 | 31 | fn residual_func_f64(&self, params: &[na::DVector]) -> na::DVector { 32 | self.residual_func(params) 33 | } 34 | } 35 | 36 | #[derive(Debug, Clone)] 37 | pub struct BetweenFactorSE2 { 38 | pub dx: f64, 39 | pub dy: f64, 40 | pub dtheta: f64, 41 | } 42 | impl Factor for BetweenFactorSE2 { 43 | fn residual_func(&self, params: &[na::DVector]) -> na::DVector { 44 | let t_origin_k0 = ¶ms[0]; 45 | let t_origin_k1 = ¶ms[1]; 46 | let se2_origin_k0 = na::Isometry2::new( 47 | na::Vector2::new(t_origin_k0[1].clone(), t_origin_k0[2].clone()), 48 | t_origin_k0[0].clone(), 49 | ); 50 | let se2_origin_k1 = na::Isometry2::new( 51 | na::Vector2::new(t_origin_k1[1].clone(), t_origin_k1[2].clone()), 52 | t_origin_k1[0].clone(), 53 | ); 54 | let se2_k0_k1 = na::Isometry2::new( 55 | na::Vector2::::new(T::from_f64(self.dx).unwrap(), T::from_f64(self.dy).unwrap()), 56 | T::from_f64(self.dtheta).unwrap(), 57 | ); 58 | 59 | let se2_diff = se2_origin_k1.inverse() * se2_origin_k0 * se2_k0_k1; 60 | na::dvector![ 61 | se2_diff.translation.x.clone(), 62 | se2_diff.translation.y.clone(), 63 | se2_diff.rotation.angle() 64 | ] 65 | } 66 | } 67 | 68 | #[derive(Debug, Clone)] 69 | pub struct BetweenFactorSE3 { 70 | pub dtx: f64, 71 | pub dty: f64, 72 | pub dtz: f64, 73 | pub dqx: f64, 74 | pub dqy: f64, 75 | pub dqz: f64, 76 | pub dqw: f64, 77 | } 78 | impl Factor for BetweenFactorSE3 { 79 | fn residual_func(&self, params: &[na::DVector]) -> na::DVector { 80 | let t_origin_k0 = ¶ms[0]; 81 | let t_origin_k1 = ¶ms[1]; 82 | let se3_origin_k0 = SE3::from_vec(t_origin_k0.as_view()); 83 | let se3_origin_k1 = SE3::from_vec(t_origin_k1.as_view()); 84 | 85 | let se3_k0_k1 = SE3::from_vec( 86 | na::dvector![self.dqx, self.dqy, self.dqz, self.dqw, self.dtx, self.dty, self.dtz,] 87 | .as_view(), 88 | ) 89 | .cast::(); 90 | 91 | let se3_diff = se3_origin_k1.inverse() * se3_origin_k0 * se3_k0_k1.cast(); 92 | 93 | se3_diff.log() 94 | } 95 | } 96 | 97 | #[derive(Debug, Clone)] 98 | pub struct PriorFactor { 99 | pub v: na::DVector, 100 | } 101 | impl Factor for PriorFactor { 102 | fn residual_func(&self, params: &[na::DVector]) -> na::DVector { 103 | params[0].clone() - self.v.clone().cast() 104 | } 105 | } 106 | -------------------------------------------------------------------------------- /src/helper.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | use std::fs::read_to_string; 3 | use std::sync::Arc; 4 | 5 | use nalgebra as na; 6 | 7 | use crate::loss_functions::HuberLoss; 8 | use crate::manifold::se3::SE3Manifold; 9 | use crate::{factors, problem}; 10 | 11 | pub fn translation_quaternion_to_na( 12 | tx: &T, 13 | ty: &T, 14 | tz: &T, 15 | qx: &T, 16 | qy: &T, 17 | qz: &T, 18 | qw: &T, 19 | ) -> na::Isometry3 { 20 | let rotation = na::UnitQuaternion::from_quaternion(na::Quaternion::new( 21 | qw.clone(), 22 | qx.clone(), 23 | qy.clone(), 24 | qz.clone(), 25 | )); 26 | na::Isometry3::from_parts( 27 | na::Translation3::new(tx.clone(), ty.clone(), tz.clone()), 28 | rotation, 29 | ) 30 | } 31 | 32 | pub fn read_g2o(filename: &str) -> (problem::Problem, HashMap>) { 33 | let mut problem = problem::Problem::new(); 34 | let mut init_values = HashMap::>::new(); 35 | for line in read_to_string(filename).unwrap().lines() { 36 | let line: Vec<&str> = line.split(' ').collect(); 37 | match line[0] { 38 | "VERTEX_SE2" => { 39 | let x = line[2].parse::().unwrap(); 40 | let y = line[3].parse::().unwrap(); 41 | let theta = line[4].parse::().unwrap(); 42 | init_values.insert(format!("x{}", line[1]), na::dvector![theta, x, y]); 43 | } 44 | "EDGE_SE2" => { 45 | let id0 = format!("x{}", line[1]); 46 | let id1 = format!("x{}", line[2]); 47 | let dx = line[3].parse::().unwrap(); 48 | let dy = line[4].parse::().unwrap(); 49 | let dtheta = line[5].parse::().unwrap(); 50 | // todo add info matrix 51 | let edge = factors::BetweenFactorSE2 { dx, dy, dtheta }; 52 | problem.add_residual_block( 53 | 3, 54 | &[&id0, &id1], 55 | Box::new(edge), 56 | Some(Box::new(HuberLoss::new(1.0))), 57 | ); 58 | } 59 | "VERTEX_SE3:QUAT" => { 60 | let x = line[2].parse::().expect("Failed to parse g2o"); 61 | let y = line[3].parse::().expect("Failed to parse g2o"); 62 | let z = line[4].parse::().expect("Failed to parse g2o"); 63 | let qx = line[5].parse::().expect("Failed to parse g2o"); 64 | let qy = line[6].parse::().expect("Failed to parse g2o"); 65 | let qz = line[7].parse::().expect("Failed to parse g2o"); 66 | let qw = line[8].parse::().expect("Failed to parse g2o"); 67 | let var_name = format!("x{}", line[1]); 68 | problem.set_variable_manifold(&var_name, Arc::new(SE3Manifold)); 69 | init_values.insert(var_name, na::dvector![qx, qy, qz, qw, x, y, z]); 70 | } 71 | "EDGE_SE3:QUAT" => { 72 | let id0 = format!("x{}", line[1]); 73 | let id1 = format!("x{}", line[2]); 74 | let dtx = line[3].parse::().expect("Failed to parse g2o"); 75 | let dty = line[4].parse::().expect("Failed to parse g2o"); 76 | let dtz = line[5].parse::().expect("Failed to parse g2o"); 77 | let dqx = line[6].parse::().expect("Failed to parse g2o"); 78 | let dqy = line[7].parse::().expect("Failed to parse g2o"); 79 | let dqz = line[8].parse::().expect("Failed to parse g2o"); 80 | let dqw = line[9].parse::().expect("Failed to parse g2o"); 81 | let edge = factors::BetweenFactorSE3 { 82 | dtx, 83 | dty, 84 | dtz, 85 | dqx, 86 | dqy, 87 | dqz, 88 | dqw, 89 | }; 90 | problem.add_residual_block( 91 | 6, 92 | &[&id0, &id1], 93 | Box::new(edge), 94 | Some(Box::new(HuberLoss::new(1.0))), 95 | ); 96 | } 97 | _ => { 98 | println!("err"); 99 | break; 100 | } 101 | } 102 | } 103 | let x0 = init_values.get("x0").unwrap(); 104 | let origin_factor = factors::PriorFactor { v: x0.clone() }; 105 | problem.add_residual_block( 106 | x0.shape().0, 107 | &["x0"], 108 | Box::new(origin_factor), 109 | Some(Box::new(HuberLoss::new(1.0))), 110 | ); 111 | (problem, init_values) 112 | } 113 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | pub mod corrector; 2 | pub mod factors; 3 | pub mod helper; 4 | pub mod linear; 5 | pub mod loss_functions; 6 | pub mod manifold; 7 | pub mod optimizer; 8 | pub mod parameter_block; 9 | pub mod problem; 10 | pub mod residual_block; 11 | 12 | pub use factors::na; 13 | pub use linear::*; 14 | pub use optimizer::*; 15 | pub use problem::*; 16 | pub use residual_block::*; 17 | 18 | #[cfg(feature = "python")] 19 | pub mod python; 20 | -------------------------------------------------------------------------------- /src/linear/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod sparse; 2 | pub mod sparse_cholesky; 3 | pub mod sparse_qr; 4 | pub use sparse_cholesky::*; 5 | pub use sparse_qr::*; 6 | -------------------------------------------------------------------------------- /src/linear/sparse.rs: -------------------------------------------------------------------------------- 1 | #[derive(Default, Clone)] 2 | pub enum LinearSolverType { 3 | #[default] 4 | SparseCholesky, 5 | SparseQR, 6 | } 7 | 8 | pub trait SparseLinearSolver { 9 | fn solve( 10 | &mut self, 11 | residuals: &faer::Mat, 12 | jacobians: &faer::sparse::SparseColMat, 13 | ) -> Option>; 14 | fn solve_jtj( 15 | &mut self, 16 | jtr: &faer::Mat, 17 | jtj: &faer::sparse::SparseColMat, 18 | ) -> Option>; 19 | } 20 | -------------------------------------------------------------------------------- /src/linear/sparse_cholesky.rs: -------------------------------------------------------------------------------- 1 | use std::fmt::Debug; 2 | use std::ops::Mul; 3 | 4 | use faer::linalg::solvers::Solve; 5 | use faer::sparse::linalg::solvers; 6 | 7 | use super::sparse::SparseLinearSolver; 8 | 9 | // #[pyclass] 10 | #[derive(Debug, Clone)] 11 | pub struct SparseCholeskySolver { 12 | symbolic_pattern: Option>, 13 | } 14 | 15 | impl SparseCholeskySolver { 16 | pub fn new() -> Self { 17 | SparseCholeskySolver { 18 | symbolic_pattern: None, 19 | } 20 | } 21 | } 22 | impl Default for SparseCholeskySolver { 23 | fn default() -> Self { 24 | Self::new() 25 | } 26 | } 27 | impl SparseLinearSolver for SparseCholeskySolver { 28 | fn solve( 29 | &mut self, 30 | residuals: &faer::Mat, 31 | jacobians: &faer::sparse::SparseColMat, 32 | ) -> Option> { 33 | let jtj = jacobians 34 | .as_ref() 35 | .transpose() 36 | .to_col_major() 37 | .unwrap() 38 | .mul(jacobians.as_ref()); 39 | let jtr = jacobians.as_ref().transpose().mul(-residuals); 40 | 41 | self.solve_jtj(&jtr, &jtj) 42 | } 43 | 44 | fn solve_jtj( 45 | &mut self, 46 | jtr: &faer::Mat, 47 | jtj: &faer::sparse::SparseColMat, 48 | ) -> Option> { 49 | // initialize the pattern 50 | if self.symbolic_pattern.is_none() { 51 | self.symbolic_pattern = 52 | Some(solvers::SymbolicLlt::try_new(jtj.symbolic(), faer::Side::Lower).unwrap()); 53 | } 54 | 55 | let sym = self.symbolic_pattern.as_ref().unwrap(); 56 | if let Ok(cholesky) = 57 | solvers::Llt::try_new_with_symbolic(sym.clone(), jtj.as_ref(), faer::Side::Lower) 58 | { 59 | let dx = cholesky.solve(jtr); 60 | 61 | Some(dx) 62 | } else { 63 | None 64 | } 65 | } 66 | } 67 | -------------------------------------------------------------------------------- /src/linear/sparse_qr.rs: -------------------------------------------------------------------------------- 1 | use super::sparse::SparseLinearSolver; 2 | use faer::linalg::solvers::SolveLstsqCore; 3 | // use faer::prelude::{SpSolver, SpSolverLstsq}; 4 | use faer::sparse::linalg::solvers; 5 | 6 | #[derive(Debug, Clone)] 7 | pub struct SparseQRSolver { 8 | symbolic_pattern: Option>, 9 | } 10 | 11 | impl SparseQRSolver { 12 | pub fn new() -> Self { 13 | SparseQRSolver { 14 | symbolic_pattern: None, 15 | } 16 | } 17 | } 18 | impl Default for SparseQRSolver { 19 | fn default() -> Self { 20 | Self::new() 21 | } 22 | } 23 | impl SparseLinearSolver for SparseQRSolver { 24 | fn solve( 25 | &mut self, 26 | residuals: &faer::Mat, 27 | jacobians: &faer::sparse::SparseColMat, 28 | ) -> Option> { 29 | if self.symbolic_pattern.is_none() { 30 | self.symbolic_pattern = 31 | Some(solvers::SymbolicQr::try_new(jacobians.symbolic()).unwrap()); 32 | } 33 | 34 | let sym = self.symbolic_pattern.as_ref().unwrap(); 35 | if let Ok(qr) = solvers::Qr::try_new_with_symbolic(sym.clone(), jacobians.as_ref()) { 36 | let mut minus_residuals = -residuals; 37 | qr.solve_lstsq_in_place_with_conj(faer::Conj::No, minus_residuals.as_mut()); 38 | Some(minus_residuals) 39 | } else { 40 | None 41 | } 42 | } 43 | 44 | fn solve_jtj( 45 | &mut self, 46 | jtr: &faer::Mat, 47 | jtj: &faer::sparse::SparseColMat, 48 | ) -> Option> { 49 | if self.symbolic_pattern.is_none() { 50 | self.symbolic_pattern = Some(solvers::SymbolicQr::try_new(jtj.symbolic()).unwrap()); 51 | } 52 | 53 | let sym = self.symbolic_pattern.as_ref().unwrap(); 54 | if let Ok(qr) = solvers::Qr::try_new_with_symbolic(sym.clone(), jtj.as_ref()) { 55 | let mut minus_jtr = -jtr; 56 | qr.solve_lstsq_in_place_with_conj(faer::Conj::No, minus_jtr.as_mut()); 57 | Some(minus_jtr) 58 | } else { 59 | None 60 | } 61 | } 62 | } 63 | -------------------------------------------------------------------------------- /src/loss_functions.rs: -------------------------------------------------------------------------------- 1 | use core::f64; 2 | 3 | pub enum LossFunc { 4 | HuberLoss, 5 | } 6 | 7 | pub trait Loss: Send + Sync { 8 | fn evaluate(&self, s: f64) -> [f64; 3]; 9 | } 10 | 11 | #[derive(Debug, Clone)] 12 | pub struct HuberLoss { 13 | scale: f64, 14 | scale2: f64, 15 | } 16 | impl HuberLoss { 17 | pub fn new(scale: f64) -> Self { 18 | if scale <= 0.0 { 19 | panic!("scale needs to be larger than zero"); 20 | } 21 | HuberLoss { 22 | scale, 23 | scale2: scale * scale, 24 | } 25 | } 26 | } 27 | 28 | impl Loss for HuberLoss { 29 | fn evaluate(&self, s: f64) -> [f64; 3] { 30 | if s > self.scale2 { 31 | // Outlier region. 32 | // 'r' is always positive. 33 | let r = s.sqrt(); 34 | let rho1 = (self.scale / r).max(f64::MIN); 35 | [2.0 * self.scale * r - self.scale2, rho1, -rho1 / (2.0 * s)] 36 | } else { 37 | // Inlier region. 38 | [s, 1.0, 0.0] 39 | } 40 | } 41 | } 42 | 43 | pub struct CauchyLoss { 44 | scale2: f64, 45 | c: f64, 46 | } 47 | impl CauchyLoss { 48 | pub fn new(scale: f64) -> Self { 49 | let scale2 = scale * scale; 50 | CauchyLoss { 51 | scale2, 52 | c: 1.0 / scale2, 53 | } 54 | } 55 | } 56 | impl Loss for CauchyLoss { 57 | fn evaluate(&self, s: f64) -> [f64; 3] { 58 | let sum = 1.0 + s * self.c; 59 | let inv = 1.0 / sum; 60 | // 'sum' and 'inv' are always positive, assuming that 's' is. 61 | [ 62 | self.scale2 * sum.log2(), 63 | inv.max(f64::MIN), 64 | -self.c * (inv * inv), 65 | ] 66 | } 67 | } 68 | 69 | pub struct ArctanLoss { 70 | tolerance: f64, 71 | inv_of_squared_tolerance: f64, 72 | } 73 | 74 | impl ArctanLoss { 75 | pub fn new(tolerance: f64) -> Self { 76 | if tolerance <= 0.0 { 77 | panic!("scale needs to be larger than zero"); 78 | } 79 | ArctanLoss { 80 | tolerance, 81 | inv_of_squared_tolerance: 1.0 / (tolerance * tolerance), 82 | } 83 | } 84 | } 85 | 86 | impl Loss for ArctanLoss { 87 | fn evaluate(&self, s: f64) -> [f64; 3] { 88 | let sum = 1.0 + s * s * self.inv_of_squared_tolerance; 89 | let inv = 1.0 / sum; 90 | 91 | [ 92 | self.tolerance * s.atan2(self.tolerance), 93 | inv.max(f64::MIN), 94 | -2.0 * s * self.inv_of_squared_tolerance * (inv * inv), 95 | ] 96 | } 97 | } 98 | -------------------------------------------------------------------------------- /src/manifold/mod.rs: -------------------------------------------------------------------------------- 1 | use std::num::NonZero; 2 | 3 | use nalgebra as na; 4 | use num_dual::DualDVec64; 5 | 6 | pub mod se3; 7 | pub mod so3; 8 | 9 | pub trait AutoDiffManifold { 10 | fn plus(&self, x: na::DVectorView, delta: na::DVectorView) -> na::DVector; 11 | fn minus(&self, y: na::DVectorView, x: na::DVectorView) -> na::DVector; 12 | } 13 | 14 | pub trait Manifold: AutoDiffManifold + AutoDiffManifold { 15 | fn tangent_size(&self) -> NonZero; 16 | fn plus_f64(&self, x: na::DVectorView, delta: na::DVectorView) -> na::DVector { 17 | self.plus(x, delta) 18 | } 19 | fn plus_dual( 20 | &self, 21 | x: na::DVectorView, 22 | delta: na::DVectorView, 23 | ) -> na::DVector { 24 | self.plus(x, delta) 25 | } 26 | fn minus_f64(&self, y: na::DVectorView, x: na::DVectorView) -> na::DVector { 27 | self.minus(y, x) 28 | } 29 | fn minus_dual( 30 | &self, 31 | y: na::DVectorView, 32 | x: na::DVectorView, 33 | ) -> na::DVector { 34 | self.minus(y, x) 35 | } 36 | } 37 | -------------------------------------------------------------------------------- /src/manifold/se3.rs: -------------------------------------------------------------------------------- 1 | use std::{num::NonZero, ops::Mul}; 2 | 3 | use nalgebra as na; 4 | 5 | use super::{so3::SO3, AutoDiffManifold, Manifold}; 6 | 7 | pub struct SE3 { 8 | pub xyz: na::Vector3, 9 | pub rot: SO3, 10 | } 11 | 12 | impl SE3 { 13 | /// [qx, qy, qz, qw, tx, ty, tz] 14 | pub fn from_vec(qxyzw_txyz: na::DVectorView) -> Self { 15 | let rot = SO3::from_xyzw( 16 | qxyzw_txyz[0].clone(), 17 | qxyzw_txyz[1].clone(), 18 | qxyzw_txyz[2].clone(), 19 | qxyzw_txyz[3].clone(), 20 | ); 21 | let xyz = na::Vector3::new( 22 | qxyzw_txyz[4].clone(), 23 | qxyzw_txyz[5].clone(), 24 | qxyzw_txyz[6].clone(), 25 | ); 26 | SE3 { xyz, rot } 27 | } 28 | pub fn from_qvec_tvec(qxyzw: na::DVectorView, tvec: na::DVectorView) -> Self { 29 | let rot = SO3::from_vec(qxyzw); 30 | let xyz = na::Vector3::from_row_slice(tvec.as_slice()); 31 | SE3 { xyz, rot } 32 | } 33 | pub fn identity() -> Self { 34 | let rot = SO3::identity(); 35 | let xyz = na::Vector3::zeros(); 36 | SE3 { xyz, rot } 37 | } 38 | 39 | pub fn exp(xi: na::DVectorView) -> Self { 40 | // fake exp 41 | let rot = SO3::::exp(xi.rows(0, 3).as_view()); 42 | let xyz = na::Vector3::new(xi[3].clone(), xi[4].clone(), xi[5].clone()); 43 | SE3 { xyz, rot } 44 | } 45 | 46 | pub fn log(&self) -> na::DVector { 47 | let mut xi = na::DVector::zeros(6); 48 | let xi_theta = self.rot.log(); 49 | // let xyz = self.xyz; 50 | xi.as_mut_slice()[0..3].clone_from_slice(xi_theta.as_slice()); 51 | xi.as_mut_slice()[3..6].clone_from_slice(self.xyz.as_slice()); 52 | xi 53 | } 54 | 55 | pub fn to_dvec(&self) -> na::DVector { 56 | let quat = self.rot.to_vec(); 57 | na::dvector![ 58 | quat[0].clone(), 59 | quat[1].clone(), 60 | quat[2].clone(), 61 | quat[3].clone(), 62 | self.xyz[0].clone(), 63 | self.xyz[1].clone(), 64 | self.xyz[2].clone(), 65 | ] 66 | } 67 | 68 | // pub fn hat(xi: na::VectorView3) -> na::Matrix3 { 69 | // let mut xi_hat = na::Matrix3::zeros(); 70 | // xi_hat[(0, 1)] = -xi[2].clone(); 71 | // xi_hat[(0, 2)] = xi[1].clone(); 72 | // xi_hat[(1, 0)] = xi[2].clone(); 73 | // xi_hat[(1, 2)] = -xi[0].clone(); 74 | // xi_hat[(2, 0)] = -xi[1].clone(); 75 | // xi_hat[(2, 1)] = xi[0].clone(); 76 | 77 | // xi_hat 78 | // } 79 | 80 | pub fn cast>(&self) -> SE3 { 81 | SE3 { 82 | rot: self.rot.cast(), 83 | xyz: self.xyz.clone().cast(), 84 | } 85 | } 86 | pub fn inverse(&self) -> Self { 87 | let inv = self.rot.inverse(); 88 | let xyz = -(&inv * self.xyz.as_view()); 89 | SE3 { xyz, rot: inv } 90 | } 91 | pub fn compose(&self, rhs: &Self) -> Self { 92 | SE3 { 93 | rot: &self.rot * &rhs.rot, 94 | xyz: (&self.rot * rhs.xyz.as_view()) + self.xyz.clone(), 95 | } 96 | } 97 | } 98 | 99 | impl Mul for SE3 { 100 | type Output = Self; 101 | 102 | fn mul(self, rhs: Self) -> Self::Output { 103 | self.compose(&rhs) 104 | } 105 | } 106 | impl Mul for &SE3 { 107 | type Output = SE3; 108 | 109 | fn mul(self, rhs: Self) -> Self::Output { 110 | self.compose(rhs) 111 | } 112 | } 113 | 114 | impl Mul> for SE3 { 115 | type Output = na::Vector3; 116 | 117 | fn mul(self, rhs: na::VectorView3<'_, T>) -> Self::Output { 118 | let qv = SO3::from_xyzw(rhs[0].clone(), rhs[1].clone(), rhs[2].clone(), T::zero()); 119 | let rinv = self.rot.inverse(); 120 | let v_rot = ((self.rot * qv) * rinv).to_vec(); 121 | let v = na::Vector3::new(v_rot[0].clone(), v_rot[1].clone(), v_rot[2].clone()); 122 | v + self.xyz 123 | } 124 | } 125 | 126 | impl Mul> for &SE3 { 127 | type Output = na::Vector3; 128 | 129 | fn mul(self, rhs: na::VectorView3<'_, T>) -> Self::Output { 130 | let qv = SO3::from_xyzw(rhs[0].clone(), rhs[1].clone(), rhs[2].clone(), T::zero()); 131 | let rinv = self.rot.inverse(); 132 | let v_rot = ((&self.rot * &qv) * rinv).to_vec(); 133 | let v = na::Vector3::new(v_rot[0].clone(), v_rot[1].clone(), v_rot[2].clone()); 134 | v + self.xyz.clone() 135 | } 136 | } 137 | 138 | #[derive(Debug, Clone)] 139 | pub struct SE3Manifold; 140 | impl AutoDiffManifold for SE3Manifold { 141 | fn plus( 142 | &self, 143 | x: nalgebra::DVectorView, 144 | delta: nalgebra::DVectorView, 145 | ) -> nalgebra::DVector { 146 | let d = SE3::exp(delta); 147 | let x_se3 = SE3::from_vec(x); 148 | let x_plus = x_se3 * d; 149 | x_plus.to_dvec() 150 | } 151 | 152 | fn minus( 153 | &self, 154 | y: nalgebra::DVectorView, 155 | x: nalgebra::DVectorView, 156 | ) -> nalgebra::DVector { 157 | let y_se3 = SE3::from_vec(y); 158 | let x_se3_inv = SE3::from_vec(x).inverse(); 159 | let x_inv_y_log = (x_se3_inv * y_se3).log(); 160 | na::dvector![ 161 | x_inv_y_log[0].clone(), 162 | x_inv_y_log[1].clone(), 163 | x_inv_y_log[2].clone(), 164 | x_inv_y_log[3].clone(), 165 | x_inv_y_log[4].clone(), 166 | x_inv_y_log[5].clone() 167 | ] 168 | } 169 | } 170 | impl Manifold for SE3Manifold { 171 | fn tangent_size(&self) -> NonZero { 172 | NonZero::new(6).unwrap() 173 | } 174 | } 175 | -------------------------------------------------------------------------------- /src/manifold/so3.rs: -------------------------------------------------------------------------------- 1 | use std::{num::NonZero, ops::Mul}; 2 | 3 | use nalgebra as na; 4 | 5 | use super::{AutoDiffManifold, Manifold}; 6 | 7 | pub struct SO3 { 8 | qx: T, 9 | qy: T, 10 | qz: T, 11 | qw: T, 12 | } 13 | 14 | impl SO3 { 15 | /// [x, y, z, w] 16 | pub fn from_vec(xyzw: na::DVectorView) -> Self { 17 | SO3 { 18 | qx: xyzw[0].clone(), 19 | qy: xyzw[1].clone(), 20 | qz: xyzw[2].clone(), 21 | qw: xyzw[3].clone(), 22 | } 23 | } 24 | pub fn from_xyzw(x: T, y: T, z: T, w: T) -> Self { 25 | SO3 { 26 | qx: x, 27 | qy: y, 28 | qz: z, 29 | qw: w, 30 | } 31 | } 32 | pub fn identity() -> Self { 33 | SO3 { 34 | qx: T::zero(), 35 | qy: T::zero(), 36 | qz: T::zero(), 37 | qw: T::one(), 38 | } 39 | } 40 | 41 | pub fn exp(xi: na::DVectorView) -> Self { 42 | let mut xyzw = na::Vector4::zeros(); 43 | 44 | let theta2 = xi.norm_squared(); 45 | 46 | if theta2 < T::from_f64(1e-6).unwrap() { 47 | // cos(theta / 2) \approx 1 - theta^2 / 8 48 | xyzw.w = T::one() - theta2 / T::from_f64(8.0).unwrap(); 49 | // Complete the square so that norm is one 50 | let tmp = T::from_f64(0.5).unwrap(); 51 | xyzw.x = xi[0].clone() * tmp.clone(); 52 | xyzw.y = xi[1].clone() * tmp.clone(); 53 | xyzw.z = xi[2].clone() * tmp; 54 | } else { 55 | let theta = theta2.sqrt(); 56 | xyzw.w = (theta.clone() * T::from_f64(0.5).unwrap()).cos(); 57 | 58 | let omega = xi / theta; 59 | let sin_theta_half = (T::one() - xyzw.w.clone() * xyzw.w.clone()).sqrt(); 60 | xyzw.x = omega[0].clone() * sin_theta_half.clone(); 61 | xyzw.y = omega[1].clone() * sin_theta_half.clone(); 62 | xyzw.z = omega[2].clone() * sin_theta_half; 63 | } 64 | 65 | SO3::from_vec(xyzw.as_view()) 66 | } 67 | 68 | pub fn log(&self) -> na::DVector { 69 | const EPS: f64 = 1e-6; 70 | let ivec = na::dvector![self.qx.clone(), self.qy.clone(), self.qz.clone()]; 71 | 72 | let squared_n = ivec.norm_squared(); 73 | let w = self.qw.clone(); 74 | 75 | let near_zero = squared_n.le(&T::from_f64(EPS * EPS).unwrap()); 76 | 77 | let w_sq = w.clone() * w.clone(); 78 | let t0 = T::from_f64(2.0).unwrap() / w.clone() 79 | - T::from_f64(2.0 / 3.0).unwrap() * squared_n.clone() / (w_sq * w.clone()); 80 | 81 | let n = squared_n.sqrt(); 82 | 83 | let sign = T::from_f64(-1.0) 84 | .unwrap() 85 | .select(w.le(&T::zero()), T::one()); 86 | let atan_nbyw = sign.clone() * n.clone().atan2(sign * w); 87 | 88 | let t = T::from_f64(2.0).unwrap() * atan_nbyw / n; 89 | 90 | let two_atan_nbyd_by_n = t0.select(near_zero, t); 91 | 92 | ivec * two_atan_nbyd_by_n 93 | } 94 | 95 | pub fn hat(xi: na::VectorView3) -> na::Matrix3 { 96 | let mut xi_hat = na::Matrix3::zeros(); 97 | xi_hat[(0, 1)] = -xi[2].clone(); 98 | xi_hat[(0, 2)] = xi[1].clone(); 99 | xi_hat[(1, 0)] = xi[2].clone(); 100 | xi_hat[(1, 2)] = -xi[0].clone(); 101 | xi_hat[(2, 0)] = -xi[1].clone(); 102 | xi_hat[(2, 1)] = xi[0].clone(); 103 | 104 | xi_hat 105 | } 106 | 107 | pub fn to_vec(&self) -> na::Vector4 { 108 | na::Vector4::new( 109 | self.qx.clone(), 110 | self.qy.clone(), 111 | self.qz.clone(), 112 | self.qw.clone(), 113 | ) 114 | } 115 | pub fn to_dvec(&self) -> na::DVector { 116 | na::dvector![ 117 | self.qx.clone(), 118 | self.qy.clone(), 119 | self.qz.clone(), 120 | self.qw.clone(), 121 | ] 122 | } 123 | pub fn cast>(&self) -> SO3 { 124 | SO3::from_vec(self.to_vec().cast().as_view()) 125 | } 126 | pub fn inverse(&self) -> Self { 127 | SO3 { 128 | qx: -self.qx.clone(), 129 | qy: -self.qy.clone(), 130 | qz: -self.qz.clone(), 131 | qw: self.qw.clone(), 132 | } 133 | } 134 | pub fn compose(&self, rhs: &Self) -> Self { 135 | let x0 = self.qx.clone(); 136 | let y0 = self.qy.clone(); 137 | let z0 = self.qz.clone(); 138 | let w0 = self.qw.clone(); 139 | 140 | let x1 = rhs.qx.clone(); 141 | let y1 = rhs.qy.clone(); 142 | let z1 = rhs.qz.clone(); 143 | let w1 = rhs.qw.clone(); 144 | 145 | // Compute the product of the two quaternions, term by term 146 | let qx = w0.clone() * x1.clone() + x0.clone() * w1.clone() + y0.clone() * z1.clone() 147 | - z0.clone() * y1.clone(); 148 | let qy = w0.clone() * y1.clone() - x0.clone() * z1.clone() 149 | + y0.clone() * w1.clone() 150 | + z0.clone() * x1.clone(); 151 | let qz = w0.clone() * z1.clone() + x0.clone() * y1.clone() - y0.clone() * x1.clone() 152 | + z0.clone() * w1.clone(); 153 | let qw = w0 * w1 - x0 * x1 - y0 * y1 - z0 * z1; 154 | // println!("q {}", self.to_vec()); 155 | SO3 { qx, qy, qz, qw } 156 | } 157 | } 158 | 159 | impl Mul for SO3 { 160 | type Output = Self; 161 | 162 | fn mul(self, rhs: Self) -> Self::Output { 163 | self.compose(&rhs) 164 | } 165 | } 166 | impl Mul for &SO3 { 167 | type Output = SO3; 168 | 169 | fn mul(self, rhs: Self) -> Self::Output { 170 | self.compose(rhs) 171 | } 172 | } 173 | 174 | impl Mul> for &SO3 { 175 | type Output = na::Vector3; 176 | 177 | fn mul(self, rhs: na::VectorView3<'_, T>) -> Self::Output { 178 | let qv = SO3::from_xyzw(rhs[0].clone(), rhs[1].clone(), rhs[2].clone(), T::zero()); 179 | let inv = self.inverse(); 180 | let v_rot: SO3 = (self * &qv) * inv; 181 | na::Vector3::new(v_rot.qx, v_rot.qy, v_rot.qz) 182 | } 183 | } 184 | 185 | #[derive(Debug, Clone)] 186 | pub struct QuaternionManifold; 187 | impl AutoDiffManifold for QuaternionManifold { 188 | fn plus( 189 | &self, 190 | x: nalgebra::DVectorView, 191 | delta: nalgebra::DVectorView, 192 | ) -> nalgebra::DVector { 193 | let d: SO3 = SO3::exp(delta); 194 | let x_s03: SO3 = SO3::from_xyzw(x[0].clone(), x[1].clone(), x[2].clone(), x[3].clone()); 195 | let x_plus = x_s03 * d; 196 | na::dvector![ 197 | x_plus.qx.clone(), 198 | x_plus.qy.clone(), 199 | x_plus.qz.clone(), 200 | x_plus.qw.clone(), 201 | ] 202 | } 203 | 204 | fn minus( 205 | &self, 206 | y: nalgebra::DVectorView, 207 | x: nalgebra::DVectorView, 208 | ) -> nalgebra::DVector { 209 | let y_so3 = SO3::from_vec(y); 210 | let x_so3_inv = SO3::from_vec(x).inverse(); 211 | let x_inv_y_log = (x_so3_inv * y_so3).log(); 212 | na::dvector![ 213 | x_inv_y_log[0].clone(), 214 | x_inv_y_log[1].clone(), 215 | x_inv_y_log[2].clone() 216 | ] 217 | } 218 | } 219 | impl Manifold for QuaternionManifold { 220 | fn tangent_size(&self) -> NonZero { 221 | NonZero::new(3).unwrap() 222 | } 223 | } 224 | -------------------------------------------------------------------------------- /src/optimizer/common.rs: -------------------------------------------------------------------------------- 1 | use std::collections::{HashMap, HashSet}; 2 | use std::ops::Add; 3 | 4 | use nalgebra as na; 5 | 6 | use crate::parameter_block::ParameterBlock; 7 | use crate::problem; 8 | use crate::sparse::LinearSolverType; 9 | 10 | pub trait Optimizer { 11 | fn optimize( 12 | &self, 13 | problem: &problem::Problem, 14 | initial_values: &HashMap>, 15 | optimizer_option: Option, 16 | ) -> Option>>; 17 | fn apply_dx( 18 | &self, 19 | dx: &na::DVector, 20 | params: &mut HashMap>, 21 | variable_name_to_col_idx_dict: &HashMap, 22 | fixed_var_indexes: &HashMap>, 23 | variable_bounds: &HashMap>, 24 | ) { 25 | for (key, param) in params.iter_mut() { 26 | if let Some(col_idx) = variable_name_to_col_idx_dict.get(key) { 27 | let var_size = param.shape().0; 28 | let mut updated_param = param.clone().add(dx.rows(*col_idx, var_size)); 29 | if let Some(indexes_to_fix) = fixed_var_indexes.get(key) { 30 | for &idx in indexes_to_fix { 31 | log::debug!("Fix {} {}", key, idx); 32 | updated_param[idx] = param[idx]; 33 | } 34 | } 35 | if let Some(indexes_to_bound) = variable_bounds.get(key) { 36 | for (&idx, &(lower, upper)) in indexes_to_bound { 37 | let old = updated_param[idx]; 38 | updated_param[idx] = updated_param[idx].max(lower).min(upper); 39 | log::debug!("bound {} {} {} -> {}", key, idx, old, updated_param[idx]); 40 | } 41 | } 42 | param.copy_from(&updated_param); 43 | } 44 | } 45 | } 46 | fn apply_dx2( 47 | &self, 48 | dx: &na::DVector, 49 | params: &mut HashMap, 50 | variable_name_to_col_idx_dict: &HashMap, 51 | ) { 52 | params.iter_mut().for_each(|(key, param)| { 53 | if let Some(col_idx) = variable_name_to_col_idx_dict.get(key) { 54 | let var_size = param.tangent_size(); 55 | param.update_params(param.plus_f64(dx.rows(*col_idx, var_size))); 56 | } 57 | }); 58 | // for (key, param) in params.par_iter_mut() { 59 | // } 60 | } 61 | fn compute_error( 62 | &self, 63 | problem: &problem::Problem, 64 | params: &HashMap, 65 | ) -> f64 { 66 | problem 67 | .compute_residuals(params, true) 68 | .as_ref() 69 | .squared_norm_l2() 70 | } 71 | } 72 | 73 | #[derive(PartialEq, Debug)] 74 | pub enum SolverStatus { 75 | Running, 76 | // Resulting solution may be OK to use. 77 | GradientTooSmall, // eps > max(J'*f(x)) 78 | RelativeStepSizeTooSmall, // eps > ||dx|| / ||x|| 79 | ErrorTooSmall, // eps > ||f(x)|| 80 | HitMaxIterations, 81 | // Numerical issues 82 | // FAILED_TO_EVALUATE_COST_FUNCTION, 83 | // FAILED_TO_SOLVER_LINEAR_SYSTEM, 84 | } 85 | 86 | #[derive(Clone)] 87 | pub struct OptimizerOptions { 88 | pub max_iteration: usize, 89 | pub linear_solver_type: LinearSolverType, 90 | pub verbosity_level: usize, 91 | pub min_abs_error_decrease_threshold: f64, 92 | pub min_rel_error_decrease_threshold: f64, 93 | pub min_error_threshold: f64, 94 | // pub relative_step_threshold: 1e-16, 95 | } 96 | 97 | impl Default for OptimizerOptions { 98 | fn default() -> Self { 99 | OptimizerOptions { 100 | max_iteration: 100, 101 | linear_solver_type: LinearSolverType::SparseCholesky, 102 | verbosity_level: 0, 103 | min_abs_error_decrease_threshold: 1e-5, 104 | min_rel_error_decrease_threshold: 1e-5, 105 | min_error_threshold: 1e-10, 106 | } 107 | } 108 | } 109 | -------------------------------------------------------------------------------- /src/optimizer/gauss_newton_optimizer.rs: -------------------------------------------------------------------------------- 1 | use log::trace; 2 | use std::{collections::HashMap, time::Instant}; 3 | 4 | use faer_ext::IntoNalgebra; 5 | 6 | use crate::common::OptimizerOptions; 7 | use crate::linear; 8 | use crate::optimizer; 9 | use crate::parameter_block::ParameterBlock; 10 | use crate::sparse::LinearSolverType; 11 | use crate::sparse::SparseLinearSolver; 12 | 13 | #[derive(Debug)] 14 | pub struct GaussNewtonOptimizer {} 15 | impl GaussNewtonOptimizer { 16 | pub fn new() -> Self { 17 | Self {} 18 | } 19 | } 20 | impl Default for GaussNewtonOptimizer { 21 | fn default() -> Self { 22 | Self::new() 23 | } 24 | } 25 | 26 | impl optimizer::Optimizer for GaussNewtonOptimizer { 27 | fn optimize( 28 | &self, 29 | problem: &crate::problem::Problem, 30 | initial_values: &std::collections::HashMap>, 31 | optimizer_option: Option, 32 | ) -> Option>> { 33 | let mut parameter_blocks: HashMap = 34 | problem.initialize_parameter_blocks(initial_values); 35 | 36 | let variable_name_to_col_idx_dict = 37 | problem.get_variable_name_to_col_idx_dict(¶meter_blocks); 38 | let total_variable_dimension = parameter_blocks.values().map(|p| p.tangent_size()).sum(); 39 | 40 | let opt_option = optimizer_option.unwrap_or_default(); 41 | let mut linear_solver: Box = match opt_option.linear_solver_type { 42 | LinearSolverType::SparseCholesky => Box::new(linear::SparseCholeskySolver::new()), 43 | LinearSolverType::SparseQR => Box::new(linear::SparseQRSolver::new()), 44 | }; 45 | 46 | let symbolic_structure = problem.build_symbolic_structure( 47 | ¶meter_blocks, 48 | total_variable_dimension, 49 | &variable_name_to_col_idx_dict, 50 | ); 51 | 52 | let mut last_err; 53 | let mut current_error = self.compute_error(problem, ¶meter_blocks); 54 | 55 | for i in 0..opt_option.max_iteration { 56 | last_err = current_error; 57 | let mut start = Instant::now(); 58 | 59 | let (residuals, jac) = problem.compute_residual_and_jacobian( 60 | ¶meter_blocks, 61 | &variable_name_to_col_idx_dict, 62 | &symbolic_structure, 63 | ); 64 | let residual_and_jacobian_duration = start.elapsed(); 65 | 66 | start = Instant::now(); 67 | let solving_duration; 68 | if let Some(dx) = linear_solver.solve(&residuals, &jac) { 69 | solving_duration = start.elapsed(); 70 | let dx_na = dx.as_ref().into_nalgebra().column(0).clone_owned(); 71 | self.apply_dx2( 72 | &dx_na, 73 | &mut parameter_blocks, 74 | &variable_name_to_col_idx_dict, 75 | ); 76 | } else { 77 | log::debug!("solve ax=b failed"); 78 | return None; 79 | } 80 | 81 | current_error = self.compute_error(problem, ¶meter_blocks); 82 | trace!( 83 | "iter:{}, total err:{}, residual + jacobian duration: {:?}, solving duration: {:?}", 84 | i, 85 | current_error, 86 | residual_and_jacobian_duration, 87 | solving_duration 88 | ); 89 | 90 | if current_error < opt_option.min_error_threshold { 91 | trace!("error too low"); 92 | break; 93 | } else if current_error.is_nan() { 94 | log::debug!("solve ax=b failed, current error is nan"); 95 | return None; 96 | } 97 | 98 | if (last_err - current_error).abs() < opt_option.min_abs_error_decrease_threshold { 99 | trace!("absolute error decrease low"); 100 | break; 101 | } else if (last_err - current_error).abs() / last_err 102 | < opt_option.min_rel_error_decrease_threshold 103 | { 104 | trace!("relative error decrease low"); 105 | break; 106 | } 107 | } 108 | let params = parameter_blocks 109 | .iter() 110 | .map(|(k, v)| (k.to_owned(), v.params.clone())) 111 | .collect(); 112 | Some(params) 113 | } 114 | } 115 | -------------------------------------------------------------------------------- /src/optimizer/levenberg_marquardt_optimizer.rs: -------------------------------------------------------------------------------- 1 | use log::trace; 2 | use std::ops::Mul; 3 | use std::{collections::HashMap, time::Instant}; 4 | 5 | use faer::sparse::Triplet; 6 | use faer_ext::IntoNalgebra; 7 | 8 | use crate::common::OptimizerOptions; 9 | use crate::linear; 10 | use crate::optimizer; 11 | use crate::parameter_block::ParameterBlock; 12 | use crate::sparse::LinearSolverType; 13 | use crate::sparse::SparseLinearSolver; 14 | 15 | const DEFAULT_MIN_DIAGONAL: f64 = 1e-6; 16 | const DEFAULT_MAX_DIAGONAL: f64 = 1e32; 17 | const DEFAULT_INITIAL_TRUST_REGION_RADIUS: f64 = 1e4; 18 | 19 | #[derive(Debug)] 20 | pub struct LevenbergMarquardtOptimizer { 21 | min_diagonal: f64, 22 | max_diagonal: f64, 23 | initial_trust_region_radius: f64, 24 | } 25 | 26 | impl LevenbergMarquardtOptimizer { 27 | pub fn new(min_diagonal: f64, max_diagonal: f64, initial_trust_region_radius: f64) -> Self { 28 | Self { 29 | min_diagonal, 30 | max_diagonal, 31 | initial_trust_region_radius, 32 | } 33 | } 34 | } 35 | 36 | impl Default for LevenbergMarquardtOptimizer { 37 | fn default() -> Self { 38 | Self { 39 | min_diagonal: DEFAULT_MIN_DIAGONAL, 40 | max_diagonal: DEFAULT_MAX_DIAGONAL, 41 | initial_trust_region_radius: DEFAULT_INITIAL_TRUST_REGION_RADIUS, 42 | } 43 | } 44 | } 45 | 46 | impl optimizer::Optimizer for LevenbergMarquardtOptimizer { 47 | fn optimize( 48 | &self, 49 | problem: &crate::problem::Problem, 50 | initial_values: &std::collections::HashMap>, 51 | optimizer_option: Option, 52 | ) -> Option>> { 53 | let mut parameter_blocks: HashMap = 54 | problem.initialize_parameter_blocks(initial_values); 55 | 56 | let variable_name_to_col_idx_dict = 57 | problem.get_variable_name_to_col_idx_dict(¶meter_blocks); 58 | let total_variable_dimension = parameter_blocks.values().map(|p| p.tangent_size()).sum(); 59 | 60 | let opt_option = optimizer_option.unwrap_or_default(); 61 | let mut linear_solver: Box = match opt_option.linear_solver_type { 62 | LinearSolverType::SparseCholesky => Box::new(linear::SparseCholeskySolver::new()), 63 | LinearSolverType::SparseQR => Box::new(linear::SparseQRSolver::new()), 64 | }; 65 | 66 | // On the first iteration, we'll generate a diagonal matrix of the jacobian. 67 | // Its shape will be (total_variable_dimension, total_variable_dimension). 68 | // With LM, rather than solving A * dx = b for dx, we solve for (A + lambda * diag(A)) dx = b. 69 | let mut jacobi_scaling_diagonal: Option> = None; 70 | 71 | let symbolic_structure = problem.build_symbolic_structure( 72 | ¶meter_blocks, 73 | total_variable_dimension, 74 | &variable_name_to_col_idx_dict, 75 | ); 76 | 77 | // Damping parameter (a.k.a lambda / Marquardt parameter) 78 | let mut u = 1.0 / self.initial_trust_region_radius; 79 | 80 | let mut last_err; 81 | let mut current_error = self.compute_error(problem, ¶meter_blocks); 82 | for i in 0..opt_option.max_iteration { 83 | last_err = current_error; 84 | 85 | let (residuals, mut jac) = problem.compute_residual_and_jacobian( 86 | ¶meter_blocks, 87 | &variable_name_to_col_idx_dict, 88 | &symbolic_structure, 89 | ); 90 | 91 | if i == 0 { 92 | // On the first iteration, generate the diagonal of the jacobian. 93 | let cols = jac.shape().1; 94 | let jacobi_scaling_vec: Vec> = (0..cols) 95 | .map(|c| { 96 | let v = jac.val_of_col(c).iter().map(|&i| i * i).sum::().sqrt(); 97 | Triplet::new(c, c, 1.0 / (1.0 + v)) 98 | }) 99 | .collect(); 100 | 101 | jacobi_scaling_diagonal = Some( 102 | faer::sparse::SparseColMat::::try_new_from_triplets( 103 | cols, 104 | cols, 105 | &jacobi_scaling_vec, 106 | ) 107 | .unwrap(), 108 | ); 109 | } 110 | 111 | // Scale the current jacobian by the diagonal matrix 112 | jac = jac * jacobi_scaling_diagonal.as_ref().unwrap(); 113 | 114 | // J^T * J = Matrix of shape (total_variable_dimension, total_variable_dimension) 115 | let jtj = jac 116 | .as_ref() 117 | .transpose() 118 | .to_col_major() 119 | .unwrap() 120 | .mul(jac.as_ref()); 121 | 122 | // J^T * -r = Matrix of shape (total_variable_dimension, 1) 123 | let jtr = jac.as_ref().transpose().mul(-&residuals); 124 | 125 | // Regularize the diagonal of jtj between the min and max diagonal values. 126 | let mut jtj_regularized = jtj.clone(); 127 | for i in 0..total_variable_dimension { 128 | jtj_regularized[(i, i)] += 129 | u * (jtj[(i, i)].max(self.min_diagonal)).min(self.max_diagonal); 130 | } 131 | 132 | let start = Instant::now(); 133 | if let Some(lm_step) = linear_solver.solve_jtj(&jtr, &jtj_regularized) { 134 | let duration = start.elapsed(); 135 | let dx = jacobi_scaling_diagonal.as_ref().unwrap() * &lm_step; 136 | 137 | trace!("Time elapsed in solve Ax=b is: {:?}", duration); 138 | 139 | let dx_na = dx.as_ref().into_nalgebra().column(0).clone_owned(); 140 | 141 | let mut new_param_blocks = parameter_blocks.clone(); 142 | 143 | self.apply_dx2( 144 | &dx_na, 145 | &mut new_param_blocks, 146 | &variable_name_to_col_idx_dict, 147 | ); 148 | 149 | // Compute residuals of (x + dx) 150 | let new_residuals = problem.compute_residuals(&new_param_blocks, true); 151 | 152 | // rho is the ratio between the actual reduction in error and the reduction 153 | // in error if the problem were linear. 154 | let actual_residual_change = 155 | residuals.as_ref().squared_norm_l2() - new_residuals.as_ref().squared_norm_l2(); 156 | trace!("actual_residual_change {}", actual_residual_change); 157 | let linear_residual_change: faer::Mat = 158 | lm_step.transpose().mul(2.0 * &jtr - &jtj * &lm_step); 159 | let rho = actual_residual_change / linear_residual_change[(0, 0)]; 160 | 161 | if rho > 0.0 { 162 | // The linear model appears to be fitting, so accept (x + dx) as the new x. 163 | parameter_blocks = new_param_blocks; 164 | 165 | // Increase the trust region by reducing u 166 | let tmp = 2.0 * rho - 1.0; 167 | u *= (1.0_f64 / 3.0).max(1.0 - tmp * tmp * tmp); 168 | } else { 169 | // If there's too much divergence, reduce the trust region and try again with the same parameters. 170 | u *= 2.0; 171 | trace!("u {}", u); 172 | } 173 | } else { 174 | log::debug!("solve ax=b failed"); 175 | return None; 176 | } 177 | 178 | current_error = self.compute_error(problem, ¶meter_blocks); 179 | trace!("iter:{} total err:{}", i, current_error); 180 | 181 | if current_error < opt_option.min_error_threshold { 182 | trace!("error too low"); 183 | break; 184 | } else if current_error.is_nan() { 185 | log::debug!("solve ax=b failed, current error is nan"); 186 | return None; 187 | } 188 | 189 | if (last_err - current_error).abs() < opt_option.min_abs_error_decrease_threshold { 190 | trace!("absolute error decrease low"); 191 | break; 192 | } else if (last_err - current_error).abs() / last_err 193 | < opt_option.min_rel_error_decrease_threshold 194 | { 195 | trace!("relative error decrease low"); 196 | break; 197 | } 198 | } 199 | let params = parameter_blocks 200 | .iter() 201 | .map(|(k, v)| (k.to_owned(), v.params.clone())) 202 | .collect(); 203 | Some(params) 204 | } 205 | } 206 | -------------------------------------------------------------------------------- /src/optimizer/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod common; 2 | pub mod gauss_newton_optimizer; 3 | pub mod levenberg_marquardt_optimizer; 4 | 5 | pub use common::*; 6 | pub use gauss_newton_optimizer::*; 7 | pub use levenberg_marquardt_optimizer::*; 8 | -------------------------------------------------------------------------------- /src/parameter_block.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | collections::{HashMap, HashSet}, 3 | sync::Arc, 4 | }; 5 | 6 | use nalgebra as na; 7 | use num_dual::DualDVec64; 8 | 9 | use crate::manifold::Manifold; 10 | 11 | #[derive(Clone)] 12 | pub struct ParameterBlock { 13 | pub params: na::DVector, 14 | pub fixed_variables: HashSet, 15 | pub variable_bounds: HashMap, 16 | pub manifold: Option>, 17 | } 18 | 19 | impl ParameterBlock { 20 | pub fn from_vec(params: na::DVector) -> Self { 21 | ParameterBlock { 22 | params, 23 | fixed_variables: HashSet::new(), 24 | variable_bounds: HashMap::new(), 25 | manifold: None, 26 | } 27 | } 28 | pub fn set_manifold(&mut self, manifold: Arc) { 29 | self.manifold = Some(manifold); 30 | } 31 | pub fn ambient_size(&self) -> usize { 32 | self.params.shape().0 33 | } 34 | pub fn tangent_size(&self) -> usize { 35 | if let Some(m) = &self.manifold { 36 | m.tangent_size().get() 37 | } else { 38 | self.ambient_size() 39 | } 40 | } 41 | pub fn plus_f64(&self, dx: na::DVectorView) -> na::DVector { 42 | let mut new_param = na::DVector::zeros(self.ambient_size()); 43 | if let Some(m) = &self.manifold { 44 | new_param = m.plus_f64(self.params.as_view(), dx); 45 | } else { 46 | self.params.add_to(&dx, &mut new_param); 47 | } 48 | new_param 49 | } 50 | pub fn plus_dual(&self, dx: na::DVectorView) -> na::DVector { 51 | let mut new_param = na::DVector::zeros(self.ambient_size()); 52 | if let Some(m) = &self.manifold { 53 | new_param = m.plus_dual(self.params.clone().cast::().as_view(), dx); 54 | } else { 55 | self.params.clone().cast().add_to(&dx, &mut new_param); 56 | } 57 | new_param 58 | } 59 | pub fn y_minus_f64(&self, y: na::DVectorView) -> na::DVector { 60 | let mut delta_x = na::DVector::zeros(self.tangent_size()); 61 | if let Some(m) = &self.manifold { 62 | delta_x = m.minus_f64(y, self.params.as_view()); 63 | } else { 64 | y.sub_to(&self.params, &mut delta_x); 65 | } 66 | delta_x 67 | } 68 | pub fn y_minus_dual(&self, y: na::DVectorView) -> na::DVector { 69 | let mut delta_x = na::DVector::zeros(self.tangent_size()); 70 | if let Some(m) = &self.manifold { 71 | delta_x = m.minus_dual(y, self.params.clone().cast().as_view()); 72 | } else { 73 | y.sub_to(&self.params.clone().cast(), &mut delta_x); 74 | } 75 | delta_x 76 | } 77 | pub fn update_params(&mut self, mut new_param: na::DVector) { 78 | // bound 79 | for (&idx, &(lower, upper)) in &self.variable_bounds { 80 | new_param[idx] = new_param[idx].max(lower).min(upper); 81 | } 82 | 83 | // fix 84 | for &index_to_fix in &self.fixed_variables { 85 | new_param[index_to_fix] = self.params[index_to_fix]; 86 | } 87 | self.params = new_param; 88 | } 89 | } 90 | -------------------------------------------------------------------------------- /src/problem.rs: -------------------------------------------------------------------------------- 1 | use std::collections::{HashMap, HashSet}; 2 | use std::sync::{Arc, Mutex}; 3 | 4 | use faer::sparse::{Argsort, Pair, SparseColMat, SymbolicSparseColMat}; 5 | use faer_ext::IntoFaer; 6 | use nalgebra as na; 7 | use rayon::prelude::*; 8 | 9 | use crate::manifold::Manifold; 10 | use crate::parameter_block::ParameterBlock; 11 | use crate::{factors, loss_functions, residual_block}; 12 | 13 | type ResidualBlockId = usize; 14 | 15 | pub struct Problem { 16 | pub total_residual_dimension: usize, 17 | residual_id_count: usize, 18 | residual_blocks: HashMap, 19 | pub fixed_variable_indexes: HashMap>, 20 | pub variable_bounds: HashMap>, 21 | pub variable_manifold: HashMap>, 22 | } 23 | impl Default for Problem { 24 | fn default() -> Self { 25 | Self::new() 26 | } 27 | } 28 | 29 | pub struct SymbolicStructure { 30 | pattern: SymbolicSparseColMat, 31 | order: Argsort, 32 | } 33 | 34 | type JacobianValue = f64; 35 | 36 | impl Problem { 37 | pub fn new() -> Problem { 38 | Problem { 39 | total_residual_dimension: 0, 40 | residual_id_count: 0, 41 | residual_blocks: HashMap::new(), 42 | fixed_variable_indexes: HashMap::new(), 43 | variable_bounds: HashMap::new(), 44 | variable_manifold: HashMap::new(), 45 | } 46 | } 47 | 48 | pub fn build_symbolic_structure( 49 | &self, 50 | parameter_blocks: &HashMap, 51 | total_variable_dimension: usize, 52 | variable_name_to_col_idx_dict: &HashMap, 53 | ) -> SymbolicStructure { 54 | let mut indices = Vec::>::new(); 55 | 56 | self.residual_blocks.iter().for_each(|(_, residual_block)| { 57 | let mut variable_local_idx_size_list = Vec::<(usize, usize)>::new(); 58 | let mut count_variable_local_idx: usize = 0; 59 | for var_key in &residual_block.variable_key_list { 60 | if let Some(param) = parameter_blocks.get(var_key) { 61 | variable_local_idx_size_list 62 | .push((count_variable_local_idx, param.tangent_size())); 63 | count_variable_local_idx += param.tangent_size(); 64 | }; 65 | } 66 | for (i, var_key) in residual_block.variable_key_list.iter().enumerate() { 67 | if let Some(variable_global_idx) = variable_name_to_col_idx_dict.get(var_key) { 68 | let (_, var_size) = variable_local_idx_size_list[i]; 69 | for row_idx in 0..residual_block.dim_residual { 70 | for col_idx in 0..var_size { 71 | let global_row_idx = residual_block.residual_row_start_idx + row_idx; 72 | let global_col_idx = variable_global_idx + col_idx; 73 | indices.push(Pair::new(global_row_idx, global_col_idx)); 74 | } 75 | } 76 | } 77 | } 78 | }); 79 | let start = std::time::Instant::now(); 80 | let (s, o) = SymbolicSparseColMat::try_new_from_indices( 81 | self.total_residual_dimension, 82 | total_variable_dimension, 83 | &indices, 84 | ) 85 | .unwrap(); 86 | log::trace!("Built symbolic matrix: {:?}", start.elapsed()); 87 | SymbolicStructure { 88 | pattern: s, 89 | order: o, 90 | } 91 | } 92 | 93 | pub fn get_variable_name_to_col_idx_dict( 94 | &self, 95 | parameter_blocks: &HashMap, 96 | ) -> HashMap { 97 | let mut count_col_idx = 0; 98 | let mut variable_name_to_col_idx_dict = HashMap::new(); 99 | parameter_blocks 100 | .iter() 101 | .for_each(|(param_name, param_block)| { 102 | variable_name_to_col_idx_dict.insert(param_name.to_owned(), count_col_idx); 103 | count_col_idx += param_block.tangent_size(); 104 | }); 105 | variable_name_to_col_idx_dict 106 | } 107 | pub fn add_residual_block( 108 | &mut self, 109 | dim_residual: usize, 110 | variable_key_size_list: &[&str], 111 | factor: Box, 112 | loss_func: Option>, 113 | ) -> ResidualBlockId { 114 | self.residual_blocks.insert( 115 | self.residual_id_count, 116 | residual_block::ResidualBlock::new( 117 | self.residual_id_count, 118 | dim_residual, 119 | self.total_residual_dimension, 120 | variable_key_size_list, 121 | factor, 122 | loss_func, 123 | ), 124 | ); 125 | let block_id = self.residual_id_count; 126 | self.residual_id_count += 1; 127 | 128 | self.total_residual_dimension += dim_residual; 129 | 130 | block_id 131 | } 132 | pub fn remove_residual_block( 133 | &mut self, 134 | block_id: ResidualBlockId, 135 | ) -> Option { 136 | if let Some(residual_block) = self.residual_blocks.remove(&block_id) { 137 | self.total_residual_dimension -= residual_block.dim_residual; 138 | Some(residual_block) 139 | } else { 140 | None 141 | } 142 | } 143 | pub fn fix_variable(&mut self, var_to_fix: &str, idx: usize) { 144 | if let Some(var_mut) = self.fixed_variable_indexes.get_mut(var_to_fix) { 145 | var_mut.insert(idx); 146 | } else { 147 | self.fixed_variable_indexes 148 | .insert(var_to_fix.to_owned(), HashSet::from([idx])); 149 | } 150 | } 151 | pub fn unfix_variable(&mut self, var_to_unfix: &str) { 152 | self.fixed_variable_indexes.remove(var_to_unfix); 153 | } 154 | pub fn set_variable_bounds( 155 | &mut self, 156 | var_to_bound: &str, 157 | idx: usize, 158 | lower_bound: f64, 159 | upper_bound: f64, 160 | ) { 161 | if lower_bound > upper_bound { 162 | log::error!("lower bound is larger than upper bound"); 163 | } else if let Some(var_mut) = self.variable_bounds.get_mut(var_to_bound) { 164 | var_mut.insert(idx, (lower_bound, upper_bound)); 165 | } else { 166 | self.variable_bounds.insert( 167 | var_to_bound.to_owned(), 168 | HashMap::from([(idx, (lower_bound, upper_bound))]), 169 | ); 170 | } 171 | } 172 | pub fn set_variable_manifold( 173 | &mut self, 174 | var_name: &str, 175 | manifold: Arc, 176 | ) { 177 | self.variable_manifold 178 | .insert(var_name.to_string(), manifold); 179 | } 180 | pub fn remove_variable_bounds(&mut self, var_to_unbound: &str) { 181 | self.variable_bounds.remove(var_to_unbound); 182 | } 183 | pub fn initialize_parameter_blocks( 184 | &self, 185 | initial_values: &HashMap>, 186 | ) -> HashMap { 187 | let parameter_blocks: HashMap = initial_values 188 | .iter() 189 | .map(|(k, v)| { 190 | let mut p_block = ParameterBlock::from_vec(v.clone()); 191 | if let Some(indexes) = self.fixed_variable_indexes.get(k) { 192 | p_block.fixed_variables = indexes.clone(); 193 | } 194 | if let Some(bounds) = self.variable_bounds.get(k) { 195 | p_block.variable_bounds = bounds.clone(); 196 | } 197 | if let Some(manifold) = self.variable_manifold.get(k) { 198 | p_block.manifold = Some(manifold.clone()) 199 | } 200 | 201 | (k.to_owned(), p_block) 202 | }) 203 | .collect(); 204 | parameter_blocks 205 | } 206 | 207 | pub fn compute_residuals( 208 | &self, 209 | parameter_blocks: &HashMap, 210 | with_loss_fn: bool, 211 | ) -> faer::Mat { 212 | let total_residual = Arc::new(Mutex::new(na::DVector::::zeros( 213 | self.total_residual_dimension, 214 | ))); 215 | self.residual_blocks 216 | .par_iter() 217 | .for_each(|(_, residual_block)| { 218 | self.compute_residual_impl( 219 | residual_block, 220 | parameter_blocks, 221 | &total_residual, 222 | with_loss_fn, 223 | ) 224 | }); 225 | let total_residual = Arc::try_unwrap(total_residual) 226 | .unwrap() 227 | .into_inner() 228 | .unwrap(); 229 | 230 | total_residual.view_range(.., ..).into_faer().to_owned() 231 | } 232 | 233 | pub fn compute_residual_and_jacobian( 234 | &self, 235 | parameter_blocks: &HashMap, 236 | variable_name_to_col_idx_dict: &HashMap, 237 | symbolic_structure: &SymbolicStructure, 238 | ) -> (faer::Mat, SparseColMat) { 239 | // multi 240 | let total_residual = Arc::new(Mutex::new(na::DVector::::zeros( 241 | self.total_residual_dimension, 242 | ))); 243 | 244 | let jacobian_lists: Vec = self 245 | .residual_blocks 246 | .par_iter() 247 | .map(|(_, residual_block)| { 248 | self.compute_residual_and_jacobian_impl( 249 | residual_block, 250 | parameter_blocks, 251 | variable_name_to_col_idx_dict, 252 | &total_residual, 253 | ) 254 | }) 255 | .flatten() 256 | .collect(); 257 | 258 | let total_residual = Arc::try_unwrap(total_residual) 259 | .unwrap() 260 | .into_inner() 261 | .unwrap(); 262 | 263 | let residual_faer = total_residual.view_range(.., ..).into_faer().to_owned(); 264 | let jacobian_faer = SparseColMat::new_from_argsort( 265 | symbolic_structure.pattern.clone(), 266 | &symbolic_structure.order, 267 | jacobian_lists.as_slice(), 268 | ) 269 | .unwrap(); 270 | (residual_faer, jacobian_faer) 271 | } 272 | 273 | fn compute_residual_impl( 274 | &self, 275 | residual_block: &crate::ResidualBlock, 276 | parameter_blocks: &HashMap, 277 | total_residual: &Arc>>, 278 | with_loss_fn: bool, 279 | ) { 280 | let mut params = Vec::new(); 281 | for var_key in &residual_block.variable_key_list { 282 | if let Some(param) = parameter_blocks.get(var_key) { 283 | params.push(param); 284 | }; 285 | } 286 | let res = residual_block.residual(¶ms, with_loss_fn); 287 | 288 | { 289 | let mut total_residual = total_residual.lock().unwrap(); 290 | total_residual 291 | .rows_mut( 292 | residual_block.residual_row_start_idx, 293 | residual_block.dim_residual, 294 | ) 295 | .copy_from(&res); 296 | } 297 | } 298 | 299 | fn compute_residual_and_jacobian_impl( 300 | &self, 301 | residual_block: &crate::ResidualBlock, 302 | parameter_blocks: &HashMap, 303 | variable_name_to_col_idx_dict: &HashMap, 304 | total_residual: &Arc>>, 305 | ) -> Vec { 306 | let mut params = Vec::new(); 307 | let mut variable_local_idx_size_list = Vec::<(usize, usize)>::new(); 308 | let mut count_variable_local_idx: usize = 0; 309 | for var_key in &residual_block.variable_key_list { 310 | if let Some(param) = parameter_blocks.get(var_key) { 311 | params.push(param); 312 | variable_local_idx_size_list.push((count_variable_local_idx, param.tangent_size())); 313 | count_variable_local_idx += param.tangent_size(); 314 | }; 315 | } 316 | let (res, jac) = residual_block.residual_and_jacobian(¶ms); 317 | { 318 | let mut total_residual = total_residual.lock().unwrap(); 319 | total_residual 320 | .rows_mut( 321 | residual_block.residual_row_start_idx, 322 | residual_block.dim_residual, 323 | ) 324 | .copy_from(&res); 325 | } 326 | 327 | let mut local_jacobian_list = Vec::new(); 328 | 329 | for (i, var_key) in residual_block.variable_key_list.iter().enumerate() { 330 | if variable_name_to_col_idx_dict.contains_key(var_key) { 331 | let (variable_local_idx, var_size) = variable_local_idx_size_list[i]; 332 | let variable_jac = jac.view((0, variable_local_idx), (jac.shape().0, var_size)); 333 | for row_idx in 0..jac.shape().0 { 334 | for col_idx in 0..var_size { 335 | local_jacobian_list.push(variable_jac[(row_idx, col_idx)]); 336 | } 337 | } 338 | } else { 339 | panic!( 340 | "Missing key {} in variable-to-column-index mapping", 341 | var_key 342 | ); 343 | } 344 | } 345 | 346 | local_jacobian_list 347 | } 348 | } 349 | -------------------------------------------------------------------------------- /src/python/mod.rs: -------------------------------------------------------------------------------- 1 | use pyo3::prelude::*; 2 | 3 | mod py_factors; 4 | mod py_linear; 5 | mod py_loss_functions; 6 | mod py_optimizer; 7 | mod py_problem; 8 | use self::py_factors::*; 9 | use self::py_linear::*; 10 | use self::py_loss_functions::*; 11 | use self::py_optimizer::*; 12 | use self::py_problem::*; 13 | 14 | fn register_child_module(parent_module: &Bound<'_, PyModule>) -> PyResult<()> { 15 | // For factors submodule 16 | let factors_module = PyModule::new_bound(parent_module.py(), "factors")?; 17 | factors_module.add_class::()?; 18 | factors_module.add_class::()?; 19 | factors_module.add_class::()?; 20 | parent_module.add_submodule(&factors_module)?; 21 | parent_module 22 | .py() 23 | .import_bound("sys")? 24 | .getattr("modules")? 25 | .set_item("tiny_solver.factors", factors_module)?; 26 | 27 | let loss_functions_module = PyModule::new_bound(parent_module.py(), "loss_functions")?; 28 | loss_functions_module.add_class::()?; 29 | parent_module.add_submodule(&loss_functions_module)?; 30 | parent_module 31 | .py() 32 | .import_bound("sys")? 33 | .getattr("modules")? 34 | .set_item("tiny_solver.loss_functions", loss_functions_module)?; 35 | Ok(()) 36 | } 37 | 38 | /// A Python module implemented in Rust. 39 | #[pymodule] 40 | pub fn tiny_solver(m: &Bound<'_, PyModule>) -> PyResult<()> { 41 | // pyo3_log::init(); 42 | m.add("__version__", env!("CARGO_PKG_VERSION"))?; 43 | m.add_class::()?; 44 | m.add_class::()?; 45 | m.add_class::()?; 46 | m.add_class::()?; 47 | register_child_module(m)?; 48 | 49 | Ok(()) 50 | } 51 | -------------------------------------------------------------------------------- /src/python/py_factors.rs: -------------------------------------------------------------------------------- 1 | use nalgebra as na; 2 | use num_dual::python::PyDual64Dyn; 3 | use numpy::PyReadonlyArray1; 4 | use pyo3::prelude::*; 5 | use pyo3::types::PyTuple; 6 | 7 | use crate::factors::*; 8 | 9 | #[derive(Debug, Clone)] 10 | #[pyclass(name = "BetweenFactorSE2")] 11 | pub struct PyBetweenFactorSE2(pub BetweenFactorSE2); 12 | 13 | #[pymethods] 14 | impl PyBetweenFactorSE2 { 15 | #[new] 16 | pub fn new(x: f64, y: f64, theta: f64) -> Self { 17 | PyBetweenFactorSE2(BetweenFactorSE2 { 18 | dx: x, 19 | dy: y, 20 | dtheta: theta, 21 | }) 22 | } 23 | } 24 | 25 | #[derive(Debug, Clone)] 26 | #[pyclass(name = "PriorFactor")] 27 | pub struct PyPriorFactor(pub PriorFactor); 28 | 29 | #[pymethods] 30 | impl PyPriorFactor { 31 | #[new] 32 | pub fn new(x: PyReadonlyArray1) -> Self { 33 | PyPriorFactor(PriorFactor { 34 | v: x.as_matrix().column(0).into(), 35 | }) 36 | } 37 | } 38 | 39 | #[pyclass] 40 | #[derive(Debug, Clone)] 41 | pub struct PyFactor { 42 | pub func: Py, 43 | } 44 | 45 | #[pymethods] 46 | impl PyFactor { 47 | #[new] 48 | pub fn new(f: Py) -> Self { 49 | PyFactor { func: f } 50 | } 51 | } 52 | 53 | impl Factor for PyFactor { 54 | fn residual_func( 55 | &self, 56 | params: &[na::DVector], 57 | ) -> na::DVector { 58 | // this can not be called with par_iter 59 | // TODO find a way to deal with multi threading 60 | let residual_py = Python::with_gil(|py| -> PyResult> { 61 | let py_params: Vec> = params 62 | .iter() 63 | .map(|param| { 64 | param 65 | .data 66 | .as_vec() 67 | .iter() 68 | .map(|x| PyDual64Dyn::from(x.clone())) 69 | .collect::>() 70 | }) 71 | .map(|x| x.into_py(py)) 72 | .collect(); 73 | let args = PyTuple::new_bound(py, py_params); 74 | let result = self.func.call1(py, args); 75 | let residual_py = result.unwrap().extract::>(py); 76 | residual_py 77 | }); 78 | let residual_py: Vec = residual_py 79 | .unwrap() 80 | .iter() 81 | .map(|x| ::clone(x).into()) 82 | .collect(); 83 | na::DVector::from_vec(residual_py) 84 | } 85 | } 86 | -------------------------------------------------------------------------------- /src/python/py_linear.rs: -------------------------------------------------------------------------------- 1 | use crate::linear::LinearSolver; 2 | use pyo3::prelude::*; 3 | 4 | #[derive(Clone)] 5 | #[pyclass(name = "LinearSolver")] 6 | pub struct PyLinearSolver(pub LinearSolver); 7 | -------------------------------------------------------------------------------- /src/python/py_loss_functions.rs: -------------------------------------------------------------------------------- 1 | use pyo3::prelude::*; 2 | 3 | use crate::loss_functions::*; 4 | 5 | #[pyclass(name = "HuberLoss")] 6 | #[derive(Clone)] 7 | pub struct PyHuberLoss(pub HuberLoss); 8 | 9 | #[pymethods] 10 | impl PyHuberLoss { 11 | #[new] 12 | #[pyo3(signature=(scale=1.0))] 13 | pub fn new_py(scale: f64) -> Self { 14 | PyHuberLoss(HuberLoss::new(scale)) 15 | } 16 | } 17 | -------------------------------------------------------------------------------- /src/python/py_optimizer.rs: -------------------------------------------------------------------------------- 1 | use log::info; 2 | use std::collections::HashMap; 3 | 4 | use numpy::{PyArray2, PyReadonlyArray1, ToPyArray}; 5 | use pyo3::prelude::*; 6 | use pyo3::types::PyDict; 7 | 8 | use super::py_linear::PyLinearSolver; 9 | use super::py_problem::PyProblem; 10 | use crate::optimizer::Optimizer; 11 | use crate::{GaussNewtonOptimizer, LinearSolver, OptimizerOptions}; 12 | 13 | #[pyclass(name = "GaussNewtonOptimizer")] 14 | pub struct PyGaussNewtonOptimizer(GaussNewtonOptimizer); 15 | 16 | #[pymethods] 17 | impl PyGaussNewtonOptimizer { 18 | #[new] 19 | pub fn new() -> Self { 20 | info!("init GaussNewtonOptimizer"); 21 | PyGaussNewtonOptimizer(GaussNewtonOptimizer {}) 22 | } 23 | 24 | #[pyo3(name = "optimize")] 25 | #[pyo3(signature=(problem, initial_values, optimizer_options=None))] 26 | pub fn optimize_py( 27 | &self, 28 | py: Python<'_>, 29 | problem: &PyProblem, 30 | initial_values: &Bound<'_, PyDict>, 31 | optimizer_options: Option, 32 | ) -> PyResult>>> { 33 | let init_values: HashMap> = initial_values.extract().unwrap(); 34 | let init_values: HashMap> = init_values 35 | .iter() 36 | .map(|(k, v)| (k.to_string(), v.as_matrix().column(0).into())) 37 | .collect(); 38 | let result = self 39 | .0 40 | .optimize(&problem.0, &init_values, Some(optimizer_options.unwrap().0)) 41 | .unwrap(); 42 | 43 | let output_d: HashMap>> = result 44 | .iter() 45 | .map(|(k, v)| (k.to_string(), v.to_pyarray_bound(py).to_owned().into())) 46 | .collect(); 47 | Ok(output_d) 48 | } 49 | } 50 | 51 | #[pyclass(name = "OptimizerOptions")] 52 | #[derive(Clone)] 53 | pub struct PyOptimizerOptions(pub OptimizerOptions); 54 | 55 | #[pymethods] 56 | impl PyOptimizerOptions { 57 | #[new] 58 | #[pyo3(signature = ( 59 | max_iteration=100, 60 | linear_solver_type=PyLinearSolver(LinearSolver::SparseCholesky), 61 | verbosity_level=0, 62 | min_abs_error_decrease_threshold=1e-5, 63 | min_rel_error_decrease_threshold=1e-5, 64 | min_error_threshold=1e-8, 65 | ))] 66 | pub fn new( 67 | max_iteration: usize, 68 | linear_solver_type: PyLinearSolver, 69 | verbosity_level: usize, 70 | min_abs_error_decrease_threshold: f64, 71 | min_rel_error_decrease_threshold: f64, 72 | min_error_threshold: f64, 73 | ) -> Self { 74 | PyOptimizerOptions(OptimizerOptions { 75 | max_iteration, 76 | linear_solver_type: linear_solver_type.0, 77 | verbosity_level, 78 | min_abs_error_decrease_threshold, 79 | min_rel_error_decrease_threshold, 80 | min_error_threshold, 81 | }) 82 | } 83 | } 84 | -------------------------------------------------------------------------------- /src/python/py_problem.rs: -------------------------------------------------------------------------------- 1 | use pyo3::prelude::*; 2 | 3 | use crate::factors::*; 4 | use crate::loss_functions::*; 5 | use crate::problem::Problem; 6 | 7 | use super::py_factors::*; 8 | use super::py_loss_functions::*; 9 | 10 | fn convert_pyany_to_factor(py_any: &Bound<'_, PyAny>) -> PyResult<(bool, Box)> { 11 | let factor_name: String = py_any.get_type().getattr("__name__")?.extract()?; 12 | match factor_name.as_str() { 13 | "BetweenFactorSE2" => { 14 | let factor: PyBetweenFactorSE2 = py_any.extract().unwrap(); 15 | Ok((false, Box::new(factor.0))) 16 | } 17 | "PriorFactor" => { 18 | let factor: PyPriorFactor = py_any.extract().unwrap(); 19 | Ok((false, Box::new(factor.0))) 20 | } 21 | "PyFactor" => { 22 | let factor: PyFactor = py_any.extract().unwrap(); 23 | Ok((true, Box::new(factor))) 24 | } 25 | _ => Err(PyErr::new::( 26 | "Unknown factor type", 27 | )), 28 | } 29 | } 30 | fn convert_pyany_to_loss_function( 31 | py_any: &Bound<'_, PyAny>, 32 | ) -> PyResult>> { 33 | let factor_name: String = py_any.get_type().getattr("__name__")?.extract()?; 34 | match factor_name.as_str() { 35 | "HuberLoss" => { 36 | let loss_func: PyHuberLoss = py_any.extract().unwrap(); 37 | Ok(Some(Box::new(loss_func.0))) 38 | } 39 | "NoneType" => Ok(None), 40 | _ => Err(PyErr::new::( 41 | "Unknown factor type", 42 | )), 43 | } 44 | } 45 | 46 | #[pyclass(name = "Problem")] 47 | pub struct PyProblem(pub Problem); 48 | 49 | #[pymethods] 50 | impl PyProblem { 51 | #[new] 52 | pub fn new_py() -> Self { 53 | PyProblem(Problem::new()) 54 | } 55 | 56 | #[pyo3(name = "add_residual_block")] 57 | pub fn add_residual_block_py( 58 | &mut self, 59 | dim_residual: usize, 60 | variable_key_size_list: Vec<(String, usize)>, 61 | pyfactor: &Bound<'_, PyAny>, 62 | pyloss_func: &Bound<'_, PyAny>, 63 | ) -> PyResult<()> { 64 | let (is_pyfactor, factor) = convert_pyany_to_factor(pyfactor).unwrap(); 65 | self.0.add_residual_block( 66 | dim_residual, 67 | variable_key_size_list, 68 | factor, 69 | convert_pyany_to_loss_function(pyloss_func).unwrap(), 70 | ); 71 | if is_pyfactor { 72 | self.0.has_py_factor() 73 | } 74 | Ok(()) 75 | } 76 | } 77 | -------------------------------------------------------------------------------- /src/residual_block.rs: -------------------------------------------------------------------------------- 1 | use nalgebra as na; 2 | use rayon::prelude::*; 3 | 4 | use crate::corrector::Corrector; 5 | use crate::factors::FactorImpl; 6 | use crate::loss_functions::Loss; 7 | use crate::parameter_block::ParameterBlock; 8 | 9 | pub struct ResidualBlock { 10 | pub residual_block_id: usize, 11 | pub dim_residual: usize, 12 | pub residual_row_start_idx: usize, 13 | pub variable_key_list: Vec, 14 | pub factor: Box, 15 | pub loss_func: Option>, 16 | } 17 | impl ResidualBlock { 18 | pub fn new( 19 | residual_block_id: usize, 20 | dim_residual: usize, 21 | residual_row_start_idx: usize, 22 | variable_key_size_list: &[&str], 23 | factor: Box, 24 | loss_func: Option>, 25 | ) -> Self { 26 | ResidualBlock { 27 | residual_block_id, 28 | dim_residual, 29 | residual_row_start_idx, 30 | variable_key_list: variable_key_size_list 31 | .iter() 32 | .map(|s| s.to_string()) 33 | .collect(), 34 | factor, 35 | loss_func, 36 | } 37 | } 38 | 39 | pub fn residual(&self, params: &[&ParameterBlock], with_loss_fn: bool) -> na::DVector { 40 | let param_vec: Vec<_> = params.iter().map(|p| p.params.clone()).collect(); 41 | let mut residual = self.factor.residual_func_f64(¶m_vec); 42 | let squared_norm = residual.norm_squared(); 43 | if with_loss_fn { 44 | if let Some(loss_func) = self.loss_func.as_ref() { 45 | let rho = loss_func.evaluate(squared_norm); 46 | // let cost = 0.5 * rho[0]; 47 | let corrector = Corrector::new(squared_norm, &rho); 48 | corrector.correct_residuals(&mut residual); 49 | } 50 | } else { 51 | // let cost = 0.5 * squared_norm; 52 | } 53 | residual 54 | } 55 | pub fn residual_and_jacobian( 56 | &self, 57 | params: &[&ParameterBlock], 58 | ) -> (na::DVector, na::DMatrix) { 59 | let variable_rows: Vec = params.iter().map(|x| x.tangent_size()).collect(); 60 | let dim_variable = variable_rows.iter().sum::(); 61 | let variable_row_idx_vec = get_variable_rows(&variable_rows); 62 | let indentity_mat = na::DMatrix::::identity(dim_variable, dim_variable); 63 | 64 | // ambient size 65 | let params_plus_tangent_dual: Vec> = params 66 | .par_iter() 67 | .enumerate() 68 | .map(|(param_idx, param)| { 69 | let zeros_with_dual = na::DVector::from_row_iterator( 70 | param.tangent_size(), 71 | (0..param.tangent_size()).map(|j| { 72 | num_dual::DualDVec64::new( 73 | 0.0, 74 | num_dual::Derivative::some(na::DVector::from( 75 | indentity_mat.column(variable_row_idx_vec[param_idx][j]), 76 | )), 77 | ) 78 | }), 79 | ); 80 | let param_plus_dual = param.plus_dual(zeros_with_dual.as_view()); 81 | param_plus_dual 82 | }) 83 | .collect(); 84 | 85 | // tangent size 86 | let residual_with_jacobian = self.factor.residual_func_dual(¶ms_plus_tangent_dual); 87 | let mut residual = residual_with_jacobian.map(|x| x.re); 88 | let jacobian = residual_with_jacobian 89 | .map(|x| x.eps.unwrap_generic(na::Dyn(dim_variable), na::Const::<1>)); 90 | let mut jacobian = 91 | na::DMatrix::::from_fn(residual_with_jacobian.nrows(), dim_variable, |r, c| { 92 | jacobian[r][c] 93 | }); 94 | let squared_norm = residual.norm_squared(); 95 | if let Some(loss_func) = self.loss_func.as_ref() { 96 | let rho = loss_func.evaluate(squared_norm); 97 | // let cost = 0.5 * rho[0]; 98 | let corrector = Corrector::new(squared_norm, &rho); 99 | corrector.correct_jacobian(&residual, &mut jacobian); 100 | corrector.correct_residuals(&mut residual); 101 | } else { 102 | // let cost = 0.5 * squared_norm; 103 | } 104 | (residual, jacobian) 105 | } 106 | } 107 | 108 | fn get_variable_rows(variable_rows: &[usize]) -> Vec> { 109 | let mut result = Vec::with_capacity(variable_rows.len()); 110 | let mut current = 0; 111 | for &num in variable_rows { 112 | let next = current + num; 113 | let range = (current..next).collect(); 114 | result.push(range); 115 | current = next; 116 | } 117 | result 118 | } 119 | -------------------------------------------------------------------------------- /tests/test_factors.rs: -------------------------------------------------------------------------------- 1 | #[cfg(test)] 2 | mod tests { 3 | use nalgebra as na; 4 | use tiny_solver::factors::*; 5 | 6 | #[test] 7 | fn prior_factor() { 8 | let prior_factor = PriorFactor { 9 | v: na::dvector![3.0, 1.0], 10 | }; 11 | 12 | let params = na::dvector![1.0, 2.0]; 13 | 14 | let residual = prior_factor.residual_func(&[params]); 15 | assert_eq!(residual[0], -2.0); 16 | assert_eq!(residual[1], 1.0); 17 | // the parameters contains two values [a, b] 18 | let dual_params = na::dvector![ 19 | // dual vec has [1.0, 0.0] means it's for tracking the derivative of the first parameter a 20 | num_dual::DualDVec64::new(1.0, num_dual::Derivative::some(na::dvector![1.0, 0.0])), 21 | // dual vec has [0.0, 1.0] is for tracking the derivative of the second parameter b 22 | num_dual::DualDVec64::new(2.0, num_dual::Derivative::some(na::dvector![0.0, 1.0])) 23 | ]; 24 | // there are two residuals r0, r1 25 | let residual_with_jacobian = prior_factor.residual_func_dual(&[dual_params]); 26 | assert!(residual_with_jacobian[0].re == -2.0); 27 | assert!(residual_with_jacobian[1].re == 1.0); 28 | 29 | let jacobian = 30 | residual_with_jacobian.map(|x| x.eps.unwrap_generic(na::Dyn(2), na::Const::<1>)); 31 | 32 | // partial derivative of r0 with respect to a 33 | let d_r0_d_a: f64 = jacobian[0][0]; 34 | // partial derivative of r0 with respect to b 35 | let d_r0_d_b: f64 = jacobian[0][1]; 36 | 37 | // partial derivative of r1 with respect to a 38 | let d_r1_d_a: f64 = jacobian[1][0]; 39 | // partial derivative of r1 with respect to b 40 | let d_r1_d_b: f64 = jacobian[1][1]; 41 | 42 | // a only contribute to r0 and b only contribute to r1 43 | assert!(d_r0_d_a == 1.0); 44 | assert!(d_r0_d_b == 0.0); 45 | assert!(d_r1_d_a == 0.0); 46 | assert!(d_r1_d_b == 1.0); 47 | } 48 | 49 | #[test] 50 | fn between_factor_se2() { 51 | let factor = BetweenFactorSE2 { 52 | dtheta: 1.0, 53 | dx: 2.0, 54 | dy: 3.0, 55 | }; 56 | 57 | let params = [na::dvector![1.0, 2.0, 3.0], na::dvector![1.0, 2.0, 3.0]]; 58 | 59 | let residual = factor.residual_func(¶ms); 60 | assert_eq!(residual, na::dvector![2.0, 3.0, 1.0]); 61 | } 62 | } 63 | -------------------------------------------------------------------------------- /tests/test_loss.rs: -------------------------------------------------------------------------------- 1 | #[cfg(test)] 2 | mod tests { 3 | use core::f64; 4 | use tiny_solver::loss_functions::*; 5 | 6 | #[test] 7 | fn arctan_loss() { 8 | let tolerance = 100.0; 9 | let asymptote = tolerance * f64::consts::PI / 2.0; 10 | 11 | let arctan_loss = ArctanLoss::new(tolerance); 12 | 13 | let rho1 = arctan_loss.evaluate(1.0); 14 | let rho2 = arctan_loss.evaluate(30.0); 15 | 16 | // Test that rho[0] grows linearly with the scale 17 | assert!(rho1[0] < rho2[0]); 18 | 19 | // Test that rho[1] decreases linearly with the scale 20 | assert!(rho1[1] > rho2[1]); 21 | 22 | // Test that scales largely above the tolerance are asymptotically bounded 23 | assert!(arctan_loss.evaluate(tolerance * tolerance * tolerance)[0] < asymptote); 24 | } 25 | } 26 | -------------------------------------------------------------------------------- /tests/test_manifold.rs: -------------------------------------------------------------------------------- 1 | #[cfg(test)] 2 | mod tests { 3 | use nalgebra as na; 4 | use std::f64::consts::PI; 5 | 6 | use tiny_solver::manifold::so3::SO3; 7 | 8 | fn equal_to_na(so3: &SO3) -> bool { 9 | let q = so3.to_vec(); 10 | let na_so3 = 11 | na::UnitQuaternion::from_quaternion(na::Quaternion::new(q[3], q[0], q[1], q[2])); 12 | let diff = so3.log() - na_so3.scaled_axis(); 13 | diff.norm() < 1e-6 14 | } 15 | 16 | #[test] 17 | fn test_so3() { 18 | for _ in 0..10000 { 19 | let mut rvec = na::DVector::new_random(3); 20 | rvec /= rvec.norm(); 21 | rvec *= rand::random::() * PI; 22 | let r = SO3::exp(rvec.as_view()); 23 | assert!(equal_to_na(&r)); 24 | } 25 | } 26 | } 27 | -------------------------------------------------------------------------------- /tests/test_problem.rs: -------------------------------------------------------------------------------- 1 | #[cfg(test)] 2 | mod tests { 3 | use std::collections::HashMap; 4 | 5 | use nalgebra as na; 6 | use tiny_solver; 7 | 8 | #[test] 9 | fn new_problem() { 10 | let problem = tiny_solver::Problem::new(); 11 | assert_eq!(problem.total_residual_dimension, 0); 12 | assert_eq!(problem.fixed_variable_indexes.len(), 0); 13 | assert_eq!(problem.variable_bounds.len(), 0); 14 | } 15 | 16 | #[test] 17 | fn add_residual_block() { 18 | let mut problem = tiny_solver::Problem::new(); 19 | let block_id1 = problem.add_residual_block( 20 | 1, 21 | &["x"], 22 | Box::new(tiny_solver::factors::PriorFactor { 23 | v: na::dvector![3.0], 24 | }), 25 | None, 26 | ); 27 | 28 | assert_eq!(problem.total_residual_dimension, 1); 29 | 30 | let block_id2 = problem.add_residual_block( 31 | 1, 32 | &["y"], 33 | Box::new(tiny_solver::factors::PriorFactor { 34 | v: na::dvector![3.0], 35 | }), 36 | None, 37 | ); 38 | 39 | assert!(block_id1 != block_id2); 40 | assert_eq!(problem.total_residual_dimension, 2); 41 | } 42 | 43 | #[test] 44 | fn remove_residual_block() { 45 | let mut problem = tiny_solver::Problem::new(); 46 | let block_id = problem.add_residual_block( 47 | 1, 48 | &["x"], 49 | Box::new(tiny_solver::factors::PriorFactor { 50 | v: na::dvector![3.0], 51 | }), 52 | None, 53 | ); 54 | 55 | assert_eq!(problem.total_residual_dimension, 1); 56 | 57 | let mut block = problem.remove_residual_block(block_id); 58 | assert!(block.is_some()); 59 | assert_eq!(problem.total_residual_dimension, 0); 60 | 61 | block = problem.remove_residual_block(block_id); 62 | assert!(block.is_none()); 63 | assert_eq!(problem.total_residual_dimension, 0); 64 | } 65 | 66 | #[test] 67 | fn fix_variable() { 68 | let mut problem = tiny_solver::Problem::new(); 69 | problem.fix_variable("x", 0); 70 | 71 | assert_eq!(problem.fixed_variable_indexes.len(), 1); 72 | assert_eq!(problem.fixed_variable_indexes["x"].len(), 1); 73 | assert!(problem.fixed_variable_indexes["x"].contains(&0)); 74 | 75 | problem.fix_variable("x", 1); 76 | 77 | assert_eq!(problem.fixed_variable_indexes.len(), 1); 78 | assert_eq!(problem.fixed_variable_indexes["x"].len(), 2); 79 | assert!(problem.fixed_variable_indexes["x"].contains(&1)); 80 | } 81 | 82 | #[test] 83 | fn unfix_variable() { 84 | let mut problem = tiny_solver::Problem::new(); 85 | problem.fix_variable("x", 0); 86 | 87 | assert_eq!(problem.fixed_variable_indexes.len(), 1); 88 | assert_eq!(problem.fixed_variable_indexes["x"].len(), 1); 89 | assert!(problem.fixed_variable_indexes["x"].contains(&0)); 90 | 91 | problem.unfix_variable("x"); 92 | 93 | assert_eq!(problem.fixed_variable_indexes.len(), 0); 94 | } 95 | 96 | #[test] 97 | fn set_variable_bounds() { 98 | let mut problem = tiny_solver::Problem::new(); 99 | 100 | problem.set_variable_bounds("x", 0, 0.0, 1.0); 101 | assert_eq!(problem.variable_bounds.len(), 1); 102 | assert_eq!(problem.variable_bounds["x"].len(), 1); 103 | assert!(problem.variable_bounds["x"].contains_key(&0)); 104 | assert_eq!(problem.variable_bounds["x"][&0], (0.0, 1.0)); 105 | } 106 | 107 | #[test] 108 | fn remove_variable_bounds() { 109 | let mut problem = tiny_solver::Problem::new(); 110 | 111 | problem.set_variable_bounds("x", 0, 0.0, 1.0); 112 | assert_eq!(problem.variable_bounds.len(), 1); 113 | 114 | problem.remove_variable_bounds("x"); 115 | assert_eq!(problem.variable_bounds.len(), 0); 116 | } 117 | 118 | #[test] 119 | fn compute_residual_and_jacobian() { 120 | let mut problem = tiny_solver::Problem::new(); 121 | problem.add_residual_block( 122 | 1, 123 | &["x"], 124 | Box::new(tiny_solver::factors::PriorFactor { 125 | v: na::dvector![3.0], 126 | }), 127 | None, 128 | ); 129 | 130 | struct CustomFactor {} 131 | impl tiny_solver::factors::Factor for CustomFactor { 132 | fn residual_func(&self, params: &[nalgebra::DVector]) -> nalgebra::DVector { 133 | println!("residual function: {:?}", params.len()); 134 | let x = ¶ms[0][0]; 135 | let y = ¶ms[1][0]; 136 | let z = ¶ms[1][1]; 137 | 138 | na::dvector![ 139 | x.clone() 140 | + y.clone() * T::from_f64(2.0).unwrap() 141 | + z.clone() * T::from_f64(4.0).unwrap(), 142 | y.clone() * z.clone() 143 | ] 144 | } 145 | } 146 | 147 | problem.add_residual_block(2, &["x", "yz"], Box::new(CustomFactor {}), None); 148 | 149 | // the initial values for x is 0.7 and yz is [-30.2, 123.4] 150 | let initial_values = HashMap::>::from([ 151 | ("x".to_string(), na::dvector![0.7]), 152 | ("yz".to_string(), na::dvector![-30.2, 123.4]), 153 | ]); 154 | let parameter_blocks = problem.initialize_parameter_blocks(&initial_values); 155 | let variable_name_to_col_idx_dict = 156 | problem.get_variable_name_to_col_idx_dict(¶meter_blocks); 157 | let total_variable_dimension = parameter_blocks.values().map(|p| p.tangent_size()).sum(); 158 | let symbolic_structure = problem.build_symbolic_structure( 159 | ¶meter_blocks, 160 | total_variable_dimension, 161 | &variable_name_to_col_idx_dict, 162 | ); 163 | 164 | let (residuals, jac) = problem.compute_residual_and_jacobian( 165 | ¶meter_blocks, 166 | &variable_name_to_col_idx_dict, 167 | &symbolic_structure, 168 | ); 169 | 170 | assert_eq!(residuals.nrows(), 3); 171 | assert_eq!(residuals.ncols(), 1); 172 | assert_eq!(jac.nrows(), 3); 173 | assert_eq!(jac.ncols(), 3); 174 | } 175 | } 176 | -------------------------------------------------------------------------------- /tests/test_tiny_solver.rs: -------------------------------------------------------------------------------- 1 | // use num_dual::DualNum; 2 | // use tiny_solver::TinySolver; 3 | // extern crate nalgebra as na; 4 | // use std::ops::Mul; 5 | 6 | // struct TestProblem; 7 | 8 | // impl TinySolver<3, 2> for TestProblem { 9 | // fn cost_function( 10 | // params: nalgebra::SVector, 3>, 11 | // ) -> nalgebra::SVector, 2> { 12 | // let x = params[0]; 13 | // let y = params[1]; 14 | // let z = params[2]; 15 | // return nalgebra::SVector::from([x + y.mul(3.0) + z.powf(1.1), y * z]); 16 | // } 17 | // } 18 | 19 | // #[test] 20 | // fn test_residual() { 21 | // let xvec = na::SVector::from([5.0, 3.0, 2.0]); 22 | // let residual = TestProblem::cost_function(xvec.map(num_dual::DualVec::from_re)); 23 | // assert_eq!(residual[0].re, 16.143546925072584); 24 | // assert_eq!(residual[1].re, 6.0); 25 | // } 26 | -------------------------------------------------------------------------------- /tiny_solver/__init__.py: -------------------------------------------------------------------------------- 1 | from .tiny_solver import * 2 | -------------------------------------------------------------------------------- /tiny_solver/factors/__init__.pyi: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | class Factor: ... 4 | 5 | class BetweenFactorSE2(Factor): 6 | def __init__(self, x: float, y: float, theta: float) -> None: ... 7 | 8 | class PriorFactor(Factor): 9 | def __init__(self, x: np.ndarray) -> None: ... 10 | 11 | class PyFactor(Factor): 12 | def __init__(self, func: callable) -> None: ... 13 | -------------------------------------------------------------------------------- /tiny_solver/loss_functions/__init__.pyi: -------------------------------------------------------------------------------- 1 | class Loss: ... 2 | 3 | class HuberLoss: 4 | def __init__(self, scale: float = 1.0) -> None: ... 5 | -------------------------------------------------------------------------------- /tiny_solver/py.typed: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/powei-lin/tiny-solver-rs/094690931a475fcd17a6d5fe9968daf7ece00cf0/tiny_solver/py.typed -------------------------------------------------------------------------------- /tiny_solver/tiny_solver.pyi: -------------------------------------------------------------------------------- 1 | from typing import List, Tuple, Dict, Optional 2 | from enum import Enum 3 | 4 | import numpy as np 5 | 6 | from tiny_solver.factors import Factor 7 | from tiny_solver.loss_functions import Loss 8 | 9 | class Problem: 10 | def __init__(self) -> None: ... 11 | def add_residual_block( 12 | self, dim_residual: int, variable_key_size_list: List[Tuple[str, int]], factor: Factor, loss: Optional[Loss] 13 | ) -> None: ... 14 | 15 | class GaussNewtonOptimizer: 16 | def __init__(self) -> None: ... 17 | def optimize( 18 | self, problem: Problem, init_values: Dict[str, np.ndarray], optimizer_options: Optional[OptimizerOptions] = None 19 | ) -> None: ... 20 | 21 | class LinearSolver(Enum): 22 | SparseCholesky = ... 23 | SparseQR = ... 24 | 25 | class OptimizerOptions: 26 | def __init__( 27 | self, 28 | max_iteration: int = 100, 29 | linear_solver_type: LinearSolver = LinearSolver.SparseCholesky, 30 | verbosity_level: int = 0, 31 | min_abs_error_decrease_threshold: float = 1e-5, 32 | min_rel_error_decrease_threshold: float = 1e-5, 33 | min_error_threshold: float = 1e-8, 34 | ) -> None: ... 35 | --------------------------------------------------------------------------------