Add cell-centric discrete spaces (experimental)

benchmarks/Schelling/schelling.py

@@ -1,12 +1,14 @@
-import mesa
+from mesa import Model
+from mesa.experimental.cell_space import CellAgent, OrthogonalGrid
+from mesa.time import RandomActivation
 
 
-class SchellingAgent(mesa.Agent):
+class SchellingAgent(CellAgent):
     """
     Schelling segregation agent
     """
 
-    def __init__(self, unique_id, model, agent_type):
+    def __init__(self, unique_id, model, agent_type, radius):
         """
         Create a new Schelling agent.
         Args:
@@ -16,23 +18,22 @@ def __init__(self, unique_id, model, agent_type):
         """
         super().__init__(unique_id, model)
         self.type = agent_type
+        self.radius = radius
 
     def step(self):
         similar = 0
-        for neighbor in self.model.grid.iter_neighbors(
-            self.pos, moore=True, radius=self.model.radius
-        ):
+        for neighbor in self.cell.neighborhood(radius=self.radius).agents:
             if neighbor.type == self.type:
                 similar += 1
 
         # If unhappy, move:
         if similar < self.model.homophily:
-            self.model.grid.move_to_empty(self)
+            self.move_to(self.model.grid.select_random_empty_cell())
         else:
             self.model.happy += 1
 
 
-class Schelling(mesa.Model):
+class Schelling(Model):
     """
     Model class for the Schelling segregation model.
     """
@@ -47,22 +48,23 @@ def __init__(
         self.density = density
         self.minority_pc = minority_pc
         self.homophily = homophily
-        self.radius = radius
 
-        self.schedule = mesa.time.RandomActivation(self)
-        self.grid = mesa.space.SingleGrid(height, width, torus=True)
+        self.schedule = RandomActivation(self)
+        self.grid = OrthogonalGrid(
+            height, width, torus=True, capacity=1, random=self.random
+        )
 
         self.happy = 0
 
         # Set up agents
         # We use a grid iterator that returns
         # the coordinates of a cell as well as
         # its contents. (coord_iter)
-        for _cont, pos in self.grid.coord_iter():
+        for cell in self.grid:
             if self.random.random() < self.density:
                 agent_type = 1 if self.random.random() < self.minority_pc else 0
-                agent = SchellingAgent(self.next_id(), self, agent_type)
-                self.grid.place_agent(agent, pos)
+                agent = SchellingAgent(self.next_id(), self, agent_type, radius)
+                agent.move_to(cell)
                 self.schedule.add(agent)
 
     def step(self):

benchmarks/WolfSheep/wolf_sheep.py

@@ -8,42 +8,40 @@
     Center for Connected Learning and Computer-Based Modeling,
     Northwestern University, Evanston, IL.
 """
+import math
 
-import mesa
+from mesa import Model
+from mesa.experimental.cell_space import CellAgent, OrthogonalGrid
+from mesa.time import RandomActivationByType
 
 
-class Animal(mesa.Agent):
-    def __init__(self, unique_id, model, moore, energy, p_reproduce, energy_from_food):
+class Animal(CellAgent):
+    def __init__(self, unique_id, model, energy, p_reproduce, energy_from_food):
         super().__init__(unique_id, model)
         self.energy = energy
         self.p_reproduce = p_reproduce
         self.energy_from_food = energy_from_food
-        self.moore = moore
 
     def random_move(self):
-        next_moves = self.model.grid.get_neighborhood(self.pos, self.moore, True)
-        next_move = self.random.choice(next_moves)
-        # Now move:
-        self.model.grid.move_agent(self, next_move)
+        self.move_to(self.cell.neighborhood().select_random_cell())
 
     def spawn_offspring(self):
         self.energy /= 2
         offspring = self.__class__(
             self.model.next_id(),
             self.model,
-            self.moore,
             self.energy,
             self.p_reproduce,
             self.energy_from_food,
         )
-        self.model.grid.place_agent(offspring, self.pos)
+        offspring.move_to(self.cell)
         self.model.schedule.add(offspring)
 
     def feed(self):
         ...
 
     def die(self):
-        self.model.grid.remove_agent(self)
+        self.cell.remove_agent(self)
         self.remove()
 
     def step(self):
@@ -67,8 +65,9 @@ class Sheep(Animal):
 
     def feed(self):
         # If there is grass available, eat it
-        agents = self.model.grid.get_cell_list_contents(self.pos)
-        grass_patch = next(obj for obj in agents if isinstance(obj, GrassPatch))
+        grass_patch = next(
+            obj for obj in self.cell.agents if isinstance(obj, GrassPatch)
+        )
         if grass_patch.fully_grown:
             self.energy += self.energy_from_food
             grass_patch.fully_grown = False
@@ -80,8 +79,7 @@ class Wolf(Animal):
     """
 
     def feed(self):
-        agents = self.model.grid.get_cell_list_contents(self.pos)
-        sheep = [obj for obj in agents if isinstance(obj, Sheep)]
+        sheep = [obj for obj in self.cell.agents if isinstance(obj, Sheep)]
         if len(sheep) > 0:
             sheep_to_eat = self.random.choice(sheep)
             self.energy += self.energy
@@ -90,7 +88,7 @@ def feed(self):
             sheep_to_eat.die()
 
 
-class GrassPatch(mesa.Agent):
+class GrassPatch(CellAgent):
     """
     A patch of grass that grows at a fixed rate and it is eaten by sheep
     """
@@ -100,7 +98,7 @@ def __init__(self, unique_id, model, fully_grown, countdown):
         Creates a new patch of grass
 
         Args:
-            grown: (boolean) Whether the patch of grass is fully grown or not
+            fully_grown: (boolean) Whether the patch of grass is fully grown or not
             countdown: Time for the patch of grass to be fully grown again
         """
         super().__init__(unique_id, model)
@@ -117,7 +115,7 @@ def step(self):
                 self.countdown -= 1
 
 
-class WolfSheep(mesa.Model):
+class WolfSheep(Model):
     """
     Wolf-Sheep Predation Model
 
@@ -161,24 +159,27 @@ def __init__(
         self.initial_wolves = initial_wolves
         self.grass_regrowth_time = grass_regrowth_time
 
-        self.schedule = mesa.time.RandomActivationByType(self)
-        self.grid = mesa.space.MultiGrid(self.height, self.width, torus=False)
+        self.schedule = RandomActivationByType(self)
+        self.grid = OrthogonalGrid(
+            self.height,
+            self.width,
+            moore=moore,
+            torus=False,
+            capacity=math.inf,
+            random=self.random,
+        )
 
+        # Create sheep:
         for _ in range(self.initial_sheep):
             pos = (
                 self.random.randrange(self.width),
                 self.random.randrange(self.height),
             )
             energy = self.random.randrange(2 * sheep_gain_from_food)
             sheep = Sheep(
-                self.next_id(),
-                self,
-                moore,
-                energy,
-                sheep_reproduce,
-                sheep_gain_from_food,
+                self.next_id(), self, energy, sheep_reproduce, sheep_gain_from_food
             )
-            self.grid.place_agent(sheep, pos)
+            sheep.move_to(self.grid.cells[pos])
             self.schedule.add(sheep)
 
         # Create wolves
@@ -189,21 +190,21 @@ def __init__(
             )
             energy = self.random.randrange(2 * wolf_gain_from_food)
             wolf = Wolf(
-                self.next_id(), self, moore, energy, wolf_reproduce, wolf_gain_from_food
+                self.next_id(), self, energy, wolf_reproduce, wolf_gain_from_food
             )
-            self.grid.place_agent(wolf, pos)
+            wolf.move_to(self.grid.cells[pos])
             self.schedule.add(wolf)
 
         # Create grass patches
         possibly_fully_grown = [True, False]
-        for _agent, pos in self.grid.coord_iter():
+        for cell in self.grid:
             fully_grown = self.random.choice(possibly_fully_grown)
             if fully_grown:
                 countdown = self.grass_regrowth_time
             else:
                 countdown = self.random.randrange(self.grass_regrowth_time)
             patch = GrassPatch(self.next_id(), self, fully_grown, countdown)
-            self.grid.place_agent(patch, pos)
+            patch.move_to(cell)
             self.schedule.add(patch)
 
     def step(self):

mesa/experimental/__init__.py

@@ -1 +1,5 @@
 from .jupyter_viz import JupyterViz, make_text, Slider  # noqa
+from mesa.experimental import cell_space
+
+
+__all__ = ["JupyterViz", "make_text", "Slider", "cell_space"]

mesa/experimental/cell_space/__init__.py

@@ -0,0 +1,17 @@
+from mesa.experimental.cell_space.cell import Cell
+from mesa.experimental.cell_space.cell_agent import CellAgent
+from mesa.experimental.cell_space.cell_collection import CellCollection
+from mesa.experimental.cell_space.discrete_space import DiscreteSpace
+from mesa.experimental.cell_space.grid import Grid, HexGrid, OrthogonalGrid
+from mesa.experimental.cell_space.network import Network
+
+__all__ = [
+    "CellCollection",
+    "Cell",
+    "CellAgent",
+    "DiscreteSpace",
+    "Grid",
+    "HexGrid",
+    "OrthogonalGrid",
+    "Network",
+]

mesa/experimental/cell_space/cell.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+from functools import cache
+from random import Random
+from typing import TYPE_CHECKING, Any
+
+from mesa.experimental.cell_space.cell_collection import CellCollection
+
+if TYPE_CHECKING:
+    from mesa.experimental.cell_space.cell_agent import CellAgent
+
+
+class Cell:
+    __slots__ = [
+        "coordinate",
+        "_connections",
+        "agents",
+        "capacity",
+        "properties",
+        "random",
+    ]
+
+    def __init__(
+        self, coordinate: Any, capacity: int | None = None, random: Random | None = None
+    ) -> None:
+        """ "
+
+        Args:
+            coordinate:
+            capacity (int) : the capacity of the cell. If None, the capacity is infinite
+            random (Random) : the random number generator to use
+
+        """
+        super().__init__()
+        self.coordinate = coordinate
+        self._connections: list[Cell] = []  # TODO: change to CellCollection?
+        self.agents = []  # TODO:: change to AgentSet or weakrefs? (neither is very performant, )
+        self.capacity = capacity
+        self.properties: dict[str, object] = {}
+        self.random = random
+
+    def connect(self, other) -> None:
+        """Connects this cell to another cell.
+
+        Args:
+            other (Cell): other cell to connect to
+
+        """
+        self._connections.append(other)
+
+    def disconnect(self, other) -> None:
+        """Disconnects this cell from another cell.
+
+        Args:
+            other (Cell): other cell to remove from connections
+
+        """
+        self._connections.remove(other)
+
+    def add_agent(self, agent: CellAgent) -> None:
+        """Adds an agent to the cell.
+
+        Args:
+            agent (CellAgent): agent to add to this Cell
+
+        """
+        n = len(self.agents)
+
+        if self.capacity and n >= self.capacity:
+            raise Exception(
+                "ERROR: Cell is full"
+            )  # FIXME we need MESA errors or a proper error
+
+        self.agents.append(agent)
+
+    def remove_agent(self, agent: CellAgent) -> None:
+        """Removes an agent from the cell.
+
+        Args:
+            agent (CellAgent): agent to remove from this cell
+
+        """
+        self.agents.remove(agent)
+        agent.cell = None
+
+    @property
+    def is_empty(self) -> bool:
+        """Returns a bool of the contents of a cell."""
+        return len(self.agents) == 0
+
+    @property
+    def is_full(self) -> bool:
+        """Returns a bool of the contents of a cell."""
+        return len(self.agents) == self.capacity
+
+    def __repr__(self):
+        return f"Cell({self.coordinate}, {self.agents})"
+
+    @cache
+    def neighborhood(self, radius=1, include_center=False):
+        return CellCollection(
+            self._neighborhood(radius=radius, include_center=include_center),
+            random=self.random,
+        )
+
+    @cache
+    def _neighborhood(self, radius=1, include_center=False):
+        # if radius == 0:
+        #     return {self: self.agents}
+        if radius < 1:
+            raise ValueError("radius must be larger than one")
+        if radius == 1:
+            neighborhood = {neighbor: neighbor.agents for neighbor in self._connections}
+            if not include_center:
+                return neighborhood
+            else:
+                neighborhood[self] = self.agents
+                return neighborhood
+        else:
+            neighborhood = {}
+            for neighbor in self._connections:
+                neighborhood.update(
+                    neighbor._neighborhood(radius - 1, include_center=True)
+                )
+            if not include_center:
+                neighborhood.pop(self, None)
+            return neighborhood