season-2/cadcad-hack-18-numba.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Table of Contents\n",
    "\n",
    "<ul>\n",
    "    <li><a href='#Dependencies'>Dependencies</a></li>\n",
    "    </li>\n",
    "    <li><a href='#Modelling'>Modelling</a>\n",
    "        <ol style='margin-top: 0em;'>\n",
    "            <li><a href='#1.-State-Variables'>State Variables</a></li>\n",
    "            <li><a href='#2.-System-Parameters'>System Parameters</a></li>\n",
    "            <li><a href='#3.-Policy-Functions'>Policy Functions</a></li>\n",
    "            <li><a href='#4.-State-Update-Functions'>State Update Functions</a></li>\n",
    "            <li><a href='#5.-Partial-State-Update-Blocks'>Partial State Update Blocks</a></li>\n",
    "        </ol>\n",
    "    </li>\n",
    "    <li><a href='#Simulation'>Simulation</a>\n",
    "        <ol style='margin-top: 0em;' start=\"6\">\n",
    "            <li><a href='#6.-Configuration'>Configuration</a></li>\n",
    "            <li><a href='#7.-Execution'>Execution</a></li>\n",
    "            <li><a href='#8.-Output-Preparation'>Output Preparation</a></li>\n",
    "            <li><a href='#9.-Analysis'>Analysis</a></li>\n",
    "        </ol>\n",
    "    </li>\n",
    "</ul>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cadCAD standard dependencies\n",
    "\n",
    "# cadCAD configuration modules\n",
    "from cadCAD.configuration.utils import config_sim\n",
    "from cadCAD.configuration import Experiment\n",
    "\n",
    "# cadCAD simulation engine modules\n",
    "from cadCAD.engine import ExecutionMode, ExecutionContext\n",
    "from cadCAD.engine import Executor\n",
    "\n",
    "# cadCAD global simulation configuration list\n",
    "from cadCAD import configs\n",
    "\n",
    "# Included with cadCAD\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Additional dependencies\n",
    "import numba\n",
    "\n",
    "# For analytics\n",
    "import numpy as np\n",
    "# For visualization\n",
    "import plotly.express as px"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modelling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. State Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_state = {\n",
    "    'predator_population': 15, # Initial number of predators\n",
    "    'prey_population': 100, # Initial number of preys\n",
    "}\n",
    "initial_state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. System Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_params = {\n",
    "    # Parameters describing the interaction between the two populations:\n",
    "    \n",
    "    # A parameter used to calculate the rate of predator birth\n",
    "    \"predator_birth_parameter\": [0.01], \n",
    "    # A parameter used to calculate the rate of predator death\n",
    "    \"predator_death_parameter\": [1.0],\n",
    "    # A parameter used to calculate the rate of prey birth\n",
    "    \"prey_birth_parameter\": [0.6, 1.0],\n",
    "    # A parameter used to calculate the rate of prey death\n",
    "    \"prey_death_parameter\": [0.03],\n",
    "\n",
    "    # Parameters used for random Numpy variable tuning\n",
    "    # These parameters scale the random variable used to create the random birth / death rate\n",
    "    \n",
    "    \"random_predator_birth\": [0.0002],\n",
    "    \"random_predator_death\": [0.005],\n",
    "    \"random_prey_birth\": [0.1],\n",
    "    \"random_prey_death\": [0.1],\n",
    "\n",
    "    # Parameter used as conversion factor between 1 unit of time and 1 timestep\n",
    "    # 10 timesteps == 1 unit of time, i.e. every 10 cadCAD model timesteps, 1 unit of actual model time passes\n",
    "    \n",
    "    \"dt\": [0.1],\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Policy Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the njit decorator for telling numba\n",
    "# that the function will be accelerated.\n",
    "# It is a good practice to type hint\n",
    "# all the arguments and output.\n",
    "@numba.njit\n",
    "def calculate_population(rate: float,\n",
    "                         population_1: float,\n",
    "                         population_2: float,\n",
    "                         dt: float) -> float:\n",
    "    \n",
    "    return rate * population_1 * population_2 * dt\n",
    "\n",
    "\n",
    "def p_predator_births(params, substep, state_history, previous_state):\n",
    "    '''Predator Births Policy Function\n",
    "    The predator birth rate (rate of predators born per unit of time) is a product of\n",
    "    the prey population and the predator birth parameter plus a random variable.\n",
    "    \n",
    "    i.e. the larger the prey population, the higher the predator birth rate\n",
    "    '''\n",
    "    # Parameters   \n",
    "    dt = params['dt']\n",
    "    predator_birth_parameter = params['predator_birth_parameter']\n",
    "    random_predator_birth = params['random_predator_birth']\n",
    "\n",
    "    # State Variables\n",
    "    predator_population = previous_state['predator_population']\n",
    "    prey_population = previous_state['prey_population']\n",
    "\n",
    "    # Calculate the predator birth rate\n",
    "    birth_rate = prey_population * (predator_birth_parameter + np.random.random() * random_predator_birth)\n",
    "    # Calculate change in predator population\n",
    "    births = calculate_population(\n",
    "        birth_rate,\n",
    "        predator_population,\n",
    "        1,\n",
    "        dt\n",
    "    )\n",
    "\n",
    "    return {'add_to_predator_population': births}\n",
    "\n",
    "\n",
    "def p_predator_deaths(params, substep, state_history, previous_state):\n",
    "    '''Predator Deaths Policy Function\n",
    "    The predator death rate (rate of predators that die per unit of time) is a function of\n",
    "    the predator death parameter plus a random variable.\n",
    "    \n",
    "    i.e. the larger the predator death parameter, the higher the predator death rate\n",
    "    '''\n",
    "    # Parameters\n",
    "    dt = params['dt']\n",
    "    predator_death_parameter = params['predator_death_parameter']\n",
    "    random_predator_death = params['random_predator_death']\n",
    "\n",
    "    # State Variables\n",
    "    predator_population = previous_state['predator_population']\n",
    "\n",
    "    # Calculate the predator death rate\n",
    "    death_rate = predator_death_parameter + np.random.random() * random_predator_death\n",
    "    # Calculate change in predator population\n",
    "    deaths = calculate_population(\n",
    "        death_rate,\n",
    "        predator_population,\n",
    "        1,\n",
    "        dt\n",
    "    )\n",
    "\n",
    "    return {'add_to_predator_population': -1.0 * deaths}\n",
    "\n",
    "\n",
    "def p_prey_births(params, substep, state_history, previous_state):\n",
    "    '''Prey Births Policy Function\n",
    "    The prey birth rate (rate of preys born per unit of time) is a function of\n",
    "    the prey birth parameter plus a random variable.\n",
    "    \n",
    "    i.e. the larger the prey birth parameter, the higher the prey birth rate\n",
    "    '''\n",
    "    # Parameters\n",
    "    dt = params['dt']\n",
    "    prey_birth_parameter = params['prey_birth_parameter']\n",
    "    random_prey_birth = params['random_prey_birth']\n",
    "\n",
    "    # State Variables\n",
    "    prey_population = previous_state['prey_population']\n",
    "\n",
    "    # Calculate the prey birth rate\n",
    "    birth_rate = prey_birth_parameter + np.random.random() * random_prey_birth\n",
    "    # Calculate change in prey population\n",
    "    births = calculate_population(\n",
    "        birth_rate,\n",
    "        1,\n",
    "        prey_population,\n",
    "        dt\n",
    "    )\n",
    "\n",
    "    return {'add_to_prey_population': births}\n",
    "\n",
    "\n",
    "def p_prey_deaths(params, substep, state_history, previous_state):\n",
    "    '''Prey Deaths Policy Function\n",
    "    The prey death rate (rate of preys that die per unit of time) is a product of\n",
    "    the predator population and the prey death parameter plus a random variable.\n",
    "    \n",
    "    i.e. the larger the predator population, the higher the prey death rate\n",
    "    '''\n",
    "    # Parameters\n",
    "    dt = params['dt']\n",
    "    prey_death_parameter = params['prey_death_parameter']\n",
    "    random_prey_death = params['random_prey_death']\n",
    "\n",
    "    # State Variables\n",
    "    prey_population = previous_state['prey_population']\n",
    "    predator_population = previous_state['predator_population']\n",
    "\n",
    "    # Calculate the prey death rate\n",
    "    death_rate = predator_population * (prey_death_parameter + np.random.random() * random_prey_death)\n",
    "    # Calculate change in prey population\n",
    "    deaths = calculate_population(\n",
    "        death_rate,\n",
    "        1,\n",
    "        prey_population,\n",
    "        dt\n",
    "    )\n",
    "\n",
    "    return {'add_to_prey_population': -1.0 * deaths}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. State Update Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def s_predator_population(params, substep, state_history, previous_state, policy_input):\n",
    "    '''Predator Population State Update Function\n",
    "    Take the Policy Input `add_to_predator_population`\n",
    "    (the net predator births and deaths)\n",
    "    and add to the `predator_population` State Variable.\n",
    "    '''\n",
    "    # Policy Inputs\n",
    "    add_to_predator_population = policy_input['add_to_predator_population']\n",
    "\n",
    "    # State Variables\n",
    "    predator_population = previous_state['predator_population']\n",
    "\n",
    "    # Calculate updated predator population\n",
    "    updated_predator_population = predator_population + add_to_predator_population\n",
    "\n",
    "    return 'predator_population', updated_predator_population\n",
    "\n",
    "\n",
    "def s_prey_population(params, substep, state_history, previous_state, policy_input):\n",
    "    '''Prey Population State Update Function\n",
    "    Take the Policy Input `add_to_prey_population`\n",
    "    (the net prey births and deaths)\n",
    "    and add to the `prey_population` State Variable.\n",
    "    '''\n",
    "    # Policy Inputs\n",
    "    add_to_prey_population = policy_input['add_to_prey_population']\n",
    "\n",
    "    # State Variables\n",
    "    prey_population = previous_state['prey_population']\n",
    "\n",
    "    # Calculate updated prey population\n",
    "    updated_prey_population = prey_population + add_to_prey_population\n",
    "\n",
    "    return 'prey_population', updated_prey_population"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Partial State Update Blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "partial_state_update_blocks = [\n",
    "    {   \n",
    "        # Configure the model Policy Functions\n",
    "        'policies': {\n",
    "            # Calculate the predator birth rate and number of births\n",
    "            'predator_births': p_predator_births,\n",
    "            # Calculate the predator death rate and number of deaths\n",
    "            'predator_deaths': p_predator_deaths,\n",
    "            # Calculate the prey birth rate and number of births\n",
    "            'prey_births': p_prey_births,\n",
    "            # Calculate the prey death rate and number of deaths\n",
    "            'prey_deaths': p_prey_deaths,\n",
    "        },\n",
    "        # Configure the model State Update Functions\n",
    "        'variables': {\n",
    "            # Update the predator population\n",
    "            'predator_population': s_prey_population,\n",
    "            # Update the prey population\n",
    "            'prey_population': s_predator_population\n",
    "        }\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_config = config_sim({\n",
    "    \"N\": 20, # the number of times we'll run the simulation (\"Monte Carlo runs\")\n",
    "    \"T\": range(400), # the number of timesteps the simulation will run for\n",
    "    \"M\": system_params # the parameters of the system\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del configs[:] # Clear any prior configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment = Experiment()\n",
    "experiment.append_configs(\n",
    "    initial_state = initial_state,\n",
    "    partial_state_update_blocks = partial_state_update_blocks,\n",
    "    sim_configs = sim_config\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exec_context = ExecutionContext()\n",
    "simulation = Executor(exec_context=exec_context, configs=configs)\n",
    "raw_result, tensor_field, sessions = simulation.execute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Output Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert raw results to a Pandas DataFrame\n",
    "df = pd.DataFrame(raw_result)\n",
    "\n",
    "# Insert cadCAD parameters for each configuration into DataFrame\n",
    "for config in configs:\n",
    "    # Get parameters from configuration\n",
    "    parameters = config.sim_config['M']\n",
    "    # Get subset index from configuration\n",
    "    subset_index = config.subset_id\n",
    "    \n",
    "    # For each parameter key value pair\n",
    "    for (key, value) in parameters.items():\n",
    "        # Select all DataFrame indices where subset == subset_index\n",
    "        dataframe_indices = df.eval(f'subset == {subset_index}')\n",
    "        # Assign each parameter key value pair to the DataFrame for the corresponding subset\n",
    "        df.loc[dataframe_indices, key] = value\n",
    "\n",
    "df.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.options.plotting.backend = \"plotly\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize how the predator and prey populations change over time\n",
    "\n",
    "# Notice that the populations are more chaotic when the prey birth rate is higher,\n",
    "# and the system is more stable when it is lower.\n",
    "\n",
    "px.line(\n",
    "    df,\n",
    "    x='timestep', # Variable on the horizontal axis\n",
    "    y=['predator_population', 'prey_population'], # Variables on the vertical axis\n",
    "    line_group='run', # One line for each MC run\n",
    "    facet_row='prey_birth_parameter', # Create a figure for each `prey_birth_parameter` parameter sweep\n",
    "    log_y=True, # Use log scale on the vertical axis\n",
    "    height=800,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize how the predator and prey populations compare\n",
    "\n",
    "# Notice that they tend to have cycles around a fixed point\n",
    "# which means that their populations are never static,\n",
    "# but rather cyclic with ups and downs.\n",
    "\n",
    "px.line(\n",
    "    df,\n",
    "    x='predator_population', # Variable on the horizontal axis\n",
    "    y='prey_population', # Variable on the vertical axis\n",
    "    color='run', # Color by MC run\n",
    "    facet_row='prey_birth_parameter', # Create a figure for each `prey_birth_parameter` parameter sweep\n",
    "    log_x=True, # Use log scale on the horizontal axis\n",
    "    log_y=True, # Use log scale on the vertical axis\n",
    "    height=800,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:percent"
  },
  "kernelspec": {
   "display_name": "Python (cadCAD Hacks)",
   "language": "python",
   "name": "python-cadcad-hacks"
  },
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}