Skip to content

Commit d63ce06

Browse files
authored
refactor: Simplify Schelling code (#222)
* refactor: Simplify Schelling code 1. Remove unused model attributes 2. Make `similar` calculation more natural language readable * Remove unused argument doc * Add type hints to agent class * refactor: Simplify self.running expression
1 parent 5739e84 commit d63ce06

File tree

1 file changed

+8
-16
lines changed

1 file changed

+8
-16
lines changed

examples/schelling/model.py

Lines changed: 8 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -6,24 +6,21 @@ class SchellingAgent(mesa.Agent):
66
Schelling segregation agent
77
"""
88

9-
def __init__(self, model, agent_type):
9+
def __init__(self, model: mesa.Model, agent_type: int) -> None:
1010
"""
1111
Create a new Schelling agent.
1212
1313
Args:
14-
x, y: Agent initial location.
1514
agent_type: Indicator for the agent's type (minority=1, majority=0)
1615
"""
1716
super().__init__(model)
1817
self.type = agent_type
1918

20-
def step(self):
21-
similar = 0
22-
for neighbor in self.model.grid.iter_neighbors(
19+
def step(self) -> None:
20+
neighbors = self.model.grid.iter_neighbors(
2321
self.pos, moore=True, radius=self.model.radius
24-
):
25-
if neighbor.type == self.type:
26-
similar += 1
22+
)
23+
similar = sum(1 for neighbor in neighbors if neighbor.type == self.type)
2724

2825
# If unhappy, move:
2926
if similar < self.model.homophily:
@@ -60,10 +57,6 @@ def __init__(
6057
"""
6158

6259
super().__init__(seed=seed)
63-
self.height = height
64-
self.width = width
65-
self.density = density
66-
self.minority_pc = minority_pc
6760
self.homophily = homophily
6861
self.radius = radius
6962

@@ -79,8 +72,8 @@ def __init__(
7972
# the coordinates of a cell as well as
8073
# its contents. (coord_iter)
8174
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
75+
if self.random.random() < density:
76+
agent_type = 1 if self.random.random() < minority_pc else 0
8477
agent = SchellingAgent(self, agent_type)
8578
self.grid.place_agent(agent, pos)
8679

@@ -95,5 +88,4 @@ def step(self):
9588

9689
self.datacollector.collect(self)
9790

98-
if self.happy == len(self.agents):
99-
self.running = False
91+
self.running = self.happy != len(self.agents)

0 commit comments

Comments
 (0)