Fix registering an AgentSet that has no agents yet

examples/boltzmann_wealth/performance_plot.py

@@ -214,7 +214,7 @@ def main():
         mesa_frames_polars_concise,
         mesa_frames_polars_native,
     ]
-    n_range_0 = [k for k in range(0, 100001, 10000)]
+    n_range_0 = [*range(0, 100001, 10000)]
     title_0 = "100 steps of the Boltzmann Wealth model:\n" + " vs ".join(labels_0)
     image_path_0 = "boltzmann_with_mesa.png"
 
@@ -228,7 +228,7 @@ def main():
         mesa_frames_polars_concise,
         mesa_frames_polars_native,
     ]
-    n_range_1 = [k for k in range(100000, 1000001, 100000)]
+    n_range_1 = [*range(100000, 1000001, 100000)]
     title_1 = "100 steps of the Boltzmann Wealth model:\n" + " vs ".join(labels_1)
     image_path_1 = "boltzmann_no_mesa.png"
 

examples/sugarscape_ig/performance_comparison.py

@@ -186,7 +186,7 @@ def main():
         mesa_frames_polars_numba_parallel,
         mesa_implementation,
     ]
-    n_range_0 = [k for k in range(10**5, 5 * 10**5 + 2, 10**5)]
+    n_range_0 = [*range(10**5, 5 * 10**5 + 2, 10**5)]
     title_0 = "100 steps of the SugarScape IG model:\n" + " vs ".join(labels_0)
     image_path_0 = "mesa_comparison.png"
     plot_and_print_benchmark(labels_0, kernels_0, n_range_0, title_0, image_path_0)
@@ -207,7 +207,7 @@ def main():
         mesa_frames_polars_numba_parallel,
         mesa_frames_polars_numba_gpu,
     ]
-    n_range_1 = [k for k in range(10**6, 3 * 10**6 + 2, 10**6)]
+    n_range_1 = [*range(10**6, 3 * 10**6 + 2, 10**6)]
     title_1 = "100 steps of the SugarScape IG model:\n" + " vs ".join(labels_1)
     image_path_1 = "polars_comparison.png"
     plot_and_print_benchmark(

mesa_frames/concrete/agentset.py

@@ -669,6 +669,10 @@ def inactive_agents(self) -> pl.DataFrame:
 
     @property
     def index(self) -> pl.Series:
+        # A freshly created AgentSet has no columns at all (see ``__init__``),
+        # so there is no ``unique_id`` column to read yet.
+        if "unique_id" not in self._df.columns:
+            return pl.Series(name="unique_id", dtype=pl.UInt64)
         return self._df["unique_id"]
 
     @property

mesa_frames/concrete/agentsetregistry.py

@@ -105,7 +105,7 @@ def add(
                 agentset.name = name
             existing_names.add(name)
         new_ids = pl.concat(
-            [obj._ids] + [pl.Series(agentset["unique_id"]) for agentset in other_list]
+            [obj._ids] + [pl.Series(agentset.index) for agentset in other_list]
         )
         if new_ids.is_duplicated().any():
             raise ValueError("Some of the agent IDs are not unique.")
@@ -617,7 +617,7 @@ def _recompute_ids(self) -> None:
         Ensures `self._ids` stays a `pl.UInt64` Series and empty when no sets.
         """
         if self._agentsets:
-            cols = [pl.Series(s["unique_id"]) for s in self._agentsets]
+            cols = [pl.Series(s.index) for s in self._agentsets]
             self._ids = (
                 pl.concat(cols)
                 if cols

tests/test_agentset.py

@@ -557,6 +557,15 @@ def test_inactive_agents(self, fix1_AgentSet: ExampleAgentSet):
         agents.select(agents.df["wealth"] > 2, inplace=True)
         assert all(agents.inactive_agents["unique_id"] == agents["unique_id"][0, 1])
 
+    def test_index_empty(self):
+        # A new AgentSet has no agents and therefore no unique_id column yet.
+        agents = ExampleAgentSet(Model())
+        index = agents.index
+        assert isinstance(index, pl.Series)
+        assert index.len() == 0
+        assert index.dtype == pl.UInt64
+        assert index.name == "unique_id"
+
     def test_pos(self, fix1_AgentSet_with_pos: ExampleAgentSet):
         pos = fix1_AgentSet_with_pos.pos
         assert isinstance(pos, pl.DataFrame)

tests/test_agentsetregistry.py

@@ -96,6 +96,29 @@ def test_add(self, fix_model: Model) -> None:
         with pytest.raises(ValueError, match="agent IDs are not unique"):
             reg.add(a3)
 
+    def test_add_empty_agentset(self, fix_model: Model) -> None:
+        # An AgentSet without agents has no unique_id column yet; it must still
+        # be possible to register it (for example in ``Model.__init__``).
+        class EmptyAgentSet(AgentSet):
+            def step(self) -> None:
+                pass
+
+        reg = AgentSetRegistry(fix_model)
+        empty = EmptyAgentSet(fix_model)
+        reg.add(empty)
+        assert len(reg) == 1
+        assert empty in reg
+        assert reg.ids.len() == 0
+        # Registering a populated set next to the empty one keeps only its ids
+        populated = EmptyAgentSet(fix_model, name="populated")
+        populated.add({"wealth": [1, 2, 3]})
+        reg.add(populated)
+        assert reg.ids.len() == len(populated) == 3
+        # Removing the populated set rebuilds the ids with the empty set left
+        reg.remove(populated)
+        assert len(reg) == 1
+        assert reg.ids.len() == 0
+
     # Public: contains
     def test_contains(self, fix_registry_with_two: AgentSetRegistry) -> None:
         reg = fix_registry_with_two