SIR example

.gitignore

@@ -320,4 +320,7 @@ pip-log.txt
 .mr.developer.cfg
 
 #Visual studio code
-.vscode
+.vscode
+
+# Idea file
+.idea

examples/GeoSIR/README.md

@@ -0,0 +1,22 @@
+# GeoSIR Epidemics Model
+
+## Summary
+
+This is a geoversion of a simple agent-based pandemic SIR model, as an example to show the capabilities of mesa-geo.
+
+It uses geographical data of Toronto's regions on top of a an Leaflet map to show the location of agents (in a continuous space).
+
+Person agents are initially located in random positions in the city, then start moving around unless they die.
+A fraction of agents start with an infection and may recover or die in each step.
+Susceptible agents (those who have never been infected) who come in proximity with an infected agent may become infected.
+
+Neighbourhood agents represent neighbourhoods in the Toronto, and become hot-spots (colored red) if there are infected agents inside them.
+Data obtained from [this link](http://adamw523.com/toronto-geojson/).
+
+## How to run
+
+To run the model interactively, run ```mesa runserver``` in this directory. e.g.
+```
+mesa runserver
+```
+Then open your browser to [http://127.0.0.1:8521/](http://127.0.0.1:8521/), set the desired parameters, press Reset, and then Run.

examples/GeoSIR/model.py

@@ -0,0 +1,240 @@
+from mesa.datacollection import DataCollector
+from mesa import Model
+from mesa.time import BaseScheduler
+from mesa_geo.geoagent import GeoAgent, AgentCreator
+from mesa_geo import GeoSpace
+from shapely.geometry import Point
+
+
+class PersonAgent(GeoAgent):
+    """Person Agent."""
+
+    def __init__(
+        self,
+        unique_id,
+        model,
+        shape,
+        agent_type="susceptible",
+        mobility_range=100,
+        recovery_rate=0.2,
+        death_risk=0.1,
+        init_infected=0.1,
+    ):
+        """
+        Create a new person agent.
+        :param unique_id:   Unique identifier for the agent
+        :param model:       Model in which the agent runs
+        :param shape:       Shape object for the agent
+        :param agent_type:  Indicator if agent is infected ("infected", "susceptible", "recovered" or "dead")
+        :param mobility_range:  Range of distance to move in one step
+        """
+        super().__init__(unique_id, model, shape)
+        # Agent parameters
+        self.atype = agent_type
+        self.mobility_range = mobility_range
+        self.recovery_rate = recovery_rate
+        self.death_risk = death_risk
+
+        # Random choose if infected
+        if self.random.random() < init_infected:
+            self.atype = "infected"
+            self.model.counts["infected"] += 1  # Adjust initial counts
+            self.model.counts["susceptible"] -= 1
+
+    def move_point(self, dx, dy):
+        """
+        Move a point by creating a new one
+        :param dx:  Distance to move in x-axis
+        :param dy:  Distance to move in y-axis
+        """
+        return Point(self.shape.x + dx, self.shape.y + dy)
+
+    def step(self):
+        """Advance one step."""
+        # If susceptible, check if exposed
+        if self.atype == "susceptible":
+            neighbors = self.model.grid.get_neighbors_within_distance(
+                self, self.model.exposure_distance
+            )
+            for neighbor in neighbors:
+                if (
+                    neighbor.atype == "infected"
+                    and self.random.random() < self.model.infection_risk
+                ):
+                    self.atype = "infected"
+                    break
+
+        # If infected, check if it recovers or if it dies
+        elif self.atype == "infected":
+            if self.random.random() < self.recovery_rate:
+                self.atype = "recovered"
+            elif self.random.random() < self.death_risk:
+                self.atype = "dead"
+
+        # If not dead, move
+        if self.atype != "dead":
+            move_x = self.random.randint(-self.mobility_range, self.mobility_range)
+            move_y = self.random.randint(-self.mobility_range, self.mobility_range)
+            self.shape = self.move_point(move_x, move_y)  # Reassign shape
+
+        self.model.counts[self.atype] += 1  # Count agent type
+
+    def __repr__(self):
+        return "Person " + str(self.unique_id)
+
+
+class NeighbourhoodAgent(GeoAgent):
+    """Neighbourhood agent. Changes color according to number of infected inside it."""
+
+    def __init__(self, unique_id, model, shape, agent_type="safe", hotspot_threshold=1):
+        """
+        Create a new Neighbourhood agent.
+        :param unique_id:   Unique identifier for the agent
+        :param model:       Model in which the agent runs
+        :param shape:       Shape object for the agent
+        :param agent_type:  Indicator if agent is infected ("infected", "susceptible", "recovered" or "dead")
+        :param hotspot_threshold:   Number of infected agents in region to be considered a hot-spot
+        """
+        super().__init__(unique_id, model, shape)
+        self.atype = agent_type
+        self.hotspot_threshold = (
+            hotspot_threshold
+        )  # When a neighborhood is considered a hot-spot
+        self.color_hotspot()
+
+    def step(self):
+        """Advance agent one step."""
+        self.color_hotspot()
+        self.model.counts[self.atype] += 1  # Count agent type
+
+    def color_hotspot(self):
+        # Decide if this region agent is a hot-spot (if more than threshold person agents are infected)
+        neighbors = self.model.grid.get_intersecting_agents(self)
+        infected_neighbors = [
+            neighbor for neighbor in neighbors if neighbor.atype == "infected"
+        ]
+        if len(infected_neighbors) >= self.hotspot_threshold:
+            self.atype = "hotspot"
+        else:
+            self.atype = "safe"
+
+    def __repr__(self):
+        return "Neighborhood " + str(self.unique_id)
+
+
+class InfectedModel(Model):
+    """Model class for a simplistic infection model."""
+
+    # Geographical parameters for desired map
+    MAP_COORDS = [43.741667, -79.373333]  # Toronto
+    geojson_regions = "TorontoNeighbourhoods.geojson"
+    unique_id = "HOODNUM"
+
+    def __init__(self, pop_size, init_infected, exposure_distance, infection_risk=0.2):
+        """
+        Create a new InfectedModel
+        :param pop_size:        Size of population
+        :param init_infected:   Probability of a person agent to start as infected
+        :param exposure_distance:   Proximity distance between agents to be exposed to each other
+        :param infection_risk:      Probability of agent to become infected, if it has been exposed to another infected
+        """
+        self.schedule = BaseScheduler(self)
+        self.grid = GeoSpace()
+        self.steps = 0
+        self.counts = None
+        self.reset_counts()
+
+        # SIR model parameters
+        self.pop_size = pop_size
+        self.counts["susceptible"] = pop_size
+        self.exposure_distance = exposure_distance
+        self.infection_risk = infection_risk
+
+        self.running = True
+        self.datacollector = DataCollector(
+            {
+                "infected": get_infected_count,
+                "susceptible": get_susceptible_count,
+                "recovered": get_recovered_count,
+                "dead": get_dead_count,
+            }
+        )
+
+        # Set up the Neighbourhood patches for every region in file (add to schedule later)
+        AC = AgentCreator(NeighbourhoodAgent, {"model": self})
+        neighbourhood_agents = AC.from_file(
+            self.geojson_regions, unique_id=self.unique_id
+        )
+        self.grid.add_agents(neighbourhood_agents)
+
+        # Generate PersonAgent population
+        ac_population = AgentCreator(
+            PersonAgent, {"model": self, "init_infected": init_infected}
+        )
+        # Generate random location, add agent to grid and scheduler
+        for i in range(pop_size):
+            this_neighbourhood = self.random.randint(
+                0, len(neighbourhood_agents) - 1
+            )  # Region where agent starts
+            center_x, center_y = neighbourhood_agents[
+                this_neighbourhood
+            ].shape.centroid.coords.xy
+            this_bounds = neighbourhood_agents[this_neighbourhood].shape.bounds
+            spread_x = int(
+                this_bounds[2] - this_bounds[0]
+            )  # Heuristic for agent spread in region
+            spread_y = int(this_bounds[3] - this_bounds[1])
+            this_x = center_x[0] + self.random.randint(0, spread_x) - spread_x / 2
+            this_y = center_y[0] + self.random.randint(0, spread_y) - spread_y / 2
+            this_person = ac_population.create_agent(
+                Point(this_x, this_y), "P" + str(i)
+            )
+            self.grid.add_agents(this_person)
+            self.schedule.add(this_person)
+
+        # Add the neighbourhood agents to schedule AFTER person agents,
+        # to allow them to update their color by using BaseScheduler
+        for agent in neighbourhood_agents:
+            self.schedule.add(agent)
+
+        self.datacollector.collect(self)
+
+    def reset_counts(self):
+        self.counts = {
+            "susceptible": 0,
+            "infected": 0,
+            "recovered": 0,
+            "dead": 0,
+            "safe": 0,
+            "hotspot": 0,
+        }
+
+    def step(self):
+        """Run one step of the model."""
+        self.steps += 1
+        self.reset_counts()
+        self.schedule.step()
+        self.grid._recreate_rtree()  # Recalculate spatial tree, because agents are moving
+
+        self.datacollector.collect(self)
+
+        # Run until no one is infected
+        if self.counts["infected"] == 0:
+            self.running = False
+
+
+# Functions needed for datacollector
+def get_infected_count(model):
+    return model.counts["infected"]
+
+
+def get_susceptible_count(model):
+    return model.counts["susceptible"]
+
+
+def get_recovered_count(model):
+    return model.counts["recovered"]
+
+
+def get_dead_count(model):
+    return model.counts["dead"]

examples/GeoSIR/run.py

@@ -0,0 +1,3 @@
+from server import server
+
+server.launch()

examples/GeoSIR/server.py

@@ -0,0 +1,65 @@
+from mesa_geo.visualization.ModularVisualization import ModularServer
+from mesa.visualization.modules import ChartModule, TextElement
+from mesa.visualization.UserParam import UserSettableParameter
+from model import InfectedModel, PersonAgent
+from mesa_geo.visualization.MapModule import MapModule
+
+
+class InfectedText(TextElement):
+    """
+    Display a text count of how many steps have been taken
+    """
+
+    def __init__(self):
+        pass
+
+    def render(self, model):
+        return "Steps: " + str(model.steps)
+
+
+model_params = {
+    "pop_size": UserSettableParameter("slider", "Population size", 30, 10, 100, 10),
+    "init_infected": UserSettableParameter(
+        "slider", "Fraction initial infection", 0.2, 0.00, 1.0, 0.05
+    ),
+    "exposure_distance": UserSettableParameter(
+        "slider", "Exposure distance", 500, 100, 1000, 100
+    ),
+}
+
+
+def infected_draw(agent):
+    """
+    Portrayal Method for canvas
+    """
+    portrayal = dict()
+    if isinstance(agent, PersonAgent):
+        portrayal["radius"] = "2"
+    if agent.atype in ["hotspot", "infected"]:
+        portrayal["color"] = "Red"
+    elif agent.atype in ["safe", "susceptible"]:
+        portrayal["color"] = "Green"
+    elif agent.atype in ["recovered"]:
+        portrayal["color"] = "Blue"
+    elif agent.atype in ["dead"]:
+        portrayal["color"] = "Black"
+    return portrayal
+
+
+infected_text = InfectedText()
+map_element = MapModule(infected_draw, InfectedModel.MAP_COORDS, 10, 500, 500)
+infected_chart = ChartModule(
+    [
+        {"Label": "infected", "Color": "Red"},
+        {"Label": "susceptible", "Color": "Green"},
+        {"Label": "recovered", "Color": "Blue"},
+        {"Label": "dead", "Color": "Black"},
+    ]
+)
+server = ModularServer(
+    InfectedModel,
+    [map_element, infected_text, infected_chart],
+    "Basic agent-based SIR model",
+    model_params,
+)
+server.launch()