Repository: projectmesa/mesa

Issue:
Adding many agents to a ContinuousSpace is slow
Dear *,

I am working on speeding up a [program that uses Mesa](https://github.com/ODISSEI-School-Choice/school-choice/).
One of the bottlenecks i found, was that adding lots of agents (60K+) to a ContinuousSpace is very slow.

https://github.com/projectmesa/mesa/blob/60bc768f5eede2dc4806684ae4b0b3046eb07d27/mesa/space.py#L796

Mesa is keeping track of agents on the grid using a numpy.array() ; however, when adding an agent numpy.array.append() has to reallocate a new array, and copy over the old one (see the note [here](https://numpy.org/doc/stable/reference/generated/numpy.append.html#numpy.append) ). This can lead to serious slow downs.

Using regular python arrays would already be orders of magnitude faster. See for instance [this stack overflow question](https://stackoverflow.com/questions/7133885/fastest-way-to-grow-a-numpy-numeric-array) where growing a numpy.array takes ~30 seconds, and growing a python list take 0.3 seconds.

If a lookup of agents is a performance bottleneck, an option would be to convert the python list to a numpy array when all the agents are added; or when the looking up an agent.

A quick fix to replace the numpy array with a list is [here](https://github.com/ODISSEI-School-Choice/school-choice/tree/jisk-v2/mesa-fixes) (based on mesa 0.9.0 from pip)



Discussion:
That sounds great! Are you open to making a PR for this (including some benchmark, for science)? If not, I will eventually do it this month.

Resolving patch:
diff --git a/mesa/space.py b/mesa/space.py
--- a/mesa/space.py
+++ b/mesa/space.py
@@ -747,8 +747,11 @@ class ContinuousSpace:
     """Continuous space where each agent can have an arbitrary position.
 
     Assumes that all agents are point objects, and have a pos property storing
-    their position as an (x, y) tuple. This class uses a numpy array internally
-    to store agent objects, to speed up neighborhood lookups.
+    their position as an (x, y) tuple.
+
+    This class uses a numpy array internally to store agent objects, to speed
+    up neighborhood lookups. This array is calculated on the first neighborhood
+    lookup, and is reused (and updated) until agents are added or removed.
     """
 
     _grid = None
@@ -782,7 +785,23 @@ def __init__(
 
         self._agent_points: npt.NDArray[FloatCoordinate] | None = None
         self._index_to_agent: Dict[int, Agent] = {}
-        self._agent_to_index: Dict[Agent, int] = {}
+        self._agent_to_index: Dict[Agent, int | None] = {}
+
+    def _build_agent_cache(self):
+        """Cache Agent positions to speed up neighbors calculations."""
+        self._index_to_agent = {}
+        agents = self._agent_to_index.keys()
+        for idx, agent in enumerate(agents):
+            self._agent_to_index[agent] = idx
+            self._index_to_agent[idx] = agent
+        self._agent_points = np.array(
+            [self._index_to_agent[idx].pos for idx in range(len(agents))]
+        )
+
+    def _invalidate_agent_cache(self):
+        """Clear cached data of Agents and positions in the space."""
+        self._agent_points = None
+        self._index_to_agent = {}
 
     def place_agent(self, agent: Agent, pos: FloatCoordinate) -> None:
         """Place a new agent in the space.
@@ -791,13 +810,9 @@ def place_agent(self, agent: Agent, pos: FloatCoordinate) -> None:
             agent: Agent object to place.
             pos: Coordinate tuple for where to place the agent.
         """
+        self._invalidate_agent_cache()
+        self._agent_to_index[agent] = None
         pos = self.torus_adj(pos)
-        if self._agent_points is None:
-            self._agent_points = np.array([pos])
-        else:
-            self._agent_points = np.append(self._agent_points, np.array([pos]), axis=0)
-        self._index_to_agent[self._agent_points.shape[0] - 1] = agent
-        self._agent_to_index[agent] = self._agent_points.shape[0] - 1
         agent.pos = pos
 
     def move_agent(self, agent: Agent, pos: FloatCoordinate) -> None:
@@ -808,11 +823,15 @@ def move_agent(self, agent: Agent, pos: FloatCoordinate) -> None:
             pos: Coordinate tuple to move the agent to.
         """
         pos = self.torus_adj(pos)
-        idx = self._agent_to_index[agent]
-        self._agent_points[idx, 0] = pos[0]
-        self._agent_points[idx, 1] = pos[1]
         agent.pos = pos
 
+        if self._agent_points is not None:
+            # instead of invalidating the full cache,
+            # apply the move to the cached values
+            idx = self._agent_to_index[agent]
+            self._agent_points[idx, 0] = pos[0]
+            self._agent_points[idx, 1] = pos[1]
+
     def remove_agent(self, agent: Agent) -> None:
         """Remove an agent from the simulation.
 
@@ -821,17 +840,9 @@ def remove_agent(self, agent: Agent) -> None:
         """
         if agent not in self._agent_to_index:
             raise Exception("Agent does not exist in the space")
-        idx = self._agent_to_index[agent]
-        del self._agent_to_index[agent]
-        max_idx = max(self._index_to_agent.keys())
-        # Delete the agent's position and decrement the index/agent mapping
-        self._agent_points = np.delete(self._agent_points, idx, axis=0)
-        for a, index in self._agent_to_index.items():
-            if index > idx:
-                self._agent_to_index[a] = index - 1
-                self._index_to_agent[index - 1] = a
-        # The largest index is now redundant
-        del self._index_to_agent[max_idx]
+        self._agent_to_index.pop(agent)
+
+        self._invalidate_agent_cache()
         agent.pos = None
 
     def get_neighbors(
@@ -847,6 +858,9 @@ def get_neighbors(
                             neighbors of a given agent, True will include that
                             agent in the results.
         """
+        if self._agent_points is None:
+            self._build_agent_cache()
+
         deltas = np.abs(self._agent_points - np.array(pos))
         if self.torus:
             deltas = np.minimum(deltas, self.size - deltas)


Test patch:
diff --git a/tests/test_space.py b/tests/test_space.py
--- a/tests/test_space.py
+++ b/tests/test_space.py
@@ -24,6 +24,30 @@
     (31, 41),
     (55, 32),
 ]
+TEST_AGENTS_PERF = 200000
+
+
+@pytest.mark.skip(reason="a perf test will slow down the CI")
+class TestSpacePerformance(unittest.TestCase):
+    """
+    Testing adding many agents for a continuous space.
+    """
+
+    def setUp(self):
+        """
+        Create a test space and populate with Mock Agents.
+        """
+        self.space = ContinuousSpace(10, 10, True, -10, -10)
+
+    def test_agents_add_many(self):
+        """
+        Add many agents
+        """
+        positions = np.random.rand(TEST_AGENTS_PERF, 2)
+        for i in range(TEST_AGENTS_PERF):
+            a = MockAgent(i, None)
+            pos = [positions[i, 0], positions[i, 1]]
+            self.space.place_agent(a, pos)
 
 
 class TestSpaceToroidal(unittest.TestCase):