如何用numpy.where实现列依赖的交易信号计算（无需循环）

我有一个仅含z_score列的小型DataFrame，需要生成sell、position、short_exit三列。

原始数据如下：

z_score
2024-05-23 21:00:00  0.639482
2024-05-23 22:00:00  1.133366
2024-05-23 23:00:00 -0.267677
2024-05-24 00:00:00  0.681130
2024-05-24 01:00:00  2.283615
2024-05-24 02:00:00  1.916469
2024-05-24 03:00:00  2.284099
2024-05-24 04:00:00  2.283456
2024-05-24 05:00:00  2.699485
2024-05-24 06:00:00  2.393399
2024-05-24 07:00:00  1.871835
2024-05-24 08:00:00  1.824887

交易规则：

• 当z_score大于2且当前无持仓（position为0）时，sell列设为1，position列设为1（表示持仓）并保持该状态
• 直到z_score低于0时平仓，此时short_exit列设为1，position列重置为0

原代码问题 #

我尝试用以下代码实现但未成功，position列出现不应为0的错误值：

import pandas as pd
import numpy as np

df = pd.read_excel("gate.xlsx", sheet_name="Sheet3")
# 初始化sell、positions、short_exit列为0
df['sell'] = 0
df['positions'] = 0
df['short_exit'] = 0

df['sell'] = np.where((df['z_score'] > 2) & (df['positions'].shift(1) == 0), 1, 0)
df['positions'] = np.where(((df['sell'] == 1) | (df['positions'].shift(1) == 1)) & (df['short_exit'] == 0), 1, df['positions'])
df['short_exit'] = np.where((df['z_score'] < 0) & (df['positions'] == 1), 1 ,0)

# 输出结果DataFrame
print(df)

同时有疑问：有人说这类计算需要递归，无法用numpy实现，请问是否真的如此？

你的代码问题在于列的计算顺序和依赖关系处理错误：numpy的np.where是向量化计算，会一次性对整列求值，无法处理逐行的状态依赖（比如position的当前值依赖上一行的状态，而short_exit又依赖当前行的position）。

方案1：逐行迭代（直观易维护） #

因为持仓状态是累计依赖的，逐行迭代是最直接的实现方式，逻辑清晰不易出错：

import pandas as pd

df = pd.read_excel("gate.xlsx", sheet_name="Sheet3")
# 初始化列
df['sell'] = 0
df['position'] = 0
df['short_exit'] = 0

# 记录当前持仓状态
current_position = 0

for idx, row in df.iterrows():
    if current_position == 0:
        # 无持仓时检查开仓条件
        if row['z_score'] > 2:
            df.loc[idx, 'sell'] = 1
            current_position = 1
            df.loc[idx, 'position'] = 1
    else:
        # 有持仓时保持状态，检查平仓条件
        df.loc[idx, 'position'] = 1
        if row['z_score'] < 0:
            df.loc[idx, 'short_exit'] = 1
            current_position = 0

print(df)

方案2：向量化实现（高效无迭代） #

如果想避免逐行迭代提升效率，可以用pandas的cumsum和布尔掩码构造持仓区间，实现无迭代的向量化计算：

import pandas as pd

df = pd.read_excel("gate.xlsx", sheet_name="Sheet3")

# 1. 标记初始开仓信号：z_score>2且之前未进入持仓
df['sell'] = (df['z_score'] > 2) & (df['z_score'].shift(1).le(2).fillna(True))

# 2. 计算持仓状态：用开仓与平仓信号的累计差判断是否持仓
open_signals = df['sell'].cumsum()
close_signals = df['z_score'].lt(0).cumsum()
df['position'] = ((open_signals - close_signals) > 0).astype(int)

# 3. 标记平仓信号：z_score<0且上一行处于持仓状态
df['short_exit'] = (df['z_score'] < 0) & (df['position'].shift(1) == 1)
df['short_exit'] = df['short_exit'].astype(int)

# 修正sell列：仅保留无持仓时的有效开仓信号
df['sell'] = df['sell'] & (df['position'].shift(1).fillna(0) == 0)
df['sell'] = df['sell'].astype(int)

print(df)

这类计算不需要递归，但确实无法用单纯的numpy向量化直接实现——因为numpy的向量化操作是无状态的，无法处理“当前值依赖上一行结果”的累计状态。不过pandas提供了cumsum、shift等工具，可以通过构造状态区间的方式实现向量化计算，避免逐行迭代。

逐行迭代逻辑直观，但数据量大时效率较低；向量化方案效率更高，但逻辑稍复杂，可以根据数据规模选择合适的方式。

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