Provide a cleaner interface.

Author: nymath

examples/main.py

@@ -11,7 +11,7 @@
 from torchqtm.utils.universe import StaticUniverse, IndexComponents
 from torchqtm.utils.warnings import catch_warnings
 from torchqtm.utils.benchmark import BenchMark
-from torchqtm.vbt.backtest import GroupTester01
+from torchqtm.vbt.backtest import GroupTester01, GroupTester02
 from torchqtm.alphas.alpha101 import *
 import torchqtm.op as op
 import torchqtm.op.functional as F
@@ -34,7 +34,6 @@ def __init__(self, env):
 
     def forward(self):
         self.data = F.divide(1, self.env.PE)
-        self.data = self.data.astype(np.float64)
         self.data = F.winsorize(self.data, 'std', 4)
         self.data = F.normalize(self.data)
         self.data = F.group_neutralize(self.data, self.env.Sector)
@@ -103,8 +102,8 @@ def aux_func(data_slice):
         self.data = F.winsorize(self.data, 'std', 4)
         self.data = F.normalize(self.data)
         self.data = pd.DataFrame(self.data, index=self.env.Close.index, columns=self.env.Close.columns)
-        cond = F.geq(F.ts_mean(np.squeeze(Close, -1), 5), F.ts_mean(np.squeeze(Close, -1), 22))
-        self.data = F.trade_when(cond, self.data, False)
+        cond = F.geq(F.ts_mean(self.close, 5), F.ts_mean(self.close, 22))
+        # self.data = F.trade_when(cond, self.data, False)
         self.data = F.group_neutralize(self.data, self.env.Sector)
         self.data = F.regression_neut(self.data, self.env.MktVal)
         return self.data
@@ -134,15 +133,18 @@ def load_data():
     # Create alpha
     # alphas = Momentum01(env=btEnv0)
     # alphas = NeutralizePE(env=btEnv0)
-    alphas = Alpha060(env=btEnv0)
+    alphas = Alpha055(env=btEnv0)
     # alphas = Ross(env=btEnv0)
     # alphas.forward(btEnv.match_env(dfs['PE']))
     with Timer():
         with catch_warnings():
             alphas.forward()
     # run backtest
-    bt = GroupTester01(env=btEnv,
-                       n_groups=5)
+    bt = GroupTester02(env=btEnv,
+                       n_groups=5,
+                       weighting='equal',
+                       exclude_suspended=False,
+                       exclude_limits=False)
 
     with Timer():
         bt.run_backtest(bt.env.match_env(F.purify(alphas.data)))

quant/backtest.py

@@ -5,6 +5,7 @@
 from typing import Dict, Hashable
 from abc import ABCMeta, abstractmethod
 from typing import Iterable
+from collections import OrderedDict
 
 
 class BackTestEnv(object):
@@ -39,9 +40,10 @@ def __init__(self,
         self.MktVal = None
         self.PE = None
         self.Sector = None
-        self._FutureReturn = None
+        self.forward_returns = None
         self._create_datas()
         self._create_features()
+        self._create_forward_returns()
 
     def _check_dfs(self):
         """
@@ -53,30 +55,32 @@ def _check_dfs(self):
         assert 'Sector' in self.dfs
 
     def _create_datas(self):
-        self.datas = {}
+        self.data = {}
         for key in self.dfs:
             if isinstance(self.dfs[key], pd.DataFrame):
-                self.datas[key] = self.dfs[key].loc[self.dates, self.symbols]
+                self.data[key] = self.dfs[key].loc[self.dates, self.symbols]
             else:
-                self.datas[key] = self.dfs[key]
-        self.datas['_FutureReturn'] = self.datas['Close'].pct_change().shift(-1)
+                self.data[key] = self.dfs[key]
+
+    def _create_forward_returns(self, D=1):
+        forward_returns = self.data['Close'].pct_change().shift(-1)
+        setattr(self, 'forward_returns', forward_returns)
 
     def _create_features(self):
         """
         Create the reference to the dict values
         :return:
         """
-        for key in self.datas.keys():
-            setattr(self, key, self.datas[key])
-        setattr(self, '_FutureReturn', self.datas['_FutureReturn'])
+        for key in self.data.keys():
+            setattr(self, key, self.data[key])
 
     def __getitem__(self, item):
         """
         Keep the operator[]
         :param item:
         :return:
         """
-        return self.datas[item]
+        return self.data[item]
 
     def __setitem__(self, item, value):
         """
@@ -85,15 +89,15 @@ def __setitem__(self, item, value):
         :return:
         """
         assert isinstance(value, pd.DataFrame)
-        self.datas[item] = value
+        self.data[item] = value
         setattr(self, item, value)
 
     def __delitem__(self, item):
-        del self.datas[item]
+        del self.data[item]
         delattr(self, item)
 
     def __contains__(self, item):
-        return item in self.datas
+        return item in self.data
 
     def match_env(self, factor):
         return factor.loc[self.dates, self.symbols]

torchqtm/base.py

@@ -812,7 +812,8 @@ def tail(x, lower=0, upper=0.5, newval=np.nan):
     return if_else(cond, x, np.nan)
 
 
-# TODO: 保持变量类型封闭
+# TODO: 保持变量类型封闭, 目前有一堆BUG
+# BUG:
 def trade_when(trigger: np.ndarray[bool],
                alpha: np.ndarray[np.float64],
                exit_cond: np.ndarray[bool]):

torchqtm/op/functional.py

@@ -47,10 +47,12 @@ class BaseGroupTester(BaseTester, TesterMixin):
     def __init__(self,
                  env: BackTestEnv = None,
                  n_groups: int = 5,
+                 weighting: str = 'equal',
                  exclude_suspended: bool = False,
                  exclude_limits: bool = False):
         super().__init__(env)
         self.n_groups = n_groups
+        self.weighting = weighting
         self.exclude_limits = exclude_limits
         self.exclude_suspended = exclude_suspended
         self.returns = None
@@ -105,9 +107,12 @@ def plot(self):
 
 class GroupTester01(BaseGroupTester):
     def __init__(self,
-                 env: BackTestEnv,
-                 n_groups: int = 5):
-        super().__init__(env, n_groups)
+                 env: BackTestEnv = None,
+                 n_groups: int = 5,
+                 weighting: str = 'equal',
+                 exclude_suspended: bool = False,
+                 exclude_limits: bool = False):
+        super().__init__(env, n_groups, weighting, exclude_suspended, exclude_limits)
 
     def run_backtest(self, modified_factor) -> None:
         assert modified_factor.shape == self.env['Close'].shape
@@ -133,9 +138,12 @@ def run_backtest(self, modified_factor) -> None:
 
                 def temp(x):
                     # TODO: develop a weight_scheme class
-                    weight = x['MktVal'] / x['MktVal'].sum()
-                    # weight = 1 / len(x['MktVal'])
-                    # weights.append(weight)
+                    if self.weighting == 'equal':
+                        weight = 1 / len(x['MktVal'])
+                    elif self.weighting == 'market_cap':
+                        weight = x['MktVal'] / x['MktVal'].sum()
+                    else:
+                        raise ValueError('Invalid weight scheme')
                     ret = x['_FutureReturn']
                     return (weight * ret).sum()
                 group_return = temp_data.groupby('group').apply(temp)
@@ -148,6 +156,55 @@ def temp(x):
         self.returns.columns.name = "group"
 
 
+class GroupTester02(BaseGroupTester):
+    def __init__(self,
+                 env: BackTestEnv = None,
+                 n_groups: int = 5,
+                 weighting: str = 'equal',
+                 exclude_suspended: bool = False,
+                 exclude_limits: bool = False):
+        super().__init__(env, n_groups, weighting, exclude_suspended, exclude_limits)
+
+    def run_backtest(self, modified_factor) -> None:
+        assert modified_factor.shape == self.env['Close'].shape
+        self._reset()
+        labels = ["group_" + str(i + 1) for i in range(self.n_groups)]
+        returns = []
+        for i in range(len(modified_factor)-1):
+            # If you are confused about concat series, you apply use the following way
+            # 1. series.unsqueeze(1) to generate an additional axes
+            # 2. concat these series along axis1
+            temp_data = pd.concat([self.env.forward_returns.iloc[i],
+                                   self.env.MktVal.iloc[i],
+                                   modified_factor.iloc[i]], axis=1)
+            temp_data.columns = ['forward_returns', 'MktVal', 'modified_factor']
+            # na stands for stocks that we you not insterested in
+            # We can develop a class to better represent this process.
+            temp_data = temp_data.loc[~np.isnan(temp_data['modified_factor'])]
+            if len(temp_data) == 0:
+                group_return = pd.Series(0, index=labels)
+            else:
+                temp_data['group'] = pd.qcut(temp_data['modified_factor'], self.n_groups, labels=labels)
+
+                def temp(x):
+                    # TODO: develop a weight_scheme class
+                    if self.weighting == 'equal':
+                        weight = 1 / len(x['MktVal'])
+                    elif self.weighting == 'market_cap':
+                        weight = x['MktVal'] / x['MktVal'].sum()
+                    else:
+                        raise ValueError('Invalid weight scheme')
+                    ret = x['forward_returns']
+                    return (weight * ret).sum()
+                group_return = temp_data.groupby('group').apply(temp)
+            returns.append(group_return)
+        returns.append(pd.Series(np.repeat(0, self.n_groups), index=group_return.index))
+        self.returns = pd.concat(returns, axis=1).T
+        # Here we need to transpose the return, since the rows are stocks.
+        self.returns.index = self.rebalance_dates
+        self.returns.index.name = "trade_date"
+        self.returns.columns.name = "group"
+
 # class QuickBackTesting02(BaseTester):
 #     def __init__(self,
 #                  env: BackTestEnv = None,