Skip to content

Commit

Permalink
Enhance agent visualization with additional customization options
Browse files Browse the repository at this point in the history
- Add support for edgecolors, linewidths, and alpha transparency in agent portrayal
- Implement backward compatibility for old keyword arguments
- Refactor _get_agent_data function to handle new visualization parameters
- Update _draw_continuous_space and _draw_voronoi functions to use new agent data format
- Improve code reusability and consistency across different space types

Co-Authored-By: Robert Hopkins <[email protected]>
  • Loading branch information
EwoutH and rmhopkins4 committed Oct 5, 2024
1 parent 47b3bd9 commit 32aec47
Showing 1 changed file with 48 additions and 51 deletions.
99 changes: 48 additions & 51 deletions mesa/visualization/components/matplotlib.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ def SpaceMatplotlib(
elif isinstance(space, mesa.space.NetworkGrid):
_draw_network_grid(space, space_ax, agent_portrayal)
elif isinstance(space, VoronoiGrid):
_draw_voronoi(space, space_ax, agent_portrayal)
_draw_voronoi(space, space_ax, agent_portrayal, model)
elif space is None and propertylayer_portrayal:
draw_property_layers(space_ax, space, propertylayer_portrayal, model)

Expand Down Expand Up @@ -157,21 +157,47 @@ def _draw_grid(space, space_ax, agent_portrayal, propertylayer_portrayal, model)

def _get_agent_data(space, agent_portrayal):
"""Helper function to get agent data for visualization."""
x, y, s, c, m = [], [], [], [], []
x, y, s, c, m, edgecolors, linewidths, alpha = [], [], [], [], [], [], [], []
translation_dict = {

Check warning on line 161 in mesa/visualization/components/matplotlib.py

View check run for this annotation

Codecov / codecov/patch

mesa/visualization/components/matplotlib.py#L160-L161

Added lines #L160 - L161 were not covered by tests
"c": "color",
"s": "size",
"marker": "shape",
"edgecolors": "edgecolor",
"linewidths": "linewidth",
}

for agents, pos in space.coord_iter():
if not agents:
continue
if not isinstance(agents, list):
agents = [agents] # noqa PLW2901
for agent in agents:
data = agent_portrayal(agent)
# Translate old keywords to new keywords if they exist
for key, value in translation_dict.items():
if value in data:
data[key] = data.pop(value)

Check warning on line 179 in mesa/visualization/components/matplotlib.py

View check run for this annotation

Codecov / codecov/patch

mesa/visualization/components/matplotlib.py#L179

Added line #L179 was not covered by tests

x.append(pos[0] + 0.5) # Center the agent in the cell
y.append(pos[1] + 0.5) # Center the agent in the cell
default_size = (180 / max(space.width, space.height)) ** 2
s.append(data.get("size", default_size))
c.append(data.get("color", "tab:blue"))
m.append(data.get("shape", "o"))
return {"x": x, "y": y, "s": s, "c": c, "m": m}
s.append(data.get("s", default_size))
c.append(data.get("c", "tab:blue"))
m.append(data.get("marker", "o"))
edgecolors.append(data.get("edgecolors", "none"))
linewidths.append(data.get("linewidths", 0))
alpha.append(data.get("alpha", 1.0))

Check warning on line 189 in mesa/visualization/components/matplotlib.py

View check run for this annotation

Codecov / codecov/patch

mesa/visualization/components/matplotlib.py#L184-L189

Added lines #L184 - L189 were not covered by tests

return {

Check warning on line 191 in mesa/visualization/components/matplotlib.py

View check run for this annotation

Codecov / codecov/patch

mesa/visualization/components/matplotlib.py#L191

Added line #L191 was not covered by tests
"x": x,
"y": y,
"s": s,
"c": c,
"m": m,
"edgecolors": edgecolors,
"linewidths": linewidths,
"alpha": alpha,
}


def _split_and_scatter(portray_data, space_ax):
Expand All @@ -184,6 +210,11 @@ def _split_and_scatter(portray_data, space_ax):
s=[s for s, show in zip(portray_data["s"], mask) if show],
c=[c for c, show in zip(portray_data["c"], mask) if show],
marker=marker,
edgecolors=[e for e, show in zip(portray_data["edgecolors"], mask) if show],
linewidths=[
lw for lw, show in zip(portray_data["linewidths"], mask) if show
],
alpha=[a for a, show in zip(portray_data["alpha"], mask) if show],
)


Expand All @@ -198,28 +229,9 @@ def _draw_network_grid(space, space_ax, agent_portrayal):
)


def _draw_continuous_space(space, space_ax, agent_portrayal, model):
def portray(space):
x = []
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)

# This is matplotlib's default marker size
default_size = 20
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)
return {"x": x, "y": y, "s": s, "c": c, "m": m}
def _draw_continuous_space(space, space_ax, agent_portrayal):
"""Draw agents in a continuous space."""
agent_data = _get_agent_data(space, agent_portrayal)

Check warning on line 234 in mesa/visualization/components/matplotlib.py

View check run for this annotation

Codecov / codecov/patch

mesa/visualization/components/matplotlib.py#L234

Added line #L234 was not covered by tests

# Determine border style based on space.torus
border_style = "solid" if not space.torus else (0, (5, 10))
Expand All @@ -238,32 +250,14 @@ def portray(space):
space_ax.set_ylim(space.y_min - y_padding, space.y_max + y_padding)

# Portray and scatter the agents in the space
_split_and_scatter(portray(space), space_ax)
_split_and_scatter(agent_data, space_ax)

Check warning on line 253 in mesa/visualization/components/matplotlib.py

View check run for this annotation

Codecov / codecov/patch

mesa/visualization/components/matplotlib.py#L253

Added line #L253 was not covered by tests


def _draw_voronoi(space, space_ax, agent_portrayal):
def portray(g):
x = []
y = []
s = [] # size
c = [] # color

for cell in g.all_cells:
for agent in cell.agents:
data = agent_portrayal(agent)
x.append(cell.coordinate[0])
y.append(cell.coordinate[1])
if "size" in data:
s.append(data["size"])
if "color" in data:
c.append(data["color"])
out = {"x": x, "y": y}
out["s"] = s
if len(c) > 0:
out["c"] = c

return out
"""Draw agents in a Voronoi space."""
agent_data = _get_agent_data(space, agent_portrayal)

Check warning on line 258 in mesa/visualization/components/matplotlib.py

View check run for this annotation

Codecov / codecov/patch

mesa/visualization/components/matplotlib.py#L258

Added line #L258 was not covered by tests

# Set plot limits based on Voronoi centroids
x_list = [i[0] for i in space.centroids_coordinates]
y_list = [i[1] for i in space.centroids_coordinates]
x_max = max(x_list)
Expand All @@ -277,8 +271,11 @@ def portray(g):
y_padding = height / 20
space_ax.set_xlim(x_min - x_padding, x_max + x_padding)
space_ax.set_ylim(y_min - y_padding, y_max + y_padding)
space_ax.scatter(**portray(space))

# Scatter the agent data
_split_and_scatter(agent_data, space_ax)

Check warning on line 276 in mesa/visualization/components/matplotlib.py

View check run for this annotation

Codecov / codecov/patch

mesa/visualization/components/matplotlib.py#L276

Added line #L276 was not covered by tests

# Draw Voronoi cells as polygons
for cell in space.all_cells:
polygon = cell.properties["polygon"]
space_ax.fill(
Expand Down

0 comments on commit 32aec47

Please sign in to comment.