Repository: projectmesa/mesa
Pull request #2336 (2024-10-03 08:17:15)

Visualize PropertyLayers
This PR adds support for visualizing PropertyLayers in the Matplotlib-based space visualization component. It allows users to overlay PropertyLayer data on top of the existing grid and agent visualizations, or on its own.

It introduces a new `propertylayer_portrayal` parameter to customize the appearance of PropertyLayers and refactors the existing space visualization code for better modularity and flexibility.

### Motive
The PropertyLayer is becoming a key feature in Mesa, so it should be visualized easily. This feature enables users to easily visualize multiple layers of data simultaneously, such as environmental factors or agent properties, alongside the agents themselves.

### Implementation
- Added a new `propertylayer_portrayal` parameter to the `make_space_matplotlib` function.
- Implemented a `draw_property_layers` function to render PropertyLayer data using Matplotlib's `imshow`.
- Modified the `_draw_grid` function to incorporate PropertyLayer visualization before drawing agents.
- Added support for both color-based and colormap-based PropertyLayer rendering.
- Implemented alpha blending for PropertyLayer visualization to allow overlaying multiple layers.
- Added colorbar support for PropertyLayer visualization.

### Usage Examples
You can visualize multiple PropertyLayers ("grass" and "temperature" in this case) along with agents in a grid space:
```python
def agent_portrayal(agent):
    return {"color": "red", "shape": "o", "size": 20}

propertylayer_portrayal = {
    "grass": {
        "colormap": "Greens",
        "alpha": 1,
    },
    "temperature": {
        "colormap": "coolwarm",
        "alpha": 0.333,
        "vmin": 0,
        "vmax": 100
    },
}

model = ExampleModel(10, 10)
page = SolaraViz(
    model,
    [make_space_matplotlib(agent_portrayal, propertylayer_portrayal)],
)
```
![image](https://github.com/user-attachments/assets/45a7eab4-3225-451b-aa98-6fd5702c1c22)

You can also use a single color for a PropertyLayer. Here we render the "cell_layer" PropertyLayer using shades of black, with the highest value having 0% opacity and the lowest value having 100% opacity. For Game of Life that means alive cells are fully black, and dead cells transparent.

```python
propertylayer_portrayal = {
    "cell_layer": {
        "color": "Black",
        "alpha": 1,
        "colorbar": False,
    },
}
```
![image](https://github.com/user-attachments/assets/b8b576be-219e-4711-8dc8-112631947093)

Any [matplotlib color](https://matplotlib.org/stable/api/colors_api.html#module-matplotlib.colors) is valid. The highest value will have that color with 0% opacity, the lowest with 100% opacity (which can again be scaled with modifying `vmin` and `vmax`).

## Additional Notes
- Currently supported spaces: Grid and standalone PropertyLayers.
  - ContinuousSpace, NetworkGrid, and VoronoiGrid can be added in future PR(s).
- Future work could include adding support for additional space types and improving the efficiency of rendering for large grids.
- The PR includes updates to handle different sizes between the PropertyLayer and the grid space, issuing warnings when there's a mismatch.
- Colorbar display can be toggled using the `colorbar` parameter in the portrayal dictionary.
- The coordinate system has been validated to ensure consistent axis directions between PropertyLayers and the grid space.
- A demo using the Game of Life model has been implemented in a separate PR (https://github.com/projectmesa/mesa-examples/pull/214) to showcase these new visualization capabilities.

### TODO
- [x] Scale colors over alpha (aside from colormap)
- [x] Add switch for colorbars and fix warning
- [x] Validate coordinate system
  - [x] Handle different sizes (just warn)
  - [x] Axis in same direction
- [x] Demo GOL (https://github.com/projectmesa/mesa-examples/pull/214)
- [ ] Add tests
- [x] Write up PR

github-actions[bot]: Performance benchmarks:

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

EwoutH: This is ready for a first round of review, API wise its done. Any conceptual, feature or API feedback is welcome. If no major things come up, I will finish the tests and such.

The GoL example is also available here: https://github.com/projectmesa/mesa-examples/pull/214
EwoutH: Implemented the last features (turning color bars on and off), validated the PropertyLayer coordinate system was visualized on the same axis as the Agent one, and updated the PR description.

![image](https://github.com/user-attachments/assets/be61e88e-b897-41c9-96cd-5e9c0925ac46)

Ready for review.
github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]       |
|-----------------|--------|-------------------------|-------------------------|
| BoltzmannWealth | small  | 🔵 -1.6% [-2.1%, -0.9%] | 🔵 -0.4% [-0.4%, -0.3%] |
| BoltzmannWealth | large  | 🔵 -1.5% [-1.8%, -1.2%] | 🔵 -0.1% [-0.6%, +0.3%] |
| Schelling       | small  | 🟢 -3.9% [-4.0%, -3.7%] | 🔵 -1.3% [-1.6%, -1.1%] |
| Schelling       | large  | 🟢 -3.9% [-4.4%, -3.4%] | 🔵 -3.5% [-5.4%, -2.2%] |
| WolfSheep       | small  | 🔵 -2.4% [-2.5%, -2.2%] | 🔵 -1.0% [-1.2%, -0.9%] |
| WolfSheep       | large  | 🔵 -2.6% [-3.4%, -2.0%] | 🔵 -1.7% [-2.2%, -1.2%] |
| BoidFlockers    | small  | 🔵 -2.2% [-2.7%, -1.6%] | 🔵 -2.2% [-2.9%, -1.6%] |
| BoidFlockers    | large  | 🔵 -1.3% [-1.7%, -0.9%] | 🔵 -2.7% [-3.3%, -2.1%] |


Patch:
diff --git a/benchmarks/WolfSheep/wolf_sheep.py b/benchmarks/WolfSheep/wolf_sheep.py
index 9fd71846..f085ce42 100644
--- a/benchmarks/WolfSheep/wolf_sheep.py
+++ b/benchmarks/WolfSheep/wolf_sheep.py
@@ -10,7 +10,7 @@ Replication of the model found in NetLogo:
 import math
 
 from mesa import Model
-from mesa.experimental.cell_space import CellAgent, OrthogonalVonNeumannGrid
+from mesa.experimental.cell_space import CellAgent, FixedAgent, OrthogonalVonNeumannGrid
 from mesa.experimental.devs import ABMSimulator
 
 
@@ -87,7 +87,7 @@ class Wolf(Animal):
             sheep_to_eat.remove()
 
 
-class GrassPatch(CellAgent):
+class GrassPatch(FixedAgent):
     """A patch of grass that grows at a fixed rate and it is eaten by sheep."""
 
     @property
diff --git a/mesa/experimental/cell_space/__init__.py b/mesa/experimental/cell_space/__init__.py
index 792dde61..69386a4c 100644
--- a/mesa/experimental/cell_space/__init__.py
+++ b/mesa/experimental/cell_space/__init__.py
@@ -6,7 +6,11 @@ expressive that the default grids available in `mesa.space`.
 """
 
 from mesa.experimental.cell_space.cell import Cell
-from mesa.experimental.cell_space.cell_agent import CellAgent
+from mesa.experimental.cell_space.cell_agent import (
+    CellAgent,
+    FixedAgent,
+    Grid2DMovingAgent,
+)
 from mesa.experimental.cell_space.cell_collection import CellCollection
 from mesa.experimental.cell_space.discrete_space import DiscreteSpace
 from mesa.experimental.cell_space.grid import (
@@ -22,6 +26,8 @@ __all__ = [
     "CellCollection",
     "Cell",
     "CellAgent",
+    "Grid2DMovingAgent",
+    "FixedAgent",
     "DiscreteSpace",
     "Grid",
     "HexGrid",
diff --git a/mesa/experimental/cell_space/cell.py b/mesa/experimental/cell_space/cell.py
index 8c7fb9bb..6c92afbe 100644
--- a/mesa/experimental/cell_space/cell.py
+++ b/mesa/experimental/cell_space/cell.py
@@ -7,12 +7,12 @@ from functools import cache, cached_property
 from random import Random
 from typing import TYPE_CHECKING, Any
 
+from mesa.experimental.cell_space.cell_agent import CellAgent
 from mesa.experimental.cell_space.cell_collection import CellCollection
 from mesa.space import PropertyLayer
 
 if TYPE_CHECKING:
     from mesa.agent import Agent
-    from mesa.experimental.cell_space.cell_agent import CellAgent
 
 Coordinate = tuple[int, ...]
 
@@ -69,7 +69,7 @@ class Cell:
         self.agents: list[
             Agent
         ] = []  # TODO:: change to AgentSet or weakrefs? (neither is very performant, )
-        self.capacity: int = capacity
+        self.capacity: int | None = capacity
         self.properties: dict[Coordinate, object] = {}
         self.random = random
         self._mesa_property_layers: dict[str, PropertyLayer] = {}
@@ -136,7 +136,7 @@ class Cell:
         return f"Cell({self.coordinate}, {self.agents})"
 
     @cached_property
-    def neighborhood(self) -> CellCollection:
+    def neighborhood(self) -> CellCollection[Cell]:
         """Returns the direct neighborhood of the cell.
 
         This is equivalent to cell.get_neighborhood(radius=1)
@@ -148,7 +148,7 @@ class Cell:
     @cache  # noqa: B019
     def get_neighborhood(
         self, radius: int = 1, include_center: bool = False
-    ) -> CellCollection:
+    ) -> CellCollection[Cell]:
         """Returns a list of all neighboring cells for the given radius.
 
         For getting the direct neighborhood (i.e., radius=1) you can also use
diff --git a/mesa/experimental/cell_space/cell_agent.py b/mesa/experimental/cell_space/cell_agent.py
index fea71b02..5dc967ee 100644
--- a/mesa/experimental/cell_space/cell_agent.py
+++ b/mesa/experimental/cell_space/cell_agent.py
@@ -2,18 +2,24 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Protocol
 
-from mesa import Agent
+from mesa.agent import Agent
 
 if TYPE_CHECKING:
-    from mesa.experimental.cell_space.cell import Cell
+    from mesa.experimental.cell_space import Cell
+
+
+class HasCellProtocol(Protocol):
+    """Protocol for discrete space cell holders."""
+
+    cell: Cell
 
 
 class HasCell:
     """Descriptor for cell movement behavior."""
 
-    _mesa_cell: Cell = None
+    _mesa_cell: Cell | None = None
 
     @property
     def cell(self) -> Cell | None:  # noqa: D102
@@ -33,17 +39,95 @@ class HasCell:
             cell.add_agent(self)
 
 
-class CellAgent(Agent, HasCell):
+class BasicMovement:
+    """Mixin for moving agents in discrete space."""
+
+    def move_to(self: HasCellProtocol, cell: Cell) -> None:
+        """Move to a new cell."""
+        self.cell = cell
+
+    def move_relative(self: HasCellProtocol, direction: tuple[int, ...]):
+        """Move to a cell relative to the current cell.
+
+        Args:
+            direction: The direction to move in.
+        """
+        new_cell = self.cell.connections.get(direction)
+        if new_cell is not None:
+            self.cell = new_cell
+        else:
+            raise ValueError(f"No cell in direction {direction}")
+
+
+class FixedCell(HasCell):
+    """Mixin for agents that are fixed to a cell."""
+
+    @property
+    def cell(self) -> Cell | None:  # noqa: D102
+        return self._mesa_cell
+
+    @cell.setter
+    def cell(self, cell: Cell) -> None:
+        if self.cell is not None:
+            raise ValueError("Cannot move agent in FixedCell")
+        self._mesa_cell = cell
+
+        cell.add_agent(self)
+
+
+class CellAgent(Agent, HasCell, BasicMovement):
     """Cell Agent is an extension of the Agent class and adds behavior for moving in discrete spaces.
 
     Attributes:
-        unique_id (int): A unique identifier for this agent.
-        model (Model): The model instance to which the agent belongs
-        pos: (Position | None): The position of the agent in the space
-        cell: (Cell | None): the cell which the agent occupies
+        cell (Cell): The cell the agent is currently in.
     """
 
     def remove(self):
         """Remove the agent from the model."""
         super().remove()
         self.cell = None  # ensures that we are also removed from cell
+
+
+class FixedAgent(Agent, FixedCell):
+    """A patch in a 2D grid."""
+
+    def remove(self):
+        """Remove the agent from the model."""
+        super().remove()
+
+        # fixme we leave self._mesa_cell on the original value
+        #  so you cannot hijack remove() to move patches
+        self.cell.remove_agent(self)
+
+
+class Grid2DMovingAgent(CellAgent):
+    """Mixin for moving agents in 2D grids."""
+
+    # fmt: off
+    DIRECTION_MAP = {
+        "n": (-1, 0), "north": (-1, 0), "up": (-1, 0),
+        "s": (1, 0), "south": (1, 0), "down": (1, 0),
+        "e": (0, 1), "east": (0, 1), "right": (0, 1),
+        "w": (0, -1), "west": (0, -1), "left": (0, -1),
+        "ne": (-1, 1), "northeast": (-1, 1), "upright": (-1, 1),
+        "nw": (-1, -1), "northwest": (-1, -1), "upleft": (-1, -1),
+        "se": (1, 1), "southeast": (1, 1), "downright": (1, 1),
+        "sw": (1, -1), "southwest": (1, -1), "downleft": (1, -1)
+    }
+    # fmt: on
+
+    def move(self, direction: str, distance: int = 1):
+        """Move the agent in a cardinal direction.
+
+        Args:
+            direction: The cardinal direction to move in.
+            distance: The distance to move.
+        """
+        direction = direction.lower()  # Convert direction to lowercase
+
+        if direction not in self.DIRECTION_MAP:
+            raise ValueError(f"Invalid direction: {direction}")
+
+        move_vector = self.DIRECTION_MAP[direction]
+        for _ in range(distance):
+            self.move_relative(move_vector)
diff --git a/mesa/experimental/devs/examples/wolf_sheep.py b/mesa/experimental/devs/examples/wolf_sheep.py
index 8d7d16d6..74318ef8 100644
--- a/mesa/experimental/devs/examples/wolf_sheep.py
+++ b/mesa/experimental/devs/examples/wolf_sheep.py
@@ -1,6 +1,7 @@
 """Example of using ABM simulator for Wolf-Sheep Predation Model."""
 
 import mesa
+from mesa.experimental.cell_space import FixedAgent
 from mesa.experimental.devs.simulator import ABMSimulator
 
 
@@ -90,7 +91,7 @@ class Wolf(Animal):
             sheep_to_eat.die()
 
 
-class GrassPatch(mesa.Agent):
+class GrassPatch(FixedAgent):
     """A patch of grass that grows at a fixed rate and it is eaten by sheep."""
 
     @property
diff --git a/mesa/visualization/components/matplotlib.py b/mesa/visualization/components/matplotlib.py
index 6061356a..771978c1 100644
--- a/mesa/visualization/components/matplotlib.py
+++ b/mesa/visualization/components/matplotlib.py
@@ -1,116 +1,190 @@
 """Matplotlib based solara components for visualization MESA spaces and plots."""
 
-from collections import defaultdict
+import warnings
 
+import matplotlib.pyplot as plt
 import networkx as nx
+import numpy as np
 import solara
+from matplotlib.cm import ScalarMappable
+from matplotlib.colors import LinearSegmentedColormap, Normalize, to_rgba
 from matplotlib.figure import Figure
-from matplotlib.ticker import MaxNLocator
 
 import mesa
 from mesa.experimental.cell_space import VoronoiGrid
+from mesa.space import PropertyLayer
 from mesa.visualization.utils import update_counter
 
 
-def make_space_matplotlib(agent_portrayal=None):  # noqa: D103
+def make_space_matplotlib(agent_portrayal=None, propertylayer_portrayal=None):
+    """Create a Matplotlib-based space visualization component.
+
+    Args:
+        agent_portrayal (function): Function to portray agents
+        propertylayer_portrayal (dict): Dictionary of PropertyLayer portrayal specifications
+
+    Returns:
+        function: A function that creates a SpaceMatplotlib component
+    """
     if agent_portrayal is None:
 
         def agent_portrayal(a):
             return {"id": a.unique_id}
 
     def MakeSpaceMatplotlib(model):
-        return SpaceMatplotlib(model, agent_portrayal)
+        return SpaceMatplotlib(model, agent_portrayal, propertylayer_portrayal)
 
     return MakeSpaceMatplotlib
 
 
 @solara.component
-def SpaceMatplotlib(model, agent_portrayal, dependencies: list[any] | None = None):  # noqa: D103
+def SpaceMatplotlib(
+    model,
+    agent_portrayal,
+    propertylayer_portrayal,
+    dependencies: list[any] | None = None,
+):
+    """Create a Matplotlib-based space visualization component."""
     update_counter.get()
     space_fig = Figure()
     space_ax = space_fig.subplots()
     space = getattr(model, "grid", None)
     if space is None:
-        # Sometimes the space is defined as model.space instead of model.grid
-        space = model.space
-    if isinstance(space, mesa.space.NetworkGrid):
-        _draw_network_grid(space, space_ax, agent_portrayal)
+        space = getattr(model, "space", None)
+
+    if isinstance(space, mesa.space._Grid):
+        _draw_grid(space, space_ax, agent_portrayal, propertylayer_portrayal, model)
     elif isinstance(space, mesa.space.ContinuousSpace):
-        _draw_continuous_space(space, space_ax, agent_portrayal)
+        _draw_continuous_space(space, space_ax, agent_portrayal, model)
+    elif isinstance(space, mesa.space.NetworkGrid):
+        _draw_network_grid(space, space_ax, agent_portrayal)
     elif isinstance(space, VoronoiGrid):
         _draw_voronoi(space, space_ax, agent_portrayal)
-    else:
-        _draw_grid(space, space_ax, agent_portrayal)
-    solara.FigureMatplotlib(space_fig, format="png", dependencies=dependencies)
-
+    elif space is None and propertylayer_portrayal:
+        draw_property_layers(space_ax, space, propertylayer_portrayal, model)
 
-# matplotlib scatter does not allow for multiple shapes in one call
-def _split_and_scatter(portray_data, space_ax):
-    grouped_data = defaultdict(lambda: {"x": [], "y": [], "s": [], "c": []})
-
-    # Extract data from the dictionary
-    x = portray_data["x"]
-    y = portray_data["y"]
-    s = portray_data["s"]
-    c = portray_data["c"]
-    m = portray_data["m"]
-
-    if not (len(x) == len(y) == len(s) == len(c) == len(m)):
-        raise ValueError(
-            "Length mismatch in portrayal data lists: "
-            f"x: {len(x)}, y: {len(y)}, size: {len(s)}, "
-            f"color: {len(c)}, marker: {len(m)}"
-        )
-
-    # Group the data by marker
-    for i in range(len(x)):
-        marker = m[i]
-        grouped_data[marker]["x"].append(x[i])
-        grouped_data[marker]["y"].append(y[i])
-        grouped_data[marker]["s"].append(s[i])
-        grouped_data[marker]["c"].append(c[i])
+    solara.FigureMatplotlib(space_fig, format="png", dependencies=dependencies)
 
-    # Plot each group with the same marker
-    for marker, data in grouped_data.items():
-        space_ax.scatter(data["x"], data["y"], s=data["s"], c=data["c"], marker=marker)
 
+def draw_property_layers(ax, space, propertylayer_portrayal, model):
+    """Draw PropertyLayers on the given axes.
+
+    Args:
+        ax (matplotlib.axes.Axes): The axes to draw on.
+        space (mesa.space._Grid): The space containing the PropertyLayers.
+        propertylayer_portrayal (dict): Dictionary of PropertyLayer portrayal specifications.
+        model (mesa.Model): The model instance.
+    """
+    for layer_name, portrayal in propertylayer_portrayal.items():
+        layer = getattr(model, layer_name, None)
+        if not isinstance(layer, PropertyLayer):
+            continue
+
+        data = layer.data.astype(float) if layer.data.dtype == bool else layer.data
+        width, height = data.shape if space is None else (space.width, space.height)
+
+        if space and data.shape != (width, height):
+            warnings.warn(
+                f"Layer {layer_name} dimensions ({data.shape}) do not match space dimensions ({width}, {height}).",
+                UserWarning,
+                stacklevel=2,
+            )
+
+        # Get portrayal properties, or use defaults
+        alpha = portrayal.get("alpha", 1)
+        vmin = portrayal.get("vmin", np.min(data))
+        vmax = portrayal.get("vmax", np.max(data))
+        colorbar = portrayal.get("colorbar", True)
+
+        # Draw the layer
+        if "color" in portrayal:
+            rgba_color = to_rgba(portrayal["color"])
+            normalized_data = (data - vmin) / (vmax - vmin)
+            rgba_data = np.full((*data.shape, 4), rgba_color)
+            rgba_data[..., 3] *= normalized_data * alpha
+            rgba_data = np.clip(rgba_data, 0, 1)
+            cmap = LinearSegmentedColormap.from_list(
+                layer_name, [(0, 0, 0, 0), (*rgba_color[:3], alpha)]
+            )
+            im = ax.imshow(
+                rgba_data.transpose(1, 0, 2),
+                extent=(0, width, 0, height),
+                origin="lower",
+            )
+            if colorbar:
+                norm = Normalize(vmin=vmin, vmax=vmax)
+                sm = ScalarMappable(norm=norm, cmap=cmap)
+                sm.set_array([])
+                ax.figure.colorbar(sm, ax=ax, orientation="vertical")
+
+        elif "colormap" in portrayal:
+            cmap = portrayal.get("colormap", "viridis")
+            if isinstance(cmap, list):
+                cmap = LinearSegmentedColormap.from_list(layer_name, cmap)
+            im = ax.imshow(
+                data.T,
+                cmap=cmap,
+                alpha=alpha,
+                vmin=vmin,
+                vmax=vmax,
+                extent=(0, width, 0, height),
+                origin="lower",
+            )
+            if colorbar:
+                plt.colorbar(im, ax=ax, label=layer_name)
+        else:
+            raise ValueError(
+                f"PropertyLayer {layer_name} portrayal must include 'color' or 'colormap'."
+            )
+
+
+def _draw_grid(space, space_ax, agent_portrayal, propertylayer_portrayal, model):
+    if propertylayer_portrayal:
+        draw_property_layers(space_ax, space, propertylayer_portrayal, model)
+
+    agent_data = _get_agent_data(space, agent_portrayal)
+
+    space_ax.set_xlim(0, space.width)
+    space_ax.set_ylim(0, space.height)
+    _split_and_scatter(agent_data, space_ax)
+
+    # Draw grid lines
+    for x in range(space.width + 1):
+        space_ax.axvline(x, color="gray", linestyle=":")
+    for y in range(space.height + 1):
+        space_ax.axhline(y, color="gray", linestyle=":")
+
+
+def _get_agent_data(space, agent_portrayal):
+    """Helper function to get agent data for visualization."""
+    x, y, s, c, m = [], [], [], [], []
+    for agents, pos in space.coord_iter():
+        if not agents:
+            continue
+        if not isinstance(agents, list):
+            agents = [agents]  # noqa PLW2901
+        for agent in agents:
+            data = agent_portrayal(agent)
+            x.append(pos[0] + 0.5)  # Center the agent in the cell
+            y.append(pos[1] + 0.5)  # Center the agent in the cell
+            default_size = (180 / max(space.width, space.height)) ** 2
+            s.append(data.get("size", default_size))
+            c.append(data.get("color", "tab:blue"))
+            m.append(data.get("shape", "o"))
+    return {"x": x, "y": y, "s": s, "c": c, "m": m}
 
-def _draw_grid(space, space_ax, agent_portrayal):
-    def portray(g):
-        x = []
-        y = []
-        s = []  # size
-        c = []  # color
-        m = []  # shape
-        for i in range(g.width):
-            for j in range(g.height):
-                content = g._grid[i][j]
-                if not content:
-                    continue
-                if not hasattr(content, "__iter__"):
-                    # Is a single grid
-                    content = [content]
-                for agent in content:
-                    data = agent_portrayal(agent)
-                    x.append(i)
-                    y.append(j)
-
-                    # This is the default value for the marker size, which auto-scales
-                    # according to the grid area.
-                    default_size = (180 / max(g.width, g.height)) ** 2
-                    # establishing a default prevents misalignment if some agents are not given size, color, etc.
-                    size = data.get("size", default_size)
-                    s.append(size)
-                    color = data.get("color", "tab:blue")
-                    c.append(color)
-                    mark = data.get("shape", "o")
-                    m.append(mark)
-        out = {"x": x, "y": y, "s": s, "c": c, "m": m}
-        return out
 
-    space_ax.set_xlim(-1, space.width)
-    space_ax.set_ylim(-1, space.height)
-    _split_and_scatter(portray(space), space_ax)
+def _split_and_scatter(portray_data, space_ax):
+    """Helper function to split and scatter agent data."""
+    for marker in set(portray_data["m"]):
+        mask = [m == marker for m in portray_data["m"]]
+        space_ax.scatter(
+            [x for x, show in zip(portray_data["x"], mask) if show],
+            [y for y, show in zip(portray_data["y"], mask) if show],
+            s=[s for s, show in zip(portray_data["s"], mask) if show],
+            c=[c for c, show in zip(portray_data["c"], mask) if show],
+            marker=marker,
+        )
 
 
 def _draw_network_grid(space, space_ax, agent_portrayal):
@@ -124,7 +198,7 @@ def _draw_network_grid(space, space_ax, agent_portrayal):
     )
 
 
-def _draw_continuous_space(space, space_ax, agent_portrayal):
+def _draw_continuous_space(space, space_ax, agent_portrayal, model):
     def portray(space):
         x = []
         y = []
@@ -139,15 +213,13 @@ def _draw_continuous_space(space, space_ax, agent_portrayal):
 
             # This is matplotlib's default marker size
             default_size = 20
-            # establishing a default prevents misalignment if some agents are not given size, color, etc.
             size = data.get("size", default_size)
             s.append(size)
             color = data.get("color", "tab:blue")
             c.append(color)
             mark = data.get("shape", "o")
             m.append(mark)
-        out = {"x": x, "y": y, "s": s, "c": c, "m": m}
-        return out
+        return {"x": x, "y": y, "s": s, "c": c, "m": m}
 
     # Determine border style based on space.torus
     border_style = "solid" if not space.torus else (0, (5, 10))
@@ -186,8 +258,6 @@ def _draw_voronoi(space, space_ax, agent_portrayal):
                 if "color" in data:
                     c.append(data["color"])
         out = {"x": x, "y": y}
-        # This is the default value for the marker size, which auto-scales
-        # according to the grid area.
         out["s"] = s
         if len(c) > 0:
             out["c"] = c
@@ -216,10 +286,19 @@ def _draw_voronoi(space, space_ax, agent_portrayal):
             alpha=min(1, cell.properties[space.cell_coloring_property]),
             c="red",
         )  # Plot filled polygon
-        space_ax.plot(*zip(*polygon), color="black")  # Plot polygon edges in red
+        space_ax.plot(*zip(*polygon), color="black")  # Plot polygon edges in black
+
+
+def make_plot_measure(measure: str | dict[str, str] | list[str] | tuple[str]):
+    """Create a plotting function for a specified measure.
 
+    Args:
+        measure (str | dict[str, str] | list[str] | tuple[str]): Measure(s) to plot.
+
+    Returns:
+        function: A function that creates a PlotMatplotlib component.
+    """
 
-def make_plot_measure(measure: str | dict[str, str] | list[str] | tuple[str]):  # noqa: D103
     def MakePlotMeasure(model):
         return PlotMatplotlib(model, measure)
 
@@ -227,7 +306,17 @@ def make_plot_measure(measure: str | dict[str, str] | list[str] | tuple[str]):
 
 
 @solara.component
-def PlotMatplotlib(model, measure, dependencies: list[any] | None = None):  # noqa: D103
+def PlotMatplotlib(model, measure, dependencies: list[any] | None = None):
+    """Create a Matplotlib-based plot for a measure or measures.
+
+    Args:
+        model (mesa.Model): The model instance.
+        measure (str | dict[str, str] | list[str] | tuple[str]): Measure(s) to plot.
+        dependencies (list[any] | None): Optional dependencies for the plot.
+
+    Returns:
+        solara.FigureMatplotlib: A component for rendering the plot.
+    """
     update_counter.get()
     fig = Figure()
     ax = fig.subplots()
@@ -244,5 +333,5 @@ def PlotMatplotlib(model, measure, dependencies: list[any] | None = None):  # no
             ax.plot(df.loc[:, m], label=m)
         fig.legend()
     # Set integer x axis
-    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
+    ax.xaxis.set_major_locator(plt.MaxNLocator(integer=True))
     solara.FigureMatplotlib(fig, dependencies=dependencies)


Test patch:
diff --git a/tests/test_cell_space.py b/tests/test_cell_space.py
index a8e4abad..4d52e159 100644
--- a/tests/test_cell_space.py
+++ b/tests/test_cell_space.py
@@ -10,6 +10,8 @@ from mesa.experimental.cell_space import (
     Cell,
     CellAgent,
     CellCollection,
+    FixedAgent,
+    Grid2DMovingAgent,
     HexGrid,
     Network,
     OrthogonalMooreGrid,
@@ -641,3 +643,51 @@ def test_cell_agent():  # noqa: D103
     assert agent not in model._all_agents
     assert agent not in cell1.agents
     assert agent not in cell2.agents
+
+    model = Model()
+    agent = CellAgent(model)
+    agent.cell = cell1
+    agent.move_to(cell2)
+    assert agent not in cell1.agents
+    assert agent in cell2.agents
+
+
+def test_grid2DMovingAgent():  # noqa: D103
+    # we first test on a moore grid because all directions are defined
+    grid = OrthogonalMooreGrid((10, 10), torus=False)
+
+    model = Model()
+    agent = Grid2DMovingAgent(model)
+
+    agent.cell = grid[4, 4]
+    agent.move("up")
+    assert agent.cell == grid[3, 4]
+
+    grid = OrthogonalVonNeumannGrid((10, 10), torus=False)
+
+    model = Model()
+    agent = Grid2DMovingAgent(model)
+    agent.cell = grid[4, 4]
+
+    with pytest.raises(ValueError):  # test for invalid direction
+        agent.move("upright")
+
+    with pytest.raises(ValueError):  # test for unknown direction
+        agent.move("back")
+
+
+def test_patch():  # noqa: D103
+    cell1 = Cell((1,), capacity=None, random=random.Random())
+    cell2 = Cell((2,), capacity=None, random=random.Random())
+
+    # connect
+    # add_agent
+    model = Model()
+    agent = FixedAgent(model)
+    agent.cell = cell1
+
+    with pytest.raises(ValueError):
+        agent.cell = cell2
+
+    agent.remove()
+    assert agent not in model._agents
diff --git a/tests/test_solara_viz.py b/tests/test_solara_viz.py
index 301294f2..a0d2b449 100644
--- a/tests/test_solara_viz.py
+++ b/tests/test_solara_viz.py
@@ -100,11 +100,14 @@ def test_call_space_drawer(mocker):  # noqa: D103
         "Shape": "circle",
         "color": "gray",
     }
+    propertylayer_portrayal = None
     # initialize with space drawer unspecified (use default)
     # component must be rendered for code to run
     solara.render(SolaraViz(model, components=[make_space_matplotlib(agent_portrayal)]))
     # should call default method with class instance and agent portrayal
-    mock_space_matplotlib.assert_called_with(model, agent_portrayal)
+    mock_space_matplotlib.assert_called_with(
+        model, agent_portrayal, propertylayer_portrayal
+    )
 
     # specify no space should be drawn
     mock_space_matplotlib.reset_mock()