Skip to content

Commit

Permalink
solara: Implement drawer for continuous space
Browse files Browse the repository at this point in the history
  • Loading branch information
rht authored and tpike3 committed Oct 15, 2023
1 parent 3cbe56f commit 477778c
Showing 1 changed file with 36 additions and 5 deletions.
41 changes: 36 additions & 5 deletions mesa/experimental/jupyter_viz.py
Original file line number Diff line number Diff line change
Expand Up @@ -265,16 +265,22 @@ def portray(g):

space_fig = Figure()
space_ax = space_fig.subplots()
if isinstance(model.grid, mesa.space.NetworkGrid):
_draw_network_grid(model, space_ax, agent_portrayal)
space = getattr(model, "grid", None)
if space is None:
# Sometimes the space is defined as model.space instead of model.grid
space = model.space
if isinstance(space, mesa.space.NetworkGrid):
_draw_network_grid(space, space_ax, agent_portrayal)
elif isinstance(space, mesa.space.ContinuousSpace):
_draw_continuous_space(space, space_ax, agent_portrayal)
else:
space_ax.scatter(**portray(model.grid))
space_ax.scatter(**portray(space))
space_ax.set_axis_off()
solara.FigureMatplotlib(space_fig, format="png")


def _draw_network_grid(model, space_ax, agent_portrayal):
graph = model.grid.G
def _draw_network_grid(space, space_ax, agent_portrayal):
graph = space.G
pos = nx.spring_layout(graph, seed=0)
nx.draw(
graph,
Expand All @@ -284,6 +290,31 @@ def _draw_network_grid(model, space_ax, agent_portrayal):
)


def _draw_continuous_space(space, space_ax, agent_portrayal):
def portray(space):
x = []
y = []
s = [] # size
c = [] # color
for agent in space._agent_to_index:
data = agent_portrayal(agent)
_x, _y = agent.pos
x.append(_x)
y.append(_y)
if "size" in data:
s.append(data["size"])
if "color" in data:
c.append(data["color"])
out = {"x": x, "y": y}
if len(s) > 0:
out["s"] = s
if len(c) > 0:
out["c"] = c
return out

space_ax.scatter(**portray(space))


def make_plot(model, measure):
fig = Figure()
ax = fig.subplots()
Expand Down

0 comments on commit 477778c

Please sign in to comment.