Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for alpha, linewidths and edgecolors to agent_portrayal #2468

Merged
merged 9 commits into from
Nov 8, 2024
3 changes: 1 addition & 2 deletions mesa/visualization/components/matplotlib_components.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
45 changes: 40 additions & 5 deletions mesa/visualization/mpl_space_drawing.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

"""

import contextlib
import itertools
import math
import warnings
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -110,24 +124,32 @@ 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:
fig, ax = plt.subplots()

# 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)
Expand Down Expand Up @@ -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],
Expand Down
83 changes: 81 additions & 2 deletions tests/test_components_matplotlib.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
draw_network,
draw_orthogonal_grid,
draw_property_layers,
draw_space,
draw_voronoi_grid,
)

Expand All @@ -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)
Expand All @@ -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
Expand Down
Loading