update rank build tests

alphamind/tests/portfolio/test_rankbuild.py

@@ -18,17 +18,20 @@ class TestRankBuild(unittest.TestCase):
         n_samples = 3000
         n_included = 300
 
-        x = np.random.randn(n_samples, 2)
+        n_portfolios = range(10)
 
-        calc_weights = rank_build(x, n_included)
+        for n_portfolio in n_portfolios:
+            x = np.random.randn(n_samples, n_portfolio)
 
-        expected_weights = np.zeros((len(x), 2))
-        masks = (-x).argsort(axis=0).argsort(axis=0) < n_included
+            calc_weights = rank_build(x, n_included)
 
-        for j in range(x.shape[1]):
-            expected_weights[masks[:, j], j] = 1. / n_included
+            expected_weights = np.zeros((len(x), n_portfolio))
+            masks = (-x).argsort(axis=0).argsort(axis=0) < n_included
 
-        np.testing.assert_array_almost_equal(calc_weights, expected_weights)
+            for j in range(x.shape[1]):
+                expected_weights[masks[:, j], j] = 1. / n_included
+
+            np.testing.assert_array_almost_equal(calc_weights, expected_weights)
 
     def test_rank_build_with_group(self):
 
@@ -36,19 +39,23 @@ class TestRankBuild(unittest.TestCase):
         n_include = 10
         n_groups = 30
 
-        x = np.random.randn(n_samples, 2)
-        groups = np.random.randint(n_groups, size=n_samples)
+        n_portfolios = range(10)
+
+        for n_portfolio in n_portfolios:
+
+            x = np.random.randn(n_samples, n_portfolio)
+            groups = np.random.randint(n_groups, size=n_samples)
 
-        calc_weights = rank_build(x, n_include, groups)
+            calc_weights = rank_build(x, n_include, groups)
 
-        grouped_ordering = pd.DataFrame(-x).groupby(groups).rank()
-        expected_weights = np.zeros((len(x), 2))
-        masks = (grouped_ordering <= n_include).values
-        choosed = masks.sum(axis=0)
-        for j in range(x.shape[1]):
-            expected_weights[masks[:, j], j] = 1. / choosed[j]
+            grouped_ordering = pd.DataFrame(-x).groupby(groups).rank()
+            expected_weights = np.zeros((len(x), n_portfolio))
+            masks = (grouped_ordering <= n_include).values
+            choosed = masks.sum(axis=0)
+            for j in range(x.shape[1]):
+                expected_weights[masks[:, j], j] = 1. / choosed[j]
 
-        np.testing.assert_array_almost_equal(calc_weights, expected_weights)
+            np.testing.assert_array_almost_equal(calc_weights, expected_weights)
 
 
 if __name__ == '__main__':