This script takes a list of hard-coded stock symbols, fetches historical bar data on them, and then creates a heatmap illustrating Pearson's correlation between various pairs.

alpaca_correlation_viz.py

# import libraries
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from dotenv import load_dotenv
from alpaca_trade_api import REST, TimeFrame
import plotly.graph_objects as go
import plotly.express as px
import os
from datetime import datetime, timedelta

# Import symbols from symbols.py
from symbols import categories

# Load environment variables
load_dotenv()

# Get API and SECRET keys
API_KEY = os.getenv("API_KEY")
SECRET_KEY = os.getenv("SECRET_KEY")

# Create REST client
rest_client = REST(API_KEY, SECRET_KEY)

# Use symbols from a specific category
category_to_use = 'Banking'  # Change this value to use a different category
symbols = categories[category_to_use]
data = pd.DataFrame()

now = datetime.now()
yesterday = now - timedelta(days=1, minutes=16)
start_date = now - timedelta(days=90)

start_date_str = start_date.strftime('%Y-%m-%d')
yesterday_str = yesterday.strftime('%Y-%m-%d')

for symbol in symbols:
    bars = rest_client.get_bars(symbol, TimeFrame.Day, start_date_str, yesterday_str).df
    bars['symbol'] = symbol
    data = pd.concat([data, bars])

# Calculate daily returns
data['returns'] = data.groupby('symbol')['close'].pct_change()

# Calculate correlations
def correlation(df):
    corr_matrix = df.pivot_table(values='returns', index=df.index, columns='symbol').corr()
    return corr_matrix

# Calculate number of days
num_days = (yesterday - start_date).days

# Plot heatmap
plt.figure(figsize=(16, 10))
sns_heatmap = sns.heatmap(correlation(data), annot=True, cmap="Blues", annot_kws={"size": 6})
sns_heatmap.set_title(f"Pearson correlation of {category_to_use.lower()} symbols: {start_date_str} to {yesterday_str} ({num_days}D)", fontsize=14, pad=20)

# Generate PNG of heatmap
os.makedirs("output", exist_ok=True)

# Save heatmap as a PNG file
heatmap_file = f"output/heatmap.png"
plt.savefig(heatmap_file, dpi=600, bbox_inches='tight')
plt.close()

symbols.py

# Categorize symbols into industry sectors
tech_symbols = ['AAPL', 'AMZN', 'GOOGL', 'NVDA', 'MSFT', 'ADBE', 'SHOP', 'TWLO', 'PYPL', 'INTC', 'NVDA', 'DOCN', 'ROKU', 'SNAP', 'META', 'CRM', 'PLTR', 'AFRM', 'T', 'V', 'EBAY', 'DDOG', 'HUBS']
automotive_symbols = ['TSLA', 'F', 'GM']
transportation_symbols = ['BA', 'AAL', 'UPS']
retail_symbols = ['TGT', 'WMT', 'NKE', 'BBY', 'MCD', 'HD', 'SBUX']
communication_symbols = ['VZ', 'TMUS']
financial_symbols = ['COIN', 'HOOD', 'BLK', 'BRK.A', 'V', 'COIN', 'MA', 'PYPL', 'SQ', 'JPM']
bank_symbols = ['C', 'WFC', 'JPM', 'WAL', 'FRC', 'TFC', 'EWBC', 'CMA', 'FITB', 'USG', 'CFG', 'HBAN', 'KEY', 'ZION', 'RF', 'MTB']
media_symbols = ['DIS', 'SPOT', 'LYFT', 'UBER', 'HOOD']
consumer_goods_symbols = ['XOM', 'CVNA', 'ABNB']
healthcare_symbols = ['JNJ', 'MMM']
miscellaneous_symbols = ['Q']
etf_symbols = ['QQQ', 'XLK', 'VGT']

# Group the symbols into a dictionary
categories = {
    'Technology': tech_symbols,
    'Automotive': automotive_symbols,
    'Transportation': transportation_symbols,
    'Retail': retail_symbols,
    'Communication': communication_symbols,
    'Financial': financial_symbols,
    'Media': media_symbols,
    'Consumer Goods': consumer_goods_symbols,
    'Healthcare': healthcare_symbols,
    'Miscellaneous': miscellaneous_symbols,
    'ETFs': etf_symbols,
    'Banking': bank_symbols
}