Boltzmann wealth model: Refactor code

projectmesa · Jan 7, 2024 · e7023dc · e7023dc
1 parent 0bc6097
commit e7023dc
Showing 1 changed file with 12 additions and 11 deletions.
diff --git a/examples/boltzmann_wealth_model_experimental/model.py b/examples/boltzmann_wealth_model_experimental/model.py
@@ -18,6 +18,7 @@ class BoltzmannWealthModel(mesa.Model):
     """
 
     def __init__(self, N=100, width=10, height=10):
+        super().__init__()
         self.num_agents = N
         self.grid = mesa.space.MultiGrid(width, height, True)
         self.datacollector = mesa.DataCollector(
@@ -31,12 +32,10 @@ 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)
 
     def step(self):
         self.agents.shuffle().do("step")
-        # collect data
         self.datacollector.collect(self)
 
     def run_model(self, n):
@@ -52,23 +51,25 @@ def __init__(self, unique_id, model):
         self.wealth = 1
 
     def move(self):
-        possible_steps = self.model.grid.get_neighborhood(
+        possible_positions = self.model.grid.get_neighborhood(
             self.pos, moore=True, include_center=False
         )
-        new_position = self.random.choice(possible_steps)
-        self.model.grid.move_agent(self, new_position)
+        self.model.grid.move_agent_to_one_of(possible_positions)
 
     def give_money(self):
-        cellmates = self.model.grid.get_cell_list_contents([self.pos])
-        cellmates.pop(
-            cellmates.index(self)
-        )  # Ensure agent is not giving money to itself
+        if self.wealth <= 0:
+            return
+        cellmates = [
+            c
+            for c in self.model.grid.get_cell_list_contents([self.pos])
+            # Ensure agent is not giving money to itself
+            if c is not self
+        ]
         if len(cellmates) > 0:
             other = self.random.choice(cellmates)
             other.wealth += 1
             self.wealth -= 1
 
     def step(self):
         self.move()
-        if self.wealth > 0:
-            self.give_money()
+        self.give_money()