dev-ocp_trial-vector2-crtl-out-struct0.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T15:34:15.891586Z",
     "iopub.status.busy": "2020-11-27T15:34:15.891357Z",
     "iopub.status.idle": "2020-11-27T15:34:26.730679Z",
     "shell.execute_reply": "2020-11-27T15:34:26.729951Z",
     "shell.execute_reply.started": "2020-11-27T15:34:15.891525Z"
    }
   },
   "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 COVIDParametersOCP 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": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T15:34:26.732264Z",
     "iopub.status.busy": "2020-11-27T15:34:26.732073Z",
     "iopub.status.idle": "2020-11-27T15:34:26.751715Z",
     "shell.execute_reply": "2020-11-27T15:34:26.750976Z",
     "shell.execute_reply.started": "2020-11-27T15:34:26.732238Z"
    }
   },
   "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",
    "    vaccrate = 0\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T15:34:26.753680Z",
     "iopub.status.busy": "2020-11-27T15:34:26.753515Z",
     "iopub.status.idle": "2020-11-27T15:34:26.780598Z",
     "shell.execute_reply": "2020-11-27T15:34:26.780068Z",
     "shell.execute_reply.started": "2020-11-27T15:34:26.753660Z"
    }
   },
   "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(['mobility_t','betaratio_t'])\n",
    "[mobility_t, betaratio_t] = 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",
    "# The rhs is at time zero, the time is also no used in the equation so that exoplain\n",
    "rhs, rhs_ell = rhs_py(0, states.cat, controls.cat, covar.cat, params.cat, mob, pop_nodeSX)\n",
    "rhs = ca.veccat(*rhs)\n",
    "rhs_ell = ca.veccat(*rhs_ell)  # mod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T15:34:26.782525Z",
     "iopub.status.busy": "2020-11-27T15:34:26.782345Z",
     "iopub.status.idle": "2020-11-27T15:34:26.785397Z",
     "shell.execute_reply": "2020-11-27T15:34:26.784872Z",
     "shell.execute_reply.started": "2020-11-27T15:34:26.782503Z"
    }
   },
   "outputs": [],
   "source": [
    "frhs = ca.Function('frhs', [states, controls, covar, params, pop_nodeSX],\n",
    "                           [rhs, rhs_ell[1]])#scale_ell * rhs_ell[1] + scale_v * v * v])# mod ICI juste ell[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T15:34:26.786606Z",
     "iopub.status.busy": "2020-11-27T15:34:26.786409Z",
     "iopub.status.idle": "2020-11-27T15:34:26.791966Z",
     "shell.execute_reply": "2020-11-27T15:34:26.791444Z",
     "shell.execute_reply.started": "2020-11-27T15:34:26.786580Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Function(frhs:(i0[9],i1[2],i2[2],i3[15],i4)->(o0[9],o1) SXFunction)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "frhs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T15:34:26.792914Z",
     "iopub.status.busy": "2020-11-27T15:34:26.792770Z",
     "iopub.status.idle": "2020-11-27T15:34:26.799683Z",
     "shell.execute_reply": "2020-11-27T15:34:26.799185Z",
     "shell.execute_reply.started": "2020-11-27T15:34:26.792894Z"
    }
   },
   "outputs": [],
   "source": [
    "dt = T / N / n_int_steps  # length of an integration interval\n",
    "# ---- dynamic constraints --------\n",
    "k1, k1ell = frhs(states, controls, covar, params, pop_nodeSX)\n",
    "k2, k2ell = frhs(states + dt / 2 * k1, controls, covar, params, pop_nodeSX)\n",
    "k3, k3ell = frhs(states + dt / 2 * k2, controls, covar, params, pop_nodeSX)\n",
    "k4, k4ell = frhs(states + dt * k3, controls, covar, params, pop_nodeSX)\n",
    "x_next = states + dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4)\n",
    "ell_next = dt / 6 * (\n",
    "        k1ell + 2 * k2ell + 2 * k3ell + k4ell)  # No need for because we sum it week by week few lines below.\n",
    "\n",
    "rk4_step = ca.Function('rk4_step', [states, controls, covar, params, pop_nodeSX], [x_next, ell_next])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T15:34:26.800860Z",
     "iopub.status.busy": "2020-11-27T15:34:26.800632Z",
     "iopub.status.idle": "2020-11-27T15:34:26.804986Z",
     "shell.execute_reply": "2020-11-27T15:34:26.804415Z",
     "shell.execute_reply.started": "2020-11-27T15:34:26.800839Z"
    }
   },
   "outputs": [],
   "source": [
    "# Overwrite last cell with eurler.\n",
    "dt = T / N / n_int_steps  # length of an integration interval\n",
    "a, b = frhs(states, controls, covar, params, pop_nodeSX)\n",
    "x_next = states + dt*a\n",
    "ell_next = dt*b\n",
    "rk4_step = ca.Function('rk4_step', [states, controls, covar, params, pop_nodeSX], [x_next, ell_next])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T15:34:26.805991Z",
     "iopub.status.busy": "2020-11-27T15:34:26.805843Z",
     "iopub.status.idle": "2020-11-27T15:34:26.812952Z",
     "shell.execute_reply": "2020-11-27T15:34:26.812365Z",
     "shell.execute_reply.started": "2020-11-27T15:34:26.805970Z"
    }
   },
   "outputs": [],
   "source": [
    "x_ = ca.veccat(*states[...])\n",
    "u_ = ca.veccat(*controls[...])\n",
    "ell = 0.\n",
    "VacPpl = states['S'] + states['E'] + states['P'] + states['A'] + states['R']\n",
    "vaccrate = controls['v']/(VacPpl+1e-10)\n",
    "x_[0] -= vaccrate * states['S']\n",
    "x_[8] += vaccrate * states['S']\n",
    "for k in range(n_int_steps):\n",
    "    x_, ell_ = rk4_step(x_, u_,covar, params, pop_nodeSX)\n",
    "    ell += ell_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T15:34:26.815948Z",
     "iopub.status.busy": "2020-11-27T15:34:26.815772Z",
     "iopub.status.idle": "2020-11-27T15:34:26.821117Z",
     "shell.execute_reply": "2020-11-27T15:34:26.820625Z",
     "shell.execute_reply.started": "2020-11-27T15:34:26.815927Z"
    }
   },
   "outputs": [],
   "source": [
    "rk4_int = ca.Function('rk4_int', [states, ca.veccat(controls, covar, params, pop_nodeSX)], [x_, ell],\n",
    "                      ['x0', 'p'], ['xf', 'qf'])\n",
    "\n",
    "# BUG TODO Isn't this a double multiplication by the scale parameter since ell is already multiplied ?\n",
    "ell = ca.Function('ell', [states, controls, covar, params, pop_nodeSX],\n",
    "                  [scale_ell * ell + scale_v * v * v, scale_ell * ell,\n",
    "                   scale_v * v * v])  # Very dependent on regularization factor\n",
    "#ell = ca.Function('ell', [states, controls, covar, params, pop_nodeSX],\n",
    "#                  [ell, scale_ell * ell,\n",
    "#                   scale_v * v * v])  # Very dependent on regularization factor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T15:34:26.823305Z",
     "iopub.status.busy": "2020-11-27T15:34:26.823139Z",
     "iopub.status.idle": "2020-11-27T15:34:29.260678Z",
     "shell.execute_reply": "2020-11-27T15:34:29.260099Z",
     "shell.execute_reply.started": "2020-11-27T15:34:26.823285Z"
    }
   },
   "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": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T15:34:29.261693Z",
     "iopub.status.busy": "2020-11-27T15:34:29.261543Z",
     "iopub.status.idle": "2020-11-27T15:35:47.557846Z",
     "shell.execute_reply": "2020-11-27T15:35:47.557350Z",
     "shell.execute_reply.started": "2020-11-27T15:34:29.261675Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pruning 10921 non-diagonal mobility elements of 11342.\n",
      "nnz before: 4088, after: 528\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",
    "Sgeq0 = [None] * N\n",
    "\n",
    "mob_prun = 0.0006\n",
    "\n",
    "# Pruning mobility once:\n",
    "k=0\n",
    "mobility_history = []\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",
    "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",
    "mobility_history.append(C)\n",
    "\n",
    "mobmat_pr = np.copy(mobmat)\n",
    "\n",
    "\n",
    "mobmat_pr[mobility_history[0] == 0] = 0\n",
    "print(f'nnz before: {np.count_nonzero(mobmat)}, after: {np.count_nonzero(mobmat_pr)}')\n",
    "\n",
    "for k in range(N):\n",
    "    mobK = Params['cov',:,k,'mobility_t'] #mobintime.to_numpy().T[:,k]\n",
    "    betaR =  Params['cov',:,k,'betaratio_t']# betaratiointime.to_numpy().T[:,k]\n",
    "    C = r*mobfrac.flatten()*mobK*mobmat_pr\n",
    "    np.fill_diagonal(C,1-C.sum(axis=1)+ C.diagonal())\n",
    "    \n",
    "    dyn[k] = []\n",
    "    spatial[k] = []\n",
    "    Sgeq0[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 or k == N-1:\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",
    "    #mobility_history.append(C)\n",
    "    \n",
    "    foi_sup = []\n",
    "    foi_inf = []\n",
    "    for n in range(M):\n",
    "        foi_sup.append(betaP0*betaR[n]*(Pk[n]+epsilonA*Ak[n]))\n",
    "        foi_inf.append(Sk[n]+Ek[n]+Pk[n]+Rk[n]+Ak[n])\n",
    "    \n",
    "    foi = []\n",
    "    for m in range(M):\n",
    "        foi.append((sum(C[n, m] * foi_sup[n] for n in range(M)) + epsilonI*betaP0*betaR[m]*Ik[m]) /\n",
    "                   (sum(C[l, m] * foi_inf[l] for l in range(M)) + Ik[m]))\n",
    "        \n",
    "\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",
    "\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",
    "        #mob_ik = sum(C[i, m] * (\n",
    "        #            (sum(C[n, m] * (betaP0*betaR[n]*(Pk[n]+epsilonA*Ak[n])) for n in range(M)) + \n",
    "        #             epsilonI*betaP0*betaR[m]*Ik[m]) /\n",
    "        #            (sum(C[l, m]*(Sk[l]+Ek[l]+Pk[l]+Rk[l]+Ak[l]) for l in range(M)) + Ik[m]))\n",
    "        #             for m in range(M))\n",
    "        mob_ik = sum(C[i, m] * foi[m] for m in range(M))\n",
    "        \n",
    "        spatial[k].append(Vars['u', i, k, 'mob'] - mob_ik) # spatial, vaccines and dyn are put in g(x), with constraints that spatial and dyn are equal to zero\n",
    "        VacPpl = sum(Vars['x', i, k,name] for name in ['S','E','P','A','R'])\n",
    "        Sgeq0[k].append( Vars['x', i, k,'S'] - Vars['u', i, k, 'v']/(VacPpl+1e-10) )\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": 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": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T15:35:47.558984Z",
     "iopub.status.busy": "2020-11-27T15:35:47.558834Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing constraints, ...DONE\n",
      "Building NLP function...DONE\n",
      "Building Jacobian function...DONE\n",
      "Building Solver..."
     ]
    }
   ],
   "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",
    "    cat.entry(\"Sgeq0\", expr=Sgeq0)\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'] = 1e6#2000*(T*.6)*M #8e6 #*M\n",
    "lbg['vaccines'] = -np.inf\n",
    "\n",
    "ubg['Sgeq0'] = np.inf\n",
    "\n",
    "optimize = 1\n",
    "lbx['u', :, :, 'v'] = 0.\n",
    "ubx['u', :, :, 'v'] = 2000 * 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": {},
   "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",
    "\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\n",
    "\n",
    "for i in range(M):\n",
    "    for k in range(N):\n",
    "        arg['p']['cov',i,k,'mobility_t'] =  mobintime.to_numpy().T[i,k]\n",
    "        arg['p']['cov',i,k,'betaratio_t'] = betaratiointime.to_numpy().T[i,k]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#arg['ubx']['u', :, :, 'v']  = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If restart:\n",
    "#for i, name in enumerate(states.keys()):\n",
    "#    for k in range(N + 1):\n",
    "#        for nd in range(M):\n",
    "#            init['x', nd, k, name] = opt['x',nd, k, name]\n",
    "#for k in range(N):\n",
    "#    for nd in range(M):       \n",
    "#        init['u', nd, k, 'v'] = opt['u',nd,k,'v']        \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#pnum.append(scale_ell)\n",
    "#pnum.append(scale_If)\n",
    "#pnum.append(scale_v)\n",
    "pnum[-3] = 1e5\n",
    "pnum[-2] = 0\n",
    "pnum[-1] = 1e-10\n",
    "pnum[-4] = 1/(9*30)  #gammaV \n",
    "arg['p']['p'] = pnum\n",
    "arg['x0'] = init"
   ]
  },
  {
   "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": [
    "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)   # 2times ?\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": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "fig, axes = plt.subplots(2,5, figsize = (20,10))\n",
    "fig.patch.set_facecolor('white')\n",
    "\n",
    "node = 9\n",
    "til =  T\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": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "fig, axes = plt.subplots(5,2, figsize = (10,10))\n",
    "node = 1\n",
    "til =  -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(len(np.array(ca.veccat(ca.veccat(*opt['u',k,:til,'v']))))),\n",
    "                                         np.array(ca.veccat(ca.veccat(*opt['u',k,:til,'v'])))\n",
    "                          )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import networkx\n",
    "import geopandas as gpd"
   ]
  },
  {
   "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": [
    "G = networkx.Graph()\n",
    "G.position = {}\n",
    "G.population = {}\n",
    "G.comp = {}\n",
    "G.epi = {}\n",
    "s.shp['vacc'] = np.nan\n",
    "s.shp['Rend'] = np.nan\n",
    "for i, node in enumerate(s.ind2name):\n",
    "    G.add_node(node)\n",
    "    G.position[node] = (s.pos_node[i,0], s.pos_node[i,1])\n",
    "    G.population[node] = s.pop_node[i]\n",
    "    #G.comp[node] = (ocp.ic['S'][i], ocp.ic['I'][i],ocp.ic['R'][i])\n",
    "    try:\n",
    "        G.epi[node] = {'vacc': sum(np.array(ca.veccat(ca.veccat(*opt['u',i,:,'v']))))[0],\n",
    "                   'Rend':                            float(opt['x',i,-1,'R'])}\n",
    "        s.shp.loc[i, 'vacc'] = sum(np.array(ca.veccat(ca.veccat(*opt['u',i,:,'v']))))[0]\n",
    "        s.shp.loc[i, 'Rend'] = float(opt['x',i,-1,'R'])\n",
    "    except NameError as e:\n",
    "        #print(f'epi data failed, {e}')\n",
    "        G.epi[node] = {'vacc': np.nan,\n",
    "                   'Rend':     np.nan}\n",
    "        s.shp.loc[i, 'vacc'] = np.nan\n",
    "        s.shp.loc[i, 'Rend'] = np.nan\n",
    "        \n",
    "    s.shp.loc[i, 'population'] =  s.pop_node[i] # overwrite \n",
    "    for j, connection in enumerate(mobility_history[0][i]):\n",
    "        if connection != 0:\n",
    "            G.add_edge(node, s.ind2name[j], weight=connection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G.number_of_edges()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1,figsize=(20, 20))\n",
    "\n",
    "\n",
    "networkx.draw(G, \n",
    "         G.position, \n",
    "         node_size=1000/max(s.pop_node) * np.array([G.population[v] for v in G]),\n",
    "          #node_color=[float(G.degree(v)) for v in G],\n",
    "        #node_color=[G.population[v] for v in G],\n",
    "          #node_color=[G.epi[v]['vacc']/G.population[v] for v in G],\n",
    "          node_color= [G.epi[v]['vacc'] for v in G],\n",
    "          width = 150* np.array([max(a['weight'],0.003) for u,v,a in G.edges(data=True)]),\n",
    "          edge_color= 10 * np.array([a['weight'] for u,v,a in G.edges(data=True)]),\n",
    "          edge_cmap = mpl.cm.viridis,\n",
    "          ax = ax,\n",
    "         with_labels=False\n",
    "       )\n",
    "\n",
    "#     # scale the axes equally\n",
    "#plt.xlim(min(s.pos_node[:,0]) - 100000, max(s.pos_node[:,0])+ 100000)\n",
    "#plt.ylim(min(s.pos_node[:,1]) - 100000, max(s.pos_node[:,1])+ 100000)\n",
    "\n",
    "#s.shp.plot(ax = ax, column='' cmap='OrRd', facecolor=\"none\", edgecolor=\"black\")\n",
    "\n",
    "s.shp.boundary.plot(ax = ax,  edgecolor=\"grey\", linewidth = .5)\n",
    "\n",
    "\n",
    "cmap= mpl.cm.viridis\n",
    "vmin = min([G.epi[v]['vacc'] for v in G])\n",
    "vmax = max([G.epi[v]['vacc'] for v in G])\n",
    "sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin = vmin, vmax=vmax))\n",
    "#sm._A = []\n",
    "plt.colorbar(sm)\n",
    "\n",
    "plt.draw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1,figsize=(10, 10))\n",
    "s.shp.plot(ax = ax, column='vacc', cmap='OrRd',  edgecolor=\"black\") #facecolor=\"none\","
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "fig, ax = plt.subplots(1,1,figsize=(4, 4))\n",
    "plt.scatter(s.shp['vacc']/s.shp['population'], s.shp['Rend']/s.shp['population'], c= s.shp['population'])\n",
    "ax.set_xlabel(\"prop. vaccinated\")\n",
    "ax.set_ylabel(\"prop. recovered\");\n",
    "ax.set_xlim(0)\n",
    "ax.set_ylim(0, 0.0002)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.scatterplot(s.shp['vacc'], s.shp['population']*100, hue= s.shp['population'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grph = networkx.from_numpy_matrix(mobility_history[0])\n",
    "networkx.draw(grph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find cutoff:\n",
    "for u,v,a in G.edges('Sassari', data=True):\n",
    "    print(a['weight']) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mobintime.plot(legend=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "betaratiointime.plot(legend=False)"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}