{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ba62c4ae",
   "metadata": {},
   "source": [
    "<div style=\"background-color:#000;\"><img src=\"pqn.png\"></img></div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d26bf77",
   "metadata": {},
   "source": [
    "This notebook performs momentum-based stock trading simulation using vectorbt. It downloads historical stock prices for selected symbols, calculates momentum, and filters the top stocks based on their momentum. The code then applies various exit strategies like Stop Loss (SL), Trailing Stop (TS), and Take Profit (TP) to the selected stocks. Finally, it creates a portfolio with the selected stocks and exit strategies and visualizes the total returns using histograms and boxplots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e804ab6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytz\n",
    "from datetime import datetime, timedelta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fd91c63",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90a345fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import vectorbt as vbt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1890d17",
   "metadata": {},
   "source": [
    "Define stock symbols to be analyzed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7229afd",
   "metadata": {},
   "outputs": [],
   "source": [
    "symbols = [\n",
    "    \"META\",\n",
    "    \"AMZN\",\n",
    "    \"AAPL\",\n",
    "    \"NFLX\",\n",
    "    \"GOOG\",\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8e0df97",
   "metadata": {},
   "source": [
    "Define the start and end dates for historical data download"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff36497f",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_date = datetime(2018, 1, 1, tzinfo=pytz.utc)\n",
    "end_date = datetime(2021, 1, 1, tzinfo=pytz.utc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab953913",
   "metadata": {},
   "source": [
    "Define various simulation parameters including number of stocks to trade, window length, seed, window count, exit types, and stop values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0417b3ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "traded_count = 3\n",
    "window_len = timedelta(days=12 * 21)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbb72d7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 42\n",
    "window_count = 400\n",
    "exit_types = [\"SL\", \"TS\", \"TP\"]\n",
    "stops = np.arange(0.01, 1 + 0.01, 0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42525ce7",
   "metadata": {},
   "source": [
    "Download historical stock data using vectorbt's YFData module and concatenate the data into a single DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a4b6d14",
   "metadata": {},
   "outputs": [],
   "source": [
    "yfdata = vbt.YFData.download(symbols, start=start_date, end=end_date)\n",
    "ohlcv = yfdata.concat()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7f07cd7",
   "metadata": {},
   "source": [
    "Split the OHLCV data into windows defined by the window length and count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5ee17e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "split_ohlcv = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54c69e07",
   "metadata": {},
   "outputs": [],
   "source": [
    "for k, v in ohlcv.items():\n",
    "    split_df, split_indexes = v.vbt.range_split(\n",
    "        range_len=window_len.days, n=window_count\n",
    "    )\n",
    "    split_ohlcv[k] = split_df\n",
    "ohlcv = split_ohlcv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cae16ab7",
   "metadata": {},
   "source": [
    "Calculate the momentum as the mean percentage change of the closing prices, then select the top stocks based on momentum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8762ce5",
   "metadata": {},
   "outputs": [],
   "source": [
    "momentum = ohlcv[\"Close\"].pct_change().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e62896e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted_momentum = (\n",
    "    momentum\n",
    "    .groupby(\n",
    "        \"split_idx\", \n",
    "        group_keys=False, \n",
    "        sort=False\n",
    "    )\n",
    "    .apply(\n",
    "        pd.Series.sort_values\n",
    "    )\n",
    "    .groupby(\"split_idx\")\n",
    "    .head(traded_count)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83807d67",
   "metadata": {},
   "source": [
    "Select the OHLCV data for the stocks with the highest momentum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a07939d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "selected_open = ohlcv[\"Open\"][sorted_momentum.index]\n",
    "selected_high = ohlcv[\"High\"][sorted_momentum.index]\n",
    "selected_low = ohlcv[\"Low\"][sorted_momentum.index]\n",
    "selected_close = ohlcv[\"Close\"][sorted_momentum.index]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64fb1e8f",
   "metadata": {},
   "source": [
    "Initialize entry signals to be true on the first day of each window"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fa6acb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "entries = pd.DataFrame.vbt.signals.empty_like(selected_open)\n",
    "entries.iloc[0, :] = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dd7d916",
   "metadata": {},
   "source": [
    "Define stop loss exits using vectorbt's OHLCSTX module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24bf8d0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sl_exits = vbt.OHLCSTX.run(\n",
    "    entries,\n",
    "    selected_open,\n",
    "    selected_high,\n",
    "    selected_low,\n",
    "    selected_close,\n",
    "    sl_stop=list(stops),\n",
    "    stop_type=None,\n",
    "    stop_price=None,\n",
    ").exits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ec6eab3",
   "metadata": {},
   "source": [
    "Define trailing stop exits using vectorbt's OHLCSTX module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b038fae0",
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_exits = vbt.OHLCSTX.run(\n",
    "    entries,\n",
    "    selected_open,\n",
    "    selected_high,\n",
    "    selected_low,\n",
    "    selected_close,\n",
    "    sl_stop=list(stops),\n",
    "    sl_trail=True,\n",
    "    stop_type=None,\n",
    "    stop_price=None,\n",
    ").exits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad721230",
   "metadata": {},
   "source": [
    "Define take profit exits using vectorbt's OHLCSTX module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ff261d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "tp_exits = vbt.OHLCSTX.run(\n",
    "    entries,\n",
    "    selected_open,\n",
    "    selected_high,\n",
    "    selected_low,\n",
    "    selected_close,\n",
    "    tp_stop=list(stops),\n",
    "    stop_type=None,\n",
    "    stop_price=None,\n",
    ").exits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2a0bd88",
   "metadata": {},
   "source": [
    "Rename and drop levels for the different exit types to standardize the DataFrame structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5674fff",
   "metadata": {},
   "outputs": [],
   "source": [
    "sl_exits.vbt.rename_levels({\"ohlcstx_sl_stop\": \"stop_value\"}, inplace=True)\n",
    "ts_exits.vbt.rename_levels({\"ohlcstx_sl_stop\": \"stop_value\"}, inplace=True)\n",
    "tp_exits.vbt.rename_levels({\"ohlcstx_tp_stop\": \"stop_value\"}, inplace=True)\n",
    "ts_exits.vbt.drop_levels(\"ohlcstx_sl_trail\", inplace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4ecca94",
   "metadata": {},
   "source": [
    "Ensure the last day in the window is always an exit signal for all exit types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89d7940d",
   "metadata": {},
   "outputs": [],
   "source": [
    "sl_exits.iloc[-1, :] = True\n",
    "ts_exits.iloc[-1, :] = True\n",
    "tp_exits.iloc[-1, :] = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ec14fef",
   "metadata": {},
   "source": [
    "Convert exits into first exit signals based on entries, allowing gaps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09e6acb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "sl_exits = sl_exits.vbt.signals.first(reset_by=entries, allow_gaps=True)\n",
    "ts_exits = ts_exits.vbt.signals.first(reset_by=entries, allow_gaps=True)\n",
    "tp_exits = tp_exits.vbt.signals.first(reset_by=entries, allow_gaps=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0570cee9",
   "metadata": {},
   "source": [
    "Concatenate all exit signals into a single DataFrame for further analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e06ca90c",
   "metadata": {},
   "outputs": [],
   "source": [
    "exits = pd.DataFrame.vbt.concat(\n",
    "    sl_exits,\n",
    "    ts_exits,\n",
    "    tp_exits,\n",
    "    keys=pd.Index(exit_types, name=\"exit_type\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a246e7b2",
   "metadata": {},
   "source": [
    "Create a portfolio using the selected close prices, entries, and exits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "657130c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "portfolio = vbt.Portfolio.from_signals(selected_close, entries, exits)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e7d40f2",
   "metadata": {},
   "source": [
    "Calculate the total return of the portfolio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76791b7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_return = portfolio.total_return()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b92b0fe",
   "metadata": {},
   "source": [
    "Unstack the total returns by exit type for visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f9c0ffc",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_return_by_type = total_return.unstack(level=\"exit_type\")[exit_types]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a1fd9bc",
   "metadata": {},
   "source": [
    "Plot histograms of the total returns for each exit type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb0943a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_return_by_type[exit_types].vbt.histplot(\n",
    "    xaxis_title=\"Total return\",\n",
    "    xaxis_tickformat=\"%\",\n",
    "    yaxis_title=\"Count\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85d76d8c",
   "metadata": {},
   "source": [
    "Plot boxplots of the total returns for each exit type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f05fdd98",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_return_by_type.vbt.boxplot(\n",
    "    yaxis_title='Total return',\n",
    "    yaxis_tickformat='%'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c88d4f8a",
   "metadata": {},
   "source": [
    "Provide descriptive statistics for the total returns by exit type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca8a1a68",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_return_by_type.describe(percentiles=[])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da178273",
   "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
}
