Scheduler: track agents with dict instead of list

examples/sugarscape_cg/sugarscape/schedule.py

@@ -1,5 +1,5 @@
 import random
-from collections import defaultdict
+from collections import OrderedDict
 
 from mesa.time import RandomActivation
 
@@ -14,11 +14,11 @@ class RandomActivationByBreed(RandomActivation):
 
     Assumes that all agents have a step() method.
     '''
-    agents_by_breed = defaultdict(list)
+    agents_by_breed = OrderedDict(dict)
 
     def __init__(self, model):
         super().__init__(model)
-        self.agents_by_breed = defaultdict(list)
+        self.agents_by_breed = OrderedDict(dict)
 
     def add(self, agent):
         '''

examples/wolf_sheep/wolf_sheep/agents.py

@@ -14,8 +14,8 @@ class Sheep(RandomWalker):
 
     energy = None
 
-    def __init__(self, pos, model, moore, energy=None):
-        super().__init__(pos, model, moore=moore)
+    def __init__(self, unique_id, pos, model, moore, energy=None):
+        super().__init__(unique_id, pos, model, moore=moore)
         self.energy = energy
 
     def step(self):
@@ -47,7 +47,8 @@ def step(self):
             # Create a new sheep:
             if self.model.grass:
                 self.energy /= 2
-            lamb = Sheep(self.pos, self.model, self.moore, self.energy)
+            lamb = Sheep(self.model.next_id(), self.pos, self.model,
+                         self.moore, self.energy)
             self.model.grid.place_agent(lamb, self.pos)
             self.model.schedule.add(lamb)
 
@@ -59,8 +60,8 @@ class Wolf(RandomWalker):
 
     energy = None
 
-    def __init__(self, pos, model, moore, energy=None):
-        super().__init__(pos, model, moore=moore)
+    def __init__(self, unique_id, pos, model, moore, energy=None):
+        super().__init__(unique_id, pos, model, moore=moore)
         self.energy = energy
 
     def step(self):
@@ -87,7 +88,8 @@ def step(self):
             if random.random() < self.model.wolf_reproduce:
                 # Create a new wolf cub
                 self.energy /= 2
-                cub = Wolf(self.pos, self.model, self.moore, self.energy)
+                cub = Wolf(self.model.next_id(), self.pos, self.model,
+                           self.moore, self.energy)
                 self.model.grid.place_agent(cub, cub.pos)
                 self.model.schedule.add(cub)
 
@@ -97,17 +99,18 @@ class GrassPatch(Agent):
     A patch of grass that grows at a fixed rate and it is eaten by sheep
     '''
 
-    def __init__(self, pos, model, fully_grown, countdown):
+    def __init__(self, unique_id, pos, 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__(pos, model)
+        super().__init__(unique_id, model)
         self.fully_grown = fully_grown
         self.countdown = countdown
+        self.pos = pos
 
     def step(self):
         if not self.fully_grown:

examples/wolf_sheep/wolf_sheep/model.py

@@ -62,7 +62,7 @@ def __init__(self, height=20, width=20,
                                  once it is eaten
             sheep_gain_from_food: Energy sheep gain from grass, if enabled.
         '''
-
+        super().__init__()
         # Set parameters
         self.height = height
         self.width = width
@@ -86,7 +86,7 @@ def __init__(self, height=20, width=20,
             x = random.randrange(self.width)
             y = random.randrange(self.height)
             energy = random.randrange(2 * self.sheep_gain_from_food)
-            sheep = Sheep((x, y), self, True, energy)
+            sheep = Sheep(self.next_id(), (x, y), self, True, energy)
             self.grid.place_agent(sheep, (x, y))
             self.schedule.add(sheep)
 
@@ -95,7 +95,7 @@ def __init__(self, height=20, width=20,
             x = random.randrange(self.width)
             y = random.randrange(self.height)
             energy = random.randrange(2 * self.wolf_gain_from_food)
-            wolf = Wolf((x, y), self, True, energy)
+            wolf = Wolf(self.next_id(), (x, y), self, True, energy)
             self.grid.place_agent(wolf, (x, y))
             self.schedule.add(wolf)
 
@@ -110,7 +110,8 @@ def __init__(self, height=20, width=20,
                 else:
                     countdown = random.randrange(self.grass_regrowth_time)
 
-                patch = GrassPatch((x, y), self, fully_grown, countdown)
+                patch = GrassPatch(self.next_id(), (x, y), self,
+                                   fully_grown, countdown)
                 self.grid.place_agent(patch, (x, y))
                 self.schedule.add(patch)
 

examples/wolf_sheep/wolf_sheep/random_walk.py

@@ -21,15 +21,15 @@ class RandomWalker(Agent):
     y = None
     moore = True
 
-    def __init__(self, pos, model, moore=True):
+    def __init__(self, unique_id, pos, model, moore=True):
         '''
         grid: The MultiGrid object in which the agent lives.
         x: The agent's current x coordinate
         y: The agent's current y coordinate
         moore: If True, may move in all 8 directions.
                 Otherwise, only up, down, left, right.
         '''
-        super().__init__(pos, model)
+        super().__init__(unique_id, model)
         self.pos = pos
         self.moore = moore
 

examples/wolf_sheep/wolf_sheep/schedule.py

@@ -1,5 +1,5 @@
 import random
-from collections import defaultdict
+from collections import OrderedDict
 
 from mesa.time import RandomActivation
 
@@ -14,11 +14,11 @@ class RandomActivationByBreed(RandomActivation):
 
     Assumes that all agents have a step() method.
     '''
-    agents_by_breed = defaultdict(list)
+    agents_by_breed = OrderedDict(dict)
 
     def __init__(self, model):
         super().__init__(model)
-        self.agents_by_breed = defaultdict(list)
+        self.agents_by_breed = OrderedDict(dict)
 
     def add(self, agent):
         '''
@@ -28,21 +28,19 @@ def add(self, agent):
             agent: An Agent to be added to the schedule.
         '''
 
-        self.agents.append(agent)
+        self.agents[agent.unique_id] = agent
         agent_class = type(agent)
-        self.agents_by_breed[agent_class].append(agent)
+        self.agents_by_breed[agent_class][agent.unique_id] = agent
 
     def remove(self, agent):
         '''
         Remove all instances of a given agent from the schedule.
         '''
 
-        while agent in self.agents:
-            self.agents.remove(agent)
+        del self.agents[agent.unique_id]
 
         agent_class = type(agent)
-        while agent in self.agents_by_breed[agent_class]:
-            self.agents_by_breed[agent_class].remove(agent)
+        del self.agents_by_breed[agent_class][agent.unique_id]
 
     def step(self, by_breed=True):
         '''
@@ -67,13 +65,13 @@ def step_breed(self, breed):
         Args:
             breed: Class object of the breed to run.
         '''
-        agents = self.agents_by_breed[breed]
-        random.shuffle(agents)
-        for agent in agents:
-            agent.step()
+        agent_keys = list(self.agents_by_breed[breed].keys())
+        random.shuffle(agent_keys)
+        for agent_key in agent_keys:
+            self.agents_by_breed[breed][agent_key].step()
 
     def get_breed_count(self, breed_class):
         '''
         Returns the current number of agents of certain breed in the queue.
         '''
-        return len(self.agents_by_breed[breed_class])
+        return len(self.agents_by_breed[breed_class].values())

mesa/batchrunner.py

@@ -154,7 +154,7 @@ def collect_model_vars(self, model):
     def collect_agent_vars(self, model):
         """ Run reporters and collect agent-level variables. """
         agent_vars = {}
-        for agent in model.schedule.agents:
+        for agent in model.schedule.agents.values():
             agent_record = {}
             for var, reporter in self.agent_reporters.items():
                 agent_record[var] = reporter(agent)

mesa/datacollection.py

@@ -136,7 +136,7 @@ def collect(self, model):
         if self.agent_reporters:
             for var, reporter in self.agent_reporters.items():
                 agent_records = []
-                for agent in model.schedule.agents:
+                for agent in model.schedule.agents.values():
                     agent_records.append((agent.unique_id, reporter(agent)))
                 self.agent_vars[var].append(agent_records)
 

mesa/model.py

@@ -34,6 +34,7 @@ def __init__(self, seed=None):
 
         self.running = True
         self.schedule = None
+        self.current_id = 0
 
     def run_model(self):
         """ Run the model until the end condition is reached. Overload as
@@ -46,3 +47,8 @@ def run_model(self):
     def step(self):
         """ A single step. Fill in here. """
         pass
+
+    def next_id(self):
+        """ Return the next unique ID for agents, increment current_id"""
+        self.current_id += 1
+        return self.current_id

mesa/time.py

@@ -43,7 +43,7 @@ def __init__(self, model):
         self.model = model
         self.steps = 0
         self.time = 0
-        self.agents = []
+        self.agents = {}
 
     def add(self, agent):
         """ Add an Agent object to the schedule.
@@ -53,7 +53,7 @@ def add(self, agent):
             have a step() method.
 
         """
-        self.agents.append(agent)
+        self.agents[agent.unique_id] = agent
 
     def remove(self, agent):
         """ Remove all instances of a given agent from the schedule.
@@ -62,19 +62,19 @@ def remove(self, agent):
             agent: An agent object.
 
         """
-        while agent in self.agents:
-            self.agents.remove(agent)
+        del self.agents[agent.unique_id]
 
     def step(self):
         """ Execute the step of all the agents, one at a time. """
-        for agent in self.agents[:]:
-            agent.step()
+        agent_keys = list(self.agents.keys())
+        for agent_key in agent_keys:
+            self.agents[agent_key].step()
         self.steps += 1
         self.time += 1
 
     def get_agent_count(self):
         """ Returns the current number of agents in the queue. """
-        return len(self.agents)
+        return len(self.agents.keys())
 
 
 class RandomActivation(BaseScheduler):
@@ -92,9 +92,11 @@ def step(self):
         random order.
 
         """
-        random.shuffle(self.agents)
-        for agent in self.agents[:]:
-            agent.step()
+        agent_keys = list(self.agents.keys())
+        random.shuffle(agent_keys)
+
+        for agent_key in agent_keys:
+            self.agents[agent_key].step()
         self.steps += 1
         self.time += 1
 
@@ -109,10 +111,11 @@ class SimultaneousActivation(BaseScheduler):
     """
     def step(self):
         """ Step all agents, then advance them. """
-        for agent in self.agents[:]:
-            agent.step()
-        for agent in self.agents[:]:
-            agent.advance()
+        agent_keys = list(self.agents.keys())
+        for agent_key in agent_keys:
+            self.agents[agent_key].step()
+        for agent_key in agent_keys:
+            self.agents[agent_key].advance()
         self.steps += 1
         self.time += 1
 
@@ -151,13 +154,14 @@ def __init__(self, model, stage_list=None, shuffle=False,
 
     def step(self):
         """ Executes all the stages for all agents. """
+        agent_keys = list(self.agents.keys())
         if self.shuffle:
-            random.shuffle(self.agents)
+            random.shuffle(agent_keys)
         for stage in self.stage_list:
-            for agent in self.agents[:]:
-                getattr(agent, stage)()  # Run stage
+            for agent_key in agent_keys:
+                getattr(self.agents[agent_key], stage)()  # Run stage
             if self.shuffle_between_stages:
-                random.shuffle(self.agents)
+                random.shuffle(agent_keys)
             self.time += self.stage_time
 
         self.steps += 1

tests/test_batchrunner.py

@@ -161,3 +161,6 @@ def test_model_with_variable_and_fixed_kwargs(self):
         self.assertEqual(model_vars.shape, (self.model_runs, expected_cols))
         self.assertEqual(model_vars['reported_fixed_param'].iloc[0],
                 self.fixed_params['fixed_name'])
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_datacollector.py

@@ -61,7 +61,7 @@ def setUp(self):
         for i in range(7):
             self.model.step()
         # Write to table:
-        for agent in self.model.schedule.agents:
+        for agent in self.model.schedule.agents.values():
             agent.write_final_values()
 
     def test_model_vars(self):
@@ -117,3 +117,6 @@ def test_exports(self):
 
         with self.assertRaises(Exception):
             table_df = data_collector.get_table_dataframe("not a real table")
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_grid.py

@@ -341,3 +341,6 @@ def test_neighbors(self):
 
         neighbors = self.grid.get_neighbors((1, 3), moore=False, radius=2)
         assert len(neighbors) == 11
+
+if __name__ == '__main__':
+    unittest.main()