Move Schelling to discrete space

mesa/examples/basic/schelling/agents.py

@@ -1,23 +1,28 @@
-from mesa import Agent
+from mesa.discrete_space import CellAgent
 
 
-class SchellingAgent(Agent):
+class SchellingAgent(CellAgent):
     """Schelling segregation agent."""
 
-    def __init__(self, model, agent_type: int) -> None:
+    def __init__(
+        self, model, cell, agent_type: int, homophily: float = 0.4, radius: int = 1
+    ) -> None:
         """Create a new Schelling agent.
         Args:
             model: The model instance the agent belongs to
             agent_type: Indicator for the agent's type (minority=1, majority=0)
+            homophily: Minimum number of similar neighbors needed for happiness
+            radius: Search radius for checking neighbor similarity
         """
         super().__init__(model)
+        self.cell = cell
         self.type = agent_type
+        self.homophily = homophily
+        self.radius = radius
 
     def step(self) -> None:
         """Determine if agent is happy and move if necessary."""
-        neighbors = self.model.grid.get_neighbors(
-            self.pos, moore=True, radius=self.model.radius
-        )
+        neighbors = list(self.cell.get_neighborhood(radius=self.radius).agents)
 
         # Count similar neighbors
         similar_neighbors = len([n for n in neighbors if n.type == self.type])
@@ -30,7 +35,7 @@ def step(self) -> None:
             similarity_fraction = 0.0
 
         # Move if unhappy
-        if similarity_fraction < self.model.homophily:
-            self.model.grid.move_to_empty(self)
+        if similarity_fraction < self.homophily:
+            self.cell = self.model.grid.select_random_empty_cell()
         else:
             self.model.happy += 1

mesa/examples/basic/schelling/model.py

@@ -1,7 +1,7 @@
 from mesa import Model
 from mesa.datacollection import DataCollector
+from mesa.discrete_space import OrthogonalMooreGrid
 from mesa.examples.basic.schelling.agents import SchellingAgent
-from mesa.space import SingleGrid
 
 
 class Schelling(Model):
@@ -31,15 +31,11 @@ def __init__(
         super().__init__(seed=seed)
 
         # Model parameters
-        self.height = height
-        self.width = width
         self.density = density
         self.minority_pc = minority_pc
-        self.homophily = homophily
-        self.radius = radius
 
         # Initialize grid
-        self.grid = SingleGrid(width, height, torus=True)
+        self.grid = OrthogonalMooreGrid((width, height), random=self.random, capacity=1)
 
         # Track happiness
         self.happy = 0
@@ -64,11 +60,12 @@ def __init__(
         )
 
         # Create agents and place them on the grid
-        for _, pos in self.grid.coord_iter():
+        for cell in self.grid.all_cells:
             if self.random.random() < self.density:
                 agent_type = 1 if self.random.random() < minority_pc else 0
-                agent = SchellingAgent(self, agent_type)
-                self.grid.place_agent(agent, pos)
+                SchellingAgent(
+                    self, cell, agent_type, homophily=homophily, radius=radius
+                )
 
         # Collect initial state
         self.datacollector.collect(self)