Dynamic programming class.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving the HJB\n",
    "The HJB equation, is used in dynamic programming to solve optimisation problem. Optimisation problems occur in all walks of life and some can even be solved. And some of those that can be solved are best solved with dynamic programming, a recursive evaluation of the best descission.\n",
    "\n",
    "This notebook aims to explain by way of demonstration how to solve the HJB and fidn the optimal set of descissions, but also provide an easy to use base to formuate and solve your optimisation problems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import time\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "v(s) = \\min_x(cost(x,s) + v(new state(x,s)))\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class DynamicProgram(object):\n",
    "    \"\"\"\n",
    "    Generate a dynamic program to find a set of optimal descissions using the HJB.\n",
    "    \n",
    "    define the program by:\n",
    "    \n",
    "    Setting intial states via:            set_inital_state(list or int)\n",
    "    Setting the number of steps via:      set_step_number(int) \n",
    "    \n",
    "    Add a set of descissions:             add_decisions_set(set)\n",
    "    Add a cost function:                  add_cost_function(function in terms of state )\n",
    "    \n",
    "    Add a state change equation:          add_state_eq(function(state))\n",
    "    Add an expression for the last value: add_final_value_expression(function(state,settings))\n",
    "    Add limits on the states:             add_state_limits(lower=list or int,upper = list or int)\n",
    "\n",
    "    See below for examples:\n",
    "    \n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        self.settings = {\n",
    "                        'Lower state limits' : [],\n",
    "                        'Upper state limits' : [],    \n",
    "                        'x_set' : set(),\n",
    "                        'cache' : {},}\n",
    "        \n",
    "    def add_state_eq(self,function):\n",
    "        \"\"\"Returns a tuple describing the states. \n",
    "        Remember to increment the first state, b convention the number of steps.\n",
    "        Load additional parameters (usually needed for cost and state value with global variables)\n",
    "        \n",
    "        Example of a function that changes the state by the decission:\n",
    "        def new_state(x,s):\n",
    "            return (s[0]+1,s[1]+x) #Return a tuple, use (s[:-1]+(5,)) to slice tuples.  \n",
    "                \n",
    "        \"\"\"\n",
    "        self.settings['State eq.'] = function\n",
    "        \n",
    "    def add_cost_function(self,function):\n",
    "        \"\"\"\n",
    "        Returns a float or integer describing the cost value.\n",
    "        Load additional parameters (usually needed for cost and state value with global variables)\n",
    "        \n",
    "        Example is a function that simply returns the decision as cost:\n",
    "        def cost(x,s):\n",
    "            return x\n",
    "            \n",
    "        \"\"\"\n",
    "        self.settings['Cost eq.'] = function\n",
    "    \n",
    "    def add_final_value_expression(self, function,):\n",
    "        \"\"\"\n",
    "        Returns a float or integer as the final value:\n",
    "        \n",
    "        Example is a function that returns the ratio of the initial state and the final state:\n",
    "        \n",
    "        def val_T(s,settings):\n",
    "            return s[1]/float(settings['Initial state'][1])\n",
    "        \"\"\"\n",
    "        self.settings['Final value'] = function\n",
    "        \n",
    "        \n",
    "    def set_step_number(self,step_number):\n",
    "        \"\"\"Number of stages / steps. Integer\"\"\"\n",
    "        self.settings['T'] = step_number\n",
    "        \n",
    "    def set_inital_state(self,intial_values):\n",
    "        \"\"\"Provide the inital state of the states other than the stage number\"\"\"\n",
    "        if type(intial_values) is list:\n",
    "            self.settings['Initial state'] = intial_values\n",
    "            self.settings['Initial state'].insert(0,0)\n",
    "        elif type(intial_values) is int:\n",
    "            self.settings['Initial state'] = [intial_values]\n",
    "            self.settings['Initial state'].insert(0,0)\n",
    "           \n",
    "        self.settings['Initial state'] = tuple(self.settings['Initial state'])\n",
    "\n",
    "        \n",
    "    def add_state_limits(self,lower=[],upper=[]):\n",
    "        \"\"\"Add the limits on the state other than the stage number, leave empty if none\"\"\"\n",
    "        if type(lower) is list:\n",
    "            self.settings['Lower state limits'].extend(lower)\n",
    "            self.settings['Upper state limits'].extend(upper)\n",
    "        elif type(lower) is int:\n",
    "            self.settings['Lower state limits'] = [lower]\n",
    "            self.settings['Upper state limits'] = [upper]\n",
    "        \n",
    "\n",
    "        \n",
    "    def solve(self):\n",
    "        \"\"\"\n",
    "        Solves the HJB. Returns the optimal value.\n",
    "        \n",
    "        Path and further info is stored in the cache. Access it via \n",
    "        retrieve_decisions()\n",
    "        \"\"\"\n",
    "        self.settings['cache'] ={}\n",
    "        return self._hjb_(self.settings['Initial state'])\n",
    "    \n",
    "    def retrieve_decisions(self):\n",
    "        \"\"\"\n",
    "        Retrieve the decisions that led to the optimal value\n",
    "        \n",
    "        Returns the cost for the different states, the optimal schedule and the states\n",
    "        that the schedule results in.\n",
    "        \"\"\"\n",
    "        sched = np.ones(self.settings['T'])*np.nan\n",
    "        cost_calc= 0\n",
    "        states = []\n",
    "\n",
    "        s = self.settings['Initial state']\n",
    "        t = 0\n",
    "\n",
    "        while t < self.settings['T']:\n",
    "            sched[t] = self.settings['cache'][s][1]\n",
    "\n",
    "            cost_calc += self.settings['Cost eq.'](sched[t],s)\n",
    "            states.append(s[1:])\n",
    "\n",
    "            s = self.settings['State eq.'](sched[t],s)\n",
    "            t += 1\n",
    "\n",
    "        states.append(s[1:])\n",
    "\n",
    "        return cost_calc, sched, states\n",
    "    \n",
    "    def return_settings(self):\n",
    "        return self.settings\n",
    "    \n",
    "    def return_cache(self):\n",
    "        return self.settings['cache']\n",
    "        \n",
    "    def add_decisions_set(self,set_of_decisions):\n",
    "        \"\"\"\n",
    "        Add a set of permissable decissions. Must be set of unique integers.\n",
    "        \"\"\"\n",
    "        if set(set_of_decisions) != set_of_decisions:\n",
    "            raise TypeError('Expected a set unique values, use set() to declare a set')\n",
    "        self.settings['x_set'] = set(set_of_decisions)\n",
    "        \n",
    "    def add_setting(self,key,value):\n",
    "        self.settings[key] = value\n",
    "       \n",
    "    def _hjb_(self,s):\n",
    "        if self.settings['cache'].has_key(s):\n",
    "            return self.settings['cache'][s][0]\n",
    "\n",
    "        # check state bounds\n",
    "        for c,i in enumerate(s[1:]):\n",
    "            if i < self.settings['Lower state limits'][c] or i > self.settings['Upper state limits'][c]:\n",
    "                return float('inf')\n",
    "            \n",
    "        #Check if reached time step limit:\n",
    "        if s[0] == self.settings['T']:            \n",
    "            m = self.settings['Final value'](s,self.settings)\n",
    "            self.settings['cache'][s] = [m, np.nan]\n",
    "            return m\n",
    "\n",
    "        # Else enter recursion\n",
    "        else:\n",
    "            ### Make decision variable vector ###\n",
    "            ######################################################################################\n",
    "            # Faster but only with integer decisions\n",
    "            p=[]\n",
    "            for x in self.settings['x_set']:\n",
    "                p.append(self.settings['Cost eq.'](x,s)+self._hjb_(self.settings['State eq.'](x,s)))\n",
    "                \n",
    "            m = min(p)\n",
    "\n",
    "                \n",
    "            ### Slower but with any imutable decisions, uncomment if desired: ###\n",
    "#             p ={}\n",
    "#             for x in self.settings['x_set']:\n",
    "#                 p[x] = self.settings['Cost eq.'](x,s)+self._hjb_(self.settings['State eq.'](x,s))\n",
    "               \n",
    "#             m = min(p, key=p.get)\n",
    "            ########################################################################################\n",
    "            \n",
    "             \n",
    "            ##############################################################    \n",
    "            ## Finding the index of the best solution ##\n",
    "            for x in self.settings['x_set']:\n",
    "                if m == p[x]:\n",
    "                    pp = x\n",
    "            \n",
    "            ################################################################\n",
    "\n",
    "\n",
    "            self.settings['cache'][s] = [m, pp]\n",
    "\n",
    "            return m\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "We can also solve classic dynamic programming problems such as the knapsack problem, hannoi towers or the fibonacci number calculation. Blank functions are outlined below. \n",
    "\n",
    "The functions must fulfill a range of conditions:\n",
    "$$\n",
    "f:\\mathcal{S}^n\\times x \\rightarrow  \\mathbb{R}\n",
    "$$\n",
    "where $\\mathcal{S}$ is the set of permissable states, $x$ the decision variable and $n$ the number of dimensions of the state. These are defined by $x \\in \\mathcal{X}$ the \\texttt{set()} of permissable decisions. The state $s \\in \\mathcal{S}^n$ where $n \\geq 2$ and $\\mathcal{S} \\subset \\mathbb{R} $ and is finite. \n",
    "\n",
    "The new state is defined by a function such that: \n",
    "$$\n",
    "f:\\mathcal{S}^n\\times x \\rightarrow  \\mathcal{S}^n\n",
    "$$\n",
    "\n",
    "The value of the final stage is defined as:\n",
    "$$\n",
    "f:\\mathcal{S}^n \\rightarrow  \\mathbb{R}\n",
    "$$\n",
    "which as an impimentation detail has the required argument settings, to which features can be added through `add_setting(key,value)` where key must be a new and unique dictionary key and value may be any permissable dictionary entry.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#help(DynamicProgram)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def cost(x,s):\n",
    "    \"\"\"Return a float or integer\"\"\"\n",
    "    pass\n",
    "\n",
    "def new_state(x,s):\n",
    "    \"\"\"Return a tuple\"\"\"\n",
    "    pass\n",
    "\n",
    "def val_T(s,settings):\n",
    "    \"\"\"Return a float or int\"\"\"\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can solve a very simple pump optimsiation where the state of water in a tank is given by h and described by:\n",
    "$$\n",
    "s_{new} = \\begin{cases} (t+1,h-1) & \\text{if } x = 0  \\\\ (t+1,h+1) & \\text{if } x = 1 \\\\ (t+1,h+1.5) & \\text{if } x = 2\\end{cases}\n",
    "$$\n",
    "The operating cost are described by:\n",
    "$$\n",
    "cost = tarrif(t)\\times x\n",
    "$$\n",
    "where $x$ is the descission variable.\n",
    "The final value is given by:\n",
    "$$\n",
    "V_T = \\begin{cases} 0 & \\text{if: } h_T \\geq h_0 \\\\ Inf &\\text{otherwise} \\end{cases}\n",
    "$$\n",
    "\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def simple_cost(x,s):\n",
    "    tariff = [19,  8,  20,  3, 12, 14,  0,  4,  3, 13, 11, 13, 13, 11, 16, 14, 16,\n",
    "       19,  1,  8,  0,  4,   0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,\n",
    "       12,  3, 18, 15,  3, 10, 12,  6,  3,  5, 11,  0, 11,  8, 10, 11,  5,\n",
    "       15,  8,  2,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,\n",
    "        9, 10, 13,  7,  7,  1, 12,  2,  2,  1,  5,  8,  4,  0, 11,  2,  5,\n",
    "       16,  8,  1, 17, 16,  3,  0,  4, 16,  0,  7]\n",
    "    return tariff[s[0]]*x\n",
    "\n",
    "def val_T(s,settings):\n",
    "    if s[1] < settings['Initial state'][1]:\n",
    "        return float('inf')\n",
    "    else:\n",
    "        return 0\n",
    "    \n",
    "def simple_state(x,s):\n",
    "    #print s\n",
    "    if x == 0:\n",
    "        return (s[0]+1,s[1]-1)\n",
    "    elif x == 1:\n",
    "        return (s[0]+1,s[1]+1)\n",
    "    elif x == 2:\n",
    "        return (s[0]+1,s[1]+1.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Cost eq.': <function __main__.simple_cost>,\n",
       " 'Final value': <function __main__.val_T>,\n",
       " 'Initial state': (0, 100),\n",
       " 'Lower state limits': [0],\n",
       " 'State eq.': <function __main__.simple_state>,\n",
       " 'T': 96,\n",
       " 'Upper state limits': [200],\n",
       " 'cache': {},\n",
       " 'x_set': {0, 1, 2}}"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pumping = DynamicProgram()\n",
    "pumping.set_step_number(96)\n",
    "pumping.add_decisions_set({0,1,2})\n",
    "pumping.add_cost_function(simple_cost)\n",
    "pumping.add_state_eq(simple_state)\n",
    "pumping.add_final_value_expression(val_T)\n",
    "pumping.add_state_limits(lower=0,upper = 200)\n",
    "pumping.set_inital_state(100)\n",
    "pumping.return_settings()\n",
    "\n",
    "#pumping.retrieve_decisions()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 loop, best of 3: 279 ms per loop\n"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(32.0, array([ 0.,  1.,  0.,  1.,  0.,  0.,  1.,  1.,  1.,  0.,  0.,  0.,  0.,\n",
       "         0.,  0.,  0.,  0.,  0.,  2.,  1.,  2.,  1.,  2.,  2.]), [(199,),\n",
       "  (198,),\n",
       "  (199,),\n",
       "  (198,),\n",
       "  (199,),\n",
       "  (198,),\n",
       "  (197,),\n",
       "  (198,),\n",
       "  (199,),\n",
       "  (200,),\n",
       "  (199,),\n",
       "  (198,),\n",
       "  (197,),\n",
       "  (196,),\n",
       "  (195,),\n",
       "  (194,),\n",
       "  (193,),\n",
       "  (192,),\n",
       "  (191,),\n",
       "  (192.5,),\n",
       "  (193.5,),\n",
       "  (195.0,),\n",
       "  (196.0,),\n",
       "  (197.5,),\n",
       "  (199.0,)])"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can have more than one state variable. For example we can add a second tank and now pump to either of them:\n",
    "Here they have similar equations, but they can be completly independant, of each other. Any cost and state function that meets the requirments above is allowed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "s_{new} = \\begin{cases} (t+1,h_1-1,h_2-1) & \\text{if } x = 0  \\\\ (t+1,h_1+1,h_2) & \\text{if } x = 1 \\\\ (t+1,h_1,h_2+2) & \\text{if } x = 2\\end{cases}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def simple_state2(x,s):\n",
    "    if x == 0:\n",
    "        return (s[0]+1,s[1]-1,s[2]-1)\n",
    "    elif x == 1:\n",
    "        return (s[0]+1,s[1]+1,s[2])\n",
    "    elif x == 2:\n",
    "        return (s[0]+1,s[1]  ,s[2]+2)\n",
    "\n",
    "# We also need to update the final value function.\n",
    "def val_T2(s,settings):\n",
    "    if s[1] < settings['Initial state'][1] or s[2] < settings['Initial state'][2]:\n",
    "        return float('inf') \n",
    "    else:\n",
    "        return 0\n",
    "    \n",
    "# For now we leave the cost function the same, but that could also be changed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "One final example is the checkerboard problem as outlined here: https://en.wikipedia.org/wiki/Dynamic_programming#Checkerboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(6, array([ 1.,  0., -1.]), [(2,), (3,), (3,), (2,)])\n"
     ]
    }
   ],
   "source": [
    "board = np.array([[1,3,4,6],\n",
    "                  [2,6,2,1],\n",
    "                  [7,3,2,1],\n",
    "                  [0,4,2,9]])\n",
    "\n",
    "def cost(x,s):\n",
    "    \"\"\"Return a float or integer\"\"\"\n",
    "    global board\n",
    "    return board[s[0],s[1]]\n",
    "    \n",
    "\n",
    "def new_state(x,s):\n",
    "    \"\"\"Return a tuple\"\"\"\n",
    "    global board\n",
    "    return (int(s[0]+1),int(s[1]+x)) \n",
    "\n",
    "def val_T(s,settings):\n",
    "    \"\"\"Return a float or int\"\"\"\n",
    "    global board\n",
    "    return board[s[0],s[1]]\n",
    "\n",
    "Checkerboard = DynamicProgram()\n",
    "Checkerboard.add_cost_function(cost)\n",
    "Checkerboard.add_decisions_set({-1,0,1})\n",
    "Checkerboard.add_final_value_expression(val_T)\n",
    "Checkerboard.add_state_eq(new_state)\n",
    "Checkerboard.add_state_limits(lower=-1,upper=3)\n",
    "Checkerboard.set_inital_state(2)\n",
    "Checkerboard.set_step_number(3)\n",
    "Checkerboard.solve()\n",
    "print Checkerboard.retrieve_decisions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "An example we can solve is the water allocation problem from the tutorial sheets:\n",
    "\n",
    "Consider a water supply allocation problem. Suppose that a quantity Q can be\n",
    "allocated to three water users (indices j=1, 2 and 3), what is allocation x4 which\n",
    "maximises the total net benefits?\n",
    "The gross benefit resulting from the allocation of $x_j$ to j is:\n",
    "\n",
    "$$\n",
    "a_j(1-exp(-b_jx_j)) \n",
    "$$\n",
    "Subject to $ (a_j, b_j > 0). $\n",
    "Moreover, the costs of this allocation are:\n",
    "\n",
    "$c_jx_j^{d_j}$ ($c_j$, $d_j >0$ and $d_j<1$, because of economy of scale).\n",
    "\n",
    "\n",
    "\n",
    "The values of the constants are:\n",
    "Q=5\n",
    "\n",
    "|   | $a_j$|$b_j$|$c_j$|$d_j$|\n",
    "|---|---|---|---|---|\n",
    "|j=1 |100|0.1| 10|0.6|\n",
    "|j=2 | 50|0.4| 10|0.8|\n",
    "|j=3 |100|0.2| 25|0.4|\n",
    "\n",
    "Solve the problem in the discrete case (i.e. assuming that $x_j$ is an integer).\n",
    "\n",
    "The optimisation problem is therefore:\n",
    "$$\n",
    "Maximise  \\sum_j (a_j(1-exp(-b_jx_j)) - c_jx_j^{d_j}\n",
    "$$\n",
    "$$\n",
    "subject to: \\sum_j x_j \\leq Q\n",
    "$$\n",
    "$$\n",
    "x_j \\geq 0\n",
    "$$\n",
    "\n",
    "This problem is a non-linear mixed integer problem and quite expensive to solve as branch and bound problem. But with DP it is much easier. (The continous case, may be much easier as NLP than DP, this shows how important it is to pay attenttion ot the problem type)\n",
    "\n",
    "The decisions are 0 ... 5\n",
    "The states are the amount of water allocated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "Costs = np.array([[100,0.1, 10,0.6],\n",
    "                              [50, 0.4, 10,0.8],\n",
    "                              [100,0.2, 25,0.4]])\n",
    "\n",
    "def value(x,s):\n",
    "    global Costs\n",
    "    return Costs[s[0],0]*(1-math.exp(-Costs[s[0],1]*x))\n",
    "    \n",
    "\n",
    "def cost(x,s):\n",
    "    \"\"\"Return a float or integer\"\"\"\n",
    "    global Costs\n",
    "    if \n",
    "    return   -value(x,s) +Costs[s[0],2]*x**Costs[s[0],3]\n",
    "    \n",
    "\n",
    "def new_state(x,s):\n",
    "    \"\"\"Return a tuple\"\"\"\n",
    "    return s[1]-x\n",
    "    \n",
    "    \n",
    "def val_T(s,settings):\n",
    "    \"\"\"Return a float or int\"\"\"\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Stochastic Programming"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "v(s,i) = \\min_x( cost(x,s,i)+\\sum_j (p_{i,j} \\times v(newstate(x,s,j))) )\n",
    "$$\n",
    "\n",
    "Where the probability $p_{i,j}$ is the probaility of jumping from state $i$ to state $j$. Currently the transition matrix is invariate, however this can be easily implimented with P as a list of lists.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class StochasticProgram(DynamicProgram):\n",
    "    \"\"\"\n",
    "    Adds a stochastic component to the dynamic program.\n",
    "    \n",
    "    state now is: s where s[0] is the step s[1:-1] is the states of the system and s[-1] is the stochastic state\n",
    "    \n",
    "    The transition matrix for the markov chain describing the stochastic bhavior is added by:\n",
    "    add_transition_matrix(P) with P as a list of lists.\n",
    "    \n",
    "    \n",
    "    \"\"\"\n",
    "    def add_transition_matrix(self,P):\n",
    "        \"\"\"\n",
    "        Add the transition matrix as list of lists\n",
    "        \n",
    "        eg. P = [[0.4,0.5,0.1],\n",
    "                 [0.2,0.6,0.2],\n",
    "                 [0.1,0.5,0.4]]\n",
    "        \n",
    "        \"\"\"\n",
    "        self.settings['P'] = np.array(P)\n",
    "        self.settings['Len P'] = len(P)\n",
    "        \n",
    "    def retrieve_decisions(self):\n",
    "        \"\"\"\n",
    "        Retrieve the decisions that led to the optimal value\n",
    "        \n",
    "        Returns the cost for the different states, the optimal schedule and the states that the schedule results in.\n",
    "        \"\"\"\n",
    "        schedule = []\n",
    "        cost_calc= np.zeros(self.settings['Len P'])\n",
    "        states = []\n",
    "        \n",
    "        for i in range(self.settings['Len P']):\n",
    "            schedule_part = []\n",
    "            states_part = []\n",
    "            s = self.settings['Initial state']\n",
    "            t = 0\n",
    "\n",
    "            while t < self.settings['T']:\n",
    "                schedule_part.append(self.settings['cache'][s][1])\n",
    "                cost_calc[i] += self.settings['Cost eq.'](schedule_part[t],s)\n",
    "                states_part.append(s[1:])\n",
    "\n",
    "                s = self.settings['State eq.'](schedule_part[t],(s[:-1]+(i,) ) )\n",
    "                t += 1\n",
    "\n",
    "            states_part.append(s[1:])\n",
    "            states.append(states_part)\n",
    "            schedule.append(schedule_part)\n",
    "\n",
    "        return cost_calc, schedule, states\n",
    "\n",
    "\n",
    "    def solve(self):\n",
    "        \"\"\"\n",
    "        Solves the HJB. Returns the optimal value.\n",
    "        \n",
    "        Path and further info is stored in the cache. Access it via \n",
    "        retrieve_decisions()\n",
    "        \"\"\"\n",
    "        return self._hjb_stoch_(self.settings['Initial state'])\n",
    "    \n",
    "        \n",
    "    def _hjb_stoch_(self,s):\n",
    "        if self.settings['cache'].has_key(s):\n",
    "            return self.settings['cache'][s][0]\n",
    "\n",
    "        # check state bounds\n",
    "        for c,i in enumerate(s[1:-1]):\n",
    "            if i < self.settings['Lower state limits'][c] or i > self.settings['Upper state limits'][c]:\n",
    "                return float('inf')\n",
    "            \n",
    "        #Check if reached time step limit:\n",
    "        if s[0] == self.settings['T']:            \n",
    "            m = self.settings['Final value'](s,self.settings)\n",
    "            self.settings['cache'][s] = [m, np.nan]\n",
    "            return m\n",
    "\n",
    "        # Else enter recursion\n",
    "        else:\n",
    "            \n",
    "            p=[]\n",
    "            for x in self.settings['x_set']:\n",
    "                #future = 0\n",
    "#                 for i in range(self.settings['Len P']):\n",
    "#                     #print self.hjb_stoch(self.settings['State eq.'](x,(s[0:-1]+(i,)))), self.settings['P'][s[-1],i]\n",
    "#                     future += self.hjb_stoch(self.settings['State eq.'](x,(s[0:-1]+(i,))))*self.settings['P'][s[-1],i]\n",
    "                \n",
    "                \n",
    "                \n",
    "                future = sum(self._hjb_stoch_(self.settings['State eq.'](x,(s[0:-1]+ (i,))))\n",
    "                          *self.settings['P'][s[-1],i] for i in range(self.settings['Len P']))\n",
    "            \n",
    "                p.append(self.settings['Cost eq.'](x,s) + future)\n",
    "        \n",
    "            m = min(p)\n",
    "\n",
    "            for x in self.settings['x_set']:\n",
    "                if m == p[x]:\n",
    "                    pp = x\n",
    "\n",
    "            self.settings['cache'][s] = [m, pp]\n",
    "\n",
    "            return m\n",
    "\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on class StochasticProgram in module __main__:\n",
      "\n",
      "class StochasticProgram(DynamicProgram)\n",
      " |  Adds a stochastic component to the dynamic program.\n",
      " |  \n",
      " |  state now is: s where s[0] is the step s[1:-1] is the states of the system and s[-1] is the stochastic state\n",
      " |  \n",
      " |  The transition matrix for the markov chain describing the stochastic bhavior is added by:\n",
      " |  add_transition_matrix(P) with P as a list of lists.\n",
      " |  \n",
      " |  Method resolution order:\n",
      " |      StochasticProgram\n",
      " |      DynamicProgram\n",
      " |      __builtin__.object\n",
      " |  \n",
      " |  Methods defined here:\n",
      " |  \n",
      " |  add_transition_matrix(self, P)\n",
      " |      Add the transition matrix as list of lists\n",
      " |      \n",
      " |      eg. P = [[0.4,0.5,0.1],\n",
      " |               [0.2,0.6,0.2],\n",
      " |               [0.1,0.5,0.4]]\n",
      " |  \n",
      " |  retrieve_decisions(self)\n",
      " |      Retrieve the decisions that led to the optimal value\n",
      " |      \n",
      " |      Returns the cost for the different states, the optimal schedule and the states that the schedule results in.\n",
      " |  \n",
      " |  solve(self)\n",
      " |      Solves the HJB. Returns the optimal value.\n",
      " |      \n",
      " |      Path and further info is stored in the cache. Access it via \n",
      " |      retrieve_decisions()\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Methods inherited from DynamicProgram:\n",
      " |  \n",
      " |  __init__(self)\n",
      " |  \n",
      " |  add_cost_function(self, function)\n",
      " |      Returns a float or integer describing the cost value.\n",
      " |      Load additional parameters (usually needed for cost and state value with global variables)\n",
      " |      \n",
      " |      Example is a function that simply returns the decision as cost:\n",
      " |      def cost(x,s):\n",
      " |          return x\n",
      " |  \n",
      " |  add_decisions_set(self, set_of_decisions)\n",
      " |      Add a set of permissable decissions. Must be set of unique integers.\n",
      " |  \n",
      " |  add_final_value_expression(self, function)\n",
      " |      Returns a float or integer as the final value:\n",
      " |      \n",
      " |      Example is a function that returns the ratio of the initial state and the final state:\n",
      " |      \n",
      " |      def val_T(s,settings):\n",
      " |          return s[1]/float(settings['Initial state'][1])\n",
      " |  \n",
      " |  add_setting(self, key, value)\n",
      " |  \n",
      " |  add_state_eq(self, function)\n",
      " |      Returns a tuple describing the states. \n",
      " |      Remember to increment the first state, b convention the number of steps.\n",
      " |      Load additional parameters (usually needed for cost and state value with global variables)\n",
      " |      \n",
      " |      Example of a function that changes the state by the decission:\n",
      " |      def new_state(x,s):\n",
      " |          return (s[0]+1,s[1]+x) #Return a tuple, use (s[:-1]+(5,)) to slice tuples.\n",
      " |  \n",
      " |  add_state_limits(self, lower=[], upper=[])\n",
      " |      Add the limits on the state other than the stage number, leave empty if none\n",
      " |  \n",
      " |  return_cache(self)\n",
      " |  \n",
      " |  return_settings(self)\n",
      " |  \n",
      " |  set_inital_state(self, intial_values)\n",
      " |      Provide the inital state of the states other than the stage number\n",
      " |  \n",
      " |  set_step_number(self, step_number)\n",
      " |      Number of stages / steps. Integer\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data descriptors inherited from DynamicProgram:\n",
      " |  \n",
      " |  __dict__\n",
      " |      dictionary for instance variables (if defined)\n",
      " |  \n",
      " |  __weakref__\n",
      " |      list of weak references to the object (if defined)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "help(StochasticProgram)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cost of operating a pump with a given wind turbine power input given by a certain state is given by:\n",
    "$$\n",
    "cost(x,t,h,j) := \\begin{cases} T(t) \\times (x \\times P_p - W(t,j)) & \\text{if} +ve \\\\ \n",
    "E_{xp} \\times (x \\times P_p - W(t,j)) & \\text{if} -ve\\end{cases}\n",
    "$$\n",
    "where $W(t,j)$ is the wind power output at time $t$ with an error state $j$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Convention s = t,h,j\n",
    "def stoch_simple_state(x,s):\n",
    "    #print s\n",
    "    if x == 0:\n",
    "        return (s[0]+1,s[1]-1,s[2])\n",
    "    elif x == 1:\n",
    "        return (s[0]+1,s[1]+1,s[2])\n",
    "    elif x == 2:        \n",
    "        return (s[0]+1,s[1]+1.5,s[2])\n",
    "    \n",
    "\n",
    "def err_corr_wind_power_cost(x,s):\n",
    "    Tariff = [5,5,5,5,5,8,8,8,8,8,12,12,12,12,12,50,50,50,50,20,20,6,5,5]\n",
    "    Wind = [46,  1,  3, 36, 30, 19,  9, 26, 35,  5, 49,  3,  6, 36, 43, 36, 14,\n",
    "       34,  2,  0,  0, 30, 13, 36]\n",
    "    \n",
    "    diff = np.array([-1,0,1])*3\n",
    "    \n",
    "    Export_price = 5.5\n",
    "    \n",
    "    wind_out = Wind[s[0]]+diff[s[2]]\n",
    "    if wind_out <= 0:\n",
    "        wind_out = 0\n",
    "    \n",
    "    power_con = x*60-wind_out\n",
    "    if power_con >= 0:\n",
    "        return power_con*Tariff[s[0]]\n",
    "    else:\n",
    "        return power_con*Export_price\n",
    "    \n",
    "    \n",
    "def val_T(s,settings):\n",
    "    if s[1] < settings['Initial state'][1]:\n",
    "        return float('inf') \n",
    "    else:\n",
    "        return 0\n",
    "\n",
    "    \n",
    "transition = np.array([[0.4,0.5,0.1],[0.2,0.6,0.2],[0.1,0.5,0.4]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1342.4444444444446"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pumping_stoch = StochasticProgram()\n",
    "pumping_stoch.add_decisions_set({0,1,2})\n",
    "pumping_stoch.add_cost_function(err_corr_wind_power_cost)\n",
    "pumping_stoch.add_state_eq(stoch_simple_state)\n",
    "pumping_stoch.add_final_value_expression(val_T)\n",
    "pumping_stoch.add_state_limits(lower=[0,0],upper = [200,3])\n",
    "pumping_stoch.set_inital_state([100,1])\n",
    "pumping_stoch.set_step_number(24)\n",
    "pumping_stoch.add_transition_matrix(transition)\n",
    "#print pumping_stoch.settings\n",
    "pumping_stoch.solve()  \n",
    "#pumping_stoch.retrieve_decisions()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0, 100, 1)\n",
      "(1, 101, 0)\n",
      "(2, 102, 0)\n",
      "(3, 103, 0)\n",
      "(4, 104, 0)\n",
      "(5, 105, 0)\n",
      "(6, 106, 0)\n",
      "(7, 105, 0)\n",
      "(8, 106, 0)\n",
      "(9, 107, 0)\n",
      "(10, 106, 0)\n",
      "(11, 107, 0)\n",
      "(12, 106, 0)\n",
      "(13, 105, 0)\n",
      "(14, 104, 0)\n",
      "(15, 103, 0)\n",
      "(16, 102, 0)\n",
      "(17, 101, 0)\n",
      "(18, 100, 0)\n",
      "(19, 99, 0)\n",
      "(20, 98, 0)\n",
      "(21, 97, 0)\n",
      "(22, 98, 0)\n",
      "(23, 99, 0)\n",
      "(0, 100, 1)\n",
      "(1, 101, 1)\n",
      "(2, 102, 1)\n",
      "(3, 103, 1)\n",
      "(4, 104, 1)\n",
      "(5, 105, 1)\n",
      "(6, 106, 1)\n",
      "(7, 105, 1)\n",
      "(8, 106, 1)\n",
      "(9, 107, 1)\n",
      "(10, 106, 1)\n",
      "(11, 107, 1)\n",
      "(12, 106, 1)\n",
      "(13, 105, 1)\n",
      "(14, 104, 1)\n",
      "(15, 103, 1)\n",
      "(16, 102, 1)\n",
      "(17, 101, 1)\n",
      "(18, 100, 1)\n",
      "(19, 99, 1)\n",
      "(20, 98, 1)\n",
      "(21, 97, 1)\n",
      "(22, 98, 1)\n",
      "(23, 99, 1)\n",
      "(0, 100, 1)\n",
      "(1, 101, 2)\n",
      "(2, 102, 2)\n",
      "(3, 103, 2)\n",
      "(4, 104, 2)\n",
      "(5, 105, 2)\n",
      "(6, 106, 2)\n",
      "(7, 105, 2)\n",
      "(8, 106, 2)\n",
      "(9, 107, 2)\n",
      "(10, 106, 2)\n",
      "(11, 107, 2)\n",
      "(12, 106, 2)\n",
      "(13, 105, 2)\n",
      "(14, 104, 2)\n",
      "(15, 103, 2)\n",
      "(16, 102, 2)\n",
      "(17, 101, 2)\n",
      "(18, 100, 2)\n",
      "(19, 99, 2)\n",
      "(20, 98, 2)\n",
      "(21, 97, 2)\n",
      "(22, 98, 2)\n",
      "(23, 99, 2)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(array([ 1718.5,  1353. ,   939. ]),\n",
       " [[1, 1, 1, 1, 1, 1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1],\n",
       "  [1, 1, 1, 1, 1, 1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1],\n",
       "  [1, 1, 1, 1, 1, 1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1]],\n",
       " [[(100, 1),\n",
       "   (101, 0),\n",
       "   (102, 0),\n",
       "   (103, 0),\n",
       "   (104, 0),\n",
       "   (105, 0),\n",
       "   (106, 0),\n",
       "   (105, 0),\n",
       "   (106, 0),\n",
       "   (107, 0),\n",
       "   (106, 0),\n",
       "   (107, 0),\n",
       "   (106, 0),\n",
       "   (105, 0),\n",
       "   (104, 0),\n",
       "   (103, 0),\n",
       "   (102, 0),\n",
       "   (101, 0),\n",
       "   (100, 0),\n",
       "   (99, 0),\n",
       "   (98, 0),\n",
       "   (97, 0),\n",
       "   (98, 0),\n",
       "   (99, 0),\n",
       "   (100, 0)],\n",
       "  [(100, 1),\n",
       "   (101, 1),\n",
       "   (102, 1),\n",
       "   (103, 1),\n",
       "   (104, 1),\n",
       "   (105, 1),\n",
       "   (106, 1),\n",
       "   (105, 1),\n",
       "   (106, 1),\n",
       "   (107, 1),\n",
       "   (106, 1),\n",
       "   (107, 1),\n",
       "   (106, 1),\n",
       "   (105, 1),\n",
       "   (104, 1),\n",
       "   (103, 1),\n",
       "   (102, 1),\n",
       "   (101, 1),\n",
       "   (100, 1),\n",
       "   (99, 1),\n",
       "   (98, 1),\n",
       "   (97, 1),\n",
       "   (98, 1),\n",
       "   (99, 1),\n",
       "   (100, 1)],\n",
       "  [(100, 1),\n",
       "   (101, 2),\n",
       "   (102, 2),\n",
       "   (103, 2),\n",
       "   (104, 2),\n",
       "   (105, 2),\n",
       "   (106, 2),\n",
       "   (105, 2),\n",
       "   (106, 2),\n",
       "   (107, 2),\n",
       "   (106, 2),\n",
       "   (107, 2),\n",
       "   (106, 2),\n",
       "   (105, 2),\n",
       "   (104, 2),\n",
       "   (103, 2),\n",
       "   (102, 2),\n",
       "   (101, 2),\n",
       "   (100, 2),\n",
       "   (99, 2),\n",
       "   (98, 2),\n",
       "   (97, 2),\n",
       "   (98, 2),\n",
       "   (99, 2),\n",
       "   (100, 2)]])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pumping_stoch.retrieve_decisions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cost of operating a pump with a given wind turbine power input is given by:\n",
    "$$\n",
    "cost(x,t,h) := \\begin{cases} T(t) \\times (x \\times P_p - W(t)) & \\text{if} +ve \\\\ \n",
    "E_{xp} \\times (x \\times P_p - W(t)) & \\text{if} -ve\\end{cases}\n",
    "$$\n",
    "where $x$ is the descision variable, $W(t)$ is the wind turbine output in time step $t$. $P_p$ is the pump power, $E_{xp}$ is the export price."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def wind_power_cost(x,s):  \n",
    "    \"\"\"Very simple cost function for a pump with wind turbine power\"\"\"\n",
    "    Tariff = [5,5,5,5,5,8,8,8,8,8,12,12,12,12,12,50,50,50,50,20,20,6,5,5]\n",
    "    Wind = [46,  1,  3, 36, 30, 19,  9, 26, 35,  5, 49,  3,  6, 36, 43, 36, 14,\n",
    "       34,  2,  0,  0, 30, 13, 36]\n",
    "    \n",
    "    Export_price = 5.5\n",
    "        \n",
    "    power_con = x*60-Wind[s[0]]\n",
    "    if power_con >= 0:\n",
    "        return power_con*Tariff[s[0]]\n",
    "    else:\n",
    "        return power_con*Export_price\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def stoch_wind_power_cost(x,s):\n",
    "    \n",
    "    Tariff = [5,5,5,5,5,8,8,8,8,8,12,12,12,12,12,50,50,50,50,20,20,6,5,5]\n",
    "    Wind = [46,  1,  3, 36, 30, 19,  9, 26, 35,  5, 49,  3,  6, 36, 43, 36, 14,\n",
    "       34,  2,  0,  0, 30, 13, 36]\n",
    "    \n",
    "    Export_price = 5.5\n",
    "    \n",
    "    wind_out = sum(Wind[s[0]]*i for i in settings['P'][s[2]])\n",
    "    \n",
    "    power_con = x*60-wind_out\n",
    "    if power_con >= 0:\n",
    "        return power_con*Tariff[s[0]]\n",
    "    else:\n",
    "        return power_con*Export_price"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "s_{new} = \\begin{cases} (t+1,h-1,i) & \\text{if } x = 0  \\\\ (t+1,h+1,i) & \\text{if } x = 1 \\\\ (t+1,h+1.5,i) & \\text{if } x = 2\\end{cases}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "assert(hjb((settings['T'],settings['H_init']-1)) == 10000)\n",
    "assert(hjb((settings['T'],settings['H_init'])) == 0)\n",
    "assert(hjb((settings['T']-1,settings['H_min']-1)) == 10000)\n",
    "assert(hjb((settings['T']-1,settings['H_max']+1)) == 10000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "state is given by $s = (t,h,i)$\n",
    "\n",
    "$$\n",
    "v(s) = \\min_x(cost(x,s) + \\sum_j p_{ij} v(new state(x,s)))\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def make_schedule(settings):\n",
    "    sched = np.ones(settings['T'])*np.nan\n",
    "    cost_calc= 0\n",
    "    elev = np.ones(settings['T']+1)*np.nan\n",
    "\n",
    "    s = settings['Initial state']\n",
    "    t = 0\n",
    "    \n",
    "    while t < settings['T']:\n",
    "        sched[t] = settings['cache'][s][1]\n",
    "        print sched[t]\n",
    "\n",
    "        cost_calc += settings['Cost eq.'](sched[t],s)\n",
    "        elev[t] = s[1]\n",
    "\n",
    "        s = settings['State eq.'](sched[t],s)\n",
    "        t += 1\n",
    "\n",
    "    elev[settings['T']] = s[1]\n",
    "    \n",
    "    return cost_calc, sched, elev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def make_schedule2(settings):\n",
    "    \n",
    "    sched_stack =  []\n",
    "    cost_summary = []\n",
    "    string_stack = []\n",
    "    elev = np.ones(settings['T']+1)*np.nan\n",
    "    for ij in [0,1,2]:\n",
    "        cost_calc = 0\n",
    "        string_stack.insert(i,[])\n",
    "        s = settings['Initial state']\n",
    "        t = 0\n",
    "        #string_stack[ij].insert(0,[])\n",
    "        #string_stack[ij].insert(0,'{0:2} {1} {2}'.format(t,settings['cache'][s][1], s[1:]))\n",
    "        string_stack[ij].insert(0,'{0}'.format(s))\n",
    "        while t < settings['T']:\n",
    "            \n",
    "            state = tuple(sk if sk in s[:-1] else ij for sk in s  )\n",
    "            print s, s[:-1], state\n",
    "            x = settings['cache'][state][1]\n",
    "            print x\n",
    "            cost_calc += settings['Cost eq.'](x,s)\n",
    "            elev[t] = s[1]\n",
    "\n",
    "            s = settings['State eq.'](x,s)\n",
    "            t += 1\n",
    "            string_stack[ij].insert(t,'{0} {1} {2}'.format(x,s[1:],ij) )\n",
    "\n",
    "\n",
    "    elev[settings['T']] = s[1]\n",
    "    \n",
    "    return string_stack, cost_calc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0, 100, 0) (0, 100) (0, 100, 0)\n",
      "0\n",
      "(1, 99, 0) (1, 99) (1, 99, 0)\n",
      "1\n",
      "(2, 100, 0) (2, 100) (2, 100, 0)\n",
      "1\n",
      "(0, 100, 0) (0, 100) (0, 100, 0)\n",
      "0\n",
      "(1, 99, 0) (1, 99) (1, 99, 1)\n",
      "1\n",
      "(2, 100, 0) (2, 100) (2, 100, 1)\n",
      "1\n",
      "(0, 100, 0) (0, 100) (0, 100, 0)\n",
      "0\n",
      "(1, 99, 0) (1, 99) (1, 99, 2)\n",
      "1\n",
      "(2, 100, 0) (2, 100) (2, 100, 2)\n",
      "1\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "([['(0, 100, 0)', '0 (99, 0) 0', '1 (100, 0) 0', '1 (101, 0) 0'],\n",
       "  ['(0, 100, 0)', '0 (99, 0) 1', '1 (100, 0) 1', '1 (101, 0) 1'],\n",
       "  ['(0, 100, 0)', '0 (99, 0) 2', '1 (100, 0) 2', '1 (101, 0) 2']],\n",
       " 363.5)"
      ]
     },
     "execution_count": 113,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "make_schedule2(settings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{(0, 100, 80): [35, 1],\n",
       " (1, 99, 79): [36, 2],\n",
       " (1, 100, 81.5): [8, 1],\n",
       " (1, 101, 80): [16, 2],\n",
       " (2, 98, 78): [10000, 0],\n",
       " (2, 99, 80.5): [20, 1],\n",
       " (2, 100, 79): [40, 2],\n",
       " (2, 100, 83.0): [20, 1],\n",
       " (2, 101, 81.5): [0, 0],\n",
       " (2, 102, 80): [20, 1]}"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "settings['cache']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def stoch_hjb(s):\n",
    "    global settings\n",
    "    if settings['cache'].has_key(s):\n",
    "        return settings['cache'][s][0]\n",
    "    \n",
    "\n",
    "  \n",
    "    \n",
    "    if s[0] == settings['T'] and s[1] < settings['H_init']:\n",
    "        return 10000\n",
    "    \n",
    "    elif s[0] == settings['T'] and s[1] >= settings['H_init']:\n",
    "        return 0\n",
    "    \n",
    "    elif s[1] < settings['H_min'] or s[1] > settings['H_max']:\n",
    "        return 10000\n",
    "    \n",
    "    else:\n",
    "        p=[]\n",
    "        for x in settings['x_set']:\n",
    "            future = sum(stoch_hjb(settings['State eq.'](x,(s[0],s[1],i)))\n",
    "                          *settings['P'][s[2]][i] for i in [0,1,2])\n",
    "            \n",
    "            p.append(settings['Cost eq.'](x,s) + future)\n",
    "                    \n",
    "        \n",
    "        m = min(p)\n",
    "        \n",
    "        for x in settings['x_set']:\n",
    "            if m == p[x]:\n",
    "                pp = x\n",
    "                \n",
    "        settings['cache'][s] = [m, pp]\n",
    "        \n",
    "        return m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[]]\n",
      "[[0, 0, 0, 0, 0], []]\n",
      "[[0, 0, 0, 0, 0], [0, 1, 4, 9, 16], []]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[[0, 0, 0, 0, 0], [0, 1, 4, 9, 16], [0, 2, 8, 18, 32]]"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = []\n",
    "for i in [0,1,2]:\n",
    "    A.insert(i,[])\n",
    "    print A\n",
    "    for t in range(5):\n",
    "        A[i].insert(t,i*t**2)\n",
    "\n",
    "        \n",
    "A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.0\n",
      "1.0\n",
      "1.0\n",
      "1.0\n",
      "1.0\n",
      "1.0\n",
      "0.0\n",
      "1.0\n",
      "1.0\n",
      "0.0\n",
      "1.0\n",
      "0.0\n",
      "0.0\n",
      "0.0\n",
      "0.0\n",
      "0.0\n",
      "0.0\n",
      "0.0\n",
      "0.0\n",
      "0.0\n",
      "0.0\n",
      "1.0\n",
      "1.0\n",
      "1.0\n",
      "[ 1.  1.  1.  1.  1.  1.  0.  1.  1.  0.  1.  0.  0.  0.  0.  0.  0.  0.\n",
      "  0.  0.  0.  1.  1.  1.]\n",
      "[ 100.  101.  102.  103.  104.  105.  106.  105.  106.  107.  106.  107.\n",
      "  106.  105.  104.  103.  102.  101.  100.   99.   98.   97.   98.   99.\n",
      "  100.]\n",
      "1353.0\n"
     ]
    }
   ],
   "source": [
    "cost_calc, sched, elev = make_schedule(settings)\n",
    "print sched\n",
    "print elev\n",
    "print cost_calc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[28, 0]"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "settings['cache'][settings['Initial state']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.  1.  0.]\n",
      "[ 100.  101.  102.  101.]\n",
      "27.0\n"
     ]
    }
   ],
   "source": [
    "cost_calc, sched, elev = make_schedule(settings)\n",
    "print sched\n",
    "print elev\n",
    "print cost_calc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.simple_cost>"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dic = {'blub': simple_cost\n",
    "      }\n",
    "dic['blub']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "19"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dic['blub'](1,(0,1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "24"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len([5,5,5,5,5,8,8,8,8,8,12,12,12,12,12,50,50,50,50,20,20,6,5,5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([46,  1,  3, 36, 30, 19,  9, 26, 35,  5, 49,  3,  6, 36, 43, 36, 14,\n",
       "       34,  2,  0,  0, 30, 13, 36])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.random.randint(50,size=24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1000\n"
     ]
    }
   ],
   "source": [
    "x = [100,7,90,787]\n",
    "x_lim_min = [0,10,0,0]\n",
    "x_lim_max = [100,10,100,1000]\n",
    "for c,i in enumerate(x):\n",
    "    if i < x_lim_min[c] or i > x_lim_max[c]:\n",
    "        print 1000\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}