examples: Convert virus_on_network to simple namespace

RichardScottOZ · May 24, 2022 · e2ccd6d · e2ccd6d
1 parent 8077cdf
commit e2ccd6d
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 25 deletions.
diff --git a/examples/virus_on_network/virus_on_network/model.py b/examples/virus_on_network/virus_on_network/model.py
@@ -2,10 +2,7 @@
 from enum import Enum
 import networkx as nx
 
-from mesa import Agent, Model
-from mesa.time import RandomActivation
-from mesa.datacollection import DataCollector
-from mesa.space import NetworkGrid
+import mesa
 
 
 class State(Enum):
@@ -30,7 +27,7 @@ def number_resistant(model):
     return number_state(model, State.RESISTANT)
 
 
-class VirusOnNetwork(Model):
+class VirusOnNetwork(mesa.Model):
     """A virus model with some number of agents"""
 
     def __init__(
@@ -47,8 +44,8 @@ def __init__(
         self.num_nodes = num_nodes
         prob = avg_node_degree / self.num_nodes
         self.G = nx.erdos_renyi_graph(n=self.num_nodes, p=prob)
-        self.grid = NetworkGrid(self.G)
-        self.schedule = RandomActivation(self)
+        self.grid = mesa.space.NetworkGrid(self.G)
+        self.schedule = mesa.time.RandomActivation(self)
         self.initial_outbreak_size = (
             initial_outbreak_size if initial_outbreak_size <= num_nodes else num_nodes
         )
@@ -57,7 +54,7 @@ def __init__(
         self.recovery_chance = recovery_chance
         self.gain_resistance_chance = gain_resistance_chance
 
-        self.datacollector = DataCollector(
+        self.datacollector = mesa.DataCollector(
             {
                 "Infected": number_infected,
                 "Susceptible": number_susceptible,
@@ -106,7 +103,7 @@ def run_model(self, n):
             self.step()
 
 
-class VirusAgent(Agent):
+class VirusAgent(mesa.Agent):
     def __init__(
         self,
         unique_id,

diff --git a/examples/virus_on_network/virus_on_network/server.py b/examples/virus_on_network/virus_on_network/server.py
@@ -1,10 +1,7 @@
 import math
 
-from mesa.visualization.ModularVisualization import ModularServer
-from mesa.visualization.UserParam import UserSettableParameter
-from mesa.visualization.modules import ChartModule
-from mesa.visualization.modules import NetworkModule
-from mesa.visualization.modules import TextElement
+import mesa
+
 from .model import VirusOnNetwork, State, number_infected
 
 
@@ -52,8 +49,8 @@ def get_agents(source, target):
     return portrayal
 
 
-network = NetworkModule(network_portrayal, 500, 500)
-chart = ChartModule(
+network = mesa.visualization.NetworkModule(network_portrayal, 500, 500)
+chart = mesa.visualization.ChartModule(
     [
         {"Label": "Infected", "Color": "#FF0000"},
         {"Label": "Susceptible", "Color": "#008000"},
@@ -62,7 +59,7 @@ def get_agents(source, target):
 )
 
 
-class MyTextElement(TextElement):
+class MyTextElement(mesa.visualization.TextElement):
     def render(self, model):
         ratio = model.resistant_susceptible_ratio()
         ratio_text = "&infin;" if ratio is math.inf else f"{ratio:.2f}"
@@ -74,7 +71,7 @@ def render(self, model):
 
 
 model_params = {
-    "num_nodes": UserSettableParameter(
+    "num_nodes": mesa.visualization.UserSettableParameter(
         "slider",
         "Number of agents",
         10,
@@ -83,10 +80,10 @@ def render(self, model):
         1,
         description="Choose how many agents to include in the model",
     ),
-    "avg_node_degree": UserSettableParameter(
+    "avg_node_degree": mesa.visualization.UserSettableParameter(
         "slider", "Avg Node Degree", 3, 3, 8, 1, description="Avg Node Degree"
     ),
-    "initial_outbreak_size": UserSettableParameter(
+    "initial_outbreak_size": mesa.visualization.UserSettableParameter(
         "slider",
         "Initial Outbreak Size",
         1,
@@ -95,7 +92,7 @@ def render(self, model):
         1,
         description="Initial Outbreak Size",
     ),
-    "virus_spread_chance": UserSettableParameter(
+    "virus_spread_chance": mesa.visualization.UserSettableParameter(
         "slider",
         "Virus Spread Chance",
         0.4,
@@ -104,7 +101,7 @@ def render(self, model):
         0.1,
         description="Probability that susceptible neighbor will be infected",
     ),
-    "virus_check_frequency": UserSettableParameter(
+    "virus_check_frequency": mesa.visualization.UserSettableParameter(
         "slider",
         "Virus Check Frequency",
         0.4,
@@ -113,7 +110,7 @@ def render(self, model):
         0.1,
         description="Frequency the nodes check whether they are infected by " "a virus",
     ),
-    "recovery_chance": UserSettableParameter(
+    "recovery_chance": mesa.visualization.UserSettableParameter(
         "slider",
         "Recovery Chance",
         0.3,
@@ -122,7 +119,7 @@ def render(self, model):
         0.1,
         description="Probability that the virus will be removed",
     ),
-    "gain_resistance_chance": UserSettableParameter(
+    "gain_resistance_chance": mesa.visualization.UserSettableParameter(
         "slider",
         "Gain Resistance Chance",
         0.5,
@@ -134,7 +131,7 @@ def render(self, model):
     ),
 }
 
-server = ModularServer(
+server = mesa.visualization.ModularServer(
     VirusOnNetwork, [network, MyTextElement(), chart], "Virus Model", model_params
 )
 server.port = 8521