Python计算Heikin Ashi蜡烛线并返回带指定列DataFrame的最优实现咨询

现有代码的错误点 #

• 返回值接收不匹配：heikinashi函数返回5个对象（原dataframe + 4个计算结果），但主脚本仅声明4个变量接收，直接触发结构异常
• 执行逻辑错误：函数返回的ha_close等是已经计算完成的Series对象，不是可调用函数，不需要再加(dataframe)执行
• 极值计算规则错误：
  1. Heikin Ashi的ha_low应为三个值的最小值、ha_high应为最大值，现有代码逻辑写反
  2. 内置max()/min()无法对Series做逐元素计算，需使用Pandas提供的逐元素极值方法
• 边界值缺失：首行的ha_open没有初始值，标准规则为首行ha_open等于原K线的open值

修正后的实现代码 #

indicators.py文件 #

from pandas import DataFrame
import pandas as pd

def heikinashi(dataframe: DataFrame) -> DataFrame:
    # 提取原K线数据
    open_ = dataframe['open']
    high = dataframe['high']
    low = dataframe['low']
    close = dataframe['close']
    
    # 计算HA指标
    ha_close = 0.25 * (open_ + high + low + close)
    # 首行ha_open赋值为原开盘价，后续行用前一天的开收盘均值
    ha_open = pd.Series(0.0, index=dataframe.index)
    ha_open.iloc[0] = open_.iloc[0]
    ha_open.iloc[1:] = 0.5 * (open_.shift(1) + close.shift(1)).iloc[1:]
    
    # 逐元素计算极值
    ha_low = pd.concat([low, ha_open, ha_close], axis=1).min(axis=1)
    ha_high = pd.concat([high, ha_open, ha_close], axis=1).max(axis=1)
    
    # 直接将结果写入原DataFrame返回
    dataframe['ha_close'] = ha_close
    dataframe['ha_open'] = ha_open
    dataframe['ha_low'] = ha_low
    dataframe['ha_high'] = ha_high
    
    return dataframe

主脚本调用代码 #

import indicators as ata

# 仅需调用1次函数，直接得到带HA列的DataFrame
dataframe = ata.heikinashi(dataframe)

方案优势 #

• 仅需1次函数调用即可完成所有计算和列写入，调用成本最低
• 全部使用Pandas向量化计算，无循环遍历，运行效率最高
• 符合标准Heikin Ashi蜡烛线计算规则，结果准确

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