diff --git a/mesa/visualization/components/matplotlib_components.py b/mesa/visualization/components/matplotlib_components.py index 6c7bb1ae040..9e8464f22fa 100644 --- a/mesa/visualization/components/matplotlib_components.py +++ b/mesa/visualization/components/matplotlib_components.py @@ -38,8 +38,7 @@ def make_mpl_space_component( the functions for drawing the various spaces for further details. ``agent_portrayal`` is called with an agent and should return a dict. Valid fields in this dict are "color", - "size", "marker", and "zorder". Other field are ignored and will result in a user warning. - + "size", "marker", "zorder", alpha, linewidths, and edgecolors. Other field are ignored and will result in a user warning. Returns: function: A function that creates a SpaceMatplotlib component diff --git a/mesa/visualization/mpl_space_drawing.py b/mesa/visualization/mpl_space_drawing.py index cfa1fa23a8b..ea7687c4210 100644 --- a/mesa/visualization/mpl_space_drawing.py +++ b/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) @@ -110,7 +124,7 @@ def draw_space( Returns the Axes object with the plot drawn onto it. ``agent_portrayal`` is called with an agent and should return a dict. Valid fields in this dict are "color", - "size", "marker", and "zorder". Other field are ignored and will result in a user warning. + "size", "marker", "zorder", alpha, linewidths, and edgecolors. Other field are ignored and will result in a user warning. """ if ax is None: @@ -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_voronoi_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], diff --git a/tests/test_components_matplotlib.py b/tests/test_components_matplotlib.py index 78e3c33639e..f258b58d90b 100644 --- a/tests/test_components_matplotlib.py +++ b/tests/test_components_matplotlib.py @@ -23,6 +23,7 @@ draw_network, draw_orthogonal_grid, draw_property_layers, + draw_space, draw_voronoi_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) @@ -62,8 +141,8 @@ def test_draw_hex_grid(): draw_hex_grid(grid, agent_portrayal, ax) -def test_draw_voroinoi_grid(): - """Test drawing voroinoi grids.""" +def test_draw_voronoi_grid(): + """Test drawing voronoi grids.""" model = Model(seed=42) coordinates = model.rng.random((100, 2)) * 10