Exploration_Stochastic.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "76a56be8-e192-4808-bb70-c7753aebc85e",
   "metadata": {},
   "source": [
    "# Gaussian Process Emulation of the Deterministic BusSim Model\n",
    "\n",
    "This notebook extends the work from the previous notebook, moving from the deterministic `BusSim` model to the stochastic one.\n",
    "The stochastic `BusSim` model introduces random variation into the number of passengers who want to board the bus at each stop - it is now a Poisson process rather than a constant rate.\n",
    "\n",
    "This poses an interesting challenge for model emulation.\n",
    "We could now conceivably have different outputs from the model for the same inputs.\n",
    "In this notebook, we will explore ways in which we can still use a GP to emulate this type of system.\n",
    "\n",
    "## Aims\n",
    "\n",
    "- Explore the impact of changing traffic speed on the stochastic bus model\n",
    "- Fit a GP emulator to the model and investigate different ways of handling the stochasticity\n",
    "- Evaluate the quality of the GP emulator\n",
    "- Discuss potential extensions to these approaches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a684402-346a-4fd6-baa4-8e156be3849f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from BusSim_stochastic import run_model as run_model_s\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import mogp_emulator\n",
    "\n",
    "# Increase default figure size\n",
    "plt.rcParams[\"figure.figsize\"] = [12, 6]\n",
    "plt.rcParams[\"font.size\"] = 16\n",
    "\n",
    "np.random.seed(0)  # for reproducibility, if running entire notebook from the beginning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7bf75ec-a157-423c-81cd-8dc48e527bf1",
   "metadata": {},
   "source": [
    "## Set up the model\n",
    "\n",
    "We will move fairly quickly through this section as it is rather similar to the analogous section in the notebook that describes the deterministic case.\n",
    "\n",
    "First, we set up a function to interface with the original `BusSim` code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dde86c46-fee5-4293-876e-6b1fe662fa15",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_stochastic_model(ArrivalRate=None, TrafficSpeed=14, bus_ids_out=[0, 1, 2], maxDemand=2, DEBUG=False):\n",
    "    \n",
    "    NumberOfStop = 20\n",
    "    minDemand = 0.5\n",
    "     \n",
    "    # Initialise the remaining model parameters\n",
    "    model_params = {\n",
    "        \"dt\": 10,\n",
    "        \"minDemand\": minDemand,        \n",
    "        \"NumberOfStop\": NumberOfStop,\n",
    "        \"LengthBetweenStop\": 2000,\n",
    "        \"EndTime\": 6000,\n",
    "        \"Headway\": 5 * 60,\n",
    "        \"BurnIn\": 1 * 60,\n",
    "        \"AlightTime\": 1,\n",
    "        \"BoardTime\": 3,\n",
    "        \"StoppingTime\": 3,\n",
    "        \"BusAcceleration\": 3\n",
    "    }\n",
    "    \n",
    "    if ArrivalRate is None:\n",
    "        ArrivalRate = np.random.uniform(minDemand / 60, maxDemand / 60, NumberOfStop) \n",
    "    else:\n",
    "        ArrivalRate = ArrivalRate * np.ones(NumberOfStop)\n",
    "        \n",
    "    DepartureRate = np.sort(np.random.uniform(0.05, 0.5, NumberOfStop)) # Sorted as more passengers get off near the end of the route\n",
    "\n",
    "    model, model_params, ArrivalData, StateData, GroundTruth, GPSData = \\\n",
    "        run_model_s(model_params, TrafficSpeed, ArrivalRate, DepartureRate, False, False, False, True)\n",
    "            \n",
    "    bus_pos = np.array([bus.trajectory for bus in model.buses])    \n",
    "    time = np.arange(0, model.EndTime, model.dt)\n",
    "    total_distance = model_params[\"NumberOfStop\"] * model_params[\"LengthBetweenStop\"]\n",
    "\n",
    "    bus_t_end = [time[np.argmax(bp>=total_distance)] for bp in bus_pos]\n",
    "    bus_end_times = [bus_t_end[id] for id in bus_ids_out]\n",
    "    \n",
    "    if DEBUG:\n",
    "        print(bus_t_end)\n",
    "    \n",
    "    return bus_end_times, GPSData"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0cf3fe7-d01a-411c-8e2e-87cffc7b7c15",
   "metadata": {},
   "source": [
    "## Initial look at variation\n",
    "\n",
    "We'll start by taking a look at the variability in bus arrival times for two different values of traffic speed.\n",
    "This will let us see how much variability arises from the model's stochasticity as opposed to the traffic speed parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cc86f6f-f457-40f8-be32-7b4f4c376574",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "bus_ids = range(8)\n",
    "repeats = 10\n",
    "\n",
    "md = 6\n",
    "ts = 14\n",
    "ts_fast = 28\n",
    "\n",
    "end_times = np.zeros((len(bus_ids), repeats))\n",
    "gps_locs = []\n",
    "for r in range(repeats):\n",
    "    t, g = run_stochastic_model(TrafficSpeed=ts, bus_ids_out=bus_ids, maxDemand=md)\n",
    "    end_times[:,r] = t\n",
    "    gps_locs.append(g)\n",
    "    \n",
    "end_times_fast = np.zeros((len(bus_ids), repeats))\n",
    "gps_locs_fast = []\n",
    "for r in range(repeats):\n",
    "    t, g = run_stochastic_model(TrafficSpeed=ts_fast, bus_ids_out=bus_ids, maxDemand=md)\n",
    "    end_times_fast[:,r] = t\n",
    "    gps_locs_fast.append(g)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48a4f955-fe2e-429d-84db-9efe0cf366ff",
   "metadata": {},
   "source": [
    "First let's check out what happens with the default speed of 14."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df24b380-6f02-41ee-b5a7-4ca68896cf03",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(bus_ids, end_times, marker=\"o\", linestyle=\"-\", linewidth=0.7)\n",
    "plt.ylim((3000, 7000))\n",
    "plt.xlabel(\"Bus ID\")\n",
    "plt.ylabel(\"Arrival time at final stop\")\n",
    "plt.title(\"Arrival times of buses at final stop using stochastic model\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d15bf01b-c55a-4136-bb8a-2c268d837749",
   "metadata": {},
   "source": [
    "In the plot above, we can see that the introduction of stocasticity causes changes in the arrival time which are appreciable (although not large) when compared to the value of the arrival time itself.\n",
    "\n",
    "Simulations contain a mix of \"early\" and \"late\" buses, and we note that some of the buses that depart later do not reach their final stop within the simulated time period.\n",
    "\n",
    "We will also compare those buses against others travelling at a much faster speed.\n",
    "We have dropped colours for clarity in the next; slow buses (with the default speed of 14) are shown in black and fast ones in blue."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7855cdb9-8a12-40b2-93ab-1122406c9091",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(bus_ids, end_times, color=\"k\", marker=\"o\", linestyle=\"-\", linewidth=0.7);\n",
    "plt.plot(bus_ids, end_times_fast, color=\"b\", marker=\"o\", linestyle=\"-\");\n",
    "plt.ylim((2000, 7000));\n",
    "plt.xlabel(\"Bus ID\");\n",
    "plt.ylabel(\"Arrival time at final stop\")\n",
    "plt.title(\"Arrival times of buses at final stop using stochastic model\")\n",
    "plt.legend(loc=\"center left\", bbox_to_anchor=(1.05, 0.5),\n",
    "           handles=[matplotlib.patches.Patch(facecolor=\"k\", edgecolor=None, label=\"Default speed, {}\".format(ts)),\n",
    "                    matplotlib.patches.Patch(facecolor=\"b\", edgecolor=None, label=\"Fast speed, {}\".format(ts_fast))]);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7afcd113-2696-4cb0-8772-aa54b212649c",
   "metadata": {},
   "source": [
    "## Using a GP to emulate bus arrival times with variable traffic speed\n",
    "\n",
    "Before we start, what might we expect to see as an output from the GP this time?\n",
    "\n",
    "In the previous notebook, we used a GP to predict final arrival times for different traffic speeds.\n",
    "That GP generally showed very low uncertainty, apart from at the limits of the traffic speed parameter space and sometimes also for Bus 0, which was often overtaken by the bus that followed it when the traffic speeds were higher."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db816e17-1d0f-44db-9883-128fbc3cd08f",
   "metadata": {},
   "source": [
    "### Generate training data for the GP\n",
    "\n",
    "We'll consider the same set-up as last time: we will vary traffic speed between 15 and 50.\n",
    "The difference this time is that we will run repeated simulations at each of those training points, so that we have `n_repeats` (by default here, 5) repeats per value of traffic speed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24e6d9d3-6ee9-44e2-9c53-ea8bd6c44b6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_bus_ids = [0, 1, 2, 3, 6, 9]\n",
    "\n",
    "lhd = mogp_emulator.LatinHypercubeDesign([(15.0, 50.0)])\n",
    "\n",
    "n_simulations = 20\n",
    "n_repeats = 5\n",
    "n_sims = n_simulations * n_repeats\n",
    "lhd_traffic_speeds = lhd.sample(n_simulations)\n",
    "\n",
    "plt.figure(figsize=[12, 2])\n",
    "plt.eventplot(lhd_traffic_speeds[:,0], lineoffsets=1)\n",
    "plt.ylim((0.4,1.6))\n",
    "plt.yticks([])\n",
    "plt.xlabel(\"traffic speed\")\n",
    "plt.title(\"Sampled values of traffic speed used to fit GP\");\n",
    "\n",
    "# Now modify the array so that repeated values are included\n",
    "lhd_traffic_speeds = np.repeat(lhd_traffic_speeds, n_repeats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d99ec434-d604-462a-a6ed-219406cea45b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "\n",
    "ts_simulation_output = np.zeros((len(ts_bus_ids), n_sims))\n",
    "\n",
    "for ind, ts in enumerate(lhd_traffic_speeds):\n",
    "    sim_out, _ = run_stochastic_model(TrafficSpeed=ts, bus_ids_out=ts_bus_ids)\n",
    "    ts_simulation_output[:,ind] = np.array(sim_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1fee730-7979-4607-94d8-5cae29659869",
   "metadata": {},
   "source": [
    "### Fit the GP\n",
    "\n",
    "We'll actually fit several different models here to try out the effects of changing the kernel.\n",
    "\n",
    "For a more thorough introduction to the various options available for kernels, see Murphy's [“Probabilistic Machine Learning: An Introduction”]( https://probml.github.io/pml-book/) (2021 edition).\n",
    "Briefly, the two kernels we consider here are Mercer kernels (or positive definite kernels, where the Gram matrix is positive definite for a set of distinct inputs).\n",
    "\n",
    "When we make a prediction via a GP emulator, we need a measure of how similar our input is to the existing input values, as our output will be a weighted combination of the existing inputs.\n",
    "To determine the weights, we need a similarity metric - this is known as a kernel.\n",
    "\n",
    "The squared exponential kernel is one of the most commonly used:\n",
    "  $$ \\mathcal{K}(\\mathbf{x}, \\mathbf{x'}) = \\exp \\left( - \\frac{|| \\mathbf{x} - \\mathbf{x'} ||^2}{2\\ell^2}\\right)$$\n",
    "There is one hyperparameter, the lengthscale or bandwidth, $\\ell$.\n",
    "For new input $\\mathbf{x'}$ which is similar to $\\mathbf{x}$, the value of $\\mathcal{K}(\\mathbf{x}, \\mathbf{x'})$ will also be large, so the predicted output will be heavily weighted towards that of $\\mathbf{x}$.\n",
    "\n",
    "We also consider Matern kernels, which give rise to functions which are less smooth than the squared exponential (which is infinitely differentiable).\n",
    "\n",
    "They have two hyperparameters: $\\nu$ (from Bessel function order) and $\\ell$.\n",
    "- $\\nu=1/2$ gives velocity of a particle undergoing Brownian motion, which yields a very jagged function\n",
    "- Higher $\\nu$ values give smoother functions\n",
    "\n",
    "We will use the 5/2 Matern kernel, as it produces functions that are less smooth than the (default) squared exponential.\n",
    "Therefore, it may be better at capturing the variability in arrival times.\n",
    "    \n",
    "There are still further kernels available, some of which cater for more specific cases such as periodic functions.\n",
    "In addition to the book recommended above, Duvenaud's [Kernel Cookbook](https://www.cs.toronto.edu/~duvenaud/cookbook/) is another useful resource on kernel choice for GPs.\n",
    "\n",
    "For each kernel, we fit two GPs.\n",
    "This is to let us investigate the effect of the nugget, which is a small term added to the diagonal of the covariance matrix.\n",
    "Normally, this would help improve the efficiency of the covariance matrix inversion, but the nugget can also be used as a proxy for a small amount of stochastic variation in the simulation.\n",
    "\n",
    "The two options that we try are `adaptive`, where the value of the nugget is increased over again until the system can be solved, and `fit`, where the value of the nugget is treated as a hyperparameter which should be fully optimised.\n",
    "These options (along with alternatives) are discussed in the documentation of the `mogp_emulator` package: see [nugget definition](https://mogp-emulator.readthedocs.io/en/latest/methods/def/DefNugget.html) and notes on its [implementation](https://mogp-emulator.readthedocs.io/en/latest/implementation/GaussianProcess.html#gaussianprocess) in the `GaussianProcess` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82ab9625-f7eb-4e1c-afb5-305960924cc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_gp_se = mogp_emulator.MultiOutputGP(lhd_traffic_speeds, ts_simulation_output, kernel=mogp_emulator.Kernel.SquaredExponential(), nugget=\"adaptive\")\n",
    "ts_gp_se = mogp_emulator.fit_GP_MAP(ts_gp_se)\n",
    "\n",
    "ts_gp_se_nfit = mogp_emulator.MultiOutputGP(lhd_traffic_speeds, ts_simulation_output, kernel=mogp_emulator.Kernel.SquaredExponential(), nugget=\"fit\")\n",
    "ts_gp_se_nfit = mogp_emulator.fit_GP_MAP(ts_gp_se_nfit)\n",
    "\n",
    "ts_gp_m52 = mogp_emulator.MultiOutputGP(lhd_traffic_speeds, ts_simulation_output, kernel=mogp_emulator.Kernel.Matern52(), nugget=\"adaptive\")\n",
    "ts_gp_m52 = mogp_emulator.fit_GP_MAP(ts_gp_m52)\n",
    "\n",
    "ts_gp_m52_nfit = mogp_emulator.MultiOutputGP(lhd_traffic_speeds, ts_simulation_output, kernel=mogp_emulator.Kernel.Matern52(), nugget=\"fit\")\n",
    "ts_gp_m52_nfit = mogp_emulator.fit_GP_MAP(ts_gp_m52_nfit)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e21c538-b1ec-492e-a91f-48a5ad6f5f12",
   "metadata": {},
   "source": [
    "We will see some warnings with the following message as we run the code above: `Matrix not positive definite, skipping this iteration`\n",
    "\n",
    "This is generated because there are different values of the final arrival time for the same input value of traffic speed.\n",
    "\n",
    "### Try out the GP\n",
    "\n",
    "We will now make some predictions from each of our trained Gaussian Process emulators.\n",
    "\n",
    "We set the `include_nugget` option to `False` below so that the nugget is not included in our uncertainty bounds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28df787d-d93d-4f57-9101-8c8f00852e8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "n_validation = 50\n",
    "\n",
    "validation_traffic_speeds = lhd.sample(n_validation)\n",
    "\n",
    "nugget_setting = False\n",
    "\n",
    "# Generate a set of predictions from each of out GP emulators\n",
    "predicted_ts_se = ts_gp_se.predict(validation_traffic_speeds, include_nugget=nugget_setting)\n",
    "predicted_ts_se_nfit = ts_gp_se_nfit.predict(validation_traffic_speeds, include_nugget=nugget_setting)\n",
    "predicted_ts_m52 = ts_gp_m52.predict(validation_traffic_speeds, include_nugget=nugget_setting)\n",
    "predicted_ts_m52_nfit = ts_gp_m52_nfit.predict(validation_traffic_speeds, include_nugget=nugget_setting)\n",
    "\n",
    "# Perform simulations based on the same input values\n",
    "actual_ts = np.zeros((len(ts_bus_ids), n_validation))\n",
    "for ind, ts in enumerate(validation_traffic_speeds):\n",
    "    sim_out, _ = run_stochastic_model(TrafficSpeed=ts[0], bus_ids_out=ts_bus_ids)\n",
    "    actual_ts[:,ind] = np.array(sim_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46fa3b51-f010-4abb-b303-dc9e8c71239a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function for plotting data used to train a gp, the simulated outputs and emulated outputs\n",
    "def plot_gp_train_and_pred(train_inputs, test_inputs, train_outputs, simulated_outputs, predicted_outputs, bus_ids, plt_bus_index=None):\n",
    "    \n",
    "    # Set the markers used in the plot\n",
    "    plt_mk = { \"train\": \".\",      # Used to train the GP\n",
    "               \"pred\": \"x\",       # Predicted by the GP\n",
    "               \"sim\": \"+\" }       # Same TS as \"pred\", full simulation\n",
    "    plt_mk_legend = [matplotlib.lines.Line2D([0], [0], marker=plt_mk[\"train\"], color=\"w\", label=\"Training\",\n",
    "                     markerfacecolor=\"k\", markeredgecolor=\"k\"), \n",
    "                     matplotlib.lines.Line2D([0], [0], marker=plt_mk[\"pred\"], color=\"w\", label=\"Predicted (GP)\",\n",
    "                     markerfacecolor=\"k\", markeredgecolor=\"k\"),\n",
    "                     matplotlib.lines.Line2D([0], [0], marker=plt_mk[\"sim\"], color=\"w\", label=\"Simulated\",\n",
    "                     markerfacecolor=\"k\", markeredgecolor=\"k\")]\n",
    "    \n",
    "    legend_entries = []\n",
    "\n",
    "    # Easier to look at one bus at a time\n",
    "    if plt_bus_index:\n",
    "        plt.title(\"Bus {}\".format(bus_ids[plt_bus_index]))\n",
    "\n",
    "    for ind, (a, p, u) in enumerate(zip(simulated_outputs, predicted_outputs.mean, predicted_outputs.unc)):\n",
    "\n",
    "        if plt_bus_index is None or ind == plt_bus_index:\n",
    "            plt.plot(test_inputs, a, marker=plt_mk[\"sim\"], linestyle=\"\", color=plt.cm.Set1(ind))\n",
    "            plt.plot(test_inputs, p, marker=plt_mk[\"pred\"], linestyle=\"\", color=plt.cm.Set1(ind))\n",
    "            plt.plot(train_inputs, train_outputs[ind,], marker=plt_mk[\"train\"], linestyle=\"\", color=plt.cm.Set1(ind))\n",
    "\n",
    "            legend_entries.append(matplotlib.patches.Patch(facecolor=plt.cm.Set1(ind), edgecolor=None, label=\"Bus {}\".format(bus_ids[ind])))\n",
    "\n",
    "            o = np.argsort(validation_traffic_speeds, axis=None)\n",
    "            plt.fill_between(validation_traffic_speeds.flatten()[o], p[o]-u[o], p[o]+u[o],\n",
    "                             color=plt.cm.Set1(ind), alpha=0.2, linewidth=0)\n",
    "\n",
    "    plt.xlabel(\"Traffic speed\")\n",
    "    plt.ylabel(\"Time at final\\ndestination\")\n",
    "    plt.legend(handles=legend_entries+plt_mk_legend, fontsize=\"small\", loc=\"center right\", bbox_to_anchor=(1.3,0.5))\n",
    "    plt.ylim((0, 8000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d828b133-294e-4bf7-95ee-e0f05e8ef3b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_bus_index = 0  # Change to examine GP emulator for different buses\n",
    "\n",
    "fig, axs = plt.subplots(2,2,figsize=(25,15))\n",
    "fig.subplots_adjust(wspace=0.5, hspace=0.3)\n",
    "\n",
    "plt.sca(axs[0,0])\n",
    "plot_gp_train_and_pred(lhd_traffic_speeds, validation_traffic_speeds, \n",
    "                      ts_simulation_output, actual_ts, predicted_ts_se, ts_bus_ids, plt_bus_index);\n",
    "plt.title(\"Squared exponential kernel (adaptive)\");\n",
    "\n",
    "plt.sca(axs[1,0])\n",
    "plot_gp_train_and_pred(lhd_traffic_speeds, validation_traffic_speeds, \n",
    "                      ts_simulation_output, actual_ts, predicted_ts_se_nfit, ts_bus_ids, plt_bus_index);\n",
    "plt.title(\"Squared exponential kernel (fit)\");\n",
    "\n",
    "plt.sca(axs[0,1])\n",
    "plot_gp_train_and_pred(lhd_traffic_speeds, validation_traffic_speeds, \n",
    "                      ts_simulation_output, actual_ts, predicted_ts_m52, ts_bus_ids, plt_bus_index);\n",
    "plt.title(\"Matern 5-2 kernel (adaptive)\");\n",
    "\n",
    "plt.sca(axs[1,1])\n",
    "plot_gp_train_and_pred(lhd_traffic_speeds, validation_traffic_speeds, \n",
    "                      ts_simulation_output, actual_ts, predicted_ts_m52_nfit, ts_bus_ids, plt_bus_index);\n",
    "plt.title(\"Matern 5-2 kernel (fit)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1b47868-6295-4f8d-8be4-e77f5443bc91",
   "metadata": {},
   "source": [
    "In this task, the squared exponential and Matern kernels seem to perform relatively comparably, with the larger differences being due to the nugget choice.\n",
    "\n",
    "As expected, the Matern kernel does better capture sharp changes in uncertainty in the predictions where there are gaps in the training points, but these uncertainties do generally seem higher than necessary.\n",
    "\n",
    "Optimising the nugget using the `fit` method yields uncertainties comparable with the variation from the model's overall trend, so this approach seems to be working well."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e80f492-c4da-488b-8fcb-8ce9c507a1a1",
   "metadata": {},
   "source": [
    "### Investigate the nugget parameter (adaptive or fit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7af855d8-87ff-4725-a727-d426be28645e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Traffic speed emulator with squared exponential kernel and adaptive nugget\")\n",
    "for bus, em in zip(ts_bus_ids, ts_gp_se.emulators):\n",
    "    print(\"Bus {}\\t{}\\t({})\".format(bus, em.nugget, em.nugget_type))\n",
    "\n",
    "print(\"\\nTraffic speed emulator with squared exponential kernel and fit nugget\")\n",
    "for bus, em in zip(ts_bus_ids, ts_gp_se_nfit.emulators):\n",
    "    print(\"Bus {}\\t{}\\t({})\".format(bus, em.nugget, em.nugget_type))\n",
    "    \n",
    "print(\"\\nTraffic speed emulator with Matern 5-2 kernel and adaptive nugget\")\n",
    "for bus, em in zip(ts_bus_ids, ts_gp_m52.emulators):\n",
    "    print(\"Bus {}\\t{}\\t({})\".format(bus, em.nugget, em.nugget_type))\n",
    "    \n",
    "print(\"\\nTraffic speed emulator with Matern 5-2 kernel and fit nugget\")\n",
    "for bus, em in zip(ts_bus_ids, ts_gp_m52_nfit.emulators):\n",
    "    print(\"Bus {}\\t{}\\t({})\".format(bus, em.nugget, em.nugget_type))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4994e80-f18b-4dcd-a50d-a859e14ca56a",
   "metadata": {},
   "source": [
    "## Accuracy of the uncertainty estimate\n",
    "\n",
    "With the `include_nugget` option turned off, we get much more reasonable results from the GP emulators than if it was set to `True`. (Can test this out above by changing the `nugget_setting` variable.)\n",
    "\n",
    "Now let's get try to understand how well the uncertainty in the model is quantified, by adding further simulation runs to the above plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca53bfcf-420a-446e-a9ed-d6b4a283a7f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "n_validation_repeats = 10\n",
    "\n",
    "further_ts = np.zeros((len(ts_bus_ids), n_validation, n_validation_repeats))\n",
    "for ind, ts in enumerate(validation_traffic_speeds):\n",
    "    for ind_rep in range(n_validation_repeats):\n",
    "        sim_out, _ = run_stochastic_model(TrafficSpeed=ts[0], bus_ids_out=ts_bus_ids)\n",
    "        further_ts[:,ind,ind_rep] = np.array(sim_out)  # dim: bus_ind, validation ts, repeat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9502764b-f346-4e26-8205-7da7914f8939",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function for plotting data used to train a gp, the simulated outputs and emulated outputs\n",
    "def plot_gp_and_sims(inputs, simulated_outputs, predicted_outputs, bus_ids, plt_bus_index=None):\n",
    "    \n",
    "    legend_entries = []\n",
    "    \n",
    "    input_order = np.argsort(inputs, axis=None)\n",
    "\n",
    "    # Easier to look at one bus at a time\n",
    "    if plt_bus_index:\n",
    "        plt.title(\"Bus {}\".format(bus_ids[plt_bus_index]))\n",
    "\n",
    "    for ind, (a, p, u) in enumerate(zip(simulated_outputs, predicted_outputs.mean, predicted_outputs.unc)):\n",
    "\n",
    "        if plt_bus_index is None or ind == plt_bus_index:\n",
    "            plt.plot(inputs[input_order], p[input_order], marker=\"\", linestyle=\"-\", color=plt.cm.Set1(ind))\n",
    "            \n",
    "            plt.fill_between(inputs.flatten()[input_order], p[input_order]-u[input_order], p[input_order]+u[input_order],\n",
    "                             color=plt.cm.Set1(ind), alpha=0.2, linewidth=0)\n",
    "            \n",
    "            for inp, sim in zip(inputs, simulated_outputs[plt_bus_index]):\n",
    "                plt.scatter(inp*np.ones((len(sim),1)), sim, marker=\".\", color=\"k\", s=5)\n",
    "                sim_std = np.std(sim)\n",
    "                sim_mean = np.mean(sim)\n",
    "                plt.scatter([inp, inp], sim_mean+[-2*sim_std, 2*sim_std], marker=\"_\", color=\"k\", s=50)\n",
    "\n",
    "    plt.xlabel(\"Traffic speed\")\n",
    "    plt.ylabel(\"Time at final\\ndestination\")\n",
    "    plt.ylim((0, 8000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66519ffa-7344-4e6d-ad14-581fcb35d62e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a dict of options to make the following plot easier to generate\n",
    "plt_gps = {\"sead\": {\"preds\": predicted_ts_se, \"title\": \"Squared exponential kernel, adaptive nugget\"},\n",
    "           \"sefit\": {\"preds\": predicted_ts_se_nfit, \"title\": \"Squared exponential, fit nugget\"},\n",
    "           \"m52ad\": {\"preds\": predicted_ts_m52, \"title\": \"Matern 5/2 kernel, adaptive nugget\"},\n",
    "           \"m52fit\": {\"preds\": predicted_ts_m52_nfit, \"title\": \"Matern 5/2 kernel, fit nugget\"}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7a0b644-282c-4c1e-ae75-ad8058175bf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_bus_index = 2\n",
    "plt_gp = \"sefit\"    # change to one of the keys in the dict above to switch between different GP emulators\n",
    "\n",
    "plot_gp_and_sims(validation_traffic_speeds, further_ts, plt_gps[plt_gp][\"preds\"], ts_bus_ids, plt_bus_index)\n",
    "plt.title(\"{}, Bus {}\".format(plt_gps[plt_gp][\"title\"], ts_bus_ids[plt_bus_index]));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5369e16e-c04b-42d2-973a-1242160182e6",
   "metadata": {},
   "source": [
    "The plot above shows how the predictions and uncertainties from each GP emulator compare with the outputs of several repeated simulations at the same value of traffic speed. The horizontal bars show the mean plus/minus two standard deviations of the simulated outputs.\n",
    "\n",
    "As we noted from the initial plots, the GP emulators where the nugget was fit show uncertainty bounds far more in keeping with the stochastic variation of the model than the adaptive nuggets do.\n",
    "However, the uncertainty bounds are larger for the final buses (try e.g. `plt_bus_index=5`, `plt_gp=\"sefit\"`).\n",
    "This is potentially due to some of the final buses not arriving at their destination before the end of the simulation, yet the uncertainty remains high even for speeds when all buses do reach that final stop.\n",
    "The nugget values (printed several cells above) are of the order of 10x greater for these bus/emulator combinations, which means we should seek to understand this further before adding to the complexity of the system we seek to emulate."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "923377e1-34ba-446e-b96f-c43059b9ed75",
   "metadata": {},
   "source": [
    "## Potential extensions\n",
    "\n",
    "In this notebook we have investigated an accessible method for implementing Gaussian Process emulation for stochastic systems.\n",
    "The treatment of the stochastic term is relatively simple here, and will not allow for detailed capture of stochasticity that is, for example, dependent on other model parameters.\n",
    "\n",
    "At the time of writing, software packages for GP emulation were lacking this feature, however, several approaches have been proposed.\n",
    "A natural extension of the work above would be to implement one or more of these type of approaches in the `mogp_emulator` package, and see how the GP emulator's performance is affected. \n",
    "A good starting point would be _Managing Uncertainty in Complex Stochastic Models: Design and Emulation of a Rabies Model_, Boukouvalas _et al._ (2009, preprint available [here](https://www.mucm.ac.uk/Pages/Downloads/Technical%20Reports/08-06.pdf)) which builds upon prior work ([Kersting _et al._, 2007](https://people.csail.mit.edu/kersting/papers/kersting07icml_mlHetGP.pdf)) to emulate the stochastic term using an additional GP emulator."
   ]
  }
 ],
 "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": 5
}