Skip to content

Commit 7d3c22e

Browse files
committed
refactor: Simplify Schelling code
1. Remove unused model attributes 2. Make `similar` calculation more natural language readable
1 parent 5739e84 commit 7d3c22e

File tree

1 file changed

+5
-11
lines changed

1 file changed

+5
-11
lines changed

examples/schelling/model.py

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,10 @@ def __init__(self, model, agent_type):
1818
self.type = agent_type
1919

2020
def step(self):
21-
similar = 0
22-
for neighbor in self.model.grid.iter_neighbors(
21+
neighbors = self.model.grid.iter_neighbors(
2322
self.pos, moore=True, radius=self.model.radius
24-
):
25-
if neighbor.type == self.type:
26-
similar += 1
23+
)
24+
similar = sum([1 for neighbor in neighbors if neighbor.type == self.type])
2725

2826
# If unhappy, move:
2927
if similar < self.model.homophily:
@@ -60,10 +58,6 @@ def __init__(
6058
"""
6159

6260
super().__init__(seed=seed)
63-
self.height = height
64-
self.width = width
65-
self.density = density
66-
self.minority_pc = minority_pc
6761
self.homophily = homophily
6862
self.radius = radius
6963

@@ -79,8 +73,8 @@ def __init__(
7973
# the coordinates of a cell as well as
8074
# its contents. (coord_iter)
8175
for _, pos in self.grid.coord_iter():
82-
if self.random.random() < self.density:
83-
agent_type = 1 if self.random.random() < self.minority_pc else 0
76+
if self.random.random() < density:
77+
agent_type = 1 if self.random.random() < minority_pc else 0
8478
agent = SchellingAgent(self, agent_type)
8579
self.grid.place_agent(agent, pos)
8680

0 commit comments

Comments
 (0)