自定义数据

PyBroker 自带了为 Yahoo Finance、Alpaca 和 AKShare 构建的 DataSource，你可以立即使用。如果你有特定需求或想使用其他数据提供商，你可以创建自己的 DataSource 类。

扩展 DataSource

以下示例实现了一个自定义的 CSVDataSource，用于从 CSV 文件加载数据。它将名为 prices.csv 的文件读取到 Pandas 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 返回数据时，必须包含 PyBroker 所需的列：symbol、date、open、high、low 和 close。

现在我们可以从 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:00

[3]:

	type	symbol	date	created	order_type	intent	shares	limit_price	market_price	fill_price	fees
id
1	buy	NKE	2021-09-21	2021-09-20	market	buy_to_open	100	NaN	154.86	154.86	0.0
2	sell	NKE	2021-11-04	2021-11-03	market	sell_to_close	100	NaN	173.82	173.82	0.0
3	buy	DIS	2021-11-16	2021-11-15	market	buy_to_open	100	NaN	159.40	159.40	0.0

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

使用 Pandas DataFrame

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

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

[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	created	order_type	intent	shares	limit_price	market_price	fill_price	fees
id
1	buy	NKE	2021-09-21	2021-09-20	market	buy_to_open	100	NaN	154.86	154.86	0.0
2	sell	NKE	2021-11-04	2021-11-03	market	sell_to_close	100	NaN	173.82	173.82	0.0
3	buy	DIS	2021-11-16	2021-11-15	market	buy_to_open	100	NaN	159.40	159.40	0.0