indicators/bollinger_bands.py

import pandas as pd

def calculate_sma(data, window):
    """
    Calculate the Simple Moving Average (SMA) for the given data.
    
    Parameters:
    - data (pd.Series): The stock data (typically closing prices).
    - window (int): The period over which to calculate the SMA.
    
    Returns:
    - pd.Series: The calculated SMA values.
    """
    return data.rolling(window=window, min_periods=1).mean()

def calculate_bollinger_bands(data, window=21, std_multiplier=1.7):
    """
    Calculate Bollinger Bands for the given stock data.
    
    Parameters:
    - data (pd.DataFrame): The stock data, expected to have a 'Close' column.
    - window (int): The SMA period for the middle band. Defaults to 21.
    - std_multiplier (float): The standard deviation multiplier for the upper and lower bands. Defaults to 1.7.
    
    Returns:
    - pd.DataFrame: The input data frame with added columns for the Bollinger Bands ('BB_Middle', 'BB_Upper', 'BB_Lower').
    """
    if 'Close' not in data.columns:
        raise ValueError("Data frame must contain a 'Close' column.")
    
    # Calculate the middle band (SMA)
    data['BB_Middle'] = calculate_sma(data['Close'], window)
    
    # Calculate the standard deviation
    std_dev = data['Close'].rolling(window=window).std()
    
    # Calculate the upper and lower bands
    data['BB_Upper'] = data['BB_Middle'] + (std_multiplier * std_dev)
    data['BB_Lower'] = data['BB_Middle'] - (std_multiplier * std_dev)
    
    return data

if __name__ == "__main__":
    # Example usage
    # Generate a sample DataFrame with 'Close' prices
    import numpy as np
    dates = pd.date_range(start='2023-01-01', periods=100, freq='D')
    close_prices = pd.Series((100 + np.random.randn(100).cumsum()), index=dates)
    sample_data = pd.DataFrame({'Close': close_prices})
    
    # Calculate Bollinger Bands
    bb_data = calculate_bollinger_bands(sample_data)
    
    print(bb_data.head())  # Print the first few rows to verify