From d7a3834c99a3be809abe2edc8b83610f3d4438ba Mon Sep 17 00:00:00 2001 From: Ewout ter Hoeven Date: Tue, 15 Oct 2024 17:50:56 +0200 Subject: [PATCH] examples/virus: Restructure - VirusAgent in agents.py - Flatten structure - Remove old visualisation code --- examples/basic/virus_on_network/README.md | 4 +- examples/basic/virus_on_network/agents.py | 70 +++++++++ examples/basic/virus_on_network/app.py | 2 +- .../{virus_on_network => }/model.py | 72 +-------- .../basic/virus_on_network/requirements.txt | 2 - examples/basic/virus_on_network/run.py | 3 - .../virus_on_network/server.py | 140 ------------------ 7 files changed, 76 insertions(+), 217 deletions(-) create mode 100644 examples/basic/virus_on_network/agents.py rename examples/basic/virus_on_network/{virus_on_network => }/model.py (57%) delete mode 100644 examples/basic/virus_on_network/requirements.txt delete mode 100644 examples/basic/virus_on_network/run.py delete mode 100644 examples/basic/virus_on_network/virus_on_network/server.py diff --git a/examples/basic/virus_on_network/README.md b/examples/basic/virus_on_network/README.md index f6a51fd580b..484a4c6ffdf 100644 --- a/examples/basic/virus_on_network/README.md +++ b/examples/basic/virus_on_network/README.md @@ -43,9 +43,9 @@ Directly run the file ``run.py`` in the terminal. e.g. ## Files -* ``run.py``: Launches a model visualization server. * ``model.py``: Contains the agent class, and the overall model class. -* ``server.py``: Defines classes for visualizing the model (network layout) in the browser via Mesa's modular server, and instantiates a visualization server. +* ``agents.py``: Contains the agent class. +* ``app.py``: Contains the code for the interactive Solara visualization. ## Further Reading diff --git a/examples/basic/virus_on_network/agents.py b/examples/basic/virus_on_network/agents.py new file mode 100644 index 00000000000..3d9486aa4eb --- /dev/null +++ b/examples/basic/virus_on_network/agents.py @@ -0,0 +1,70 @@ +from mesa import Agent +from enum import Enum + + +class State(Enum): + SUSCEPTIBLE = 0 + INFECTED = 1 + RESISTANT = 2 + + +class VirusAgent(Agent): + """ + Individual Agent definition and its properties/interaction methods + """ + + def __init__( + self, + model, + initial_state, + virus_spread_chance, + virus_check_frequency, + recovery_chance, + gain_resistance_chance, + ): + super().__init__(model) + + self.state = initial_state + + self.virus_spread_chance = virus_spread_chance + self.virus_check_frequency = virus_check_frequency + self.recovery_chance = recovery_chance + self.gain_resistance_chance = gain_resistance_chance + + def try_to_infect_neighbors(self): + neighbors_nodes = self.model.grid.get_neighborhood( + self.pos, include_center=False + ) + susceptible_neighbors = [ + agent + for agent in self.model.grid.get_cell_list_contents(neighbors_nodes) + if agent.state is State.SUSCEPTIBLE + ] + for a in susceptible_neighbors: + if self.random.random() < self.virus_spread_chance: + a.state = State.INFECTED + + def try_gain_resistance(self): + if self.random.random() < self.gain_resistance_chance: + self.state = State.RESISTANT + + def try_remove_infection(self): + # Try to remove + if self.random.random() < self.recovery_chance: + # Success + self.state = State.SUSCEPTIBLE + self.try_gain_resistance() + else: + # Failed + self.state = State.INFECTED + + def try_check_situation(self): + if (self.random.random() < self.virus_check_frequency) and ( + self.state is State.INFECTED + ): + self.try_remove_infection() + + def step(self): + if self.state is State.INFECTED: + self.try_to_infect_neighbors() + self.try_check_situation() diff --git a/examples/basic/virus_on_network/app.py b/examples/basic/virus_on_network/app.py index 4ed96b79b0b..89c377b25d2 100644 --- a/examples/basic/virus_on_network/app.py +++ b/examples/basic/virus_on_network/app.py @@ -4,7 +4,7 @@ from matplotlib.figure import Figure from matplotlib.ticker import MaxNLocator from mesa.visualization import SolaraViz, Slider, make_space_matplotlib -from virus_on_network.model import State, VirusOnNetwork, number_infected +from model import State, VirusOnNetwork, number_infected def agent_portrayal(graph): diff --git a/examples/basic/virus_on_network/virus_on_network/model.py b/examples/basic/virus_on_network/model.py similarity index 57% rename from examples/basic/virus_on_network/virus_on_network/model.py rename to examples/basic/virus_on_network/model.py index d892a0c4c06..d0ad600b8d7 100644 --- a/examples/basic/virus_on_network/virus_on_network/model.py +++ b/examples/basic/virus_on_network/model.py @@ -1,14 +1,10 @@ import math -from enum import Enum import mesa +from mesa import Model import networkx as nx - -class State(Enum): - SUSCEPTIBLE = 0 - INFECTED = 1 - RESISTANT = 2 +from agents import VirusAgent, State def number_state(model, state): @@ -27,7 +23,7 @@ def number_resistant(model): return number_state(model, State.RESISTANT) -class VirusOnNetwork(mesa.Model): +class VirusOnNetwork(Model): """ A virus model with some number of agents """ @@ -102,65 +98,3 @@ def step(self): def run_model(self, n): for i in range(n): self.step() - - -class VirusAgent(mesa.Agent): - """ - Individual Agent definition and its properties/interaction methods - """ - - def __init__( - self, - model, - initial_state, - virus_spread_chance, - virus_check_frequency, - recovery_chance, - gain_resistance_chance, - ): - super().__init__(model) - - self.state = initial_state - - self.virus_spread_chance = virus_spread_chance - self.virus_check_frequency = virus_check_frequency - self.recovery_chance = recovery_chance - self.gain_resistance_chance = gain_resistance_chance - - def try_to_infect_neighbors(self): - neighbors_nodes = self.model.grid.get_neighborhood( - self.pos, include_center=False - ) - susceptible_neighbors = [ - agent - for agent in self.model.grid.get_cell_list_contents(neighbors_nodes) - if agent.state is State.SUSCEPTIBLE - ] - for a in susceptible_neighbors: - if self.random.random() < self.virus_spread_chance: - a.state = State.INFECTED - - def try_gain_resistance(self): - if self.random.random() < self.gain_resistance_chance: - self.state = State.RESISTANT - - def try_remove_infection(self): - # Try to remove - if self.random.random() < self.recovery_chance: - # Success - self.state = State.SUSCEPTIBLE - self.try_gain_resistance() - else: - # Failed - self.state = State.INFECTED - - def try_check_situation(self): - if (self.random.random() < self.virus_check_frequency) and ( - self.state is State.INFECTED - ): - self.try_remove_infection() - - def step(self): - if self.state is State.INFECTED: - self.try_to_infect_neighbors() - self.try_check_situation() diff --git a/examples/basic/virus_on_network/requirements.txt b/examples/basic/virus_on_network/requirements.txt deleted file mode 100644 index 03e3c237233..00000000000 --- a/examples/basic/virus_on_network/requirements.txt +++ /dev/null @@ -1,2 +0,0 @@ -networkx>=2.0 -mesa~=2.0 \ No newline at end of file diff --git a/examples/basic/virus_on_network/run.py b/examples/basic/virus_on_network/run.py deleted file mode 100644 index c911c372a9a..00000000000 --- a/examples/basic/virus_on_network/run.py +++ /dev/null @@ -1,3 +0,0 @@ -from virus_on_network.server import server - -server.launch(open_browser=True) diff --git a/examples/basic/virus_on_network/virus_on_network/server.py b/examples/basic/virus_on_network/virus_on_network/server.py deleted file mode 100644 index dcc7643f080..00000000000 --- a/examples/basic/virus_on_network/virus_on_network/server.py +++ /dev/null @@ -1,140 +0,0 @@ -import math - -import mesa - -from .model import State, VirusOnNetwork, number_infected - - -def network_portrayal(G): - # The model ensures there is always 1 agent per node - - def node_color(agent): - return {State.INFECTED: "#FF0000", State.SUSCEPTIBLE: "#008000"}.get( - agent.state, "#808080" - ) - - def edge_color(agent1, agent2): - if State.RESISTANT in (agent1.state, agent2.state): - return "#000000" - return "#e8e8e8" - - def edge_width(agent1, agent2): - if State.RESISTANT in (agent1.state, agent2.state): - return 3 - return 2 - - def get_agents(source, target): - return G.nodes[source]["agent"][0], G.nodes[target]["agent"][0] - - portrayal = {} - portrayal["nodes"] = [ - { - "size": 6, - "color": node_color(agents[0]), - "tooltip": f"id: {agents[0].unique_id}
state: {agents[0].state.name}", - } - for (_, agents) in G.nodes.data("agent") - ] - - portrayal["edges"] = [ - { - "source": source, - "target": target, - "color": edge_color(*get_agents(source, target)), - "width": edge_width(*get_agents(source, target)), - } - for (source, target) in G.edges - ] - - return portrayal - - -network = mesa.visualization.NetworkModule( - portrayal_method=network_portrayal, - canvas_height=500, - canvas_width=500, -) -chart = mesa.visualization.ChartModule( - [ - {"Label": "Infected", "Color": "#FF0000"}, - {"Label": "Susceptible", "Color": "#008000"}, - {"Label": "Resistant", "Color": "#808080"}, - ] -) - - -def get_resistant_susceptible_ratio(model): - ratio = model.resistant_susceptible_ratio() - ratio_text = "∞" if ratio is math.inf else f"{ratio:.2f}" - infected_text = str(number_infected(model)) - - return f"Resistant/Susceptible Ratio: {ratio_text}
Infected Remaining: {infected_text}" - - -model_params = { - "num_nodes": mesa.visualization.Slider( - name="Number of agents", - value=10, - min_value=10, - max_value=100, - step=1, - description="Choose how many agents to include in the model", - ), - "avg_node_degree": mesa.visualization.Slider( - name="Avg Node Degree", - value=3, - min_value=3, - max_value=8, - step=1, - description="Avg Node Degree", - ), - "initial_outbreak_size": mesa.visualization.Slider( - name="Initial Outbreak Size", - value=1, - min_value=1, - max_value=10, - step=1, - description="Initial Outbreak Size", - ), - "virus_spread_chance": mesa.visualization.Slider( - name="Virus Spread Chance", - value=0.4, - min_value=0.0, - max_value=1.0, - step=0.1, - description="Probability that susceptible neighbor will be infected", - ), - "virus_check_frequency": mesa.visualization.Slider( - name="Virus Check Frequency", - value=0.4, - min_value=0.0, - max_value=1.0, - step=0.1, - description="Frequency the nodes check whether they are infected by a virus", - ), - "recovery_chance": mesa.visualization.Slider( - name="Recovery Chance", - value=0.3, - min_value=0.0, - max_value=1.0, - step=0.1, - description="Probability that the virus will be removed", - ), - "gain_resistance_chance": mesa.visualization.Slider( - name="Gain Resistance Chance", - value=0.5, - min_value=0.0, - max_value=1.0, - step=0.1, - description="Probability that a recovered agent will become " - "resistant to this virus in the future", - ), -} - -server = mesa.visualization.ModularServer( - model_cls=VirusOnNetwork, - visualization_elements=[network, get_resistant_susceptible_ratio, chart], - name="Virus on Network Model", - model_params=model_params, -) -server.port = 8521