Mesa 3.0.0a2__py3-none-any.whl → 3.0.0a4__py3-none-any.whl

mesa/experimental/solara_viz.py

@@ -0,0 +1,462 @@
+"""
+Mesa visualization module for creating interactive model visualizations.
+
+This module provides components to create browser- and Jupyter notebook-based visualizations of
+Mesa models, allowing users to watch models run step-by-step and interact with model parameters.
+
+Key features:
+    - SolaraViz: Main component for creating visualizations, supporting grid displays and plots
+    - ModelController: Handles model execution controls (step, play, pause, reset)
+    - UserInputs: Generates UI elements for adjusting model parameters
+    - Card: Renders individual visualization elements (space, measures)
+
+The module uses Solara for rendering in Jupyter notebooks or as standalone web applications.
+It supports various types of visualizations including matplotlib plots, agent grids, and
+custom visualization components.
+
+Usage:
+    1. Define an agent_portrayal function to specify how agents should be displayed
+    2. Set up model_params to define adjustable parameters
+    3. Create a SolaraViz instance with your model, parameters, and desired measures
+    4. Display the visualization in a Jupyter notebook or run as a Solara app
+
+See the Visualization Tutorial and example models for more details.
+"""
+
+import threading
+
+import reacton.ipywidgets as widgets
+import solara
+from solara.alias import rv
+
+import mesa.experimental.components.altair as components_altair
+import mesa.experimental.components.matplotlib as components_matplotlib
+from mesa.experimental.UserParam import Slider
+
+
+# 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, dependencies, color, layout_type
+):
+    """
+    Create a card component for visualizing model space or measures.
+
+    Args:
+        model: The Mesa model instance
+        measures: List of measures to be plotted
+        agent_portrayal: Function to define agent appearance
+        space_drawer: Method to render agent space
+        dependencies: List of dependencies for updating the visualization
+        color: Background color of the card
+        layout_type: Type of layout (Space or Measure)
+
+    Returns:
+        rv.Card: A card component containing the visualization
+    """
+    with rv.Card(
+        style_=f"background-color: {color}; width: 100%; height: 100%"
+    ) as main:
+        if "Space" in layout_type:
+            rv.CardTitle(children=["Space"])
+            if space_drawer == "default":
+                # draw with the default implementation
+                components_matplotlib.SpaceMatplotlib(
+                    model, agent_portrayal, dependencies=dependencies
+                )
+            elif space_drawer == "altair":
+                components_altair.SpaceAltair(
+                    model, agent_portrayal, dependencies=dependencies
+                )
+            elif space_drawer:
+                # if specified, draw agent space with an alternate renderer
+                space_drawer(model, agent_portrayal, dependencies=dependencies)
+        elif "Measure" in layout_type:
+            rv.CardTitle(children=["Measure"])
+            measure = measures[layout_type["Measure"]]
+            if callable(measure):
+                # Is a custom object
+                measure(model)
+            else:
+                components_matplotlib.PlotMatplotlib(
+                    model, measure, dependencies=dependencies
+                )
+    return main
+
+
+@solara.component
+def SolaraViz(
+    model_class,
+    model_params,
+    measures=None,
+    name=None,
+    agent_portrayal=None,
+    space_drawer="default",
+    play_interval=150,
+    seed=None,
+):
+    """
+    Initialize a component to visualize a model.
+
+    Args:
+        model_class: Class of the model to instantiate
+        model_params: Parameters for initializing the model
+        measures: List of callables or data attributes to plot
+        name: Name for display
+        agent_portrayal: Options for rendering agents (dictionary);
+            Default drawer supports custom `"size"`, `"color"`, and `"shape"`.
+        space_drawer: Method to render the agent space for
+            the model; default implementation is the `SpaceMatplotlib` component;
+            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__
+
+    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()}}
+    )
+
+    # 2. Set up Model
+    def make_model():
+        """Create a new model instance with current parameters and seed."""
+        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,
+            reactive_seed.value,
+        ],
+    )
+
+    def handle_change_model_params(name: str, value: any):
+        """Update model parameters when user input changes."""
+        set_model_parameters({**model_parameters, name: value})
+
+    # 3. Set up UI
+
+    with solara.AppBar():
+        solara.AppBarTitle(name)
+
+    # render layout and plot
+    def do_reseed():
+        """Update the random seed for the model."""
+        reactive_seed.value = model.random.random()
+
+    dependencies = [
+        *list(model_parameters.values()),
+        current_step.value,
+        reactive_seed.value,
+    ]
+
+    # if space drawer is disabled, do not include it
+    layout_types = [{"Space": "default"}] if space_drawer else []
+
+    if measures:
+        layout_types += [{"Measure": elem} for elem in range(len(measures))]
+
+    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)
+            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(
+            model,
+            measures,
+            agent_portrayal,
+            space_drawer,
+            dependencies,
+            color="white",
+            layout_type=layout_types[i],
+        )
+        for i in range(len(layout_types))
+    ]
+    solara.GridDraggable(
+        items=items,
+        grid_layout=grid_layout,
+        resizable=True,
+        draggable=True,
+        on_grid_layout=set_grid_layout,
+    )
+
+
+JupyterViz = SolaraViz
+
+
+@solara.component
+def ModelController(model, play_interval, current_step, reset_counter):
+    """
+    Create controls for model execution (step, play, pause, reset).
+
+    Args:
+        model: The model being visualized
+        play_interval: Interval between steps during play
+        current_step: Reactive value for the current step
+        reset_counter: Counter to trigger model reset
+    """
+    playing = solara.use_reactive(False)
+    thread = solara.use_reactive(None)
+    # We track the previous step to detect if user resets the model via
+    # clicking the reset button or changing the parameters. If previous_step >
+    # current_step, it means a model reset happens while the simulation is
+    # still playing.
+    previous_step = solara.use_reactive(0)
+
+    def on_value_play(change):
+        """Handle play/pause state changes."""
+        if previous_step.value > current_step.value and current_step.value == 0:
+            # We add extra checks for current_step.value == 0, just to be sure.
+            # We automatically stop the playing if a model is reset.
+            playing.value = False
+        elif model.running:
+            do_step()
+        else:
+            playing.value = False
+
+    def do_step():
+        """Advance the model by one step."""
+        model.step()
+        previous_step.value = current_step.value
+        current_step.value = model.steps
+
+    def do_play():
+        """Run the model continuously."""
+        model.running = True
+        while model.running:
+            do_step()
+
+    def threaded_do_play():
+        """Start a new thread for continuous model execution."""
+        if thread is not None and thread.is_alive():
+            return
+        thread.value = threading.Thread(target=do_play)
+        thread.start()
+
+    def do_pause():
+        """Pause the model execution."""
+        if (thread is None) or (not thread.is_alive()):
+            return
+        model.running = False
+        thread.join()
+
+    def do_reset():
+        """Reset the model."""
+        reset_counter.value += 1
+
+    def do_set_playing(value):
+        """Set the playing state."""
+        if current_step.value == 0:
+            # This means the model has been recreated, and the step resets to
+            # 0. We want to avoid triggering the playing.value = False in the
+            # on_value_play function.
+            previous_step.value = current_step.value
+        playing.set(value)
+
+    with solara.Row():
+        solara.Button(label="Step", color="primary", on_click=do_step)
+        # This style is necessary so that the play widget has almost the same
+        # height as typical Solara buttons.
+        solara.Style(
+            """
+        .widget-play {
+            height: 35px;
+        }
+        .widget-play button {
+            color: white;
+            background-color: #1976D2;  // Solara blue color
+        }
+        """
+        )
+        widgets.Play(
+            value=0,
+            interval=play_interval,
+            repeat=True,
+            show_repeat=False,
+            on_value=on_value_play,
+            playing=playing.value,
+            on_playing=do_set_playing,
+        )
+        solara.Button(label="Reset", color="primary", on_click=do_reset)
+        # threaded_do_play is not used for now because it
+        # doesn't work in Google colab. We use
+        # ipywidgets.Play until it is fixed. The threading
+        # version is definite a much better implementation,
+        # if it works.
+        # solara.Button(label="▶", color="primary", on_click=viz.threaded_do_play)
+        # solara.Button(label="⏸︎", color="primary", on_click=viz.do_pause)
+        # solara.Button(label="Reset", color="primary", on_click=do_reset)
+
+
+def split_model_params(model_params):
+    """
+    Split model parameters into user-adjustable and fixed parameters.
+
+    Args:
+        model_params: Dictionary of all model parameters
+
+    Returns:
+        tuple: (user_adjustable_params, fixed_params)
+    """
+    model_params_input = {}
+    model_params_fixed = {}
+    for k, v in model_params.items():
+        if check_param_is_fixed(v):
+            model_params_fixed[k] = v
+        else:
+            model_params_input[k] = v
+    return model_params_input, model_params_fixed
+
+
+def check_param_is_fixed(param):
+    """
+    Check if a parameter is fixed (not user-adjustable).
+
+    Args:
+        param: Parameter to check
+
+    Returns:
+        bool: True if parameter is fixed, False otherwise
+    """
+    if isinstance(param, Slider):
+        return False
+    if not isinstance(param, dict):
+        return True
+    if "type" not in param:
+        return True
+
+
+@solara.component
+def UserInputs(user_params, on_change=None):
+    """
+    Initialize user inputs for configurable model parameters.
+    Currently supports :class:`solara.SliderInt`, :class:`solara.SliderFloat`,
+    :class:`solara.Select`, and :class:`solara.Checkbox`.
+
+    Args:
+        user_params: Dictionary with options for the input, including label,
+        min and max values, and other fields specific to the input type.
+        on_change: Function to be called with (name, value) when the value of an input changes.
+    """
+
+    for name, options in user_params.items():
+
+        def change_handler(value, name=name):
+            on_change(name, value)
+
+        if isinstance(options, Slider):
+            slider_class = (
+                solara.SliderFloat if options.is_float_slider else solara.SliderInt
+            )
+            slider_class(
+                options.label,
+                value=options.value,
+                on_value=change_handler,
+                min=options.min,
+                max=options.max,
+                step=options.step,
+            )
+            continue
+
+        # label for the input is "label" from options or name
+        label = options.get("label", name)
+        input_type = options.get("type")
+        if input_type == "SliderInt":
+            solara.SliderInt(
+                label,
+                value=options.get("value"),
+                on_value=change_handler,
+                min=options.get("min"),
+                max=options.get("max"),
+                step=options.get("step"),
+            )
+        elif input_type == "SliderFloat":
+            solara.SliderFloat(
+                label,
+                value=options.get("value"),
+                on_value=change_handler,
+                min=options.get("min"),
+                max=options.get("max"),
+                step=options.get("step"),
+            )
+        elif input_type == "Select":
+            solara.Select(
+                label,
+                value=options.get("value"),
+                on_value=change_handler,
+                values=options.get("values"),
+            )
+        elif input_type == "Checkbox":
+            solara.Checkbox(
+                label=label,
+                on_value=change_handler,
+                value=options.get("value"),
+            )
+        else:
+            raise ValueError(f"{input_type} is not a supported input type")
+
+
+def make_text(renderer):
+    """
+    Create a function that renders text using Markdown.
+
+    Args:
+        renderer: Function that takes a model and returns a string
+
+    Returns:
+        function: A function that renders the text as Markdown
+    """
+
+    def function(model):
+        solara.Markdown(renderer(model))
+
+    return function
+
+
+def make_initial_grid_layout(layout_types):
+    """
+    Create an initial grid layout for visualization components.
+
+    Args:
+        layout_types: List of layout types (Space or Measure)
+
+    Returns:
+        list: Initial grid layout configuration
+    """
+    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))
+    ]

mesa/model.py

@@ -8,9 +8,8 @@ Core Objects: Model
 # Remove this __future__ import once the oldest supported Python is 3.10
 from __future__ import annotations
 
-import itertools
 import random
-from collections import defaultdict
+import warnings
 
 # mypy
 from typing import Any
@@ -18,8 +17,6 @@ from typing import Any
 from mesa.agent import Agent, AgentSet
 from mesa.datacollection import DataCollector
 
-TimeT = float | int
-
 
 class Model:
     """Base class for models in the Mesa ABM library.
@@ -31,21 +28,30 @@ class Model:
     Attributes:
         running: A boolean indicating if the model should continue running.
         schedule: An object to manage the order and execution of agent steps.
-        current_id: A counter for assigning unique IDs to agents.
-        agents_: A defaultdict mapping each agent type to a dict of its instances.
-                 This private attribute is used internally to manage agents.
 
     Properties:
-        agents: An AgentSet containing all agents in the model, generated from the _agents attribute.
+        agents: An AgentSet containing all agents in the model
         agent_types: A list of different agent types present in the model.
+        agents_by_type: A dictionary where the keys are agent types and the values are the corresponding AgentSets.
+        steps: An integer representing the number of steps the model has taken.
+               It increases automatically at the start of each step() call.
 
     Methods:
         get_agents_of_type: Returns an AgentSet of agents of the specified type.
+            Deprecated: Use agents_by_type[agenttype] instead.
         run_model: Runs the model's simulation until a defined end condition is reached.
         step: Executes a single step of the model's simulation process.
         next_id: Generates and returns the next unique identifier for an agent.
         reset_randomizer: Resets the model's random number generator with a new or existing seed.
         initialize_data_collector: Sets up the data collector for the model, requiring an initialized scheduler and agents.
+        register_agent : register an agent with the model
+        deregister_agent : remove an agent from the model
+
+    Notes:
+        Model.agents returns the AgentSet containing all agents registered with the model. Changing
+        the content of the AgentSet directly can result in strange behavior. If you want change the
+        composition of this AgentSet, ensure you operate on a copy.
+
     """
 
     def __new__(cls, *args: Any, **kwargs: Any) -> Any:
@@ -57,9 +63,6 @@ class Model:
             # advance.
             obj._seed = random.random()
         obj.random = random.Random(obj._seed)
-        # TODO: Remove these 2 lines just before Mesa 3.0
-        obj._steps = 0
-        obj._time = 0
         return obj
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
@@ -70,40 +73,117 @@ class Model:
 
         self.running = True
         self.schedule = None
-        self.current_id = 0
-        self.agents_: defaultdict[type, dict] = defaultdict(dict)
+        self.steps: int = 0
+
+        self._setup_agent_registration()
 
-        self._steps: int = 0
-        self._time: TimeT = 0  # the model's clock
+        # Wrap the user-defined step method
+        self._user_step = self.step
+        self.step = self._wrapped_step
+
+    def _wrapped_step(self, *args: Any, **kwargs: Any) -> None:
+        """Automatically increments time and steps after calling the user's step method."""
+        # Automatically increment time and step counters
+        self.steps += 1
+        # Call the original user-defined step method
+        self._user_step(*args, **kwargs)
+
+    def next_id(self) -> int:
+        warnings.warn(
+            "using model.next_id() is deprecated. Agents track their unique ID automatically",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return 0
 
     @property
     def agents(self) -> AgentSet:
         """Provides an AgentSet of all agents in the model, combining agents from all types."""
-
-        if hasattr(self, "_agents"):
-            return self._agents
-        else:
-            all_agents = itertools.chain.from_iterable(self.agents_.values())
-            return AgentSet(all_agents, self)
+        return self._all_agents
 
     @agents.setter
     def agents(self, agents: Any) -> None:
         raise AttributeError(
-            "You are trying to set model.agents. In Mesa 3.0 and higher, this attribute will be "
+            "You are trying to set model.agents. In Mesa 3.0 and higher, this attribute is "
             "used by Mesa itself, so you cannot use it directly anymore."
             "Please adjust your code to use a different attribute name for custom agent storage."
         )
 
-        self._agents = agents
-
     @property
     def agent_types(self) -> list[type]:
-        """Return a list of different agent types."""
-        return list(self.agents_.keys())
+        """Return a list of all unique agent types registered with the model."""
+        return list(self._agents_by_type.keys())
+
+    @property
+    def agents_by_type(self) -> dict[type[Agent], AgentSet]:
+        """A dictionary where the keys are agent types and the values are the corresponding AgentSets."""
+        return self._agents_by_type
 
     def get_agents_of_type(self, agenttype: type[Agent]) -> AgentSet:
-        """Retrieves an AgentSet containing all agents of the specified type."""
-        return AgentSet(self.agents_[agenttype].keys(), self)
+        """Deprecated: Retrieves an AgentSet containing all agents of the specified type."""
+        warnings.warn(
+            f"Model.get_agents_of_type() is deprecated, please replace get_agents_of_type({agenttype})"
+            f"with the property agents_by_type[{agenttype}].",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.agents_by_type[agenttype]
+
+    def _setup_agent_registration(self):
+        """helper method to initialize the agent registration datastructures"""
+        self._agents = {}  # the hard references to all agents in the model
+        self._agents_by_type: dict[
+            type[Agent], AgentSet
+        ] = {}  # a dict with an agentset for each class of agents
+        self._all_agents = AgentSet([], self)  # an agenset with all agents
+
+    def register_agent(self, agent):
+        """Register the agent with the model
+
+        Args:
+            agent: The agent to register.
+
+        Notes:
+            This method is called automatically by ``Agent.__init__``, so there is no need to use this
+            if you are subclassing Agent and calling its super in the ``__init__`` method.
+
+        """
+        if not hasattr(self, "_agents"):
+            self._setup_agent_registration()
+
+            warnings.warn(
+                "The Mesa Model class was not initialized. In the future, you need to explicitly initialize "
+                "the Model by calling super().__init__() on initialization.",
+                FutureWarning,
+                stacklevel=2,
+            )
+
+        self._agents[agent] = None
+
+        # because AgentSet requires model, we cannot use defaultdict
+        # tricks with a function won't work because model then cannot be pickled
+        try:
+            self._agents_by_type[type(agent)].add(agent)
+        except KeyError:
+            self._agents_by_type[type(agent)] = AgentSet(
+                [
+                    agent,
+                ],
+                self,
+            )
+
+        self._all_agents.add(agent)
+
+    def deregister_agent(self, agent):
+        """Deregister the agent with the model
+
+        Notes::
+        This method is called automatically by ``Agent.remove``
+
+        """
+        del self._agents[agent]
+        self._agents_by_type[type(agent)].remove(agent)
+        self._all_agents.remove(agent)
 
     def run_model(self) -> None:
         """Run the model until the end condition is reached. Overload as
@@ -115,16 +195,6 @@ class Model:
     def step(self) -> None:
         """A single step. Fill in here."""
 
-    def _advance_time(self, deltat: TimeT = 1):
-        """Increment the model's steps counter and clock."""
-        self._steps += 1
-        self._time += deltat
-
-    def next_id(self) -> int:
-        """Return the next unique ID for agents, increment current_id"""
-        self.current_id += 1
-        return self.current_id
-
     def reset_randomizer(self, seed: int | None = None) -> None:
         """Reset the model random number generator.
 

mesa/space.py

@@ -459,15 +459,21 @@ class _Grid:
             elif selection == "closest":
                 current_pos = agent.pos
                 # Find the closest position without sorting all positions
-                closest_pos = None
+                # TODO: See if this method can be optimized further
+                closest_pos = []
                 min_distance = float("inf")
                 agent.random.shuffle(pos)
                 for p in pos:
                     distance = self._distance_squared(p, current_pos)
                     if distance < min_distance:
                         min_distance = distance
-                        closest_pos = p
-                chosen_pos = closest_pos
+                        closest_pos.clear()
+                        closest_pos.append(p)
+                    elif distance == min_distance:
+                        closest_pos.append(p)
+
+                chosen_pos = agent.random.choice(closest_pos)
+
             else:
                 raise ValueError(
                     f"Invalid selection method {selection}. Choose 'random' or 'closest'."