{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a43648c6",
   "metadata": {},
   "source": [
    "<div style=\"background-color:#000;\"><img src=\"pqn.png\"></img></div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ed53dc4",
   "metadata": {},
   "source": [
    "This code downloads and processes stock price data, trains an autoencoder to compress and reconstruct the data, and then uses the learned embeddings for clustering and visualization. First, it fetches historical stock price data from Yahoo Finance. It then computes log returns, moving averages, and volatility to form features. A PyTorch autoencoder is trained on these features to learn compressed representations (embeddings). Finally, it uses K-Means clustering on the embeddings and visualizes the clusters using PCA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5dc8b2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import yfinance as yf\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from sklearn.cluster import KMeans\n",
    "from sklearn.decomposition import PCA\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import warnings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "165d5792",
   "metadata": {},
   "outputs": [],
   "source": [
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14014259",
   "metadata": {},
   "source": [
    "Define a list of stock symbols to fetch data for"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00cf6af0",
   "metadata": {},
   "outputs": [],
   "source": [
    "symbols = [\n",
    "    \"AAPL\", \"MSFT\", \"GOOGL\", \"AMZN\", \"META\",\n",
    "    \"TSLA\", \"BRK-B\", \"V\", \"JNJ\", \"WMT\", \"JPM\",\n",
    "    \"MA\", \"PG\", \"UNH\", \"DIS\", \"NVDA\", \"HD\", \n",
    "    \"PYPL\", \"BAC\", \"VZ\", \"ADBE\", \"CMCSA\", \"NFLX\",\n",
    "    \"KO\", \"NKE\", \"MRK\", \"PEP\", \"T\", \"PFE\", \"INTC\",\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbc075cc",
   "metadata": {},
   "source": [
    "Download adjusted close prices for the specified symbols from Yahoo Finance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b5e8366",
   "metadata": {},
   "outputs": [],
   "source": [
    "stock_data = yf.download(\n",
    "    symbols, \n",
    "    start=\"2020-01-01\", \n",
    "    end=\"2023-12-31\"\n",
    ")[\"Adj Close\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0283e34",
   "metadata": {},
   "source": [
    "Calculate log returns, moving averages, and volatility for the stock data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba5af677",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_returns = np.log(stock_data / stock_data.shift(1))\n",
    "moving_avg = stock_data.rolling(window=22).mean()\n",
    "volatility = stock_data.rolling(window=22).std()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da7d71aa",
   "metadata": {},
   "source": [
    "Concatenate the features and standardize them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a77301f",
   "metadata": {},
   "outputs": [],
   "source": [
    "features = pd.concat([log_returns, moving_avg, volatility], axis=1).dropna()\n",
    "processed_data = (features - features.mean()) / features.std()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0064e95",
   "metadata": {},
   "source": [
    "Convert the features into PyTorch tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c80d578a",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "tensor = torch.tensor(processed_data.values, dtype=torch.float32)\n",
    "dataset = TensorDataset(tensor)\n",
    "data_loader = DataLoader(dataset, batch_size=32, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac181de3",
   "metadata": {},
   "source": [
    "Define an autoencoder neural network for stock data embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d962be87",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class StockAutoencoder(nn.Module):\n",
    "    \"\"\"Autoencoder neural network for stock data embedding\n",
    "    \n",
    "    This class defines an autoencoder with an encoder and decoder\n",
    "    to compress and reconstruct stock data.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    feature_dim : int\n",
    "        The dimensionality of the input features\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, feature_dim):\n",
    "        super(StockAutoencoder, self).__init__()\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Linear(feature_dim, 64),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(64, 32),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(32, 10),  # Latent space\n",
    "        )\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Linear(10, 32),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(32, 64),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(64, feature_dim),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.encoder(x)\n",
    "        x = self.decoder(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc7f8754",
   "metadata": {},
   "source": [
    "Train the autoencoder on the stock data using MSE loss and Adam optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "413d6d54",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def train(model, data_loader, epochs=100):\n",
    "    \"\"\"Train the autoencoder model\n",
    "    \n",
    "    This function trains the autoencoder using MSE loss and Adam optimizer\n",
    "    over a specified number of epochs.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    model : nn.Module\n",
    "        The autoencoder model to be trained\n",
    "    data_loader : DataLoader\n",
    "        DataLoader object to iterate through the dataset\n",
    "    epochs : int, optional\n",
    "        Number of epochs to train the model (default is 100)\n",
    "    \"\"\"\n",
    "    \n",
    "    criterion = nn.MSELoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "    model.train()\n",
    "    for epoch in range(epochs):\n",
    "        for data in data_loader:\n",
    "            inputs = data[0]\n",
    "            optimizer.zero_grad()\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, inputs)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        print(f\"Epoch {epoch+1}, Loss: {loss.item()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e48b0a76",
   "metadata": {},
   "source": [
    "Initialize and train the autoencoder model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca7a96ba",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "feature_dim = processed_data.shape[1]\n",
    "model = StockAutoencoder(feature_dim)\n",
    "train(model, data_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07873939",
   "metadata": {},
   "source": [
    "Extract embeddings from the trained autoencoder model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5768c886",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def extract_embeddings(model, data_loader):\n",
    "    \"\"\"Extract embeddings from the trained autoencoder model\n",
    "    \n",
    "    This function extracts embeddings by passing data through the encoder\n",
    "    part of the autoencoder.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    model : nn.Module\n",
    "        The trained autoencoder model\n",
    "    data_loader : DataLoader\n",
    "        DataLoader object to iterate through the dataset\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    embeddings : torch.Tensor\n",
    "        Tensor containing the extracted embeddings\n",
    "    \"\"\"\n",
    "    \n",
    "    model.eval()\n",
    "    embeddings = []\n",
    "    with torch.no_grad():\n",
    "        for data in data_loader:\n",
    "            inputs = data[0]\n",
    "            encoded = model.encoder(inputs)\n",
    "            embeddings.append(encoded)\n",
    "    return torch.vstack(embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e6f9d1e",
   "metadata": {},
   "source": [
    "Extract embeddings from the trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91040de9",
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = extract_embeddings(model, data_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28747eff",
   "metadata": {},
   "source": [
    "Apply K-Means clustering on the embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a6b69b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "kmeans = KMeans(n_clusters=5, random_state=42).fit(embeddings.numpy())\n",
    "clusters = kmeans.labels_"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ac98fcd",
   "metadata": {},
   "source": [
    "Reduce the dimensionality of embeddings using PCA for visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35295d99",
   "metadata": {},
   "outputs": [],
   "source": [
    "pca = PCA(n_components=2)\n",
    "embeddings_2d = pca.fit_transform(embeddings.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef42a090",
   "metadata": {},
   "source": [
    "Plot the clusters in 2D space using PCA components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7026ea6",
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.scatterplot(\n",
    "    x=embeddings_2d[:, 0],\n",
    "    y=embeddings_2d[:, 1],\n",
    "    hue=clusters,\n",
    "    palette=sns.color_palette(\"hsv\", len(set(clusters))),\n",
    ")\n",
    "plt.xlabel(\"PCA Dimension 1\")\n",
    "plt.ylabel(\"PCA Dimension 2\")\n",
    "plt.legend(title=\"Cluster\")\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1044a6a3",
   "metadata": {},
   "source": [
    "<a href=\"https://pyquantnews.com/\">PyQuant News</a> is where finance practitioners level up with Python for quant finance, algorithmic trading, and market data analysis. Looking to get started? Check out the fastest growing, top-selling course to <a href=\"https://gettingstartedwithpythonforquantfinance.com/\">get started with Python for quant finance</a>. For educational purposes. Not investment advise. Use at your own risk."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
