Skip to content

Commit

Permalink
Add support for alpha, linewidths and edgecolors to agent_portrayal (#…
Browse files Browse the repository at this point in the history
…2468)

* add some more support to agent_portrayal

* Update mpl_space_drawing.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update mpl_space_drawing.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update mpl_space_drawing.py

* some additional docs

* Update test_components_matplotlib.py

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
quaquel and pre-commit-ci[bot] authored Nov 8, 2024
1 parent 46ff9fb commit 60fc1c8
Show file tree
Hide file tree
Showing 3 changed files with 122 additions and 9 deletions.
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

0 comments on commit 60fc1c8

Please sign in to comment.