Update indicators/sma.py

indicators/sma.py

@@ -1,42 +1,45 @@
-# indicators/sma.py
-
 import pandas as pd
 
-def calculate_sma(data, period):
+def calculate_sma(data, window):
     """
-    Calculates the Simple Moving Average (SMA) for a given period.
+    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 (int).
+    - data (pd.Series): The stock data (typically closing prices).
+    - window (int): The period over which to calculate the SMA.
 
     Returns:
-    - sma: The calculated SMA as a Series (pd.Series).
+    - pd.Series: The calculated SMA values.
     """
-    sma = data['Close'].rolling(window=period, min_periods=1).mean()
-    return sma
+    return data.rolling(window=window, min_periods=1).mean()
 
-def add_sma_columns(data):
+def calculate_21_50_sma(data):
     """
-    Adds SMA columns for the 21 and 50 periods to the input DataFrame.
+    Calculate both the 21-period and 50-period SMAs for the given stock data.
 
     Parameters:
-    - data: DataFrame containing stock prices. Must include a 'Close' column (DataFrame).
+    - data (pd.DataFrame): The stock data, expected to have a 'Close' column.
 
-    Modifies:
-    - data: The input DataFrame is modified in-place, adding two new columns: 'SMA_21' and 'SMA_50'.
+    Returns:
+    - pd.DataFrame: The input data frame with added columns for the 21-period and 50-period SMAs.
     """
-    data['SMA_21'] = calculate_sma(data, 21)
-    data['SMA_50'] = calculate_sma(data, 50)
+    if 'Close' not in data.columns:
+        raise ValueError("Data frame must contain a 'Close' column.")
+
+    # Calculate the SMAs
+    data['SMA_21'] = calculate_sma(data['Close'], 21)
+    data['SMA_50'] = calculate_sma(data['Close'], 50)
+
+    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([i * 0.01 for i in range(len(dates))], index=dates)
-    data = pd.DataFrame(prices, columns=['Close'])
-    
-    # Add SMA columns
-    add_sma_columns(data)
-    print(data.head())  # Display the first few rows to verify the SMA calculations
+    # Example usage
+    # Generate a sample DataFrame with 'Close' prices
+    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 the 21 and 50 period SMAs
+    sma_data = calculate_21_50_sma(sample_data)
+
+    print(sma_data.head())  # Print the first few rows to verify