Attempt to Optimize Wolf-Sheep Performance

mesa/examples/advanced/wolf_sheep/agents.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from mesa.discrete_space import CellAgent, FixedAgent
 
 
@@ -22,6 +24,8 @@ def __init__(
         self.energy_from_food = energy_from_food
         self.cell = cell
 
+        self.w, self.h = self.model.grid.width, self.model.grid.height
+
     def spawn_offspring(self):
         """Create offspring by splitting energy and creating new instance."""
         self.energy /= 2
@@ -67,29 +71,46 @@ def feed(self):
 
     def move(self):
         """Move towards a cell where there isn't a wolf, and preferably with grown grass."""
-        cells_without_wolves = self.cell.neighborhood.select(
-            lambda cell: not any(isinstance(obj, Wolf) for obj in cell.agents)
-        )
+
+        # Get current radius-1 neighborhood cell coords
+        x, y = self.cell.coordinate
+        neighbor_x = np.array([(x - 1) % self.w, x, (x + 1) % self.w])
+        neighbor_y = np.array([(y - 1) % self.h, y, (y + 1) % self.h])
+        neighbors = np.dstack(np.meshgrid(neighbor_x, neighbor_y, indexing="xy"))
+
+        # Get property layers
+        wolf_pos = self.model.grid.wolf_pos.data[neighbor_y[:, None], neighbor_x]
+        grass_pos = self.model.grid.grass_pos.data[neighbor_y[:, None], neighbor_x]
+
+        wolf_pos[[0, 0, 1, 2, 2], [0, 2, 1, 0, 2]] = 1
+        coords_without_wolves = wolf_pos == 0
         # If all surrounding cells have wolves, stay put
-        if len(cells_without_wolves) == 0:
+        safe_coords = neighbors[coords_without_wolves]
+        if safe_coords.size == 0:
             return
 
         # Among safe cells, prefer those with grown grass
-        cells_with_grass = cells_without_wolves.select(
-            lambda cell: any(
-                isinstance(obj, GrassPatch) and obj.fully_grown for obj in cell.agents
-            )
-        )
+        grass_pos[[0, 0, 1, 2, 2], [0, 2, 1, 0, 2]] = 0
+        coords_with_grass = neighbors[
+            np.logical_and(coords_without_wolves, grass_pos == 1)
+        ]
+
         # Move to a cell with grass if available, otherwise move to any safe cell
-        target_cells = (
-            cells_with_grass if len(cells_with_grass) > 0 else cells_without_wolves
-        )
-        self.cell = target_cells.select_random_cell()
+        target_coords = coords_with_grass if coords_with_grass.size > 0 else safe_coords
+        random_idx = self.random.randrange(len(target_coords))
+
+        self.cell = self.model.grid[tuple(target_coords[random_idx])]
 
 
 class Wolf(Animal):
     """A wolf that walks around, reproduces (asexually) and eats sheep."""
 
+    def __init__(
+        self, model, energy=8, p_reproduce=0.04, energy_from_food=4, cell=None
+    ):
+        super().__init__(model, energy, p_reproduce, energy_from_food, cell)
+        self.update_wolf_layer(1)
+
     def feed(self):
         """If possible, eat a sheep at current location."""
         sheep = [obj for obj in self.cell.agents if isinstance(obj, Sheep)]
@@ -106,7 +127,20 @@ def move(self):
         target_cells = (
             cells_with_sheep if len(cells_with_sheep) > 0 else self.cell.neighborhood
         )
+        self.update_wolf_layer(0)
         self.cell = target_cells.select_random_cell()
+        self.update_wolf_layer(1)
+
+    def remove(self):
+        """Ensure that when a wolf is removed, its mark is cleared from the grid."""
+        self.update_wolf_layer(0)
+        super().remove()
+
+    def update_wolf_layer(self, val):
+        """Update the wolf property layer"""
+        self.model.grid.wolf_pos.data[
+            self.cell.coordinate[1], self.cell.coordinate[0]
+        ] = val
 
 
 class GrassPatch(FixedAgent):
@@ -121,6 +155,7 @@ def fully_grown(self):
     def fully_grown(self, value: bool) -> None:
         """Set grass growth state and schedule regrowth if eaten."""
         self._fully_grown = value
+        self.model.changed_grass[self.cell.coordinate] = value
 
         if not value:  # If grass was just eaten
             self.model.simulator.schedule_event_relative(
@@ -142,6 +177,7 @@ def __init__(self, model, countdown, grass_regrowth_time, cell):
         self._fully_grown = countdown == 0
         self.grass_regrowth_time = grass_regrowth_time
         self.cell = cell
+        self.model.changed_grass[self.cell.coordinate] = self._fully_grown
 
         # Schedule initial growth if not fully grown
         if not self.fully_grown:

mesa/examples/advanced/wolf_sheep/model.py

@@ -11,6 +11,8 @@
 
 import math
 
+import numpy as np
+
 from mesa import Model
 from mesa.datacollection import DataCollector
 from mesa.discrete_space import OrthogonalVonNeumannGrid
@@ -68,6 +70,7 @@ def __init__(
         self.height = height
         self.width = width
         self.grass = grass
+        self.changed_grass = {}
 
         # Create grid using experimental cell space
         self.grid = OrthogonalVonNeumannGrid(
@@ -77,6 +80,10 @@ def __init__(
             random=self.random,
         )
 
+        # Set up property layers
+        self.grid.create_property_layer("grass_pos", default_value=0, dtype=int)
+        self.grid.create_property_layer("wolf_pos", default_value=0, dtype=int)
+
         # Set up data collection
         model_reporters = {
             "Wolves": lambda m: len(m.agents_by_type[Wolf]),
@@ -118,13 +125,26 @@ def __init__(
                 )
                 GrassPatch(self, countdown, grass_regrowth_time, cell)
 
+        self.update_grass_layer()
+
         # Collect initial data
         self.running = True
         self.datacollector.collect(self)
 
+    def update_grass_layer(self):
+        """Update grass_layer on cells where grass state has changed."""
+        if self.changed_grass:
+            coords = np.array(list(self.changed_grass.keys()))
+            states = np.array(list(self.changed_grass.values()), dtype=int)
+            self.grid.grass_pos.data[coords[:, 1], coords[:, 0]] = states
+            self.changed_grass.clear()
+
     def step(self):
         """Execute one step of the model."""
-        # First activate all sheep, then all wolves, both in random order
+        # Update the property layers before agents act.
+        self.update_grass_layer()
+
+        # Activate all sheep, then all wolves, both in random order
         self.agents_by_type[Sheep].shuffle_do("step")
         self.agents_by_type[Wolf].shuffle_do("step")