Fix underdetermined regularized least squares

gnm/shape/fitting_utils/regularized_least_squares.py

@@ -57,11 +57,11 @@ def regularized_least_squares(
   """
 
   regularization_sqrt = np.sqrt(regularization + _EPSILON)
+  num_coefficients = left_hand_side_data_term.shape[1]
 
   # Append a diagonal array multiplied with the regularization weight.
   left_hand_side_regularization_term = regularization_sqrt * np.eye(
-      left_hand_side_data_term.shape[0],
-      left_hand_side_data_term.shape[1],
+      num_coefficients,
       dtype=left_hand_side_data_term.dtype,
   )
   left_hand_side = np.concatenate(
@@ -70,7 +70,10 @@ def regularized_least_squares(
 
   # Append an array full of zeros to the right-hand side for the regularization
   # term.
-  right_hand_side_regularization_term = np.zeros_like(right_hand_side_data_term)
+  right_hand_side_regularization_term = np.zeros(
+      (num_coefficients, *right_hand_side_data_term.shape[1:]),
+      dtype=right_hand_side_data_term.dtype,
+  )
   right_hand_side = np.concatenate(
       [right_hand_side_data_term, right_hand_side_regularization_term]
   )
@@ -93,11 +96,11 @@ def __init__(
       regularization: The scalar regularization that will be used.
     """
     regularization_sqrt = np.sqrt(regularization + _EPSILON)
+    self._num_coefficients = left_hand_side_data_term.shape[1]
 
     # Append a diagonal array multiplied with the regularization weight.
     left_hand_side_regularization_term = regularization_sqrt * np.eye(
-        left_hand_side_data_term.shape[0],
-        left_hand_side_data_term.shape[1],
+        self._num_coefficients,
         dtype=left_hand_side_data_term.dtype,
     )
     left_hand_side = np.concatenate(
@@ -121,8 +124,9 @@ def solve(
     Returns:
       The least-squares solution, x.
     """
-    right_hand_side_regularization_term = np.zeros_like(
-        right_hand_side_data_term
+    right_hand_side_regularization_term = np.zeros(
+        (self._num_coefficients, *right_hand_side_data_term.shape[1:]),
+        dtype=right_hand_side_data_term.dtype,
     )
     right_hand_side = np.concatenate(
         [right_hand_side_data_term, right_hand_side_regularization_term],

gnm/shape/fitting_utils/regularized_least_squares_test.py

@@ -74,6 +74,39 @@ def test_solver(self, dim: int, weight: float):
     )
     np.testing.assert_allclose(actual_solution, expected_solution)
 
+  @parameterized.parameters(False, True)
+  def test_underdetermined_system(self, multiple_targets):
+    """Tests regularization when there are fewer equations than unknowns."""
+    array = np.array([[1.0, 1.0]])
+    target = np.array([[1.0, 2.0]]) if multiple_targets else np.array([1.0])
+    weight = 1.0
+
+    actual_solution = _regularized_least_squares(array, target, weight)[0]
+    expected_solution = (
+        np.linalg.inv(array.T @ array + weight * np.eye(array.shape[1]))
+        @ array.T
+        @ target
+    )
+
+    np.testing.assert_allclose(actual_solution, expected_solution)
+
+  @parameterized.parameters(False, True)
+  def test_solver_underdetermined_system(self, multiple_targets):
+    """Tests the solver when there are fewer equations than unknowns."""
+    array = np.array([[1.0, 1.0]])
+    target = np.array([[1.0, 2.0]]) if multiple_targets else np.array([1.0])
+    weight = 1.0
+
+    solver = regularized_least_squares.RegularizedLeastSquares(array, weight)
+    actual_solution = solver.solve(target)
+    expected_solution = (
+        np.linalg.inv(array.T @ array + weight * np.eye(array.shape[1]))
+        @ array.T
+        @ target
+    )
+
+    np.testing.assert_allclose(actual_solution, expected_solution)
+
 
 if __name__ == "__main__":
   absltest.main()