Repository: rpiRobotics/ik-geo (language: python)

Issue:
Rust Subproblem 2 Extended



Discussion / hints:

Fix patch:
diff --git a/rust/src/subproblems.rs b/rust/src/subproblems.rs
index 8764fe1..1afb789 100644
--- a/rust/src/subproblems.rs
+++ b/rust/src/subproblems.rs
@@ -1,9 +1,47 @@
-use nalgebra::{ Vector3, Matrix3x2, Matrix4x3, Matrix4x1, DVector };
+use nalgebra::{ Vector3, Matrix3x2, Matrix4x3, Matrix4x1, Matrix3x4, Matrix3 };
 
-/// Solves for theta where rot(k, theta) * p1 = p2 if possible.
-/// If not, minimizes || rot(k, theta) * p1 - p2 ||
-/// Also returns a boolean of whether or not theta is a least-squares solution
-pub fn subproblem1(p1: Vector3<f64>, p2: Vector3<f64>, k: Vector3<f64>) -> (f64, bool) {
+#[derive(Clone)]
+pub enum SolutionSet2<T> {
+    One(T),
+    Two(T, T),
+}
+
+impl<T> SolutionSet2<T> where T: Copy {
+    pub fn expect_one(&self) -> T {
+        match self {
+            Self::One(s) => *s,
+            Self::Two(..) => panic!("Found two solutions where one was expected"),
+        }
+    }
+
+    pub fn expect_two(&self) -> [T; 2] {
+        match self {
+            Self::One(_) => panic!("Found one solution where two were expected"),
+            Self::Two(s1, s2) => [*s1, *s2],
+        }
+    }
+
+    pub fn get_first(&self) -> T {
+        match self {
+            Self::One(s) => *s,
+            Self::Two(s, _) => *s,
+        }
+    }
+
+    pub fn get_all(&self) -> Vec<T> {
+        match self {
+            Self::One(s) => vec![*s],
+            Self::Two(s1, s2) => vec![*s1, *s2],
+        }
+    }
+}
+
+/**
+Solves for `theta` where `rot(k, theta) * p1 = p2` if possible.
+If not, minimizes `|| rot(k, theta) * p1 - p2 ||`.
+Also returns a boolean of whether or not `theta` is a least-squares solution.
+*/
+pub fn subproblem1(p1: &Vector3<f64>, p2: &Vector3<f64>, k: &Vector3<f64>) -> (f64, bool) {
     let kxp = k.cross(&p1);
     let a = Matrix3x2::from_columns(&[kxp, -k.cross(&kxp)]);
     let x = a.transpose() * p2;
@@ -14,11 +52,13 @@ pub fn subproblem1(p1: Vector3<f64>, p2: Vector3<f64>, k: Vector3<f64>) -> (f64,
     (theta, is_ls)
 }
 
-/// Solves for theta1 and theta2 where rot(k1, theta1)*p1 = rot(k2, theta2)*p2 if possible
-/// If not, minimizes || rot(k1, theta1)*p1 - rot(k2, theta2)*p2 ||
-/// Also returns a boolean of whether or not { theta1, theta2 } is a least-squares solution
-/// There may be 1 or 2 solutions for theta1 and theta2 so they are returned as 1 or 2 sized dynamic vectors
-pub fn subproblem2(p1: Vector3<f64>, p2: Vector3<f64>, k1: Vector3<f64>, k2: Vector3<f64>) -> (DVector<f64>, DVector<f64>, bool) {
+/**
+Solves for `theta1` and `theta2` where `rot(k1, theta1) * p1 = rot(k2, theta2) * p2` if possible.
+If not, minimizes `|| rot(k1, theta1) * p1 - rot(k2, theta2) * p2 ||`.
+Also returns a boolean of whether or not `{ theta1, theta2 }` is a least-squares solution.
+There may be 1 or 2 solutions for `theta1` and `theta2`.
+*/
+pub fn subproblem2(p1: &Vector3<f64>, p2: &Vector3<f64>, k1: &Vector3<f64>, k2: &Vector3<f64>) -> (SolutionSet2<f64>, SolutionSet2<f64>, bool) {
     let is_ls = (p1.norm() - p2.norm()).abs() >= 1e-8;
 
     let p1 = p1.normalize();
@@ -55,15 +95,15 @@ pub fn subproblem2(p1: Vector3<f64>, p2: Vector3<f64>, k1: Vector3<f64>, k2: Vec
         let sc_1 = x_ls + xi * a_perp_tilde;
         let sc_2 = x_ls - xi * a_perp_tilde;
 
-        let theta1 = DVector::from_vec(vec![
+        let theta1 = SolutionSet2::Two(
             sc_1[0].atan2(sc_1[1]),
             sc_2[0].atan2(sc_2[1]),
-        ]);
+        );
 
-        let theta2 = DVector::from_vec(vec![
+        let theta2 = SolutionSet2::Two(
             sc_1[2].atan2(sc_1[3]),
             sc_2[2].atan2(sc_2[3]),
-        ]);
+        );
 
         (theta1, theta2, is_ls)
     }
@@ -71,6 +111,46 @@ pub fn subproblem2(p1: Vector3<f64>, p2: Vector3<f64>, k1: Vector3<f64>, k2: Vec
         let theta1 = x_ls[0].atan2(x_ls[1]);
         let theta2 = x_ls[2].atan2(x_ls[3]);
 
-        (DVector::from_vec(vec![theta1]), DVector::from_vec(vec![theta2]), true)
+        (SolutionSet2::One(theta1), SolutionSet2::One(theta2), true)
     }
 }
+
+/**
+Solves for `theta1` and `theta2` where `p0 + rot(k1, theta1) * p1 = rot(k2, theta2) * p2`.
+Assumes only one solution. If there could be two, `subproblem2` should be used.
+*/
+pub fn subproblem2extended(p0: &Vector3<f64>, p1: &Vector3<f64>, p2: &Vector3<f64>, k1: &Vector3<f64>, k2: &Vector3<f64>) -> (f64, f64) {
+    let kxp1 = k1.cross(&p1);
+    let kxp2 = k2.cross(&p2);
+
+    let a1  = Matrix3x2::from_columns(&[kxp1, -k1.cross(&kxp1)]);
+    let a2  = Matrix3x2::from_columns(&[kxp2, -k2.cross(&kxp2)]);
+    let a2_neg = -a2;
+
+    let a = Matrix3x4::from_columns(&[a1.column(0), a1.column(1), a2_neg.column(0), a2_neg.column(1)]);
+
+    let p = -k1 * k1.dot(&p1) + k2 * k2.dot(p2) - p0;
+
+    let radius1_sq = kxp1.dot(&kxp1);
+    let radius2_sq = kxp2.dot(&kxp2);
+
+    let alpha = radius1_sq / (radius1_sq + radius2_sq);
+    let beta = radius2_sq / (radius1_sq + radius2_sq);
+    let m_inv = Matrix3::identity() + k1 * k1.transpose() * (alpha / (1.0 - alpha));
+    let aat_inv = 1.0 / (radius1_sq + radius2_sq) * (m_inv + m_inv * k2 * k2.transpose() * m_inv * beta / (1.0 - (k2.transpose() * m_inv * k2 * beta)[0]));
+    let x_ls = a.transpose() * aat_inv * p;
+
+    let n_sym = k1.cross(&k2);
+    let pinv_a1 = a1.transpose() / radius1_sq;
+    let pinv_a2 = a2.transpose() / radius2_sq;
+    let a_perp_tilde = Matrix4x3::from_rows(&[pinv_a1.row(0), pinv_a1.row(1), pinv_a2.row(0), pinv_a2.row(1)]) * n_sym;
+
+    let num = (x_ls.fixed_rows::<2>(2).norm_squared() - 1.0) * a_perp_tilde.fixed_rows::<2>(0).norm_squared() - (x_ls.fixed_rows::<2>(0).norm_squared() - 1.0) * a_perp_tilde.fixed_rows::<2>(2).norm_squared();
+    let den = 2.0 * (x_ls.fixed_rows::<2>(0).transpose() * a_perp_tilde.fixed_rows::<2>(0) * a_perp_tilde.fixed_rows::<2>(2).norm_squared() - x_ls.fixed_rows::<2>(2).transpose() * a_perp_tilde.fixed_rows::<2>(2) * a_perp_tilde.fixed_rows::<2>(0).norm_squared())[0];
+
+    let xi = num / den;
+
+    let sc = x_ls + xi * a_perp_tilde;
+
+    (sc[0].atan2(sc[1]), sc[2].atan2(sc[3]))
+}


Test patch:
diff --git a/rust/src/tests/correctness.rs b/rust/src/tests/correctness.rs
index 629289a..cdb5cc0 100644
--- a/rust/src/tests/correctness.rs
+++ b/rust/src/tests/correctness.rs
@@ -1,14 +1,15 @@
-use super::setups::{ Setup, Subproblem1Setup, Subproblem2Setup };
+use super::setups::{ Setup, Subproblem1Setup, Subproblem2Setup, Subproblem2ExtendedSetup };
 
 const TEST_ITERATIONS: usize = 1000;
 
 #[test]
 fn run_tests() {
-    const ERROR_THRESHOLD: f64 = 1e-12;
+    const ERROR_THRESHOLD: f64 = 1e-9;
 
     let setups: Vec<Box<dyn Setup>> = vec![
         Box::new(Subproblem1Setup::new()),
         Box::new(Subproblem2Setup::new()),
+        Box::new(Subproblem2ExtendedSetup::new()),
     ];
 
     for mut setup in setups {
@@ -19,11 +20,11 @@ fn run_tests() {
             setup.run();
 
             let error = setup.error();
-            assert!(error <= ERROR_THRESHOLD);
+            assert!(error <= ERROR_THRESHOLD, "{} error was at: {}", setup.name(), error);
             total_error += error;
         }
 
-        println!("{} Avg Error:\t{}", setup.name(), total_error / TEST_ITERATIONS as f64);
+        println!("{}\tAvg Error:\t{}", setup.name(), total_error / TEST_ITERATIONS as f64);
     }
 }
 
@@ -39,7 +40,7 @@ fn run_tests_ls() {
             setup.setup_ls();
             setup.run();
 
-            assert!(setup.is_at_local_min());
+            assert!(setup.is_at_local_min(), "{} solution was not at local minimum", setup.name());
         }
     }
 }
diff --git a/rust/src/tests/setups.rs b/rust/src/tests/setups.rs
index 2106c4c..2838a4b 100644
--- a/rust/src/tests/setups.rs
+++ b/rust/src/tests/setups.rs
@@ -7,19 +7,15 @@ use {
             random_norm_vector3
         },
 
-        subproblems::{ subproblem1, subproblem2 }
+        subproblems::{ SolutionSet2, subproblem1, subproblem2, subproblem2extended }
     },
 
-    nalgebra::{ Vector3, DVector },
+    nalgebra::Vector3,
 };
 
 /// An interface for setting up subproblem testing. Resposible for generating parameters, running
 /// the function, and calculating data such as the error.
 pub trait Setup {
-    fn name(&self) -> &'static str;
-    fn is_at_local_min(&self) -> bool;
-    fn error(&self) -> f64;
-
     /// Initialize parameters to test the case where theta is solved for exactly
     fn setup(&mut self);
 
@@ -27,6 +23,10 @@ pub trait Setup {
     fn setup_ls(&mut self);
 
     fn run(&mut self);
+
+    fn error(&self) -> f64;
+    fn is_at_local_min(&self) -> bool;
+    fn name(&self) -> &'static str;
 }
 
 pub struct Subproblem1Setup {
@@ -44,8 +44,19 @@ pub struct Subproblem2Setup {
     k1: Vector3<f64>,
     k2: Vector3<f64>,
 
-    theta1: DVector<f64>,
-    theta2: DVector<f64>,
+    theta1: SolutionSet2<f64>,
+    theta2: SolutionSet2<f64>,
+}
+
+pub struct Subproblem2ExtendedSetup {
+    p0: Vector3<f64>,
+    p1: Vector3<f64>,
+    p2: Vector3<f64>,
+    k1: Vector3<f64>,
+    k2: Vector3<f64>,
+
+    theta1: f64,
+    theta2: f64,
 }
 
 impl Subproblem1Setup {
@@ -72,12 +83,12 @@ impl Subproblem2Setup {
             k1: Vector3::zeros(),
             k2: Vector3::zeros(),
 
-            theta1: DVector::zeros(0),
-            theta2: DVector::zeros(0),
+            theta1: SolutionSet2::One(0.0),
+            theta2: SolutionSet2::One(0.0),
         }
     }
 
-    fn calculate_error(&self, theta1: &DVector<f64>, theta2: &DVector<f64>) -> f64 {
+    fn calculate_error(&self, theta1: &Vec<f64>, theta2: &Vec<f64>) -> f64 {
         theta1.iter()
             .zip(theta2.iter())
             .map(|(&t1, &t2)| {
@@ -87,6 +98,21 @@ impl Subproblem2Setup {
     }
 }
 
+impl Subproblem2ExtendedSetup {
+    pub fn new() -> Self {
+        Self {
+            p0: Vector3::zeros(),
+            p1: Vector3::zeros(),
+            p2: Vector3::zeros(),
+            k1: Vector3::zeros(),
+            k2: Vector3::zeros(),
+
+            theta1: 0.0,
+            theta2: 0.0,
+        }
+    }
+}
+
 impl Setup for Subproblem1Setup {
     fn setup(&mut self) {
         self.p1 = random_vector3();
@@ -118,7 +144,7 @@ impl Setup for Subproblem1Setup {
     }
 
     fn run(&mut self) {
-        (self.theta, _) = subproblem1(self.p1, self.p2, self.k);
+        (self.theta, _) = subproblem1(&self.p1, &self.p2, &self.k);
     }
 
     fn name(&self) -> &'static str {
@@ -128,13 +154,16 @@ impl Setup for Subproblem1Setup {
 
 impl Setup for Subproblem2Setup {
     fn setup(&mut self) {
+        let theta1 = random_angle();
+        let theta2 = random_angle();
+
         self.p1 = random_vector3();
         self.k1 = random_norm_vector3();
         self.k2 = random_norm_vector3();
-        self.theta1 = DVector::from_vec(vec![random_angle()]);
-        self.theta2 = DVector::from_vec(vec![random_angle()]);
+        self.theta1 = SolutionSet2::One(theta1);
+        self.theta2 = SolutionSet2::One(theta2);
 
-        self.p2 = rot(self.k2, -self.theta2[0]) * rot(self.k1,self.theta1[0]) * self.p1;
+        self.p2 = rot(self.k2, -theta2) * rot(self.k1, theta1) * self.p1;
     }
 
     fn setup_ls(&mut self) {
@@ -143,16 +172,16 @@ impl Setup for Subproblem2Setup {
         self.k1 = random_norm_vector3();
         self.k2 = random_norm_vector3();
         self.k2 = random_norm_vector3();
-        self.theta1 = DVector::from_vec(vec![random_angle()]);
-        self.theta2 = DVector::from_vec(vec![random_angle()]);
+        self.theta1 = SolutionSet2::One(random_angle());
+        self.theta2 = SolutionSet2::One(random_angle());
     }
 
     fn run(&mut self) {
-        (self.theta1, self.theta2, _) = subproblem2(self.p1, self.p2, self.k1, self.k2);
+        (self.theta1, self.theta2, _) = subproblem2(&self.p1, &self.p2, &self.k1, &self.k2);
     }
 
     fn error(&self) -> f64 {
-        self.calculate_error(&self.theta1, &self.theta2)
+        self.calculate_error(&self.theta1.get_all(), &self.theta2.get_all())
     }
 
     fn is_at_local_min(&self) -> bool {
@@ -161,10 +190,10 @@ impl Setup for Subproblem2Setup {
         let error = self.error();
         let error_check = error - DELTA;
 
-        let mut solutions = [self.theta1.clone(), self.theta2.clone()];
+        let mut solutions = [self.theta1.get_all(), self.theta2.get_all()];
 
         for i in 0..solutions.len() {
-            for j in 0..self.theta1.len() {
+            for j in 0..solutions[i].len() {
                 for sign in [-1.0, 1.0] {
                     let solution_prev = solutions[i][j];
                     solutions[i][j] += sign;
@@ -185,3 +214,38 @@ impl Setup for Subproblem2Setup {
         "Subproblem 2"
     }
 }
+
+impl Setup for Subproblem2ExtendedSetup {
+    fn setup(&mut self) {
+        self.p0 = random_vector3();
+        self.p1 = random_vector3();
+
+        self.k1 = random_norm_vector3();
+        self.k2 = random_norm_vector3();
+
+        self.theta1 = random_angle();
+        self.theta2 = random_angle();
+
+        self.p2 = rot(self.k2, -self.theta2) * (self.p0 + rot(self.k1, self.theta1) * self.p1);
+    }
+
+    fn setup_ls(&mut self) {
+        unimplemented!();
+    }
+
+    fn run(&mut self) {
+        (self.theta1, self.theta2) = subproblem2extended(&self.p0, &self.p1, &self.p2, &self.k1, &self.k2);
+    }
+
+    fn error(&self) -> f64 {
+        (self.p0 + rot(self.k1, self.theta1) * self.p1 - rot(self.k2, self.theta2) * self.p2).norm()
+    }
+
+    fn is_at_local_min(&self) -> bool {
+        unimplemented!();
+    }
+
+    fn name(&self) -> &'static str {
+        "Subproblem 2 Ex"
+    }
+}