From 88a394ac055aa612768eaa6e421a37476050d303 Mon Sep 17 00:00:00 2001 From: rht Date: Mon, 21 Aug 2023 09:03:38 -0400 Subject: [PATCH] solara: Implement decorator to simplify custom space drawer --- mesa/experimental/__init__.py | 2 +- mesa/experimental/jupyter_viz.py | 17 ++++++++++++----- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/mesa/experimental/__init__.py b/mesa/experimental/__init__.py index 964dc5d19a3..71d6a5dd53f 100644 --- a/mesa/experimental/__init__.py +++ b/mesa/experimental/__init__.py @@ -1 +1 @@ -from .jupyter_viz import JupyterViz, make_text # noqa +from .jupyter_viz import JupyterViz, make_text, prepare_matplotlib_space # noqa diff --git a/mesa/experimental/jupyter_viz.py b/mesa/experimental/jupyter_viz.py index 7a49d7aa64f..c092ef7e241 100644 --- a/mesa/experimental/jupyter_viz.py +++ b/mesa/experimental/jupyter_viz.py @@ -105,15 +105,22 @@ def _draw_network_grid(viz, space_ax): ) -def make_space(viz): - space_fig = Figure() - space_ax = space_fig.subplots() +def prepare_matplotlib_space(drawer): + def wrapped_drawer(viz): + space_fig = Figure() + space_ax = space_fig.subplots() + drawer(viz, space_fig, space_ax) + space_ax.set_axis_off() + solara.FigureMatplotlib(space_fig, dependencies=[viz.model, viz.df]) + return wrapped_drawer + + +@prepare_matplotlib_space +def make_space(viz, space_fig, space_ax): if isinstance(viz.model.grid, mesa.space.NetworkGrid): _draw_network_grid(viz, space_ax) else: space_ax.scatter(**viz.portray(viz.model.grid)) - space_ax.set_axis_off() - solara.FigureMatplotlib(space_fig, dependencies=[viz.model, viz.df]) def make_plot(viz, measure):