Skip to content

Commit f4c3da6

Browse files
committed
refactor: Simplify Schelling code
Ported from projectmesa/mesa-examples#222.
1 parent e255a2d commit f4c3da6

File tree

1 file changed

+3
-7
lines changed

1 file changed

+3
-7
lines changed

benchmarks/Schelling/schelling.py

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -34,11 +34,8 @@ def __init__(
3434

3535
def step(self):
3636
"""Run one step of the agent."""
37-
similar = 0
38-
neighborhood = self.cell.get_neighborhood(radius=self.radius)
39-
for neighbor in neighborhood.agents:
40-
if neighbor.type == self.type:
41-
similar += 1
37+
neighbors = self.cell.get_neighborhood(radius=self.radius).agents
38+
similar = sum(1 for neighbor in neighbors if neighbor.type == self.type)
4239

4340
# If unhappy, move:
4441
if similar < self.homophily:
@@ -76,7 +73,6 @@ def __init__(
7673
"""
7774
super().__init__(seed=seed)
7875
self.simulator = simulator
79-
self.minority_pc = minority_pc
8076
self.happy = 0
8177

8278
self.grid = OrthogonalMooreGrid(
@@ -92,7 +88,7 @@ def __init__(
9288
# its contents. (coord_iter)
9389
for cell in self.grid:
9490
if self.random.random() < density:
95-
agent_type = 1 if self.random.random() < self.minority_pc else 0
91+
agent_type = 1 if self.random.random() < minority_pc else 0
9692
SchellingAgent(self, agent_type, radius, homophily, cell)
9793

9894
def step(self):

0 commit comments

Comments
 (0)