[PERF] Improve move_to_empty performance

rht · rht · commit c24b61288aa6 · 2022-01-16T22:23:13.000-05:00
diff --git a/mesa/space.py b/mesa/space.py
@@ -39,6 +39,12 @@
 FloatCoordinate = Union[Tuple[float, float], np.ndarray]
 
 
+def clamp(x: float, lowest: float, highest: float) -> float:
+    # This should be faster than np.clip for a scalar x.
+    # TODO: measure how much faster this function is.
+    return max(lowest, min(x, highest))
+
+
 def accept_tuple_argument(wrapped_function):
     """Decorator to allow grid methods that take a list of (x, y) coord tuples
     to also handle a single position, by automatically wrapping tuple in
@@ -410,12 +416,46 @@ def is_cell_empty(self, pos: Coordinate) -> bool:
         x, y = pos
         return self.grid[x][y] == self.default_val()
 
-    def move_to_empty(self, agent: Agent) -> None:
+    def move_to_empty(
+        self, agent: Agent, cutoff: float = 0.998, num_agents: Optional[int] = None
+    ) -> None:
         """Moves agent to a random empty cell, vacating agent's old cell."""
         pos = agent.pos
         if len(self.empties) == 0:
             raise Exception("ERROR: No empty cells")
-        new_pos = agent.random.choice(sorted(self.empties))
+        if num_agents is None:
+            try:
+                num_agents = agent.model.schedule.get_agent_count()
+            except AttributeError:
+                raise Exception(
+                    "Your agent is not attached to a model, and so Mesa is unable\n"
+                    "to figure out the total number of agents you have created.\n"
+                    "This number is required in order to calculate the threshold\n"
+                    "for using a much faster algorithm to find an empty cell.\n"
+                    "In this case, you must specify `num_agents`."
+                )
+        new_pos = (0, 0)  # Initialize it with a starting value.
+        # This method is based on Agents.jl's random_empty() implementation.
+        # See https://github.com/JuliaDynamics/Agents.jl/pull/541.
+        # For the discussion, see
+        # https://github.com/projectmesa/mesa/issues/1052.
+        # This switch assumes the worst case (for this algorithm) of one
+        # agent per position, which is not true in general but is appropriate
+        # here.
+        if clamp(num_agents / (self.width * self.height), 0.0, 1.0) < cutoff:
+            # The default cutoff value provided is the break-even comparison
+            # with the time taken in the else branching point.
+            # The number is measured to be 0.998 in Agents.jl, but since Mesa
+            # run under different environment, the number is different here.
+            while True:
+                new_pos = (
+                    agent.random.randrange(self.width),
+                    agent.random.randrange(self.height),
+                )
+                if self.is_cell_empty(new_pos):
+                    break
+        else:
+            new_pos = agent.random.choice(sorted(self.empties))
         self._place_agent(new_pos, agent)
         agent.pos = new_pos
         self._remove_agent(pos, agent)
diff --git a/tests/test_grid.py b/tests/test_grid.py
@@ -229,6 +229,7 @@ def setUp(self):
                 a = MockAgent(counter, None)
                 self.agents.append(a)
                 self.grid.place_agent(a, (x, y))
+        self.num_agents = len(self.agents)
 
     def test_enforcement(self):
         """
@@ -242,25 +243,29 @@ def test_enforcement(self):
 
         # Place the agent in an empty cell
         self.grid.position_agent(a)
+        self.num_agents += 1
         # Test whether after placing, the empty cells are reduced by 1
         assert a.pos not in self.grid.empties
         assert len(self.grid.empties) == 8
         for i in range(10):
-            self.grid.move_to_empty(a)
+            # Since the agents and the grid are not associated with a model, we
+            # must explicitly tell move_to_empty the number of agents.
+            self.grid.move_to_empty(a, num_agents=self.num_agents)
         assert len(self.grid.empties) == 8
 
         # Place agents until the grid is full
         empty_cells = len(self.grid.empties)
         for i in range(empty_cells):
             a = MockAgent(101 + i, None)
             self.grid.position_agent(a)
+            self.num_agents += 1
         assert len(self.grid.empties) == 0
 
         a = MockAgent(110, None)
         with self.assertRaises(Exception):
             self.grid.position_agent(a)
         with self.assertRaises(Exception):
-            self.move_to_empty(self.agents[0])
+            self.move_to_empty(self.agents[0], num_agents=self.num_agents)
 
 
 # Number of agents at each position for testing
