Repository: projectmesa/mesa
Pull request #2214 (2024-08-15 16:57:17)

Visualisation: Allow specifying Agent shapes in agent_portrayal
This PR allows specifying an `"shape"` in the `agent_portrayal` dictionary used by matplotlib component of the Solara visualisation. In short, it allows you represent an Agent in any [matplotlib marker](https://matplotlib.org/stable/api/markers_api.html), by adding a "shape" key-value pair to the `agent_portrayal` dictionary.

This is especially useful when you're using the default shape drawer for grid or continuous space.

For example:
```Python
def agent_portrayal(cell):
    return {
        "color": "blue"
        "size": 5,
        "shape": "h"  # marker is a hexagon!
    }
```
Partial resolution of #2164.

### Examples

In Conway's Game of Life life, you can now use black squares for the agents, as how it's traditionally displayed:

```Python
def agent_portrayal(cell):
    return {
        "color": "white" if cell.alive else "black",
        "size": 45,
        "shape": "s" # marker is a square
    }
```
![image](https://github.com/user-attachments/assets/e80dac7e-a98c-4670-a8e9-b22cc674fecc)

You can also conditionally modify the shape, just like you can with the color and size:

```Python
COP_COLOR = "#000000"
AGENT_QUIET_COLOR = "#648FFF"
AGENT_REBEL_COLOR = "#FE6100"
JAIL_COLOR = "#808080"
JAIL_SHAPE = "s"


def citizen_cop_portrayal(agent):
    portrayal = {
        "Shape": "circle",
        "x": agent.pos[0],
        "y": agent.pos[1],
        "Filled": "true",
    }

    if isinstance(agent, Citizen):
        color = (
            AGENT_QUIET_COLOR if agent.condition == "Quiescent" else AGENT_REBEL_COLOR
        )
        if agent.jail_sentence:
            color = JAIL_COLOR
            shape = JAIL_SHAPE
        else:
            shape = "."
        color = JAIL_COLOR if agent.jail_sentence else color
        portrayal["color"] = color
        portrayal["shape"] = shape
        portrayal["size"] = 15.0

    else:
        assert isinstance(agent, Cop)
        portrayal["color"] = COP_COLOR
        portrayal["shape"] = "."
        portrayal["size"] = 25.0

    return portrayal
```
![image](https://github.com/user-attachments/assets/8ef8ab4e-8bbd-4c55-a571-cbfc2826e082)

github-actions[bot]: Performance benchmarks:

| Model        | Size   | Init time [95% CI]      | Run time [95% CI]       |
|--------------|--------|-------------------------|-------------------------|
| Schelling    | small  | 🔵 -0.0% [-0.4%, +0.3%] | 🔵 +0.7% [+0.5%, +0.9%] |
| Schelling    | large  | 🔵 +0.5% [-0.4%, +1.4%] | 🔴 +5.5% [+3.2%, +8.0%] |
| WolfSheep    | small  | 🔵 -0.1% [-1.2%, +1.1%] | 🔵 +0.4% [+0.1%, +0.7%] |
| WolfSheep    | large  | 🔵 +0.3% [-0.5%, +1.4%] | 🔵 +2.8% [+1.4%, +4.2%] |
| BoidFlockers | small  | 🔵 +0.0% [-0.7%, +0.9%] | 🔵 +0.5% [-0.3%, +1.3%] |
| BoidFlockers | large  | 🔵 -0.0% [-0.4%, +0.3%] | 🔵 -0.6% [-1.2%, -0.1%] |

rht: Implementation LGTM, thank you for implementing this feature. I could test it on the Epstein civil violence model if I have the time. But if you don't mind, showing that the jailed citizens as square in that example would be helpful. You can used the code from https://github.com/projectmesa/mesa-examples/pull/158 and undo the experimental data collector part.
rmhopkins4: ![image](https://github.com/user-attachments/assets/eeb3aee7-c24f-4587-a4d9-705ed401b761)
Here is an example of jailed citizens as squares.
While working on this I made matplotlib.py provide default visualization values so that users only need to provide overrides 
rmhopkins4: > ![image](https://private-user-images.githubusercontent.com/115746285/358376282-eeb3aee7-c24f-4587-a4d9-705ed401b761.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjM3NTM5OTIsIm5iZiI6MTcyMzc1MzY5MiwicGF0aCI6Ii8xMTU3NDYyODUvMzU4Mzc2MjgyLWVlYjNhZWU3LWMyNGYtNDU4Ny1hNGQ5LTcwNWVkNDAxYjc2MS5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjQwODE1JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI0MDgxNVQyMDI4MTJaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT1mNDllNTU0OGI2ZTQ3MmRjMzE3MmIwOWEyNGM2MjJmMWY4NmM4M2I3M2RiYTM4NzU3NTg4MGJlMjA4ZmVmMDRhJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCZhY3Rvcl9pZD0wJmtleV9pZD0wJnJlcG9faWQ9MCJ9.hKNsgQhamwuCDZKgZAqoVrh99tUlwAIiwMhVgeJpqNg) Here is an example of jailed citizens as squares. While working on this I made matplotlib.py provide default visualization values so that users only need to provide overrides

Matplotlib's 's' and 'c' (size, color) accept no values (all default), one value (all the same), and lists exactly as long as the number of points provided. 
When the list lengths were not aligned problems arose.
github-actions[bot]: Performance benchmarks:

| Model        | Size   | Init time [95% CI]      | Run time [95% CI]       |
|--------------|--------|-------------------------|-------------------------|
| Schelling    | small  | 🔵 +0.5% [+0.1%, +0.9%] | 🔵 +0.3% [+0.1%, +0.5%] |
| Schelling    | large  | 🔵 -0.0% [-0.4%, +0.4%] | 🔴 +4.8% [+3.0%, +6.5%] |
| WolfSheep    | small  | 🔵 +0.8% [-0.4%, +2.0%] | 🔵 +0.5% [+0.3%, +0.8%] |
| WolfSheep    | large  | 🔵 +0.9% [+0.5%, +1.2%] | 🔵 +2.2% [+1.5%, +2.8%] |
| BoidFlockers | small  | 🔵 +2.3% [+1.5%, +2.9%] | 🔵 -0.2% [-0.9%, +0.5%] |
| BoidFlockers | large  | 🔵 +2.1% [+1.4%, +2.7%] | 🔵 -1.1% [-1.5%, -0.6%] |

EwoutH: Thanks a lot for this PR! I will review it somewhere tomorrow.
EwoutH: @rht could you give a short summary what you already have sufficiently reviewed and what could use a second pair of eyes?
rht: Actually, this PR is already LGTM, and I didn't merge because you still wanted to review. If any, maybe the perf consideration of an array of defaults instead of a single default value scalar, but this is minor because any loop done without calling Matplotlib many times is going to be fast.
EwoutH: Thanks. I would have expected this feature taking less complexity / lines of code, let me take a closer look tomorrow morning.
EwoutH: Okay, I'm largely good with this PR. It's a bit weird https://github.com/matplotlib/matplotlib/issues/11155 was never implemented, but since the only thing changes for our users is that they can now _optionally_ pass a matplotlib shape in a `portrayal` dict (right @rht?), we can always implement the back-end later in a more (performance) efficient way.

I would like the documentation around this PR to improve a little:
- Update the docstring to note that a `"shape"` can be passed, and add a link to the list with allowed values (I think it's [this one](https://matplotlib.org/stable/api/markers_api.html))
- Update the first comment of this PR with the screenshot, one or two more examples, and a link to 
- Mention that you can pass "shape" in the [visualization_tutorial](https://mesa.readthedocs.io/en/latest/tutorials/visualization_tutorial.html) ([source](https://github.com/projectmesa/mesa/blob/main/docs/tutorials/visualization_tutorial.ipynb)).
rht: > Okay, I'm largely good with this PR. It's a bit weird https://github.com/matplotlib/matplotlib/issues/11155 was never implemented, but since the only thing changes for our users is that they can now optionally pass a matplotlib shape in a portrayal dict (right @rht?), we can always implement the back-end later in a more (performance) efficient way.

Yes, it's optional and is not a breaking change.

While it's good to update the docstring now, I think the tutorial change can happen in a separate PR. Otherwise it makes Mesa too bureaucratic to contribute to (that every new feature needs to be in the tutorial in the same PR).
EwoutH: Yeah I see your perspective and also considered that. But otherwise I'm afraid it doesn't get done. Especially with our user base changes need to be (very) visible.

@rmhopkins4 let me know if you want to do these things yourself, or otherwise allow us to do it within this PR.

Patch:
diff --git a/docs/tutorials/visualization_tutorial.ipynb b/docs/tutorials/visualization_tutorial.ipynb
index 14d9ba4d..ba8d7a92 100644
--- a/docs/tutorials/visualization_tutorial.ipynb
+++ b/docs/tutorials/visualization_tutorial.ipynb
@@ -142,7 +142,8 @@
    "source": [
     "#### Changing the agents\n",
     "\n",
-    "In the visualization above, all we could see is the agents moving around -- but not how much money they had, or anything else of interest. Let's change it so that agents who are broke (wealth 0) are drawn in red, smaller. (TODO: currently, we can't predict the drawing order of the circles, so a broke agent may be overshadowed by a wealthy agent. We should fix this by doing a hollow circle instead)\n",
+    "In the visualization above, all we could see is the agents moving around -- but not how much money they had, or anything else of interest. Let's change it so that agents who are broke (wealth 0) are drawn in red, smaller. (TODO: Currently, we can't predict the drawing order of the circles, so a broke agent may be overshadowed by a wealthy agent. We should fix this by doing a hollow circle instead)\n",
+    "In addition to size and color, an agent's shape can also be customized when using the default drawer. The allowed values for shapes can be found [here](https://matplotlib.org/stable/api/markers_api.html).\n",
     "\n",
     "To do this, we go back to our `agent_portrayal` code and add some code to change the portrayal based on the agent properties and launch the server again."
    ]
diff --git a/mesa/agent.py b/mesa/agent.py
index 2f0e842d..e65b25f6 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -11,9 +11,7 @@ from __future__ import annotations
 import contextlib
 import copy
 import operator
-import warnings
 import weakref
-from collections import defaultdict
 from collections.abc import Callable, Iterable, Iterator, MutableSet, Sequence
 from random import Random
 
@@ -52,17 +50,11 @@ class Agent:
         # register agent
         try:
             self.model.agents_[type(self)][self] = None
-        except AttributeError:
+        except AttributeError as err:
             # model super has not been called
-            self.model.agents_ = defaultdict(dict)
-            self.model.agents_[type(self)][self] = None
-            self.model.agentset_experimental_warning_given = False
-
-            warnings.warn(
-                "The Mesa Model class was not initialized. In the future, you need to explicitly initialize the Model by calling super().__init__() on initialization.",
-                FutureWarning,
-                stacklevel=2,
-            )
+            raise RuntimeError(
+                "The Mesa Model class was not initialized. You must explicitly initialize the Model by calling super().__init__() on initialization."
+            ) from err
 
     def remove(self) -> None:
         """Remove and delete the agent from the model."""
@@ -100,8 +92,6 @@ class AgentSet(MutableSet, Sequence):
         which means that agents not referenced elsewhere in the program may be automatically removed from the AgentSet.
     """
 
-    agentset_experimental_warning_given = False
-
     def __init__(self, agents: Iterable[Agent], model: Model):
         """
         Initializes the AgentSet with a collection of agents and a reference to the model.
@@ -227,26 +217,37 @@ class AgentSet(MutableSet, Sequence):
         return self
 
     def do(
-        self, method_name: str, *args, return_results: bool = False, **kwargs
+        self, method: str | Callable, *args, return_results: bool = False, **kwargs
     ) -> AgentSet | list[Any]:
         """
-        Invoke a method on each agent in the AgentSet.
+        Invoke a method or function on each agent in the AgentSet.
 
         Args:
-            method_name (str): The name of the method to call on each agent.
+            method (str, callable): the callable to do on each agents
+
+                                        * in case of str, the name of the method to call on each agent.
+                                        * in case of callable, the function to be called with each agent as first argument
+
             return_results (bool, optional): If True, returns the results of the method calls; otherwise, returns the AgentSet itself. Defaults to False, so you can chain method calls.
-            *args: Variable length argument list passed to the method being called.
-            **kwargs: Arbitrary keyword arguments passed to the method being called.
+            *args: Variable length argument list passed to the callable being called.
+            **kwargs: Arbitrary keyword arguments passed to the callable being called.
 
         Returns:
-            AgentSet | list[Any]: The results of the method calls if return_results is True, otherwise the AgentSet itself.
+            AgentSet | list[Any]: The results of the callable calls if return_results is True, otherwise the AgentSet itself.
         """
         # we iterate over the actual weakref keys and check if weakref is alive before calling the method
-        res = [
-            getattr(agent, method_name)(*args, **kwargs)
-            for agentref in self._agents.keyrefs()
-            if (agent := agentref()) is not None
-        ]
+        if isinstance(method, str):
+            res = [
+                getattr(agent, method)(*args, **kwargs)
+                for agentref in self._agents.keyrefs()
+                if (agent := agentref()) is not None
+            ]
+        else:
+            res = [
+                method(agent, *args, **kwargs)
+                for agentref in self._agents.keyrefs()
+                if (agent := agentref()) is not None
+            ]
 
         return res if return_results else self
 
diff --git a/mesa/space.py b/mesa/space.py
index a6d070f0..39433cb3 100644
--- a/mesa/space.py
+++ b/mesa/space.py
@@ -586,7 +586,7 @@ class PropertyLayer:
         aggregate_property(operation): Performs an aggregate operation over all cells.
     """
 
-    agentset_experimental_warning_given = False
+    propertylayer_experimental_warning_given = False
 
     def __init__(
         self, name: str, width: int, height: int, default_value, dtype=np.float64
@@ -633,14 +633,14 @@ class PropertyLayer:
 
         self.data = np.full((width, height), default_value, dtype=dtype)
 
-        if not self.__class__.agentset_experimental_warning_given:
+        if not self.__class__.propertylayer_experimental_warning_given:
             warnings.warn(
                 "The new PropertyLayer and _PropertyGrid classes experimental. It may be changed or removed in any and all future releases, including patch releases.\n"
                 "We would love to hear what you think about this new feature. If you have any thoughts, share them with us here: https://github.com/projectmesa/mesa/discussions/1932",
                 FutureWarning,
                 stacklevel=2,
             )
-            self.__class__.agentset_experimental_warning_given = True
+            self.__class__.propertylayer_experimental_warning_given = True
 
     def set_cell(self, position: Coordinate, value):
         """
diff --git a/mesa/visualization/components/matplotlib.py b/mesa/visualization/components/matplotlib.py
index aadfa206..83d0e3d8 100644
--- a/mesa/visualization/components/matplotlib.py
+++ b/mesa/visualization/components/matplotlib.py
@@ -1,3 +1,5 @@
+from collections import defaultdict
+
 import networkx as nx
 import solara
 from matplotlib.figure import Figure
@@ -23,12 +25,44 @@ def SpaceMatplotlib(model, agent_portrayal, dependencies: list[any] | None = Non
     solara.FigureMatplotlib(space_fig, format="png", dependencies=dependencies)
 
 
+# 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])
+
+    # 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_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]
@@ -41,23 +75,23 @@ def _draw_grid(space, space_ax, agent_portrayal):
                     data = agent_portrayal(agent)
                     x.append(i)
                     y.append(j)
-                    if "size" in data:
-                        s.append(data["size"])
-                    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"] = (180 / max(g.width, g.height)) ** 2
-        if len(s) > 0:
-            out["s"] = s
-        if len(c) > 0:
-            out["c"] = c
+
+                    # 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)
-    space_ax.scatter(**portray(space))
+    _split_and_scatter(portray(space), space_ax)
 
 
 def _draw_network_grid(space, space_ax, agent_portrayal):
@@ -77,20 +111,23 @@ def _draw_continuous_space(space, space_ax, agent_portrayal):
         y = []
         s = []  # size
         c = []  # color
+        m = []  # shape
         for agent in space._agent_to_index:
             data = agent_portrayal(agent)
             _x, _y = agent.pos
             x.append(_x)
             y.append(_y)
-            if "size" in data:
-                s.append(data["size"])
-            if "color" in data:
-                c.append(data["color"])
-        out = {"x": x, "y": y}
-        if len(s) > 0:
-            out["s"] = s
-        if len(c) > 0:
-            out["c"] = c
+
+            # 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
 
     # Determine border style based on space.torus
@@ -110,7 +147,7 @@ def _draw_continuous_space(space, space_ax, agent_portrayal):
     space_ax.set_ylim(space.y_min - y_padding, space.y_max + y_padding)
 
     # Portray and scatter the agents in the space
-    space_ax.scatter(**portray(space))
+    _split_and_scatter(portray(space), space_ax)
 
 
 @solara.component
diff --git a/mesa/visualization/solara_viz.py b/mesa/visualization/solara_viz.py
index 4f159884..6ec33231 100644
--- a/mesa/visualization/solara_viz.py
+++ b/mesa/visualization/solara_viz.py
@@ -70,7 +70,7 @@ def Card(
                 )
             elif space_drawer:
                 # if specified, draw agent space with an alternate renderer
-                space_drawer(model, agent_portrayal)
+                space_drawer(model, agent_portrayal, dependencies=dependencies)
         elif "Measure" in layout_type:
             rv.CardTitle(children=["Measure"])
             measure = measures[layout_type["Measure"]]
@@ -103,7 +103,8 @@ def SolaraViz(
         model_params: Parameters for initializing the model
         measures: List of callables or data attributes to plot
         name: Name for display
-        agent_portrayal: Options for rendering agents (dictionary)
+        agent_portrayal: Options for rendering agents (dictionary);
+            Default drawer supports custom `"size"`, `"color"`, and `"shape"`.
         space_drawer: Method to render the agent space for
             the model; default implementation is the `SpaceMatplotlib` component;
             simulations with no space to visualize should


Test patch:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 0cd21112..f4e64ce5 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -176,7 +176,7 @@ def test_agentset_get_item():
         _ = agentset[20]
 
 
-def test_agentset_do_method():
+def test_agentset_do_str():
     model = Model()
     agents = [TestAgent(model.next_id(), model) for _ in range(10)]
     agentset = AgentSet(agents, model)
@@ -210,6 +210,72 @@ def test_agentset_do_method():
     assert len(agentset) == 0
 
 
+def test_agentset_do_callable():
+    model = Model()
+    agents = [TestAgent(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+
+    # Test callable with non-existent function
+    with pytest.raises(AttributeError):
+        agentset.do(lambda agent: agent.non_existing_method())
+
+    # tests for addition and removal in do using callables
+    # do iterates, so no error should be raised to change size while iterating
+    # related to issue #1595
+
+    # setup for lambda function tests
+    n = 10
+    model = Model()
+    agents = [TestAgentDo(model.next_id(), model) for _ in range(n)]
+    agentset = AgentSet(agents, model)
+    for agent in agents:
+        agent.agent_set = agentset
+
+    # Lambda for addition
+    agentset.do(lambda agent: agent.do_add())
+    assert len(agentset) == 2 * n
+
+    # setup again for lambda function tests
+    model = Model()
+    agents = [TestAgentDo(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+    for agent in agents:
+        agent.agent_set = agentset
+
+    # Lambda for removal
+    agentset.do(lambda agent: agent.do_remove())
+    assert len(agentset) == 0
+
+    # setup for actual function tests
+    def add_function(agent):
+        agent.do_add()
+
+    def remove_function(agent):
+        agent.do_remove()
+
+    # setup again for actual function tests
+    model = Model()
+    agents = [TestAgentDo(model.next_id(), model) for _ in range(n)]
+    agentset = AgentSet(agents, model)
+    for agent in agents:
+        agent.agent_set = agentset
+
+    # Actual function for addition
+    agentset.do(add_function)
+    assert len(agentset) == 2 * n
+
+    # setup again for actual function tests
+    model = Model()
+    agents = [TestAgentDo(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+    for agent in agents:
+        agent.agent_set = agentset
+
+    # Actual function for removal
+    agentset.do(remove_function)
+    assert len(agentset) == 0
+
+
 def test_agentset_get_attribute():
     model = Model()
     agents = [TestAgent(model.next_id(), model) for _ in range(10)]
diff --git a/tests/test_solara_viz.py b/tests/test_solara_viz.py
index f3422b87..5fef7e0b 100644
--- a/tests/test_solara_viz.py
+++ b/tests/test_solara_viz.py
@@ -136,7 +136,9 @@ def test_call_space_drawer(mocker):
             space_drawer=altspace_drawer,
         )
     )
-    altspace_drawer.assert_called_with(model, agent_portrayal)
+    altspace_drawer.assert_called_with(
+        model, agent_portrayal, dependencies=dependencies
+    )
 
 
 def test_slider():