事件驱动回测

@njit
def cooldown_signal_nb(c, entries, exits, cooldown):  
    is_exit = vbt.pf_nb.select_nb(c, exits)
    if c.i < c.in_outputs.blocked_until[c.col]:
        return False, is_exit, False, False
    return vbt.pf_nb.select_nb(c, entries), is_exit, False, False

@njit
def cooldown_post_order_nb(c, cooldown):  
    if vbt.pf_nb.order_closed_position_nb(c):
        if vbt.pf_nb.get_last_order_nb(c)["stop_type"] >= 0:
            wait = vbt.pf_nb.select_nb(c, cooldown)
            c.in_outputs.blocked_until[c.col] = c.i + 1 + wait

data = vbt.YFData.pull("BTC-USD", start="2022-01-01", end="2024-01-01")
sma = data.close.rolling(20).mean()
pf = vbt.PF.from_signals(
    data,
    signal_func_nb=cooldown_signal_nb,
    signal_args=(vbt.Rep("entries"), vbt.Rep("exits"), vbt.Rep("cooldown")),
    post_order_func_nb=cooldown_post_order_nb,
    post_order_args=(vbt.Rep("cooldown"),),
    broadcast_named_args=dict(
        entries=data.close > sma,
        exits=data.close < sma,
        cooldown=vbt.Param([0, 10]),  
    ),
    in_outputs=dict(
        blocked_until=vbt.RepEval("np.zeros(wrapper.shape_2d[1], dtype=np.int_)")
    ),
    sl_stop=0.05,
)
pf.stats(["total_trades", "win_rate", "total_return", "max_dd"], agg_func=None)

pf[10].orders.readable[["Fill Index", "Side", "Stop Type"]].iloc[2:6]

symbols = [f"S{i:02d}" for i in range(20)]
data = vbt.GBMOHLCData.pull(symbols, start="2024-01-01", end="2025-01-01", seed=42)
sma = data.close.rolling(20).mean()

@njit
def max_positions_nb(c, entries, exits, max_positions):
    if c.in_outputs.bar_entries[0] != c.i:  
        c.in_outputs.bar_entries[:] = (c.i, 0)
    if c.last_position[c.col] > 0:
        return False, vbt.pf_nb.select_nb(c, exits), False, False
    if vbt.pf_nb.select_nb(c, entries):
        n_open = vbt.pf_nb.get_n_active_positions_nb(c)
        if n_open + c.in_outputs.bar_entries[1] < max_positions:
            c.in_outputs.bar_entries[1] += 1
            return True, False, False, False
    return False, False, False, False

pf = vbt.PF.from_signals(
    data,
    signal_func_nb=max_positions_nb,
    signal_args=(vbt.Rep("entries"), vbt.Rep("exits"), 3),
    broadcast_named_args=dict(
        entries=data.close.vbt.crossed_above(sma),
        exits=data.close.vbt.crossed_below(sma),
    ),
    in_outputs=dict(bar_entries=vbt.RepEval("np.full(2, -1, dtype=np.int_)")),
    size=0.3,
    size_type="valuepercent",
    group_by=True,  
    cash_sharing=True,
)
print((pf.assets > 0).sum(axis=1).max())

@njit
def open_to_close_nb(c, gap_down):
    col = c.from_col
    if not gap_down[c.i, col]:
        return -1, vbt.pf_enums.NoOrder
    if c.call_idx == 0:  
        return col, vbt.pf_nb.order_nb(size=1, price=c.open[c.i, col])
    if c.call_idx == 1:
        return col, vbt.pf_nb.close_position_nb(price=c.close[c.i, col])
    return -1, vbt.pf_enums.NoOrder  

data = vbt.YFData.pull("SPY", start="2024-01-01", end="2024-02-01")
gap_down = data.open < data.close.shift(1) * 0.995
pf = vbt.PF.from_order_func(
    data,
    flex_order_func_nb=open_to_close_nb,
    flex_order_args=(vbt.Rep("gap_down"),),
    broadcast_named_args=dict(gap_down=gap_down),
    max_order_records=len(data.index) * 2,  
)
pf.orders.readable[["Index", "Side", "Price"]]

@njit
def dip_buyer_nb(close, init_cash, dip, rebound):
    order_records = np.empty(close.shape, dtype=vbt.pf_enums.order_dt)
    order_counts = np.zeros(close.shape[1], dtype=np.int_)
    for col in range(close.shape[1]):
        state = vbt.pf_enums.ExecState(  
            init_cash, 0.0, 0.0, 0.0, init_cash, np.nan, np.nan
        )
        peak = close[0, col]
        for i in range(close.shape[0]):
            price = close[i, col]
            peak = max(peak, price)
            if state.position == 0 and price <= peak * (1 - dip):
                order = vbt.pf_nb.order_nb(size=np.inf, price=price)
            elif state.position > 0 and price >= peak * (1 - rebound):
                order = vbt.pf_nb.close_position_nb(price=price)
                peak = price
            else:
                continue
            _, state = vbt.pf_nb.process_order_nb(  
                col, col, i, state, order,
                order_records=order_records, order_counts=order_counts
            )
    return vbt.nb.repartition_nb(order_records, order_counts)

data = vbt.YFData.pull("BTC-USD", start="2023-01-01", end="2024-01-01")
records = dip_buyer_nb(data.close.values[:, None], 100.0, 0.1, 0.02)
pf = vbt.PF(
    data.symbol_wrapper,
    close=data.close,
    order_records=records,
    init_cash=100.0
)
pf.orders.readable[["Index", "Side", "Price"]]

DCAMode = namedtuple("DCAMode", ["In", "Out"])(0, 1)  

@njit
def pre_group_func_nb(ctx):  
    total_deposited = np.full(1, 0.0)
    dca_mode = np.full(1, DCAMode.In)
    return (total_deposited, dca_mode)

@njit
def pre_segment_func_nb(ctx, total_deposited, dca_mode, cash_deposits, dca_amount):  
    if ctx.i == 0 or vbt.dt_nb.month_nb(ctx.index[ctx.i - 1]) != vbt.dt_nb.month_nb(ctx.index[ctx.i]):
        dca_amount_now = vbt.pf_nb.select_from_group_nb(ctx, ctx.group, dca_amount)
        if dca_mode[0] == DCAMode.In and ctx.track_cash_deposits:
            cash_deposits[ctx.i, ctx.group] = dca_amount_now
            total_deposited[0] += dca_amount_now
    else:
        dca_amount_now = 0.0
    return (total_deposited, dca_mode, dca_amount_now)

@njit
def signal_func_nb(ctx, total_deposited, dca_mode, dca_amount_now, size):  
    if dca_amount_now > 0:
        size[ctx.i, ctx.col] = dca_amount_now
        if dca_mode[0] == DCAMode.In:
            return True, False, False, False
        return False, True, False, False
    return False, False, False, False

@njit
def post_order_func_nb(ctx, total_deposited, dca_mode, dca_amount_now):  
    if dca_mode[0] == DCAMode.In:
        if vbt.pf_nb.order_increased_position_nb(ctx):
            tp_info = ctx.last_tp_info[ctx.col]
            tp_info["stop"] = 1.0
            tp_info["init_price"] = ctx.last_pos_info[ctx.col]["entry_price"]
            tp_info["exit_size"] = total_deposited[0]
            tp_info["exit_size_type"] = vbt.pf_enums.SizeType.Value
        elif vbt.pf_nb.get_last_order_nb(ctx)["stop_type"] == vbt.pf_enums.StopType.TP:
            dca_mode[0] = DCAMode.Out

pf = vbt.PF.from_signals(
    vbt.YFData.pull("AAPL", start="2018"),
    pre_group_func_nb=pre_group_func_nb,
    pre_segment_func_nb=pre_segment_func_nb,
    pre_segment_args=(
        vbt.Rep("cash_deposits"),
        vbt.Rep("dca_amount")
    ),
    signal_func_nb=signal_func_nb,
    signal_args=(
        vbt.Rep("size"),
    ),
    post_order_func_nb=post_order_func_nb,
    broadcast_named_args=dict(dca_amount=100),
    arg_config=dict(
        cash_deposits=dict(full_shape=True),  
        size=dict(full_shape=True)
    ),
    accumulate=True,
    size_type="value",
    cash_sharing=True
)
pf.plot_orders().show()

data = vbt.YFData.pull(["BTC-USD", "ETH-USD"], missing_index="drop")
size = data.symbol_wrapper.fill(np.nan)
np.random.seed(42)
rand_indices = np.random.choice(np.arange(len(size)), 10)
size.iloc[rand_indices[0::2]] = -np.inf
size.iloc[rand_indices[1::2]] = np.inf

@njit
def post_segment_func_nb(ctx):
    for col in range(ctx.from_col, ctx.to_col):
        col_debt = ctx.last_debt[col]
        ctx.in_outputs.debt[ctx.i, col] = col_debt
        if col_debt > ctx.in_outputs.max_debt[col]:
            ctx.in_outputs.max_debt[col] = col_debt

pf = vbt.PF.from_def_order_func(
    data.close,
    size=size,
    post_segment_func_nb=post_segment_func_nb,
    in_outputs=dict(
        debt=vbt.RepEval("np.empty_like(close)"),
        max_debt=vbt.RepEval("np.full(close.shape[1], 0.)")
    )  
)
pf.get_in_output("debt")  

pf.get_in_output("max_debt")  

InOutputs = namedtuple("InOutputs", ["fast_sma", "slow_sma"])

def initialize_in_outputs(target_shape):
    return InOutputs(
        fast_sma=np.full(target_shape, np.nan),
        slow_sma=np.full(target_shape, np.nan)
    )

@njit
def signal_func_nb(ctx, fast_window, slow_window):
    fast_sma = ctx.in_outputs.fast_sma
    slow_sma = ctx.in_outputs.slow_sma
    fast_start_i = ctx.i - fast_window + 1
    slow_start_i = ctx.i - slow_window + 1
    if fast_start_i >= 0 and slow_start_i >= 0:
        fast_sma[ctx.i, ctx.col] = np.nanmean(ctx.close[fast_start_i : ctx.i + 1])
        slow_sma[ctx.i, ctx.col] = np.nanmean(ctx.close[slow_start_i : ctx.i + 1])
        is_entry = vbt.pf_nb.iter_crossed_above_nb(ctx, fast_sma, slow_sma)
        is_exit = vbt.pf_nb.iter_crossed_below_nb(ctx, fast_sma, slow_sma)
        return is_entry, is_exit, False, False
    return False, False, False, False

pf = vbt.PF.from_signals(
    vbt.YFData.pull("BTC-USD"),
    signal_func_nb=signal_func_nb,
    signal_args=(50, 200),
    in_outputs=vbt.RepFunc(initialize_in_outputs),
)
fig = pf.get_in_output("fast_sma").vbt.plot()
pf.get_in_output("slow_sma").vbt.plot(fig=fig)
pf.orders.plot(plot_ohlc=False, plot_close=False, fig=fig)
fig.show()