fix rank function and test_rank

alphamind/data/rank.py

@@ -23,7 +23,7 @@ def rank(x: np.ndarray, groups: Optional[np.ndarray]=None) -> np.ndarray:
         start = 0
         for diff_loc in index_diff:
             curr_idx = order[start:diff_loc + 1]
-            res[curr_idx] = rankdata(x[curr_idx]).astype(float) - 1.
+            res[curr_idx] = (rankdata(x[curr_idx]).astype(float) - 1.).reshape((-1, 1))
             start = diff_loc + 1
         return res
     else:

alphamind/tests/data/test_rank.py

@@ -29,10 +29,14 @@ class TestRank(unittest.TestCase):
         np.testing.assert_array_less(0, arr_diff)
 
     def test_rank_with_groups(self):
-        data_rank = rank(self.x, groups=self.groups)
-
-        df = pd.DataFrame(self.x, index=self.groups)
-        expected_rank = df.groupby(level=0).apply(lambda x: x.values.argsort(axis=0).argsort(axis=0))
-        print(expected_rank)
-
-
+        data = pd.DataFrame(data={'raw': self.x.tolist()}, index=self.groups)
+        data['rank'] = rank(data['raw'], groups=data.index)
+        groups = dict(list(data['rank'].groupby(level=0)))
+        ret = []
+        for index in range(10):
+            ret.append(groups[index].values)
+        ret = np.concatenate(ret).reshape(-1, 1)
+
+        expected_rank = data['raw'].groupby(level=0).apply(lambda x: x.values.argsort(axis=0).argsort(axis=0))
+        expected_rank = np.concatenate(expected_rank).reshape(-1, 1)
+        np.testing.assert_array_equal(ret, expected_rank)