diff --git a/examples/boltzmann_wealth_model_experimental/model.py b/examples/boltzmann_wealth_model_experimental/model.py index 11a3e958..2d7714c1 100644 --- a/examples/boltzmann_wealth_model_experimental/model.py +++ b/examples/boltzmann_wealth_model_experimental/model.py @@ -1,10 +1,10 @@ import mesa +from mesa.experimental import DataCollector def compute_gini(model): - agent_wealths = [agent.wealth for agent in model.schedule.agents] - x = sorted(agent_wealths) - N = model.num_agents + x = sorted(model.agents.get("wealth")) + N = len(x) B = sum(xi * (N - i) for i, xi in enumerate(x)) / (N * sum(x)) return 1 + (1 / N) - 2 * B @@ -22,9 +22,6 @@ def __init__(self, N=100, width=10, height=10): self.num_agents = N self.grid = mesa.space.MultiGrid(width, height, True) self.schedule = mesa.time.RandomActivation(self) - self.datacollector = mesa.DataCollector( - model_reporters={"Gini": compute_gini}, agent_reporters={"Wealth": "wealth"} - ) # Create agents for i in range(self.num_agents): a = MoneyAgent(i, self) @@ -34,13 +31,19 @@ def __init__(self, N=100, width=10, height=10): y = self.random.randrange(self.grid.height) self.grid.place_agent(a, (x, y)) - self.running = True - self.datacollector.collect(self) + self.datacollector = DataCollector( + self, + { + "model": {"Gini": compute_gini}, + "agents": {"Wealth": "wealth"}, + }, + ) + self.datacollector.collect() def step(self): self.schedule.step() # collect data - self.datacollector.collect(self) + self.datacollector.collect() def run_model(self, n): for i in range(n):