Mesa 3.0.0b0__py3-none-any.whl → 3.0.0b1__py3-none-any.whl

examples/basic/virus_on_network/agents.py

@@ -0,0 +1,69 @@
+from enum import Enum
+
+from mesa import Agent
+
+
+class State(Enum):
+    SUSCEPTIBLE = 0
+    INFECTED = 1
+    RESISTANT = 2
+
+
+class VirusAgent(Agent):
+    """Individual Agent definition and its properties/interaction methods."""
+
+    def __init__(
+        self,
+        model,
+        initial_state,
+        virus_spread_chance,
+        virus_check_frequency,
+        recovery_chance,
+        gain_resistance_chance,
+    ):
+        super().__init__(model)
+
+        self.state = initial_state
+
+        self.virus_spread_chance = virus_spread_chance
+        self.virus_check_frequency = virus_check_frequency
+        self.recovery_chance = recovery_chance
+        self.gain_resistance_chance = gain_resistance_chance
+
+    def try_to_infect_neighbors(self):
+        neighbors_nodes = self.model.grid.get_neighborhood(
+            self.pos, include_center=False
+        )
+        susceptible_neighbors = [
+            agent
+            for agent in self.model.grid.get_cell_list_contents(neighbors_nodes)
+            if agent.state is State.SUSCEPTIBLE
+        ]
+        for a in susceptible_neighbors:
+            if self.random.random() < self.virus_spread_chance:
+                a.state = State.INFECTED
+
+    def try_gain_resistance(self):
+        if self.random.random() < self.gain_resistance_chance:
+            self.state = State.RESISTANT
+
+    def try_remove_infection(self):
+        # Try to remove
+        if self.random.random() < self.recovery_chance:
+            # Success
+            self.state = State.SUSCEPTIBLE
+            self.try_gain_resistance()
+        else:
+            # Failed
+            self.state = State.INFECTED
+
+    def try_check_situation(self):
+        if (self.random.random() < self.virus_check_frequency) and (
+            self.state is State.INFECTED
+        ):
+            self.try_remove_infection()
+
+    def step(self):
+        if self.state is State.INFECTED:
+            self.try_to_infect_neighbors()
+        self.try_check_situation()

examples/basic/virus_on_network/app.py

@@ -0,0 +1,133 @@
+import math
+
+import solara
+from matplotlib.figure import Figure
+from matplotlib.ticker import MaxNLocator
+
+from mesa.visualization import Slider, SolaraViz, make_space_matplotlib
+
+from .model import State, VirusOnNetwork, number_infected
+
+
+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)
+    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,
+    }
+
+
+def get_resistant_susceptible_ratio(model):
+    ratio = model.resistant_susceptible_ratio()
+    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)
+
+
+model_params = {
+    "num_nodes": Slider(
+        label="Number of agents",
+        value=10,
+        min=10,
+        max=100,
+        step=1,
+    ),
+    "avg_node_degree": Slider(
+        label="Avg Node Degree",
+        value=3,
+        min=3,
+        max=8,
+        step=1,
+    ),
+    "initial_outbreak_size": Slider(
+        label="Initial Outbreak Size",
+        value=1,
+        min=1,
+        max=10,
+        step=1,
+    ),
+    "virus_spread_chance": Slider(
+        label="Virus Spread Chance",
+        value=0.4,
+        min=0.0,
+        max=1.0,
+        step=0.1,
+    ),
+    "virus_check_frequency": Slider(
+        label="Virus Check Frequency",
+        value=0.4,
+        min=0.0,
+        max=1.0,
+        step=0.1,
+    ),
+    "recovery_chance": Slider(
+        label="Recovery Chance",
+        value=0.3,
+        min=0.0,
+        max=1.0,
+        step=0.1,
+    ),
+    "gain_resistance_chance": Slider(
+        label="Gain Resistance Chance",
+        value=0.5,
+        min=0.0,
+        max=1.0,
+        step=0.1,
+    ),
+}
+
+SpacePlot = make_space_matplotlib(agent_portrayal)
+
+model1 = VirusOnNetwork()
+
+page = SolaraViz(
+    model1,
+    [
+        SpacePlot,
+        make_plot,
+        # get_resistant_susceptible_ratio,  # TODO: Fix and uncomment
+    ],
+    model_params=model_params,
+    name="Virus Model",
+)
+page  # noqa

examples/basic/virus_on_network/model.py

@@ -0,0 +1,99 @@
+import math
+
+import networkx as nx
+
+import mesa
+from mesa import Model
+
+from .agents import State, VirusAgent
+
+
+def number_state(model, state):
+    return sum(1 for a in model.grid.get_all_cell_contents() if a.state is state)
+
+
+def number_infected(model):
+    return number_state(model, State.INFECTED)
+
+
+def number_susceptible(model):
+    return number_state(model, State.SUSCEPTIBLE)
+
+
+def number_resistant(model):
+    return number_state(model, State.RESISTANT)
+
+
+class VirusOnNetwork(Model):
+    """A virus model with some number of agents."""
+
+    def __init__(
+        self,
+        num_nodes=10,
+        avg_node_degree=3,
+        initial_outbreak_size=1,
+        virus_spread_chance=0.4,
+        virus_check_frequency=0.4,
+        recovery_chance=0.3,
+        gain_resistance_chance=0.5,
+    ):
+        super().__init__()
+        self.num_nodes = num_nodes
+        prob = avg_node_degree / self.num_nodes
+        self.G = nx.erdos_renyi_graph(n=self.num_nodes, p=prob)
+        self.grid = mesa.space.NetworkGrid(self.G)
+
+        self.initial_outbreak_size = (
+            initial_outbreak_size if initial_outbreak_size <= num_nodes else num_nodes
+        )
+        self.virus_spread_chance = virus_spread_chance
+        self.virus_check_frequency = virus_check_frequency
+        self.recovery_chance = recovery_chance
+        self.gain_resistance_chance = gain_resistance_chance
+
+        self.datacollector = mesa.DataCollector(
+            {
+                "Infected": number_infected,
+                "Susceptible": number_susceptible,
+                "Resistant": number_resistant,
+            }
+        )
+
+        # Create agents
+        for node in self.G.nodes():
+            a = VirusAgent(
+                self,
+                State.SUSCEPTIBLE,
+                self.virus_spread_chance,
+                self.virus_check_frequency,
+                self.recovery_chance,
+                self.gain_resistance_chance,
+            )
+
+            # Add the agent to the node
+            self.grid.place_agent(a, node)
+
+        # Infect some nodes
+        infected_nodes = self.random.sample(list(self.G), self.initial_outbreak_size)
+        for a in self.grid.get_cell_list_contents(infected_nodes):
+            a.state = State.INFECTED
+
+        self.running = True
+        self.datacollector.collect(self)
+
+    def resistant_susceptible_ratio(self):
+        try:
+            return number_state(self, State.RESISTANT) / number_state(
+                self, State.SUSCEPTIBLE
+            )
+        except ZeroDivisionError:
+            return math.inf
+
+    def step(self):
+        self.agents.shuffle_do("step")
+        # collect data
+        self.datacollector.collect(self)
+
+    def run_model(self, n):
+        for _ in range(n):
+            self.step()

mesa/__init__.py

@@ -5,6 +5,8 @@ Core Objects: Model, and Agent.
 
 import datetime
 
+import mesa.examples as examples
+import mesa.experimental as experimental
 import mesa.space as space
 import mesa.time as time
 from mesa.agent import Agent
@@ -20,10 +22,11 @@ __all__ = [
     "DataCollector",
     "batch_run",
     "experimental",
+    "examples",
 ]
 
 __title__ = "mesa"
-__version__ = "3.0.0b0"
+__version__ = "3.0.0b1"
 __license__ = "Apache 2.0"
 _this_year = datetime.datetime.now(tz=datetime.timezone.utc).date().year
 __copyright__ = f"Copyright {_this_year} Project Mesa Team"

mesa/agent.py

@@ -62,9 +62,9 @@ class Agent:
         super().__init__(*args, **kwargs)
 
         self.model: Model = model
-        self.model.register_agent(self)
         self.unique_id: int = next(self._ids[model])
         self.pos: Position | None = None
+        self.model.register_agent(self)
 
     def remove(self) -> None:
         """Remove and delete the agent from the model."""
@@ -99,14 +99,18 @@ class AgentSet(MutableSet, Sequence):
         which means that agents not referenced elsewhere in the program may be automatically removed from the AgentSet.
     """
 
-    def __init__(self, agents: Iterable[Agent], model: Model):
+    def __init__(self, agents: Iterable[Agent], random: Random | None = None):
         """Initializes the AgentSet with a collection of agents and a reference to the model.
 
         Args:
             agents (Iterable[Agent]): An iterable of Agent objects to be included in the set.
-            model (Model): The ABM model instance to which this AgentSet belongs.
+            random (Random): the random number generator
         """
-        self.model = model
+        if random is None:
+            random = (
+                Random()
+            )  # FIXME see issue 1981, how to get the central rng from model
+        self.random = random
         self._agents = weakref.WeakKeyDictionary({agent: None for agent in agents})
 
     def __len__(self) -> int:
@@ -177,7 +181,7 @@ class AgentSet(MutableSet, Sequence):
 
         agents = agent_generator(filter_func, agent_type, at_most)
 
-        return AgentSet(agents, self.model) if not inplace else self._update(agents)
+        return AgentSet(agents, self.random) if not inplace else self._update(agents)
 
     def shuffle(self, inplace: bool = False) -> AgentSet:
         """Randomly shuffle the order of agents in the AgentSet.
@@ -200,7 +204,7 @@ class AgentSet(MutableSet, Sequence):
             return self
         else:
             return AgentSet(
-                (agent for ref in weakrefs if (agent := ref()) is not None), self.model
+                (agent for ref in weakrefs if (agent := ref()) is not None), self.random
             )
 
     def sort(
@@ -225,7 +229,7 @@ class AgentSet(MutableSet, Sequence):
         sorted_agents = sorted(self._agents.keys(), key=key, reverse=not ascending)
 
         return (
-            AgentSet(sorted_agents, self.model)
+            AgentSet(sorted_agents, self.random)
             if not inplace
             else self._update(sorted_agents)
         )
@@ -477,7 +481,7 @@ class AgentSet(MutableSet, Sequence):
         Returns:
             dict: A dictionary representing the state of the AgentSet.
         """
-        return {"agents": list(self._agents.keys()), "model": self.model}
+        return {"agents": list(self._agents.keys()), "random": self.random}
 
     def __setstate__(self, state):
         """Set the state of the AgentSet during deserialization.
@@ -485,18 +489,9 @@ class AgentSet(MutableSet, Sequence):
         Args:
             state (dict): A dictionary representing the state to restore.
         """
-        self.model = state["model"]
+        self.random = state["random"]
         self._update(state["agents"])
 
-    @property
-    def random(self) -> Random:
-        """Provide access to the model's random number generator.
-
-        Returns:
-            Random: The random number generator associated with the model.
-        """
-        return self.model.random
-
     def groupby(self, by: Callable | str, result_type: str = "agentset") -> GroupBy:
         """Group agents by the specified attribute or return from the callable.
 
@@ -529,7 +524,7 @@ class AgentSet(MutableSet, Sequence):
 
         if result_type == "agentset":
             return GroupBy(
-                {k: AgentSet(v, model=self.model) for k, v in groups.items()}
+                {k: AgentSet(v, random=self.random) for k, v in groups.items()}
             )
         else:
             return GroupBy(groups)

mesa/examples.py

@@ -0,0 +1,3 @@
+"""This module is a collection of example models built using the Mesa framework."""
+
+__path__ = ["examples"]

mesa/experimental/__init__.py

@@ -2,6 +2,12 @@
 
 from mesa.experimental import cell_space
 
-from .solara_viz import JupyterViz, Slider, SolaraViz, make_text
+try:
+    from .solara_viz import JupyterViz, Slider, SolaraViz, make_text
 
-__all__ = ["cell_space", "JupyterViz", "SolaraViz", "make_text", "Slider"]
+    __all__ = ["cell_space", "JupyterViz", "Slider", "SolaraViz", "make_text"]
+except ImportError:
+    print(
+        "Could not import SolaraViz. If you need it, install with 'pip install --pre mesa[viz]'"
+    )
+    __all__ = ["cell_space"]

mesa/experimental/cell_space/cell.py

@@ -211,3 +211,12 @@ class Cell:
         self._mesa_property_layers[property_name].modify_cell(
             self.coordinate, operation, value
         )
+
+    def __getstate__(self):
+        """Return state of the Cell with connections set to empty."""
+        # fixme, once we shift to 3.11, replace this with super. __getstate__
+        state = (self.__dict__, {k: getattr(self, k) for k in self.__slots__})
+        state[1][
+            "connections"
+        ] = {}  # replace this with empty connections to avoid infinite recursion error in pickle/deepcopy
+        return state

mesa/experimental/cell_space/discrete_space.py

@@ -22,7 +22,7 @@ class DiscreteSpace(Generic[T]):
         all_cells (CellCollection): The cells composing the discrete space
         random (Random): The random number generator
         cell_klass (Type) : the type of cell class
-        empties (CellCollection) : collecction of all cells that are empty
+        empties (CellCollection) : collection of all cells that are empty
         property_layers (dict[str, PropertyLayer]): the property layers of the discrete space
     """
 
@@ -55,6 +55,7 @@ class DiscreteSpace(Generic[T]):
     def cutoff_empties(self):  # noqa
         return 7.953 * len(self._cells) ** 0.384
 
+    def _connect_cells(self): ...
     def _connect_single_cell(self, cell: T): ...
 
     @cached_property
@@ -134,3 +135,8 @@ class DiscreteSpace(Generic[T]):
             condition: a function that takes a cell and returns a boolean (used to filter cells)
         """
         self.property_layers[property_name].modify_cells(operation, value, condition)
+
+    def __setstate__(self, state):
+        """Set the state of the discrete space and rebuild the connections."""
+        self.__dict__ = state
+        self._connect_cells()

mesa/experimental/cell_space/grid.py

@@ -22,8 +22,21 @@ class Grid(DiscreteSpace[T], Generic[T]):
         random (Random): the random number generator
         _try_random (bool): whether to get empty cell be repeatedly trying random cell
 
+    Notes:
+        width and height are accessible via properties, higher dimensions can be retrieved via dimensions
+
     """
 
+    @property
+    def width(self) -> int:
+        """Convenience access to the width of the grid."""
+        return self.dimensions[0]
+
+    @property
+    def height(self) -> int:
+        """Convenience access to the height of the grid."""
+        return self.dimensions[1]
+
     def __init__(
         self,
         dimensions: Sequence[int],

mesa/experimental/cell_space/network.py

@@ -34,6 +34,9 @@ class Network(DiscreteSpace[Cell]):
                 node_id, capacity, random=self.random
             )
 
+        self._connect_cells()
+
+    def _connect_cells(self) -> None:
         for cell in self.all_cells:
             self._connect_single_cell(cell)
 

mesa/model.py

@@ -8,14 +8,21 @@ Core Objects: Model
 from __future__ import annotations
 
 import random
+import sys
 import warnings
+from collections.abc import Sequence
 
 # mypy
 from typing import Any
 
+import numpy as np
+
 from mesa.agent import Agent, AgentSet
 from mesa.datacollection import DataCollector
 
+SeedLike = int | np.integer | Sequence[int] | np.random.SeedSequence
+RNGLike = np.random.Generator | np.random.BitGenerator
+
 
 class Model:
     """Base class for models in the Mesa ABM library.
@@ -28,7 +35,8 @@ class Model:
         running: A boolean indicating if the model should continue running.
         schedule: An object to manage the order and execution of agent steps.
         steps: the number of times `model.step()` has been called.
-        random: a seeded random number generator.
+        random: a seeded python.random number generator.
+        rng : a seeded numpy.random.Generator
 
     Notes:
         Model.agents returns the AgentSet containing all agents registered with the model. Changing
@@ -37,7 +45,13 @@ class Model:
 
     """
 
-    def __init__(self, *args: Any, seed: float | None = None, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        *args: Any,
+        seed: float | None = None,
+        rng: RNGLike | SeedLike | None = None,
+        **kwargs: Any,
+    ) -> None:
         """Create a new model.
 
         Overload this method with the actual code to initialize the model. Always start with super().__init__()
@@ -46,25 +60,51 @@ class Model:
         Args:
             args: arguments to pass onto super
             seed: the seed for the random number generator
+            rng : Pseudorandom number generator state. When `rng` is None, a new `numpy.random.Generator` is created
+                  using entropy from the operating system. Types other than `numpy.random.Generator` are passed to
+                  `numpy.random.default_rng` to instantiate a `Generator`.
             kwargs: keyword arguments to pass onto super
+
+        Notes:
+            you have to pass either seed or rng, but not both.
+
         """
         super().__init__(*args, **kwargs)
         self.running = True
         self.steps: int = 0
 
-        self._setup_agent_registration()
-
-        self._seed = seed
-        if self._seed is None:
-            # We explicitly specify the seed here so that we know its value in
-            # advance.
-            self._seed = random.random()
-        self.random = random.Random(self._seed)
+        if (seed is not None) and (rng is not None):
+            raise ValueError("you have to pass either rng or seed, not both")
+        elif seed is None:
+            self.rng: np.random.Generator = np.random.default_rng(rng)
+            self._rng = (
+                self.rng.bit_generator.state
+            )  # this allows for reproducing the rng
+
+            try:
+                self.random = random.Random(rng)
+            except TypeError:
+                seed = int(self.rng.integers(np.iinfo(np.int32).max))
+                self.random = random.Random(seed)
+            self._seed = seed  # this allows for reproducing stdlib.random
+        elif rng is None:
+            self.random = random.Random(seed)
+            self._seed = seed  # this allows for reproducing stdlib.random
+
+            try:
+                self.rng: np.random.Generator = np.random.default_rng(rng)
+            except TypeError:
+                rng = self.random.randint(0, sys.maxsize)
+                self.rng: np.random.Generator = np.random.default_rng(rng)
+            self._rng = self.rng.bit_generator.state
 
         # Wrap the user-defined step method
         self._user_step = self.step
         self.step = self._wrapped_step
 
+        # setup agent registration data structures
+        self._setup_agent_registration()
+
     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
@@ -119,7 +159,9 @@ class 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
+        self._all_agents = AgentSet(
+            [], random=self.random
+        )  # an agenset with all agents
 
     def register_agent(self, agent):
         """Register the agent with the model.
@@ -153,7 +195,7 @@ class Model:
                 [
                     agent,
                 ],
-                self,
+                random=self.random,
             )
 
         self._all_agents.add(agent)
@@ -194,6 +236,15 @@ class Model:
         self.random.seed(seed)
         self._seed = seed
 
+    def reset_rng(self, rng: RNGLike | SeedLike | None = None) -> None:
+        """Reset the model random number generator.
+
+        Args:
+            rng: A new seed for the RNG; if None, reset using the current seed
+        """
+        self.rng = np.random.default_rng(rng)
+        self._rng = self.rng.bit_generator.state
+
     def initialize_data_collector(
         self,
         model_reporters=None,

mesa/time.py

@@ -77,7 +77,7 @@ class BaseScheduler:
         if agents is None:
             agents = []
 
-        self._agents: AgentSet = AgentSet(agents, model)
+        self._agents: AgentSet = AgentSet(agents, model.random)
 
         self._remove_warning_given = False
         self._agents_key_warning_given = False
@@ -312,7 +312,9 @@ class RandomActivationByType(BaseScheduler):
                 try:
                     self._agents_by_type[type(agent)].add(agent)
                 except KeyError:
-                    self._agents_by_type[type(agent)] = AgentSet([agent], self.model)
+                    self._agents_by_type[type(agent)] = AgentSet(
+                        [agent], self.model.random
+                    )
 
     def add(self, agent: Agent) -> None:
         """Add an Agent object to the schedule.
@@ -325,7 +327,7 @@ class RandomActivationByType(BaseScheduler):
         try:
             self._agents_by_type[type(agent)].add(agent)
         except KeyError:
-            self._agents_by_type[type(agent)] = AgentSet([agent], self.model)
+            self._agents_by_type[type(agent)] = AgentSet([agent], self.model.random)
 
     def remove(self, agent: Agent) -> None:
         """Remove all instances of a given agent from the schedule.