More pythonic implementation of wolf sheep

benchmarks/WolfSheep/__init__.py

@@ -1,14 +0,0 @@
-from .wolf_sheep import WolfSheep
-
-if __name__ == "__main__":
-    # for profiling this benchmark model
-    import time
-
-    # model = WolfSheep(15, 25, 25, 60, 40, 0.2, 0.1, 20)
-    model = WolfSheep(15, 100, 100, 1000, 500, 0.4, 0.2, 20)
-
-    start_time = time.perf_counter()
-    for _ in range(100):
-        model.step()
-
-    print(time.perf_counter() - start_time)

benchmarks/WolfSheep/wolf_sheep.py

@@ -9,14 +9,117 @@
     Northwestern University, Evanston, IL.
 """
 
-import mesa
+from mesa import Agent, Model
 from mesa.space import MultiGrid
 from mesa.time import RandomActivationByType
 
-from .agents import GrassPatch, Sheep, Wolf
 
+class Animal(Agent):
+    def __init__(self, unique_id, model, moore, 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
 
-class WolfSheep(mesa.Model):
+    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)
+
+    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)
+        self.model.schedule.add(offspring)
+
+    def feed(self):
+        ...
+
+    def die(self):
+        self.model.grid.remove_agent(self)
+        self.remove()
+
+    def step(self):
+        self.random_move()
+        self.energy -= 1
+
+        self.feed()
+
+        if self.energy < 0:
+            self.die()
+        elif self.random.random() < self.p_reproduce:
+            self.spawn_offspring()
+
+
+class Sheep(Animal):
+    """
+    A sheep that walks around, reproduces (asexually) and gets eaten.
+
+    The init is the same as the RandomWalker.
+    """
+
+    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))
+        if grass_patch.fully_grown:
+            self.energy += self.energy_from_food
+            grass_patch.fully_grown = False
+
+
+class Wolf(Animal):
+    """
+    A wolf that walks around, reproduces (asexually) and eats sheep.
+    """
+
+    def feed(self):
+        agents = self.model.grid.get_cell_list_contents(self.pos)
+        sheep = [obj for obj in agents if isinstance(obj, Sheep)]
+        if len(sheep) > 0:
+            sheep_to_eat = self.random.choice(sheep)
+            self.energy += self.energy
+
+            # Kill the sheep
+            sheep_to_eat.die()
+
+
+class GrassPatch(Agent):
+    """
+    A patch of grass that grows at a fixed rate and it is eaten by sheep
+    """
+
+    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
+            countdown: Time for the patch of grass to be fully grown again
+        """
+        super().__init__(unique_id, model)
+        self.fully_grown = fully_grown
+        self.countdown = countdown
+
+    def step(self):
+        if not self.fully_grown:
+            if self.countdown <= 0:
+                # Set as fully grown
+                self.fully_grown = True
+                self.countdown = self.model.grass_regrowth_time
+            else:
+                self.countdown -= 1
+
+
+class WolfSheep(Model):
     """
     Wolf-Sheep Predation Model
 
@@ -35,6 +138,7 @@ def __init__(
         grass_regrowth_time,
         wolf_gain_from_food=13,
         sheep_gain_from_food=5,
+        moore=False,
     ):
         """
         Create a new Wolf-Sheep model with the given parameters.
@@ -49,41 +153,46 @@ def __init__(
             grass_regrowth_time: How long it takes for a grass patch to regrow
                                  once it is eaten
             sheep_gain_from_food: Energy sheep gain from grass, if enabled.
+            moore:
         """
         super().__init__(seed=seed)
         # Set parameters
         self.height = height
         self.width = width
         self.initial_sheep = initial_sheep
         self.initial_wolves = initial_wolves
-        self.sheep_reproduce = sheep_reproduce
-        self.wolf_reproduce = wolf_reproduce
-        self.wolf_gain_from_food = wolf_gain_from_food
         self.grass_regrowth_time = grass_regrowth_time
-        self.sheep_gain_from_food = sheep_gain_from_food
 
         self.schedule = RandomActivationByType(self)
         self.grid = MultiGrid(self.height, self.width, torus=False)
 
-        # Create sheep:
-        for _i in range(self.initial_sheep):
+        for _ in range(self.initial_sheep):
             pos = (
                 self.random.randrange(self.width),
                 self.random.randrange(self.height),
             )
-            energy = self.random.randrange(2 * self.sheep_gain_from_food)
-            sheep = Sheep(self.next_id(), pos, self, True, energy)
+            energy = self.random.randrange(2 * sheep_gain_from_food)
+            sheep = Sheep(
+                self.next_id(),
+                self,
+                moore,
+                energy,
+                sheep_reproduce,
+                sheep_gain_from_food,
+            )
             self.grid.place_agent(sheep, pos)
             self.schedule.add(sheep)
 
         # Create wolves
-        for _i in range(self.initial_wolves):
+        for _ in range(self.initial_wolves):
             pos = (
                 self.random.randrange(self.width),
                 self.random.randrange(self.height),
             )
-            energy = self.random.randrange(2 * self.wolf_gain_from_food)
-            wolf = Wolf(self.next_id(), pos, self, True, energy)
+            energy = self.random.randrange(2 * wolf_gain_from_food)
+            wolf = Wolf(
+                self.next_id(), self, moore, energy, wolf_reproduce, wolf_gain_from_food
+            )
             self.grid.place_agent(wolf, pos)
             self.schedule.add(wolf)
 
@@ -95,9 +204,20 @@ def __init__(
                 countdown = self.grass_regrowth_time
             else:
                 countdown = self.random.randrange(self.grass_regrowth_time)
-            patch = GrassPatch(self.next_id(), pos, self, fully_grown, countdown)
+            patch = GrassPatch(self.next_id(), self, fully_grown, countdown)
             self.grid.place_agent(patch, pos)
             self.schedule.add(patch)
 
     def step(self):
         self.schedule.step()
+
+
+if __name__ == "__main__":
+    import time
+
+    model = WolfSheep(15, 25, 25, 60, 40, 0.2, 0.1, 20)
+
+    start_time = time.perf_counter()
+    for _ in range(100):
+        model.step()
+    print("Time:", time.perf_counter() - start_time)