diff --git a/mesa/experimental/jupyter_viz.py b/mesa/experimental/jupyter_viz.py index 9d29e08fb60..a4dec7f6a9a 100644 --- a/mesa/experimental/jupyter_viz.py +++ b/mesa/experimental/jupyter_viz.py @@ -17,7 +17,7 @@ # TODO: Turn this function into a Solara component once the current_step.value # dependency is passed to measure() def Card( - model, measures, agent_portrayal, space_drawer, current_step, color, layout_type + model, measures, agent_portrayal, space_drawer, dependencies, color, layout_type ): with rv.Card( style_=f"background-color: {color}; width: 100%; height: 100%" @@ -27,11 +27,11 @@ def Card( if space_drawer == "default": # draw with the default implementation components_matplotlib.SpaceMatplotlib( - model, agent_portrayal, dependencies=[current_step.value] + model, agent_portrayal, dependencies=dependencies ) elif space_drawer == "altair": components_altair.SpaceAltair( - model, agent_portrayal, dependencies=[current_step.value] + model, agent_portrayal, dependencies=dependencies ) elif space_drawer: # if specified, draw agent space with an alternate renderer @@ -44,7 +44,7 @@ def Card( measure(model) else: components_matplotlib.PlotMatplotlib( - model, measure, dependencies=[current_step.value] + model, measure, dependencies=dependencies ) return main @@ -58,6 +58,7 @@ def JupyterViz( agent_portrayal=None, space_drawer="default", play_interval=150, + seed=None, ): """Initialize a component to visualize a model. Args: @@ -71,6 +72,7 @@ def JupyterViz( simulations with no space to visualize should specify `space_drawer=False` play_interval: play interval (default: 150) + seed: the random seed used to initialize the model """ if name is None: name = model_class.__name__ @@ -78,6 +80,7 @@ def JupyterViz( current_step = solara.use_reactive(0) # 1. Set up model parameters + reactive_seed = solara.use_reactive(0) user_params, fixed_params = split_model_params(model_params) model_parameters, set_model_parameters = solara.use_state( {**fixed_params, **{k: v.get("value") for k, v in user_params.items()}} @@ -85,13 +88,21 @@ def JupyterViz( # 2. Set up Model def make_model(): - model = model_class(**model_parameters) + model = model_class.__new__( + model_class, **model_parameters, seed=reactive_seed.value + ) + model.__init__(**model_parameters) current_step.value = 0 return model reset_counter = solara.use_reactive(0) model = solara.use_memo( - make_model, dependencies=[*list(model_parameters.values()), reset_counter.value] + make_model, + dependencies=[ + *list(model_parameters.values()), + reset_counter.value, + reactive_seed.value, + ], ) def handle_change_model_params(name: str, value: any): @@ -103,8 +114,12 @@ def handle_change_model_params(name: str, value: any): solara.AppBarTitle(name) # render layout and plot + def do_reseed(): + reactive_seed.value = model.random.random() # jupyter + dependencies = [current_step.value, reactive_seed.value] + def render_in_jupyter(): with solara.GridFixed(columns=2): UserInputs(user_params, on_change=handle_change_model_params) @@ -116,11 +131,11 @@ def render_in_jupyter(): if space_drawer == "default": # draw with the default implementation components_matplotlib.SpaceMatplotlib( - model, agent_portrayal, dependencies=[current_step.value] + model, agent_portrayal, dependencies=dependencies ) elif space_drawer == "altair": components_altair.SpaceAltair( - model, agent_portrayal, dependencies=[current_step.value] + model, agent_portrayal, dependencies=dependencies ) elif space_drawer: # if specified, draw agent space with an alternate renderer @@ -134,7 +149,7 @@ def render_in_jupyter(): measure(model) else: components_matplotlib.PlotMatplotlib( - model, measure, dependencies=[current_step.value] + model, measure, dependencies=dependencies ) def render_in_browser(): @@ -144,15 +159,21 @@ def render_in_browser(): if measures: layout_types += [{"Measure": elem} for elem in range(len(measures))] - grid_layout_initial = get_initial_grid_layout(layout_types=layout_types) + grid_layout_initial = make_initial_grid_layout(layout_types=layout_types) grid_layout, set_grid_layout = solara.use_state(grid_layout_initial) with solara.Sidebar(): with solara.Card("Controls", margin=1, elevation=2): + solara.InputText( + label="Seed", + value=reactive_seed, + continuous_update=True, + ) UserInputs(user_params, on_change=handle_change_model_params) ModelController(model, play_interval, current_step, reset_counter) - with solara.Card("Progress", margin=1, elevation=2): - solara.Markdown(md_text=f"####Step - {current_step}") + solara.Button(label="Reseed", color="primary", on_click=do_reseed) + with solara.Card("Information", margin=1, elevation=2): + solara.Markdown(md_text=f"Step - {current_step}") items = [ Card( @@ -160,7 +181,7 @@ def render_in_browser(): measures, agent_portrayal, space_drawer, - current_step, + dependencies, color="white", layout_type=layout_types[i], ) @@ -241,7 +262,11 @@ def do_set_playing(value): solara.Style( """ .widget-play { - height: 30px; + height: 35px; + } + .widget-play button { + color: white; + background-color: #1976D2; // Solara blue color } """ ) @@ -361,20 +386,15 @@ def function(model): return function -def get_initial_grid_layout(layout_types): - grid_lay = [] - y_coord = 0 - for ii in range(len(layout_types)): - template_layout = {"h": 10, "i": 0, "moved": False, "w": 6, "y": 0, "x": 0} - if ii == 0: - grid_lay.append(template_layout) - else: - template_layout.update({"i": ii}) - if ii % 2 == 0: - template_layout.update({"x": 0}) - y_coord += 16 - else: - template_layout.update({"x": 6}) - template_layout.update({"y": y_coord}) - grid_lay.append(template_layout) - return grid_lay +def make_initial_grid_layout(layout_types): + return [ + { + "i": i, + "w": 6, + "h": 10, + "moved": False, + "x": 6 * (i % 2), + "y": 16 * (i - i % 2), + } + for i in range(len(layout_types)) + ] diff --git a/pyproject.toml b/pyproject.toml index 0233bc5874b..087bf4cb5d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,7 @@ dev = [ "pytest >= 4.6", "pytest-cov", "sphinx", + "pytest-mock", ] docs = [ "sphinx", diff --git a/tests/test_jupyter_viz.py b/tests/test_jupyter_viz.py index 248e8c319d5..f67daa09518 100644 --- a/tests/test_jupyter_viz.py +++ b/tests/test_jupyter_viz.py @@ -1,9 +1,10 @@ import unittest -from unittest.mock import Mock, patch +from unittest.mock import Mock import ipyvuetify as vw import solara +import mesa from mesa.experimental.jupyter_viz import JupyterViz, Slider, UserInputs @@ -81,58 +82,61 @@ def Test(user_params): assert slider_int.step is None -class TestJupyterViz(unittest.TestCase): - @patch("mesa.experimental.components.matplotlib.SpaceMatplotlib") - def test_call_space_drawer(self, mock_space_matplotlib): - mock_model_class = Mock() - mock_model_class.__name__ = "MockModelClass" - agent_portrayal = { - "Shape": "circle", - "color": "gray", - } - current_step = 0 - dependencies = [current_step] - # initialize with space drawer unspecified (use default) - # component must be rendered for code to run - solara.render( - JupyterViz( - model_class=mock_model_class, - model_params={}, - agent_portrayal=agent_portrayal, - ) +def test_call_space_drawer(mocker): + mock_space_matplotlib = mocker.patch( + "mesa.experimental.components.matplotlib.SpaceMatplotlib" + ) + + model = mesa.Model() + mocker.patch.object(mesa.Model, "__new__", return_value=model) + mocker.patch.object(mesa.Model, "__init__", return_value=None) + + agent_portrayal = { + "Shape": "circle", + "color": "gray", + } + current_step = 0 + seed = 0 + dependencies = [current_step, seed] + # initialize with space drawer unspecified (use default) + # component must be rendered for code to run + solara.render( + JupyterViz( + model_class=mesa.Model, + model_params={}, + agent_portrayal=agent_portrayal, ) - # should call default method with class instance and agent portrayal - mock_space_matplotlib.assert_called_with( - mock_model_class.return_value, agent_portrayal, dependencies=dependencies - ) - - # specify no space should be drawn; any false value should work - for falsy_value in [None, False, 0]: - mock_space_matplotlib.reset_mock() - solara.render( - JupyterViz( - model_class=mock_model_class, - model_params={}, - agent_portrayal=agent_portrayal, - space_drawer=falsy_value, - ) - ) - # should call default method with class instance and agent portrayal - assert mock_space_matplotlib.call_count == 0 - - # specify a custom space method - altspace_drawer = Mock() + ) + # should call default method with class instance and agent portrayal + mock_space_matplotlib.assert_called_with( + model, agent_portrayal, dependencies=dependencies + ) + + # specify no space should be drawn; any false value should work + for falsy_value in [None, False, 0]: + mock_space_matplotlib.reset_mock() solara.render( JupyterViz( - model_class=mock_model_class, + model_class=mesa.Model, model_params={}, agent_portrayal=agent_portrayal, - space_drawer=altspace_drawer, + space_drawer=falsy_value, ) ) - altspace_drawer.assert_called_with( - mock_model_class.return_value, agent_portrayal + # should call default method with class instance and agent portrayal + assert mock_space_matplotlib.call_count == 0 + + # specify a custom space method + altspace_drawer = Mock() + solara.render( + JupyterViz( + model_class=mesa.Model, + model_params={}, + agent_portrayal=agent_portrayal, + space_drawer=altspace_drawer, ) + ) + altspace_drawer.assert_called_with(model, agent_portrayal) def test_slider():