diff --git a/examples/boltzmann_wealth/performance_plot.py b/examples/boltzmann_wealth/performance_plot.py index e565bda3..aa11e517 100644 --- a/examples/boltzmann_wealth/performance_plot.py +++ b/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" diff --git a/examples/sugarscape_ig/performance_comparison.py b/examples/sugarscape_ig/performance_comparison.py index d8d2f196..e253a87d 100644 --- a/examples/sugarscape_ig/performance_comparison.py +++ b/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( diff --git a/mesa_frames/concrete/agentset.py b/mesa_frames/concrete/agentset.py index 3f30104b..e265ce0f 100644 --- a/mesa_frames/concrete/agentset.py +++ b/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 diff --git a/mesa_frames/concrete/agentsetregistry.py b/mesa_frames/concrete/agentsetregistry.py index 9cc3879f..b494e0ba 100644 --- a/mesa_frames/concrete/agentsetregistry.py +++ b/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 diff --git a/tests/test_agentset.py b/tests/test_agentset.py index e25b8163..6f82423e 100644 --- a/tests/test_agentset.py +++ b/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) diff --git a/tests/test_agentsetregistry.py b/tests/test_agentsetregistry.py index b2bb8d91..202d8cae 100644 --- a/tests/test_agentsetregistry.py +++ b/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