Repository: projectmesa/mesa
Pull request #2307 (2024-09-21 08:18:44)

Update to CellCollection.select
This PR updates CellCollect.select to reflect the updated API of AgentSet.select. So, we switch to `at_most` instead of using `n`. Since `n` was not used anyway before and this is still experimental, I have just removed `n` altogether.

github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]          |
|-----------------|--------|-------------------------|----------------------------|
| BoltzmannWealth | small  | 🟢 -4.3% [-5.1%, -3.5%] | 🔴 +8.5% [+8.3%, +8.7%]    |
| BoltzmannWealth | large  | 🔵 -2.0% [-2.5%, -1.4%] | 🔴 +16.0% [+11.7%, +21.6%] |
| Schelling       | small  | 🔵 -0.0% [-0.5%, +0.4%] | 🔵 +0.6% [+0.4%, +0.9%]    |
| Schelling       | large  | 🔵 -0.1% [-0.9%, +0.5%] | 🔵 +1.5% [-0.6%, +3.8%]    |
| WolfSheep       | small  | 🔵 +0.0% [-0.4%, +0.4%] | 🔵 +2.9% [-0.3%, +6.2%]    |
| WolfSheep       | large  | 🔵 +1.0% [-0.2%, +2.2%] | 🔴 +6.6% [+4.8%, +8.4%]    |
| BoidFlockers    | small  | 🔵 +0.2% [-0.2%, +0.6%] | 🔵 -0.3% [-1.0%, +0.4%]    |
| BoidFlockers    | large  | 🔵 -0.0% [-0.5%, +0.4%] | 🔵 -0.6% [-1.1%, -0.1%]    |

quaquel: these performance benchmarks are strange. Neither one uses `CellCollection.select`.
github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]          |
|-----------------|--------|-------------------------|----------------------------|
| BoltzmannWealth | small  | 🔵 -3.5% [-4.6%, -2.5%] | 🔵 +0.4% [+0.2%, +0.6%]    |
| BoltzmannWealth | large  | 🔵 -1.3% [-1.6%, -0.9%] | 🔵 -0.2% [-1.3%, +1.0%]    |
| Schelling       | small  | 🔵 -3.3% [-3.7%, -2.9%] | 🔵 -0.1% [-0.4%, +0.2%]    |
| Schelling       | large  | 🟢 -5.6% [-6.0%, -5.2%] | 🟢 -12.5% [-13.4%, -11.8%] |
| WolfSheep       | small  | 🟢 -3.4% [-3.7%, -3.2%] | 🔵 -1.0% [-1.2%, -0.8%]    |
| WolfSheep       | large  | 🟢 -6.3% [-7.5%, -5.3%] | 🟢 -11.5% [-12.4%, -10.7%] |
| BoidFlockers    | small  | 🟢 -4.2% [-4.6%, -3.7%] | 🔵 -1.4% [-1.9%, -0.8%]    |
| BoidFlockers    | large  | 🔵 -3.1% [-3.7%, -2.6%] | 🔵 -1.3% [-1.9%, -0.6%]    |

EwoutH: Nice to align them. Reminds me I wanted to add a `sample` method.

We merged some stuff, I think that's causing this PR branch to have more differences with main. Maybe rebase it.

Patch:
diff --git a/mesa/experimental/cell_space/cell.py b/mesa/experimental/cell_space/cell.py
index 381d0f6c..08a37102 100644
--- a/mesa/experimental/cell_space/cell.py
+++ b/mesa/experimental/cell_space/cell.py
@@ -48,7 +48,7 @@ class Cell:
     def __init__(
         self,
         coordinate: Coordinate,
-        capacity: float | None = None,
+        capacity: int | None = None,
         random: Random | None = None,
     ) -> None:
         """Initialise the cell.
@@ -65,7 +65,7 @@ class Cell:
         self.agents: list[
             Agent
         ] = []  # TODO:: change to AgentSet or weakrefs? (neither is very performant, )
-        self.capacity = capacity
+        self.capacity: int = capacity
         self.properties: dict[Coordinate, object] = {}
         self.random = random
 
diff --git a/mesa/experimental/cell_space/cell_collection.py b/mesa/experimental/cell_space/cell_collection.py
index 14832d51..a8c36f51 100644
--- a/mesa/experimental/cell_space/cell_collection.py
+++ b/mesa/experimental/cell_space/cell_collection.py
@@ -83,25 +83,36 @@ class CellCollection(Generic[T]):
         """
         return self.random.choice(list(self.agents))
 
-    def select(self, filter_func: Callable[[T], bool] | None = None, n=0):
+    def select(
+        self,
+        filter_func: Callable[[T], bool] | None = None,
+        at_most: int | float = float("inf"),
+    ):
         """Select cells based on filter function.
 
         Args:
             filter_func: filter function
-            n: number of cells to select
+            at_most: The maximum amount of cells to select. Defaults to infinity.
+              - If an integer, at most the first number of matching cells is selected.
+              - If a float between 0 and 1, at most that fraction of original number of cells
 
         Returns:
             CellCollection
 
         """
-        # FIXME: n is not considered
-        if filter_func is None and n == 0:
+        if filter_func is None and at_most == float("inf"):
             return self
 
-        return CellCollection(
-            {
-                cell: agents
-                for cell, agents in self._cells.items()
-                if filter_func is None or filter_func(cell)
-            }
-        )
+        if at_most <= 1.0 and isinstance(at_most, float):
+            at_most = int(len(self) * at_most)  # Note that it rounds down (floor)
+
+        def cell_generator(filter_func, at_most):
+            count = 0
+            for cell in self:
+                if count >= at_most:
+                    break
+                if not filter_func or filter_func(cell):
+                    yield cell
+                    count += 1
+
+        return CellCollection(cell_generator(filter_func, at_most))


Test patch:
diff --git a/tests/test_cell_space.py b/tests/test_cell_space.py
index 5ef4c44e..64d27566 100644
--- a/tests/test_cell_space.py
+++ b/tests/test_cell_space.py
@@ -500,3 +500,15 @@ def test_cell_collection():
 
     agents = collection[cells[0]]
     assert agents == cells[0].agents
+
+    cell = collection.select(at_most=1)
+    assert len(cell) == 1
+
+    cells = collection.select(at_most=2)
+    assert len(cells) == 2
+
+    cells = collection.select(at_most=0.5)
+    assert len(cells) == 5
+
+    cells = collection.select()
+    assert len(cells) == len(collection)