From d3fa55f9e431fa1a4f2e5d2228968a93122c637f Mon Sep 17 00:00:00 2001
From: pcain-usgs <pcain@usgs.gov>
Date: Wed, 10 Feb 2021 11:29:15 -0700
Subject: [PATCH] Refactor QRFactorization

---
 geomagio/adjusted/transform/LeastSq.py        | 21 ++++--
 .../adjusted/transform/QRFactorization.py     | 65 +++++++------------
 geomagio/adjusted/transform/Rescale3D.py      |  2 +-
 .../transform/RotationTranslationXY.py        |  2 +-
 geomagio/adjusted/transform/SVD.py            | 28 ++++----
 geomagio/adjusted/transform/ShearYZ.py        |  2 +-
 geomagio/adjusted/transform/Transform.py      |  1 +
 .../adjusted/transform/TranslateOrigins.py    |  4 +-
 .../adjusted/transform/ZRotationHScale.py     |  2 +-
 .../transform/ZRotationHScaleZBaseline.py     |  4 +-
 geomagio/adjusted/transform/ZRotationShear.py |  2 +-
 11 files changed, 62 insertions(+), 71 deletions(-)

diff --git a/geomagio/adjusted/transform/LeastSq.py b/geomagio/adjusted/transform/LeastSq.py
index 45c6cb8bc..f33eace15 100644
--- a/geomagio/adjusted/transform/LeastSq.py
+++ b/geomagio/adjusted/transform/LeastSq.py
@@ -8,7 +8,7 @@ from .Transform import Transform
 class LeastSq(Transform):
     """Intance of Transform. Applies least squares to generate matrices"""
 
-    def get_stacked_values(self, absolutes, ordinates):
+    def get_stacked_values(self, absolutes, ordinates, weights=None):
         # LHS, or dependent variables
         # [A[0,0], A[1,0], A[2,0], A[0,1], A[1,1], A[2,1], ...]
         abs_stacked = self.get_stacked_absolutes(absolutes)
@@ -54,6 +54,11 @@ class LeastSq(Transform):
 
         return ord_stacked
 
+    def valid(self, rank):
+        if rank < self.ndims:
+            return False
+        return True
+
     def get_weighted_values(
         self,
         values: Tuple[List[float], List[float], List[float]],
@@ -83,17 +88,19 @@ class LeastSq(Transform):
         weights: List[float],
     ) -> np.array:
         """Calculates affine with no constraints using least squares."""
-        abs_stacked, ord_stacked = self.get_stacked_values(absolutes, ordinates)
+        abs_stacked, ord_stacked = self.get_stacked_values(
+            absolutes, ordinates, weights
+        )
         ord_stacked = self.get_weighted_values(ord_stacked, weights)
         abs_stacked = self.get_weighted_values(abs_stacked, weights)
         # regression matrix M that minimizes L2 norm
         matrix, res, rank, sigma = spl.lstsq(ord_stacked.T, abs_stacked.T)
-        if rank < 3:
-            print("Poorly conditioned or singular matrix, returning NaNs")
-            return np.nan * np.ones((4, 4))
-        return self.format_matrix(matrix)
+        if self.valid(rank):
+            return self.get_matrix(matrix, absolutes, ordinates, weights)
+        print("Poorly conditioned or singular matrix, returning NaNs")
+        return np.nan * np.ones((4, 4))
 
-    def format_matrix(self, matrix):
+    def get_matrix(self, matrix, absolutes=None, ordinates=None, weights=None):
         return np.array(
             [
                 [matrix[0], matrix[1], matrix[2], matrix[3]],
diff --git a/geomagio/adjusted/transform/QRFactorization.py b/geomagio/adjusted/transform/QRFactorization.py
index fa8187382..b633ffb66 100644
--- a/geomagio/adjusted/transform/QRFactorization.py
+++ b/geomagio/adjusted/transform/QRFactorization.py
@@ -2,15 +2,17 @@ import numpy as np
 import scipy.linalg as spl
 from typing import List, Optional, Tuple
 
+from .LeastSq import LeastSq
 from .SVD import SVD
 
 
-class QRFactorization(SVD):
+class QRFactorization(LeastSq):
     """Calculates affine using singular value decomposition with QR factorization"""
 
     ndims = 2
+    svd = SVD(ndims=ndims)
 
-    def get_weighted_values_lstsq(
+    def get_weighted_values(
         self,
         values: Tuple[List[float], List[float], List[float]],
         weights: Optional[List[float]],
@@ -19,57 +21,34 @@ class QRFactorization(SVD):
         if weights is None:
             return values
         weights = np.sqrt(weights)
-        return np.array(
-            [
-                values[0] * weights,
-                values[1] * weights,
-            ]
-        )
-
-    def calculate(
-        self,
-        ordinates: Tuple[List[float], List[float], List[float]],
-        absolutes: Tuple[List[float], List[float], List[float]],
-        weights: List[float],
-    ) -> np.array:
+        return np.array([values[i] * weights for i in range(self.ndims)])
 
-        if weights is None:
-            weights = np.ones_like(ordinates[0])
-
-        weighted_absolutes = self.get_weighted_values(values=absolutes, weights=weights)
-        weighted_ordinates = self.get_weighted_values(values=ordinates, weights=weights)
+    def get_stacked_values(self, absolutes, ordinates, weights):
+        weighted_absolutes = self.svd.get_weighted_values(
+            values=absolutes, weights=weights
+        )
+        weighted_ordinates = self.svd.get_weighted_values(
+            values=ordinates, weights=weights
+        )
         # LHS, or dependent variables
-        abs_stacked = self.get_stacked_values(
+        abs_stacked = self.svd.get_stacked_values(
             values=absolutes,
             weighted_values=weighted_absolutes,
         )
 
         # RHS, or independent variables
-        ord_stacked = self.get_stacked_values(
+        ord_stacked = self.svd.get_stacked_values(
             values=ordinates,
             weighted_values=weighted_ordinates,
         )
+        return abs_stacked, ord_stacked
 
-        abs_stacked = self.get_weighted_values_lstsq(
-            values=abs_stacked,
-            weights=weights,
-        )
-        ord_stacked = self.get_weighted_values_lstsq(
-            values=ord_stacked,
-            weights=weights,
-        )
-
-        # regression matrix M that minimizes L2 norm
-        M_r, res, rank, sigma = spl.lstsq(ord_stacked.T, abs_stacked.T)
-        if rank < 2:
-            print("Poorly conditioned or singular matrix, returning NaNs")
-            return np.nan * np.ones((4, 4))
-
+    def get_matrix(self, matrix, absolutes, ordinates, weights):
         # QR fatorization
         # NOTE: forcing the diagonal elements of Q to be positive
         #       ensures that the determinant is 1, not -1, and is
         #       therefore a rotation, not a reflection
-        Q, R = np.linalg.qr(M_r.T)
+        Q, R = np.linalg.qr(matrix.T)
         neg = np.diag(Q) < 0
         Q[:, neg] = -1 * Q[:, neg]
         R[neg, :] = -1 * R[neg, :]
@@ -81,10 +60,11 @@ class QRFactorization(SVD):
         # combine shear and rotation
         QH = np.dot(Q, H)
 
+        weighted_absolutes = self.svd.get_weighted_values(absolutes, weights)
+        weighted_ordinates = self.svd.get_weighted_values(ordinates, weights)
+
         # now get translation using weighted centroids and R
-        T = np.array([weighted_absolutes[0], weighted_absolutes[1]]) - np.dot(
-            QH, [weighted_ordinates[0], weighted_ordinates[1]]
-        )
+        T = self.svd.get_translation_matrix(QH, weighted_absolutes, weighted_ordinates)
 
         return [
             [QH[0, 0], QH[0, 1], 0.0, T[0]],
@@ -93,7 +73,8 @@ class QRFactorization(SVD):
                 0.0,
                 0.0,
                 1.0,
-                np.array(weighted_absolutes[2]) - np.array(weighted_ordinates[2]),
+                np.array(weighted_absolutes[self.ndims])
+                - np.array(weighted_ordinates[self.ndims]),
             ],
             [0.0, 0.0, 0.0, 1.0],
         ]
diff --git a/geomagio/adjusted/transform/Rescale3D.py b/geomagio/adjusted/transform/Rescale3D.py
index c55ca8f1e..c74a9bb72 100644
--- a/geomagio/adjusted/transform/Rescale3D.py
+++ b/geomagio/adjusted/transform/Rescale3D.py
@@ -21,7 +21,7 @@ class Rescale3D(LeastSq):
         ord_stacked[2, 2::3] = ordinates[2]
         return ord_stacked
 
-    def format_matrix(self, matrix):
+    def get_matrix(self, matrix, absolutes=None, ordinates=None, weights=None):
         return [
             [matrix[0], 0.0, 0.0, 0.0],
             [0.0, matrix[1], 0.0, 0.0],
diff --git a/geomagio/adjusted/transform/RotationTranslationXY.py b/geomagio/adjusted/transform/RotationTranslationXY.py
index ade36e195..f140a3845 100644
--- a/geomagio/adjusted/transform/RotationTranslationXY.py
+++ b/geomagio/adjusted/transform/RotationTranslationXY.py
@@ -14,7 +14,7 @@ class RotationTranslationXY(SVD):
     def get_rotation_matrix(self, U, Vh):
         return np.dot(Vh.T, np.dot(np.diag([1, np.linalg.det(np.dot(Vh.T, U.T))]), U.T))
 
-    def format_matrix(self, R, T, weighted_absolutes, weighted_ordinates):
+    def get_matrix(self, R, T, weighted_absolutes, weighted_ordinates):
         return [
             [R[0, 0], R[0, 1], 0.0, T[0]],
             [R[1, 0], R[1, 1], 0.0, T[1]],
diff --git a/geomagio/adjusted/transform/SVD.py b/geomagio/adjusted/transform/SVD.py
index d1b8505dc..f9ff49056 100644
--- a/geomagio/adjusted/transform/SVD.py
+++ b/geomagio/adjusted/transform/SVD.py
@@ -7,8 +7,6 @@ from .Transform import Transform
 class SVD(Transform):
     """Instance of Transform. Applies singular value decomposition to generate matrices"""
 
-    ndims = 3
-
     def get_covariance_matrix(self, absolutes, ordinates, weights):
         weighted_ordinates = self.get_weighted_values(values=ordinates, weights=weights)
         weighted_absolutes = self.get_weighted_values(values=absolutes, weights=weights)
@@ -67,6 +65,11 @@ class SVD(Transform):
             R, [weighted_ordinates[i] for i in range(self.ndims)]
         )
 
+    def valid(self, singular_values):
+        if np.sum(singular_values) < self.ndims:
+            return False
+        return True
+
     def calculate(
         self,
         ordinates: Tuple[List[float], List[float], List[float]],
@@ -80,21 +83,20 @@ class SVD(Transform):
         # Singular value decomposition, then rotation matrix from L&R eigenvectors
         # (the determinant guarantees a rotation, and not a reflection)
         U, S, Vh = np.linalg.svd(H)
-        if np.sum(S) < 3:
-            print("Poorly conditioned or singular matrix, returning NaNs")
-            return np.nan * np.ones((4, 4))
-        R = self.get_rotation_matrix(U, Vh)
-        # now get translation using weighted centroids and R
-        T = self.get_translation_matrix(R, weighted_absolutes, weighted_ordinates)
-
-        return self.format_matrix(R, T, weighted_absolutes, weighted_ordinates)
+        if self.valid(S):
+            R = self.get_rotation_matrix(U, Vh)
+            # now get translation using weighted centroids and R
+            T = self.get_translation_matrix(R, weighted_absolutes, weighted_ordinates)
+            return self.get_matrix(R, T, weighted_absolutes, weighted_ordinates)
+        print("Poorly conditioned or singular matrix, returning NaNs")
+        return np.nan * np.ones((4, 4))
 
-    def format_matrix(
+    def get_matrix(
         self,
         R,
         T,
-        weighted_absolutes,
-        weighted_ordinates,
+        weighted_absolutes=None,
+        weighted_ordinates=None,
     ):
         return [
             [R[0, 0], R[0, 1], R[0, 2], T[0]],
diff --git a/geomagio/adjusted/transform/ShearYZ.py b/geomagio/adjusted/transform/ShearYZ.py
index dc2e43c21..7873c7b29 100644
--- a/geomagio/adjusted/transform/ShearYZ.py
+++ b/geomagio/adjusted/transform/ShearYZ.py
@@ -23,7 +23,7 @@ class ShearYZ(LeastSq):
         ord_stacked[2, 2::3] = 1.0
         return ord_stacked
 
-    def format_matrix(self, matrix):
+    def get_matrix(self, matrix, absolutes=None, ordinates=None, weights=None):
         return [
             [1.0, 0.0, 0.0, 0.0],
             [matrix[0], 1.0, 0.0, 0.0],
diff --git a/geomagio/adjusted/transform/Transform.py b/geomagio/adjusted/transform/Transform.py
index 5fdbad7c6..a365d9f9f 100644
--- a/geomagio/adjusted/transform/Transform.py
+++ b/geomagio/adjusted/transform/Transform.py
@@ -16,6 +16,7 @@ class Transform(BaseModel):
 
     acausal: bool = False
     memory: Optional[float] = None
+    ndims = 3
 
     def get_weights(self, times: UTCDateTime, time: int = None) -> List[float]:
         """
diff --git a/geomagio/adjusted/transform/TranslateOrigins.py b/geomagio/adjusted/transform/TranslateOrigins.py
index 69c21915f..fbb08b8f9 100644
--- a/geomagio/adjusted/transform/TranslateOrigins.py
+++ b/geomagio/adjusted/transform/TranslateOrigins.py
@@ -8,7 +8,7 @@ from .LeastSq import LeastSq
 class TranslateOrigins(LeastSq):
     """Calculates affine using using least squares, constrained to tanslate origins"""
 
-    def get_stacked_values(self, absolutes, ordinates):
+    def get_stacked_values(self, absolutes, ordinates, weights=None):
         # LHS, or dependent variables
         abs_stacked = self.get_stacked_absolutes(absolutes)
         # subtract ords from abs to force simple translation
@@ -40,7 +40,7 @@ class TranslateOrigins(LeastSq):
         ord_stacked[2, 2::3] = 1.0
         return ord_stacked
 
-    def format_matrix(self, matrix):
+    def get_matrix(self, matrix, absolutes=None, ordinates=None, weights=None):
         return [
             [1.0, 0.0, 0.0, matrix[0]],
             [0.0, 1.0, 0.0, matrix[1]],
diff --git a/geomagio/adjusted/transform/ZRotationHScale.py b/geomagio/adjusted/transform/ZRotationHScale.py
index a5536bce8..4e56a6e83 100644
--- a/geomagio/adjusted/transform/ZRotationHScale.py
+++ b/geomagio/adjusted/transform/ZRotationHScale.py
@@ -26,7 +26,7 @@ class ZRotationHscale(LeastSq):
         ord_stacked[5, 2::3] = 1.0
         return ord_stacked
 
-    def format_matrix(self, matrix):
+    def get_matrix(self, matrix, absolutes=None, ordinates=None, weights=None):
         return [
             [matrix[0], matrix[1], 0.0, matrix[2]],
             [-matrix[1], matrix[0], 0.0, matrix[3]],
diff --git a/geomagio/adjusted/transform/ZRotationHScaleZBaseline.py b/geomagio/adjusted/transform/ZRotationHScaleZBaseline.py
index 569786afc..13393ac38 100644
--- a/geomagio/adjusted/transform/ZRotationHScaleZBaseline.py
+++ b/geomagio/adjusted/transform/ZRotationHScaleZBaseline.py
@@ -9,7 +9,7 @@ class ZRotationHscaleZbaseline(LeastSq):
     """Calculates affine using least squares, constrained to rotate about the Z axis,
     apply a uniform horizontal scaling, and apply a baseline shift for the Z axis."""
 
-    def get_stacked_values(self, absolutes, ordinates):
+    def get_stacked_values(self, absolutes, ordinates, weights=None):
         # LHS, or dependent variables
         abs_stacked = self.get_stacked_absolutes(absolutes)
         # subtract z_o from z_a to force simple z translation
@@ -34,7 +34,7 @@ class ZRotationHscaleZbaseline(LeastSq):
         ord_stacked[2, 2::3] = 1.0
         return ord_stacked
 
-    def format_matrix(self, matrix):
+    def get_matrix(self, matrix, absolutes=None, ordinates=None, weights=None):
         return [
             [matrix[0], matrix[1], 0.0, 0.0],
             [-matrix[1], matrix[0], 0.0, 0.0],
diff --git a/geomagio/adjusted/transform/ZRotationShear.py b/geomagio/adjusted/transform/ZRotationShear.py
index 693cbb415..d1872c0e6 100644
--- a/geomagio/adjusted/transform/ZRotationShear.py
+++ b/geomagio/adjusted/transform/ZRotationShear.py
@@ -22,7 +22,7 @@ class ZRotationShear(LeastSq):
         ord_stacked[7, 2::3] = 1.0
         return ord_stacked
 
-    def format_matrix(self, matrix):
+    def get_matrix(self, matrix, absolutes=None, ordinates=None, weights=None):
         return [
             [matrix[0], matrix[1], 0.0, matrix[2]],
             [matrix[3], matrix[4], 0.0, matrix[5]],
-- 
GitLab