Refactor code

1 files changed, 10 insertions, 13 deletions

diff --git a/ml_exp/math.py b/ml_exp/math.py
index e7c8dcabc..da1fdf08f 100644
--- a/ml_exp/math.py
+++ b/ml_exp/math.py
@@ -23,23 +23,23 @@ SOFTWARE.
 import numpy as np
 
 
-def cholesky_solve(K, y):
+def cholesky_solve(K,
+                   y):
     """
-    Applies Cholesky decomposition to obtain the 'alpha coeficients'.
+    Applies Cholesky decomposition to solve Ka=y. Where a are the alpha
+        coeficients.
     K: kernel.
     y: known parameters.
     """
-    # The initial mathematical problem is to solve Ka=y.
-
-    # First, add a small lambda value.
+    # Add a small lambda value.
     K[np.diag_indices_from(K)] += 1e-8
 
     # Get the Cholesky decomposition of the kernel.
     L = np.linalg.cholesky(K)
-    size = len(L)
+    size = K.shape[0]
 
     # Solve Lx=y for x.
-    x = np.zeros(size)
+    x = np.zeros(size, dtype=float)
     x[0] = y[0] / L[0, 0]
     for i in range(1, size):
         temp_sum = 0.0
@@ -49,12 +49,9 @@ def cholesky_solve(K, y):
 
     # Now, solve LTa=x for a.
     L2 = L.T
-    a = np.zeros(size)
-    a_ms = size - 1
-    a[a_ms] = x[a_ms] / L2[a_ms, a_ms]
-    # Because of the form of L2 (upper triangular matriz), an inversion of
-    # range() needs to be done.
-    for i in range(0, a_ms)[::-1]:
+    a = np.zeros(size, dtype=float)
+    a[size - 1] = x[size - 1] / L2[size - 1, size - 1]
+    for i in range(0, size - 1)[::-1]:
         temp_sum = 0.0
         for j in range(i, size)[::-1]:
             temp_sum += L2[i, j] * a[j]