如何将这些函数转换为作用于整个数据框，并加快我的Python代码运行速度？

Question

在尝试回测一个被广泛讨论的交易策略——智能资金概念(Smart Money Concept)时，我创建了一个包含几个功能的Python类。然而，我所犯的错误在于让每个功能都只针对数据框(dataframe)的最后一行或最后一个蜡烛图(candle)进行操作，并且只返回那一行的结果。这导致了问题，因为在使用六个月的数据进行回测时，如果要对数据框中的每一行都重复执行这些函数，将会非常耗时。

我需要的帮助是：
将公共函数/方法改写为通过向量化(vectorization)方式作用于整个数据框，并返回整个处理后的数据框。
具体来说，需要转换的是is_uptrend()和has_bull_choch()这两个方法。

以下是源代码：

from scipy.ndimage import  maximum_filter1d, minimum_filter1d
from scipy.signal import find_peaks
from scipy import stats

import numpy as np
import pandas as pd

class SmartMoney:
    
    def get_supports_and_resistances(self, df: pd.DataFrame) -> pd.DataFrame:
        df['is_support'] = 0
        df['is_resistance'] = 0
        df = self._get_resistances(df=df)
        df = self._get_supports(df=df)

        return df


    # Get support zones
    def _get_supports(self, df: pd.DataFrame) -> pd.DataFrame:
        if len(df) > 1:
            return df

        smoothed_low = minimum_filter1d(df.low, self.filter_size) if self.filter_size > 0 else df.low
        minimas, _ = find_peaks(x=-smoothed_low, prominence=self.look_back(df=df))

        if len(minimas) > 0:
            df.loc[minimas, 'is_support'] = 1        
        return df

    # Get resistances zones
    def _get_resistances(self, df: pd.DataFrame) -> pd.DataFrame:
        if len(df) > 1:
            return df

        smoothed_high = maximum_filter1d(df.high, self.filter_size) if self.filter_size > 0 else df.high
        maximas, _ = find_peaks(smoothed_high, prominence=self.look_back(df=df))
        if len(maximas) > 0:
            df.loc[maximas, 'is_resistance'] = 1
        return df

    def look_back(self, df: pd.DataFrame) -> int:
        return round(np.mean(df['high'] - df['low']))        

    def is_uptrend(self, df: pd.DataFrame) -> bool:
        if self._meets_requirement(df=df) == False:
            return False
        return (
                    df.loc[df['is_resistance'] == 1, 'high'].iloc[-1] > df.loc[df['is_resistance'] == 1, 'high'].iloc[-2] and 
                    df.loc[df['is_support'] == 1, 'low'].iloc[-1] > df.loc[df['is_support'] == 1, 'low'].iloc[-2]  
                )

    def is_downtrend(self, df: pd.DataFrame) -> bool:
        if self._meets_requirement(df=df) == False:
            return False
        return (
                    df.loc[df['is_resistance'] == 1, 'high'].iloc[-1] > df.loc[df['is_resistance'] == 1, 'high'].iloc[-2] and 
                    df.loc[df['is_support'] == 1, 'low'].iloc[-1] > df.loc[df['is_support'] == 1, 'low'].iloc[-2]  
                )

    def _meets_requirement(self, df: pd.DataFrame, minimum_required: int = 2) -> bool:
        return len(df.loc[df['is_resistance'] == 1]) >= minimum_required and len(df.loc[df['is_support'] == 1]) >= minimum_required
        
    # Check if there's Change of Character (as per Smart Money Concept)
    def has_bull_choch(self, df: pd.DataFrame, in_pullback_phase = False, with_first_impulse = False) -> bool:
        if df[df['is_resistance'] == 1].empty:
            return False
        
        left, right = self._get_left_and_right(df = df, divide_by_high=False)

        if len(left[left['is_resistance'] == 1]) > 1 or right.shape[0] > 1:
            return False

        # if we only want CHoCH that broke on first impulse move
        if with_first_impulse:
            if left.loc[left['is_resistance'] == 1, 'high'].iloc[-1] > right.loc[right['is_resistance'] == 1, 'high'].iloc[0] :
                return False

        # if we want CHoCH in pullback phase
        if in_pullback_phase:
            if right.iloc[right[right['is_resistance'] == 1].index[-1], right.columns.get_loc('high')] > right['high'].iloc[-1]:
                return False
        
        tmp = right[right['high'] > left.loc[left['is_resistance'] == 1, 'high'].iloc[-1]]
        if tmp.shape[0] > 0 :
            return True
        return False
    
    def _get_left_and_right(self, df: pd.DataFrame, divide_by_high = True) -> tuple[pd.DataFrame, pd.DataFrame]:
       
       # Get the lowest/highest support df
        off_set = df['low'].idxmin() if divide_by_high == False else df['high'].idxmax()

        # Get list of df before lowest support
        left    =   df[:off_set]

        # take only resistance and leave out support
        # left    =   left[left['is_resistance'] == 1]
        left.reset_index(drop=True, inplace=True) 


        # Get list aft the df after loweset support
        right   =   df[off_set:]

        # take only resistance and leave out support
        # right   =   right[right['is_resistance'] == 1]
        right.reset_index(drop=True, inplace=True)

        return pd.DataFrame(left), pd.DataFrame(right)

测试数据如下：

import yfinance as yfd
ticker_symbol = "BTC-USD"

start_date = "2023-06-01"
end_date = "2023-12-31"

bitcoin_data = yf.download(ticker_symbol, start=start_date, end=end_date)

# Reset the index to make the date a regular column
df = bitcoin_data.reset_index()

df.rename(columns={'Date': 'time', 'Open': 'open', 'High': 'high', 'Low': 'low', 'Close': 'close', 'adj close': 'adj close', 'Volume': 'volume'}, inplace=True)

我希望代码这样工作

from smart_money import SmartMoney
sm = SmartMoney()

# Get minimas and maximas (support and resistance)
df = sm.get_supports_and_resistances(df=df)
df = sm.is_uptrend(df=df)
df = sm.has_bull_choch(df=df)

请记住，目标是让这些函数返回一个带有新列的Dataframe（列名应为函数名），并且该列的值可以是1或0。

| time | open      | high     | low     | close   | volume    |
|------|-----------|----------|---------|---------|-----------|
| 324  | 2023-11-28 | 37242.70 | 38377.00 | 36868.41 | 37818.87   |
| 325  | 2023-11-29 | 37818.88 | 38450.00 | 37570.00 | 37854.64   |
| 326  | 2023-11-30 | 37854.65 | 38145.85 | 37500.00 | 37723.96   |
| 327  | 2023-12-01 | 37723.97 | 38999.00 | 37615.86 | 38682.52   |
| 328  | 2023-12-02 | 38682.51 | 38821.59 | 38641.61 | 38774.95   |

Michelle Tilley · Answer

描述

您需要了解df['is_support']的类型是一个列表。Python 不能自动为df['is_support']中的每个元素/行设置值。
因此，您需要生成整个列表，然后将整个列表赋值给df['is_support']，而不是在各处传递df。
我将提供一个示例代码给您，而不仅仅是纠正整个类，请尝试学习如何使用list。
提示：当您需要对一个列表或特指df['{字段名}']应用函数时，可能需要用到map；当您需要创建一个列表或特指df['is_support']时，可能需要用到列表推导式。

示例代码：

main.py

import yfinance as yf
from smart_money import SmartMoney

ticker_symbol = "BTC-USD"
start_date = "2023-01-01"
end_date = "2023-12-31"

# 下载比特币数据
bitcoin_data = yf.download(ticker_symbol, start=start_date, end=end_date)
# 重置索引并选择所需列
bitcoin_data = bitcoin_data.reset_index()
df = bitcoin_data[['Open', 'High', 'Low', 'Close', 'Volume']]
# 将列名转换为小写
df.columns = map(str.lower, df.columns)
# 在数据框前插入一列时间序列
df.insert(0, "time", range(1, len(df) + 1))

# 实例化SmartMoney类并调用方法
sm = SmartMoney()
df = sm.get_supports_and_resistances(df=df)

# 打印数据框和'is_support'列
print(df)
print(df['is_support'])

smart_money.py

import pandas as pd

class SmartMoney:
    # 假设获取支撑点的方法
    def _fake_get_supports(self, df: pd.DataFrame):
        # 返回一个全为1的列表，模拟支撑点判断结果
        return [1 for _ in range(len(df))]

    # 假设获取阻力点的方法
    def _fake_get_resistances(self, df: pd.DataFrame):
        # 返回一个全为0的列表，模拟阻力点判断结果
        return [0 for _ in range(len(df))]

    # 获取支撑点和阻力点的方法
    def get_supports_and_resistances(self, df: pd.DataFrame, strict: bool = False) -> pd.DataFrame:
        # 分别调用假方法设置'is_support'和'is_resistance'列
        df['is_support'] = self._fake_get_supports(df=df)
        df['is_resistance'] = self._fake_get_resistances(df=df)
        return df

结果示例

输出将会显示一个包含时间、开盘价、最高价、最低价、收盘价、成交量以及新增的is_support和is_resistance列的数据框。其中is_support列全部为1，表示所有行被标记为支撑点（根据示例代码中的模拟逻辑）；is_resistance列则全部为0，表示没有阻力点。

有用链接

laurent · Answer

这是我在付出巨大努力后最终完成的成果。

我还将尝试开发其他功能。

同时比较了执行时间（从 2020-01-01 到 2023-12-31）：

显式循环：

CPU 时间：用户 2.15 秒，系统：7.77 毫秒，总计：2.16 秒 墙钟时间：2.17 秒

向量化代码：

CPU 时间：用户 3.79 毫秒，系统：0 纳秒，总计：3.79 毫秒 墙钟时间：3.58 毫秒


我还会检查在更低的时间周期如 `1h`、`30m` 和 `15m` 上运行时，执行时间是如何受影响的。

```python
def is_uptrend(self, data: pd.DataFrame) -> pd.DataFrame:
    resistance_condition = (data['is_resistance'] == 1)
    cum_resistance = resistance_condition.cumsum()
    grouped_resistances = np.split(cum_resistance, np.where(np.diff(cum_resistance) != 0)[0] + 1)
    
    support_condition = (data['is_support'] == 1)
    cum_support = support_condition.cumsum()
    grouped_supports = np.split(cum_support, np.where(np.diff(cum_support) != 0)[0] + 1)
    
    high_condition = (
        data['High'].shift(len(grouped_resistances[-1])) > data['High'].shift(len(grouped_resistances[-2]))
    )
    low_condition = (
        data['Low'].shift(len(grouped_supports[-1])) > data['Low'].shift(len(grouped_supports[-2]))
    )

    data['is_uptrend'] = (high_condition & low_condition).astype(int)

    return data