diff --git a/.gitignore b/.gitignore index ea8c4bf..806d6d2 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,6 @@ /target +*.pyc +__pycache__ +*.so +.DS_Store +.vscode/* \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 54f5cd9..001da62 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -114,7 +114,7 @@ dependencies = [ "js-sys", "num-traits", "wasm-bindgen", - "windows-targets", + "windows-targets 0.52.4", ] [[package]] @@ -695,6 +695,12 @@ dependencies = [ "png", ] +[[package]] +name = "indoc" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e186cfbae8084e513daff4240b4797e342f988cecda4fb6c939150f96315fd8" + [[package]] name = "jpeg-decoder" version = "0.3.1" @@ -729,7 +735,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c2a198fb6b0eada2a8df47933734e6d35d350665a33a3593d7164fa52c75c19" dependencies = [ "cfg-if", - "windows-targets", + "windows-targets 0.52.4", ] [[package]] @@ -749,6 +755,16 @@ dependencies = [ "redox_syscall", ] +[[package]] +name = "lock_api" +version = "0.4.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c168f8615b12bc01f9c17e2eb0cc07dcae1940121185446edc3744920e8ef45" +dependencies = [ + "autocfg", + "scopeguard", +] + [[package]] name = "log" version = "0.4.21" @@ -787,6 +803,15 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "523dc4f511e55ab87b694dc30d0f820d60906ef06413f93d4d7a1385599cc149" +[[package]] +name = "memoffset" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a634b1c61a95585bd15607c6ab0c4e5b226e695ff2800ba0cdccddf208c406c" +dependencies = [ + "autocfg", +] + [[package]] name = "miniz_oxide" version = "0.7.2" @@ -824,6 +849,19 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "ndarray" +version = "0.15.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb12d4e967ec485a5f71c6311fe28158e9d6f4bc4a447b474184d0f91a8fa32" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "rawpointer", +] + [[package]] name = "npyz" version = "0.8.3" @@ -898,12 +936,51 @@ dependencies = [ "libm", ] +[[package]] +name = "numpy" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef41cbb417ea83b30525259e30ccef6af39b31c240bda578889494c5392d331" +dependencies = [ + "libc", + "nalgebra", + "ndarray", + "num-complex", + "num-integer", + "num-traits", + "pyo3", + "rustc-hash", +] + [[package]] name = "once_cell" version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" +[[package]] +name = "parking_lot" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c42a9226546d68acdd9c0a280d17ce19bfe27a46bf68784e4066115788d008e" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-targets 0.48.5", +] + [[package]] name = "paste" version = "1.0.14" @@ -1039,6 +1116,12 @@ dependencies = [ "miniz_oxide", ] +[[package]] +name = "portable-atomic" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7170ef9988bc169ba16dd36a7fa041e5c4cbeb6a35b76d4c03daded371eae7c0" + [[package]] name = "proc-macro2" version = "1.0.78" @@ -1073,6 +1156,69 @@ dependencies = [ "pest_derive", ] +[[package]] +name = "pyo3" +version = "0.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53bdbb96d49157e65d45cc287af5f32ffadd5f4761438b527b055fb0d4bb8233" +dependencies = [ + "cfg-if", + "indoc", + "libc", + "memoffset", + "parking_lot", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "deaa5745de3f5231ce10517a1f5dd97d53e5a2fd77aa6b5842292085831d48d7" +dependencies = [ + "once_cell", + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62b42531d03e08d4ef1f6e85a2ed422eb678b8cd62b762e53891c05faf0d4afa" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7305c720fa01b8055ec95e484a6eca7a83c841267f0dd5280f0c8b8551d2c158" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn 2.0.52", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c7e9b68bb9c3149c5b0cade5d07f953d6d125eb4337723c4ccdb665f1f96185" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn 2.0.52", +] + [[package]] name = "quote" version = "1.0.35" @@ -1168,6 +1314,12 @@ dependencies = [ "thiserror", ] +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + [[package]] name = "rustc_version" version = "0.4.0" @@ -1195,6 +1347,12 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + [[package]] name = "semver" version = "1.0.22" @@ -1257,6 +1415,12 @@ version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d66dc143e6b11c1eddc06d5c423cfc97062865baf299914ab64caa38182078fe" +[[package]] +name = "smallvec" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6ecd384b10a64542d77071bd64bd7b231f4ed5940fba55e98c3de13824cf3d7" + [[package]] name = "syn" version = "1.0.109" @@ -1293,6 +1457,12 @@ dependencies = [ "walkdir", ] +[[package]] +name = "target-lexicon" +version = "0.12.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1fc403891a21bcfb7c37834ba66a547a8f402146eba7265b5a6d88059c9ff2f" + [[package]] name = "thiserror" version = "1.0.57" @@ -1315,14 +1485,17 @@ dependencies = [ [[package]] name = "tiny-solver" -version = "0.2.0" +version = "0.3.1" dependencies = [ "faer", "faer-ext", "nalgebra", "num-dual", "num-traits", + "numpy", "plotters", + "pyo3", + "rayon", ] [[package]] @@ -1349,6 +1522,12 @@ version = "1.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" +[[package]] +name = "unindent" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7de7d73e1754487cb58364ee906a499937a0dfabd86bcb980fa99ec8c8fa2ce" + [[package]] name = "version_check" version = "0.9.4" @@ -1488,7 +1667,22 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" dependencies = [ - "windows-targets", + "windows-targets 0.52.4", +] + +[[package]] +name = "windows-targets" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" +dependencies = [ + "windows_aarch64_gnullvm 0.48.5", + "windows_aarch64_msvc 0.48.5", + "windows_i686_gnu 0.48.5", + "windows_i686_msvc 0.48.5", + "windows_x86_64_gnu 0.48.5", + "windows_x86_64_gnullvm 0.48.5", + "windows_x86_64_msvc 0.48.5", ] [[package]] @@ -1497,51 +1691,93 @@ version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7dd37b7e5ab9018759f893a1952c9420d060016fc19a472b4bb20d1bdd694d1b" dependencies = [ - "windows_aarch64_gnullvm", - "windows_aarch64_msvc", - "windows_i686_gnu", - "windows_i686_msvc", - "windows_x86_64_gnu", - "windows_x86_64_gnullvm", - "windows_x86_64_msvc", + "windows_aarch64_gnullvm 0.52.4", + "windows_aarch64_msvc 0.52.4", + "windows_i686_gnu 0.52.4", + "windows_i686_msvc 0.52.4", + "windows_x86_64_gnu 0.52.4", + "windows_x86_64_gnullvm 0.52.4", + "windows_x86_64_msvc 0.52.4", ] +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" + [[package]] name = "windows_aarch64_gnullvm" version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bcf46cf4c365c6f2d1cc93ce535f2c8b244591df96ceee75d8e83deb70a9cac9" +[[package]] +name = "windows_aarch64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" + [[package]] name = "windows_aarch64_msvc" version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "da9f259dd3bcf6990b55bffd094c4f7235817ba4ceebde8e6d11cd0c5633b675" +[[package]] +name = "windows_i686_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" + [[package]] name = "windows_i686_gnu" version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b474d8268f99e0995f25b9f095bc7434632601028cf86590aea5c8a5cb7801d3" +[[package]] +name = "windows_i686_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" + [[package]] name = "windows_i686_msvc" version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1515e9a29e5bed743cb4415a9ecf5dfca648ce85ee42e15873c3cd8610ff8e02" +[[package]] +name = "windows_x86_64_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" + [[package]] name = "windows_x86_64_gnu" version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5eee091590e89cc02ad514ffe3ead9eb6b660aedca2183455434b93546371a03" +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" + [[package]] name = "windows_x86_64_gnullvm" version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77ca79f2451b49fa9e2af39f0747fe999fcda4f5e241b2898624dca97a1f2177" +[[package]] +name = "windows_x86_64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" + [[package]] name = "windows_x86_64_msvc" version = "0.52.4" diff --git a/Cargo.toml b/Cargo.toml index 2ed4124..8aa804d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,23 +1,39 @@ [package] name = "tiny-solver" -version = "0.2.0" +version = "0.3.1" edition = "2021" +authors = ["Powei Lin "] +readme = "README.md" +license = "MIT OR Apache-2.0" +description = "Factor graph solver" +homepage = "https://github.com/powei-lin/tiny-solver-rs" +repository = "https://github.com/powei-lin/tiny-solver-rs" +keywords = ["factor graph", "ceres-solver", "minisam"] +categories = ["data-structures", "science", "mathematics"] +exclude = ["/.github/*", "*.ipynb", "./scripts/*", "examples/*"] # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] faer = "0.18.2" -faer-ext = {version = "0.1.0", features = ["nalgebra"] } +faer-ext = { version = "0.1.0", features = ["nalgebra"] } nalgebra = "0.32.4" num-dual = "0.8.1" num-traits = "0.2.18" +numpy = { version = "0.20.0", features = ["nalgebra"] } +pyo3 = { version = "0.20.3", features = ["abi3", "abi3-py38"] } +rayon = "1.9.0" [[example]] -name = "example_solve" +name = "m3500_benchmark" crate-type = ["bin"] [dev-dependencies] plotters = "0.3.5" [profile.dev.package.faer] -opt-level = 3 \ No newline at end of file +opt-level = 3 + +[lib] +name = "tiny_solver" +# crate-type = ["cdylib"] diff --git a/README.md b/README.md index 7c0c861..a01f6da 100644 --- a/README.md +++ b/README.md @@ -1,20 +1,21 @@ # tiny-solver-rs Inspired by ceres-solver and tiny-solver. This is a re-implemented version of [tiny-solver](https://github.com/keir/tinysolver/tree/master) in rust. +### Installation +#### rust + +#### python + ### Usage -You need to implement the cost function of `TinySolver`. -Just like ceres-solver. +Under development. + +### Example +#### rust +```sh +cargo run -r --example m3500_benchmar +``` -See full example in `examples/example_solve.rs`. -```rust -impl TinySolver<3, 2> for ExampleStatic { - fn cost_function( - params: nalgebra::SVector, 3>, - ) -> nalgebra::SVector, 2> { - let x = params[0]; - let y = params[1]; - let z = params[2]; - return nalgebra::SVector::from([x + y.mul(2.0) + z.mul(4.0), y * z]); - } -} +#### python ``` +python3 examples/python/m3500.py +``` \ No newline at end of file diff --git a/examples/example_solve.rs b/examples/example_solve.rs deleted file mode 100644 index 620975c..0000000 --- a/examples/example_solve.rs +++ /dev/null @@ -1,28 +0,0 @@ -use std::ops::Mul; - -use tiny_solver::TinySolver; - -pub type Matrix3x1d = nalgebra::SMatrix; -pub struct ExampleStatic {} - -impl TinySolver<3, 2> for ExampleStatic { - fn cost_function( - params: nalgebra::SVector, 3>, - ) -> nalgebra::SVector, 2> { - let x = params[0]; - let y = params[1]; - let z = params[2]; - return nalgebra::SVector::from([x + y.mul(2.0) + z.mul(4.0), y * z]); - } -} -fn main() { - // Example usage - let mut x0 = Matrix3x1d::new(0.76026643, -30.01799744, 0.55192142); - let result = ExampleStatic::solve_inplace(&mut x0); - println!("auto grad: {}", x0); - println!( - "residaul: {}", - ExampleStatic::cost_function(x0.map(num_dual::DualSVec64::from_re)) - ); - println!("{:#?}", result); -} diff --git a/examples/m3500_benchmark.rs b/examples/m3500_benchmark.rs index 61f7dbf..895e800 100644 --- a/examples/m3500_benchmark.rs +++ b/examples/m3500_benchmark.rs @@ -1,58 +1,12 @@ -use std::time::Instant; -extern crate nalgebra as na; -use plotters::prelude::*; - use std::collections::HashMap; use std::fs::read_to_string; -use tiny_solver::{gauss_newton_optimizer, optimizer::Optimizer, problem, residual_block}; +use std::time::Instant; -#[derive(Default)] -pub struct CostFactorSE2 { - dx: f64, - dy: f64, - dtheta: f64, -} -impl residual_block::Factor for CostFactorSE2 { - fn residual_func( - &self, - params: &Vec>, - ) -> na::DVector { - let t_origin_k0 = ¶ms[0]; - let t_origin_k1 = ¶ms[1]; - let se2_origin_k0 = na::Isometry2::new( - na::Vector2::new(t_origin_k0[1].clone(), t_origin_k0[2].clone()), - t_origin_k0[0].clone(), - ); - let se2_origin_k1 = na::Isometry2::new( - na::Vector2::new(t_origin_k1[1].clone(), t_origin_k1[2].clone()), - t_origin_k1[0].clone(), - ); - let se2_k0_k1 = na::Isometry2::new( - na::Vector2::::new( - num_dual::DualDVec64::from_re(self.dx), - num_dual::DualDVec64::from_re(self.dy), - ), - num_dual::DualDVec64::from_re(self.dtheta), - ); +use nalgebra as na; +use plotters::prelude::*; - let se2_diff = se2_origin_k1.inverse() * se2_origin_k0 * se2_k0_k1; - return na::dvector![ - se2_diff.translation.x.clone(), - se2_diff.translation.y.clone(), - se2_diff.rotation.angle() - ]; - } -} -#[derive(Default)] -pub struct BetweenFactor {} -impl residual_block::Factor for BetweenFactor { - fn residual_func( - &self, - params: &Vec>, - ) -> na::DVector { - return params[0].clone(); - } -} +use tiny_solver::{factors, gauss_newton_optimizer, optimizer::Optimizer, problem}; +// use tiny_solver::{gauss_newton_optimizer, optimizer::Optimizer, problem, residual_block, factors}; fn read_g2o(filename: &str) -> (problem::Problem, HashMap>) { let mut problem = problem::Problem::new(); @@ -73,7 +27,7 @@ fn read_g2o(filename: &str) -> (problem::Problem, HashMap().unwrap(); let dtheta = line[5].parse::().unwrap(); // todo add info matrix - let edge = CostFactorSE2 { + let edge = factors::CostFactorSE2 { dx: dx, dy: dy, dtheta: dtheta, @@ -86,7 +40,9 @@ fn read_g2o(filename: &str) -> (problem::Problem, HashMap (problem::Problem, HashMap = init_values.iter().map(|(_, v)| (v[1], v[2])).collect(); - let root_drawing_area = BitMapBackend::new("m3500_resul.png", (1024, 1024)).into_drawing_area(); + let root_drawing_area = BitMapBackend::new("m3500_rs.png", (1024, 1024)).into_drawing_area(); root_drawing_area.fill(&WHITE).unwrap(); @@ -118,7 +74,7 @@ fn main() { .unwrap(); let gn = gauss_newton_optimizer::GaussNewtonOptimizer {}; let start = Instant::now(); - let result = gn.optimize(problem, &init_values); + let result = gn.optimize(&problem, &init_values); let duration = start.elapsed(); println!("Time elapsed in total is: {:?}", duration); let result_points: Vec<(f64, f64)> = result.iter().map(|(_, v)| (v[1], v[2])).collect(); diff --git a/examples/python/m3500.py b/examples/python/m3500.py new file mode 100644 index 0000000..62cf42a --- /dev/null +++ b/examples/python/m3500.py @@ -0,0 +1,88 @@ +from time import perf_counter + +import matplotlib.pyplot as plt +import numpy as np + +from tiny_solver import GaussNewtonOptimizer, Problem +from tiny_solver.factors import PriorFactor, CostFactorSE2 + + +def load_g2o(file_path: str): + init_values = {} + factor_graph = Problem() + vertex_num = 4000 + with open(file_path) as ifile: + for line in ifile.readlines(): + items = line[:-1].split(" ") + if items[0] == "EDGE_SE2": + if int(items[1]) > vertex_num or int(items[2]) > vertex_num: + continue + point_id0 = f"x{int(items[1])}" + point_id1 = f"x{int(items[2])}" + items_float = [float(i) for i in items[3:]] + dx = items_float[0] + dy = items_float[1] + dtheta = items_float[2] + dpose = np.array([dtheta, dx, dy]) + + # if point_id0 == "x10" and point_id1 == "x11": + # print(dpose) + i11, i12, i13, i22, i23, i33 = items_float[3:] + matrix_i = np.array([[i11, i12, i13], [i12, i22, i23], [i13, i23, i33]]) + # loss = np.linalg.cholesky(matrix_i) + factor = CostFactorSE2(dx, dy, dtheta) + factor_graph.add_residual_block(3, [(point_id0, 3), (point_id1, 3)], factor) + elif items[0] == "VERTEX_SE2": + if int(items[1]) > vertex_num: + continue + point_id = f"x{int(items[1])}" + x = float(items[2]) + y = float(items[3]) + theta = float(items[4]) + # if point_id == "x10" or point_id == "x11": + # print(point_id, theta, x, y) + + init_values[point_id] = np.array([theta, x, y], dtype=np.float64) + else: + print(items) + break + return factor_graph, init_values + # show_pose(init_values=init_values) + # print(init_values) + + +def show_pose(init_values, color): + data_x = [x[1] for x in init_values.values()] + data_y = [x[2] for x in init_values.values()] + plt.scatter(data_x, data_y, s=1, c=color) + + +def main(): + file_path = "tests/data/input_M3500_g2o.g2o" + factor_graph, init_values = load_g2o(file_path) + + prior_factor = PriorFactor(np.zeros(3)) + factor_graph.add_residual_block(3, [("x0", 3)], prior_factor) + solver = GaussNewtonOptimizer() + # gn = LevenbergMarquardtOptimizer() + draw = True + if draw: + plt.figure(figsize=(8, 8)) + show_pose(init_values, "red") + + start_time = perf_counter() + init_values = solver.optimize(factor_graph, init_values) + end_time = perf_counter() + print(f"{solver.__class__.__name__} takes {end_time-start_time:.3f} sec") + if draw: + show_pose(init_values, "blue") + ax = plt.gca() + ax.set_xlim((-50, 50)) + ax.set_ylim((-80, 20)) + plt.tight_layout() + plt.savefig("m3500_py.png") + print("end") + + +if __name__ == "__main__": + main() diff --git a/examples/python/try_import.py b/examples/python/try_import.py new file mode 100644 index 0000000..6a29ba4 --- /dev/null +++ b/examples/python/try_import.py @@ -0,0 +1,35 @@ +import tiny_solver +from tiny_solver import GaussNewtonOptimizer, Problem +from tiny_solver.factors import PriorFactor, CostFactorSE2 +import numpy as np + +def main(): + print(f"{tiny_solver.__version__=}") + print(dir(tiny_solver)) + # print(tiny_solver.sum_as_string(1, 2)) + # tiny_solver.mult(np.zeros((1, 2))) + # a = tiny_solver.Dual64() + # print(a.first_derivative) + b = CostFactorSE2(1.0, 2.0, 3.0) + # print("factor module\n", dir(factors)) + # b = Costf(1.0, 2.0, 3.0) + + print(type(b)) + print(dir(b)) + print(b) + problem = Problem() + print(dir(problem)) + problem.num = 200 + print(problem.num) + d = PriorFactor(np.array([1.0, 2.0,3.0])) + problem.add_residual_block(1, [("aa", 1)], d) + problem.add_residual_block(1, [("aaa", 1)], b) + # c.add_residual_block(1, [("aaa", 1)]) + # c.add_residual_block(1, [("aa", 1)]) + # d = tiny_solver.BetweenFactor() + # d.ttt() + # tiny_solver.te(d) + optimizer = GaussNewtonOptimizer() + optimizer.optimize(problem, {"aa": np.array([123, 2, 3, 4], dtype=np.float64)}) +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/examples/test_residual_block.rs b/examples/test_residual_block.rs deleted file mode 100644 index e6f3199..0000000 --- a/examples/test_residual_block.rs +++ /dev/null @@ -1,75 +0,0 @@ -use std::vec; -use std::{ - collections::HashMap, - ops::{Mul, Sub}, -}; - -extern crate nalgebra as na; -use na::{Const, Dyn}; -use num_dual::{DualDVec64, DualVec}; -use tiny_solver::gauss_newton_optimizer::GaussNewtonOptimizer; -use tiny_solver::{gauss_newton_optimizer, optimizer::Optimizer, problem, residual_block}; - -#[derive(Default)] -pub struct CostFactor0 {} -#[derive(Default)] -pub struct CostFactor1 {} -impl residual_block::Factor for CostFactor0 { - fn residual_func( - &self, - params: &Vec>, - ) -> na::DVector { - let x = ¶ms[0][0]; - let y = ¶ms[1][0]; - let z = ¶ms[1][1]; - return na::dvector![x + y.clone().mul(2.0) + z.clone().mul(4.0), y * z]; - } -} - -impl residual_block::Factor for CostFactor1 { - fn residual_func( - &self, - params: &Vec>, - ) -> na::DVector { - let x = ¶ms[0][0]; - return na::dvector![x.clone().sub(5.0)]; - } -} - -fn rows(aa: &Vec) -> Vec> { - let mut result = Vec::new(); - let mut current = 0; - for &num in aa { - let next = current + num; - let range = (current..next).collect::>(); - result.push(range); - current = next; - } - result -} - -fn main() { - println!("rs block"); - // let (r, j) = rsb.jacobian(&vec![na::dvector![1.0], na::dvector![-2.0, 3.0]]); - // println!("{},{}", r, j); - let mut problem = problem::Problem::new(); - problem.add_residual_block( - 2, - vec![("x".to_string(), 1), ("yz".to_string(), 2)], - Box::new(CostFactor0::default()), - ); - problem.add_residual_block( - 1, - vec![("x".to_string(), 1)], - Box::new(CostFactor1::default()), - ); - let initial_values = HashMap::from([ - ("x".to_string(), na::dvector![0.76026643]), - ("yz".to_string(), na::dvector![-30.01799744, 0.55192142]), - ]); - let gn = GaussNewtonOptimizer {}; - let result = gn.optimize(problem, &initial_values); - for k in result { - println!("{} {}", k.0, k.1); - } -} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..39c90bd --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,18 @@ +[build-system] +requires = ["maturin>=1.4,<2.0"] +build-backend = "maturin" + +[project] +name = "tiny_solver" +requires-python = ">=3.8" +classifiers = [ + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", + "Programming Language :: Python :: Implementation :: PyPy", +] +dynamic = ["version"] +dependencies = ["numpy"] + + +[tool.maturin] +features = ["pyo3/extension-module"] diff --git a/scripts/compile_aarch64.sh b/scripts/compile_aarch64.sh new file mode 100755 index 0000000..7674536 --- /dev/null +++ b/scripts/compile_aarch64.sh @@ -0,0 +1 @@ +maturin build --release --target aarch64-unknown-linux-gnu --zig \ No newline at end of file diff --git a/scripts/compile_local.sh b/scripts/compile_local.sh new file mode 100755 index 0000000..b8057e8 --- /dev/null +++ b/scripts/compile_local.sh @@ -0,0 +1 @@ +maturin build --release \ No newline at end of file diff --git a/src/factors.rs b/src/factors.rs new file mode 100644 index 0000000..c411975 --- /dev/null +++ b/src/factors.rs @@ -0,0 +1,62 @@ +use nalgebra as na; +use pyo3::prelude::*; + +pub trait Factor: Send + Sync { + fn residual_func( + &self, + params: &Vec>, + ) -> na::DVector; +} + +#[pyclass] +#[derive(Debug, Clone)] +pub struct CostFactorSE2 { + pub dx: f64, + pub dy: f64, + pub dtheta: f64, +} +impl Factor for CostFactorSE2 { + fn residual_func( + &self, + params: &Vec>, + ) -> na::DVector { + let t_origin_k0 = ¶ms[0]; + let t_origin_k1 = ¶ms[1]; + let se2_origin_k0 = na::Isometry2::new( + na::Vector2::new(t_origin_k0[1].clone(), t_origin_k0[2].clone()), + t_origin_k0[0].clone(), + ); + let se2_origin_k1 = na::Isometry2::new( + na::Vector2::new(t_origin_k1[1].clone(), t_origin_k1[2].clone()), + t_origin_k1[0].clone(), + ); + let se2_k0_k1 = na::Isometry2::new( + na::Vector2::::new( + num_dual::DualDVec64::from_re(self.dx), + num_dual::DualDVec64::from_re(self.dy), + ), + num_dual::DualDVec64::from_re(self.dtheta), + ); + + let se2_diff = se2_origin_k1.inverse() * se2_origin_k0 * se2_k0_k1; + return na::dvector![ + se2_diff.translation.x.clone(), + se2_diff.translation.y.clone(), + se2_diff.rotation.angle() + ]; + } +} + +#[pyclass] +#[derive(Debug, Clone)] +pub struct PriorFactor { + pub v: na::DVector, +} +impl Factor for PriorFactor { + fn residual_func( + &self, + params: &Vec>, + ) -> na::DVector { + return params[0].clone() - self.v.map(num_dual::DualDVec64::from_re); + } +} diff --git a/src/gauss_newton_optimizer.rs b/src/gauss_newton_optimizer.rs index 91755ad..b8b4a02 100644 --- a/src/gauss_newton_optimizer.rs +++ b/src/gauss_newton_optimizer.rs @@ -1,18 +1,23 @@ use std::time::Instant; use faer_ext::IntoNalgebra; +use pyo3::prelude::*; use crate::{linear::sparse_cholesky, optimizer}; + +#[pyclass] +#[derive(Debug, Clone)] pub struct GaussNewtonOptimizer {} + impl optimizer::Optimizer for GaussNewtonOptimizer { fn optimize( &self, - problem: crate::problem::Problem, + problem: &crate::problem::Problem, initial_values: &std::collections::HashMap>, ) -> std::collections::HashMap> { let mut params = initial_values.clone(); - for i in 0..10 { + for i in 0..50 { println!("{}", i); let (residuals, jac) = problem.compute_residual_and_jacobian(¶ms); @@ -21,7 +26,7 @@ impl optimizer::Optimizer for GaussNewtonOptimizer { let duration = start.elapsed(); println!("Time elapsed in solve() is: {:?}", duration); - if dx.norm_l1() < 1e-16 { + if dx.norm_l1() < 1e-8 { println!("grad too low"); break; } diff --git a/src/lib.rs b/src/lib.rs index 774b36f..2202954 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,15 @@ +pub mod factors; pub mod gauss_newton_optimizer; pub mod linear; pub mod optimizer; pub mod problem; +pub mod python; pub mod residual_block; -pub mod tiny_solver; -pub use crate::tiny_solver::*; +pub mod tiny_solver_old; + +pub use gauss_newton_optimizer::*; +pub use linear::*; +pub use optimizer::*; +pub use problem::*; +pub use python::*; +pub use residual_block::*; diff --git a/src/optimizer.rs b/src/optimizer.rs index 6c3dc9c..4b51ef3 100644 --- a/src/optimizer.rs +++ b/src/optimizer.rs @@ -8,7 +8,7 @@ use crate::problem; pub trait Optimizer { fn optimize( &self, - problem: problem::Problem, + problem: &problem::Problem, initial_values: &HashMap>, ) -> HashMap>; fn apply_dx( diff --git a/src/problem.rs b/src/problem.rs index efe1832..1e12bed 100644 --- a/src/problem.rs +++ b/src/problem.rs @@ -1,16 +1,21 @@ use std::collections::HashMap; +use std::sync::{Arc, Mutex}; use faer::sparse::SparseColMat; use faer_ext::IntoFaer; use nalgebra as na; +use pyo3::prelude::*; +use rayon::prelude::*; -use crate::residual_block::{self, Factor}; +use crate::{factors, residual_block}; + +#[pyclass] pub struct Problem { pub total_variable_dimension: usize, pub total_residual_dimension: usize, residual_blocks: Vec, pub variable_name_to_col_idx_dict: HashMap, - // col_idx_to_variable_dict: HashMap, + pub thread_num: usize, } impl Problem { pub fn new() -> Problem { @@ -19,22 +24,23 @@ impl Problem { total_residual_dimension: 0, residual_blocks: Vec::::new(), variable_name_to_col_idx_dict: HashMap::::new(), + thread_num: 1, } } pub fn add_residual_block( &mut self, dim_residual: usize, variable_key_size_list: Vec<(String, usize)>, - factor: Box, + factor: Box, ) { self.residual_blocks.push(residual_block::ResidualBlock { - dim_residual: dim_residual, + dim_residual, residual_row_start_idx: self.total_residual_dimension, variable_key_list: variable_key_size_list .iter() .map(|(x, _)| x.to_string()) .collect(), - factor: factor, + factor, }); for (key, variable_dimesion) in variable_key_size_list { if !self.variable_name_to_col_idx_dict.contains_key(&key) { @@ -63,10 +69,13 @@ impl Problem { &self, variable_key_value_map: &HashMap>, ) -> (faer::Mat, SparseColMat) { - let mut total_residual = na::DVector::::zeros(self.total_residual_dimension); - let mut jacobian_list = Vec::<(usize, usize, f64)>::new(); + // multi + let total_residual = Arc::new(Mutex::new(na::DVector::::zeros( + self.total_residual_dimension, + ))); + let jacobian_list = Arc::new(Mutex::new(Vec::<(usize, usize, f64)>::new())); - for residual_block in &self.residual_blocks { + self.residual_blocks.par_iter().for_each(|residual_block| { let mut params = Vec::>::new(); let mut variable_local_idx_size_list = Vec::<(usize, usize)>::new(); let mut count_variable_local_idx: usize = 0; @@ -78,27 +87,46 @@ impl Problem { }; } let (res, jac) = residual_block.jacobian(¶ms); - total_residual - .rows_mut( - residual_block.residual_row_start_idx, - residual_block.dim_residual, - ) - .copy_from(&res); + + { + let mut total_residual = total_residual.lock().unwrap(); + total_residual + .rows_mut( + residual_block.residual_row_start_idx, + residual_block.dim_residual, + ) + .copy_from(&res); + } + for (i, vk) in residual_block.variable_key_list.iter().enumerate() { if let Some(variable_global_idx) = self.variable_name_to_col_idx_dict.get(vk) { let (variable_local_idx, var_size) = variable_local_idx_size_list[i]; let variable_jac = jac.view((0, variable_local_idx), (jac.shape().0, var_size)); + let mut local_jacobian_list = Vec::new(); for row_idx in 0..jac.shape().0 { for col_idx in 0..var_size { - let globle_row_idx = residual_block.residual_row_start_idx + row_idx; - let globle_col_idx = variable_global_idx + col_idx; + let global_row_idx = residual_block.residual_row_start_idx + row_idx; + let global_col_idx = variable_global_idx + col_idx; let value = variable_jac[(row_idx, col_idx)]; - jacobian_list.push((globle_row_idx, globle_col_idx, value)); + local_jacobian_list.push((global_row_idx, global_col_idx, value)); } } + let mut jacobian_list = jacobian_list.lock().unwrap(); + jacobian_list.extend(local_jacobian_list); } } - } + }); + + let total_residual = Arc::try_unwrap(total_residual) + .unwrap() + .into_inner() + .unwrap(); + let jacobian_list = Arc::try_unwrap(jacobian_list) + .unwrap() + .into_inner() + .unwrap(); + // end + let residual_faer = total_residual.view_range(.., ..).into_faer().to_owned(); let jacobian_faer = SparseColMat::try_new_from_triplets( self.total_residual_dimension, diff --git a/src/python/mod.rs b/src/python/mod.rs new file mode 100644 index 0000000..ce60f06 --- /dev/null +++ b/src/python/mod.rs @@ -0,0 +1,30 @@ +use pyo3::prelude::*; + +use crate::factors::*; +use crate::*; + +mod py_factors; +mod py_optimizer; +mod py_problem; + +fn register_child_module(py: Python<'_>, parent_module: &PyModule) -> PyResult<()> { + let child_module = PyModule::new(py, "factors")?; + child_module.add_class::()?; + child_module.add_class::()?; + parent_module.add_submodule(child_module)?; + py.import("sys")? + .getattr("modules")? + .set_item("tiny_solver.factors", child_module)?; + Ok(()) +} + +/// A Python module implemented in Rust. +#[pymodule] +pub fn tiny_solver<'py>(_py: Python<'py>, m: &'py PyModule) -> PyResult<()> { + m.add("__version__", env!("CARGO_PKG_VERSION"))?; + m.add_class::()?; + m.add_class::()?; + register_child_module(_py, m)?; + + Ok(()) +} diff --git a/src/python/py_factors.rs b/src/python/py_factors.rs new file mode 100644 index 0000000..cddfcdc --- /dev/null +++ b/src/python/py_factors.rs @@ -0,0 +1,26 @@ +use numpy::PyReadonlyArray1; +use pyo3::prelude::*; + +use crate::factors::*; + +#[pymethods] +impl CostFactorSE2 { + #[new] + pub fn new(x: f64, y: f64, theta: f64) -> Self { + CostFactorSE2 { + dx: x, + dy: y, + dtheta: theta, + } + } +} + +#[pymethods] +impl PriorFactor { + #[new] + pub fn new(x: PyReadonlyArray1) -> Self { + PriorFactor { + v: x.as_matrix().column(0).into(), + } + } +} diff --git a/src/python/py_optimizer.rs b/src/python/py_optimizer.rs new file mode 100644 index 0000000..994098c --- /dev/null +++ b/src/python/py_optimizer.rs @@ -0,0 +1,40 @@ +use std::collections::HashMap; + +use numpy::{PyArray2, PyReadonlyArray1, ToPyArray}; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +use crate::optimizer::Optimizer; +use crate::problem::Problem; +use crate::GaussNewtonOptimizer; + +#[pymethods] +impl GaussNewtonOptimizer { + #[new] + pub fn new() -> Self { + println!("init GaussNewtonOptimizer"); + GaussNewtonOptimizer {} + } + + #[pyo3(name = "optimize")] + pub fn optimize_py( + &self, + py: Python<'_>, + problem: &Problem, + initial_values: &PyDict, + ) -> PyResult>>> { + let init_values: HashMap> = initial_values.extract().unwrap(); + let init_values: HashMap> = init_values + .iter() + .map(|(k, v)| (k.to_string(), v.as_matrix().column(0).into())) + .collect(); + // println!("{}", initial_values); + let result = self.optimize(problem, &init_values); + + let output_d: HashMap>> = result + .iter() + .map(|(k, v)| (k.to_string(), v.to_pyarray(py).to_owned().into())) + .collect(); + Ok(output_d) + } +} diff --git a/src/python/py_problem.rs b/src/python/py_problem.rs new file mode 100644 index 0000000..26450bd --- /dev/null +++ b/src/python/py_problem.rs @@ -0,0 +1,60 @@ +use pyo3::prelude::*; + +use crate::factors::*; +use crate::problem::Problem; + +fn convert_pyany_to_factor(py_any: &PyAny) -> PyResult> { + let factor_name: String = py_any.get_type().getattr("__name__")?.extract()?; + match factor_name.as_str() { + "CostFactorSE2" => { + let factor: CostFactorSE2 = py_any.extract().unwrap(); + Ok(Box::new(factor)) + } + "PriorFactor" => { + let factor: PriorFactor = py_any.extract().unwrap(); + Ok(Box::new(factor)) + } + _ => Err(PyErr::new::( + "Unknown factor type", + )), + } +} + +#[pymethods] +impl Problem { + #[new] + pub fn new_py() -> Self { + Problem::new() + } + + #[pyo3(name = "add_residual_block")] + pub fn add_residual_block_py( + &mut self, + dim_residual: usize, + variable_key_size_list: Vec<(String, usize)>, + pyfactor: &PyAny, + ) -> PyResult<()> { + self.add_residual_block( + dim_residual, + variable_key_size_list, + convert_pyany_to_factor(pyfactor).unwrap(), + ); + + // println!( + // "total residual {}, total keys {}", + // self.total_residual_dimension, + // self.variable_name_to_col_idx_dict.len() + // ); + Ok(()) + } + + #[getter] + pub fn get_num(&self) -> PyResult { + Ok(self.thread_num) + } + #[setter] + pub fn set_num(&mut self, value: usize) -> PyResult<()> { + self.thread_num = value; + Ok(()) + } +} diff --git a/src/residual_block.rs b/src/residual_block.rs index 8a29c8d..c45e178 100644 --- a/src/residual_block.rs +++ b/src/residual_block.rs @@ -1,16 +1,12 @@ use nalgebra as na; -pub trait Factor { - fn residual_func( - &self, - params: &Vec>, - ) -> na::DVector; -} +use crate::factors::Factor; + pub struct ResidualBlock { pub dim_residual: usize, pub residual_row_start_idx: usize, pub variable_key_list: Vec, - pub factor: Box, + pub factor: Box, } impl ResidualBlock { pub fn jacobian(&self, params: &Vec>) -> (na::DVector, na::DMatrix) { diff --git a/src/tiny_solver.rs b/src/tiny_solver_old.rs similarity index 98% rename from src/tiny_solver.rs rename to src/tiny_solver_old.rs index a03ba42..02be7cd 100644 --- a/src/tiny_solver.rs +++ b/src/tiny_solver_old.rs @@ -1,8 +1,9 @@ +use std::ops::Mul; + use nalgebra as na; use num_dual; -use std::ops::Mul; -struct SolverParameters { +pub struct SolverParameters { gradient_threshold: f64, relative_step_threshold: f64, error_threshold: f64, @@ -10,7 +11,7 @@ struct SolverParameters { max_iterations: usize, } impl SolverParameters { - fn defualt() -> SolverParameters { + pub fn defualt() -> SolverParameters { SolverParameters { gradient_threshold: 1e-16, relative_step_threshold: 1e-16, diff --git a/tiny_solver/__init__.py b/tiny_solver/__init__.py new file mode 100644 index 0000000..6ffc03b --- /dev/null +++ b/tiny_solver/__init__.py @@ -0,0 +1 @@ +from .tiny_solver import * \ No newline at end of file diff --git a/tiny_solver/factors/__init__.pyi b/tiny_solver/factors/__init__.pyi new file mode 100644 index 0000000..6519317 --- /dev/null +++ b/tiny_solver/factors/__init__.pyi @@ -0,0 +1,7 @@ +import numpy as np + +class CostFactorSE2: + def __init__(self, x: float, y: float, theta: float) -> None: ... + +class PriorFactor: + def __init__(self, x: np.ndarray, /) -> None: ... \ No newline at end of file diff --git a/tiny_solver/py.typed b/tiny_solver/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/tiny_solver/tiny_solver.pyi b/tiny_solver/tiny_solver.pyi new file mode 100644 index 0000000..a300015 --- /dev/null +++ b/tiny_solver/tiny_solver.pyi @@ -0,0 +1,11 @@ +from typing import List, Tuple, Dict +import numpy as np + + +class Problem: + def __init__(self) -> None: ... + def add_residual_block(self, dim_residual: int, variable_key_size_list: List[Tuple[str, int]], factor) -> None: ... + +class GaussNewtonOptimizer: + def __init__(self) -> None: ... + def optimize(self, problem: Problem, init_values: Dict[str, np.ndarray]) -> None: ... \ No newline at end of file