Merge remote-tracking branch 'upstream/main' into new_spaces

Author: quaquel

examples/schelling/model.py

@@ -7,22 +7,19 @@ class SchellingAgent(CellAgent):
     Schelling segregation agent
     """
 
-    def __init__(self, model, agent_type):
+    def __init__(self, model: mesa.Model, agent_type: int) -> None:
         """
         Create a new Schelling agent.
 
         Args:
-           x, y: Agent initial location.
            agent_type: Indicator for the agent's type (minority=1, majority=0)
         """
         super().__init__(model)
         self.type = agent_type
 
     def step(self):
-        similar = 0
-        for neighbor in self.cell.get_neighborhood(radius=self.model.radius).agents:
-            if neighbor.type == self.type:
-                similar += 1
+        neighbors = self.cell.get_neighborhood(radius=self.model.radius).agents
+        similar = len([neighbor for neighbor in neighbors if neighbor.type == self.type])
 
         # If unhappy, move:
         if similar < self.model.homophily:
@@ -59,10 +56,6 @@ def __init__(
         """
 
         super().__init__(seed=seed)
-        self.height = height
-        self.width = width
-        self.density = density
-        self.minority_pc = minority_pc
         self.homophily = homophily
         self.radius = radius
 
@@ -78,8 +71,8 @@ def __init__(
         # the coordinates of a cell as well as
         # its contents. (coord_iter)
         for cell in self.grid.all_cells:
-            if self.random.random() < self.density:
-                agent_type = 1 if self.random.random() < self.minority_pc else 0
+            if self.random.random() < density:
+                agent_type = 1 if self.random.random() < minority_pc else 0
                 agent = SchellingAgent(self, agent_type)
                 agent.cell = cell
 
@@ -94,5 +87,4 @@ def step(self):
 
         self.datacollector.collect(self)
 
-        if self.happy == len(self.agents):
-            self.running = False
+        self.running = self.happy != len(self.agents)