From e7023dcdff72c62dd50919ad784aaeecaed4d616 Mon Sep 17 00:00:00 2001
From: rht <rhtbot@protonmail.com>
Date: Sat, 6 Jan 2024 22:20:49 -0500
Subject: [PATCH] Boltzmann wealth model: Refactor code

---
 .../model.py                                  | 23 ++++++++++---------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/examples/boltzmann_wealth_model_experimental/model.py b/examples/boltzmann_wealth_model_experimental/model.py
index acfb3fc0..1660243c 100644
--- 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,17 +51,20 @@ 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
@@ -70,5 +72,4 @@ def give_money(self):
 
     def step(self):
         self.move()
-        if self.wealth > 0:
-            self.give_money()
+        self.give_money()