Create Tradingview RSI in python using numpy RMA

Question

Here are some links i am using to come up with the code

https://www.tradingview.com/pine-script-reference/v5/#fun_ta.rma

So i was searching and searching and searching to figure out how to create the rsi like tradingview and after looking at a ton of stackoverview questions i was able to find one that really helped me figure out what the problem was ... it was the answer here Pandas.Series.ewm for calculating a RSI Indicator, with Wilders Moving Average

it was the fact that the very first average is calculated using the mean of the past rsi length ... like this

avg_gain[period] = gain[1 : period + 1].mean()

then after you do that you apply the rma tradingview calculation https://www.tradingview.com/pine-script-reference/v5/#fun_ta.rma

pine_rma(src, length) =>
    alpha = 1/length
    sum = 0.0
    sum := na(sum[1]) ? ta.sma(src, length) : alpha * src + (1 - alpha) * nz(sum[1])

so the formula i came up with was this and i checked it and everything is like 98% accurate ... usually it is off by .1 or .05 ... but still super super super super close

import numpy as np
def qf_tv_calc_rsi_wilder(prices: np.array, period: int):
    prices_shift = np.roll(prices, 1)
    prices_shift[0] = np.nan
    pchg = (prices - prices_shift) / prices_shift

    alpha = 1 / period
    gain = np.where(pchg > 0, pchg, 0)
    avg_gain = np.full_like(gain, np.nan)

    loss = np.where(pchg < 0, abs(pchg), 0)
    avg_loss = np.full_like(loss, np.nan)

    avg_gain[period] = gain[1 : period + 1].mean()
    avg_loss[period] = loss[1 : period + 1].mean()

    for i in range(period + 1, gain.size):
        avg_gain[i] = alpha * gain[i] + (1 - alpha) * avg_gain[i - 1]
        avg_loss[i] = alpha * loss[i] + (1 - alpha) * avg_loss[i - 1]

    rs = avg_gain / avg_loss

    rsi = 100 - (100 / (1 + rs))
    return rsi

this is my rsi

this is tradingviews rsi ... toward the start it gets a little crazy but that is because i only downloaded 200 candles while tradingview is using like 3000

so my question here is in the for loop i tried to vectorize it and it didn't work

avg_gain[period + 1 :] = alpha * gain[period + 1 :] + (1 - alpha) * avg_gain[period]

avg_loss[period + 1 :] = alpha * loss[period + 1 :] + (1 - alpha) * avg_loss[period]

it spits out some crazy numbers ... so does anyone know how i can get rid of the for loop ?

Answer 1

This script will give you a fairly closer RSI value without having to loop through the dataset.

https://github.com/ImDakshitha/tradingview_rsi

But might need a lager dataset for more accurate results(about 100 datapoints)