refactor: Simplify Schelling code (#2353)

* refactor: Simplify Schelling code Ported from projectmesa/mesa-examples#222. * Update benchmarks/Schelling/schelling.py Co-authored-by: Ewout ter Hoeven <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Jan Kwakkel <[email protected]> Co-authored-by: Ewout ter Hoeven <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
projectmesa · Oct 18, 2024 · dea931d · dea931d
1 parent 4ed0d6d
commit dea931d
Showing 1 changed file with 5 additions and 7 deletions.
diff --git a/benchmarks/Schelling/schelling.py b/benchmarks/Schelling/schelling.py
@@ -34,11 +34,10 @@ def __init__(
 
     def step(self):
         """Run one step of the agent."""
-        similar = 0
-        neighborhood = self.cell.get_neighborhood(radius=self.radius)
-        for neighbor in neighborhood.agents:
-            if neighbor.type == self.type:
-                similar += 1
+        neighbors = self.cell.get_neighborhood(radius=self.radius).agents
+        similar = len(
+            [neighbor for neighbor in neighbors if neighbor.type == self.type]
+        )
 
         # If unhappy, move:
         if similar < self.homophily:
@@ -76,7 +75,6 @@ def __init__(
         """
         super().__init__(seed=seed)
         self.simulator = simulator
-        self.minority_pc = minority_pc
         self.happy = 0
 
         self.grid = OrthogonalMooreGrid(
@@ -92,7 +90,7 @@ def __init__(
         # its contents. (coord_iter)
         for cell in self.grid:
             if self.random.random() < density:
-                agent_type = 1 if self.random.random() < self.minority_pc else 0
+                agent_type = 1 if self.random.random() < minority_pc else 0
                 SchellingAgent(self, agent_type, radius, homophily, cell)
 
     def step(self):