From a0caa1ec199ad7390c3f0805589a97f2e3db454d Mon Sep 17 00:00:00 2001 From: Jay Naik Date: Thu, 5 Nov 2026 19:40:35 +0530 Subject: [PATCH 1/3] Fix registering an AgentSet that has no agents yet A new AgentSet is created with an empty DataFrame that has no columns, so it has no unique_id column until agents are added. AgentSetRegistry.add and _recompute_ids read agentset["unique_id"] directly and raised polars ColumnNotFoundError for such a set. AgentSet.index now returns an empty UInt64 Series named unique_id when the column does not exist yet, and the registry uses it instead of indexing the column. This also fixes AgentSet.pos for an empty set, since it is based on index. Adds tests for registering an empty AgentSet (including removing another set while it remains registered) and for the index of an empty AgentSet. Fixes #204 --- mesa_frames/concrete/agentset.py | 4 ++++ mesa_frames/concrete/agentsetregistry.py | 4 ++-- tests/test_agentset.py | 9 +++++++++ tests/test_agentsetregistry.py | 22 ++++++++++++++++++++++ 4 files changed, 37 insertions(+), 2 deletions(-) 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..4de99dfe 100644 --- a/tests/test_agentsetregistry.py +++ b/tests/test_agentsetregistry.py @@ -96,6 +96,28 @@ 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 = ExampleAgentSetA(fix_model) + reg.add(populated) + assert reg.ids.len() == len(populated) + # 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 From a44c6e17997a520ad79f2123ce085756f35bc720 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 5 Nov 2026 14:33:03 +0000 Subject: [PATCH 2/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/boltzmann_wealth/performance_plot.py | 4 ++-- examples/sugarscape_ig/performance_comparison.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) 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( From 7fbad9ce1d68d7a0cf1fccb2f85c3157c4c40fd1 Mon Sep 17 00:00:00 2001 From: Jay Naik Date: Thu, 5 Nov 2026 20:15:04 +0530 Subject: [PATCH 3/3] Make the empty AgentSet registry test independent of the shared fixtures The ExampleAgentSetA fixture sets a column on an empty AgentSet, which fails with polars 2.x (ShapeError: can't broadcast Series of length 4 to length 0). Build the populated set with AgentSet.add instead, so the new test passes with both polars 1.36.1 and 2.1.0. --- tests/test_agentsetregistry.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_agentsetregistry.py b/tests/test_agentsetregistry.py index 4de99dfe..202d8cae 100644 --- a/tests/test_agentsetregistry.py +++ b/tests/test_agentsetregistry.py @@ -110,9 +110,10 @@ def step(self) -> None: assert empty in reg assert reg.ids.len() == 0 # Registering a populated set next to the empty one keeps only its ids - populated = ExampleAgentSetA(fix_model) + populated = EmptyAgentSet(fix_model, name="populated") + populated.add({"wealth": [1, 2, 3]}) reg.add(populated) - assert reg.ids.len() == len(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