Add support for alpha, linewidths and edgecolors to agent_portrayal

mesa/visualization/mpl_space_drawing.py

@@ -6,6 +6,7 @@
 
 """
 
+import contextlib
 import itertools
 import math
 import warnings
@@ -61,10 +62,19 @@ def collect_agent_data(
         zorder: default zorder
 
     agent_portrayal should return a dict, limited to size (size of marker), color (color of marker), zorder (z-order),
-    and marker (marker style)
+    marker (marker style), alpha, linewidths, and edgecolors
 
     """
-    arguments = {"s": [], "c": [], "marker": [], "zorder": [], "loc": []}
+    arguments = {
+        "s": [],
+        "c": [],
+        "marker": [],
+        "zorder": [],
+        "loc": [],
+        "alpha": [],
+        "edgecolors": [],
+        "linewidths": [],
+    }
 
     for agent in space.agents:
         portray = agent_portrayal(agent)
@@ -78,6 +88,10 @@ def collect_agent_data(
         arguments["marker"].append(portray.pop("marker", marker))
         arguments["zorder"].append(portray.pop("zorder", zorder))
 
+        for entry in ["alpha", "edgecolors", "linewidths"]:
+            with contextlib.suppress(KeyError):
+                arguments[entry].append(portray.pop(entry))
+
         if len(portray) > 0:
             ignored_fields = list(portray.keys())
             msg = ", ".join(ignored_fields)
@@ -118,16 +132,24 @@ def draw_space(
 
     # https://stackoverflow.com/questions/67524641/convert-multiple-isinstance-checks-to-structural-pattern-matching
     match space:
-        case mesa.space._Grid() | OrthogonalMooreGrid() | OrthogonalVonNeumannGrid():
-            draw_orthogonal_grid(space, agent_portrayal, ax=ax, **space_drawing_kwargs)
+        # order matters here given the class structure of old-style grid spaces
         case HexSingleGrid() | HexMultiGrid() | mesa.experimental.cell_space.HexGrid():
             draw_hex_grid(space, agent_portrayal, ax=ax, **space_drawing_kwargs)
+        case (
+            mesa.space.SingleGrid()
+            | OrthogonalMooreGrid()
+            | OrthogonalVonNeumannGrid()
+            | mesa.space.MultiGrid()
+        ):
+            draw_orthogonal_grid(space, agent_portrayal, ax=ax, **space_drawing_kwargs)
         case mesa.space.NetworkGrid() | mesa.experimental.cell_space.Network():
             draw_network(space, agent_portrayal, ax=ax, **space_drawing_kwargs)
         case mesa.space.ContinuousSpace():
             draw_continuous_space(space, agent_portrayal, ax=ax)
         case VoronoiGrid():
             draw_voroinoi_grid(space, agent_portrayal, ax=ax)
+        case _:
+            raise ValueError(f"Unknown space type: {type(space)}")
 
     if propertylayer_portrayal:
         draw_property_layers(space, propertylayer_portrayal, ax=ax)
@@ -543,11 +565,24 @@ def _scatter(ax: Axes, arguments, **kwargs):
     marker = arguments.pop("marker")
     zorder = arguments.pop("zorder")
 
+    # we check if edgecolor, linewidth, and alpha are specified
+    # at the agent level, if not, we remove them from the arguments dict
+    # and fallback to the default value in ax.scatter / use what is passed via **kwargs
+    for entry in ["edgecolors", "linewidths", "alpha"]:
+        if len(arguments[entry]) == 0:
+            arguments.pop(entry)
+        else:
+            if entry in kwargs:
+                raise ValueError(
+                    f"{entry} is specified in agent portrayal and via plotting kwargs, you can only use one or the other"
+                )
+
     for mark in np.unique(marker):
         mark_mask = marker == mark
         for z_order in np.unique(zorder):
             zorder_mask = z_order == zorder
             logical = mark_mask & zorder_mask
+
             ax.scatter(
                 x[logical],
                 y[logical],

tests/test_components_matplotlib.py

@@ -23,6 +23,7 @@
     draw_network,
     draw_orthogonal_grid,
     draw_property_layers,
+    draw_space,
     draw_voroinoi_grid,
 )
 
@@ -41,6 +42,84 @@ def agent_portrayal(agent):
     }
 
 
+def test_draw_space():
+    """Test draw_space helper method."""
+    import networkx as nx
+
+    def my_portrayal(agent):
+        """Simple portrayal of an agent.
+
+        Args:
+            agent (Agent): The agent to portray
+
+        """
+        return {
+            "s": 10,
+            "c": "tab:blue",
+            "marker": "s" if (agent.unique_id % 2) == 0 else "o",
+            "alpha": 0.5,
+            "linewidths": 1,
+            "linecolors": "tab:orange",
+        }
+
+    # draw space for hexgrid
+    model = Model(seed=42)
+    grid = HexSingleGrid(10, 10, torus=True)
+    for _ in range(10):
+        agent = Agent(model)
+        grid.move_to_empty(agent)
+
+    fig, ax = plt.subplots()
+    draw_space(grid, my_portrayal, ax=ax)
+
+    # draw space for voroinoi
+    model = Model(seed=42)
+    coordinates = model.rng.random((100, 2)) * 10
+    grid = VoronoiGrid(coordinates.tolist(), random=model.random, capacity=1)
+    for _ in range(10):
+        agent = CellAgent(model)
+        agent.cell = grid.select_random_empty_cell()
+
+    fig, ax = plt.subplots()
+    draw_space(grid, my_portrayal, ax=ax)
+
+    # draw orthogonal grid
+    model = Model(seed=42)
+    grid = OrthogonalMooreGrid((10, 10), torus=True, random=model.random, capacity=1)
+    for _ in range(10):
+        agent = CellAgent(model)
+        agent.cell = grid.select_random_empty_cell()
+    fig, ax = plt.subplots()
+    draw_space(grid, my_portrayal, ax=ax)
+
+    # draw network
+    n = 10
+    m = 20
+    seed = 42
+    graph = nx.gnm_random_graph(n, m, seed=seed)
+
+    model = Model(seed=42)
+    grid = NetworkGrid(graph)
+    for _ in range(10):
+        agent = Agent(model)
+        pos = agent.random.randint(0, len(graph.nodes) - 1)
+        grid.place_agent(agent, pos)
+    fig, ax = plt.subplots()
+    draw_space(grid, my_portrayal, ax=ax)
+
+    # draw continuous space
+    model = Model(seed=42)
+    space = ContinuousSpace(10, 10, torus=True)
+    for _ in range(10):
+        x = model.random.random() * 10
+        y = model.random.random() * 10
+        agent = Agent(model)
+        space.place_agent(agent, (x, y))
+
+    fig, ax = plt.subplots()
+    draw_space(space, my_portrayal, ax=ax)
+
+
 def test_draw_hex_grid():
     """Test drawing hexgrids."""
     model = Model(seed=42)