docs/tutorials/intro_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4d52735e-0198-403b-bd9d-56cb12d4dde2",
   "metadata": {},
   "source": [
    "# Mesa-Geo Introductory Model \n",
    "\n",
    "To overview the critical parts of Mesa-Geo this tutorial uses the pandemic modelling approach known as a S(usceptible), I(infected) and R(ecovered) or SIR model. \n",
    "\n",
    "Components of the model are:\n",
    "\n",
    "**Agents:** Each agent in the model represents an individual in the population. Agents have states of susceptible, infected, recovered, or dead. The Agents are point agents, randomly placed into the environment.\n",
    "\n",
    "**Environment:** The environment is a set of polygons of a few Toronto neighborhoods.   \n",
    "\n",
    "**Interaction Rules:** Susceptible agents can become infected with a certain probability, if they come into contact with infected agents. Infected agents then recover after a certain period or perish based on a probability.\n",
    "\n",
    "**Parameters:** \n",
    "- Population Size (number of human agents in the model)\n",
    "- Initial Infection (percent of the population initial infected)\n",
    "- Exposure Distance (proximity suscpetible agents must be to infected agents to possibly get infected)\n",
    "- Infection Risk (probability of becoming infected)\n",
    "- Recovery Rate (time infection lasts)\n",
    "- Mobility (distance agent moves)\n",
    "\n",
    "**The tutorial then proceeds in three parts:** \n",
    "- Part 1 Create the Basic Model \n",
    "- Part 2 Add Agent Behaviors and Model Complexity \n",
    "- Part 3 Add Visualizations and Interface\n",
    "\n",
    " (You can use the table of contents button on the left side of the interface to skip to any specific part)\n",
    "\n",
    "\n",
    "Users can use Google Colab\\* (Please ensure you run the Colab dependency import cell- below) \n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/projectmesa/mesa-geo/blob/main/docs/tutorials/intro_tutorial.ipynb)\n",
    "\n",
    "\\*Based on a recent Google Colab update currently, Solara visuals are not rendering. However, the link still provides a easy way to download the jupyter file. You can see the issue on the [Solara GitHub site](https://github.com/widgetti/solara/issues/323)\n",
    "\n",
    "You can also [download the file directly from GitHub](https://github.com/projectmesa/mesa-geo/blob/main/docs/tutorials/intro_tutorial.ipynb)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b0dc4bb-cf8a-4c21-82aa-4a1cc43f1871",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this if in colab or if you need to install mesa and mesa-geo in your local environment.\n",
    "!pip install -U --pre mesa-geo --quiet\n",
    "!mkdir -p data\n",
    "!wget -P data https://raw.githubusercontent.com/projectmesa/mesa-geo/main/docs/tutorials/data/TorontoNeighbourhoods.geojson"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f1b8324-9849-4992-b13a-db9a94372d0a",
   "metadata": {},
   "source": [
    "## Part 1  Create the Basic Model\n",
    "This portion initializes the human agents, the neighborhood agents, and the model class that manages the model dynamics. \n",
    "\n",
    "First we import our dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "431d4b10-a0ee-4c03-9197-1c1266165169",
   "metadata": {
    "explanatory": true,
    "has_explanation": false
   },
   "source": [
    "\n",
    "\n",
    "This cell imports the specific libraries we need to create our model. \n",
    "\n",
    "- [Shapley](https://shapely.readthedocs.io/en/stable/) a library GIS library for object in the cartesian plane. From Shapely we specifically need the Point class to create our human agents\n",
    "- [Mesa](https://mesa.readthedocs.io/en/stable/) the parent ABM library to Mesa-Geo\n",
    "\n",
    "Then of course mesa-geo which although not strictly necessary we also specifically import the visualization part of the library \n",
    "so we do not have to write out mesa-geo.visualization modules when we call them. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7d43c11-a846-45e1-bbcf-77de58ff5033",
   "metadata": {
    "has_explanation": true
   },
   "outputs": [],
   "source": [
    "from shapely.geometry import Point\n",
    "\n",
    "import mesa\n",
    "from mesa.visualization import SolaraViz, make_plot_component\n",
    "import mesa_geo as mg\n",
    "from mesa_geo.visualization import make_geospace_component"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "040912e4-e07b-4ffa-97fa-5aeae014c844",
   "metadata": {},
   "source": [
    "### Create the person agent class\n",
    "\n",
    "The person in this model represents one human being and we initilaize each person agent with two key parts: \n",
    "1. The agent attributes, such as recovery rate and death risk\n",
    "2. The step function, actions the agent will take each model step "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d0bd901-c038-40d5-8017-df3fdedbd838",
   "metadata": {
    "explanatory": true,
    "has_explanation": false
   },
   "source": [
    "The first thing we are going to do is create the person agent class. This class has several attributes necessary to make a more holistic model. \n",
    "\n",
    "First, there are the required attributes for any GeoAgent in Mesa-Geo\n",
    "\n",
    "- **model**: Model object class that we will build later, this is a pointer to the model instance so the agent can get information from the model as it behaves\n",
    "- **geometry**: GIS geometric object in this case a GIS Point\n",
    "- **crs**: A string describing the coordinate reference system the agent is using\n",
    "\n",
    "As you can see these are inherited from the mesa-geo librarary through the \"mg.GeoAgent\" in the class instantiation. \n",
    "\n",
    "Second, the variable attributes these are unique to our SIR model:\n",
    "\n",
    "- **agent_type**: A string which describes the agent state (susceptible, infected, recovered, or dead) \n",
    "- **mobility_range**: Distance the agent can move in meters\n",
    "- **infection risk**: A float from 0.0 to 1.0 that determines the risk of the agent being infected if exposed. \n",
    "- **recovery_rate**: A float from 0.0 to 1.0 that determine how long the agents takes to recover\n",
    "- **death_risk**: A float from 0.0 to 1.0 that determines the probability the agent will die\n",
    "\n",
    "The **`__repr__`** function is a Python primitive that will print out information as directed by the code. In this case we will print out the agent ID\n",
    "\n",
    "The **step** function is a Mesa primitive that describes what action the agent takes each step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b39b18eb-16e9-4db0-8739-a050cc74d40a",
   "metadata": {
    "has_explanation": true
   },
   "outputs": [],
   "source": [
    "class PersonAgent(mg.GeoAgent):\n",
    "    \"\"\"Person Agent.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        model,\n",
    "        geometry,\n",
    "        crs,\n",
    "        agent_type,\n",
    "        mobility_range,\n",
    "        infection_risk,\n",
    "        recovery_rate,\n",
    "        death_risk,\n",
    "    ):\n",
    "        super().__init__(model, geometry, crs)\n",
    "        # Agent attributes\n",
    "        self.atype = agent_type\n",
    "        self.mobility_range = mobility_range\n",
    "        self.infection_risk = infection_risk\n",
    "        self.recovery_rate = recovery_rate\n",
    "        self.death_risk = death_risk\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f\"Person {self.unique_id}\"\n",
    "\n",
    "    def step(self):\n",
    "        print(repr(self))\n",
    "        print(self.atype, self.death_risk, self.recovery_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce98080f-209d-4476-b50c-88a232ff0bbb",
   "metadata": {},
   "source": [
    "### Create the neighborhood agent \n",
    "\n",
    "The neighborhood in this model represents one geographical area as defined by the geojson file we uploaded. \n",
    "\n",
    "Similar to the person agent, we initialize each neighborhood agent with the same two key parts. \n",
    "\n",
    "1. The agent attributes, such as geometry and state of neighborhood\n",
    "2. The step function, behaviors the agent will take during each model step. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b44f5943-ae82-4329-9255-9ec8458311aa",
   "metadata": {
    "explanatory": true
   },
   "source": [
    "\n",
    "Similar to the person agent for the neighborhood agent there are two types of attributes. \n",
    "\n",
    "The required attributes for any GeoAgent in Mesa-Geo:\n",
    "\n",
    "- **model**: Model object class that we will build later, this is a pointer to the model instance so the agent can get information from the model as it behaves\n",
    "- **geometry**: GIS geometric object in this case a polygon form the geojson defining the perimeter of the neighborhood\n",
    "- **crs**: A string describing the coordinate reference system the agent is using\n",
    "\n",
    "Similar to the person agent, \"mg.GeoAgent\" is inherited from mesa-geo. \n",
    "\n",
    "Next are the variable attributes:\n",
    "\n",
    "- **agent_type**: A string which describes the state of the neighborhood which will be either safe or hot spot  \n",
    "- **hotspot_threshold**: An integer that is the number of infected people in a neighborhood to call it a hotspot \n",
    "\n",
    "We will also use the **`__repr__`** function to print out the agent ID\n",
    "\n",
    "Then the **step** function, which is a primitive that describes what action the agent takes each step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da9c0121-9f35-4c7b-ab6a-d8d970189ec6",
   "metadata": {
    "has_explanation": true
   },
   "outputs": [],
   "source": [
    "class NeighbourhoodAgent(mg.GeoAgent):\n",
    "    \"\"\"Neighbourhood agent. Changes color according to number of infected inside it.\"\"\"\n",
    "\n",
    "    def __init__(self, model, geometry, crs, agent_type=\"safe\", hotspot_threshold=1):\n",
    "        super().__init__(model, geometry, crs)\n",
    "        self.atype = agent_type\n",
    "        self.hotspot_threshold = (\n",
    "            hotspot_threshold  # When a neighborhood is considered a hot-spot\n",
    "        )\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f\"Neighbourhood {self.HOODNUM}\"\n",
    "\n",
    "    def step(self):\n",
    "        \"\"\"Advance agent one step.\"\"\"\n",
    "        print(repr(self))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efe370a7-a27e-4c76-830d-155a8853d81c",
   "metadata": {},
   "source": [
    "### Create the Model Class\n",
    "\n",
    "The model class is the manager that instantiates the agents, then manages what is happening in the model through the step function, and collects data. \n",
    "\n",
    "We will create the model with parameters that will set the attributes of the agents as it instantiates them and a step function to call the agent step function. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c44bdd40-6ee9-4579-90fa-502ba482f157",
   "metadata": {
    "explanatory": true
   },
   "source": [
    "\n",
    "First, we name our class in this case GeoSIR and we inherit the model class from Mesa. We store the path to our GeoJSON file in the object geojson regions. As JSONs mirror Pythons dictionary structure, we store the key for the neighbourhood id (\"HOODNUM\") as attribute, and update the __repr__ function to print out the neighbourhood id.\n",
    "\n",
    "Second, we set up the python initializer to initiate our model class. To do this we will, set up key word arguments or kwargs of the parameters we want for our model. In this case we will use: \n",
    "- population size (pop_size): An integer that determines the number of person agents\n",
    "- initial infection (init_infection): A float between 0.0 and 1.0 which determines what percentage of the population is infected as the model initiates\n",
    "- exposure_distance (exposure_dist): An integer for the distance in meters a susceptible person agent must within to be infected by a person agent who is infected\n",
    "- maximum infection risk (max_infection_risk): A float between 0.0 and 1.0 of which determines the highest suscpetibility rate in the population\n",
    "\n",
    "Third, we initialize our agents. Mesa-Geo has an AgentCreator class inside is geoagent.py file that can create GeoAgents from files, GeoDataFrames, GeoJSON or Shapely objects. \n",
    "\n",
    "**Creating the NeighbourhoodAgents**\n",
    "\n",
    "In this case we will use the `torontoneighbourhoods.geojson` file located in the data folder to to create the NeighbourhoodAgents. Next, we will add them to the environment with the space.add_agents function.\n",
    "\n",
    "\n",
    "**Creating the PersonAgents**\n",
    "\n",
    "We will use Mesa-Geo AgentCreator to create the person agents. To create a heterogeneous (diverse) population we will use the [random object created as part of Mesa's base class](https://github.com/projectmesa/mesa/blob/01477fc9624b70078fe1c82634d7c9e4938de942/mesa/model.py#L60) to help initialize the population's parameters. \n",
    "\n",
    "- death_risk: A float from 0 to 1\n",
    "- agent_type: Compares the model parameter of initial infection of a random float between 0 and 1 and the initial infection parameter. If it is less than the initial infection parameter the agent is initialized as infected.\n",
    "- recover: Is an integer between 1 and the recovery rate. This determines the number of steps it takes for the agent to recover.\n",
    "- infection_risk: is a float between 0 and the parameter of max_infection_risk, which will then determine how likely a person is to get infected.\n",
    "- death_risk: Is a random float between 0 and 1 that will determine how likely a person is to die when infected.\n",
    "\n",
    "By using Python's random library to create these attributes for each agent, we can now create a diverse agent population. \n",
    "\n",
    "Passing these parameters through the `AgentCreator` class we initialize our agent object. \n",
    "\n",
    "As Mesa-Geo is an GIS based ABM, we need assign each PersonAgent a Geometry and location. To do this we will use a helper function `find_home`. This helper function first identifies a NeighbourhoodAgent where the PersonAgent will start. Next it identifies the center of the neighborhood and its boundary and then randomly moving from the center point, put staying within the bounds, it a lat and long to aissgn the PersonAgent is starting location.\n",
    "\n",
    "**Step Function**\n",
    "\n",
    "The final piece is to initialize a step function. This function is a Mesa primitive that iterates through each agent calling their step function.\n",
    "\n",
    "**The Model**\n",
    "\n",
    "We know have the pieces of our Model. A GIS layer of polygons that creates NeighbourhoodAgents from our GeoJSON file. A diverse population of GIS Point objects, with different infection, recovery and death risks. A model class that initializes these agents, a GIS space and step function to execute the simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d80d198b-6fa3-4be3-ad05-213d70e7a093",
   "metadata": {
    "has_explanation": true
   },
   "outputs": [],
   "source": [
    "class GeoSIR(mesa.Model):\n",
    "    \"\"\"Model class for a simplistic infection model.\"\"\"\n",
    "\n",
    "    # Geographical parameters for desired map\n",
    "    geojson_regions = \"data/TorontoNeighbourhoods.geojson\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        pop_size=30,\n",
    "        mobility_range=500,\n",
    "        init_infection=0.2,\n",
    "        exposure_dist=500,\n",
    "        max_infection_risk=0.2,\n",
    "        max_recovery_time=5,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.space = mg.GeoSpace(warn_crs_conversion=False)\n",
    "\n",
    "        # SIR model parameters\n",
    "        self.pop_size = pop_size\n",
    "        self.mobility_range = mobility_range\n",
    "        self.initial_infection = init_infection\n",
    "        self.exposure_distance = exposure_dist\n",
    "        self.infection_risk = max_infection_risk\n",
    "        self.recovery_rate = max_recovery_time\n",
    "\n",
    "        # Set up the Neighbourhood patches for every region in file\n",
    "        ac = mg.AgentCreator(NeighbourhoodAgent, model=self)\n",
    "        neighbourhood_agents = ac.from_file(self.geojson_regions)\n",
    "\n",
    "        # Add neighbourhood agents to space\n",
    "        self.space.add_agents(neighbourhood_agents)\n",
    "\n",
    "        # Generate random location, add agent to grid\n",
    "        for i in range(pop_size):\n",
    "            # assess if they are infected\n",
    "            if self.random.random() < self.initial_infection:\n",
    "                agent_type = \"infected\"\n",
    "            else:\n",
    "                agent_type = \"susceptible\"\n",
    "            # determine movement range\n",
    "            mobility_range = self.random.randint(0, self.mobility_range)\n",
    "            # determine agent recovery rate\n",
    "            recover = self.random.randint(1, self.recovery_rate)\n",
    "            # determine agents infection risk\n",
    "            infection_risk = self.random.uniform(0, self.infection_risk)\n",
    "            # determine agent death probability\n",
    "            death_risk = self.random.random()\n",
    "\n",
    "            # Generate PersonAgent population\n",
    "            unique_person = mg.AgentCreator(\n",
    "                PersonAgent,\n",
    "                model=self,\n",
    "                crs=self.space.crs,\n",
    "                agent_kwargs={\n",
    "                    \"agent_type\": agent_type,\n",
    "                    \"mobility_range\": mobility_range,\n",
    "                    \"recovery_rate\": recover,\n",
    "                    \"infection_risk\": infection_risk,\n",
    "                    \"death_risk\": death_risk,\n",
    "                },\n",
    "            )\n",
    "\n",
    "            x_home, y_home = self.find_home(neighbourhood_agents)\n",
    "\n",
    "            this_person = unique_person.create_agent(Point(x_home, y_home))\n",
    "            self.space.add_agents(this_person)\n",
    "\n",
    "    def find_home(self, neighbourhood_agents):\n",
    "        \"\"\"Find start location of agent\"\"\"\n",
    "\n",
    "        # identify location\n",
    "        this_neighbourhood = self.random.randint(\n",
    "            0, len(neighbourhood_agents) - 1\n",
    "        )  # Region where agent starts\n",
    "        center_x, center_y = neighbourhood_agents[\n",
    "            this_neighbourhood\n",
    "        ].geometry.centroid.coords.xy\n",
    "        this_bounds = neighbourhood_agents[this_neighbourhood].geometry.bounds\n",
    "        spread_x = int(\n",
    "            this_bounds[2] - this_bounds[0]\n",
    "        )  # Heuristic for agent spread in region\n",
    "        spread_y = int(this_bounds[3] - this_bounds[1])\n",
    "        this_x = center_x[0] + self.random.randint(0, spread_x) - spread_x / 2\n",
    "        this_y = center_y[0] + self.random.randint(0, spread_y) - spread_y / 2\n",
    "\n",
    "        return this_x, this_y\n",
    "\n",
    "    def step(self):\n",
    "        \"\"\"Run one step of the model.\"\"\"\n",
    "        # Activate PersonAgents in random order\n",
    "        self.agents_by_type[PersonAgent].shuffle_do(\"step\")\n",
    "        # For NeighbourhoodAgents the order doesn't matter, since they update independently from each other\n",
    "        self.agents_by_type[NeighbourhoodAgent].do(\"step\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb15ce7d-3537-4e23-84da-dca93ad7c3f9",
   "metadata": {},
   "source": [
    "### Run The Base Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36bdf628-b152-45ec-a568-03e6faa413e9",
   "metadata": {
    "explanatory": true
   },
   "source": [
    "#explanatory\n",
    "\n",
    "This cell is fairly simple\n",
    "\n",
    "1 - We instantiate the SIR model by call the class name \"GeoSIR\" into the object `model`.\n",
    "\n",
    "2 - Then we call the step function to see if it prints out the Agent IDs, infection status, death_risk, and recovery rate as called in the PersonAgent class. \n",
    "\n",
    "You can also see all the person agents are called and then the neighbourhood agents. This will become important later as we want to update the neighbourhood status later based on its PersonAgent status. \n",
    "\n",
    "If you are curious about the numbers for the neighbourhood agents, you can open up the GeoJSON in the data folder and see that each neighborhood gets a `HOODNUM` attribute assigned.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c181dbf7-c175-47c7-8c20-2ec8fa758c60",
   "metadata": {
    "has_explanation": true
   },
   "outputs": [],
   "source": [
    "model = GeoSIR()\n",
    "model.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ceeed343-3cdb-45af-be8e-90816f542e4d",
   "metadata": {},
   "source": [
    "## Part 2 Add Agent and Model Complexity "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfbc2440-a9a3-4e40-903e-c0d9ff98c548",
   "metadata": {},
   "source": [
    "### Increase PersonAgent Complexity\n",
    "\n",
    "In this section we add behaviors to the PersonAgent to build the necessary SIR dynamics. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26967d31-9dab-4a8f-b970-17ab2e1c7990",
   "metadata": {
    "explanatory": true
   },
   "source": [
    "\n",
    "To create the SIR dynamics we need the agents move, determine if they have been exposed and if they have process the probability of them being infected and possibly dying. \n",
    "\n",
    "To do this we will update our step function. The step function logic uses the agent's `atype` to determine what actions to process\n",
    "\n",
    "**Part 1**\n",
    "\n",
    "If the PersonAgent `atype` is susceptible, then we need to identify all PersonAgent's neighbors within the exposure distance. To do this, we will use Mesa-Geo's `get_neighbors_within_distance` function which takes 2 parameters, the agent, and a distance, which in this case is the model parameter for exposure distance in meters. This creates a list of PersonAgents within that distance. \n",
    "\n",
    "The `get_neighbors_within_distance` [function](https://github.com/projectmesa/mesa-geo/blob/56c598486e0f58f3626d9951796998d38bbab0b5/mesa_geo/geospace.py#L197) has two keyword arguments `center` and `relation`. `center` takes `True` or `False` on whether to include the center, it is set to `False` and measures as a buffer around the agent's geometry. If `True` it measures from the Center of the point. `relation` is defaulted to `intersects` but can take any common spatial relationship, such as `contains`, `within`, `touches`, `crosses`\n",
    "\n",
    "The step function then iterates through the list of neighbors to see if any agents are infected. If so it does a probabilistic comparison of a random float compared to the agents infection risk and if `True` the agent becomes infected and the iteration ends.  \n",
    "\n",
    "**Part 2**\n",
    "\n",
    "If the agent `atype` is infected, then the step function does comparisons. First, it sees how many steps the agent has been infected. To track this the PersonAgent got a new attribute counter which is `steps_infected`. If the steps are greater than or equal to their recovery rate, the agent is recovered, if not then the function does a probabilistic comparison with the agents death risk to see if the agent dies. If neither of these things happen the `steps_infected` increases by one.\n",
    "\n",
    "**Part 3**\n",
    "\n",
    "The next part is if the agent `atype` is not dead then the agent moves. For this we randomly get an integer for the x any (lat and long) between their negative `mobility_range` and positive `mobility range`. We pass these two integers into the helper function `move_point` and then update the agents geometry with this new point. \n",
    "\n",
    "Finally, we update the counts of agent types. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f6e8a2c-ec86-4cd8-8feb-86e3da6920cb",
   "metadata": {
    "has_explanation": true
   },
   "outputs": [],
   "source": [
    "class PersonAgent(mg.GeoAgent):\n",
    "    \"\"\"Person Agent.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        model,\n",
    "        geometry,\n",
    "        crs,\n",
    "        agent_type,\n",
    "        mobility_range,\n",
    "        infection_risk,\n",
    "        recovery_rate,\n",
    "        death_risk,\n",
    "    ):\n",
    "        super().__init__(model, geometry, crs)\n",
    "        # Agent attributes\n",
    "        self.atype = agent_type\n",
    "        self.mobility_range = mobility_range\n",
    "        self.infection_risk = (infection_risk,)\n",
    "        self.recovery_rate = recovery_rate\n",
    "        self.death_risk = death_risk\n",
    "        self.steps_infected = 0\n",
    "        self.steps_recovered = 0\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f\"Person {self.unique_id}\"\n",
    "\n",
    "    # Helper function for moving agent\n",
    "    def move_point(self, dx, dy):\n",
    "        \"\"\"\n",
    "        Move a point by creating a new one\n",
    "        :param dx:  Distance to move in x-axis\n",
    "        :param dy:  Distance to move in y-axis\n",
    "        \"\"\"\n",
    "        return Point(self.geometry.x + dx, self.geometry.y + dy)\n",
    "\n",
    "    def step(self):\n",
    "\n",
    "        # Part 1 - find neighbors based on infection distance\n",
    "        if self.atype == \"susceptible\":\n",
    "            neighbors = self.model.space.get_neighbors_within_distance(\n",
    "                self, self.model.exposure_distance\n",
    "            )\n",
    "            for neighbor in neighbors:\n",
    "                if (\n",
    "                    neighbor.atype == \"infected\"\n",
    "                    and self.random.random() < self.model.infection_risk\n",
    "                ):\n",
    "                    self.atype = \"infected\"\n",
    "                    break  # stop process if agent becomes infected\n",
    "\n",
    "        # Part -2 If infected, check if agent recovers or agent dies\n",
    "        elif self.atype == \"infected\":\n",
    "            if self.steps_infected >= self.recovery_rate:\n",
    "                self.atype = \"recovered\"\n",
    "                self.steps_infected = 0\n",
    "            elif self.random.random() < self.death_risk:\n",
    "                self.atype = \"dead\"\n",
    "            else:\n",
    "                self.steps_infected += 1\n",
    "\n",
    "        elif self.atype == \"recovered\":\n",
    "            self.steps_recovered += 1\n",
    "            if self.steps_recovered >= 2:\n",
    "                self.atype = \"susceptible\"\n",
    "                self.steps_recovered = 0\n",
    "\n",
    "        # Part 3 - If not dead, move\n",
    "        if self.atype != \"dead\":\n",
    "            move_x = self.random.randint(-self.mobility_range, self.mobility_range)\n",
    "            move_y = self.random.randint(-self.mobility_range, self.mobility_range)\n",
    "            self.geometry = self.move_point(move_x, move_y)  # Reassign geometry\n",
    "\n",
    "        self.model.counts[self.atype] += 1  # Count agent type"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4823f40-2837-4176-8ee5-e09276888e45",
   "metadata": {},
   "source": [
    "### Increase NeighbourhoodAgent Complexity"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1af08cc-4f1e-45b7-8282-06d30f6dbe3e",
   "metadata": {
    "explanatory": true
   },
   "source": [
    "\n",
    "For the NeighbourhoodAgent we want to change their color based on the number of infected PersonAgents in their neighbourhood. \n",
    "\n",
    "To do this we will create a helper function called `color_hotspot`. We will then use mesa-geo's `get_intersecting_agents` [function](https://github.com/projectmesa/mesa-geo/blob/56c598486e0f58f3626d9951796998d38bbab0b5/mesa_geo/geospace.py#L194). We will then iterate through that list to get the agents with `atype` infected if the list is longer than our `hotspot_threshold` equal to 1 (so if two agents in the neighborhood are infected) then the `atype` will change to `hotspot`. \n",
    "\n",
    "We then update our model counts. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50aab416-959f-4d72-abd9-844ed42c784d",
   "metadata": {
    "has_explanation": true
   },
   "outputs": [],
   "source": [
    "class NeighbourhoodAgent(mg.GeoAgent):\n",
    "    \"\"\"Neighbourhood agent. Changes color according to number of infected inside it.\"\"\"\n",
    "\n",
    "    def __init__(self, model, geometry, crs, agent_type=\"safe\", hotspot_threshold=1):\n",
    "        super().__init__(model, geometry, crs)\n",
    "        self.atype = agent_type\n",
    "        self.hotspot_threshold = (\n",
    "            hotspot_threshold  # When a neighborhood is considered a hot-spot\n",
    "        )\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f\"Neighbourhood {self.unique_id}\"\n",
    "\n",
    "    def color_hotspot(self):\n",
    "        # Decide if this region agent is a hot-spot\n",
    "        # (if more than threshold person agents are infected)\n",
    "        neighbors = self.model.space.get_intersecting_agents(self)\n",
    "        infected_neighbors = [\n",
    "            neighbor for neighbor in neighbors if neighbor.atype == \"infected\"\n",
    "        ]\n",
    "        if len(infected_neighbors) > self.hotspot_threshold:\n",
    "            self.atype = \"hotspot\"\n",
    "        else:\n",
    "            self.atype = \"safe\"\n",
    "\n",
    "    def step(self):\n",
    "        \"\"\"Advance agent one step.\"\"\"\n",
    "        self.color_hotspot()\n",
    "        self.model.counts[self.atype] += 1  # Count agent type"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee5ba9f7-db79-43a9-8a76-37dfbf6b2bd5",
   "metadata": {},
   "source": [
    "### Increase model complexity\n",
    "\n",
    "For this section will add data collection where we collect the status of the PersonAgents and the NeighbourhoodAgents but counting the different `atypes`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ccbc3cb-083d-4fc1-9516-85a71fe36701",
   "metadata": {
    "explanatory": true
   },
   "source": [
    "\n",
    "As we run our SIR model, we want to ensure we are collecting information about the status of the disease. \n",
    "\n",
    "To do this we will create helper functions that get this information. In this case we will put them in a separate cell, but depending on the developers preference they could also put them in the model class or collect the information in a handful of other ways. \n",
    "\n",
    "In this case, we set up an attribute in the model called counts and these functions just get the total number from Mesa's data collector of each of our statuses. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e93ff02-3d17-4663-a33a-5113e65c6aef",
   "metadata": {
    "has_explanation": true
   },
   "outputs": [],
   "source": [
    "# Functions needed for datacollector\n",
    "def get_infected_count(model):\n",
    "    return model.counts[\"infected\"]\n",
    "\n",
    "\n",
    "def get_susceptible_count(model):\n",
    "    return model.counts[\"susceptible\"]\n",
    "\n",
    "\n",
    "def get_recovered_count(model):\n",
    "    return model.counts[\"recovered\"]\n",
    "\n",
    "\n",
    "def get_dead_count(model):\n",
    "    return model.counts[\"dead\"]\n",
    "\n",
    "\n",
    "def get_hotspot_count(model):\n",
    "    return model.counts[\"hotspot\"]\n",
    "\n",
    "\n",
    "def get_safe_count(model):\n",
    "    return model.counts[\"safe\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ef38b13-a442-480d-ab23-0fa4366dd716",
   "metadata": {
    "explanatory": true
   },
   "source": [
    "\n",
    "Now to finish the model so we can add the interface we add datacollection and a stop condition. As these updates are interspersed throughout the class. The comment `#added` is used to make the changes easier to identify.\n",
    "\n",
    "First, we add an attribute called `self.counts` which will track our the agent types (e.g. infected). We will initialize it as None. We then initialize the counts in our next line `self.reset_counts()`. This helper function located directly above the step function, resets the counts of each type of agent so it is always based on the current situation in the Model. \n",
    "\n",
    "We are then going to add the attribute self.running so we can input the stop condition. Next we set our our data collector that call our functions from the previous cell which collects our agent types\n",
    "\n",
    "With these added we can now call `self.reset_counts` and `self.datacollector.collect` in our step function so it collect our agent states each step. \n",
    "\n",
    "Finally we add a stop condition. If no PersonAgent is infected the pandemic is over and we stop the model. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d24aa5cb-61bc-4651-9792-3ab68ff3b0f6",
   "metadata": {
    "has_explanation": true
   },
   "outputs": [],
   "source": [
    "class GeoSIR(mesa.Model):\n",
    "    \"\"\"Model class for a simplistic infection model.\"\"\"\n",
    "\n",
    "    # Geographical parameters for desired map\n",
    "    geojson_regions = \"data/TorontoNeighbourhoods.geojson\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        pop_size=30,\n",
    "        mobility_range=500,\n",
    "        init_infection=0.2,\n",
    "        exposure_dist=500,\n",
    "        max_infection_risk=0.2,\n",
    "        max_recovery_time=5,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        # Space\n",
    "        self.space = mg.GeoSpace(warn_crs_conversion=False)\n",
    "        # Data Collection\n",
    "        self.counts = None  # added\n",
    "        self.reset_counts()  # added\n",
    "\n",
    "        # SIR model parameters\n",
    "        self.pop_size = pop_size\n",
    "        self.mobility_range = mobility_range\n",
    "        self.initial_infection = init_infection\n",
    "        self.exposure_distance = exposure_dist\n",
    "        self.infection_risk = max_infection_risk\n",
    "        self.recovery_rate = max_recovery_time\n",
    "        self.running = True  # added\n",
    "        # added\n",
    "        self.datacollector = mesa.DataCollector(\n",
    "            {\n",
    "                \"infected\": get_infected_count,\n",
    "                \"susceptible\": get_susceptible_count,\n",
    "                \"recovered\": get_recovered_count,\n",
    "                \"dead\": get_dead_count,\n",
    "                \"safe\": get_safe_count,\n",
    "                \"hotspot\": get_hotspot_count,\n",
    "            }\n",
    "        )\n",
    "\n",
    "        # Set up the Neighbourhood patches for every region in file\n",
    "        ac = mg.AgentCreator(NeighbourhoodAgent, model=self)\n",
    "        neighbourhood_agents = ac.from_file(self.geojson_regions)\n",
    "\n",
    "        # Add neighbourhood agents to space\n",
    "        self.space.add_agents(neighbourhood_agents)\n",
    "\n",
    "        # Generate random location, add agent to grid\n",
    "        for i in range(pop_size):\n",
    "            # assess if they are infected\n",
    "            if self.random.random() < self.initial_infection:\n",
    "                agent_type = \"infected\"\n",
    "            else:\n",
    "                agent_type = \"susceptible\"\n",
    "            # determine movement range\n",
    "            mobility_range = self.random.randint(0, self.mobility_range)\n",
    "            # determine agent recovery rate\n",
    "            recover = self.random.randint(1, self.recovery_rate)\n",
    "            # determine agents infection risk\n",
    "            infection_risk = self.random.uniform(0, self.infection_risk)\n",
    "            # determine agent death probability\n",
    "            death_risk = self.random.uniform(0, 0.05)\n",
    "\n",
    "            # Generate PersonAgent population\n",
    "            unique_person = mg.AgentCreator(\n",
    "                PersonAgent,\n",
    "                model=self,\n",
    "                crs=self.space.crs,\n",
    "                agent_kwargs={\n",
    "                    \"agent_type\": agent_type,\n",
    "                    \"mobility_range\": mobility_range,\n",
    "                    \"recovery_rate\": recover,\n",
    "                    \"infection_risk\": infection_risk,\n",
    "                    \"death_risk\": death_risk,\n",
    "                },\n",
    "            )\n",
    "\n",
    "            x_home, y_home = self.find_home(neighbourhood_agents)\n",
    "\n",
    "            this_person = unique_person.create_agent(Point(x_home, y_home))\n",
    "            self.space.add_agents(this_person)\n",
    "\n",
    "    def find_home(self, neighbourhood_agents):\n",
    "        \"\"\"Find start location of agent\"\"\"\n",
    "\n",
    "        # identify location\n",
    "        this_neighbourhood = self.random.randint(\n",
    "            0, len(neighbourhood_agents) - 1\n",
    "        )  # Region where agent starts\n",
    "        center_x, center_y = neighbourhood_agents[\n",
    "            this_neighbourhood\n",
    "        ].geometry.centroid.coords.xy\n",
    "        this_bounds = neighbourhood_agents[this_neighbourhood].geometry.bounds\n",
    "        spread_x = int(\n",
    "            this_bounds[2] - this_bounds[0]\n",
    "        )  # Heuristic for agent spread in region\n",
    "        spread_y = int(this_bounds[3] - this_bounds[1])\n",
    "        this_x = center_x[0] + self.random.randint(0, spread_x) - spread_x / 2\n",
    "        this_y = center_y[0] + self.random.randint(0, spread_y) - spread_y / 2\n",
    "\n",
    "        return this_x, this_y\n",
    "\n",
    "    # added\n",
    "    def reset_counts(self):\n",
    "        self.counts = {\n",
    "            \"susceptible\": 0,\n",
    "            \"infected\": 0,\n",
    "            \"recovered\": 0,\n",
    "            \"dead\": 0,\n",
    "            \"safe\": 0,\n",
    "            \"hotspot\": 0,\n",
    "        }\n",
    "\n",
    "    def step(self):\n",
    "        \"\"\"Run one step of the model.\"\"\"\n",
    "\n",
    "        self.reset_counts()  # added\n",
    "        \n",
    "        # Activate PersonAgents in random order\n",
    "        self.agents_by_type[PersonAgent].shuffle_do(\"step\")\n",
    "        # For NeighbourhoodAgents the order doesn't matter, since they update independently from each other\n",
    "        self.agents_by_type[NeighbourhoodAgent].do(\"step\")\n",
    "        \n",
    "        self.datacollector.collect(self)  # added\n",
    "\n",
    "        # Run until no one is infected\n",
    "        if self.counts[\"infected\"] == 0:\n",
    "            self.running = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae5bf0be-e748-4850-9cec-ecbb5a0c2201",
   "metadata": {
    "explanatory": true
   },
   "source": [
    "\n",
    "To test our code we will run the model through 5 steps and then call the model dataframe via data collector with `get_model_vars_dataframe()`. This will show a Pandas DataFrame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "027a8065-13d7-4bc2-96ba-9c556ce16b3c",
   "metadata": {
    "has_explanation": true
   },
   "outputs": [],
   "source": [
    "model = GeoSIR()\n",
    "for i in range(5):\n",
    "    model.step()\n",
    "\n",
    "model.datacollector.get_model_vars_dataframe()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f299f22b-9115-48a4-8277-ede19628707c",
   "metadata": {},
   "source": [
    "## Part 3 - Add Interface"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1035cc55-0169-4096-8b63-d99b8259c980",
   "metadata": {},
   "source": [
    "Adding the interface requires three steps: \n",
    "\n",
    "1. Define the agent portrayal\n",
    "2. Set the sliders for the model parameters\n",
    "3. Call the model through the Mesa-Geo visualization model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5924dd0-38e6-4b16-911e-4bd93d80976e",
   "metadata": {
    "explanatory": true
   },
   "source": [
    "\n",
    "Visualizing agents is done through a function that is is passed in as a parameter. By default agents they are Point geometries are rendered as circles. However, Mesa uses [ipyleaflet](https://ipyleaflet.readthedocs.io/en/latest/layers/marker.html) Users can pass through any Point geometry for their Agent (i.e. Marker, Circle, Icon, AwesomeIcon). To show this we will use different colors for the PersonAgent base don infection status and if they die, we will use the [Font Awesome Icons](https://fontawesome.com/v4/) and represent them with an x, in the traditional ipyleaflet marker. \n",
    "\n",
    "We will also change the color of the NeighbourhoodAgent based whether or not it is a hotspot\n",
    "\n",
    "Next we will build Sliders for each of our input parameters. These use the [Solara's input approach](https://solara.dev/documentation/components/input/slider). This is stored in a dictionary of dictionaries that is then passed through in the model instantiation. \n",
    "\n",
    "If you want the model to fill the entire screen you can hit the expand button in the upper right. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d19cb45-3499-47e2-b5b8-2096bc0fc268",
   "metadata": {
    "editable": true,
    "has_explanation": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def SIR_draw(agent):\n",
    "    \"\"\"\n",
    "    Portrayal Method for canvas\n",
    "    \"\"\"\n",
    "    portrayal = {}\n",
    "    if isinstance(agent, PersonAgent):\n",
    "        if agent.atype == \"susceptible\":\n",
    "            portrayal[\"color\"] = \"Green\"\n",
    "        elif agent.atype == \"infected\":\n",
    "            portrayal[\"color\"] = \"Red\"\n",
    "        elif agent.atype == \"recovered\":\n",
    "            portrayal[\"color\"] = \"Blue\"\n",
    "        else:\n",
    "            portrayal[\"marker_type\"] = \"AwesomeIcon\"\n",
    "            portrayal[\"name\"] = \"times\"\n",
    "            portrayal[\"icon_properties\"] = {\n",
    "                \"marker_color\": \"black\",\n",
    "                \"icon_color\": \"white\",\n",
    "            }\n",
    "\n",
    "    if isinstance(agent, NeighbourhoodAgent):\n",
    "        if agent.atype == \"hotspot\":\n",
    "            portrayal[\"color\"] = \"Red\"\n",
    "        else:\n",
    "            portrayal[\"color\"] = \"Green\"\n",
    "\n",
    "    return portrayal\n",
    "\n",
    "\n",
    "model_params = {\n",
    "    \"pop_size\": {\n",
    "        \"type\": \"SliderInt\",\n",
    "        \"value\": 80,\n",
    "        \"label\": \"Population Size\",\n",
    "        \"min\": 0,\n",
    "        \"max\": 100,\n",
    "        \"step\": 1,\n",
    "    },\n",
    "    \"mobility_range\": {\n",
    "        \"type\": \"SliderInt\",\n",
    "        \"value\": 500,\n",
    "        \"label\": \"Max Possible Agent Movement\",\n",
    "        \"min\": 100,\n",
    "        \"max\": 1000,\n",
    "        \"step\": 50,\n",
    "    },\n",
    "    \"init_infection\": {\n",
    "        \"type\": \"SliderFloat\",\n",
    "        \"value\": 0.4,\n",
    "        \"label\": \"Initial Infection\",\n",
    "        \"min\": 0.0,\n",
    "        \"max\": 1.0,\n",
    "        \"step\": 0.1,\n",
    "    },\n",
    "    \"exposure_dist\": {\n",
    "        \"type\": \"SliderInt\",\n",
    "        \"value\": 800,\n",
    "        \"label\": \"Exposure Distance\",\n",
    "        \"min\": 100,\n",
    "        \"max\": 1000,\n",
    "        \"step\": 50,\n",
    "    },\n",
    "    \"max_infection_risk\": {\n",
    "        \"type\": \"SliderFloat\",\n",
    "        \"value\": 0.7,\n",
    "        \"label\": \"Maximum Infection Risk\",\n",
    "        \"min\": 0.0,\n",
    "        \"max\": 1.0,\n",
    "        \"step\": 0.1,\n",
    "    },\n",
    "    \"max_recovery_time\": {\n",
    "        \"type\": \"SliderInt\",\n",
    "        \"value\": 7,\n",
    "        \"label\": \"Maximum Number of Steps to Recover\",\n",
    "        \"min\": 1,\n",
    "        \"max\": 10,\n",
    "        \"step\": 1,\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee5f4df8-f8af-40a1-a469-843b4c043866",
   "metadata": {
    "explanatory": true
   },
   "source": [
    "To create the model with the interface we use [Mesa's GeoJupyterViz module](https://github.com/projectmesa/mesa-geo/tree/main/mesa_geo/visualization). First we pass in the model class and next the parameters. We then switch to key word arguments. First measures, in this case of list of lists, where the first list will be a chart of the PersonAgent statuses and the second chart will be the NeighbourhoodAgent statuses. We also pass in a name, our agent portrayal function a zoom level and in this case set the scroll wheel zoom to false. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "324cabe0-68be-4aab-a70d-10ffcd328114",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = GeoSIR()\n",
    "page = SolaraViz(\n",
    "    model,\n",
    "    name=\"GeoSIR\",\n",
    "    model_params=model_params,\n",
    "    components=[\n",
    "        make_geospace_component(SIR_draw, zoom=12, scroll_wheel_zoom=False),\n",
    "        make_plot_component([\"infected\", \"susceptible\", \"recovered\", \"dead\"]),\n",
    "        make_plot_component([\"safe\", \"hotspot\"]),\n",
    "    ],\n",
    ")\n",
    "# This is required to render the visualization in the Jupyter notebook\n",
    "page"
   ]
  }
 ],
 "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}