examples: Convert boltzman_wealth_model_network to simple namespace

arbehr · May 18, 2022 · 95a36e7 · 95a36e7
1 parent aaf5081
commit 95a36e7
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 19 deletions.
diff --git a/examples/boltzmann_wealth_model_network/boltzmann_wealth_model_network/model.py b/examples/boltzmann_wealth_model_network/boltzmann_wealth_model_network/model.py
@@ -1,10 +1,6 @@
-from mesa import Agent, Model
-from mesa.time import RandomActivation
-from mesa.datacollection import DataCollector
+import mesa
 import networkx as nx
 
-from mesa.space import NetworkGrid
-
 
 def compute_gini(model):
     agent_wealths = [agent.wealth for agent in model.schedule.agents]
@@ -14,17 +10,17 @@ def compute_gini(model):
     return 1 + (1 / N) - 2 * B
 
 
-class BoltzmannWealthModelNetwork(Model):
+class BoltzmannWealthModelNetwork(mesa.Model):
     """A model with some number of agents."""
 
     def __init__(self, num_agents=7, num_nodes=10):
 
         self.num_agents = num_agents
         self.num_nodes = num_nodes if num_nodes >= self.num_agents else self.num_agents
         self.G = nx.erdos_renyi_graph(n=self.num_nodes, p=0.5)
-        self.grid = NetworkGrid(self.G)
-        self.schedule = RandomActivation(self)
-        self.datacollector = DataCollector(
+        self.grid = mesa.space.NetworkGrid(self.G)
+        self.schedule = mesa.time.RandomActivation(self)
+        self.datacollector = mesa.DataCollector(
             model_reporters={"Gini": compute_gini},
             agent_reporters={"Wealth": lambda _: _.wealth},
         )
@@ -51,7 +47,7 @@ def run_model(self, n):
             self.step()
 
 
-class MoneyAgent(Agent):
+class MoneyAgent(mesa.Agent):
     """An agent with fixed initial wealth."""
 
     def __init__(self, unique_id, model):

diff --git a/examples/boltzmann_wealth_model_network/boltzmann_wealth_model_network/server.py b/examples/boltzmann_wealth_model_network/boltzmann_wealth_model_network/server.py
@@ -1,7 +1,5 @@
-from mesa.visualization.ModularVisualization import ModularServer
-from mesa.visualization.UserParam import UserSettableParameter
-from mesa.visualization.modules import ChartModule
-from mesa.visualization.modules import NetworkModule
+import mesa
+
 from .model import BoltzmannWealthModelNetwork
 
 
@@ -29,13 +27,13 @@ def network_portrayal(G):
     return portrayal
 
 
-grid = NetworkModule(network_portrayal, 500, 500)
-chart = ChartModule(
+grid = mesa.visualization.NetworkModule(network_portrayal, 500, 500)
+chart = mesa.visualization.ChartModule(
     [{"Label": "Gini", "Color": "Black"}], data_collector_name="datacollector"
 )
 
 model_params = {
-    "num_agents": UserSettableParameter(
+    "num_agents": mesa.visualization.UserSettableParameter(
         "slider",
         "Number of agents",
         7,
@@ -44,7 +42,7 @@ def network_portrayal(G):
         1,
         description="Choose how many agents to include in the model",
     ),
-    "num_nodes": UserSettableParameter(
+    "num_nodes": mesa.visualization.UserSettableParameter(
         "slider",
         "Number of nodes",
         10,
@@ -56,7 +54,7 @@ def network_portrayal(G):
     ),
 }
 
-server = ModularServer(
+server = mesa.visualization.ModularServer(
     BoltzmannWealthModelNetwork, [grid, chart], "Money Model", model_params
 )
 server.port = 8521