{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-08-15T00:40:48.444824Z",
     "start_time": "2018-08-15T00:40:44.999275Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 314720 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 51408 characters]\n",
      "text/plain": [
       "<Figure size 1080x576 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits import mplot3d\n",
    "import matplotlib.cm as cm\n",
    "from scipy.interpolate import griddata\n",
    "import os\n",
    "import sys\n",
    "import concurrent\n",
    "\n",
    "from pyqstrat.pq_utils import set_defaults, has_display\n",
    "from pyqstrat.plot import *\n",
    "\n",
    "set_defaults()\n",
    "\n",
    "class Experiment:\n",
    "    '''An Experiment stores a suggestion and its result\n",
    "    \n",
    "    Attributes:\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",
    "    def __init__(self, suggestion, cost, other_costs):\n",
    "        self.suggestion = suggestion\n",
    "        self.cost = cost\n",
    "        self.other_costs = other_costs\n",
    "        \n",
    "    def valid(self):\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):\n",
    "        return f'suggestion: {self.suggestion} cost: {self.cost} other costs: {self.other_costs}'\n",
    "\n",
    "class Optimizer:\n",
    "    '''Optimizer is used to optimize parameters for a strategy.'''\n",
    "    def __init__(self, name, generator, cost_func, max_processes = None):\n",
    "        '''\n",
    "        Args:\n",
    "            name: string used to display title in plotting, etc.\n",
    "            generator: A generator (see Python Generators) that takes no inputs and yields a list of dictionaries 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 = []\n",
    "        \n",
    "        \n",
    "    def _run_single_process(self):\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",
    "    \n",
    "    #TODO: Needs to be rewritten to send costs back to generator when we do parrallel gradient descent, etc.\n",
    "    def _run_multi_process(self, raise_on_error):\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",
    "                #print(f'completed suggestion: {suggestion} result: {cost} {other_costs}')\n",
    "                \n",
    "    \n",
    "    def run(self, raise_on_error = False):\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 = 'lowest_cost'):\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}')\n",
    "        return experiments\n",
    "    \n",
    "    def df_experiments(self, sort_column = 'cost', ascending = True):\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]\n",
    "        df = pd.DataFrame.from_records(records, columns = sugg_keys + ['cost'] + pc_keys)\n",
    "        df.sort_values(by = [sort_column], ascending = ascending, inplace = True)\n",
    "        return df\n",
    "    \n",
    "    def plot_3d(self, x, y, z = 'all', plot_type = 'surface', figsize = (15,15), interpolation = 'linear', \n",
    "             cmap = 'viridis', marker = 'X', marker_size = 50, marker_color = 'r', xlim = None, ylim = None, hspace = 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: 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",
    "        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",
    "            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[zname] for experiment in experiments])))\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",
    "            subplots.append(Subplot(data_list = [\n",
    "                XYZData(name, xvalues, yvalues, zarray, plot_type = plot_type, \n",
    "                        marker = marker, marker_size = marker_size, marker_color = marker_color, interpolation = interpolation, cmap = cmap\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, x, y = 'all', plot_type = 'line', figsize = (15,8), marker = 'X', marker_size = 50, marker_color = 'r', xlim = None, hspace = 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[zname] for experiment in experiments])))\n",
    "            \n",
    "        subplots = []\n",
    "        for tup in yvalues:\n",
    "            name = tup[0]\n",
    "            yarray = tup[1]\n",
    "            subplots.append(Subplot(data_list = [\n",
    "                XYData(name, xvalues, yarray, plot_type = plot_type, \n",
    "                        marker = marker, marker_size = marker_size, marker_color = marker_color\n",
    "                    )], xlabel = x, ylabel = name, xlim = xlim))\n",
    "        plot = Plot(subplots, figsize = figsize, title = 'Optimizer 1D Test')\n",
    "        plot.draw()\n",
    "            \n",
    "def test_optimize():\n",
    "    \n",
    "    def generator_2d():\n",
    "        for x in range(0,5):\n",
    "            for y in range(0,5):\n",
    "                costs = (yield {'x' : x, 'y' : y})\n",
    "                yield\n",
    "\n",
    "    def cost_func_2d(suggestion):\n",
    "        x = suggestion['x']\n",
    "        y = suggestion['y']\n",
    "        cost = x ** 2 + y ** 2\n",
    "        cost = np.sin(np.sqrt(x ** 2 + y ** 2))\n",
    "        return cost, {'sharpe' : cost, 'std' : -0.1 * cost}\n",
    "\n",
    "    optimizer_2d = Optimizer('test', generator_2d(), cost_func_2d, max_processes = 1)\n",
    "    optimizer_2d.run()\n",
    "    optimizer_2d.plot_3d(x = 'x', y = 'y')\n",
    "\n",
    "    def generator_1d():\n",
    "        for x in range(0,5):\n",
    "            costs = (yield {'x' : x})\n",
    "            yield\n",
    "\n",
    "    def cost_func_1d(suggestion):\n",
    "        x = suggestion['x']\n",
    "        cost = np.sin(x)\n",
    "        return cost, {'std' : -0.1 * cost}\n",
    "\n",
    "    optimizer_1d = Optimizer('test', generator_1d(), cost_func_1d, max_processes = 1)\n",
    "    optimizer_1d.run()\n",
    "    optimizer_1d.plot_2d(x = 'x', plot_type = 'line', marker = 'o', marker_color = 'blue')\n",
    "            \n",
    "if __name__ == \"__main__\":\n",
    "    test_optimize()\n"
   ]
  }
 ],
 "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
