How can I optimize this Python Pandas code?

Viewed 50

The following function takes market data (OHLCV candles) and tries to find periods where the price oscillates between two bounds (consolidation zones). I wrote it by translating, from Pine Script to Python, an open-source indicator found on TradingView.

The function works, in the sense that it finds correctly consolidation zones. But problem is performance, mainly due to the for cycle at the end: with 30K candles it takes just ~5 seconds to execute the code before the for, and then it takes over 2 minutes to run the cycle.

import numpy as np
import pandas as pd
from pandas import (DataFrame, Series)

def _find_zz(row: Series):
    if pd.notnull(row['hb']) and pd.notnull(row['lb']):
        if row['dir'] == 1:
            return row['hb']
        else:
            return row['lb']
    else:
        return row['hb'] if pd.notnull(row['hb']) else row['lb'] if pd.notnull(row['lb']) else np.NaN

def consolidation_zones(dataframe: DataFrame, timeperiod: int = 100,
        minlength: int = 20) -> DataFrame:
    rolling = dataframe.rolling(timeperiod, min_periods=1)
    idxmax = rolling['high'].apply(lambda x: x.idxmax()).astype(int)
    idxmin = rolling['low'].apply(lambda x: x.idxmin()).astype(int)
    highest = pd.concat({'value': dataframe['high'], 'offset': dataframe.index - idxmax}, axis=1)
    lowest = pd.concat({'value': dataframe['low'], 'offset': dataframe.index - idxmin}, axis=1)
    hb = highest.apply(lambda x: x['value'] if x['offset'] == 0 else np.NaN, axis=1)
    lb = lowest.apply(lambda x: x['value'] if x['offset'] == 0 else np.NaN, axis=1)
    direction = pd.concat({'hb': hb, 'lb': lb}, axis=1).apply(lambda x: 1 if pd.notnull(x['hb']) and pd.isnull(x['lb']) else -1 if pd.isnull(x['hb']) and pd.notnull(x['lb']) else np.NaN, axis=1).fillna(method='ffill').fillna(0).astype(int)
    zz = pd.concat({'hb': hb, 'lb': lb, 'dir': direction}, axis=1).apply(_find_zz, axis=1)

    group = direction.ne(direction.shift()).cumsum()
    zzdir = pd.concat({'zz': zz, 'dir': direction, 'group': group}, axis=1)
    zzdir['min'] = zzdir.groupby('group')['zz'].cummin().fillna(method='ffill')
    zzdir['max'] = zzdir.groupby('group')['zz'].cummax().fillna(method='ffill')
    zzdir['pp'] = np.NaN
    pp = Series(np.where(zzdir['dir'] == 1, zzdir['max'], np.where(zzdir['dir'] == -1, zzdir['min'], zzdir['pp'])))

    H = dataframe.rolling(minlength, min_periods=1)['high'].max()
    L = dataframe.rolling(minlength, min_periods=1)['low'].min()
    prevpp = np.NaN
    conscnt = 0
    condhigh = np.NaN
    condlow = np.NaN
    zones = DataFrame(index=dataframe.index, columns=['upper_bound', 'lower_bound'])
    indexes = [] # will keep indexes of candles that are part of the consolidation
#----------------
    for index, value in pp.items():
        # pp is a value computed before: when it changes, it *may* be the end of a consolidation zone
        if value != prevpp:
            if conscnt > 0 and value <= condhigh and value >= condlow:
                # if condlow <= pp <= condhigh, we are still in consolidation
                conscnt = conscnt + 1
                indexes.append(index)
            else: # end of consolidation
                conscnt = 0
                indexes = []
        else:
            conscnt = conscnt + 1
            indexes.append(index)

        if conscnt >= minlength:
            if conscnt == minlength:
                # initially, condhigh/low is equal to the highest/lowest value in last minlength candles
                condhigh = H.get(index)
                condlow = L.get(index)
            else:
                # update condhigh/low with new high/low
                condhigh = max(condhigh, dataframe.loc[index, 'high'])
                condlow = min(condlow, dataframe.loc[index, 'low'])
            zones.loc[zones.index.isin(indexes), 'upper_bound'] = condhigh
            zones.loc[zones.index.isin(indexes), 'lower_bound'] = condlow
        prevpp = value
#----------------

    return zones

I don't know how to write the last part of the code (delimited by comments) without iterating over all the rows.

This is the original Pine Script: Consolidation Zones - Live - TradingView

0 Answers
Related