Repository: projectmesa/mesa
Pull request #2333 (2024-09-28 19:06:02)

Move cell movement into a descriptor
This PR encapsulates the addition to and removal from a cell of an agent in a `HasCell` mixin. It updates CellAgent to use this mixin, and adds an override of `agent.remove` to ensure agents remove themselves from the cell they are occupying when they are removed from the model. The main implication of this PR is that `CellAgent.move_to` is now redundant. I have thus removed it. In #2292, we might bring it back but with more powerful behavior.



### outdated based on discussion below
This PR is for discussion in relation to #2292. It shows how it might be possible to use composition for cell-related movement by encapsulating all the movement behavior into a descriptor. So instead of having a special `Agent.move_to` method, you can do `Agent.cell = new_cell`. You can also use whatever name you want instead of `cell`.

An additional benefit that I had not realized at first is that it cleans up the API a bit more as well. A user does not have to interact with the Cell class directly. She only needs to assign the Cell instance to whatever attribute is used for the descriptor. 

github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]          |
|-----------------|--------|-------------------------|----------------------------|
| BoltzmannWealth | small  | 🟢 -4.7% [-5.7%, -3.8%] | 🔵 -0.5% [-0.7%, -0.3%]    |
| BoltzmannWealth | large  | 🔵 -2.1% [-2.5%, -1.4%] | 🟢 -6.3% [-7.0%, -5.7%]    |
| Schelling       | small  | 🔵 +0.9% [+0.4%, +1.3%] | 🔴 +8.9% [+8.6%, +9.3%]    |
| Schelling       | large  | 🔵 +1.2% [+0.4%, +1.9%] | 🔵 +2.4% [-1.1%, +5.4%]    |
| WolfSheep       | small  | 🔴 +3.9% [+3.7%, +4.2%] | 🔴 +37.8% [+37.4%, +38.3%] |
| WolfSheep       | large  | 🔴 +4.6% [+3.7%, +5.5%] | 🔴 +44.7% [+41.4%, +48.3%] |
| BoidFlockers    | small  | 🔵 -0.4% [-0.8%, +0.0%] | 🔵 +0.5% [-0.3%, +1.4%]    |
| BoidFlockers    | large  | 🔵 -1.3% [-1.7%, -0.7%] | 🔵 -0.5% [-1.6%, +0.3%]    |

github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]          |
|-----------------|--------|-------------------------|----------------------------|
| BoltzmannWealth | small  | 🔵 -1.7% [-2.7%, -0.8%] | 🔵 -0.2% [-0.4%, +0.0%]    |
| BoltzmannWealth | large  | 🟢 -3.6% [-4.1%, -3.1%] | 🟢 -9.2% [-10.4%, -8.1%]   |
| Schelling       | small  | 🔵 -0.2% [-0.7%, +0.3%] | 🔵 -2.2% [-2.9%, -1.5%]    |
| Schelling       | large  | 🔵 +0.3% [-0.4%, +1.2%] | 🟢 -13.7% [-15.5%, -12.2%] |
| WolfSheep       | small  | 🔵 -0.8% [-1.1%, -0.4%] | 🔴 +7.9% [+7.4%, +8.4%]    |
| WolfSheep       | large  | 🔵 -0.2% [-1.6%, +0.8%] | 🔵 -0.1% [-2.4%, +2.5%]    |
| BoidFlockers    | small  | 🔵 -0.4% [-1.0%, +0.0%] | 🔵 -1.3% [-2.1%, -0.5%]    |
| BoidFlockers    | large  | 🔵 -0.6% [-1.2%, -0.1%] | 🔵 -2.1% [-2.6%, -1.5%]    |

EwoutH: I think I like this, it makes it much more explicit what's happening. It's like modifying `pos` directly, but for the new spaces.

One thing I'm curious about if `cell` has now a special status. How does the space know what Agent variable to lookup / modify? What is there are multiple spaces?

I will review in more detail tomorrow!
Corvince: Ha! Given enough time I think @quaquel will turn every mesa behaviour into a descriptor 😂

No, I really like this and this somehow put me over the edge to finally try to fully grok descriptors. 

Let me write in more detail when I have more time, but I think using just a property with a setter might be better
quaquel: > One thing I'm curious about if cell has now a special status. How does the space know what Agent variable to lookup / modify? What is there are multiple spaces?

A space does not need to know anything about Agents. A space has cells. A cell can have agents. Agents and cells need to be aware of each other. A space is primarily a container for cells, nothing more. Multiple spaces should not be a problem. You just again another cell to the attribute and your done. 


> No, I really like this and this somehow put me over the edge to finally try to fully grok descriptors.

The shortest summary is "reusable properties". 

> Let me write in more detail when I have more time, but I think using just a property with a setter might be better

I think you need a getter because the agent needs to know its cell for e.g., a random walk, but curious to see what more you have to say in detail. 


Given the fact that both @Corvince and @EwoutH like this idea, I'll try to update the tests asap to make sure this is fully covered again.
EwoutH: (I just looked at it from my phone, still need to do a thorough review)
Corvince: > I think you need a getter because the agent needs to know its cell for e.g., a random walk, but curious to see what more you have to say in detail.

My main objection would be that it might make some things easier if we have a clear "cell" definition. That is with a reusable descriptor we have the advantage of allowing several "cell" descriptor from several spaces. In practice I think this isn't very common. 

In the case of a "normal" property we could have the following:

```python
class HasCell:
    _cell: Cell | None = None

    @property
    def cell(self):
        return self._cell

    @cell.setter
    def cell(self, cell: Cell):
        ... # Your logic from CellDescriptor.__set__
```

And then we can have separate movement mixin
```python
class CellMovement:
    def move_to(self, cell: Cell):
         self.cell = cell

    ... # More movement functions
```

And finally the CellAgent:

```python
class CellAgent(Agent, HasCell, CellMovement):
    ....
```

In this case its unambigious and easy. We know we have to assign to "cell". If the cell descriptor can be assigned to any name, how would we know the right name?

Further advantages are that one could theoretically overwrite the `cell.setter` function. Also I think more people know properties in comparison to descriptors (although the former is just a higher level abstraction).
quaquel: > In this case its unambigious and easy. We know we have to assign to "cell". If the cell descriptor can be assigned to any name, how would we know the right name?

Not sure I fully understand your point, but a descriptor has the [`__set_name__` magic method](https://docs.python.org/3/howto/descriptor.html#customized-names). This magic method is called during class creation with the owning class and name of the attribute to which the descriptor is assigned. So, with a descriptor, you can always know the name of the attribute. 

Also, the `__set__` method in the descriptor would make a `move_to` method completely redundant.

I like your name `HasCell` as a possible name for the descriptor. It is probably better than my current uninspiring `CellDescriptor`
Corvince: Sorry if I wasn't clear. The movement mixin was just an example (although I would like to have the clearer move_to API, even if redundant. But this is irrelevant for my point here). The point is that for any mixin or rather anything wanting to operate on the cell property it needs to know the name. Of course the user knows the name, but our framework doesn't. So in the example move_to method the line `self.cell = cell` would not work if the name is not cell, but something else. 
quaquel: > Sorry if I wasn't clear. The movement mixin was just an example (although I would like to have the clearer move_to API, even if redundant. But this is irrelevant for my point here). The point is that for any mixin or rather anything wanting to operate on the cell property it needs to know the name. Of course the user knows the name, but our framework doesn't. So in the example move_to method the line self.cell = cell would not work if the name is not cell, but something else.

Ok, that makes sense. If we want to add any additional functionality like the MoveTo example, then yes we need to lock down the name of the attribute to which the current cell is assigned.
quaquel: A quick follow-up to @Corvince point about the naming of the attribute. There is, in fact, already a case like this. It is up to the user to actively remove an agent from the cell before calling `Agent.remove`. Moving this behavior into `CellAgent.remove` would make sense, but this method needs to know the attribute to which `Cell` is assigned. So the descriptor idea, which leaves the name of the attribute to the user, is less than ideal. 

This leaves the other idea: instead of having a `move_to` method solve everything via attribute assignment: `self.cell = self.cell.neighborhood.select_random_cell()`, which can be done with properties as shown by the HasCell mixin shown by @Corvince . I do like the clean API that results from solving cell movement this way and slightly prefer it over the current move_to method. 

Of course this still leaves room for more sophisticated movement methods that for example use a heading or something.
quaquel: > One thing I would find interesting is how this would work with multiple grids.

1. I think this question is only tangentially related to the point of this PR/discussion.
2. Do we want to support that out of the box with MESA? I personally don't think so. 
EwoutH: Just making sure we don’t box ourselves in / prohibit a use case unintentionally.
Corvince: If we would want to provide strong support, it would be relevant for this PR. But I also think we must not provide that out of the box. 

For me the space should be uniquely defined. A single object can't exist in two places at once. Everything else is a special case/abstraction. And it's probably special enough not to deal with it preemptively. But of course we must observe this and if we see common patterns we can abstract over that.

And by the way the cell space is super flexible with its connections. It's actually straight forward to have connections to the local spatial neighborhood, but also connect cells to their "network neighbors". Given that we now have "named connections" it's even easy to differentiate between both types. 
quaquel: > I very much agree with this as the underlying low-level implementation. It's also very much the intuitive way to assign a new position. So from a model or space point of view that's the way to do it.
> 
> But from an agent perspective I would also like to have a move_to function that semantically explicitly moves an agent from one location/cell to another. Even if that's just an assignment under the hood. But given the different perspectives I don't think this violates the "one obvious way to do" recommendation.


Ok, I'll update this PR accordingly. Not yet sure about keeping `move_to` but I am fine with keeping it for now.
Corvince: Just to clarify, I would not make it part of HasCell, but in a separate movement mixin, probably with other movement options, thus also out of scope of this PR 
Corvince: Yes, I see this as a perfect base case for the ideas discussed in #2292. After this is merged I would update #2292 to incorporate the  named coordinates and create a "Movement" mixin class. Then CellAgent can become just a class inherting from Agent, HasCell and Movement without any MRO related headaches (because of those, only Agent has an __init__ method). I think this will come out nicely, but once implemented we can see how it plays out by adapting some examples. 
github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]         |
|-----------------|--------|-------------------------|---------------------------|
| BoltzmannWealth | small  | 🔵 +0.5% [-0.4%, +1.3%] | 🔵 -0.2% [-0.2%, -0.1%]   |
| BoltzmannWealth | large  | 🔵 +0.4% [-0.2%, +1.1%] | 🔵 +2.1% [+1.3%, +3.0%]   |
| Schelling       | small  | 🔵 +1.9% [+1.5%, +2.3%] | 🔵 +3.5% [+2.7%, +4.3%]   |
| Schelling       | large  | 🔵 +2.3% [+1.2%, +3.5%] | 🔵 +7.0% [-0.2%, +14.2%]  |
| WolfSheep       | small  | 🔵 +2.5% [+1.8%, +3.2%] | 🔴 +13.4% [+9.5%, +17.5%] |
| WolfSheep       | large  | 🔵 -1.1% [-2.9%, +0.6%] | 🔴 +11.1% [+6.8%, +15.4%] |
| BoidFlockers    | small  | 🔵 +1.4% [+1.0%, +1.7%] | 🔵 +0.1% [-0.5%, +0.7%]   |
| BoidFlockers    | large  | 🔵 +1.6% [+0.9%, +2.2%] | 🔵 +0.3% [-0.1%, +0.7%]   |

github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]          |
|-----------------|--------|-------------------------|----------------------------|
| BoltzmannWealth | small  | 🔵 -0.2% [-1.0%, +0.5%] | 🔵 +0.3% [+0.2%, +0.4%]    |
| BoltzmannWealth | large  | 🔵 -0.6% [-1.2%, -0.1%] | 🔵 +0.1% [-0.4%, +0.7%]    |
| Schelling       | small  | 🔵 +1.4% [+0.9%, +1.8%] | 🔵 +1.7% [+1.4%, +2.0%]    |
| Schelling       | large  | 🔵 +0.9% [+0.4%, +1.4%] | 🔵 +1.5% [-0.0%, +2.9%]    |
| WolfSheep       | small  | 🔵 +1.8% [+1.4%, +2.1%] | 🔴 +12.1% [+7.8%, +16.6%]  |
| WolfSheep       | large  | 🔵 +0.8% [-0.4%, +2.3%] | 🔴 +17.8% [+15.2%, +20.4%] |
| BoidFlockers    | small  | 🔵 -2.0% [-2.3%, -1.7%] | 🔵 -0.7% [-1.5%, +0.1%]    |
| BoidFlockers    | large  | 🔵 -1.6% [-2.1%, -1.1%] | 🔵 -0.6% [-1.2%, -0.1%]    |

github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]          |
|-----------------|--------|-------------------------|----------------------------|
| BoltzmannWealth | small  | 🔵 -0.1% [-1.1%, +0.9%] | 🔵 +0.1% [-0.2%, +0.3%]    |
| BoltzmannWealth | large  | 🔵 +0.3% [-0.6%, +1.0%] | 🔵 +0.5% [-0.6%, +1.8%]    |
| Schelling       | small  | 🔵 +0.7% [+0.3%, +1.2%] | 🔵 +2.3% [+2.0%, +2.6%]    |
| Schelling       | large  | 🔵 +3.0% [+1.8%, +4.3%] | 🔵 +0.7% [-0.3%, +1.7%]    |
| WolfSheep       | small  | 🔵 +1.5% [+1.2%, +1.8%] | 🔴 +12.9% [+8.9%, +17.1%]  |
| WolfSheep       | large  | 🔵 +1.7% [+0.5%, +3.0%] | 🔴 +19.4% [+15.8%, +23.8%] |
| BoidFlockers    | small  | 🔵 -0.8% [-1.3%, -0.3%] | 🔵 +1.5% [+0.9%, +2.1%]    |
| BoidFlockers    | large  | 🔵 +0.2% [-0.6%, +0.9%] | 🔵 +0.8% [+0.1%, +1.3%]    |

Corvince: There are 6 ruff errors otherwise this looks good to go from my side (but a lot of things are based on my ideas so I would like another reviewer to approve)
rht: Disclaimer: I haven't read #2292. I'm not fully sure why having a `move_to` method makes it not compose-able, but from UX perspective, I prefer for it to be not removed. It's a lot easier to search and understand the codebase if there is an actual `move_to` verb, and having to grep for `self.cell = self.cell.*` for an agent movement (this is a relatively more complex expression regexp vs just a keyword `move_to`), making it as a too Mesa-specific quirk and not as modeling in general.

While @quaquel 's stance on accommodating LLM can be found at https://github.com/projectmesa/mesa/pull/2237#issuecomment-2304032222, I find that making it less likely for LLM to fall into a gotcha / to help it to reason better actually doubles as to help human reasoners as well.

Why can't we have both `move_to` and while having the descriptor under the hood?
quaquel: > Why can't we have both move_to and while having the descriptor under the hood?

The idea is that `move_to` will be made more powerful via #2292. So yes we'll likely end up with both. The key thing here is that the logic of adding and removing agents from cells is encapsulated in the attribute assignment rather than via a separate method. 

Patch:
diff --git a/benchmarks/Schelling/schelling.py b/benchmarks/Schelling/schelling.py
index 52519efc..2b6614da 100644
--- a/benchmarks/Schelling/schelling.py
+++ b/benchmarks/Schelling/schelling.py
@@ -1,13 +1,22 @@
 """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:
@@ -15,11 +24,13 @@ class SchellingAgent(CellAgent):
             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 @@ class SchellingAgent(CellAgent):
 
         # 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,6 +52,7 @@ class Schelling(Model):
 
     def __init__(
         self,
+        simulator=None,
         height=40,
         width=40,
         homophily=3,
@@ -48,11 +60,11 @@ class Schelling(Model):
         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 @@ class Schelling(Model):
             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 @@ class Schelling(Model):
         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."""
diff --git a/benchmarks/WolfSheep/wolf_sheep.py b/benchmarks/WolfSheep/wolf_sheep.py
index aaeed135..9fd71846 100644
--- a/benchmarks/WolfSheep/wolf_sheep.py
+++ b/benchmarks/WolfSheep/wolf_sheep.py
@@ -17,7 +17,7 @@ from mesa.experimental.devs import ABMSimulator
 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:
@@ -25,43 +25,36 @@ class Animal(CellAgent):
             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 @@ class Wolf(Animal):
             self.energy += self.energy_from_food
 
             # Kill the sheep
-            sheep_to_eat.die()
+            sheep_to_eat.remove()
 
 
 class GrassPatch(CellAgent):
@@ -112,19 +105,19 @@ class GrassPatch(CellAgent):
                 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 @@ class WolfSheep(Model):
                 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 @@ class WolfSheep(Model):
                 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."""
diff --git a/mesa/experimental/cell_space/cell.py b/mesa/experimental/cell_space/cell.py
index 51b2ad75..b7c002bb 100644
--- a/mesa/experimental/cell_space/cell.py
+++ b/mesa/experimental/cell_space/cell.py
@@ -117,7 +117,6 @@ class Cell:
 
         """
         self.agents.remove(agent)
-        agent.cell = None
 
     @property
     def is_empty(self) -> bool:
diff --git a/mesa/experimental/cell_space/cell_agent.py b/mesa/experimental/cell_space/cell_agent.py
index 76a7d16b..fea71b02 100644
--- a/mesa/experimental/cell_space/cell_agent.py
+++ b/mesa/experimental/cell_space/cell_agent.py
@@ -4,13 +4,36 @@ from __future__ import annotations
 
 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


Test patch:
diff --git a/tests/test_cell_space.py b/tests/test_cell_space.py
index 01420043..0dd3855b 100644
--- a/tests/test_cell_space.py
+++ b/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