{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-06-12T21:16:53.827640Z",
     "start_time": "2020-06-12T21:16:47.287113Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running typecheck\n",
      "Success: no issues found in 1 source file\n",
      "\n",
      "running flake8\n",
      "flake8 success\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 64324 characters]\n",
      "text/plain": [
       "<Figure size 1080x576 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 466280 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 6 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%checkall\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "import sys\n",
    "import concurrent\n",
    "import concurrent.futures\n",
    "from pyqstrat.pq_utils import has_display\n",
    "from pyqstrat.plot import Plot, Subplot, XYData, XYZData, SurfacePlotAttributes, LinePlotAttributes\n",
    "from typing import Mapping, Any, Callable, Generator, Tuple, Sequence, List, Optional\n",
    "\n",
    "\n",
    "class Experiment:\n",
    "    '''An Experiment stores a suggestion and its result'''\n",
    "    def __init__(self, suggestion: Mapping[str, Any], cost: float, other_costs: Mapping[str, float]) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            suggestion: A dictionary of variable name -> value\n",
    "            cost: A float representing output of the function we are testing with this suggestion as input.\n",
    "            other_costs: A dictionary of other results we want to store and look at later.\n",
    "\n",
    "        '''\n",
    "        self.suggestion = suggestion\n",
    "        self.cost = cost\n",
    "        self.other_costs = other_costs\n",
    "        \n",
    "    def valid(self) -> bool:\n",
    "        '''\n",
    "        Returns True if all suggestions and costs are finite, i.e not NaN or +/- Infinity\n",
    "        '''\n",
    "        if not all(np.isfinite(list(self.suggestion.values()))): return False\n",
    "        if not np.isfinite(self.cost): return False\n",
    "        if not all(np.isfinite(list(self.other_costs.values()))): return False\n",
    "        return True\n",
    "    \n",
    "    def __repr__(self) -> str:\n",
    "        return f'suggestion: {self.suggestion} cost: {self.cost} other costs: {self.other_costs}'\n",
    "\n",
    "\n",
    "class Optimizer:\n",
    "    '''Optimizer is used to optimize parameters for a strategy.'''\n",
    "    def __init__(self, name: str, \n",
    "                 generator: Generator[Mapping[str, Any], Tuple[float, Mapping[str, float]], None], \n",
    "                 cost_func: Callable[[Mapping[str, Any]], Tuple[float, Mapping[str, float]]], \n",
    "                 max_processes: int = None) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            name: Display title for plotting, etc.\n",
    "            generator: A generator (see Python Generators) that takes no inputs and yields a dictionary with parameter name -> parameter value.\n",
    "            cost_func: A function that takes a dictionary of parameter name -> parameter value as input and outputs cost for that set of parameters.\n",
    "            max_processes: If not set, the Optimizer will look at the number of CPU cores on your machine to figure out how many processes to run.\n",
    "        '''\n",
    "        self.name = name\n",
    "        self.generator = generator\n",
    "        self.cost_func = cost_func\n",
    "        import sys\n",
    "        if sys.platform in ['win32', 'cygwin']:\n",
    "            if max_processes is not None and max_processes != 1:\n",
    "                raise Exception(\"max_processes must be 1 on Microsoft Windows\")\n",
    "            max_processes = 1\n",
    "        self.max_processes = max_processes\n",
    "        self.experiments: List[Experiment] = []\n",
    "        \n",
    "    def _run_single_process(self) -> None:\n",
    "        try:\n",
    "            for suggestion in self.generator:\n",
    "                if suggestion is None: continue\n",
    "                cost, other_costs = self.cost_func(suggestion)\n",
    "                self.generator.send((cost, other_costs))\n",
    "                self.experiments.append(Experiment(suggestion, cost, other_costs))\n",
    "        except StopIteration:\n",
    "            # Exhausted generator\n",
    "            return\n",
    "    \n",
    "    # TODO: Needs to be rewritten to send costs back to generator when we do parallel gradient descent, etc.\n",
    "    def _run_multi_process(self, raise_on_error: bool) -> None:\n",
    "        fut_map = {}\n",
    "        with concurrent.futures.ProcessPoolExecutor(self.max_processes) as executor:\n",
    "            for suggestion in self.generator:\n",
    "                if suggestion is None: continue\n",
    "                future = executor.submit(self.cost_func, suggestion)\n",
    "                fut_map[future] = suggestion\n",
    "                \n",
    "            for future in concurrent.futures.as_completed(fut_map):\n",
    "                try:\n",
    "                    cost, other_costs = future.result()\n",
    "                except Exception as e:\n",
    "                    new_exc = type(e)(f'Exception: {str(e)} with suggestion: {suggestion}').with_traceback(sys.exc_info()[2])\n",
    "                    if raise_on_error: raise new_exc\n",
    "                    else: print(str(new_exc))\n",
    "                    continue\n",
    "                suggestion = fut_map[future]\n",
    "                self.experiments.append(Experiment(suggestion, cost, other_costs))\n",
    "    \n",
    "    def run(self, raise_on_error: bool = False) -> None:\n",
    "        '''Run the optimizer.\n",
    "        \n",
    "        Args:\n",
    "            raise_on_error: If set to True, even if we are running a multiprocess optimization, any Exceptions will bubble up and stop the Optimizer.\n",
    "              This can be useful for debugging to see stack traces for Exceptions.\n",
    "        '''\n",
    "        if self.max_processes == 1: self._run_single_process()\n",
    "        else: self._run_multi_process(raise_on_error)\n",
    "        \n",
    "    def experiment_list(self, sort_order: str = 'lowest_cost') -> Sequence[Experiment]:\n",
    "        '''Returns the list of experiments we have run\n",
    "        \n",
    "        Args:\n",
    "            sort_order: Can be set to lowest_cost, highest_cost or sequence.  \n",
    "              If set to sequence, experiments are returned in the sequence in which they were run\n",
    "        '''\n",
    "        if sort_order == 'lowest_cost':\n",
    "            experiments = sorted(self.experiments, key=lambda x: x.cost, reverse=True)\n",
    "        elif sort_order == 'highest_cost':\n",
    "            experiments = sorted(self.experiments, key=lambda x: x.cost, reverse=False)\n",
    "        elif sort_order == 'sequence':  # in order in which experiment was run\n",
    "            experiments = self.experiments\n",
    "        else:\n",
    "            raise Exception(f'invalid sort order: {sort_order}')\n",
    "        return experiments\n",
    "    \n",
    "    def df_experiments(self, sort_column: str = 'cost', ascending: bool = True) -> pd.DataFrame:\n",
    "        '''\n",
    "        Returns a dataframe containing experiment data, sorted by sort_column (default \"cost\")\n",
    "        '''\n",
    "        if len(self.experiments) == 0: return None\n",
    "        pc_keys = list(self.experiments[0].other_costs.keys())\n",
    "        sugg_keys = list(self.experiments[0].suggestion.keys())\n",
    "        records = [[exp.suggestion[k] for k in sugg_keys] + [exp.cost] + [exp.other_costs[k] for k in pc_keys] for exp in self.experiments if exp.valid()]\n",
    "        df = pd.DataFrame.from_records(records, columns=sugg_keys + ['cost'] + pc_keys)\n",
    "        df = df.sort_values(by=[sort_column], ascending=ascending)\n",
    "        return df\n",
    "    \n",
    "    def plot_3d(self, \n",
    "                x: str, \n",
    "                y: str, \n",
    "                z: str = 'all', \n",
    "                plot_type: str = 'surface', \n",
    "                figsize: Tuple[float, float] = (15, 15), \n",
    "                interpolation: str = 'linear', \n",
    "                cmap: str = 'viridis',\n",
    "                marker: str = 'X', \n",
    "                marker_size: int = 50,\n",
    "                marker_color: str = 'r',\n",
    "                xlim: Tuple[float, float] = None,\n",
    "                ylim: Tuple[float, float] = None, \n",
    "                hspace: float = None) -> None:\n",
    "        \n",
    "        \"\"\"Creates a 3D plot of the optimization output for plotting 2 parameters and costs.\n",
    "        \n",
    "        Args:\n",
    "            x: Name of the parameter to plot on the x axis, corresponding to the same name in the generator.\n",
    "            y: Name of the parameter to plot on the y axis, corresponding to the same name in the generator.\n",
    "            z: Can be one of:\n",
    "              \"cost\" \n",
    "              The name of another cost variable corresponding to the output from the cost function\n",
    "              \"all\", which creates a subplot for cost plus all other costs\n",
    "            plot_type: surface or contour (default surface)\n",
    "            figsize: Figure size\n",
    "            interpolation: Can be ‘linear’, ‘nearest’ or ‘cubic’ for plotting z points between the ones passed in.  See scipy.interpolate.griddata for details\n",
    "            cmap: Colormap to use (default viridis).  See matplotlib colormap for details\n",
    "            marker: Adds a marker to each point in x, y, z to show the actual data used for interpolation.  You can set this to None to turn markers off.\n",
    "            hspace: Vertical space between subplots\n",
    "         \"\"\"\n",
    "        \n",
    "        if len(self.experiments) == 0: \n",
    "            print('No experiments found')\n",
    "            return\n",
    "        if not has_display(): return\n",
    "\n",
    "        # Get rid of nans since matplotlib does not like them\n",
    "        experiments = [experiment for experiment in self.experiments if experiment.valid()]\n",
    "        if not len(experiments):\n",
    "            print('No valid experiments found')\n",
    "            return\n",
    "        \n",
    "        if xlim:\n",
    "            experiments = [experiment for experiment in experiments if experiment.suggestion[x] >= xlim[0] and experiment.suggestion[x] <= xlim[1]]\n",
    "        if ylim:\n",
    "            experiments = [experiment for experiment in experiments if experiment.suggestion[y] >= ylim[0] and experiment.suggestion[y] <= ylim[1]]\n",
    "\n",
    "        xvalues = [experiment.suggestion[x] for experiment in experiments]\n",
    "        yvalues = [experiment.suggestion[y] for experiment in experiments]\n",
    "        zvalues = []\n",
    "\n",
    "        if z == 'all':\n",
    "            zvalues.append(('cost', np.array([experiment.cost for experiment in experiments])))\n",
    "            if len(experiments[0].other_costs):\n",
    "                other_cost_keys = experiments[0].other_costs.keys()\n",
    "                for key in other_cost_keys:\n",
    "                    zvalues.append((key, np.array([experiment.other_costs[key] for experiment in experiments])))\n",
    "        elif z == 'cost':\n",
    "            zvalues.append(('cost', np.array([experiment.cost for experiment in experiments])))\n",
    "        else:\n",
    "            zvalues.append((z, np.array([experiment.other_costs[z] for experiment in experiments])))\n",
    "            \n",
    "        title: Optional[str]\n",
    "            \n",
    "        subplots = []\n",
    "        for tup in zvalues:\n",
    "            name = tup[0]\n",
    "            zarray = tup[1]\n",
    "            if plot_type == 'contour':\n",
    "                zlabel = None\n",
    "                title = name\n",
    "            else:\n",
    "                zlabel = name\n",
    "                title = None\n",
    "                \n",
    "            disp = SurfacePlotAttributes(marker=marker, marker_size=marker_size, marker_color=marker_color, interpolation=interpolation)\n",
    "            \n",
    "            subplots.append(Subplot(\n",
    "                data_list=[\n",
    "                    XYZData(name, xvalues, yvalues, zarray, display_attributes=disp)],\n",
    "                title=title, xlabel=x, ylabel=y, zlabel=zlabel, xlim=xlim, ylim=ylim))\n",
    "        plot = Plot(subplots, figsize=figsize, title='Optimizer 2D Test', hspace=hspace)\n",
    "        plot.draw()\n",
    "        \n",
    "    def plot_2d(self, \n",
    "                x: str, \n",
    "                y: str = 'all',\n",
    "                plot_type: str = 'line', \n",
    "                figsize: Tuple[float, float] = (15, 8),\n",
    "                marker: str = 'X', \n",
    "                marker_size: int = 50, \n",
    "                marker_color: str = 'r', \n",
    "                xlim: Tuple[float, float] = None, \n",
    "                hspace: float = None) -> None:\n",
    "        \"\"\"Creates a 2D plot of the optimization output for plotting 1 parameter and costs.\n",
    "        \n",
    "        Args:\n",
    "            x: Name of the parameter to plot on the x axis, corresponding to the same name in the generator.\n",
    "            y: Can be one of:\n",
    "              \"cost\" \n",
    "              The name of another cost variable corresponding to the output from the cost function\n",
    "              \"all\", which creates a subplot for cost plus all other costs\n",
    "            plot_type: line or scatter (default line)\n",
    "            figsize: Figure size\n",
    "            marker: Adds a marker to each point in x, y to show the actual data used for interpolation.  You can set this to None to turn markers off.\n",
    "            hspace: Vertical space between subplots\n",
    "         \"\"\"\n",
    "        if len(self.experiments) == 0: return\n",
    "        if not has_display(): return\n",
    "\n",
    "        # Get rid of nans since matplotlib does not like them\n",
    "        experiments = [experiment for experiment in self.experiments if experiment.valid()]\n",
    "        if xlim:\n",
    "            experiments = [experiment for experiment in experiments if experiment.suggestion[x] >= xlim[0] and experiment.suggestion[x] <= xlim[1]]\n",
    "\n",
    "        xvalues = [experiment.suggestion[x] for experiment in experiments]\n",
    "        yvalues = []\n",
    "\n",
    "        if y == 'all':\n",
    "            yvalues.append(('cost', np.array([experiment.cost for experiment in experiments])))\n",
    "            other_cost_keys = experiments[0].other_costs.keys()\n",
    "            for key in other_cost_keys:\n",
    "                yvalues.append((key, np.array([experiment.other_costs[key] for experiment in experiments])))\n",
    "        elif y == 'cost':\n",
    "            yvalues.append(('cost', np.array([experiment.cost for experiment in experiments])))\n",
    "        else:\n",
    "            yvalues.append((y, np.array([experiment.other_costs[y] for experiment in experiments])))\n",
    "            \n",
    "        xvalues = np.array(xvalues)\n",
    "        x_sort_indices = np.argsort(xvalues)\n",
    "        xvalues = xvalues[x_sort_indices]\n",
    "        subplots = []\n",
    "        for tup in yvalues:\n",
    "            name = tup[0]\n",
    "            yarray = tup[1]\n",
    "            yarray = yarray[x_sort_indices]\n",
    "            \n",
    "            disp = LinePlotAttributes(marker=marker, marker_size=marker_size, marker_color=marker_color)\n",
    "           \n",
    "            subplots.append(\n",
    "                Subplot(data_list=[XYData(name, xvalues, yarray, display_attributes=disp)], xlabel=x, ylabel=name, xlim=xlim))\n",
    "            \n",
    "        plot = Plot(subplots, figsize=figsize, title='Optimizer 1D Test')\n",
    "        plot.draw()\n",
    "        \n",
    "# Functions used in unit testing\n",
    "\n",
    "\n",
    "def _generator_1d() -> Generator[Mapping[str, Any], Tuple[float, Mapping[str, float]], None]:\n",
    "    for x in np.arange(0, np.pi * 2, 0.1):\n",
    "        _ = (yield {'x': x})\n",
    "\n",
    "\n",
    "def _cost_func_1d(suggestion: Mapping[str, Any]) -> Tuple[float, Mapping[str, float]]:\n",
    "    x = suggestion['x']\n",
    "    cost = np.sin(x)\n",
    "    ret = cost, {'std': -0.1 * cost}\n",
    "    return ret\n",
    "\n",
    "\n",
    "def _generator_2d() -> Generator[Mapping[str, Any], Tuple[float, Mapping[str, float]], None]:\n",
    "    for x in np.arange(0, np.pi * 2, 0.5):\n",
    "        for y in np.arange(0, np.pi * 2, 0.5):\n",
    "            _ = (yield {'x': x, 'y': y})\n",
    "\n",
    "\n",
    "def _cost_func_2d(suggestion: Mapping[str, Any]) -> Tuple[float, Mapping[str, float]]:\n",
    "    x = suggestion['x']\n",
    "    y = suggestion['y']\n",
    "    cost = np.sin(np.sqrt(x**2 + y ** 2))\n",
    "    return cost, {'sharpe': cost, 'std': -0.1 * cost}\n",
    "\n",
    "            \n",
    "def test_optimize():\n",
    "    max_processes = 1 if os.name == 'nt' else 4\n",
    "\n",
    "    optimizer_1d = Optimizer('test', _generator_1d(), _cost_func_1d, max_processes=max_processes)\n",
    "    optimizer_1d.run(raise_on_error=True)\n",
    "    optimizer_1d.plot_2d(x='x', plot_type='line', marker='o', marker_color='blue')\n",
    "    \n",
    "    optimizer_2d = Optimizer('test', _generator_2d(), _cost_func_2d, max_processes=max_processes)\n",
    "    optimizer_2d.run()\n",
    "    optimizer_2d.plot_3d(x='x', y='y')\n",
    "            \n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_optimize()\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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
