JupyterViz: allow direct access to model instance, both from Solara interface and Python interpreter

[EwoutH]

Update: Further discussed here:

• Access and modify model from Solara visualisation. #2176

Objective

Modify the JupyterViz component to allow direct access and modification of the model instance, both from the Solara interface and the Python interpreter.

Basically allow things like:

page.model.step()
page.model.datacollector.get_model_vars_dataframe()

which ideally, also updates the interface and visa versa.

Rationale

Currently, the JupyterViz component encapsulates the model, making it difficult for users to interact with or modify the model directly. This enhancement will provide greater flexibility and control, allowing users to:

1. Access the model state directly from their Python code.
2. Modify the model attributes or parameters programmatically.
3. Trigger model actions (step, reset) from outside the Solara interface.

These capabilities will be particularly useful for debugging, advanced simulations, and integrating Mesa models with external systems or workflows.

Implementation

1. Refactor JupyterViz as a Class-Based Component

Convert JupyterViz from a function to a class-based component. This allows us to maintain state and provide methods for external interaction.

@solara.component
class JupyterViz:
    def __init__(self, model_class, model_params, ...):
        # Initialize component attributes
        
    def render(self):
        # Render the visualization

2. Use Solara State for Model Management

Utilize Solara's state management to handle the model instance, ensuring updates are reflected in the UI.

self.model, self.set_model = solara.use_state(self.make_model())

3. Implement Model Interaction Methods

Add methods for common model interactions:

def update_model(self, **kwargs):
    """Update model attributes."""
    new_model = self.model.value
    for key, value in kwargs.items():
        setattr(new_model, key, value)
    self.set_model(new_model)

def step_model(self):
    """Step the model forward."""
    new_model = self.model.value
    new_model.step()
    self.current_step.value += 1
    self.set_model(new_model)

def reset_model(self):
    """Reset the model."""
    self.set_model(self.make_model())
    self.reset_counter.value += 1

4. Expose Model Parameters

Make model parameters accessible and modifiable:

def handle_change_model_params(self, name: str, value: any):
    """Update model parameters and recreate the model."""
    new_params = {**self.model_parameters, name: value}
    self.set_model_parameters(new_params)
    self.set_model(self.model_class(**new_params, seed=self.reactive_seed.value))

5. Update Rendering Logic

Modify the render method to use the new state-based model and provide basic UI controls:

def render(self):
    with solara.AppBar():
        solara.AppBarTitle(self.name)

    with solara.Column():
        with solara.Card("Model Controls"):
            solara.Button("Step", on_click=self.step_model)
            solara.Button("Reset", on_click=self.reset_model)
        
        # Add more UI components as needed

    # Render visualization components
    if self.space_drawer == "default":
        components_matplotlib.SpaceMatplotlib(
            self.model.value, self.agent_portrayal, 
            dependencies=[self.current_step.value]
        )

    # Add more visualization components as needed

Usage Examples

Initialization

page = JupyterViz(
    LifeModel,
    pulsar_params,
    measures=None,
    name="Game of Life Pulsar Pattern",
    agent_portrayal=agent_portrayal,
)

Interacting with the Model from Python

# Access the current model state
current_model = page.model.value

# Modify a model attribute
page.update_model(grid_size=(20, 20))

# Step the model
page.step_model()

# Reset the model
page.reset_model()

# Change a model parameter
page.handle_change_model_params("initial_density", 0.4)

Extending the Component

Users can extend the JupyterViz class to add custom functionality:

class CustomJupyterViz(JupyterViz):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
    
    def custom_action(self):
        # Implement custom model interaction
        pass

    def render(self):
        super().render()
        # Add custom UI elements
        solara.Button("Custom Action", on_click=self.custom_action)

[EwoutH]

Closing this discussion in favor of:

• Access and modify model from Solara visualisation. #2176