Add Cholesky decomposition to nucal._linear_fit

hera_cal/nucal.py

@@ -761,10 +761,11 @@
     # Assert that the method is valid
     assert solver in [
         "lu_solve",
+        "cho_solve",
         "solve",
         "pinv",
         "lstsq",
-    ], "method must be one of {}".format(["lu_solve", "solve", "pinv", "lstsq"])
+    ], "method must be one of {}".format(["lu_solve", "cho_solve", "solve", "pinv", "lstsq"])
 
     # Assert that the regularization tolerance is non-negative
     assert alpha >= 0.0, "alpha must be non-negative."
@@ -789,6 +790,19 @@
         # Save info
         cached_output = {'LU': L}
 
+    elif solver == "cho_solve":
+        # Factor XTX using scipy.linalg.cho_factor
+        if "c_and_lower" in cached_input:
+            c_and_lower = cached_input.get('c_and_lower')
+        else:
+            c_and_lower = linalg.cho_factor(XTX)
+
+        # Solve the linear system of equations using scipy.linalg.cho_solve
+        beta = linalg.cho_solve(c_and_lower, Xy)
+
+        # Save info
+        cached_output = {'c_and_lower': c_and_lower}
+
     elif solver == "solve":
         # Solve the linear system of equations using np.linalg.solve
         beta = np.linalg.solve(XTX, Xy)

hera_cal/tests/test_nucal.py

@@ -426,12 +426,14 @@ def test_linear_fit():
     b2, _ = nucal._linear_fit(XTX, Xy, solver='solve')
     b3, _ = nucal._linear_fit(XTX, Xy, solver='lstsq')
     b4, cached_input = nucal._linear_fit(XTX, Xy, solver='pinv')
+    b5, _ = nucal._linear_fit(XTX, Xy, solver='cho_solve')
     assert cached_input.get('XTXinv') is not None
 
     # Show that all modes give the same result
     np.testing.assert_allclose(b1, b2, atol=1e-6)
     np.testing.assert_allclose(b1, b3, atol=1e-6)
     np.testing.assert_allclose(b1, b4, atol=1e-6)
+    np.testing.assert_allclose(b1, b5, atol=1e-6)
 
     # Test that the fit is correct
     model = np.dot(X, b4)