dev-ocp_trial-vector3-JIT.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.status.idle": "2020-11-23T09:04:39.524579Z",
     "shell.execute_reply": "2020-11-23T09:04:39.523822Z",
     "shell.execute_reply.started": "2020-11-23T09:04:21.919927Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded Italy Setup with 107 nodes.\n"
     ]
    }
   ],
   "source": [
    "import casadi as ca\n",
    "import casadi.tools as cat\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import matplotlib as mpl\n",
    "import networkx as nx\n",
    "import sys, copy\n",
    "from scipy.integrate import solve_ivp\n",
    "import matlab.engine\n",
    "\n",
    "sys.path.append('.')\n",
    "sys.path.append('covidOCP/')\n",
    "\n",
    "from ItalySetupProvinces import ItalySetupProvinces\n",
    "\n",
    "from ocp_utils import *\n",
    "\n",
    "eng = matlab.engine.start_matlab()\n",
    "eng.cd('geography-paper-master/', nargout=0)\n",
    "eng.run('single_build.m', nargout=0)\n",
    "\n",
    "model_size = 107\n",
    "\n",
    "s = ItalySetupProvinces(model_size)\n",
    "\n",
    "N = 163           # Number of control time interval\n",
    "T = 163           #52 * 7       # N * 7\n",
    "n_int_steps = 1\n",
    "\n",
    "model_days = pd.date_range(s.start_date, s.end_date, freq='1D')\n",
    "freq = '1D'  # 'W-MON'\n",
    "model_step = pd.date_range(s.start_date, s.end_date, freq=freq)\n",
    "mobintime = s.mobility_ts.resample(freq).mean()\n",
    "\n",
    "N = len(model_step) - 1\n",
    "T = len(model_days)\n",
    "\n",
    "N = 30\n",
    "T = 31\n",
    "\n",
    "nx = 9\n",
    "\n",
    "states_list = ['S', 'E', 'P', 'I', 'A', 'Q', 'H', 'R', 'V']\n",
    "states = states_list\n",
    "\n",
    "S, E, P, I, A, Q, H, R, V = np.arange(nx)\n",
    "\n",
    "integ_matlab = np.array(eng.eval('x'))\n",
    "\n",
    "p_dict, mobfrac, mobmat, betaratiointime, x0 = get_parameters_from_matlab(eng, s, model_size, model_days, freq)\n",
    "\n",
    "dt = T / N / n_int_steps\n",
    "\n",
    "obj_params = {\n",
    "    'scale_ell': 1e4,\n",
    "    'scale_If': 0,\n",
    "    'scale_v': 1e-8\n",
    "}\n",
    "obj_params = {\n",
    "    'scale_ell': 1e0,\n",
    "    'scale_If': 1e2,\n",
    "    'scale_v': 1e-6\n",
    "}\n",
    "betaP0 = p_dict['betaP0']\n",
    "epsilonA = p_dict['epsilonA']\n",
    "epsilonI = p_dict['epsilonI']\n",
    "r = p_dict['r']\n",
    "\n",
    "p_dict.pop('betaP0')\n",
    "p_dict.pop('r')\n",
    "p_dict.pop('epsilonA')\n",
    "p_dict.pop('epsilonI')\n",
    "C = mobmat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### implement the vaccintion as we did for mobility, computer outside the ODE. make like easier.\n",
    "\n",
    "```\n",
    "        x_ = ca.veccat(*states[...])\n",
    "        u_ = ca.veccat(*controls[...])\n",
    "        R_ = states['R']\n",
    "        SpR = pop_nodeSX - states['I'] - states['V']\n",
    "        x_[1] -= controls['v']*R_/(SpR+1e-10)\n",
    "        x_[3] += controls['v']*SpR/(SpR+1e-10)\n",
    "        ell = 0.\n",
    "        \n",
    "        for k in range(n_int_steps):\n",
    "          x, ell = rk4_step(x_,u_,covar, params, pop_nodeSX)\n",
    "          ell += ell_\n",
    "            \n",
    "```\n",
    "JIT compilation\n",
    "```\n",
    "dae = {'x':states, 'p':ca.veccat(controls,covar,params,pop_nodeSX), 'ode':rhs, 'quad':scale_ell*rhs_ell + scale_v*v*v}\n",
    "int_type = 'rk'\n",
    "# int_type = 'collocation'\n",
    "\n",
    "opts = {'tf':T/N, 'number_of_finite_elements':n_int_steps}\n",
    "opts[\"jit\"] = True\n",
    "opts[\"compiler\"] = \"shell\"\n",
    "opts[\"simplify\"] = True\n",
    "opts[\"expand\"] = True\n",
    "\n",
    "F = ca.integrator('F', int_type, dae, opts)\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-23T09:04:39.526395Z",
     "iopub.status.busy": "2020-11-23T09:04:39.526095Z",
     "iopub.status.idle": "2020-11-23T09:04:39.542743Z",
     "shell.execute_reply": "2020-11-23T09:04:39.542124Z",
     "shell.execute_reply.started": "2020-11-23T09:04:39.526366Z"
    }
   },
   "outputs": [],
   "source": [
    "def rhs_py(t, x, u, cov, p, mob, pop_node):\n",
    "    S, E, P, I, A, Q, H, R, V = x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7], x[8]\n",
    "\n",
    "    deltaE, deltaP,sigma,eta,gammaI,gammaA,gammaQ,gammaH,alphaI,alphaH, zeta,  gammaV = \\\n",
    "    p[0], p[1], p[2], p[3], p[4], p[5], p[6], p[7], p[8], p[9], p[10], p[11]\n",
    "   \n",
    "    v = u[0]\n",
    "    foi = mob\n",
    "    rhs = [None] * nx\n",
    "    vaccrate = v / (S + E + P + A + Q + H + R+1)  # NOT I NOT H ?\n",
    "    vaccrate = v / ca.sqrt((S+E+P+A+Q+H+R)**2+10)   # 10 (1e7 ?1e3 ?)is to tune depending on scale, approx abs. value (brings quantities non negative by vaccinating)\n",
    "    rhs[0] = -(foi + vaccrate) * S + gammaV * V                 # S\n",
    "    rhs[1] = foi * S - deltaE * E;                              # E \n",
    "    rhs[2] = deltaE * E - deltaP * P;                           # P\n",
    "    rhs[3] = sigma * deltaP * P - (eta + gammaI + alphaI) * I;  # I\n",
    "    rhs[4] = (1 - sigma) * deltaP * P - gammaA * A;             # A\n",
    "    rhs[5] = zeta * eta * I - gammaQ * Q;                       # Q\n",
    "    rhs[6] = (1 - zeta) * eta * I - (gammaH + alphaH) * H;      # H\n",
    "    rhs[7] = gammaI * I + gammaA * A + gammaH * H + gammaQ * Q;  # R\n",
    "    rhs[8] = vaccrate * S - gammaV * V                           # V\n",
    "    rhs_ell = [None] * 3\n",
    "    rhs_ell[0] = gammaH * H;  # recovered from the hospital\n",
    "    rhs_ell[1] = alphaH * H;  # total death\n",
    "    rhs_ell[2] = (1 - zeta) * eta * I;  # cumulative hospitalized cases\n",
    "\n",
    "    return rhs, rhs_ell\n",
    "\n",
    "def rhs_py2(t, x, p):\n",
    "    \n",
    "    S, E, P, I, A, Q, H, R, V = x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7], x[8]\n",
    "    \n",
    "    #u, cov, p, mob, pop_node =  p[0], p[1], p[2], p[3], p[4]\n",
    "    #print(u, cov, p, mob, pop_node)\n",
    "\n",
    "    v, mob, J, deltaE, deltaP, sigma, eta, gammaI, gammaA, gammaQ, gammaH, alphaI, alphaH, zeta, gammaV, scale_ell, scale_If, scale_v, pop_node = \\\n",
    "    p[0], p[1], p[2], p[3], p[4], p[5], p[6], p[7], p[8], p[9], p[10], p[11], p[12], p[13], p[14], p[15], p[16], p[17], p[18]\n",
    "   \n",
    "    #v = u[0]\n",
    "    foi = mob\n",
    "    rhs = [None] * nx\n",
    "    #vaccrate = v / (S + E + P + A + Q + H + R+1)  # NOT I NOT H ?\n",
    "    #vaccrate = v / ca.sqrt((S+E+P+A+Q+H+R)**2+10)   # 10 (1e7 ?1e3 ?)is to tune depending on scale, approx abs. value (brings quantities non negative by vaccinating)\n",
    "    vaccrate = v /(S+E+P+A+R+1e-10)\n",
    "    rhs[0] = -(foi + vaccrate) * S + gammaV * V                 # S\n",
    "    rhs[1] = foi * S - deltaE * E;                              # E \n",
    "    rhs[2] = deltaE * E - deltaP * P;                           # P\n",
    "    rhs[3] = sigma * deltaP * P - (eta + gammaI + alphaI) * I;  # I\n",
    "    rhs[4] = (1 - sigma) * deltaP * P - gammaA * A;             # A\n",
    "    rhs[5] = zeta * eta * I - gammaQ * Q;                       # Q\n",
    "    rhs[6] = (1 - zeta) * eta * I - (gammaH + alphaH) * H;      # H\n",
    "    rhs[7] = gammaI * I + gammaA * A + gammaH * H + gammaQ * Q;  # R\n",
    "    rhs[8] = vaccrate * S - gammaV * V                           # V\n",
    "    #rhs_ell = [None] * 3\n",
    "    #rhs_ell[0] = gammaH * H;  # recovered from the hospital\n",
    "    rhs_ell = alphaH * H;  # total death\n",
    "    #rhs_ell[2] = (1 - zeta) * eta * I;  # cumulative hospitalized cases\n",
    "\n",
    "    return rhs, rhs_ell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-23T09:04:39.544880Z",
     "iopub.status.busy": "2020-11-23T09:04:39.544666Z",
     "iopub.status.idle": "2020-11-23T09:04:39.561017Z",
     "shell.execute_reply": "2020-11-23T09:04:39.560230Z",
     "shell.execute_reply.started": "2020-11-23T09:04:39.544851Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building OCP with 107 nodes\n"
     ]
    }
   ],
   "source": [
    "setup = s\n",
    "model_params = copy.copy(p_dict)\n",
    "ind2name = setup.ind2name\n",
    "#ind_to_plot = setup.ind_to_plot\n",
    "pop_node = setup.pop_node\n",
    "\n",
    "M = setup.nnodes\n",
    "mobility = setup.mobility\n",
    "\n",
    "print(f'Building OCP with {M} nodes')\n",
    "\n",
    "c = setup.mobility\n",
    "\n",
    "scale_ell = obj_params['scale_ell']\n",
    "scale_If = obj_params['scale_If']\n",
    "scale_v = obj_params['scale_v']\n",
    "pnum = params_to_vector(model_params)\n",
    "\n",
    "pnum.append(scale_ell)\n",
    "pnum.append(scale_If)\n",
    "pnum.append(scale_v)\n",
    "\n",
    "# ---- decision variables ---------\n",
    "states = ['S', 'E', 'P', 'I', 'A', 'Q', 'H', 'R', 'V']\n",
    "states = cat.struct_symSX(states)\n",
    "[S, E, P, I, A, Q, H, R, V] = states[...]\n",
    "\n",
    "controls = cat.struct_symSX(['v', 'mob'])\n",
    "[v, mob] = controls[...]\n",
    "\n",
    "covar = cat.struct_symSX(['J'])\n",
    "[J] = covar[...]\n",
    "\n",
    "params = cat.struct_symSX(list(model_params.keys()) + ['scale_ell', 'scale_If', 'scale_v'])\n",
    "[deltaE, deltaP, sigma, eta, gammaI, gammaA, gammaQ, gammaH, alphaI, alphaH, zeta, gammaV,\n",
    " scale_ell, scale_If, scale_v] = params[...]\n",
    "\n",
    "pop_nodeSX = ca.SX.sym('pop_node')\n",
    "\n",
    "\n",
    "\n",
    "rhs, rhs_ell = rhs_py2(0, states.cat, ca.veccat(controls.cat, covar.cat, params.cat, pop_nodeSX))\n",
    "rhs = ca.veccat(*rhs)\n",
    "#rhs_ell = ca.veccat(*rhs_ell)#[1]  # mod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-23T09:04:39.563062Z",
     "iopub.status.busy": "2020-11-23T09:04:39.562800Z",
     "iopub.status.idle": "2020-11-23T09:04:39.569685Z",
     "shell.execute_reply": "2020-11-23T09:04:39.568993Z",
     "shell.execute_reply.started": "2020-11-23T09:04:39.563036Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SX((alphaH*H))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rhs_ell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-23T09:04:39.571187Z",
     "iopub.status.busy": "2020-11-23T09:04:39.570969Z",
     "iopub.status.idle": "2020-11-23T09:04:40.664047Z",
     "shell.execute_reply": "2020-11-23T09:04:40.663478Z",
     "shell.execute_reply.started": "2020-11-23T09:04:39.571161Z"
    }
   },
   "outputs": [],
   "source": [
    "dae = {'x':states,\n",
    "       'p':ca.veccat(controls,covar,params, pop_nodeSX), \n",
    "       'ode':rhs,\n",
    "       'quad':scale_ell*rhs_ell + scale_v*v*v\n",
    "      }\n",
    "\n",
    "int_type = 'rk'\n",
    "opts = {'tf':T/N, 'number_of_finite_elements':n_int_steps}\n",
    "opts[\"jit\"] = True\n",
    "opts[\"compiler\"] = \"shell\"\n",
    "opts[\"simplify\"] = True\n",
    "opts[\"expand\"] = True\n",
    "\n",
    "F = ca.integrator('F', int_type, dae, opts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-23T09:04:40.665369Z",
     "iopub.status.busy": "2020-11-23T09:04:40.665175Z",
     "iopub.status.idle": "2020-11-23T09:04:43.607911Z",
     "shell.execute_reply": "2020-11-23T09:04:43.607245Z",
     "shell.execute_reply.started": "2020-11-23T09:04:40.665344Z"
    }
   },
   "outputs": [],
   "source": [
    "Vars = cat.struct_symMX([\n",
    "    (\n",
    "        cat.entry(\"x\", struct=states, repeat=[M, N + 1]),\n",
    "        cat.entry(\"u\", struct=controls, repeat=[M, N]),\n",
    "    ),\n",
    "])\n",
    "\n",
    "Params = cat.struct_symMX([\n",
    "    (\n",
    "        cat.entry(\"cov\", struct=covar, repeat=[M, N]), \n",
    "        cat.entry(\"p\", struct=params),\n",
    "    ),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-23T09:04:43.609683Z",
     "iopub.status.busy": "2020-11-23T09:04:43.609471Z",
     "iopub.status.idle": "2020-11-23T09:04:45.111221Z",
     "shell.execute_reply": "2020-11-23T09:04:45.110394Z",
     "shell.execute_reply.started": "2020-11-23T09:04:43.609648Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pruning 10921 non-diagonal mobility elements of 11342.\n",
      "0:1:2:3:4:5:6:7:8:9:10:11:12:13:14:15:16:17:18:19:20:21:22:23:24:25:26:27:28:29:"
     ]
    }
   ],
   "source": [
    "lbx = Vars(-np.inf)\n",
    "ubx = Vars(np.inf)\n",
    "\n",
    "f = 0\n",
    "vaccines = 0\n",
    "cases = 0\n",
    "reg = 0\n",
    "cdot_T = 0\n",
    "dyn = [None] * N\n",
    "spatial = [None] * N\n",
    "\n",
    "mob_prun = 0.0006\n",
    "#mob_prun = .005\n",
    "#mob_prun = .01 # not tested\n",
    "\n",
    "for k in range(N):\n",
    "    mobK = mobintime.to_numpy().T[:,k]\n",
    "    betaR = betaratiointime.to_numpy().T[:,k]\n",
    "    C = r*mobfrac.flatten()*mobK*mobmat\n",
    "    np.fill_diagonal(C,1-C.sum(axis=1)+ C.diagonal())\n",
    "    \n",
    "    dyn[k] = []\n",
    "    spatial[k] = []\n",
    "    Sk, Ek, Pk, Rk, Ak, Ik = ca.veccat(*Vars['x', :, k, 'S']), ca.veccat(*Vars['x', :, k, 'E']), \\\n",
    "                             ca.veccat(*Vars['x', :, k, 'P']), ca.veccat(*Vars['x', :, k, 'R']), \\\n",
    "                             ca.veccat(*Vars['x', :, k, 'A']), ca.veccat(*Vars['x', :, k, 'I'])\n",
    "\n",
    "    if k == 0:\n",
    "        print(f'pruning {C[C<mob_prun].size} non-diagonal mobility elements of {C.size-M}.')  \n",
    "    C[C<mob_prun] = 0 # Prune elements\n",
    "        \n",
    "    mob_ik=C@((C.T@(betaP0*betaR*(Pk+epsilonA*Ak))+epsilonI*betaP0*betaR*Ik)/ \\\n",
    "             (C.T@(Sk+Ek+Pk+Rk+Ak)+Ik))\n",
    "    print(f'{k}:', end='')\n",
    "    for i in range(M):\n",
    "        #[X_, ell_ik] = rk4_int(Vars['x', i, k],\n",
    "        #                       ca.veccat(Vars['u', i, k], Params['cov', i, k], Params['p'], pop_node[i]))\n",
    "        Fk = F(x0=Vars['x',i,k],p=ca.veccat(Vars['u',i,k],Params['cov',i,k],Params['p'], pop_node[i]))\n",
    "        X_ = Fk['xf']\n",
    "        ell_ik = Fk['qf']\n",
    "        \n",
    "        \n",
    "        dyn[k].append(Vars['x', i, k + 1] - X_)\n",
    "        #ell_ik_, cases_ik, reg_ik = ell(Vars['x', i, k], Vars['u', i, k], Params['cov', i, k], Params['p'], pop_node[i])\n",
    "        f += ell_ik #MOD: before  ell_ik_\n",
    "        #cases += cases_ik\n",
    "        #reg += reg_ik\n",
    "        \n",
    "        spatial[k].append(Vars['u', i, k, 'mob'] - mob_ik[i]) # spatial, vaccines and dyn are put in g(x), with constraints that spatial and dyn are equal to zero\n",
    "        # thus imposing the dynamics and coupling.\n",
    "        vaccines += Vars[ 'u', i, k, 'v'] * T / N  # Number of vaccine spent = num of vaccine rate * 7 (number of days)\n",
    "\n",
    "f /= T  # Average over interval for cost ^ but not terminal cost "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-23T09:04:45.112789Z",
     "iopub.status.busy": "2020-11-23T09:04:45.112540Z",
     "iopub.status.idle": "2020-11-23T09:10:49.844454Z",
     "shell.execute_reply": "2020-11-23T09:10:49.842622Z",
     "shell.execute_reply.started": "2020-11-23T09:04:45.112764Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing constraints, ...DONE\n",
      "Building NLP function...DONE\n",
      "Building Jacobian function..."
     ]
    },
    {
     "ename": "SystemError",
     "evalue": "<built-in function Function_factory> returned a result with an error set",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m~/.local/share/virtualenvs/COVID-19_Italy-vaccination-ocp-vr3EY8D_/lib/python3.7/site-packages/casadi/casadi.py\u001b[0m in \u001b[0;36m<lambda>\u001b[0;34m(self, name, value)\u001b[0m\n\u001b[1;32m  11550\u001b[0m         \u001b[0m__swig_setmethods__\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mupdate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgetattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_s\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'__swig_setmethods__'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m> 11551\u001b[0;31m     \u001b[0m__setattr__\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mlambda\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0m_swig_setattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mFunction\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m  11552\u001b[0m     \u001b[0m__swig_getmethods__\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: ",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001b[0;31mSystemError\u001b[0m                               Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-8-0900e0007eb2>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     41\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'DONE'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     42\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Building Jacobian function...'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mend\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m''\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 43\u001b[0;31m \u001b[0mnlpJac\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnlpFun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfactory\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'nlpJac'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m'i0'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'i1'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m'jac:o1:i0'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     44\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     45\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'DONE'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/share/virtualenvs/COVID-19_Italy-vaccination-ocp-vr3EY8D_/lib/python3.7/site-packages/casadi/casadi.py\u001b[0m in \u001b[0;36mfactory\u001b[0;34m(self, *args)\u001b[0m\n\u001b[1;32m  12049\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m  12050\u001b[0m         \"\"\"\n\u001b[0;32m> 12051\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0m_casadi\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mFunction_factory\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m  12052\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m  12053\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mSystemError\u001b[0m: <built-in function Function_factory> returned a result with an error set"
     ]
    }
   ],
   "source": [
    "print('Writing constraints, ...', end='')\n",
    "g = cat.struct_MX([\n",
    "    cat.entry(\"dyn\", expr=dyn),\n",
    "    cat.entry(\"spatial\", expr=spatial),\n",
    "    cat.entry(\"vaccines\", expr=vaccines),\n",
    "])\n",
    "\n",
    "costTerms = ca.Function('costTerms', [Vars, Params], [cases, reg])\n",
    "\n",
    "\n",
    "# This initialize\n",
    "lbg = g(0)\n",
    "ubg = g(0)\n",
    "\n",
    "\n",
    "ubg['vaccines'] = 2000. * M\n",
    "lbg['vaccines'] = -np.inf\n",
    "\n",
    "optimize = 1\n",
    "lbx['u', :, :, 'v'] = 0.\n",
    "ubx['u', :, :, 'v'] = 1000 * optimize  # = 0 if we don't want to optimize\n",
    "# ubx['u',:,:,'v'] = 0\n",
    "#ubx['u', :, :1, 'v'] = 0.\n",
    "\n",
    "\n",
    "# Set initial conditions as constraints\n",
    "for cp, name in enumerate(states.keys()):\n",
    "    for i in range(M):\n",
    "        lbx['x', i, 0, name] = ubx['x', i, 0, name] = x0[i*nx+cp]\n",
    "\n",
    "print('DONE')\n",
    "print('Building NLP function...', end='')\n",
    "\n",
    "# NLOP arguments:\n",
    "# 'x' : variable names to optimize: here states and control\n",
    "# 'p' : params that can be change after\n",
    "# 'f' is the objective function of 'x', that we ought to minize\n",
    "# 'g' is a function that goes lbg < g(x,p) < ubg. If you want equality constraint then ubg = lbg = the number yoiu want it to be.\n",
    "nlp = {'x': Vars, 'p': Params, 'f': f, 'g': g}\n",
    "nlpFun = ca.Function('nlpFun', [Vars, Params], [f, g])\n",
    "print('DONE')\n",
    "print('Building Jacobian function...', end='')\n",
    "nlpJac = nlpFun.factory('nlpJac', ['i0', 'i1'], ['jac:o1:i0'])\n",
    "\n",
    "print('DONE')\n",
    "print('Building Solver...', end='')\n",
    "\n",
    "options = {}\n",
    "options['ipopt'] = {}\n",
    "options['ipopt'][\"linear_solver\"] = \"ma57\"\n",
    "options['ipopt'][\"linear_solver\"] = \"ma86\"\n",
    "# options['ipopt'][\"linear_solver\"] = \"ma86\"\n",
    "# options['ipopt'][\"linear_solver\"] = \"ma97\"\n",
    "# options['ipopt']['bound_frac'] = 1e-4\n",
    "# options['ipopt']['bound_push'] = 1e-4\n",
    "# options['ipopt']['slack_bound_frac'] = 1e-4\n",
    "# options['ipopt']['slack_bound_push'] = 1e-6\n",
    "# options['ipopt']['hessian_constant'] = 'yes'\n",
    "# options['ipopt'][\"tol\"] = 1e-8\n",
    "\n",
    "solver = ca.nlpsol('solver', \"ipopt\", nlp, options)\n",
    "\n",
    "# options = {}\n",
    "# options['ls_step_factor'] = 0.5\n",
    "# options['pr_tol'] = 1e-8\n",
    "# options['du_tol'] = 1e-8\n",
    "# options['qp_solver'] = 'nlpsol'\n",
    "# options['qp_solver_options']  = {'nlpsol':'ipopt','nlpsol_options':{'print_time':False,'ipopt.print_level':0,'ipopt.hessian_constant':'yes','ipopt.linear_solver':'ma97','ipopt.tol':1e-16}}\n",
    "# # options['regularization'] = 'regularize_reduced'\n",
    "# options['regularization'] = 'regularize_reduced_diag'\n",
    "# options['regularization_tol'] = 1e-10\n",
    "# options['linesearch'] = 'filter'\n",
    "# options[\"iteration_callback\"] = plotIterates\n",
    "# self.solver = SQPmethod(nlp,options)\n",
    "\n",
    "print('DONE')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2020-11-23T09:10:49.845182Z",
     "iopub.status.idle": "2020-11-23T09:10:49.845419Z",
     "shell.execute_reply": "2020-11-23T09:10:49.845292Z"
    }
   },
   "outputs": [],
   "source": [
    "integ_matlab = np.array(eng.eval('x'))\n",
    "\n",
    "init = Vars(0)\n",
    "for i, name in enumerate(states.keys()):\n",
    "    for k in range(N + 1):\n",
    "        # init['x',:,k,name] = sol0.y[i,k]\n",
    "        for nd in range(M):\n",
    "            init['x', nd, k, name] = integ_matlab.T[nd+107*i, k].T # p['x0']  # self.ic[name][i]\n",
    "\n",
    "init['u'] = 0.\n",
    "\n",
    "arg = {}\n",
    "arg['lbg'] = lbg\n",
    "arg['ubg'] = ubg\n",
    "\n",
    "arg['lbx'] = lbx\n",
    "arg['ubx'] = ubx\n",
    "arg['x0'] = init\n",
    "arg['p'] = Params()\n",
    "# for i in range(M):\n",
    "#    for k in range(N):\n",
    "#        arg['p']['cov',i,k] = rainfall_norm[i,k]**params(pnum)['r']\n",
    "# arg['p']['cov',i,k] = np.random.random()**params(pnum)['r']\n",
    "\n",
    "arg['p']['p'] = pnum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2020-11-23T09:10:49.846537Z",
     "iopub.status.idle": "2020-11-23T09:10:49.846778Z",
     "shell.execute_reply": "2020-11-23T09:10:49.846649Z"
    }
   },
   "outputs": [],
   "source": [
    "#pnum.append(scale_ell)\n",
    "#pnum.append(scale_If)\n",
    "#pnum.append(scale_v)\n",
    "pnum[-3] = 1e8\n",
    "pnum[-2] = 0\n",
    "pnum[-1] = 1e-30\n",
    "arg['p']['p'] = pnum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2020-11-23T09:10:49.847639Z",
     "iopub.status.idle": "2020-11-23T09:10:49.847881Z",
     "shell.execute_reply": "2020-11-23T09:10:49.847750Z"
    }
   },
   "outputs": [],
   "source": [
    "arg['p']['p']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2020-11-23T09:10:49.848713Z",
     "iopub.status.idle": "2020-11-23T09:10:49.849052Z",
     "shell.execute_reply": "2020-11-23T09:10:49.848896Z"
    }
   },
   "outputs": [],
   "source": [
    "sol = solver(**arg)\n",
    "opt = Vars(sol['x'])\n",
    "lam_g = g(sol['lam_g'])\n",
    "lam_x = Vars(sol['lam_x'])\n",
    "[fnum,gnum]=nlpFun(opt,arg['p'])\n",
    "Jgnum=nlpJac(opt, arg['p'])\n",
    "gnum = g(gnum)\n",
    "print(f\"\"\"\n",
    "Vaccines stockpile: \n",
    "    {float(arg['ubg']['vaccines']):010f} total.\n",
    "    {float(g(gnum)['vaccines']):010f} spent.\n",
    "    {float((arg['ubg']['vaccines']  - g(gnum)['vaccines'])):010f} left.\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2020-11-23T09:10:49.850075Z",
     "iopub.status.idle": "2020-11-23T09:10:49.850435Z",
     "shell.execute_reply": "2020-11-23T09:10:49.850273Z"
    }
   },
   "outputs": [],
   "source": [
    "pop_node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2020-11-23T09:10:49.851394Z",
     "iopub.status.idle": "2020-11-23T09:10:49.851687Z",
     "shell.execute_reply": "2020-11-23T09:10:49.851536Z"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "fig, axes = plt.subplots(2,5, figsize = (20,10))\n",
    "node = -1\n",
    "til =  30#-1\n",
    "\n",
    "for i, st in enumerate(states_list):\n",
    "    axes.flat[i].plot(np.array(ca.veccat(*opt['x',node,:til,st])), \n",
    "                  linestyle=':', lw = 4, color='r')\n",
    "    if st != 'V':\n",
    "        axes.flat[i].plot(np.array(integ_matlab.T[node+107*i,:til].T), \n",
    "                  linestyle='-', lw = 2, color='k')\n",
    "\n",
    "    axes.flat[i].set_title(st);\n",
    "\n",
    "axes.flat[-1].step(np.array(ca.veccat(ca.veccat(*opt['u',node,:til,'v']))),#,opt['u',node,-1,'v'])),\n",
    "                  'k',label=r\"$\\nu(t)$\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2020-11-23T09:10:49.852660Z",
     "iopub.status.idle": "2020-11-23T09:10:49.852926Z",
     "shell.execute_reply": "2020-11-23T09:10:49.852783Z"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "fig, axes = plt.subplots(5,2, figsize = (10,10))\n",
    "node = 1\n",
    "til = 30 # -1\n",
    "\n",
    "for i, st in enumerate(states_list):\n",
    "    for k in range(M):\n",
    "        axes.flat[i].plot(np.array(ca.veccat(*opt['x',k,:til,st])), \n",
    "                   lw = 2, ls = '--')\n",
    "        if st != 'V':\n",
    "            axes.flat[i].plot(np.array(integ_matlab.T[k+107*i,:til].T), \n",
    "                   lw = .5)\n",
    "        axes.flat[i].set_title(st);\n",
    "        axes.flat[-1].step(np.arange(til),\n",
    "            np.array(ca.veccat(ca.veccat(*opt['u',k,:til,'v']))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2020-11-23T09:10:49.853901Z",
     "iopub.status.idle": "2020-11-23T09:10:49.854169Z",
     "shell.execute_reply": "2020-11-23T09:10:49.854023Z"
    }
   },
   "outputs": [],
   "source": [
    "len(np.array(ca.veccat(ca.veccat(*opt['u',k,:til,'v']))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2020-11-23T09:10:49.854998Z",
     "iopub.status.idle": "2020-11-23T09:10:49.855341Z",
     "shell.execute_reply": "2020-11-23T09:10:49.855197Z"
    }
   },
   "outputs": [],
   "source": [
    "import networkx\n",
    "grph = networkx.from_numpy_matrix(C)\n",
    "networkx.draw(grph, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2020-11-23T09:10:49.856212Z",
     "iopub.status.idle": "2020-11-23T09:10:49.856471Z",
     "shell.execute_reply": "2020-11-23T09:10:49.856332Z"
    }
   },
   "outputs": [],
   "source": [
    "grph = networkx.from_numpy_matrix(C)\n",
    "networkx.draw(grph, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2020-11-23T09:10:49.857437Z",
     "iopub.status.idle": "2020-11-23T09:10:49.857707Z",
     "shell.execute_reply": "2020-11-23T09:10:49.857561Z"
    }
   },
   "outputs": [],
   "source": [
    "networkx.draw(grph, )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (ocp-covid)",
   "language": "python",
   "name": "ocp-covid"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}