{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modeling and Simulation in Python\n",
    "\n",
    "Chapter 6\n",
    "\n",
    "Copyright 2017 Allen Downey\n",
    "\n",
    "License: [Creative Commons Attribution 4.0 International](https://creativecommons.org/licenses/by/4.0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure Jupyter so figures appear in the notebook\n",
    "%matplotlib inline\n",
    "\n",
    "# Configure Jupyter to display the assigned value after an assignment\n",
    "%config InteractiveShell.ast_node_interactivity='last_expr_or_assign'\n",
    "\n",
    "# import functions from the modsim.py module\n",
    "from modsim import *\n",
    "\n",
    "from pandas import read_html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Code from the previous chapter\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = 'data/World_population_estimates.html'\n",
    "tables = read_html(filename, header=0, index_col=0, decimal='M')\n",
    "table2 = tables[2]\n",
    "table2.columns = ['census', 'prb', 'un', 'maddison', \n",
    "                  'hyde', 'tanton', 'biraben', 'mj', \n",
    "                  'thomlinson', 'durand', 'clark']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "un = table2.un / 1e9\n",
    "un.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "census = table2.census / 1e9\n",
    "census.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_0 = get_first_label(census)\n",
    "t_end = get_last_label(census)\n",
    "elapsed_time = t_end - t_0\n",
    "\n",
    "p_0 = get_first_value(census)\n",
    "p_end = get_last_value(census)\n",
    "total_growth = p_end - p_0\n",
    "\n",
    "annual_growth = total_growth / elapsed_time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### System objects"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can rewrite the code from the previous chapter using system objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "system = System(t_0=t_0, \n",
    "                t_end=t_end,\n",
    "                p_0=p_0,\n",
    "                annual_growth=annual_growth)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can encapsulate the code that runs the model in a function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_simulation1(system):\n",
    "    \"\"\"Runs the constant growth model.\n",
    "    \n",
    "    system: System object\n",
    "    \n",
    "    returns: TimeSeries\n",
    "    \"\"\"\n",
    "    results = TimeSeries()\n",
    "    results[system.t_0] = system.p_0\n",
    "    \n",
    "    for t in linrange(system.t_0, system.t_end):\n",
    "        results[t+1] = results[t] + system.annual_growth\n",
    "    \n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also encapsulate the code that plots the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_results(census, un, timeseries, title):\n",
    "    \"\"\"Plot the estimates and the model.\n",
    "    \n",
    "    census: TimeSeries of population estimates\n",
    "    un: TimeSeries of population estimates\n",
    "    timeseries: TimeSeries of simulation results\n",
    "    title: string\n",
    "    \"\"\"\n",
    "    plot(census, ':', label='US Census')\n",
    "    plot(un, '--', label='UN DESA')\n",
    "    plot(timeseries, color='gray', label='model')\n",
    "    \n",
    "    decorate(xlabel='Year', \n",
    "             ylabel='World population (billion)',\n",
    "             title=title)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's how we run it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = run_simulation1(system)\n",
    "plot_results(census, un, results, 'Constant growth model')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Proportional growth"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a more realistic model where the number of births and deaths is proportional to the current population."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_simulation2(system):\n",
    "    \"\"\"Run a model with proportional birth and death.\n",
    "    \n",
    "    system: System object\n",
    "    \n",
    "    returns: TimeSeries\n",
    "    \"\"\"\n",
    "    results = TimeSeries()\n",
    "    results[system.t_0] = system.p_0\n",
    "    \n",
    "    for t in linrange(system.t_0, system.t_end):\n",
    "        births = system.birth_rate * results[t]\n",
    "        deaths = system.death_rate * results[t]\n",
    "        results[t+1] = results[t] + births - deaths\n",
    "        \n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I picked a death rate that seemed reasonable and then adjusted the birth rate to fit the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "system.death_rate = 0.01\n",
    "system.birth_rate = 0.027"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what it looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = run_simulation2(system)\n",
    "plot_results(census, un, results, 'Proportional model')\n",
    "savefig('figs/chap06-fig01.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model fits the data pretty well for the first 20 years, but not so well after that."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Factoring out the update function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`run_simulation1` and `run_simulation2` are nearly identical except the body of the loop.  So we can factor that part out into a function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_func1(pop, t, system):\n",
    "    \"\"\"Compute the population next year.\n",
    "    \n",
    "    pop: current population\n",
    "    t: current year\n",
    "    system: system object containing parameters of the model\n",
    "    \n",
    "    returns: population next year\n",
    "    \"\"\"\n",
    "    births = system.birth_rate * pop\n",
    "    deaths = system.death_rate * pop\n",
    "    return pop + births - deaths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The name `update_func` refers to a function object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "update_func1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which we can confirm by checking its type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "type(update_func1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`run_simulation` takes the update function as a parameter and calls it just like any other function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_simulation(system, update_func):\n",
    "    \"\"\"Simulate the system using any update function.\n",
    "    \n",
    "    system: System object\n",
    "    update_func: function that computes the population next year\n",
    "    \n",
    "    returns: TimeSeries\n",
    "    \"\"\"\n",
    "    results = TimeSeries()\n",
    "    results[system.t_0] = system.p_0\n",
    "    \n",
    "    for t in linrange(system.t_0, system.t_end):\n",
    "        results[t+1] = update_func(results[t], t, system)\n",
    "        \n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's how we use it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_0 = get_first_label(census)\n",
    "t_end = get_last_label(census)\n",
    "p_0 = census[t_0]\n",
    "\n",
    "system = System(t_0=t_0, \n",
    "                t_end=t_end,\n",
    "                p_0=p_0,\n",
    "                birth_rate=0.027,\n",
    "                death_rate=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = run_simulation(system, update_func1)\n",
    "plot_results(census, un, results, 'Proportional model, factored')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remember not to put parentheses after `update_func1`.  What happens if you try?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** When you run `run_simulation`, it runs `update_func1` once for each year between `t_0` and `t_end`.  To see that for yourself, add a print statement at the beginning of `update_func1` that prints the values of `t` and `pop`, then run `run_simulation` again."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combining birth and death"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since births and deaths get added up, we don't have to compute them separately.  We can combine the birth and death rates into a single net growth rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_func2(pop, t, system):\n",
    "    \"\"\"Compute the population next year.\n",
    "    \n",
    "    pop: current population\n",
    "    t: current year\n",
    "    system: system object containing parameters of the model\n",
    "    \n",
    "    returns: population next year\n",
    "    \"\"\"\n",
    "    net_growth = system.alpha  * pop\n",
    "    return pop + net_growth"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's how it works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "system.alpha = system.birth_rate - system.death_rate\n",
    "\n",
    "results = run_simulation(system, update_func2)\n",
    "plot_results(census, un, results, 'Proportional model, combined birth and death')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercises\n",
    "\n",
    "**Exercise:** Maybe the reason the proportional model doesn't work very well is that the growth rate, `alpha`, is changing over time.  So let's try a model with different growth rates before and after 1980 (as an arbitrary choice).\n",
    "\n",
    "Write an update function that takes `pop`, `t`, and `system` as parameters.  The system object, `system`, should contain two parameters: the growth rate before 1980, `alpha1`, and the growth rate after 1980, `alpha2`.  It should use `t` to determine which growth rate to use.  Note: Don't forget the `return` statement.\n",
    "\n",
    "Test your function by calling it directly, then pass it to `run_simulation`.  Plot the results.  Adjust the parameters `alpha1` and `alpha2` to fit the data as well as you can.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}