Mesa 3.0.0b2__py3-none-any.whl → 3.0.1__py3-none-any.whl

mesa/examples/basic/boid_flockers/model.py

@@ -1,60 +1,81 @@
-"""Flockers.
-=============================================================
+"""
+Boids Flocking Model
+===================
 A Mesa implementation of Craig Reynolds's Boids flocker model.
 Uses numpy arrays to represent vectors.
 """
 
 import numpy as np
 
-import mesa
+from mesa import Model
 from mesa.examples.basic.boid_flockers.agents import Boid
+from mesa.space import ContinuousSpace
 
 
-class BoidFlockers(mesa.Model):
+class BoidFlockers(Model):
     """Flocker model class. Handles agent creation, placement and scheduling."""
 
     def __init__(
         self,
-        seed=None,
         population=100,
         width=100,
         height=100,
-        vision=10,
         speed=1,
-        separation=1,
+        vision=10,
+        separation=2,
         cohere=0.03,
         separate=0.015,
         match=0.05,
+        seed=None,
     ):
-        """Create a new Flockers model.
+        """Create a new Boids Flocking model.
 
         Args:
-            population: Number of Boids
-            width, height: Size of the space.
-            speed: How fast should the Boids move.
-            vision: How far around should each Boid look for its neighbors
-            separation: What's the minimum distance each Boid will attempt to
-                    keep from any other
-            cohere, separate, match: factors for the relative importance of
-                    the three drives.
+            population: Number of Boids in the simulation (default: 100)
+            width: Width of the space (default: 100)
+            height: Height of the space (default: 100)
+            speed: How fast the Boids move (default: 1)
+            vision: How far each Boid can see (default: 10)
+            separation: Minimum distance between Boids (default: 2)
+            cohere: Weight of cohesion behavior (default: 0.03)
+            separate: Weight of separation behavior (default: 0.015)
+            match: Weight of alignment behavior (default: 0.05)
+            seed: Random seed for reproducibility (default: None)
         """
         super().__init__(seed=seed)
+
+        # Model Parameters
         self.population = population
         self.vision = vision
         self.speed = speed
         self.separation = separation
 
-        self.space = mesa.space.ContinuousSpace(width, height, True)
+        # Set up the space
+        self.space = ContinuousSpace(width, height, torus=True)
+
+        # Store flocking weights
         self.factors = {"cohere": cohere, "separate": separate, "match": match}
+
+        # Create and place the Boid agents
         self.make_agents()
 
+        # For tracking statistics
+        self.average_heading = None
+        self.update_average_heading()
+
     def make_agents(self):
-        """Create self.population agents, with random positions and starting headings."""
+        """Create and place all Boid agents randomly in the space."""
         for _ in range(self.population):
+            # Random position
             x = self.random.random() * self.space.x_max
             y = self.random.random() * self.space.y_max
             pos = np.array((x, y))
-            direction = np.random.random(2) * 2 - 1
+
+            # Random initial direction
+            direction = np.random.random(2) * 2 - 1  # Random vector between -1 and 1
+            direction /= np.linalg.norm(direction)  # Normalize
+
+            # Create and place the Boid
             boid = Boid(
                 model=self,
                 speed=self.speed,
@@ -65,5 +86,20 @@ class BoidFlockers(mesa.Model):
             )
             self.space.place_agent(boid, pos)
 
+    def update_average_heading(self):
+        """Calculate the average heading (direction) of all Boids."""
+        if not self.agents:
+            self.average_heading = 0
+            return
+
+        headings = np.array([agent.direction for agent in self.agents])
+        mean_heading = np.mean(headings, axis=0)
+        self.average_heading = np.arctan2(mean_heading[1], mean_heading[0])
+
     def step(self):
+        """Run one step of the model.
+
+        All agents are activated in random order using the AgentSet shuffle_do method.
+        """
         self.agents.shuffle_do("step")
+        self.update_average_heading()

mesa/examples/basic/boltzmann_wealth_model/agents.py

@@ -2,13 +2,26 @@ from mesa import Agent
 
 
 class MoneyAgent(Agent):
-    """An agent with fixed initial wealth."""
+    """An agent with fixed initial wealth.
+
+    Each agent starts with 1 unit of wealth and can give 1 unit to other agents
+    if they occupy the same cell.
+
+    Attributes:
+        wealth (int): The agent's current wealth (starts at 1)
+    """
 
     def __init__(self, model):
+        """Create a new agent.
+
+        Args:
+            model (Model): The model instance that contains the agent
+        """
         super().__init__(model)
         self.wealth = 1
 
     def move(self):
+        """Move the agent to a random neighboring cell."""
         possible_steps = self.model.grid.get_neighborhood(
             self.pos, moore=True, include_center=False
         )
@@ -16,16 +29,21 @@ class MoneyAgent(Agent):
         self.model.grid.move_agent(self, new_position)
 
     def give_money(self):
+        """Give 1 unit of wealth to a random agent in the same cell."""
         cellmates = self.model.grid.get_cell_list_contents([self.pos])
-        cellmates.pop(
-            cellmates.index(self)
-        )  # Ensure agent is not giving money to itself
-        if len(cellmates) > 0:
+        # Remove self from potential recipients
+        cellmates.pop(cellmates.index(self))
+
+        if cellmates:  # Only give money if there are other agents present
             other = self.random.choice(cellmates)
             other.wealth += 1
             self.wealth -= 1
 
     def step(self):
+        """Execute one step for the agent:
+        1. Move to a neighboring cell
+        2. If wealth > 0, maybe give money to another agent in the same cell
+        """
         self.move()
         if self.wealth > 0:
             self.give_money()

mesa/examples/basic/boltzmann_wealth_model/app.py

@@ -1,18 +1,14 @@
-from mesa.examples.basic.boltzmann_wealth_model.model import BoltzmannWealthModel
+from mesa.examples.basic.boltzmann_wealth_model.model import BoltzmannWealth
 from mesa.visualization import (
     SolaraViz,
-    make_plot_measure,
-    make_space_matplotlib,
+    make_plot_component,
+    make_space_component,
 )
 
 
 def agent_portrayal(agent):
-    size = 10
-    color = "tab:red"
-    if agent.wealth > 0:
-        size = 50
-        color = "tab:blue"
-    return {"size": size, "color": color}
+    color = agent.wealth  # we are using a colormap to translate wealth to color
+    return {"color": color}
 
 
 model_params = {
@@ -24,20 +20,33 @@ model_params = {
         "max": 100,
         "step": 1,
     },
+    "seed": {
+        "type": "InputText",
+        "value": 42,
+        "label": "Random Seed",
+    },
     "width": 10,
     "height": 10,
 }
 
+
+def post_process(ax):
+    ax.get_figure().colorbar(ax.collections[0], label="wealth", ax=ax)
+
+
 # Create initial model instance
-model1 = BoltzmannWealthModel(50, 10, 10)
+model = BoltzmannWealth(50, 10, 10)
 
 # Create visualization elements. The visualization elements are solara components
 # that receive the model instance as a "prop" and display it in a certain way.
 # Under the hood these are just classes that receive the model instance.
 # You can also author your own visualization elements, which can also be functions
 # that receive the model instance and return a valid solara component.
-SpaceGraph = make_space_matplotlib(agent_portrayal)
-GiniPlot = make_plot_measure("Gini")
+
+SpaceGraph = make_space_component(
+    agent_portrayal, cmap="viridis", vmin=0, vmax=10, post_process=post_process
+)
+GiniPlot = make_plot_component("Gini")
 
 # Create the SolaraViz page. This will automatically create a server and display the
 # visualization elements in a web browser.
@@ -45,7 +54,7 @@ GiniPlot = make_plot_measure("Gini")
 # solara run app.py
 # It will automatically update and display any changes made to this file
 page = SolaraViz(
-    model1,
+    model,
     components=[SpaceGraph, GiniPlot],
     model_params=model_params,
     name="Boltzmann Wealth Model",

mesa/examples/basic/boltzmann_wealth_model/model.py

@@ -1,43 +1,78 @@
-import mesa
+"""
+Boltzmann Wealth Model
+=====================
+
+A simple model of wealth distribution based on the Boltzmann-Gibbs distribution.
+Agents move randomly on a grid, giving one unit of wealth to a random neighbor
+when they occupy the same cell.
+"""
+
+from mesa import Model
+from mesa.datacollection import DataCollector
 from mesa.examples.basic.boltzmann_wealth_model.agents import MoneyAgent
+from mesa.space import MultiGrid
 
 
-class BoltzmannWealthModel(mesa.Model):
+class BoltzmannWealth(Model):
     """A simple model of an economy where agents exchange currency at random.
 
-    All the agents begin with one unit of currency, and each time step can give
-    a unit of currency to another agent. Note how, over time, this produces a
-    highly skewed distribution of wealth.
+    All agents begin with one unit of currency, and each time step agents can give
+    a unit of currency to another agent in the same cell. Over time, this produces
+    a highly skewed distribution of wealth.
+
+    Attributes:
+        num_agents (int): Number of agents in the model
+        grid (MultiGrid): The space in which agents move
+        running (bool): Whether the model should continue running
+        datacollector (DataCollector): Collects and stores model data
     """
 
     def __init__(self, n=100, width=10, height=10, seed=None):
+        """Initialize the model.
+
+        Args:
+            n (int, optional): Number of agents. Defaults to 100.
+            width (int, optional): Grid width. Defaults to 10.
+            height (int, optional): Grid height. Defaults to 10.
+            seed (int, optional): Random seed. Defaults to None.
+        """
         super().__init__(seed=seed)
+
         self.num_agents = n
-        self.grid = mesa.space.MultiGrid(width, height, True)
+        self.grid = MultiGrid(width, height, torus=True)
 
-        self.datacollector = mesa.DataCollector(
+        # Set up data collection
+        self.datacollector = DataCollector(
             model_reporters={"Gini": self.compute_gini},
             agent_reporters={"Wealth": "wealth"},
         )
-        # Create agents
+
+        # Create and place the agents
         for _ in range(self.num_agents):
-            a = MoneyAgent(self)
+            agent = MoneyAgent(self)
 
-            # Add the agent to a random grid cell
+            # Add agent to random grid cell
             x = self.random.randrange(self.grid.width)
             y = self.random.randrange(self.grid.height)
-            self.grid.place_agent(a, (x, y))
+            self.grid.place_agent(agent, (x, y))
 
         self.running = True
         self.datacollector.collect(self)
 
     def step(self):
-        self.agents.shuffle_do("step")
-        self.datacollector.collect(self)
+        self.agents.shuffle_do("step")  # Activate all agents in random order
+        self.datacollector.collect(self)  # Collect data
 
     def compute_gini(self):
+        """Calculate the Gini coefficient for the model's current wealth distribution.
+
+        The Gini coefficient is a measure of inequality in distributions.
+        - A Gini of 0 represents complete equality, where all agents have equal wealth.
+        - A Gini of 1 represents maximal inequality, where one agent has all wealth.
+        """
         agent_wealths = [agent.wealth for agent in self.agents]
         x = sorted(agent_wealths)
         n = self.num_agents
+        # Calculate using the standard formula for Gini coefficient
         b = sum(xi * (n - i) for i, xi in enumerate(x)) / (n * sum(x))
         return 1 + (1 / n) - 2 * b

mesa/examples/basic/boltzmann_wealth_model/st_app.py

@@ -5,7 +5,7 @@ import time
 import altair as alt
 import pandas as pd
 import streamlit as st
-from model import BoltzmannWealthModel
+from model import BoltzmannWealth
 
 model = st.title("Boltzman Wealth Model")
 num_agents = st.slider(
@@ -19,7 +19,7 @@ num_ticks = st.slider(
 )
 height = st.slider("Select Grid Height", min_value=10, max_value=100, step=10, value=15)
 width = st.slider("Select Grid Width", min_value=10, max_value=100, step=10, value=20)
-model = BoltzmannWealthModel(num_agents, height, width)
+model = BoltzmannWealth(num_agents, height, width)
 
 
 status_text = st.empty()

mesa/examples/basic/conways_game_of_life/app.py

@@ -1,12 +1,22 @@
 from mesa.examples.basic.conways_game_of_life.model import ConwaysGameOfLife
 from mesa.visualization import (
     SolaraViz,
-    make_space_matplotlib,
+    make_space_component,
 )
 
 
 def agent_portrayal(agent):
-    return {"color": "white" if agent.state == 0 else "black"}
+    return {
+        "color": "white" if agent.state == 0 else "black",
+        "marker": "s",
+        "size": 25,
+    }
+
+
+def post_process(ax):
+    ax.set_aspect("equal")
+    ax.set_xticks([])
+    ax.set_yticks([])
 
 
 model_params = {
@@ -22,7 +32,9 @@ model1 = ConwaysGameOfLife(50, 50)
 # Under the hood these are just classes that receive the model instance.
 # You can also author your own visualization elements, which can also be functions
 # that receive the model instance and return a valid solara component.
-SpaceGraph = make_space_matplotlib(agent_portrayal)
+SpaceGraph = make_space_component(
+    agent_portrayal, post_process=post_process, draw_grid=False
+)
 
 
 # Create the SolaraViz page. This will automatically create a server and display the

mesa/examples/basic/schelling/agents.py

@@ -1,25 +1,29 @@
-from mesa import Agent, Model
+from mesa import Agent
 
 
 class SchellingAgent(Agent):
     """Schelling segregation agent."""
 
-    def __init__(self, model: Model, agent_type: int) -> None:
+    def __init__(self, model, agent_type: int) -> None:
         """Create a new Schelling agent.
 
         Args:
-           agent_type: Indicator for the agent's type (minority=1, majority=0)
+            model: The model instance the agent belongs to
+            agent_type: Indicator for the agent's type (minority=1, majority=0)
         """
         super().__init__(model)
         self.type = agent_type
 
     def step(self) -> None:
+        """Determine if agent is happy and move if necessary."""
         neighbors = self.model.grid.iter_neighbors(
             self.pos, moore=True, radius=self.model.radius
         )
-        similar = sum(1 for neighbor in neighbors if neighbor.type == self.type)
 
-        # If unhappy, move:
+        # Count similar neighbors
+        similar = sum(neighbor.type == self.type for neighbor in neighbors)
+
+        # If unhappy, move to a random empty cell:
         if similar < self.model.homophily:
             self.model.grid.move_to_empty(self)
         else:

mesa/examples/basic/schelling/app.py

@@ -4,8 +4,8 @@ from mesa.examples.basic.schelling.model import Schelling
 from mesa.visualization import (
     Slider,
     SolaraViz,
-    make_plot_measure,
-    make_space_matplotlib,
+    make_plot_component,
+    make_space_component,
 )
 
 
@@ -28,13 +28,13 @@ model_params = {
 
 model1 = Schelling(20, 20, 0.8, 0.2, 3)
 
-HappyPlot = make_plot_measure("happy")
+HappyPlot = make_plot_component({"happy": "tab:green"})
 
 page = SolaraViz(
     model1,
     components=[
-        make_space_matplotlib(agent_portrayal),
-        make_plot_measure("happy"),
+        make_space_component(agent_portrayal),
+        HappyPlot,
         get_happy_agents,
     ],
     model_params=model_params,

mesa/examples/basic/schelling/model.py

@@ -1,6 +1,7 @@
-import mesa
 from mesa import Model
+from mesa.datacollection import DataCollector
 from mesa.examples.basic.schelling.agents import SchellingAgent
+from mesa.space import SingleGrid
 
 
 class Schelling(Model):
@@ -8,52 +9,73 @@ class Schelling(Model):
 
     def __init__(
         self,
-        height=20,
-        width=20,
-        homophily=3,
-        radius=1,
-        density=0.8,
-        minority_pc=0.2,
+        height: int = 40,
+        width: int = 40,
+        density: float = 0.8,
+        minority_pc: float = 0.5,
+        homophily: int = 3,
+        radius: int = 1,
         seed=None,
     ):
         """Create a new Schelling model.
 
         Args:
-            width, height: Size of the space.
-            density: Initial Chance for a cell to populated
-            minority_pc: Chances for an agent to be in minority class
-            homophily: Minimum number of agents of same class needed to be happy
-            radius: Search radius for checking similarity
-            seed: Seed for Reproducibility
+            width: Width of the grid
+            height: Height of the grid
+            density: Initial chance for a cell to be populated (0-1)
+            minority_pc: Chance for an agent to be in minority class (0-1)
+            homophily: Minimum number of similar neighbors needed for happiness
+            radius: Search radius for checking neighbor similarity
+            seed: Seed for reproducibility
         """
         super().__init__(seed=seed)
+
+        # Model parameters
+        self.height = height
+        self.width = width
+        self.density = density
+        self.minority_pc = minority_pc
         self.homophily = homophily
         self.radius = radius
 
-        self.grid = mesa.space.SingleGrid(width, height, torus=True)
+        # Initialize grid
+        self.grid = SingleGrid(width, height, torus=True)
 
+        # Track happiness
         self.happy = 0
-        self.datacollector = mesa.DataCollector(
-            model_reporters={"happy": "happy"},  # Model-level count of happy agents
+
+        # Set up data collection
+        self.datacollector = DataCollector(
+            model_reporters={
+                "happy": "happy",
+                "pct_happy": lambda m: (m.happy / len(m.agents)) * 100
+                if len(m.agents) > 0
+                else 0,
+                "population": lambda m: len(m.agents),
+                "minority_pct": lambda m: (
+                    sum(1 for agent in m.agents if agent.type == 1)
+                    / len(m.agents)
+                    * 100
+                    if len(m.agents) > 0
+                    else 0
+                ),
+            },
+            agent_reporters={"agent_type": "type"},
         )
 
-        # Set up agents
-        # We use a grid iterator that returns
-        # the coordinates of a cell as well as
-        # its contents. (coord_iter)
+        # Create agents and place them on the grid
         for _, pos in self.grid.coord_iter():
-            if self.random.random() < density:
+            if self.random.random() < self.density:
                 agent_type = 1 if self.random.random() < minority_pc else 0
                 agent = SchellingAgent(self, agent_type)
                 self.grid.place_agent(agent, pos)
 
+        # Collect initial state
         self.datacollector.collect(self)
 
     def step(self):
         """Run one step of the model."""
         self.happy = 0  # Reset counter of happy agents
-        self.agents.shuffle_do("step")
-
-        self.datacollector.collect(self)
-
-        self.running = self.happy != len(self.agents)
+        self.agents.shuffle_do("step")  # Activate all agents in random order
+        self.datacollector.collect(self)  # Collect data
+        self.running = self.happy < len(self.agents)  # Continue until everyone is happy

mesa/examples/basic/virus_on_network/app.py

@@ -1,44 +1,27 @@
 import math
 
 import solara
-from matplotlib.figure import Figure
-from matplotlib.ticker import MaxNLocator
 
 from mesa.examples.basic.virus_on_network.model import (
     State,
     VirusOnNetwork,
     number_infected,
 )
-from mesa.visualization import Slider, SolaraViz, make_space_matplotlib
-
+from mesa.visualization import (
+    Slider,
+    SolaraViz,
+    make_plot_component,
+    make_space_component,
+)
 
-def agent_portrayal(graph):
-    def get_agent(node):
-        return graph.nodes[node]["agent"][0]
 
-    edge_width = []
-    edge_color = []
-    for u, v in graph.edges():
-        agent1 = get_agent(u)
-        agent2 = get_agent(v)
-        w = 2
-        ec = "#e8e8e8"
-        if State.RESISTANT in (agent1.state, agent2.state):
-            w = 3
-            ec = "black"
-        edge_width.append(w)
-        edge_color.append(ec)
+def agent_portrayal(agent):
     node_color_dict = {
         State.INFECTED: "tab:red",
         State.SUSCEPTIBLE: "tab:green",
         State.RESISTANT: "tab:gray",
     }
-    node_color = [node_color_dict[get_agent(node).state] for node in graph.nodes()]
-    return {
-        "width": edge_width,
-        "edge_color": edge_color,
-        "node_color": node_color,
-    }
+    return {"color": node_color_dict[agent.state], "size": 10}
 
 
 def get_resistant_susceptible_ratio(model):
@@ -46,25 +29,9 @@ def get_resistant_susceptible_ratio(model):
     ratio_text = r"$\infty$" if ratio is math.inf else f"{ratio:.2f}"
     infected_text = str(number_infected(model))
 
-    return f"Resistant/Susceptible Ratio: {ratio_text}<br>Infected Remaining: {infected_text}"
-
-
-def make_plot(model):
-    # This is for the case when we want to plot multiple measures in 1 figure.
-    fig = Figure()
-    ax = fig.subplots()
-    measures = ["Infected", "Susceptible", "Resistant"]
-    colors = ["tab:red", "tab:green", "tab:gray"]
-    for i, m in enumerate(measures):
-        color = colors[i]
-        df = model.datacollector.get_model_vars_dataframe()
-        ax.plot(df.loc[:, m], label=m, color=color)
-    fig.legend()
-    # Set integer x axis
-    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
-    ax.set_xlabel("Step")
-    ax.set_ylabel("Number of Agents")
-    return solara.FigureMatplotlib(fig)
+    return solara.Markdown(
+        f"Resistant/Susceptible Ratio: {ratio_text}<br>Infected Remaining: {infected_text}"
+    )
 
 
 model_params = {
@@ -119,7 +86,18 @@ model_params = {
     ),
 }
 
-SpacePlot = make_space_matplotlib(agent_portrayal)
+
+def post_process_lineplot(ax):
+    ax.set_ylim(ymin=0)
+    ax.set_ylabel("# people")
+    ax.legend(bbox_to_anchor=(1.05, 1.0), loc="upper left")
+
+
+SpacePlot = make_space_component(agent_portrayal)
+StatePlot = make_plot_component(
+    {"Infected": "tab:red", "Susceptible": "tab:green", "Resistant": "tab:gray"},
+    post_process=post_process_lineplot,
+)
 
 model1 = VirusOnNetwork()
 
@@ -127,8 +105,8 @@ page = SolaraViz(
     model1,
     [
         SpacePlot,
-        make_plot,
-        # get_resistant_susceptible_ratio,  # TODO: Fix and uncomment
+        StatePlot,
+        get_resistant_susceptible_ratio,
     ],
     model_params=model_params,
     name="Virus Model",