From 72d8f4914549956c816c6c95a95238b4d526fbd3 Mon Sep 17 00:00:00 2001
From: Matt <matthew.morley.ca@gmail.com>
Date: Sun, 13 Nov 2022 15:49:59 -0500
Subject: [PATCH] Add orthogonalizeRotationMatrix (#587)

* Add orthogonalizeRotationMatrix

* Update docstring

Co-authored-by: Chris Gerth <gerth2@users.noreply.github.com>
---
 .../common/util/math/MathUtils.java           | 36 +++++++++++++++++++
 .../vision/apriltag/DetectionResult.java      |  6 ++--
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/photon-core/src/main/java/org/photonvision/common/util/math/MathUtils.java b/photon-core/src/main/java/org/photonvision/common/util/math/MathUtils.java
index df9d711ea..e46aa5b74 100644
--- a/photon-core/src/main/java/org/photonvision/common/util/math/MathUtils.java
+++ b/photon-core/src/main/java/org/photonvision/common/util/math/MathUtils.java
@@ -18,6 +18,7 @@
 package org.photonvision.common.util.math;
 
 import edu.wpi.first.math.MatBuilder;
+import edu.wpi.first.math.Matrix;
 import edu.wpi.first.math.Nat;
 import edu.wpi.first.math.VecBuilder;
 import edu.wpi.first.math.geometry.CoordinateSystem;
@@ -25,10 +26,14 @@ import edu.wpi.first.math.geometry.Pose3d;
 import edu.wpi.first.math.geometry.Quaternion;
 import edu.wpi.first.math.geometry.Rotation3d;
 import edu.wpi.first.math.geometry.Transform3d;
+import edu.wpi.first.math.numbers.N3;
 import edu.wpi.first.math.util.Units;
 import edu.wpi.first.util.WPIUtilJNI;
 import java.util.Arrays;
 import java.util.List;
+import org.ejml.data.DMatrixRMaj;
+import org.ejml.dense.row.factory.DecompositionFactory_DDRM;
+import org.ejml.simple.SimpleMatrix;
 import org.opencv.core.Mat;
 
 public class MathUtils {
@@ -201,4 +206,35 @@ public class MathUtils {
         var axis = rotation.getAxis().times(angle);
         rvecOutput.put(0, 0, axis.getData());
     }
+
+    /**
+     * Orthogonalize an input matrix using a QR decomposition. QR decompositions decompose a
+     * rectangular matrix 'A' such that 'A=QR', where Q is the closest orthogonal matrix to the input,
+     * and R is an upper triangular matrix.
+     */
+    public static Matrix<N3, N3> orthogonalizeRotationMatrix(Matrix<N3, N3> input) {
+        var a = DecompositionFactory_DDRM.qr(3, 3);
+        if (!a.decompose(input.getStorage().getDDRM())) {
+            // best we can do is return the input
+            return input;
+        }
+
+        // Grab results (thanks for this _great_ api, EJML)
+        var Q = new DMatrixRMaj(3, 3);
+        var R = new DMatrixRMaj(3, 3);
+        a.getQ(Q, false);
+        a.getR(R, false);
+
+        // Fix signs in R if they're < 0 so it's close to an identity matrix
+        // (our QR decomposition implementation sometimes flips the signs of columns)
+        for (int colR = 0; colR < 3; ++colR) {
+            if (R.get(colR, colR) < 0) {
+                for (int rowQ = 0; rowQ < 3; ++rowQ) {
+                    Q.set(rowQ, colR, -Q.get(rowQ, colR));
+                }
+            }
+        }
+
+        return new Matrix<>(new SimpleMatrix(Q));
+    }
 }
diff --git a/photon-core/src/main/java/org/photonvision/vision/apriltag/DetectionResult.java b/photon-core/src/main/java/org/photonvision/vision/apriltag/DetectionResult.java
index 481547b14..268529b70 100644
--- a/photon-core/src/main/java/org/photonvision/vision/apriltag/DetectionResult.java
+++ b/photon-core/src/main/java/org/photonvision/vision/apriltag/DetectionResult.java
@@ -33,6 +33,8 @@ import edu.wpi.first.math.geometry.Transform3d;
 import edu.wpi.first.math.geometry.Translation3d;
 import java.util.Arrays;
 
+import org.photonvision.common.util.math.MathUtils;
+
 public class DetectionResult {
     public int getId() {
         return id;
@@ -136,12 +138,12 @@ public class DetectionResult {
         this.poseResult1 =
                 new Transform3d(
                         new Translation3d(pose1TransArr[0], pose1TransArr[1], pose1TransArr[2]),
-                        new Rotation3d(new MatBuilder<>(Nat.N3(), Nat.N3()).fill(pose1RotArr)));
+                        new Rotation3d(MathUtils.orthogonalizeRotationMatrix(new MatBuilder<>(Nat.N3(), Nat.N3()).fill(pose1RotArr))));
         this.error2 = err2;
         this.poseResult2 =
                 new Transform3d(
                         new Translation3d(pose2TransArr[0], pose2TransArr[1], pose2TransArr[2]),
-                        new Rotation3d(new MatBuilder<>(Nat.N3(), Nat.N3()).fill(pose2RotArr)));
+                        new Rotation3d(MathUtils.orthogonalizeRotationMatrix(new MatBuilder<>(Nat.N3(), Nat.N3()).fill(pose2RotArr))));
     }
 
     /**