{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "abad0a97",
   "metadata": {},
   "source": [
    "<div style=\"background-color:#000;\"><img src=\"pqn.png\"></img></div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1e037dc",
   "metadata": {},
   "source": [
    "This code implements a trading algorithm using the Zipline library in Python. It defines a custom momentum factor and a pipeline to rank stocks based on their momentum. The algorithm schedules rebalancing every week and executes trades based on the ranked stocks. The performance of the algorithm is then simulated over a specified time period. This is useful for backtesting trading strategies and analyzing their performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19b48d3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12999480",
   "metadata": {},
   "outputs": [],
   "source": [
    "from zipline import run_algorithm\n",
    "from zipline.pipeline import Pipeline, CustomFactor\n",
    "from zipline.pipeline.data import USEquityPricing\n",
    "from zipline.api import (\n",
    "    attach_pipeline,\n",
    "    calendars,\n",
    "    pipeline_output,\n",
    "    date_rules,\n",
    "    time_rules,\n",
    "    set_commission,\n",
    "    set_slippage,\n",
    "    record,\n",
    "    order_target_percent,\n",
    "    get_open_orders,\n",
    "    schedule_function\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d81d298",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60cdbb78",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7c0cdd7",
   "metadata": {},
   "source": [
    "Define the number of stocks to long and short"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a32b948a",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "N_LONGS = N_SHORTS = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40820869",
   "metadata": {},
   "source": [
    "Define a custom factor for momentum calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e85c8b94",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class Momentum(CustomFactor):\n",
    "    \"\"\"Calculate momentum as the ratio of the last to first price.\n",
    "\n",
    "    This class calculates momentum by dividing the last closing price\n",
    "    by the first closing price over a window.\n",
    "\n",
    "    Attributes\n",
    "    ----------\n",
    "    inputs : list\n",
    "        Data inputs for the factor, here the closing prices.\n",
    "    window_length : int\n",
    "        Length of the window to compute momentum.\n",
    "    \"\"\"\n",
    "\n",
    "    inputs = [USEquityPricing.close]\n",
    "\n",
    "    def compute(self, today, assets, out, close):\n",
    "        \"\"\"Compute the momentum factor.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        today : datetime\n",
    "            Current date.\n",
    "        assets : iterable\n",
    "            List of asset identifiers.\n",
    "        out : ndarray\n",
    "            Output array to store computed momentum values.\n",
    "        close : ndarray\n",
    "            Array of closing prices over the window length.\n",
    "        \"\"\"\n",
    "        out[:] = close[-1] / close[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75e25588",
   "metadata": {},
   "source": [
    "Construct a pipeline to filter and rank stocks based on momentum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8c181d4",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def make_pipeline():\n",
    "    \"\"\"Create a pipeline for ranking stocks based on momentum.\n",
    "\n",
    "    This function defines a pipeline that filters stocks with positive\n",
    "    momentum and ranks them.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Pipeline\n",
    "        A Zipline Pipeline object.\n",
    "    \"\"\"\n",
    "\n",
    "    twenty_day_momentum = Momentum(window_length=20)\n",
    "    thirty_day_momentum = Momentum(window_length=30)\n",
    "\n",
    "    positive_momentum = (\n",
    "        (twenty_day_momentum > 1) & \n",
    "        (thirty_day_momentum > 1)\n",
    "    )\n",
    "\n",
    "    return Pipeline(\n",
    "        columns={\n",
    "            'longs': thirty_day_momentum.top(N_LONGS),\n",
    "            'shorts': thirty_day_momentum.top(N_SHORTS),\n",
    "            'ranking': twenty_day_momentum.rank(ascending=False)\n",
    "        },\n",
    "        screen=positive_momentum\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0aa83e7",
   "metadata": {},
   "source": [
    "Function called before the trading day starts to fetch pipeline output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18c012bc",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def before_trading_start(context, data):\n",
    "    \"\"\"Fetch pipeline output before trading starts.\n",
    "\n",
    "    This function is called before the trading day begins to fetch\n",
    "    the output of the pipeline.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    context : object\n",
    "        An object to store global variables.\n",
    "    data : object\n",
    "        An object to fetch data.\n",
    "    \"\"\"\n",
    "    context.factor_data = pipeline_output(\"factor_pipeline\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2294da67",
   "metadata": {},
   "source": [
    "Initialize the algorithm by attaching the pipeline and scheduling rebalancing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d3b17c6",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def initialize(context):\n",
    "    \"\"\"Initialize the trading algorithm.\n",
    "\n",
    "    This function sets up the pipeline, schedules the rebalancing\n",
    "    function, and initializes other settings.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    context : object\n",
    "        An object to store global variables.\n",
    "    \"\"\"\n",
    "    attach_pipeline(make_pipeline(), \"factor_pipeline\")\n",
    "    schedule_function(\n",
    "        rebalance,\n",
    "        date_rules.week_start(),\n",
    "        time_rules.market_open(),\n",
    "        calendar=calendars.US_EQUITIES,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a56d8fe",
   "metadata": {},
   "source": [
    "Function to rebalance the portfolio based on pipeline output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96a3e1d5",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def rebalance(context, data):\n",
    "    \"\"\"Rebalance the portfolio based on pipeline output.\n",
    "\n",
    "    This function rebalances the portfolio by executing trades\n",
    "    according to the ranked stocks from the pipeline.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    context : object\n",
    "        An object to store global variables.\n",
    "    data : object\n",
    "        An object to fetch data.\n",
    "    \"\"\"\n",
    "    factor_data = context.factor_data\n",
    "    record(factor_data=factor_data.ranking)\n",
    "\n",
    "    assets = factor_data.index\n",
    "    record(prices=data.current(assets, 'price'))\n",
    "\n",
    "    longs = assets[factor_data.longs]\n",
    "    shorts = assets[factor_data.shorts]\n",
    "    divest = set(context.portfolio.positions.keys()) - set(longs.union(shorts))\n",
    "\n",
    "    exec_trades(data, assets=divest, target_percent=0)\n",
    "    exec_trades(data, assets=longs, target_percent=1 / N_LONGS)\n",
    "    exec_trades(data, assets=shorts, target_percent=-1 / N_SHORTS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8191036c",
   "metadata": {},
   "source": [
    "Function to execute trades for given assets and target percentages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c071a6b",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def exec_trades(data, assets, target_percent):\n",
    "    \"\"\"Execute trades for given assets and target percentages.\n",
    "\n",
    "    This function loops through each asset and executes trades\n",
    "    if the asset is tradeable and has no open orders.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    data : object\n",
    "        An object to fetch data.\n",
    "    assets : iterable\n",
    "        List of asset identifiers.\n",
    "    target_percent : float\n",
    "        Target portfolio weight for each asset.\n",
    "    \"\"\"\n",
    "    for asset in assets:\n",
    "        if data.can_trade(asset) and not get_open_orders(asset):\n",
    "            order_target_percent(asset, target_percent)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3078e517",
   "metadata": {},
   "source": [
    "Define the start and end dates for the backtest simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "545a10c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "start = pd.Timestamp('2016')\n",
    "end = pd.Timestamp('2018')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89fb2b32",
   "metadata": {},
   "source": [
    "Run the algorithm with the defined parameters and fetch performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8486b932",
   "metadata": {},
   "outputs": [],
   "source": [
    "perf = run_algorithm(\n",
    "    start=start,\n",
    "    end=end,\n",
    "    initialize=initialize,\n",
    "    before_trading_start=before_trading_start,\n",
    "    capital_base=100_000,\n",
    "    bundle=\"quandl\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8cb0a55",
   "metadata": {},
   "source": [
    "Output the final portfolio value after the simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76abaa0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "perf.portfolio_value"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0c61cb6",
   "metadata": {},
   "source": [
    "Construct a DataFrame with symbols as columns and dates as rows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "982e3994",
   "metadata": {},
   "outputs": [],
   "source": [
    "prices = pd.concat([df.to_frame(d) for d, df in perf.prices.dropna().items()], axis=1).T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "077d7376",
   "metadata": {},
   "source": [
    "Convert column names to strings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3941c22c",
   "metadata": {},
   "outputs": [],
   "source": [
    "prices.columns = [col.symbol for col in prices.columns]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c7c8b43",
   "metadata": {},
   "source": [
    "Normalize Timestamp to midnight, preserving TZ information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "108487c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "prices.index = prices.index.normalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31255068",
   "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
}
