indicators/sma.py

# indicators/sma.py

import pandas as pd

def calculate_sma(data, period):
    """
    Calculates the Simple Moving Average (SMA) for a given period.

    Parameters:
    - data: DataFrame containing stock prices with a 'Close' column (DataFrame).
    - period: The period over which to calculate the SMA (int).

    Returns:
    - sma: The calculated SMA as a Series (pd.Series).
    """
    sma = data['Close'].rolling(window=period, min_periods=1).mean()
    return sma

def add_sma_columns(data):
    """
    Adds SMA columns for the 21 and 50 periods to the input DataFrame.

    Parameters:
    - data: DataFrame containing stock prices. Must include a 'Close' column (DataFrame).

    Modifies:
    - data: The input DataFrame is modified in-place, adding two new columns: 'SMA_21' and 'SMA_50'.
    """
    data['SMA_21'] = calculate_sma(data, 21)
    data['SMA_50'] = calculate_sma(data, 50)

# 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