Update indicators/bollinger_bands.py

indicators/bollinger_bands.py

@@ -1,39 +1,54 @@
-# indicators/bollinger_bands.py
-
 import pandas as pd
 
-def calculate_bollinger_bands(data, period=21, std_multiplier=1.7):
+def calculate_sma(data, window):
     """
-    Calculates Bollinger Bands for a given period and standard deviation multiplier.
-
+    Calculate the Simple Moving Average (SMA) for the given data.
+    
     Parameters:
-    - data: DataFrame containing stock prices with a 'Close' column (DataFrame).
-    - period: The period over which to calculate the SMA and standard deviation (int).
-    - std_multiplier: The multiplier for the standard deviation to calculate the upper and lower bands (float).
+    - 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:
-    - A DataFrame with columns 'BB_Middle', 'BB_Upper', 'BB_Lower'.
+    - 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'] = data['Close'].rolling(window=period, min_periods=1).mean()
+    data['BB_Middle'] = calculate_sma(data['Close'], window)
     
     # Calculate the standard deviation
-    std_dev = data['Close'].rolling(window=period, min_periods=1).std()
+    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[['BB_Middle', 'BB_Upper', 'BB_Lower']]
+    return data
 
-# Example usage
 if __name__ == "__main__":
-    # Assuming 'data' is a DataFrame that contains stock price data including a 'Close' column.
-    # For the sake of example, let's create a dummy DataFrame.
-    dates = pd.date_range(start="2023-01-01", end="2023-02-28", freq='D')
-    prices = pd.Series([100 + i * 0.5 for i in range(len(dates))], index=dates)
-    data = pd.DataFrame(prices, columns=['Close'])
+    # 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
-    bollinger_bands = calculate_bollinger_bands(data)
-    print(bollinger_bands.head())  # Display the first few rows to verify the calculations
+    bb_data = calculate_bollinger_bands(sample_data)
+    
+    print(bb_data.head())  # Print the first few rows to verify