Move cell movement into a descriptor (#2333)

Author: quaquel

benchmarks/Schelling/schelling.py

@@ -1,25 +1,36 @@
 """Schelling separation for performance benchmarking."""
 
+from __future__ import annotations
+
 from mesa import Model
-from mesa.experimental.cell_space import CellAgent, OrthogonalMooreGrid
+from mesa.experimental.cell_space import Cell, CellAgent, OrthogonalMooreGrid
 
 
 class SchellingAgent(CellAgent):
     """Schelling segregation agent."""
 
-    def __init__(self, model, agent_type, radius, homophily):
+    def __init__(
+        self,
+        model: Schelling,
+        agent_type: int,
+        radius: int,
+        homophily: float,
+        cell: Cell,
+    ):
         """Create a new Schelling agent.
 
         Args:
             model: model instance
             agent_type: type of agent (minority=1, majority=0)
             radius: size of neighborhood of agent
             homophily: fraction of neighbors of the same type that triggers movement
+            cell: the cell in which the agent is located
         """
         super().__init__(model)
         self.type = agent_type
         self.radius = radius
         self.homophily = homophily
+        self.cell = cell
 
     def step(self):
         """Run one step of the agent."""
@@ -31,7 +42,7 @@ def step(self):
 
         # If unhappy, move:
         if similar < self.homophily:
-            self.move_to(self.model.grid.select_random_empty_cell())
+            self.cell = self.model.grid.select_random_empty_cell()
         else:
             self.model.happy += 1
 
@@ -41,18 +52,19 @@ class Schelling(Model):
 
     def __init__(
         self,
+        simulator=None,
         height=40,
         width=40,
         homophily=3,
         radius=1,
         density=0.8,
         minority_pc=0.5,
         seed=None,
-        simulator=None,
     ):
         """Create a new Schelling model.
 
         Args:
+            simulator: simulator instance
             height: height of the grid
             width: width of the grid
             homophily: Minimum number of agents of same class needed to be happy
@@ -63,8 +75,9 @@ def __init__(
             simulator: a simulator instance
         """
         super().__init__(seed=seed)
-        self.minority_pc = minority_pc
         self.simulator = simulator
+        self.minority_pc = minority_pc
+        self.happy = 0
 
         self.grid = OrthogonalMooreGrid(
             [height, width],
@@ -80,8 +93,7 @@ def __init__(
         for cell in self.grid:
             if self.random.random() < density:
                 agent_type = 1 if self.random.random() < self.minority_pc else 0
-                agent = SchellingAgent(self, agent_type, radius, homophily)
-                agent.move_to(cell)
+                SchellingAgent(self, agent_type, radius, homophily, cell)
 
     def step(self):
         """Run one step of the model."""

benchmarks/WolfSheep/wolf_sheep.py

@@ -17,51 +17,44 @@
 class Animal(CellAgent):
     """The base animal class."""
 
-    def __init__(self, model, energy, p_reproduce, energy_from_food):
+    def __init__(self, model, energy, p_reproduce, energy_from_food, cell):
         """Initializes an animal.
 
         Args:
             model: a model instance
             energy: starting amount of energy
             p_reproduce: probability of sexless reproduction
             energy_from_food: energy obtained from 1 unit of food
+            cell: the cell in which the animal starts
         """
         super().__init__(model)
         self.energy = energy
         self.p_reproduce = p_reproduce
         self.energy_from_food = energy_from_food
-
-    def random_move(self):
-        """Move to a random neighboring cell."""
-        self.move_to(self.cell.neighborhood.select_random_cell())
+        self.cell = cell
 
     def spawn_offspring(self):
         """Create offspring."""
         self.energy /= 2
-        offspring = self.__class__(
+        self.__class__(
             self.model,
             self.energy,
             self.p_reproduce,
             self.energy_from_food,
+            self.cell,
         )
-        offspring.move_to(self.cell)
 
     def feed(self): ...  # noqa: D102
 
-    def die(self):
-        """Die."""
-        self.cell.remove_agent(self)
-        self.remove()
-
     def step(self):
         """One step of the agent."""
-        self.random_move()
+        self.cell = self.cell.neighborhood.select_random_cell()
         self.energy -= 1
 
         self.feed()
 
         if self.energy < 0:
-            self.die()
+            self.remove()
         elif self.random.random() < self.p_reproduce:
             self.spawn_offspring()
 
@@ -91,7 +84,7 @@ def feed(self):
             self.energy += self.energy_from_food
 
             # Kill the sheep
-            sheep_to_eat.die()
+            sheep_to_eat.remove()
 
 
 class GrassPatch(CellAgent):
@@ -112,19 +105,19 @@ def fully_grown(self, value: bool) -> None:
                 function_args=[self, "fully_grown", True],
             )
 
-    def __init__(self, model, fully_grown, countdown, grass_regrowth_time):
+    def __init__(self, model, countdown, grass_regrowth_time, cell):
         """Creates a new patch of grass.
 
         Args:
             model: a model instance
-            fully_grown: (boolean) Whether the patch of grass is fully grown or not
             countdown: Time for the patch of grass to be fully grown again
             grass_regrowth_time : time to fully regrow grass
+            cell: the cell to which the patch of grass belongs
         """
-        # TODO:: fully grown can just be an int --> so one less param (i.e. countdown)
         super().__init__(model)
-        self._fully_grown = fully_grown
+        self._fully_grown = True if countdown == 0 else False  # Noqa: SIM210
         self.grass_regrowth_time = grass_regrowth_time
+        self.cell = cell
 
         if not self.fully_grown:
             self.model.simulator.schedule_event_relative(
@@ -191,13 +184,7 @@ def __init__(
                 self.random.randrange(self.height),
             )
             energy = self.random.randrange(2 * sheep_gain_from_food)
-            sheep = Sheep(
-                self,
-                energy,
-                sheep_reproduce,
-                sheep_gain_from_food,
-            )
-            sheep.move_to(self.grid[pos])
+            Sheep(self, energy, sheep_reproduce, sheep_gain_from_food, self.grid[pos])
 
         # Create wolves
         for _ in range(self.initial_wolves):
@@ -206,24 +193,14 @@ def __init__(
                 self.random.randrange(self.height),
             )
             energy = self.random.randrange(2 * wolf_gain_from_food)
-            wolf = Wolf(
-                self,
-                energy,
-                wolf_reproduce,
-                wolf_gain_from_food,
-            )
-            wolf.move_to(self.grid[pos])
+            Wolf(self, energy, wolf_reproduce, wolf_gain_from_food, self.grid[pos])
 
         # Create grass patches
         possibly_fully_grown = [True, False]
         for cell in self.grid:
             fully_grown = self.random.choice(possibly_fully_grown)
-            if fully_grown:
-                countdown = grass_regrowth_time
-            else:
-                countdown = self.random.randrange(grass_regrowth_time)
-            patch = GrassPatch(self, fully_grown, countdown, grass_regrowth_time)
-            patch.move_to(cell)
+            countdown = 0 if fully_grown else self.random.randrange(grass_regrowth_time)
+            GrassPatch(self, countdown, grass_regrowth_time, cell)
 
     def step(self):
         """Run one step of the model."""

mesa/experimental/cell_space/cell.py

@@ -117,7 +117,6 @@ def remove_agent(self, agent: CellAgent) -> None:
 
         """
         self.agents.remove(agent)
-        agent.cell = None
 
     @property
     def is_empty(self) -> bool:

mesa/experimental/cell_space/cell_agent.py

@@ -4,13 +4,36 @@
 
 from typing import TYPE_CHECKING
 
-from mesa import Agent, Model
+from mesa import Agent
 
 if TYPE_CHECKING:
     from mesa.experimental.cell_space.cell import Cell
 
 
-class CellAgent(Agent):
+class HasCell:
+    """Descriptor for cell movement behavior."""
+
+    _mesa_cell: Cell = None
+
+    @property
+    def cell(self) -> Cell | None:  # noqa: D102
+        return self._mesa_cell
+
+    @cell.setter
+    def cell(self, cell: Cell | None) -> None:
+        # remove from current cell
+        if self.cell is not None:
+            self.cell.remove_agent(self)
+
+        # update private attribute
+        self._mesa_cell = cell
+
+        # add to new cell
+        if cell is not None:
+            cell.add_agent(self)
+
+
+class CellAgent(Agent, HasCell):
     """Cell Agent is an extension of the Agent class and adds behavior for moving in discrete spaces.
 
     Attributes:
@@ -20,23 +43,7 @@ class CellAgent(Agent):
         cell: (Cell | None): the cell which the agent occupies
     """
 
-    def __init__(self, model: Model) -> None:
-        """Create a new agent.
-
-        Args:
-            model (Model): The model instance in which the agent exists.
-        """
-        super().__init__(model)
-        self.cell: Cell | None = None
-
-    def move_to(self, cell) -> None:
-        """Move agent to cell.
-
-        Args:
-            cell: cell to which agent is to move
-
-        """
-        if self.cell is not None:
-            self.cell.remove_agent(self)
-        self.cell = cell
-        cell.add_agent(self)
+    def remove(self):
+        """Remove the agent from the model."""
+        super().remove()
+        self.cell = None  # ensures that we are also removed from cell

tests/test_cell_space.py

@@ -524,3 +524,32 @@ def test_cell_collection():
 
     cells = collection.select()
     assert len(cells) == len(collection)
+
+
+def test_cell_agent():  # noqa: D103
+    cell1 = Cell((1,), capacity=None, random=random.Random())
+    cell2 = Cell((2,), capacity=None, random=random.Random())
+
+    # connect
+    # add_agent
+    model = Model()
+    agent = CellAgent(model)
+
+    agent.cell = cell1
+    assert agent in cell1.agents
+
+    agent.cell = None
+    assert agent not in cell1.agents
+
+    agent.cell = cell2
+    assert agent not in cell1.agents
+    assert agent in cell2.agents
+
+    agent.cell = cell1
+    assert agent in cell1.agents
+    assert agent not in cell2.agents
+
+    agent.remove()
+    assert agent not in model._all_agents
+    assert agent not in cell1.agents
+    assert agent not in cell2.agents