自定义数据

PyBroker 自带了为 Yahoo Finance、Alpaca 和 AKShare 构建的 DataSource，你可以立即使用，无需进行任何额外设置。但是，如果你有特定需求或想使用其他数据源，PyBroker 也允许你创建自己的 DataSource 类。

扩展 DataSource

在下面提供的示例代码中，实现了一个名为 CSVDataSource 的新数据源，它从 CSV 文件加载数据。CSVDataSource 读取名为 prices.csv 的文件到 Pandas DataFrame 中，然后根据提供的输入参数从这个 DataFrame 返回数据：

[1]:

import pandas as pd
import pybroker
from pybroker.data import DataSource

class CSVDataSource(DataSource):

    def __init__(self):
        super().__init__()
        # Register custom columns in the CSV.
        pybroker.register_columns('rsi')

    def _fetch_data(self, symbols, start_date, end_date, _timeframe, _adjust):
        df = pd.read_csv('data/prices.csv')
        df = df[df['symbol'].isin(symbols)]
        df['date'] = pd.to_datetime(df['date'])
        return df[(df['date'] >= start_date) & (df['date'] <= end_date)]

为了使 CSV 文件中自定义的 'rsi' 列对 PyBroker 可用，我们使用 pybroker.register_columns 进行注册。这允许 PyBroker 在处理数据时使用这个自定义列。

需要注意的是，在从自定义 DataSource 返回数据时，必须包含以下列：symbol、date、open、high、low 和 close，因为 PyBroker 需要这些列。

现在我们可以从 CSVDataSource 的一个实例查询 CSV 数据：

[2]:

csv_data_source = CSVDataSource()
df = csv_data_source.query(['MCD', 'NKE', 'DIS'], '6/1/2021', '12/1/2021')
df

Loading bar data...
Loaded bar data: 0:00:00

[2]:

	date	symbol	open	high	low	close	rsi
0	2021-06-01	DIS	180.179993	181.009995	178.740005	178.839996	46.321532
1	2021-06-01	MCD	235.979996	235.990005	232.740005	233.240005	46.522926
2	2021-06-01	NKE	137.850006	138.050003	134.210007	134.509995	53.308085
3	2021-06-02	DIS	179.039993	179.100006	176.929993	177.000000	42.635256
4	2021-06-02	MCD	233.970001	234.330002	232.809998	233.779999	48.051484
...	...	...	...	...	...	...	...
382	2021-11-30	MCD	247.380005	247.899994	243.949997	244.600006	40.461178
383	2021-11-30	NKE	168.789993	171.550003	167.529999	169.240005	51.505558
384	2021-12-01	DIS	146.699997	148.369995	142.039993	142.149994	16.677555
385	2021-12-01	MCD	245.759995	250.899994	244.110001	244.179993	39.853689
386	2021-12-01	NKE	170.889999	173.369995	166.679993	166.699997	46.704527

387 rows × 7 columns

要在回测中使用 CSVDataSource，我们创建一个新的 Strategy 对象，并传入自定义的 DataSource：

[3]:

from pybroker import Strategy

def buy_low_sell_high_rsi(ctx):
    pos = ctx.long_pos()
    if not pos and ctx.rsi[-1] < 30:
        ctx.buy_shares = 100
    elif pos and ctx.rsi[-1] > 70:
        ctx.sell_shares = pos.shares

strategy = Strategy(csv_data_source, '6/1/2021', '12/1/2021')
strategy.add_execution(buy_low_sell_high_rsi, ['MCD', 'NKE', 'DIS'])
result = strategy.backtest()
result.orders

Backtesting: 2021-06-01 00:00:00 to 2021-12-01 00:00:00

Loading bar data...
Loaded bar data: 0:00:00

Test split: 2021-06-01 00:00:00 to 2021-12-01 00:00:00

100% (129 of 129) |######################| Elapsed Time: 0:00:00 Time:  0:00:00


Finished backtest: 0:00:02

[3]:

	type	symbol	date	shares	limit_price	fill_price	fees
id
1	buy	NKE	2021-09-21	100	NaN	154.86	0.0
2	sell	NKE	2021-11-04	100	NaN	173.82	0.0
3	buy	DIS	2021-11-16	100	NaN	159.40	0.0

请注意，因为我们已经使用 PyBroker 注册了自定义的 rsi 列，所以可以在 ExecContext 中使用 ctx.rsi 来访问它。

使用 Pandas DataFrame

如果你不需要实现自己的 DataSource 的灵活性，那么可以将 Pandas DataFrame 传递给策略。

为了演示，可以按照以下方式重新实现前面的示例：

[4]:

df = pd.read_csv('data/prices.csv')
df['date'] = pd.to_datetime(df['date'])

pybroker.register_columns('rsi')

strategy = Strategy(df, '6/1/2021', '12/1/2021')
strategy.add_execution(buy_low_sell_high_rsi, ['MCD', 'NKE', 'DIS'])
result = strategy.backtest()
result.orders

Backtesting: 2021-06-01 00:00:00 to 2021-12-01 00:00:00

Test split: 2021-06-01 00:00:00 to 2021-12-01 00:00:00

100% (129 of 129) |######################| Elapsed Time: 0:00:00 Time:  0:00:00


Finished backtest: 0:00:00

[4]:

	type	symbol	date	shares	limit_price	fill_price	fees
id
1	buy	NKE	2021-09-21	100	NaN	154.86	0.0
2	sell	NKE	2021-11-04	100	NaN	173.82	0.0
3	buy	DIS	2021-11-16	100	NaN	159.40	0.0