more robust and more efficient

alphamind/data/neutralize.py

@@ -16,12 +16,16 @@ import alphamind.utilities as utils
 def neutralize(x: np.ndarray,
                y: np.ndarray,
                groups: np.ndarray=None,
-               detail: bool=False) \
+               detail: bool=False,
+               weights: np.ndarray = None) \
         -> Union[np.ndarray, Tuple[np.ndarray, Dict]]:
 
     if y.ndim == 1:
         y = y.reshape((-1, 1))
 
+    if weights is None:
+        weights = np.ones(len(y), dtype=float)
+
     output_dict = {}
 
     if detail:
@@ -37,17 +41,21 @@ def neutralize(x: np.ndarray,
         if detail:
             for diff_loc in index_diff:
                 curr_idx = order[start:diff_loc + 1]
-                curr_x, b = _sub_step(x, y, curr_idx, res)
+                curr_x, b = _sub_step(x, y, weights, curr_idx, res)
                 exposure[curr_idx, :, :] = b
                 explained[curr_idx] = ls_explain(curr_x, b)
                 start = diff_loc + 1
         else:
             for diff_loc in index_diff:
                 curr_idx = order[start:diff_loc + 1]
-                _sub_step(x, y, curr_idx, res)
+                _sub_step(x, y, weights, curr_idx, res)
                 start = diff_loc + 1
     else:
-        b = ls_fit(x, y)
+        try:
+            b = ls_fit(x, y, weights)
+        except ValueError:
+            b = ls_fit_pinv(x, y, weights)
+
         res = ls_res(x, y, b)
 
         if detail:
@@ -60,17 +68,26 @@ def neutralize(x: np.ndarray,
         return res
 
 
-@nb.njit(nogil=True, cache=True)
-def _sub_step(x, y, curr_idx, res) -> Tuple[np.ndarray, np.ndarray]:
-    curr_x, curr_y = x[curr_idx], y[curr_idx]
-    b = ls_fit(curr_x, curr_y)
+def _sub_step(x, y, w, curr_idx, res) -> Tuple[np.ndarray, np.ndarray]:
+    curr_x, curr_y, curr_w = x[curr_idx], y[curr_idx], w[curr_idx]
+    try:
+        b = ls_fit(curr_x, curr_y, curr_w)
+    except ValueError:
+        b = ls_fit_pinv(curr_x, curr_y, curr_w)
     res[curr_idx] = ls_res(curr_x, curr_y, b)
     return curr_x, b
 
 
 @nb.njit(nogil=True, cache=True)
-def ls_fit(x: np.ndarray, y: np.ndarray) -> np.ndarray:
-    x_bar = x.T
+def ls_fit(x: np.ndarray, y: np.ndarray, w: np.ndarray) -> np.ndarray:
+    x_bar = x.T * w
+    b = np.linalg.solve(x_bar @ x, x_bar @ y)
+    return b
+
+
+@nb.njit(nogil=True, cache=True)
+def ls_fit_pinv(x: np.ndarray, y: np.ndarray, w: np.ndarray) -> np.ndarray:
+    x_bar = x.T * w
     b = np.linalg.pinv(x_bar @ x) @ x_bar @ y
     return b