| 93 | self.datas = pd.DataFrame() |
| 94 | |
| 95 | def load_data( |
| 96 | self, |
| 97 | data: pd.DataFrame, |
| 98 | start_dt: str = None, |
| 99 | end_dt: str = None, |
| 100 | datafeed_cls: bt.feeds.PandasData = None, |
| 101 | ) -> None: |
| 102 | |
| 103 | if start_dt is not None: |
| 104 | data = data.loc[pd.to_datetime(start_dt):] |
| 105 | |
| 106 | if end_dt is not None: |
| 107 | data = data.loc[:pd.to_datetime(end_dt)] |
| 108 | |
| 109 | if (start_dt is None) and (end_dt is None): |
| 110 | start_dt, end_dt = data.index.min(), data.index.max() |
| 111 | data = data.loc[start_dt:end_dt] |
| 112 | |
| 113 | self.datas = data |
| 114 | |
| 115 | logger.info("开始加载数据...") |
| 116 | |
| 117 | for code, df in track( |
| 118 | data.groupby("code"), |
| 119 | desc="数据加载到回测引擎..." |
| 120 | ): |
| 121 | df = df.drop(columns=["code"]) |
| 122 | df = check_dataframe_cols(df, datafeed_cls) |
| 123 | assert "factor1" in df.columns, "factor1 列缺失" |
| 124 | assert "factor2" in df.columns, "factor2 列缺失" |
| 125 | assert "signal" in df.columns, "signal 列缺失" |
| 126 | assert "vwap" in df.columns, "vwap 列缺失" |
| 127 | |
| 128 | if df["close"].dropna().empty: |
| 129 | logger.warning(f"{code} close 全为 NaN，跳过...") |
| 130 | continue |
| 131 | |
| 132 | datafeed = datafeed_cls( |
| 133 | dataname=df.sort_index() |
| 134 | ) |
| 135 | self.cerebro.adddata(datafeed, name=code) |
| 136 | |
| 137 | logger.success("数据加载完毕！") |
| 138 | |
| 139 | def add_strategy(self, strategy: bt.Strategy, *args, **kwargs) -> None: |
| 140 | self.cerebro.addstrategy(strategy, *args, **kwargs) |