added masks flag fro long short builder and percent builder

alphamind/portfolio/longshortbulder.py

@@ -13,7 +13,14 @@ from alphamind.utilities import transform
 
 def long_short_build(er: np.ndarray,
                      leverage: float=1.,
-                     groups: np.ndarray=None) -> np.ndarray:
+                     groups: np.ndarray=None,
+                     masks: np.ndarray=None) -> np.ndarray:
+
+    er = er.copy()
+
+    if masks is not None:
+        er[~masks] = 0.
+
     if er.ndim == 1:
         er = er.reshape((-1, 1))
 

alphamind/portfolio/percentbuilder.py

@@ -12,7 +12,12 @@ from alphamind.utilities import groupby
 from alphamind.utilities import set_value
 
 
-def percent_build(er: np.ndarray, percent: float, groups: np.ndarray=None) -> np.ndarray:
+def percent_build(er: np.ndarray, percent: float, groups: np.ndarray=None, masks: np.ndarray=None) -> np.ndarray:
+
+    er = er.copy()
+
+    if masks is not None:
+        er[~masks] = -np.inf
 
     if er.ndim == 1 or (er.shape[0] == 1 or er.shape[1] == 1):
         # fast path methods for single column er

alphamind/tests/portfolio/test_longshortbuild.py

@@ -16,6 +16,9 @@ class TestLongShortBuild(unittest.TestCase):
     def setUp(self):
         self.x = np.random.randn(3000, 10)
         self.groups = np.random.randint(10, 40, size=3000)
+        choices = np.random.choice(3000, 100, replace=False)
+        self.masks = np.full(3000, True, dtype=bool)
+        self.masks[choices] = False
 
     def test_long_short_build(self):
         x = self.x[:, 0].flatten()
@@ -37,6 +40,16 @@ class TestLongShortBuild(unittest.TestCase):
         expected_weights = pd.DataFrame(self.x).groupby(self.groups).apply(lambda s: s / np.abs(s).sum(axis=0))
         np.testing.assert_array_almost_equal(calc_weights, expected_weights)
 
+    def test_long_short_build_with_masks(self):
+        x = self.x[:, 0].flatten()
+        masked_x = x.copy()
+        masked_x[~self.masks] = 0.
+        leverage = np.abs(masked_x).sum()
+        calc_weights = long_short_build(x, masks=self.masks, leverage=leverage).flatten()
+        expected_weights = x.copy()
+        expected_weights[~self.masks] = 0.
+        np.testing.assert_array_almost_equal(calc_weights, expected_weights)
+
 
 if __name__ == '__main__':
     unittest.main()

alphamind/tests/portfolio/test_percentbuild.py

@@ -18,6 +18,7 @@ class TestPercentBuild(unittest.TestCase):
         self.p_included = 0.1
         self.n_groups = 30
         self.n_portfolios = range(1, 10)
+        self.n_mask = 100
 
     def test_percent_build(self):
         n_include = int(self.n_samples * self.p_included)
@@ -55,6 +56,30 @@ class TestPercentBuild(unittest.TestCase):
 
             np.testing.assert_array_almost_equal(calc_weights, expected_weights)
 
+    def test_percent_build_with_masks(self):
+        for n_portfolio in self.n_portfolios:
+            x = np.random.randn(self.n_samples, n_portfolio)
+            choices = np.random.choice(self.n_samples, self.n_mask, replace=False)
+            masks = np.full(self.n_samples, True, dtype=bool)
+            masks[choices] = False
+
+            calc_weights = percent_build(x, self.p_included, masks=masks)
+
+            expected_weights = np.zeros((len(x), n_portfolio))
+
+            filtered_index = np.arange(len(x))[masks]
+            filtered_x = x[masks]
+            big_boolen = np.full(x.shape, False, dtype=bool)
+
+            n_included = int(self.p_included * len(x))
+            chosen = (-filtered_x).argsort(axis=0).argsort(axis=0) < n_included
+            big_boolen[filtered_index] = chosen
+
+            for j in range(x.shape[1]):
+                expected_weights[big_boolen[:, j], j] = 1.
+
+            np.testing.assert_array_almost_equal(calc_weights, expected_weights)
+
 
 if __name__ == '__main__':
     unittest.main()
\ No newline at end of file