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