SuperTrend code using pandas python

Viewed 16047

I am trying to code the following algorithm for SuperTrend indicator in python using pandas.

BASIC UPPERBAND = (HIGH + LOW) / 2 + Multiplier * ATR
BASIC LOWERBAND = (HIGH + LOW) / 2 - Multiplier * ATR

FINAL UPPERBAND = IF( (Current BASICUPPERBAND < Previous FINAL UPPERBAND) or (Previous Close > Previous FINAL UPPERBAND))
                    THEN (Current BASIC UPPERBAND) ELSE Previous FINALUPPERBAND)
FINAL LOWERBAND = IF( (Current BASIC LOWERBAND > Previous FINAL LOWERBAND) or (Previous Close < Previous FINAL LOWERBAND)) 
                    THEN (Current BASIC LOWERBAND) ELSE Previous FINAL LOWERBAND)

SUPERTREND = IF((Previous SUPERTREND = Previous FINAL UPPERBAND) and (Current Close <= Current FINAL UPPERBAND)) THEN
                Current FINAL UPPERBAND
            ELSE
                IF((Previous SUPERTREND = Previous FINAL UPPERBAND) and (Current Close > Current FINAL UPPERBAND)) THEN
                    Current FINAL LOWERBAND
                ELSE
                    IF((Previous SUPERTREND = Previous FINAL LOWERBAND) and (Current Close >= Current FINAL LOWERBAND)) THEN
                        Current FINAL LOWERBAND
                    ELSE
                        IF((Previous SUPERTREND = Previous FINAL LOWERBAND) and (Current Close < Current FINAL LOWERBAND)) THEN
                            Current FINAL UPPERBAND

Here is the code that I wrote and tested:

# Compute basic upper and lower bands
df['basic_ub'] = (df['high'] + df['low']) / 2 + multiplier * df[atr]
df['basic_lb'] = (df['high'] + df['low']) / 2 - multiplier * df[atr]

# Compute final upper and lower bands
for i in range(0, len(df)):
    if i < period:
        df.set_value(i, 'basic_ub', 0.00)
        df.set_value(i, 'basic_lb', 0.00)
        df.set_value(i, 'final_ub', 0.00)
        df.set_value(i, 'final_lb', 0.00)
    else:
        df.set_value(i, 'final_ub', (df.get_value(i, 'basic_ub') 
                                     if df.get_value(i, 'basic_ub') < df.get_value(i-1, 'final_ub') or df.get_value(i-1, 'close') > df.get_value(i-1, 'final_ub') 
                                     else df.get_value(i-1, 'final_ub')))
        df.set_value(i, 'final_lb', (df.get_value(i, 'basic_lb') 
                                     if df.get_value(i, 'basic_lb') > df.get_value(i-1, 'final_lb') or df.get_value(i-1, 'close') < df.get_value(i-1, 'final_lb') 
                                     else df.get_value(i-1, 'final_lb')))

# Set the Supertrend value
for i in range(0, len(df)):
    if i < period:
        df.set_value(i, st, 0.00)
    else:
        df.set_value(i, 'st', (df.get_value(i, 'final_ub')
                             if ((df.get_value(i-1, 'st') == df.get_value(i-1, 'final_ub')) and (df.get_value(i, 'close') <= df.get_value(i, 'final_ub')))
                             else (df.get_value(i, 'final_lb')
                                   if ((df.get_value(i-1, 'st') == df.get_value(i-1, 'final_ub')) and (df.get_value(i, 'close') > df.get_value(i, 'final_ub')))
                                   else (df.get_value(i, 'final_lb')
                                         if ((df.get_value(i-1, 'st') == df.get_value(i-1, 'final_lb')) and (df.get_value(i, 'close') >= df.get_value(i, 'final_lb')))
                                         else (df.get_value(i, 'final_ub')
                                               if((df.get_value(i-1, 'st') == df.get_value(i-1, 'final_lb')) and (df.get_value(i, 'close') < df.get_value(i, 'final_lb')))
                                               else 0.00
                                              )
                                        )
                                  ) 
                            )
                    )


# Mark the trend direction up/down
df['stx'] = np.where((df['st'] > 0.00), np.where((df['close'] < df['st']), 'down',  'up'), np.NaN)

I works, but I am not happy with the for loop. Can anyone help optimise it?

You can find the released code on Github!

2 Answers

SuperTrend Indicator is included in pandas_ta where you can simply:

import pandas_ta as ta

sti = ta.supertrend(df['High'], df['Low'], df['Close'], 7, 3)

given that df is a pandas DataFrame with OHLC prices. The result sti is a DataFrame with 4 columns:

  1. trend
  2. direction
  3. long
  4. short

where the trend is a concatenation of the long and short lines. Note that column captions are dynamic and contain the length and multiplier parameter values.

Use this complete code

import numpy as np
import yfinance as yf
import pandas_datareader as pdr
import pandas as pd

data =yf.download("ACC.NS", period="1d",interval="5m")
data=data.reset_index(drop=True)

data['tr0'] = abs(data["High"] - data["Low"])
data['tr1'] = abs(data["High"] - data["Close"].shift(1))
data['tr2'] = abs(data["Low"]- data["Close"].shift(1))
data["TR"] = round(data[['tr0', 'tr1', 'tr2']].max(axis=1),2)
data["ATR"]=0.00
data['BUB']=0.00
data["BLB"]=0.00
data["FUB"]=0.00
data["FLB"]=0.00
data["ST"]=0.00

# Calculating ATR 
for i, row in data.iterrows():
    if i == 0:
        data.loc[i,'ATR'] = 0.00#data['ATR'].iat[0]
    else:
        data.loc[i,'ATR'] = ((data.loc[i-1,'ATR'] * 13)+data.loc[i,'TR'])/14

data['BUB'] = round(((data["High"] + data["Low"]) / 2) + (2 * data["ATR"]),2)
data['BLB'] = round(((data["High"] + data["Low"]) / 2) - (2 * data["ATR"]),2)


# FINAL UPPERBAND = IF( (Current BASICUPPERBAND < Previous FINAL UPPERBAND) or (Previous Close > Previous FINAL UPPERBAND))
#                     THEN (Current BASIC UPPERBAND) ELSE Previous FINALUPPERBAND)


for i, row in data.iterrows():
    if i==0:
        data.loc[i,"FUB"]=0.00
    else:
        if (data.loc[i,"BUB"]<data.loc[i-1,"FUB"])|(data.loc[i-1,"Close"]>data.loc[i-1,"FUB"]):
            data.loc[i,"FUB"]=data.loc[i,"BUB"]
        else:
            data.loc[i,"FUB"]=data.loc[i-1,"FUB"]

# FINAL LOWERBAND = IF( (Current BASIC LOWERBAND > Previous FINAL LOWERBAND) or (Previous Close < Previous FINAL LOWERBAND)) 
#                     THEN (Current BASIC LOWERBAND) ELSE Previous FINAL LOWERBAND)

for i, row in data.iterrows():
    if i==0:
        data.loc[i,"FLB"]=0.00
    else:
        if (data.loc[i,"BLB"]>data.loc[i-1,"FLB"])|(data.loc[i-1,"Close"]<data.loc[i-1,"FLB"]):
            data.loc[i,"FLB"]=data.loc[i,"BLB"]
        else:
            data.loc[i,"FLB"]=data.loc[i-1,"FLB"]



# SUPERTREND = IF((Previous SUPERTREND = Previous FINAL UPPERBAND) and (Current Close <= Current FINAL UPPERBAND)) THEN
#                 Current FINAL UPPERBAND
#             ELSE
#                 IF((Previous SUPERTREND = Previous FINAL UPPERBAND) and (Current Close > Current FINAL UPPERBAND)) THEN
#                     Current FINAL LOWERBAND
#                 ELSE
#                     IF((Previous SUPERTREND = Previous FINAL LOWERBAND) and (Current Close >= Current FINAL LOWERBAND)) THEN
#                         Current FINAL LOWERBAND
#                     ELSE
#                         IF((Previous SUPERTREND = Previous FINAL LOWERBAND) and (Current Close < Current FINAL LOWERBAND)) THEN
#                             Current FINAL UPPERBAND


for i, row in data.iterrows():
    if i==0:
        data.loc[i,"ST"]=0.00
    elif (data.loc[i-1,"ST"]==data.loc[i-1,"FUB"]) & (data.loc[i,"Close"]<=data.loc[i,"FUB"]):
        data.loc[i,"ST"]=data.loc[i,"FUB"]
    elif (data.loc[i-1,"ST"]==data.loc[i-1,"FUB"])&(data.loc[i,"Close"]>data.loc[i,"FUB"]):
        data.loc[i,"ST"]=data.loc[i,"FLB"]
    elif (data.loc[i-1,"ST"]==data.loc[i-1,"FLB"])&(data.loc[i,"Close"]>=data.loc[i,"FLB"]):
        data.loc[i,"ST"]=data.loc[i,"FLB"]
    elif (data.loc[i-1,"ST"]==data.loc[i-1,"FLB"])&(data.loc[i,"Close"]<data.loc[i,"FLB"]):
        data.loc[i,"ST"]=data.loc[i,"FUB"]

# Buy Sell Indicator
for i, row in data.iterrows():
    if i==0:
        data["ST_BUY_SELL"]="NA"
    elif (data.loc[i,"ST"]<data.loc[i,"Close"]) :
        data.loc[i,"ST_BUY_SELL"]="BUY"
    else:
        data.loc[i,"ST_BUY_SELL"]="SELL"

Related