{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-16T21:24:19.733504Z",
     "start_time": "2020-09-16T21:24:18.855184Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running flake8\n",
      "flake8 success\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>gmean</th>\n",
       "      <th>amean</th>\n",
       "      <th>std</th>\n",
       "      <th>shrp</th>\n",
       "      <th>srt</th>\n",
       "      <th>k</th>\n",
       "      <th>calmar</th>\n",
       "      <th>mar</th>\n",
       "      <th>mdd_pct</th>\n",
       "      <th>mdd_dates</th>\n",
       "      <th>dd_3y_pct</th>\n",
       "      <th>dd_3y_timestamps</th>\n",
       "      <th>up_dwn</th>\n",
       "      <th>2018</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>0.0635</td>\n",
       "      <td>0.0419</td>\n",
       "      <td>0.000899</td>\n",
       "      <td>2.93</td>\n",
       "      <td>5.69</td>\n",
       "      <td>3.49</td>\n",
       "      <td>8.73</td>\n",
       "      <td>8.73</td>\n",
       "      <td>0.00479</td>\n",
       "      <td>2018-01-19/2018-01-22</td>\n",
       "      <td>0.00479</td>\n",
       "      <td>2018-01-19/2018-01-22</td>\n",
       "      <td>34/25/0.576</td>\n",
       "      <td>0.0635</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   gmean   amean       std  shrp   srt     k calmar   mar  mdd_pct  \\\n",
       "  0.0635  0.0419  0.000899  2.93  5.69  3.49   8.73  8.73  0.00479   \n",
       "\n",
       "              mdd_dates dd_3y_pct       dd_3y_timestamps       up_dwn    2018  \n",
       "  2018-01-19/2018-01-22   0.00479  2018-01-19/2018-01-22  34/25/0.576  0.0635  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 78752 characters]\n",
      "text/plain": [
       "<Figure size 1080x576 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%flake8\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import statsmodels as sm\n",
    "import statsmodels.api as smapi\n",
    "import math\n",
    "from pyqstrat.pq_utils import monotonically_increasing, infer_frequency\n",
    "from pyqstrat.plot import TimeSeries, DateLine, Subplot, HorizontalLine, BucketedValues, Plot, LinePlotAttributes\n",
    "import matplotlib as mpl\n",
    "import matplotlib.figure as mpl_fig\n",
    "from typing import Tuple, Sequence, Mapping, MutableMapping, Optional, Any, Callable, Dict\n",
    "\n",
    "\n",
    "def compute_periods_per_year(timestamps: np.ndarray) -> float:\n",
    "    \"\"\"\n",
    "    Computes trading periods per year for an array of numpy datetime64's.\n",
    "        e.g. if most of the timestamps are separated by 1 day, will return 252.\n",
    "      \n",
    "    Args:\n",
    "        timestamps: a numpy array of datetime64's\n",
    "        \n",
    "    >>> compute_periods_per_year(np.array(['2018-01-01', '2018-01-02', '2018-01-03', '2018-01-09'], dtype='M8[D]'))\n",
    "    252.0\n",
    "    >>> round(compute_periods_per_year(np.array(['2018-01-01 10:00', '2018-01-01 10:05', '2018-01-01 10:10'], dtype='M8[m]')), 2)\n",
    "    72576.05\n",
    "    \"\"\"\n",
    "    if not len(timestamps): return np.nan\n",
    "    freq = infer_frequency(timestamps)\n",
    "    if freq == 31: return 12\n",
    "    return 252. / freq if freq != 0 else np.nan\n",
    "\n",
    "\n",
    "def compute_amean(returns: np.ndarray, periods_per_year: int) -> float:\n",
    "    '''\n",
    "    Computes arithmetic mean of a return array, ignoring NaNs\n",
    "    \n",
    "    Args:\n",
    "        returns: Represents returns at any frequency\n",
    "        periods_per_year: Frequency of the returns, e.g. 252 for daily returns\n",
    "        \n",
    "    >>> compute_amean(np.array([0.003, 0.004, np.nan]), 252)\n",
    "    0.882\n",
    "    '''\n",
    "    if not len(returns): return np.nan\n",
    "    return np.nanmean(returns) * periods_per_year\n",
    "\n",
    "\n",
    "def compute_num_periods(timestamps: np.ndarray, periods_per_year: float) -> float:\n",
    "    '''\n",
    "    Given an array of timestamps, we compute how many periods there are between the first and last element, where the length\n",
    "        of a period is defined by periods_per_year.  For example, if there are 6 periods per year, \n",
    "        then each period would be approx. 2 months long.\n",
    "        \n",
    "    Args:\n",
    "        timestamps (np.ndarray of np.datetime64): a numpy array of returns, can contain nans\n",
    "        periods_per_year: number of periods between first and last return\n",
    "         \n",
    "    >>> assert(compute_num_periods(np.array(['2015-01-01', '2015-03-01', '2015-05-01'], dtype='M8[D]'), 6) == 2)\n",
    "    '''\n",
    "    if not len(timestamps): return np.nan\n",
    "    assert(monotonically_increasing(timestamps))\n",
    "    fraction_of_year = (timestamps[-1] - timestamps[0]) / (np.timedelta64(1, 's') * 365 * 24 * 60 * 60)\n",
    "    return round(fraction_of_year * periods_per_year)\n",
    "    \n",
    "\n",
    "def compute_gmean(timestamps: np.ndarray, returns: np.ndarray, periods_per_year: float) -> float:\n",
    "    \"\"\"\n",
    "    Compute geometric mean of an array of returns\n",
    "    \n",
    "    Args:\n",
    "        returns: a numpy array of returns, can contain nans\n",
    "        periods_per_year: Used for annualizing returns\n",
    "        \n",
    "    >>> round(compute_gmean(np.array(['2015-01-01', '2015-03-01', '2015-05-01'], dtype='M8[D]'), np.array([0.001, 0.002, 0.003]), 252.), 6)\n",
    "    0.018362\n",
    "    \"\"\"\n",
    "    if not len(returns): return np.nan\n",
    "    assert(len(returns) == len(timestamps))\n",
    "    assert(isinstance(timestamps, np.ndarray) and isinstance(returns, np.ndarray))\n",
    "    mask = np.isfinite(returns)\n",
    "    timestamps = timestamps[mask]\n",
    "    returns = returns[mask]\n",
    "    num_periods = compute_num_periods(timestamps, periods_per_year)\n",
    "    g_mean = ((1.0 + returns).prod())**(1.0 / num_periods)\n",
    "    g_mean = np.power(g_mean, periods_per_year) - 1.0\n",
    "    return g_mean\n",
    "\n",
    "\n",
    "def compute_std(returns: np.ndarray) -> float:\n",
    "    \"\"\" Computes standard deviation of an array of returns, ignoring nans \"\"\"\n",
    "    if not len(returns): return np.nan\n",
    "    return np.nanstd(returns)\n",
    "\n",
    "\n",
    "def compute_sortino(returns: np.ndarray, amean: float, periods_per_year: float) -> float:\n",
    "    '''\n",
    "    Note that this assumes target return is 0.\n",
    "    \n",
    "    Args:\n",
    "        returns: a numpy array of returns\n",
    "        amean: arithmetic mean of returns\n",
    "        periods_per_year: number of trading periods per year\n",
    "        \n",
    "    >>> print(round(compute_sortino(np.array([0.001, -0.001, 0.002]), 0.001, 252), 6))\n",
    "    0.133631\n",
    "    '''\n",
    "    if not len(returns) or not np.isfinite(amean) or periods_per_year <= 0: return np.nan\n",
    "    returns = np.where((~np.isfinite(returns)), 0.0, returns)\n",
    "    normalized_rets = np.where(returns > 0.0, 0.0, returns)\n",
    "    sortino_denom = np.std(normalized_rets)\n",
    "    sortino = np.nan if sortino_denom == 0 else amean / (sortino_denom * np.sqrt(periods_per_year))\n",
    "    return sortino\n",
    "\n",
    "\n",
    "def compute_sharpe(returns: np.ndarray, amean: float, periods_per_year: float) -> float:\n",
    "    '''\n",
    "    Note that this does not take into risk free returns so it's really a sharpe0, i.e. assumes risk free returns are 0\n",
    "    \n",
    "    Args:\n",
    "        returns: a numpy array of returns\n",
    "        amean: arithmetic mean of returns\n",
    "        periods_per_year: number of trading periods per year\n",
    "        \n",
    "    >>> round(compute_sharpe(np.array([0.001, -0.001, 0.002]), 0.001, 252), 6)\n",
    "    0.050508\n",
    "    '''\n",
    "    if not len(returns) or not np.isfinite(amean) or periods_per_year <= 0: return np.nan\n",
    "    returns = np.where((~np.isfinite(returns)), 0.0, returns)\n",
    "    s = np.std(returns)\n",
    "    sharpe = np.nan if s == 0 else amean / (s * np.sqrt(periods_per_year))\n",
    "    return sharpe\n",
    "\n",
    "\n",
    "def compute_k_ratio(equity: np.ndarray, periods_per_year: int, halflife_years: float = None) -> float:\n",
    "    '''\n",
    "    Compute k-ratio (2013 or original versions by Lars Kestner). See https://papers.ssrn.com/sol3/papers.cfm?abstract_id=2230949\n",
    "    We also implement a modification that allows higher weighting for more recent returns.\n",
    "    \n",
    "    Args:\n",
    "        equity: a numpy array of the equity in your account\n",
    "        periods_per_year: 252 for daily values\n",
    "        halflife_years: If set, we use weighted linear regression to give less weight to older returns.\n",
    "            In this case, we compute the original k-ratio which does not use periods per year or number of observations\n",
    "            If not set, we compute the 2013 version of the k-ratio which weights k-ratio by sqrt(periods_per_year) / nobs\n",
    "        \n",
    "    Returns:\n",
    "        weighted or unweighted k-ratio\n",
    "        \n",
    "    >>> np.random.seed(0)\n",
    "    >>> t = np.arange(1000)\n",
    "    >>> ret = np.random.normal(loc = 0.0025, scale = 0.01, size = len(t))\n",
    "    >>> equity = (1 + ret).cumprod()\n",
    "    >>> assert(math.isclose(compute_k_ratio(equity, 252, None), 3.888, abs_tol=0.001))\n",
    "    >>> assert(math.isclose(compute_k_ratio(equity, 252, 0.5), 602.140, abs_tol=0.001))\n",
    "    '''\n",
    "    equity = equity[np.isfinite(equity)]\n",
    "    equity = np.log(equity)\n",
    "    t = np.arange(len(equity))\n",
    "    if halflife_years:\n",
    "        halflife = halflife_years * periods_per_year\n",
    "        k = math.log(0.5) / halflife\n",
    "        w = np.empty(len(equity), dtype=np.float)\n",
    "        w = np.exp(k * t)\n",
    "        w = w ** 2  # Statsmodels requires square of weights\n",
    "        w = w[::-1]\n",
    "        fit = sm.regression.linear_model.WLS(endog=equity, exog=t, weights=w, hasconst=False).fit()\n",
    "        k_ratio = fit.params[0] / fit.bse[0]\n",
    "    else:\n",
    "        fit = smapi.OLS(endog=equity, exog=np.arange(len(equity)), hasconst=False).fit()\n",
    "        k_ratio = fit.params[0] * math.sqrt(periods_per_year) / (fit.bse[0] * len(equity))\n",
    "\n",
    "    return k_ratio\n",
    "\n",
    "\n",
    "def compute_equity(timestamps: np.ndarray, starting_equity: float, returns: np.ndarray) -> np.ndarray:\n",
    "    ''' Given starting equity, timestamps and returns, create a numpy array of equity at each date'''\n",
    "    return starting_equity * np.cumprod(1. + returns)\n",
    "\n",
    "\n",
    "def compute_rolling_dd(timestamps: np.ndarray, equity: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:\n",
    "    '''\n",
    "    Compute numpy array of rolling drawdown percentage\n",
    "    \n",
    "    Args:\n",
    "        timestamps: numpy array of datetime64\n",
    "        equity: numpy array of equity\n",
    "    '''\n",
    "    assert(len(timestamps) == len(equity))\n",
    "    if not len(timestamps): return np.array([], dtype='M8[ns]'), np.array([], dtype=np.float)\n",
    "    s = pd.Series(equity, index=timestamps)\n",
    "    rolling_max = s.expanding(min_periods=1).max()\n",
    "    dd = np.where(s >= rolling_max, 0.0, -(s - rolling_max) / rolling_max)\n",
    "    return timestamps, dd\n",
    "\n",
    "\n",
    "def compute_maxdd_pct(rolling_dd: np.ndarray) -> float:\n",
    "    '''Compute max drawdown percentage given a numpy array of rolling drawdowns, ignoring NaNs'''\n",
    "    if not len(rolling_dd): return np.nan\n",
    "    return np.nanmax(rolling_dd)\n",
    "\n",
    "\n",
    "def compute_maxdd_date(rolling_dd_dates: np.ndarray, rolling_dd: np.ndarray) -> float:\n",
    "    ''' Compute date of max drawdown given numpy array of timestamps, and corresponding rolling dd percentages'''\n",
    "    if not len(rolling_dd_dates): return pd.NaT\n",
    "    assert(len(rolling_dd_dates) == len(rolling_dd))\n",
    "    return rolling_dd_dates[np.argmax(rolling_dd)]\n",
    "\n",
    "\n",
    "def compute_maxdd_start(rolling_dd_dates: np.ndarray, rolling_dd: np.ndarray, mdd_date: np.datetime64) -> np.datetime64:\n",
    "    '''Compute date when max drawdown starts, given numpy array of timestamps corresponding rolling dd \n",
    "        percentages and date of the max draw down'''\n",
    "    if not len(rolling_dd_dates) or pd.isnull(mdd_date): return pd.NaT\n",
    "    assert(len(rolling_dd_dates) == len(rolling_dd))\n",
    "    maxdd_dates = rolling_dd_dates[(rolling_dd <= 0) & (rolling_dd_dates < mdd_date)]\n",
    "    if not len(maxdd_dates): return pd.NaT\n",
    "    return maxdd_dates[-1]\n",
    "\n",
    "\n",
    "def compute_mar(returns: np.ndarray, periods_per_year: float, mdd_pct: float) -> float:\n",
    "    '''Compute MAR ratio, which is annualized return divided by biggest drawdown since inception.'''\n",
    "    if not len(returns) or np.isnan(mdd_pct) or mdd_pct == 0: return np.nan\n",
    "    return np.mean(returns) * periods_per_year / mdd_pct\n",
    "\n",
    "\n",
    "def compute_dates_3yr(timestamps: np.ndarray) -> np.ndarray:\n",
    "    ''' Given an array of numpy datetimes, return those that are within 3 years of the last date in the array'''\n",
    "    if not len(timestamps): return np.array([], dtype='M8[D]')\n",
    "    last_date = timestamps[-1]\n",
    "    d = pd.to_datetime(last_date)\n",
    "    start_3yr = np.datetime64(d.replace(year=d.year - 3))\n",
    "    return timestamps[timestamps > start_3yr]\n",
    "\n",
    "\n",
    "def compute_returns_3yr(timestamps: np.ndarray, returns: np.ndarray) -> np.ndarray:\n",
    "    '''Given an array of numpy datetimes and an array of returns, return those that are within 3 years \n",
    "        of the last date in the datetime array '''\n",
    "    if not len(timestamps): return np.array([], dtype=np.float)\n",
    "    assert(len(timestamps) == len(returns))\n",
    "    timestamps_3yr = compute_dates_3yr(timestamps)\n",
    "    return returns[timestamps >= timestamps_3yr[0]]\n",
    "\n",
    "\n",
    "def compute_rolling_dd_3yr(timestamps: np.ndarray, equity: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:\n",
    "    '''Compute rolling drawdowns over the last 3 years'''\n",
    "    if not len(timestamps): return np.array([], dtype='M8[D]')\n",
    "    last_date = timestamps[-1]\n",
    "    d = pd.to_datetime(last_date)\n",
    "    start_3yr = np.datetime64(d.replace(year=d.year - 3))\n",
    "    equity = equity[timestamps >= start_3yr]\n",
    "    timestamps = timestamps[timestamps >= start_3yr]\n",
    "    return compute_rolling_dd(timestamps, equity)\n",
    "\n",
    "\n",
    "def compute_maxdd_pct_3yr(rolling_dd_3yr: np.ndarray) -> float:\n",
    "    '''Compute max drawdown percentage over the last 3 years'''\n",
    "    return compute_maxdd_pct(rolling_dd_3yr)\n",
    "\n",
    "\n",
    "def compute_maxdd_date_3yr(rolling_dd_3yr_timestamps: np.ndarray, rolling_dd_3yr: np.ndarray) -> np.datetime64:\n",
    "    '''Compute max drawdown date over the last 3 years'''\n",
    "    return compute_maxdd_date(rolling_dd_3yr_timestamps, rolling_dd_3yr)\n",
    "\n",
    "\n",
    "def compute_maxdd_start_3yr(rolling_dd_3yr_timestamps: np.ndarray, rolling_dd_3yr: np.ndarray, mdd_date_3yr: np.datetime64) -> np.datetime64:\n",
    "    '''Comput max drawdown start date over the last 3 years'''\n",
    "    return compute_maxdd_start(rolling_dd_3yr_timestamps, rolling_dd_3yr, mdd_date_3yr)\n",
    "\n",
    "\n",
    "def compute_calmar(returns_3yr: np.ndarray, periods_per_year: float, mdd_pct_3yr: float) -> float:\n",
    "    '''Compute Calmar ratio, which is the annualized return divided by max drawdown over the last 3 years'''\n",
    "    return compute_mar(returns_3yr, periods_per_year, mdd_pct_3yr)\n",
    "\n",
    "\n",
    "def compute_bucketed_returns(timestamps: np.ndarray, returns: np.ndarray) -> Tuple[Sequence[int], Sequence[np.ndarray]]:\n",
    "    '''\n",
    "    Bucket returns by year\n",
    "    \n",
    "    Returns:\n",
    "        A tuple with the first element being a list of years and the second a list of \n",
    "            numpy arrays containing returns for each corresponding year\n",
    "    '''\n",
    "    assert(len(timestamps) == len(returns))\n",
    "    if not len(timestamps): return np.array([], dtype=np.str), np.array([], dtype=np.float)\n",
    "    s = pd.Series(returns, index=timestamps)\n",
    "    years_list = []\n",
    "    rets_list = []\n",
    "    for year, rets in s.groupby(s.index.map(lambda x: x.year)):\n",
    "        years_list.append(year)\n",
    "        rets_list.append(rets.values)\n",
    "    \n",
    "    return years_list, rets_list\n",
    "\n",
    "\n",
    "def compute_annual_returns(timestamps: np.ndarray, returns: np.ndarray, periods_per_year: float) -> Tuple[np.ndarray, np.ndarray]:\n",
    "    '''Takes the output of compute_bucketed_returns and returns geometric mean of returns by year\n",
    "    \n",
    "    Returns:\n",
    "        A tuple with the first element being an array of years (integer) and the second element \n",
    "        an array of annualized returns for those years\n",
    "        \n",
    "    '''\n",
    "    assert(len(timestamps) == len(returns) and periods_per_year > 0)\n",
    "    if not len(timestamps): return np.array([], dtype=np.str), np.array([], dtype=np.float)\n",
    "    df = pd.DataFrame({'ret': returns, 'timestamp': timestamps})\n",
    "    years = []\n",
    "    gmeans = []\n",
    "    for k, g in df.groupby(df.timestamp.map(lambda x: x.year)):\n",
    "        years.append(k)\n",
    "        gmeans.append(compute_gmean(g.timestamp.values, g.ret.values, periods_per_year))\n",
    "    return np.array(years), np.array(gmeans)\n",
    "\n",
    "\n",
    "class Evaluator:\n",
    "    \"\"\"You add functions to the evaluator that are dependent on the outputs of other functions.  \n",
    "    The evaluator will call these functions in the right order\n",
    "    so dependencies are computed first before the functions that need their output.  \n",
    "    You can retrieve the output of a metric using the metric member function\n",
    "    \n",
    "    >>> evaluator = Evaluator(initial_metrics={'x': np.array([1, 2, 3]), 'y': np.array([3, 4, 5])})\n",
    "    >>> evaluator.add_metric('z', lambda x, y: sum(x, y), dependencies=['x', 'y'])\n",
    "    >>> evaluator.compute()\n",
    "    >>> evaluator.metric('z')\n",
    "    array([ 9, 10, 11])\n",
    "    \"\"\"\n",
    "    def __init__(self, initial_metrics: Dict[str, Any]) -> None:\n",
    "        \"\"\"Inits Evaluator with a dictionary of initial metrics that are used to compute subsequent metrics\n",
    "        \n",
    "        Args:\n",
    "            initial_metrics: a dictionary of string name -> metric.  metric can be any object including a scalar, \n",
    "                an array or a tuple\n",
    "        \"\"\"\n",
    "        assert(type(initial_metrics) == dict)\n",
    "        self.metric_values: Dict[str, Any] = initial_metrics.copy()\n",
    "        self._metrics: MutableMapping[str, Tuple[Callable, Sequence[str]]] = {}\n",
    "        \n",
    "    def add_metric(self, name: str, func: Callable, dependencies: Sequence[str]) -> None:\n",
    "        self._metrics[name] = (func, dependencies)\n",
    "    \n",
    "    def compute(self, metric_names: Sequence[str] = None) -> None:\n",
    "        '''Compute metrics using the internal dependency graph\n",
    "        \n",
    "        Args:\n",
    "            metric_names: an array of metric names.  If not passed in, evaluator will compute and store all metrics\n",
    "        '''\n",
    "\n",
    "        if metric_names is None: metric_names = list(self._metrics.keys())\n",
    "        for metric_name in metric_names:\n",
    "            self.compute_metric(metric_name)\n",
    "            \n",
    "    def compute_metric(self, metric_name: str) -> None:\n",
    "        '''\n",
    "        Compute and store a single metric:\n",
    "        \n",
    "        Args:\n",
    "            metric_name: string representing the metric to compute\n",
    "        '''\n",
    "        func, dependencies = self._metrics[metric_name]\n",
    "        for dependency in dependencies:\n",
    "            if dependency not in self.metric_values:\n",
    "                self.compute_metric(dependency)\n",
    "        dependency_values = {k: self.metric_values[k] for k in dependencies}\n",
    "        \n",
    "        values = func(**dependency_values)\n",
    "            \n",
    "        self.metric_values[metric_name] = values\n",
    "                \n",
    "    def metric(self, metric_name: str) -> Any:\n",
    "        '''Return the value of a single metric given its name'''\n",
    "        return self.metric_values[metric_name]\n",
    "    \n",
    "    def metrics(self) -> Mapping[str, Any]:\n",
    "        '''Return a dictionary of metric name -> metric value'''\n",
    "        return self.metric_values\n",
    "    \n",
    "\n",
    "def handle_non_finite_returns(timestamps: np.ndarray,\n",
    "                              rets: np.ndarray,\n",
    "                              leading_non_finite_to_zeros: bool,\n",
    "                              subsequent_non_finite_to_zeros: bool) -> Tuple[np.ndarray, np.ndarray]: \n",
    "    '''\n",
    "    >>> np.set_printoptions(formatter={'float': '{: .6g}'.format})\n",
    "    >>> timestamps = np.arange(np.datetime64('2019-01-01'), np.datetime64('2019-01-07'))\n",
    "    >>> rets = np.array([np.nan, np.nan, 3, 4, np.nan, 5])\n",
    "    >>> handle_non_finite_returns(timestamps, rets, leading_non_finite_to_zeros = False, subsequent_non_finite_to_zeros = True)\n",
    "    (array(['2019-01-03', '2019-01-04', '2019-01-05', '2019-01-06'], dtype='datetime64[D]'), array([ 3,  4,  0,  5]))\n",
    "    >>> handle_non_finite_returns(timestamps, rets, leading_non_finite_to_zeros = True, subsequent_non_finite_to_zeros = False)\n",
    "    (array(['2019-01-01', '2019-01-02', '2019-01-03', '2019-01-04', '2019-01-06'], dtype='datetime64[D]'), array([ 0,  0,  3,  4,  5]))\n",
    "    >>> handle_non_finite_returns(timestamps, rets, leading_non_finite_to_zeros = False, subsequent_non_finite_to_zeros = False)\n",
    "    (array(['2019-01-01', '2019-01-02', '2019-01-03', '2019-01-04', '2019-01-06'], dtype='datetime64[D]'), array([ 0,  0,  3,  4,  5]))\n",
    "    >>> rets = np.array([1, 2, 3, 4, 4.5,  5])\n",
    "    >>> handle_non_finite_returns(timestamps, rets, leading_non_finite_to_zeros = False, subsequent_non_finite_to_zeros = True)\n",
    "    (array(['2019-01-01', '2019-01-02', '2019-01-03', '2019-01-04', '2019-01-05', '2019-01-06'], \n",
    "        dtype='datetime64[D]'), array([ 1,  2,  3,  4,  4.5,  5]))\n",
    "    '''\n",
    "    \n",
    "    first_non_nan_index = np.ravel(np.nonzero(~np.isnan(rets)))\n",
    "    if len(first_non_nan_index):\n",
    "        first_non_nan_index = first_non_nan_index[0]\n",
    "    else:\n",
    "        first_non_nan_index = -1\n",
    "    \n",
    "    if first_non_nan_index > 0 and first_non_nan_index < len(rets):\n",
    "        if leading_non_finite_to_zeros:\n",
    "            rets[:first_non_nan_index] = np.nan_to_num(rets[:first_non_nan_index])\n",
    "        else:\n",
    "            timestamps = timestamps[first_non_nan_index:]\n",
    "            rets = rets[first_non_nan_index:]\n",
    "    \n",
    "    if subsequent_non_finite_to_zeros:\n",
    "        rets = np.nan_to_num(rets)\n",
    "    else:\n",
    "        timestamps = timestamps[np.isfinite(rets)]\n",
    "        rets = rets[np.isfinite(rets)]\n",
    "    \n",
    "    return timestamps, rets\n",
    "    \n",
    "\n",
    "def compute_return_metrics(timestamps: np.ndarray,\n",
    "                           rets: np.ndarray,\n",
    "                           starting_equity: float,\n",
    "                           leading_non_finite_to_zeros: bool = False, \n",
    "                           subsequent_non_finite_to_zeros: bool = True) -> Evaluator:\n",
    "    '''\n",
    "    Compute a set of common metrics using returns (for example, of an instrument or a portfolio)\n",
    "    \n",
    "    Args:\n",
    "        timestamps (np.array of datetime64): Timestamps for the returns\n",
    "        rets (nd.array of float): The returns, use 0.01 for 1%\n",
    "        starting_equity (float): Starting equity value in your portfolio\n",
    "        leading_non_finite_to_zeros (bool, optional): If set, we replace leading nan, inf, -inf returns with zeros.  \n",
    "            For example, you may need a warmup period for moving averages.  Default False\n",
    "        subsequent_non_finite_to_zeros (bool, optional): If set, we replace any nans that follow the first non nan value with zeros.\n",
    "            There may be periods where you have no prices but removing these returns would result in incorrect annualization. \n",
    "            Default True\n",
    "         \n",
    "    Returns:\n",
    "        An Evaluator object containing computed metrics off the returns passed in.  \n",
    "        If needed, you can add your own metrics to this object based on the values of existing metrics and recompute the Evaluator.\n",
    "        Otherwise, you can just use the output of the evaluator using the metrics function.\n",
    "        \n",
    "    >>> timestamps = np.array(['2015-01-01', '2015-03-01', '2015-05-01', '2015-09-01'], dtype='M8[D]')\n",
    "    >>> rets = np.array([0.01, 0.02, np.nan, -0.015])\n",
    "    >>> starting_equity = 1.e6\n",
    "    >>> ev = compute_return_metrics(timestamps, rets, starting_equity)\n",
    "    >>> metrics = ev.metrics()\n",
    "    >>> assert(round(metrics['gmean'], 6) == 0.021061)\n",
    "    >>> assert(round(metrics['sharpe'], 6) == 0.599382)\n",
    "    >>> assert(all(metrics['returns_3yr'] == np.array([0.01, 0.02, 0, -0.015])))\n",
    "    '''\n",
    "    \n",
    "    assert(starting_equity > 0.)\n",
    "    assert(type(rets) == np.ndarray and rets.dtype == np.float64)\n",
    "    assert(type(timestamps) == np.ndarray and np.issubdtype(timestamps.dtype, np.datetime64) and monotonically_increasing(timestamps))\n",
    "    \n",
    "    timestamps, rets = handle_non_finite_returns(timestamps, rets, leading_non_finite_to_zeros, subsequent_non_finite_to_zeros)\n",
    "\n",
    "    ev = Evaluator({'timestamps': timestamps, 'returns': rets, 'starting_equity': starting_equity})\n",
    "    ev.add_metric('periods_per_year', compute_periods_per_year, dependencies=['timestamps'])\n",
    "    ev.add_metric('amean', compute_amean, dependencies=['returns', 'periods_per_year'])\n",
    "    ev.add_metric('std', compute_std, dependencies=['returns'])\n",
    "    ev.add_metric('up_periods', lambda returns: len(returns[returns > 0]), dependencies=['returns'])\n",
    "    ev.add_metric('down_periods', lambda returns: len(returns[returns < 0]), dependencies=['returns'])\n",
    "    ev.add_metric('up_pct', \n",
    "                  lambda up_periods, down_periods: up_periods * 1.0 / (up_periods + down_periods) if (up_periods + down_periods) != 0 else np.nan, \n",
    "                  dependencies=['up_periods', 'down_periods'])\n",
    "    ev.add_metric('gmean', compute_gmean, dependencies=['timestamps', 'returns', 'periods_per_year'])\n",
    "    ev.add_metric('sharpe', compute_sharpe, dependencies=['returns', 'periods_per_year', 'amean'])\n",
    "    ev.add_metric('sortino', compute_sortino, dependencies=['returns', 'periods_per_year', 'amean'])\n",
    "    ev.add_metric('equity', compute_equity, dependencies=['timestamps', 'starting_equity', 'returns'])\n",
    "    ev.add_metric('k_ratio', compute_k_ratio, dependencies=['equity', 'periods_per_year'])\n",
    "    ev.add_metric('k_ratio_weighted', lambda equity, periods_per_year: compute_k_ratio(equity, periods_per_year, 3),\n",
    "                  dependencies=['equity', 'periods_per_year'])\n",
    "\n",
    "    # Drawdowns\n",
    "    ev.add_metric('rolling_dd', compute_rolling_dd, dependencies=['timestamps', 'equity'])\n",
    "    ev.add_metric('mdd_pct', lambda rolling_dd: compute_maxdd_pct(rolling_dd[1]), dependencies=['rolling_dd'])\n",
    "    ev.add_metric('mdd_date', lambda rolling_dd: compute_maxdd_date(rolling_dd[0], rolling_dd[1]), dependencies=['rolling_dd'])\n",
    "    ev.add_metric('mdd_start', lambda rolling_dd, mdd_date: compute_maxdd_start(rolling_dd[0], rolling_dd[1], mdd_date), \n",
    "                  dependencies=['rolling_dd', 'mdd_date'])\n",
    "    ev.add_metric('mar', compute_mar, dependencies=['returns', 'periods_per_year', 'mdd_pct'])\n",
    "    \n",
    "    ev.add_metric('timestamps_3yr', compute_dates_3yr, dependencies=['timestamps'])\n",
    "    ev.add_metric('returns_3yr', compute_returns_3yr, dependencies=['timestamps', 'returns'])\n",
    "\n",
    "    ev.add_metric('rolling_dd_3yr', compute_rolling_dd_3yr, dependencies=['timestamps', 'equity'])\n",
    "    ev.add_metric('mdd_pct_3yr', lambda rolling_dd_3yr: compute_maxdd_pct_3yr(rolling_dd_3yr[1]), dependencies=['rolling_dd_3yr'])\n",
    "    ev.add_metric('mdd_date_3yr', lambda rolling_dd_3yr: compute_maxdd_date_3yr(rolling_dd_3yr[0], rolling_dd_3yr[1]), \n",
    "                  dependencies=['rolling_dd_3yr'])\n",
    "    ev.add_metric('mdd_start_3yr', lambda rolling_dd_3yr, mdd_date_3yr: \n",
    "                  compute_maxdd_start_3yr(rolling_dd_3yr[0], rolling_dd_3yr[1], mdd_date_3yr), \n",
    "                  dependencies=['rolling_dd_3yr', 'mdd_date_3yr'])\n",
    "    ev.add_metric('calmar', compute_calmar, dependencies=['returns_3yr', 'periods_per_year', 'mdd_pct_3yr'])\n",
    "\n",
    "    ev.add_metric('annual_returns', compute_annual_returns, dependencies=['timestamps', 'returns', 'periods_per_year'])\n",
    "    ev.add_metric('bucketed_returns', compute_bucketed_returns, dependencies=['timestamps', 'returns'])\n",
    "\n",
    "    ev.compute()\n",
    "    return ev\n",
    "\n",
    "\n",
    "def display_return_metrics(metrics: Mapping[str, Any], float_precision: int = 3) -> pd.DataFrame:\n",
    "    '''\n",
    "    Creates a dataframe making it convenient to view the output of the metrics obtained using the compute_return_metrics function.\n",
    "    \n",
    "    Args:\n",
    "        float_precision: Change if you want to display floats with more or less significant figures than the default, \n",
    "            3 significant figures.       \n",
    "    Returns:\n",
    "        A one row dataframe with formatted metrics.\n",
    "    '''\n",
    "    from IPython.core.display import display\n",
    "    \n",
    "    _metrics = {}\n",
    "    cols = ['gmean', 'amean', 'std', 'shrp', 'srt', 'k', 'calmar', 'mar', 'mdd_pct', 'mdd_start', 'mdd_date', 'dd_3y_pct', \n",
    "            'up_periods', 'down_periods', 'up_pct', 'mdd_start_3yr', 'mdd_date_3yr']\n",
    "    \n",
    "    translate = {'shrp': 'sharpe', 'srt': 'sortino', 'dd_3y_pct': 'mdd_pct_3yr', 'k': 'k_ratio'}\n",
    "    for col in cols:\n",
    "        key = col\n",
    "        if col in translate: key = translate[col]\n",
    "        _metrics[col] = metrics[key]\n",
    "            \n",
    "    _metrics['mdd_dates'] = f'{str(metrics[\"mdd_start\"])[:10]}/{str(metrics[\"mdd_date\"])[:10]}'\n",
    "    _metrics['up_dwn'] = f'{metrics[\"up_periods\"]}/{metrics[\"down_periods\"]}/{metrics[\"up_pct\"]:.3g}'\n",
    "    _metrics['dd_3y_timestamps'] = f'{str(metrics[\"mdd_start_3yr\"])[:10]}/{str(metrics[\"mdd_date_3yr\"])[:10]}'\n",
    "    \n",
    "    years = metrics['annual_returns'][0]\n",
    "    ann_rets = metrics['annual_returns'][1]\n",
    "    for i, year in enumerate(years):\n",
    "        _metrics[str(year)] = ann_rets[i]\n",
    "        \n",
    "    format_str = '{:.' + str(float_precision) + 'g}'\n",
    "        \n",
    "    for k, v in _metrics.items():\n",
    "        if isinstance(v, np.float) or isinstance(v, float):\n",
    "            _metrics[k] = format_str.format(v)\n",
    "       \n",
    "    cols = ['gmean', 'amean', 'std', 'shrp', 'srt', 'k', 'calmar', 'mar', 'mdd_pct', 'mdd_dates', 'dd_3y_pct', 'dd_3y_timestamps', 'up_dwn'] + [\n",
    "        str(year) for year in sorted(years)]\n",
    "    \n",
    "    df = pd.DataFrame(index=[''])\n",
    "    for metric_name, metric_value in _metrics.items():\n",
    "        df.insert(0, metric_name, metric_value)\n",
    "    df = df[cols]\n",
    "        \n",
    "    display(df)\n",
    "    return df\n",
    "\n",
    "\n",
    "def plot_return_metrics(metrics: Mapping[str, Any], \n",
    "                        title: str = None, \n",
    "                        disp_attribs: LinePlotAttributes = None,\n",
    "                        drawdown_lines: bool = True, \n",
    "                        zero_line: bool = True,\n",
    "                        show_date_gaps: bool = True) -> Optional[Tuple[mpl_fig.Figure, mpl.axes.Axes]]:\n",
    "    '''\n",
    "    Plot equity, rolling drawdowns and and a boxplot of annual returns given the output of compute_return_metrics.\n",
    "    '''\n",
    "    timestamps = metrics['timestamps']\n",
    "    equity = metrics['equity']\n",
    "    equity = TimeSeries('equity', timestamps=timestamps, values=equity, display_attributes=disp_attribs)\n",
    "    mdd_date, mdd_start = metrics['mdd_start'], metrics['mdd_date']\n",
    "    mdd_date_3yr, mdd_start_3yr = metrics['mdd_start_3yr'], metrics['mdd_date_3yr']\n",
    "    \n",
    "    if drawdown_lines:\n",
    "        date_lines = [DateLine(name='max dd', date=mdd_start, color='red'),\n",
    "                      DateLine(date=mdd_date, color='red'),\n",
    "                      DateLine(name='3y dd', date=mdd_start_3yr, color='orange'),\n",
    "                      DateLine(date=mdd_date_3yr, color='orange')]\n",
    "    else:\n",
    "        date_lines = []\n",
    "        \n",
    "    horizontal_lines = [HorizontalLine(metrics['starting_equity'], color='black')] if zero_line else []\n",
    "    equity_subplot = Subplot(equity, ylabel='Equity', height_ratio=0.6, log_y=True, y_tick_format='${x:,.0f}', \n",
    "                             date_lines=date_lines, horizontal_lines=horizontal_lines)     \n",
    "\n",
    "    rolling_dd = TimeSeries('drawdowns', timestamps=metrics['rolling_dd'][0], values=metrics['rolling_dd'][1],\n",
    "                            display_attributes=disp_attribs)\n",
    "    horizontal_lines = [HorizontalLine(y=0, color='black')] if zero_line else []\n",
    "    dd_subplot = Subplot(rolling_dd, ylabel='Drawdowns', height_ratio=0.2, date_lines=date_lines, horizontal_lines=horizontal_lines)\n",
    "    years = metrics['bucketed_returns'][0]\n",
    "    ann_rets = metrics['bucketed_returns'][1]\n",
    "    ann_ret = BucketedValues('annual returns', bucket_names=years, bucket_values=ann_rets)\n",
    "    ann_ret_subplot = Subplot(ann_ret, ylabel='Annual Returns', height_ratio=0.2, horizontal_lines=horizontal_lines)\n",
    "    \n",
    "    plt = Plot([equity_subplot, dd_subplot, ann_ret_subplot], title=title, show_date_gaps=show_date_gaps)\n",
    "    return plt.draw()\n",
    " \n",
    "\n",
    "def test_evaluator() -> None:\n",
    "    from datetime import datetime, timedelta\n",
    "    np.random.seed(10)\n",
    "    timestamps = np.arange(datetime(2018, 1, 1), datetime(2018, 3, 1), timedelta(days=1))\n",
    "    rets = np.random.normal(size=len(timestamps)) / 1000\n",
    "    starting_equity = 1.e6\n",
    "    \n",
    "    ev = compute_return_metrics(timestamps, rets, starting_equity)\n",
    "    display_return_metrics(ev.metrics())\n",
    "    plot_return_metrics(ev.metrics(), zero_line=False)\n",
    "    \n",
    "    assert(round(ev.metric('sharpe'), 6) == 2.932954)\n",
    "    assert(round(ev.metric('sortino'), 6) == 5.690878)\n",
    "    assert(ev.metric('annual_returns')[0] == [2018])\n",
    "    assert(round(ev.metric('annual_returns')[1][0], 6) == [0.063530])\n",
    "    assert(ev.metric('mdd_start') == np.datetime64('2018-01-19'))\n",
    "    assert(ev.metric('mdd_date') == np.datetime64('2018-01-22'))\n",
    "    \n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_evaluator()\n",
    "    import doctest\n",
    "    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
