Generate automatically unique_id for pandas agent set

Author: luca-patrignani

mesa_frames/concrete/pandas/agentset.py

@@ -51,8 +51,13 @@ def step(self):
 refer to the class docstring.
 """
 
+from collections import defaultdict
 from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
+import functools
+import itertools
+from logging import warning
 from typing import TYPE_CHECKING
+import warnings
 
 import numpy as np
 import pandas as pd
@@ -79,6 +84,7 @@ class AgentSetPandas(AgentSetDF, PandasMixin):
     
     """
 
+    _ids = defaultdict(functools.partial(itertools.count, 0))
     _agents: pd.DataFrame
     _mask: pd.Series
     _copy_with_method: dict[str, tuple[str, list[str]]] = {
@@ -113,23 +119,26 @@ def add(  # noqa : D102
         obj = self._get_obj(inplace)
         if isinstance(agents, pd.DataFrame):
             new_agents = agents
-            if "unique_id" != agents.index.name:
-                try:
-                    new_agents.set_index("unique_id", inplace=True, drop=True)
-                except KeyError:
-                    raise KeyError("DataFrame must have a unique_id column/index.")
+            if "unique_id" == agents.index.name:
+                warnings.warn("DataFrame has 'unique_id' column, which will be ignored.")
+                new_agents = new_agents.reindex([next(self._ids[self.model]) for _ in agents.index])
         elif isinstance(agents, dict):
-            if "unique_id" not in agents:
-                raise KeyError("Dictionary must have a unique_id key.")
-            index = agents.pop("unique_id")
-            if not isinstance(index, list):
-                index = [index]
+            if "unique_id" in agents:
+                warnings.warn("Dictionary contains a 'unique_id' key, which will be ignored.")
+                if isinstance(agents["unique_id"], list):
+                    index = [next(self._ids[self.model]) for _ in agents["unique_id"]]
+                else:
+                    index = [next(self._ids[self.model])]
+            agents.pop("unique_id")
             new_agents = pd.DataFrame(agents, index=pd.Index(index, name="unique_id"))
         else:
             if len(agents) != len(obj._agents.columns) + 1:
                 raise ValueError(
                     "Length of data must match the number of columns in the AgentSet if being added as a Collection."
                 )
+            if len(agents) == len(obj._agents.columns):
+                # we suppose the first element of the list is unique_id
+                agents[0] = next(self._ids[self.model])
             columns = pd.Index(["unique_id"]).append(obj._agents.columns.copy())
             new_agents = pd.DataFrame([agents], columns=columns).set_index(
                 "unique_id", drop=True

tests/pandas/test_agentset_pandas.py

@@ -11,9 +11,9 @@
 
 @tg.typechecked
 class ExampleAgentSetPandas(AgentSetPandas):
-    def __init__(self, model: ModelDF, index: pd.Index):
+    def __init__(self, model: ModelDF):
         super().__init__(model)
-        self.starting_wealth = pd.Series([1, 2, 3, 4], name="wealth", index=index)
+        self.starting_wealth = pd.Series([1, 2, 3, 4], name="wealth")
 
     def add_wealth(self, amount: int) -> None:
         self.agents["wealth"] += amount
@@ -25,7 +25,7 @@ def step(self) -> None:
 @pytest.fixture
 def fix1_AgentSetPandas() -> ExampleAgentSetPandas:
     model = ModelDF()
-    agents = ExampleAgentSetPandas(model, pd.Index([0, 1, 2, 3], name="unique_id"))
+    agents = ExampleAgentSetPandas(model)
     agents.add({"unique_id": [0, 1, 2, 3]})
     agents["wealth"] = agents.starting_wealth
     agents["age"] = [10, 20, 30, 40]
@@ -36,7 +36,7 @@ def fix1_AgentSetPandas() -> ExampleAgentSetPandas:
 @pytest.fixture
 def fix2_AgentSetPandas() -> ExampleAgentSetPandas:
     model = ModelDF()
-    agents = ExampleAgentSetPandas(model, pd.Index([4, 5, 6, 7], name="unique_id"))
+    agents = ExampleAgentSetPandas(model)
     agents.add({"unique_id": [4, 5, 6, 7]})
     agents["wealth"] = agents.starting_wealth + 10
     agents["age"] = [100, 200, 300, 400]
@@ -55,7 +55,7 @@ def fix1_AgentSetPandas_with_pos(fix1_AgentSetPandas) -> ExampleAgentSetPandas:
 class Test_AgentSetPandas:
     def test__init__(self):
         model = ModelDF()
-        agents = ExampleAgentSetPandas(model, pd.Index([0, 1, 2, 3]))
+        agents = ExampleAgentSetPandas(model)
         assert agents.model == model
         assert isinstance(agents.agents, pd.DataFrame)
         assert agents.agents.index.name == "unique_id"
@@ -86,7 +86,6 @@ def test_add(
         # Test with a dict[str, Any]
         agents.add({"unique_id": [4, 5], "wealth": [5, 6], "age": [50, 60]})
         assert agents.agents.wealth.tolist() == [1, 2, 3, 4, 5, 6]
-        assert agents.agents.index.tolist() == [0, 1, 2, 3, 4, 5]
         assert agents.agents.age.tolist() == [10, 20, 30, 40, 50, 60]
         assert agents.agents.index.name == "unique_id"
 
@@ -287,7 +286,6 @@ def test__add__(
 
         # Test with an AgentSetPandas and a dict
         agents3 = agents + {"unique_id": 10, "wealth": 5}
-        assert agents3.agents.index.tolist() == [0, 1, 2, 3, 10]
         assert agents3.agents.wealth.tolist() == [1, 2, 3, 4, 5]
 
     def test__contains__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
@@ -360,7 +358,6 @@ def test__iadd__(
         # Test with an AgentSetPandas and a dict
         agents = deepcopy(fix1_AgentSetPandas)
         agents += {"unique_id": 10, "wealth": 5}
-        assert agents.agents.index.tolist() == [0, 1, 2, 3, 10]
         assert agents.agents.wealth.tolist() == [1, 2, 3, 4, 5]
 
     def test__iter__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
@@ -439,24 +436,24 @@ def test_agents(
 
         # Test agents.setter
         agents.agents = agents2.agents
-        assert agents.agents.index.tolist() == [4, 5, 6, 7]
+        assert len(agents.active_agents) == 4
 
     def test_active_agents(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
 
         # Test with select
         agents.select(agents["wealth"] > 2, inplace=True)
-        assert agents.active_agents.index.tolist() == [2, 3]
+        assert len(agents.active_agents) == 2
 
         # Test with active_agents.setter
         agents.active_agents = agents.agents.wealth > 2
-        assert agents.active_agents.index.to_list() == [2, 3]
+        assert len(agents.active_agents) == 2
 
     def test_inactive_agents(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
 
         agents.select(agents["wealth"] > 2, inplace=True)
-        assert agents.inactive_agents.index.to_list() == [0, 1]
+        assert len(agents.active_agents) == 2
 
     def test_pos(self, fix1_AgentSetPandas_with_pos: ExampleAgentSetPandas):
         pos = fix1_AgentSetPandas_with_pos.pos