Repository: projectmesa/mesa
Pull request #2332 (2024-09-28 18:47:02)

viz: stop running and disable buttons when model.running is False
This PR stops calling `model.step()` from the frontend when `model.running` is `False`. Play and step buttons are also disabled.

github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]       |
|-----------------|--------|-------------------------|-------------------------|
| BoltzmannWealth | small  | 🔵 -1.4% [-2.2%, -0.5%] | 🔵 -0.3% [-0.5%, -0.1%] |
| BoltzmannWealth | large  | 🔵 -0.5% [-0.8%, -0.1%] | 🔵 -0.0% [-0.6%, +0.6%] |
| Schelling       | small  | 🔵 -0.3% [-0.6%, +0.1%] | 🔵 -0.0% [-0.3%, +0.2%] |
| Schelling       | large  | 🔵 -0.3% [-0.7%, +0.1%] | 🔵 -0.3% [-0.9%, +0.2%] |
| WolfSheep       | small  | 🔵 +1.7% [+1.5%, +1.9%] | 🔵 +1.3% [+1.1%, +1.4%] |
| WolfSheep       | large  | 🔵 +1.7% [+1.1%, +2.3%] | 🔵 +1.0% [+0.5%, +1.6%] |
| BoidFlockers    | small  | 🔵 +0.6% [+0.1%, +1.2%] | 🔵 -0.1% [-0.8%, +0.5%] |
| BoidFlockers    | large  | 🔵 +1.2% [+0.3%, +1.9%] | 🔵 -0.4% [-0.9%, +0.1%] |

Corvince: Looks good to me!
EwoutH: How/when should/can I use this?

Patch:
diff --git a/benchmarks/BoltzmannWealth/boltzmann_wealth.py b/benchmarks/BoltzmannWealth/boltzmann_wealth.py
index 5423eb82..2d863efc 100644
--- a/benchmarks/BoltzmannWealth/boltzmann_wealth.py
+++ b/benchmarks/BoltzmannWealth/boltzmann_wealth.py
@@ -40,7 +40,7 @@ class BoltzmannWealth(mesa.Model):
             width: the width of the grid
             height: the height of the grid
         """
-        super().__init__(seed)
+        super().__init__(seed=seed)
         self.num_agents = n
         self.grid = mesa.space.MultiGrid(width, height, True)
         self.datacollector = mesa.DataCollector(
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..8c7fb9bb 100644
--- a/mesa/experimental/cell_space/cell.py
+++ b/mesa/experimental/cell_space/cell.py
@@ -2,11 +2,13 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from functools import cache, cached_property
 from random import Random
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from mesa.experimental.cell_space.cell_collection import CellCollection
+from mesa.space import PropertyLayer
 
 if TYPE_CHECKING:
     from mesa.agent import Agent
@@ -34,6 +36,7 @@ class Cell:
         "capacity",
         "properties",
         "random",
+        "_mesa_property_layers",
         "__dict__",
     ]
 
@@ -69,6 +72,7 @@ class Cell:
         self.capacity: int = capacity
         self.properties: dict[Coordinate, object] = {}
         self.random = random
+        self._mesa_property_layers: dict[str, PropertyLayer] = {}
 
     def connect(self, other: Cell, key: Coordinate | None = None) -> None:
         """Connects this cell to another cell.
@@ -117,7 +121,6 @@ class Cell:
 
         """
         self.agents.remove(agent)
-        agent.cell = None
 
     @property
     def is_empty(self) -> bool:
@@ -191,3 +194,20 @@ class Cell:
             if not include_center:
                 neighborhood.pop(self, None)
             return neighborhood
+
+    # PropertyLayer methods
+    def get_property(self, property_name: str) -> Any:
+        """Get the value of a property."""
+        return self._mesa_property_layers[property_name].data[self.coordinate]
+
+    def set_property(self, property_name: str, value: Any):
+        """Set the value of a property."""
+        self._mesa_property_layers[property_name].set_cell(self.coordinate, value)
+
+    def modify_property(
+        self, property_name: str, operation: Callable, value: Any = None
+    ):
+        """Modify the value of a property."""
+        self._mesa_property_layers[property_name].modify_cell(
+            self.coordinate, operation, value
+        )
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
diff --git a/mesa/experimental/cell_space/discrete_space.py b/mesa/experimental/cell_space/discrete_space.py
index 0d3af169..1cfaedeb 100644
--- a/mesa/experimental/cell_space/discrete_space.py
+++ b/mesa/experimental/cell_space/discrete_space.py
@@ -2,12 +2,14 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from functools import cached_property
 from random import Random
-from typing import Generic, TypeVar
+from typing import Any, Generic, TypeVar
 
 from mesa.experimental.cell_space.cell import Cell
 from mesa.experimental.cell_space.cell_collection import CellCollection
+from mesa.space import PropertyLayer
 
 T = TypeVar("T", bound=Cell)
 
@@ -21,7 +23,7 @@ class DiscreteSpace(Generic[T]):
         random (Random): The random number generator
         cell_klass (Type) : the type of cell class
         empties (CellCollection) : collecction of all cells that are empty
-
+        property_layers (dict[str, PropertyLayer]): the property layers of the discrete space
     """
 
     def __init__(
@@ -47,6 +49,7 @@ class DiscreteSpace(Generic[T]):
 
         self._empties: dict[tuple[int, ...], None] = {}
         self._empties_initialized = False
+        self.property_layers: dict[str, PropertyLayer] = {}
 
     @property
     def cutoff_empties(self):  # noqa
@@ -73,3 +76,61 @@ class DiscreteSpace(Generic[T]):
     def select_random_empty_cell(self) -> T:
         """Select random empty cell."""
         return self.random.choice(list(self.empties))
+
+    # PropertyLayer methods
+    def add_property_layer(
+        self, property_layer: PropertyLayer, add_to_cells: bool = True
+    ):
+        """Add a property layer to the grid.
+
+        Args:
+            property_layer: the property layer to add
+            add_to_cells: whether to add the property layer to all cells (default: True)
+        """
+        if property_layer.name in self.property_layers:
+            raise ValueError(f"Property layer {property_layer.name} already exists.")
+        self.property_layers[property_layer.name] = property_layer
+        if add_to_cells:
+            for cell in self._cells.values():
+                cell._mesa_property_layers[property_layer.name] = property_layer
+
+    def remove_property_layer(self, property_name: str, remove_from_cells: bool = True):
+        """Remove a property layer from the grid.
+
+        Args:
+            property_name: the name of the property layer to remove
+            remove_from_cells: whether to remove the property layer from all cells (default: True)
+        """
+        del self.property_layers[property_name]
+        if remove_from_cells:
+            for cell in self._cells.values():
+                del cell._mesa_property_layers[property_name]
+
+    def set_property(
+        self, property_name: str, value, condition: Callable[[T], bool] | None = None
+    ):
+        """Set the value of a property for all cells in the grid.
+
+        Args:
+            property_name: the name of the property to set
+            value: the value to set
+            condition: a function that takes a cell and returns a boolean
+        """
+        self.property_layers[property_name].set_cells(value, condition)
+
+    def modify_properties(
+        self,
+        property_name: str,
+        operation: Callable,
+        value: Any = None,
+        condition: Callable[[T], bool] | None = None,
+    ):
+        """Modify the values of a specific property for all cells in the grid.
+
+        Args:
+            property_name: the name of the property to modify
+            operation: the operation to perform
+            value: the value to use in the operation
+            condition: a function that takes a cell and returns a boolean (used to filter cells)
+        """
+        self.property_layers[property_name].modify_cells(operation, value, condition)
diff --git a/mesa/model.py b/mesa/model.py
index d9566f44..e4f70793 100644
--- a/mesa/model.py
+++ b/mesa/model.py
@@ -48,8 +48,8 @@ class Model:
             seed: the seed for the random number generator
             kwargs: keyword arguments to pass onto super
         """
+        super().__init__(*args, **kwargs)
         self.running = True
-        self.schedule = None
         self.steps: int = 0
 
         self._setup_agent_registration()
diff --git a/mesa/visualization/solara_viz.py b/mesa/visualization/solara_viz.py
index ed495732..fb6eacd9 100644
--- a/mesa/visualization/solara_viz.py
+++ b/mesa/visualization/solara_viz.py
@@ -225,6 +225,7 @@ def ModelController(model: solara.Reactive[Model], play_interval=100):
         play_interval (int, optional): Interval for playing the model steps in milliseconds.
     """
     playing = solara.use_reactive(False)
+    running = solara.use_reactive(True)
     original_model = solara.use_reactive(None)
 
     def save_initial_model():
@@ -236,19 +237,23 @@ def ModelController(model: solara.Reactive[Model], play_interval=100):
     solara.use_effect(save_initial_model, [model.value])
 
     async def step():
-        while playing.value:
+        while playing.value and running.value:
             await asyncio.sleep(play_interval / 1000)
             do_step()
 
-    solara.lab.use_task(step, dependencies=[playing.value], prefer_threaded=False)
+    solara.lab.use_task(
+        step, dependencies=[playing.value, running.value], prefer_threaded=False
+    )
 
     def do_step():
         """Advance the model by one step."""
         model.value.step()
+        running.value = model.value.running
 
     def do_reset():
         """Reset the model to its initial state."""
         playing.value = False
+        running.value = True
         model.value = copy.deepcopy(original_model.value)
 
     def do_play_pause():
@@ -261,9 +266,13 @@ def ModelController(model: solara.Reactive[Model], play_interval=100):
             label="▶" if not playing.value else "❚❚",
             color="primary",
             on_click=do_play_pause,
+            disabled=not running.value,
         )
         solara.Button(
-            label="Step", color="primary", on_click=do_step, disabled=playing.value
+            label="Step",
+            color="primary",
+            on_click=do_step,
+            disabled=playing.value or not running.value,
         )
 
 
diff --git a/pyproject.toml b/pyproject.toml
index 28a6d065..fd48a030 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -70,6 +70,7 @@ examples = [
   "scipy",
 ]
 docs = [
+  "mesa[rec]",
   "sphinx",
   "ipython",
   "pydata_sphinx_theme",


Test patch:
diff --git a/tests/test_cell_space.py b/tests/test_cell_space.py
index 01420043..a8e4abad 100644
--- a/tests/test_cell_space.py
+++ b/tests/test_cell_space.py
@@ -2,6 +2,7 @@
 
 import random
 
+import numpy as np
 import pytest
 
 from mesa import Model
@@ -15,6 +16,7 @@ from mesa.experimental.cell_space import (
     OrthogonalVonNeumannGrid,
     VoronoiGrid,
 )
+from mesa.space import PropertyLayer
 
 
 def test_orthogonal_grid_neumann():
@@ -524,3 +526,118 @@ def test_cell_collection():
 
     cells = collection.select()
     assert len(cells) == len(collection)
+
+
+### PropertyLayer tests
+def test_property_layer_integration():
+    """Test integration of PropertyLayer with DiscrateSpace and Cell."""
+    width, height = 10, 10
+    grid = OrthogonalMooreGrid((width, height), torus=False)
+
+    # Test adding a PropertyLayer to the grid
+    elevation = PropertyLayer("elevation", width, height, default_value=0)
+    grid.add_property_layer(elevation)
+    assert "elevation" in grid.property_layers
+    assert len(grid.property_layers) == 1
+
+    # Test accessing PropertyLayer from a cell
+    cell = grid._cells[(0, 0)]
+    assert "elevation" in cell._mesa_property_layers
+    assert cell.get_property("elevation") == 0
+
+    # Test setting property value for a cell
+    cell.set_property("elevation", 100)
+    assert cell.get_property("elevation") == 100
+
+    # Test modifying property value for a cell
+    cell.modify_property("elevation", lambda x: x + 50)
+    assert cell.get_property("elevation") == 150
+
+    cell.modify_property("elevation", np.add, 50)
+    assert cell.get_property("elevation") == 200
+
+    # Test modifying PropertyLayer values
+    grid.set_property("elevation", 100, condition=lambda value: value == 200)
+    assert cell.get_property("elevation") == 100
+
+    # Test modifying PropertyLayer using numpy operations
+    grid.modify_properties("elevation", np.add, 50)
+    assert cell.get_property("elevation") == 150
+
+    # Test removing a PropertyLayer
+    grid.remove_property_layer("elevation")
+    assert "elevation" not in grid.property_layers
+    assert "elevation" not in cell._mesa_property_layers
+
+
+def test_multiple_property_layers():
+    """Test initialization of DiscrateSpace with PropertyLayers."""
+    width, height = 5, 5
+    elevation = PropertyLayer("elevation", width, height, default_value=0)
+    temperature = PropertyLayer("temperature", width, height, default_value=20)
+
+    # Test initialization with a single PropertyLayer
+    grid1 = OrthogonalMooreGrid((width, height), torus=False)
+    grid1.add_property_layer(elevation)
+    assert "elevation" in grid1.property_layers
+    assert len(grid1.property_layers) == 1
+
+    # Test initialization with multiple PropertyLayers
+    grid2 = OrthogonalMooreGrid((width, height), torus=False)
+    grid2.add_property_layer(temperature, add_to_cells=False)
+    grid2.add_property_layer(elevation, add_to_cells=True)
+
+    assert "temperature" in grid2.property_layers
+    assert "elevation" in grid2.property_layers
+    assert len(grid2.property_layers) == 2
+
+    # Modify properties
+    grid2.modify_properties("elevation", lambda x: x + 10)
+    grid2.modify_properties("temperature", lambda x: x + 5)
+
+    for cell in grid2.all_cells:
+        assert cell.get_property("elevation") == 10
+        # Assert error temperature, since it was not added to cells
+        with pytest.raises(KeyError):
+            cell.get_property("temperature")
+
+
+def test_property_layer_errors():
+    """Test error handling for PropertyLayers."""
+    width, height = 5, 5
+    grid = OrthogonalMooreGrid((width, height), torus=False)
+    elevation = PropertyLayer("elevation", width, height, default_value=0)
+
+    # Test adding a PropertyLayer with an existing name
+    grid.add_property_layer(elevation)
+    with pytest.raises(ValueError, match="Property layer elevation already exists."):
+        grid.add_property_layer(elevation)
+
+
+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
diff --git a/tests/test_model.py b/tests/test_model.py
index c7d17bf8..1931e002 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -8,7 +8,6 @@ def test_model_set_up():
     """Test Model initialization."""
     model = Model()
     assert model.running is True
-    assert model.schedule is None
     assert model.steps == 0
     model.step()
     assert model.steps == 1