如何在numpy.select函数中设置条件优先级，添加最高优先级Stopped Out判定

问题描述 #

示例数据集 #

Price  SL    X  
14    13.8  100
14.5  13.8   0
15    13.8   0
14.7  13.8   0
13.6  13.8   0
15    13.8   0

原始需求 #

基于X列生成备注列，规则如下：

• X等于100时标注BUY
• BUY后续3行标注HOLD
• 再之后2行标注SELL

原有实现代码 #

cond = [   
    (df['X'] == 100),
    (df['X'].shift(1) == 100),
    (df['X'].shift(2) == 100),
    (df['X'].shift(3) == 100),
    (df['X'].shift(4) == 100),
    (df['X'].shift(4) == 100),
]

choices = ['BUY', 'HOLD', 'HOLD', 'HOLD', 'SELL','SELL']
df['remarks'] = np.select(cond, choices)

原有输出结果 #

Price  SL    X    remarks
14    13.8  100    BUY
14.5  13.8   0     HOLD
15    13.8   0     HOLD
14.7  13.8   0     HOLD
13.6  13.8   0     SELL
15    13.8   0     SELL

新增规则 #

1. 当Price < SL时标注Stopped Out，优先级高于所有已有条件
2. 出现Stopped Out后后续备注终止生成
3. 该规则仅在BUY出现后生效

预期输出 #

Price  SL    X    remarks
14    13.8  100    BUY
14.5  13.8   0     HOLD
15    13.8   0     HOLD
14.7  13.8   0     HOLD
13.6  13.8   0     Stopped Out
15    13.8   0     Stopped Out
14.5  13.8   0     
15    14.7  100    BUY
14.8  14.7   0     HOLD
13    14.7   0      Stopped Out

实现方案 #

np.select本身支持按条件顺序判断优先级，我们把止损条件放在最前面，再加分组逻辑隔离不同BUY的交易周期即可，完整代码如下：

import pandas as pd
import numpy as np

# 构造测试数据
data = [
    [14, 13.8, 100],
    [14.5, 13.8, 0],
    [15, 13.8, 0],
    [14.7, 13.8, 0],
    [13.6, 13.8, 0],
    [15, 13.8, 0],
    [14.5, 13.8, 0],
    [15, 14.7, 100],
    [14.8, 14.7, 0],
    [13, 14.7, 0]
]
df = pd.DataFrame(data, columns=['Price', 'SL', 'X'])

# 第一步：生成基础BUY/HOLD/SELL标记（修正原代码重复shift(4)的问题）
base_cond = [
    (df['X'] == 100),
    (df['X'].shift(1) == 100),
    (df['X'].shift(2) == 100),
    (df['X'].shift(3) == 100),
    (df['X'].shift(4) == 100),
    (df['X'].shift(5) == 100),
]
base_choices = ['BUY', 'HOLD', 'HOLD', 'HOLD', 'SELL','SELL']
df['base_remark'] = np.select(base_cond, base_choices, default='')

# 第二步：计算止损状态
# 按每次BUY分组，隔离不同交易周期
df['trade_group'] = (df['X'] == 100).cumsum()
# 标记当前行是否满足止损条件
df['is_stop'] = (df['Price'] < df['SL']) & (df['base_remark'] != 'BUY')
# 标记当前交易周期是否已触发止损
df['stop_triggered'] = df.groupby('trade_group')['is_stop'].cummax()
# 统计止损触发后的行数，超过2行就清空备注
df['stop_count'] = df.groupby('trade_group')['stop_triggered'].cumsum()

# 第三步：按优先级生成最终备注
final_cond = [
    df['stop_triggered'] & (df['stop_count'] <=2),
    ~df['stop_triggered']
]
final_choices = [
    'Stopped Out',
    df['base_remark']
]
df['remarks'] = np.select(final_cond, final_choices, default='')

# 清除辅助列
df = df.drop(columns=['base_remark', 'trade_group', 'is_stop', 'stop_triggered', 'stop_count'])
print(df)

内容的提问来源于stack exchange，提问作者Struggling brat