Remote_Simulation_Parameter_Sweeps.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "sys.path.insert(1, os.path.abspath(os.path.join(os.getcwd(), '../GillesPy2')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import time\n",
    "import pickle\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axes_grid1.inset_locator import (inset_axes, InsetPosition,mark_inset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gillespy2 import TauHybridCSolver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client\n",
    "from dask import delayed, compute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from Devils_DFTD_2_Stage_Infection import DevilsDFTD2StageInfection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = Client(n_workers=2, threads_per_worker=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read in observed data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "pop_data = pd.read_csv('Devils_Dataset__Population_1985-2020.csv')\n",
    "devil_pop = np.array(pop_data['Population'].iloc[:].values)\n",
    "\n",
    "obs = np.vstack([devil_pop]).reshape(1, 1, -1)\n",
    "carry_cap = int(max(devil_pop)*1.16)\n",
    "\n",
    "dates = []\n",
    "year = 1985\n",
    "while len(dates) < len(devil_pop):\n",
    "    for month in [\"Jan\", \"Feb\", \"Mar\", \"Apr\", \"May\", \"Jun\", \"Jul\", \"Aug\", \"Sept\", \"Oct\", \"Nov\", \"Dec\"]:\n",
    "        dates.append(f\"{month} {year}\")\n",
    "    year += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "sweeps={}\n",
    "mean_std_db={}\n",
    "type_variables = {\n",
    "    \"juvenile_concentration\": float,\n",
    "    \"birth_rate\": float,\n",
    "    \"maturity_rate\": float,\n",
    "    \"death_rate_juvenile\": float,\n",
    "    \"death_rate_susceptible\": float,\n",
    "    \"death_rate_over_population\": float,\n",
    "    \"infection_rate_infected\": float,\n",
    "    \"infection_rate_diseased\": float,\n",
    "    \"incubation\": float,\n",
    "    \"progression\": float,\n",
    "    \"death_rate_infected\": float,\n",
    "    \"death_rate_diseased\": float,\n",
    "    \"DFTD_start\": float,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_sweep():\n",
    "    for name, v_type in type_variables.items():\n",
    "        sweeps[name] = [variables[name]]\n",
    "        \n",
    "        sweeps[name].insert(0, variables[name] * 0.975)\n",
    "        sweeps[name].insert(0, variables[name] * 0.95)\n",
    "        sweeps[name].insert(0, variables[name] * 0.9)\n",
    "        sweeps[name].insert(0, variables[name] * 0.85)\n",
    "        sweeps[name].insert(0, variables[name] * 0.8)\n",
    "        \n",
    "        sweeps[name].append(variables[name] * 1.025)\n",
    "        sweeps[name].append(variables[name] * 1.05)\n",
    "        sweeps[name].append(variables[name] * 1.1)\n",
    "        sweeps[name].append(variables[name] * 1.15)\n",
    "        sweeps[name].append(variables[name] * 1.2)\n",
    "            \n",
    "        sweeps[name] = sorted(list(set(sweeps[name])))\n",
    "        for n,x in enumerate(sweeps[name]):\n",
    "            if type_variables[name]==int:\n",
    "                x = str(x)\n",
    "            else:\n",
    "                x = f\"{x:.4e}\"\n",
    "            if '00000' in x:\n",
    "                x=re.sub(r\"00000\\d+\",'',x)\n",
    "            elif '99999' in x:\n",
    "                x=re.sub(r\"99999\\d+\",'9',x)\n",
    "            sweeps[name][n]=x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_all_space(variables, sweeps):\n",
    "    # first add the base point\n",
    "    all_space =[variables.copy()]\n",
    "    # then go through each sweep and add the points\n",
    "    for sn, svs in sweeps.items():\n",
    "        for sv in svs:\n",
    "            if variables[sn] == sv: continue # don't re-add base point\n",
    "            v = variables.copy()\n",
    "            if sn == \"juvenile_concentration\":\n",
    "                init_Devils_pop = round(model.devil_pop[0])\n",
    "                init_J_pop = round(model.devil_pop[0] * type_variables[sn](sv))\n",
    "                init_S_pop = round(model.devil_pop[0] - init_J_pop)\n",
    "                v[\"Juvenile\"] = init_J_pop\n",
    "                v[\"Susceptible\"] = init_S_pop\n",
    "            else:\n",
    "                v[sn] = type_variables[sn](sv)\n",
    "            life_span = 1/v['death_rate_diseased'] + v['incubation'] + v['progression']\n",
    "            if life_span < 26 and life_span > 12:\n",
    "                all_space.append(v)\n",
    "    return all_space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_remote(var, sweeps, model, solver, batch_size=15):\n",
    "    jobs = []\n",
    "    keys = []\n",
    "    all_space = generate_all_space(var, sweeps)\n",
    "    for space in all_space:\n",
    "        filename = make_filename(space)\n",
    "        if os.path.exists(filename + \".m\"):\n",
    "            print(f\"Found {filename}.m, skipping\")\n",
    "        \n",
    "        job = delayed(model.run)(num_sims=1000, solver=solver, variables=space)\n",
    "        batch = len(jobs)\n",
    "        if batch == 0 or len(jobs[batch - 1]) == batch_size:\n",
    "            jobs.append([job])\n",
    "            keys.append([filename])\n",
    "        else:\n",
    "            jobs[batch-1].append(job)\n",
    "            keys[batch-1].append(filename)\n",
    "    return jobs, keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_remote(model, jobs, keys, batch_size=15):\n",
    "    total_sims = 0 if len(jobs) == 0 else batch_size * (len(jobs) - 1) + len(jobs[-1])\n",
    "    print(f\"Running {total_sims} new parameter points\", end=\" \")\n",
    "    print(f\"in {len(jobs)} batches with {batch_size} points per batch\")\n",
    "    \n",
    "    for i, batch in enumerate(jobs):\n",
    "        results = dict(zip(keys[i], compute(*batch)))\n",
    "        \n",
    "        for filename, result in results.items():\n",
    "            (m, s) = model.calculate_distance(result)\n",
    "            \n",
    "            with open(f\"{filename}.m\", \"w\") as fd:\n",
    "                fd.write(f\"{m},{s}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_filename(variables):\n",
    "    cmdpath = \"ParameterSweeps/tmp_result_state\"\n",
    "    ret = f\"{cmdpath}/ps-\"\n",
    "    for k in sorted(variables.keys()):\n",
    "        v = float(variables[k])\n",
    "        x = f\"{v:.4e}\"\n",
    "        ret += f\"{x},\"\n",
    "    return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_mean_std(var):\n",
    "    f = make_filename(var)\n",
    "    if f in mean_std_db:\n",
    "        return mean_std_db[f]\n",
    "    \n",
    "    file = f\"{f}.m\"\n",
    "    try:\n",
    "        with open(file, \"r\") as fd:\n",
    "            resp = fd.read()\n",
    "            (m,s) = resp.split(\",\",2)\n",
    "            mean_std_db[f] = (float(m),float(s))\n",
    "            return mean_std_db[f]\n",
    "    except Exception as e:\n",
    "        time.sleep(0.2)\n",
    "        raise e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_sweep(name):\n",
    "    xvals = np.zeros(len(sweeps[name]))\n",
    "    mvals = np.zeros(len(sweeps[name]))\n",
    "    svals = np.zeros(len(sweeps[name]))\n",
    "    for n,v in enumerate(sweeps[name]):\n",
    "        c = variables.copy()\n",
    "        c[name] = type_variables[name](v)\n",
    "        keepgoing=True\n",
    "        while keepgoing:\n",
    "            try:\n",
    "                (m,s) = get_mean_std(c)\n",
    "                xvals[n]=c[name]\n",
    "                mvals[n]=m\n",
    "                svals[n]=s\n",
    "                keepgoing=False\n",
    "            except Exception as e:\n",
    "                print(f\"caught e={e} while name={name} \")\n",
    "                time.sleep(1)\n",
    "    \n",
    "    plt.figure(figsize=[12, 6])\n",
    "    plt.errorbar(xvals,mvals,yerr=svals, capsize=10)\n",
    "    plt.plot([variables[name], variables[name]],[0, max(mvals)],'--')\n",
    "    plt.xlabel(name, fontsize=12)\n",
    "    plt.ylabel('Error', fontsize=12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_eresult(eresults): \n",
    "    fig, ax1 = plt.subplots(figsize=[12, 6])\n",
    "\n",
    "    plt.title(\"Tasmanian Devil Population with DFTD: Observed vs. Simulated\", fontsize=18)\n",
    "    ax1.set_xlabel(f\"Time (months) since {dates[0]}\", fontsize=16)\n",
    "    ax1.set_ylabel(\"Population of Tasmanian Devils\", fontsize=16)\n",
    "    ax1.plot(eresults[0]['time'], obs[0][0], '--', color='black', label='Observed Total')\n",
    "    ax1.plot(eresults[0]['time'], eresults[0]['Devils'], color='blue', label='Simulated Total')\n",
    "    ax1.plot(eresults[0]['time'], eresults[0]['Juvenile'], color='purple', alpha=.6, label='Juvenile')\n",
    "    ax1.plot(eresults[0]['time'], eresults[0]['Susceptible'], color='green', alpha=.6, label='Susceptible')\n",
    "    ax1.plot(eresults[0]['time'], eresults[0]['Exposed'], color='magenta', alpha=.6, label='Exposed')\n",
    "    ax1.plot(eresults[0]['time'], eresults[0]['Infected'], color='red', alpha=.6, label='Infected')\n",
    "    ax1.plot(eresults[0]['time'], eresults[0]['Diseased'], color='cyan', alpha=.6, label='Diseased')\n",
    "    ax1.plot([variables['DFTD_start'], variables['DFTD_start']], [-3000, carry_cap], '--k', alpha=0.4)\n",
    "    ax1.text(variables['DFTD_start']-7, 5000, \"DFTD Start\", rotation=\"vertical\", color=\"black\", fontsize=14)\n",
    "    ax1.text(variables['DFTD_start']-7, 24000, dates[variables['DFTD_start']], rotation=\"vertical\", color=\"black\", fontsize=14)\n",
    "    ax1.tick_params(axis='x', labelsize=12)\n",
    "    ax1.set_yticks([20000,40000,60000])\n",
    "    ax1.tick_params(axis='y',labelsize=12, labelrotation=90)\n",
    "    ax1.legend(loc='upper right', fontsize=16)\n",
    "    ax1.set_ylim([0, carry_cap])\n",
    "    ax1.set_xlim(0,eresults[0]['time'][-1])\n",
    "    for n,r in enumerate(eresults):\n",
    "        if n==0: continue\n",
    "        ax1.plot(eresults[n]['time'], eresults[n]['Devils'],'b', alpha=0.025)\n",
    "        ax1.plot(eresults[n]['time'], eresults[n]['Juvenile'], color='purple', alpha=0.025, label='Juvenile')\n",
    "        ax1.plot(eresults[n]['time'], eresults[n]['Susceptible'], color='green', alpha=0.025, label='Susceptible')\n",
    "        ax1.plot(eresults[n]['time'], eresults[n]['Exposed'], color='magenta', alpha=0.025, label='Exposed')\n",
    "        ax1.plot(eresults[n]['time'], eresults[n]['Infected'], color='red', alpha=0.025, label='Infected')\n",
    "        ax1.plot(eresults[n]['time'], eresults[n]['Diseased'], color='cyan', alpha=0.025, label='Diseased')\n",
    "\n",
    "    print(variables)\n",
    "    \n",
    "    def calculate_distance(eresults):\n",
    "        '''return mean/stddev of L2 norm distance'''\n",
    "        global obs\n",
    "        dists = np.zeros(len(eresults))\n",
    "        for n,r in enumerate(eresults):\n",
    "            dists[n] = np.linalg.norm(r['Devils']-obs[0][0],2)\n",
    "        return np.average(dists), np.std(dists)\n",
    "\n",
    "    dist_l2 = calculate_distance(eresults)\n",
    "    print(f\"L2 = {dist_l2[0]:.2f} +/-{dist_l2[1]:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run Parameter Sweep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "variables = {\n",
    "    'juvenile_concentration': 0.49534348836011316,\n",
    "    'birth_rate': 0.055, \n",
    "    'maturity_rate': 0.04,#0.04167,\n",
    "    'death_rate_juvenile': 0.007,#0.006, \n",
    "    'death_rate_over_population': 2.3e-07, \n",
    "    'death_rate_susceptible': 0.02335, \n",
    "    'incubation': 10.25,#10, #10.99687624550675, \n",
    "    'progression': 10.74,#11.015,#10.746230534983676,\n",
    "    'infection_rate_diseased': 3.84e-05,#3.2e-05, #3.4e-05 #3e-05, #4.978182435648742e-05, \n",
    "    'infection_rate_infected': 1e-05,#1.0698e-05,#1.1261e-05,#1.155e-05,#1.1e-05, #1.2e-05 #1e-5, #1.4809664001475363e-05, \n",
    "    'death_rate_diseased': 0.29017,#0.25232,#0.22938,#0.23526,#0.27678,#0.29134996217062514, \n",
    "    'death_rate_infected': 0.022609,#0.01966,#0.020695079156445997, \n",
    "    'DFTD_start': 100,#97,#95, #105, \n",
    "}\n",
    "\n",
    "variables_orig = variables.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "setup_sweep()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "sweeps['DFTD_start']=['85','90','92','94','96','98','100','102','104','106','108','110','115']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"DFTD_start\": 100,\n",
      "    \"birth_rate\": 0.055,\n",
      "    \"death_rate_diseased\": 0.29017,\n",
      "    \"death_rate_infected\": 0.022609,\n",
      "    \"death_rate_juvenile\": 0.007,\n",
      "    \"death_rate_over_population\": 2.3e-07,\n",
      "    \"death_rate_susceptible\": 0.02335,\n",
      "    \"incubation\": 10.25,\n",
      "    \"infection_rate_diseased\": 3.84e-05,\n",
      "    \"infection_rate_infected\": 1e-05,\n",
      "    \"maturity_rate\": 0.04,\n",
      "    \"progression\": 10.74\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(json.dumps(variables, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = DevilsDFTD2StageInfection()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver = delayed(TauHybridCSolver)(model=model, variable=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found ParameterSweeps/tmp_result_state/ps-1.0000e+02,5.5000e-02,2.9017e-01,2.2609e-02,7.0000e-03,2.3000e-07,2.3350e-02,1.0250e+01,3.8400e-05,1.0000e-05,4.0000e-02,1.0740e+01,.m, skipping\n",
      "Found ParameterSweeps/tmp_result_state/ps-1.0000e+02,5.5000e-02,2.9017e-01,2.2609e-02,7.0000e-03,2.3000e-07,2.3350e-02,1.0250e+01,3.8400e-05,1.0000e-05,4.0000e-02,1.0740e+01,.m, skipping\n",
      "Found ParameterSweeps/tmp_result_state/ps-1.0000e+02,5.5000e-02,2.9017e-01,2.2609e-02,7.0000e-03,2.3000e-07,2.3350e-02,1.0250e+01,3.8400e-05,1.0000e-05,4.0000e-02,1.0740e+01,.m, skipping\n",
      "Found ParameterSweeps/tmp_result_state/ps-1.0000e+02,5.5000e-02,2.9017e-01,2.2609e-02,7.0000e-03,2.3000e-07,2.3350e-02,1.0250e+01,3.8400e-05,1.0000e-05,4.0000e-02,1.0740e+01,.m, skipping\n",
      "Found ParameterSweeps/tmp_result_state/ps-1.0000e+02,5.5000e-02,2.9017e-01,2.2609e-02,7.0000e-03,2.3000e-07,2.3350e-02,1.0250e+01,3.8400e-05,1.0000e-05,4.0000e-02,1.0740e+01,.m, skipping\n",
      "Found ParameterSweeps/tmp_result_state/ps-1.0000e+02,5.5000e-02,2.9017e-01,2.2609e-02,7.0000e-03,2.3000e-07,2.3350e-02,1.0250e+01,3.8400e-05,1.0000e-05,4.0000e-02,1.0740e+01,.m, skipping\n",
      "Found ParameterSweeps/tmp_result_state/ps-1.0000e+02,5.5000e-02,2.9017e-01,2.2609e-02,7.0000e-03,2.3000e-07,2.3350e-02,1.0250e+01,3.8400e-05,1.0000e-05,4.0000e-02,1.0740e+01,.m, skipping\n",
      "Found ParameterSweeps/tmp_result_state/ps-1.0000e+02,5.5000e-02,2.9017e-01,2.2609e-02,7.0000e-03,2.3000e-07,2.3350e-02,1.0250e+01,3.8400e-05,1.0000e-05,4.0000e-02,1.0740e+01,.m, skipping\n",
      "Found ParameterSweeps/tmp_result_state/ps-1.0000e+02,5.5000e-02,2.9017e-01,2.2609e-02,7.0000e-03,2.3000e-07,2.3350e-02,1.0250e+01,3.8400e-05,1.0000e-05,4.0000e-02,1.0740e+01,.m, skipping\n",
      "Found ParameterSweeps/tmp_result_state/ps-1.0000e+02,5.5000e-02,2.9017e-01,2.2609e-02,7.0000e-03,2.3000e-07,2.3350e-02,1.0250e+01,3.8400e-05,1.0000e-05,4.0000e-02,1.0740e+01,.m, skipping\n",
      "Found ParameterSweeps/tmp_result_state/ps-1.0000e+02,5.5000e-02,2.9017e-01,2.2609e-02,7.0000e-03,2.3000e-07,2.3350e-02,1.0250e+01,3.8400e-05,1.0000e-05,4.0000e-02,1.0740e+01,.m, skipping\n",
      "Found ParameterSweeps/tmp_result_state/ps-1.0000e+02,5.5000e-02,2.9017e-01,2.2609e-02,7.0000e-03,2.3000e-07,2.3350e-02,1.0250e+01,3.8400e-05,1.0000e-05,4.0000e-02,1.0740e+01,.m, skipping\n",
      "Found ParameterSweeps/tmp_result_state/ps-1.0000e+02,5.5000e-02,2.9017e-01,2.2609e-02,7.0000e-03,2.3000e-07,2.3350e-02,1.0250e+01,3.8400e-05,1.0000e-05,4.0000e-02,1.0740e+01,.m, skipping\n",
      "CPU times: user 51.9 ms, sys: 4.94 ms, total: 56.9 ms\n",
      "Wall time: 49.3 ms\n"
     ]
    }
   ],
   "source": [
    "%time jobs, keys = load_remote(var=variables, sweeps=sweeps, model=model, solver=solver)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running 132 new parameter points in 9 batches with 15 points per batch\n"
     ]
    }
   ],
   "source": [
    "%time run_remote(model=model, jobs=jobs, keys=keys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "for k in sweeps.keys():\n",
    "    plot_sweep(k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model1 = DevilsDFTD2StageInfection(devil_fitting=True)\n",
    "job1 = delayed(model1.run)(num_sims=100, solver=solver, variables=variables)\n",
    "%time eresults1 = job1.compute()\n",
    "plot_eresult(eresults1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_devil_pop = eresults1[0]['Devils'][-1]\n",
    "\n",
    "print(f\"Number of Devils {final_devil_pop} at simulation end,\",\n",
    "      f\"{final_devil_pop/carry_cap * 100:.1f}% of target carrying capacity\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_devil_pops = []\n",
    "for n, r in enumerate(eresults1):\n",
    "    final_devil_pops.append(r['Devils'][-1])\n",
    "    \n",
    "print(f\"Number of Devils = {np.average(final_devil_pops):.0f} +/-\"\n",
    "      f\"{np.std(final_devil_pops):.0f} at simulation end,\",\n",
    "      f\"{np.average(final_devil_pops)/carry_cap*100:.1f}% of target carrying capacity\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_data1 = {\n",
    "    \"variables\": variables, \"sweeps\": sweeps, \"eresults\": eresults1, \"mean_std_db\": mean_std_db\n",
    "}\n",
    "print(sim_data1)\n",
    "# with open('figure_data/Fig3A_Fig14A-data.p','wb+') as fd:\n",
    "#     pickle.dump(sim_data1, fd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model2 = DevilsDFTD2StageInfection()\n",
    "job2 = delayed(model2.run)(num_sims=100, solver=solver, variables=variables)\n",
    "%time eresults2 = job2.compute()\n",
    "plot_eresult(eresults2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_data2 = {\n",
    "    \"variables\": variables, \"sweeps\": sweeps, \"eresults\": eresults2, \"mean_std_db\": mean_std_db\n",
    "}\n",
    "print(sim_data2)\n",
    "# with open('figure_data/Fig3B_Fig14B-C-data.p','wb+') as fd:\n",
    "#     pickle.dump(sim_data2, fd)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}