indicators/sma.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_21_50_sma(data):
    """
    Calculate both the 21-period and 50-period SMAs for the given stock data.

    Parameters:
    - data (pd.DataFrame): The stock data, expected to have a 'Close' column.

    Returns:
    - pd.DataFrame: The input data frame with added columns for the 21-period and 50-period SMAs.
    """
    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

if __name__ == "__main__":
    # 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