real_sp500_recovery.py

#!/usr/bin/env python3
'''
Generate a graph of the real returns of the S&P 500
'''

import logging

import yfinance as yf
import pandas as pd
import requests
import matplotlib.pyplot as plt

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

def main():
    '''
    Main function
    '''
    sp500 = fetch_sp500_data()
    cpi_data = fetch_cpi_data()

    real_returns = calculate_real_returns(sp500, cpi_data)

    plot_real_returns(real_returns)


def fetch_sp500_data(start_date=None, end_date=None):
    '''
    Fetch S&P 500 data from Yahoo Finance
    '''
    if not start_date:
        start_date = pd.Timestamp.today() - pd.DateOffset(years=10)

    if not end_date:
        end_date = pd.Timestamp.today()

    # Fetch S&P 500 data
    sp500 = yf.download('^GSPC', start=start_date, end=end_date)
    if isinstance(sp500.columns, pd.MultiIndex):
        sp500.columns = sp500.columns.droplevel(1)

    logging.info('Fetched S&P 500 data: %s', sp500)

    return sp500


def fetch_cpi_data(start_date=None, end_date=None):
    '''
    Fetch CPI data from FRED
    '''
    if not start_date:
        start_date = pd.Timestamp.today() - pd.DateOffset(years=10)

    if not end_date:
        end_date = pd.Timestamp.today()

    # BLS API endpoint for CPI data
    cpi_url = 'https://api.bls.gov/publicAPI/v2/timeseries/data/CUUR0000SA0'

    # Request parameters
    params = {
        'startyear': start_date.year,
        'endyear': end_date.year,
    }

    # Fetch CPI data
    response = requests.get(
        cpi_url,
        params=params,
        timeout=10
    )

    data = response.json()

    if not response.status_code == 200:
        logging.error('Failed to fetch CPI data (%s): %s', response.status_code, response.text)

    if data['status'] != 'REQUEST_SUCCEEDED':
        logging.error('Failed to fetch CPI data: %s', data)

        logging.info('Falling back to manual data entry')
        with open('cpi.csv', 'r', encoding='utf-8') as f:
            cpi_data = [
                (pd.Timestamp(line.split(',')[0]), float(line.split(',')[6])) for line in f.readlines()[1:]
            ]
    else:
        # Extract CPI values
        cpi_data = []

        for series in data['Results']['series']:
            for item in series['data']:
                date = pd.Timestamp(f"{item['year']}-{item['period'][1:]:0>2}-01")
                value = float(item['value'])
                cpi_data.append((date, value))

    # Create DataFrame
    cpi_df = pd.DataFrame(cpi_data, columns=['Date', 'CPI'])
    cpi_df.set_index('Date', inplace=True)

    # Convert monthly CPI data to daily frequency and forward-fill
    cpi_df = cpi_df.resample('D').ffill()

    # Restrict to requested date range
    cpi_df = cpi_df.loc[start_date:end_date]

    logging.info('Fetched CPI data: %s', cpi_df)

    return cpi_df


def calculate_real_returns(sp500, cpi_data):
    '''
    Calculate real returns of the S&P 500
    '''
    logging.info('Calculating real returns')

    data = sp500.merge(cpi_data, left_index=True, right_index=True)

    inflation = (data['CPI'] / data['CPI'].iloc[0])

    data['Real Close'] = data['Close'] / inflation

    data['Dollar'] = inflation

    logging.info('Calculated real returns: %s', data)

    return data


def annotate_market_corrections(ax, data, column, label, line, max_duration=180):
    """
    Annotates market corrections on a given matplotlib axis.

    Args:
        ax: The axis to annotate.
        data: The DataFrame containing the market data.
        column: The column (Close or Real Close) to analyze.
        label: A label indicating whether it's nominal or real.
    """
    peak_date  = None
    peak_value = None

    for i in range(len(data) - 1):
        current_date  = data.index[i]
        current_value = data[column].iloc[i]

        # Initialize peak
        if not peak_value:
            peak_value = current_value
            peak_date  = current_date

        # Identify local peak
        if current_value > peak_value:
            recovery_date = current_date

            duration = (recovery_date - peak_date).days

            logging.info('Detected market recovery: %s to %s (%s)', peak_date, recovery_date, duration)

            if duration > max_duration:
                logging.info('Annotating market correction: %s', data.loc[peak_date:recovery_date])

                ax.hlines(peak_value, peak_date, recovery_date, color=ax.lines[line].get_color(), linestyle='--')
                ax.axvspan(peak_date, recovery_date, color='grey', alpha=0.1)
                ax.annotate(
                    f"{duration} days ({label})",
                    xy=(peak_date, peak_value * 1.01),
                    color=ax.lines[line].get_color()
                )

            peak_value = current_value
            peak_date  = current_date

    return ax


def plot_real_returns(real_returns):
    '''
    Plot real returns of the S&P 500
    '''
    logging.info('Plotting real returns')

    # discard the fig
    _, ax1 = plt.subplots(figsize=(12, 6))

    ax1.plot(real_returns.index, real_returns['Close'], label='Nominal S&P 500')
    ax1.plot(real_returns.index, real_returns['Real Close'], label='Real S&P 500 (Adjusted for Inflation)')
    ax1.set_title('S&P 500: Nominal vs. Real Returns Over the Last 10 Years')
    ax1.set_xlabel('Date')
    ax1.set_ylabel('S&P 500 Index Level')
    ax1.grid(True)
    ax1.legend(loc='center left')

    ax2 = ax1.twinx()
    ax2.plot(real_returns.index, 1 / real_returns['Dollar'], 'r', label='USD Spending Power')
    ax2.set_ylabel('USD Spending Power')
    ax2.legend(loc='lower left')

    logging.info('Annotating market corrections')

    ax1 = annotate_market_corrections(ax1, real_returns, 'Close', 'nominal', 0)
    ax1 = annotate_market_corrections(ax1, real_returns, 'Real Close', 'real', 1)

    logging.debug('Data plotted:\n%s', real_returns.to_csv())

    plt.show()

    logging.info('Plotting complete')


if __name__ == '__main__':
    main()