functions.py

import matplotlib.pyplot as plt
import numpy as np


#  HELPER FUNCTIONS:
def plot_column(dataframe, column_name):
    plt.figure(figsize=(6, 3))
    plt.plot(dataframe[column_name])
    plt.title(f"Plot of {column_name}")
    plt.xlabel("Index")
    plt.ylabel(column_name)
    plt.grid(True)
    return plt


# Calculate EMA
def ema(close, period=20):
    return close.ewm(span=period, adjust=False).mean()


# Calculate RSI
def rsi(close, period=14):
    delta = close.diff()
    gain, loss = delta.copy(), delta.copy()
    gain[gain < 0] = 0
    loss[loss > 0] = 0
    avg_gain = gain.rolling(period).mean()
    avg_loss = abs(loss.rolling(period).mean())
    rs = avg_gain / avg_loss
    rsi = 100.0 - (100.0 / (1.0 + rs))
    return rsi


# Calculate MACD
def macd(close, fast_period=12, slow_period=26, signal_period=9):
    fast_ema = close.ewm(span=fast_period, adjust=False).mean()
    slow_ema = close.ewm(span=slow_period, adjust=False).mean()
    macd_line = fast_ema - slow_ema
    signal_line = macd_line.ewm(span=signal_period, adjust=False).mean()
    histogram = macd_line - signal_line
    return macd_line


# Calculate OBV
def obv(close, volume):
    obv = np.where(
        close > close.shift(), volume, np.where(close < close.shift(), -volume, 0)
    ).cumsum()
    return obv