Let add agents without specifing unique_id for polars and adapt AgentsDF for handling the unique_id

Author: luca-patrignani

mesa_frames/abstract/agents.py

@@ -382,6 +382,23 @@ def sort(
             A new or updated AgentContainer.
         """
 
+    @abstractmethod
+    def shift_indexes(self, first_index: int, inplace: bool = True) -> Self:
+        """Shift the indexes of the agents in the AgentContainer by the specified amount.
+
+        Parameters
+        ----------
+        first_index : int
+            The new first index to be used.
+        inplace : bool
+            Whether to shift the indexes in place.
+
+        Returns
+        -------
+        Self
+            A new or updated AgentContainer.
+        """
+
     def __add__(
         self, other: DataFrameInput | AgentSetDF | Collection[AgentSetDF]
     ) -> Self:

mesa_frames/concrete/agents.py

@@ -110,13 +110,11 @@ def add(
         other_list = obj._return_agentsets_list(agents)
         if obj._check_agentsets_presence(other_list).any():
             raise ValueError("Some agentsets are already present in the AgentsDF.")
-        new_ids = pl.concat(
-            [obj._ids] + [pl.Series(agentset["unique_id"]) for agentset in other_list]
-        )
-        if new_ids.is_duplicated().any():
-            raise ValueError("Some of the agent IDs are not unique.")
-        obj._agentsets.extend(other_list)
-        obj._ids = new_ids
+        for agentset in other_list:
+            if len(obj._agentsets) > 0: 
+                agentset.shift_indexes(obj._ids.max() + 1)
+            obj._agentsets.append(agentset)
+            obj._ids = pl.concat([obj._ids, pl.Series(agentset["unique_id"])])
         return obj
 
     @overload
@@ -607,3 +605,11 @@ def index(self) -> dict[AgentSetDF, Index]:
     @property
     def pos(self) -> dict[AgentSetDF, DataFrame]:
         return {agentset: agentset.pos for agentset in self._agentsets}
+    
+    def shift_indexes(self, first_index: int, inplace: bool = True) -> Self:
+        obj = self._get_obj(inplace)
+        obj._ids += first_index
+        for agentset in obj._agentsets:
+            agentset.shift_indexes(first_index)
+            first_index += len(agentset)
+        return obj

mesa_frames/concrete/pandas/agentset.py

@@ -279,6 +279,11 @@ def to_polars(self) -> AgentSetPolars:
         new_obj._agents = pl.DataFrame(self._agents)
         new_obj._mask = pl.Series(self._mask)
         return new_obj
+    
+    def shift_indexes(self, first_index: int, inplace: bool = True):
+        obj = self._get_obj(inplace)
+        obj._agents.index = np.arange(first_index, first_index + len(obj._agents))
+        return obj
 
     def _concatenate_agentsets(
         self,

mesa_frames/concrete/polars/agentset.py

@@ -59,6 +59,7 @@ def step(self):
 
 from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
 from typing import TYPE_CHECKING
+import warnings
 
 import polars as pl
 from polars._typing import IntoExpr
@@ -120,20 +121,28 @@ def add(
         """
         obj = self._get_obj(inplace)
         if isinstance(agents, pl.DataFrame):
-            if "unique_id" not in agents.columns:
-                raise KeyError("DataFrame must have a unique_id column.")
+            if "unique_id" in agents.columns:
+                warnings.warn("Dataframe should not have a unique_id index/column. It will be ignored.")
             new_agents = agents
         elif isinstance(agents, dict):
-            if "unique_id" not in agents:
-                raise KeyError("Dictionary must have a unique_id key.")
+            if "unique_id" in agents:
+                warnings.warn("Dictionary should not have a unique_id key. It will be ignored.")
             new_agents = pl.DataFrame(agents)
         else:
-            if len(agents) != len(obj._agents.columns):
+            if len(agents) not in {len(obj._agents.columns), 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):
+                warnings.warn("Length of data should have the number of columns in the AgentSet," +
+                    "we suppose the first element is the unique_id. It will be ignored.")
             new_agents = pl.DataFrame([agents], schema=obj._agents.schema)
 
+        if len(self.agents) == 0:
+            unique_ids = pl.arange(len(new_agents))
+        else:
+            unique_ids = pl.arange(self.index.max() + 1, self.index.max() + 1 + len(new_agents))
+        new_agents = new_agents.with_columns(unique_ids.alias("unique_id"))
         if new_agents["unique_id"].dtype != pl.Int64:
             raise TypeError("unique_id column must be of type int64.")
 
@@ -293,7 +302,14 @@ def to_pandas(self) -> "AgentSetPandas":
                 .to_pandas()
             )
         return new_obj
-
+    
+    def shift_indexes(self, first_index: int, inplace: bool = True):
+        obj = self._get_obj(inplace)
+        obj._agents = obj._agents.with_columns(
+            pl.arange(first_index, first_index + len(obj._agents)).alias("unique_id")
+        )
+        return obj
+    
     def _concatenate_agentsets(
         self,
         agentsets: Iterable[Self],

tests/pandas/test_agentset_pandas.py

@@ -513,3 +513,8 @@ def test_pos(self, fix1_AgentSetPandas_with_pos: ExampleAgentSetPandas):
         assert all(math.isnan(val) for val in pos["dim_0"].tolist()[2:])
         assert pos["dim_1"].tolist()[:2] == [0, 1]
         assert all(math.isnan(val) for val in pos["dim_1"].tolist()[2:])
+
+    def test_shift_indexes(self, fix1_AgentSetPandas_with_unique_id: ExampleAgentSetPandas):
+        agents = fix1_AgentSetPandas_with_unique_id
+        agents.shift_indexes(10, inplace=True)
+        assert agents.agents.index.tolist() == [10, 11, 12, 13]

tests/polars/test_agentset_polars.py

@@ -43,7 +43,7 @@ def fix2_AgentSetPolars_with_unique_id() -> ExampleAgentSetPolars:
     model.agents.add(agents)
     space = GridPandas(model, dimensions=[3, 3], capacity=2)
     model.space = space
-    space.place_agents(agents=[4, 5], pos=[[2, 1], [1, 2]])
+    space.place_agents(agents=[0, 1], pos=[[2, 1], [1, 2]])
     return agents
 
 
@@ -83,7 +83,7 @@ def test_add_with_unique_id(
 
         # Test with a list (Sequence[Any])
         result = agents.add([10, 5, 10], inplace=False)
-        assert result.agents["unique_id"].to_list() == [0, 1, 2, 3, 10]
+        assert result.agents["unique_id"].to_list() == [0, 1, 2, 3, 4]
         assert result.agents["wealth"].to_list() == [1, 2, 3, 4, 5]
         assert result.agents["age"].to_list() == [10, 20, 30, 40, 10]
 
@@ -290,7 +290,7 @@ def test__add__(
 
         # Test with an AgentSetPolars and a dict
         agents3 = agents + {"unique_id": 10, "wealth": 5}
-        assert agents3.agents["unique_id"].to_list() == [0, 1, 2, 3, 10]
+        assert agents3.agents["unique_id"].to_list() == [0, 1, 2, 3, 4]
         assert agents3.agents["wealth"].to_list() == [1, 2, 3, 4, 5]
 
     def test__contains__(self, fix1_AgentSetPolars_with_unique_id: ExampleAgentSetPolars):
@@ -362,7 +362,7 @@ def test__iadd__(
         # Test with an AgentSetPolars and a dict
         agents = deepcopy(fix1_AgentSetPolars_with_unique_id)
         agents += {"unique_id": 10, "wealth": 5}
-        assert agents.agents["unique_id"].to_list() == [0, 1, 2, 3, 10]
+        assert agents.agents["unique_id"].to_list() == [0, 1, 2, 3, 4]
         assert agents.agents["wealth"].to_list() == [1, 2, 3, 4, 5]
 
     def test__iter__(self, fix1_AgentSetPolars_with_unique_id: ExampleAgentSetPolars):
@@ -441,7 +441,7 @@ def test_agents(
 
         # Test agents.setter
         agents.agents = agents2.agents
-        assert agents.agents["unique_id"].to_list() == [4, 5, 6, 7]
+        assert agents.agents["unique_id"].to_list() == [0, 1, 2, 3]
 
     def test_active_agents(self, fix1_AgentSetPolars_with_unique_id: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars_with_unique_id
@@ -467,3 +467,8 @@ def test_pos(self, fix1_AgentSetPolars_with_pos: ExampleAgentSetPolars):
         assert pos.columns == ["unique_id", "dim_0", "dim_1"]
         assert pos["dim_0"].to_list() == [0, 1, None, None]
         assert pos["dim_1"].to_list() == [0, 1, None, None]
+
+    def test_shift_indexes(self, fix1_AgentSetPolars_with_unique_id: ExampleAgentSetPolars):
+        agents = fix1_AgentSetPolars_with_unique_id
+        agents.shift_indexes(10, inplace=True)
+        assert agents.agents["unique_id"].to_list() == [10, 11, 12, 13]

tests/test_agents.py

@@ -77,10 +77,6 @@ def test_add(
             + agentset_polars._agents["unique_id"].to_list()
         )
 
-        # Test if adding the same AgentSetDF raises ValueError
-        with pytest.raises(ValueError):
-            agents.add(agentset_pandas, inplace=False)
-
     def test_contains(
         self, fix2_AgentSetPandas_with_unique_id: ExampleAgentSetPandas, fix_AgentsDF: AgentsDF
     ):
@@ -629,10 +625,6 @@ def test___add__(
             + agentset_polars._agents["unique_id"].to_list()
         )
 
-        # Test if adding the same AgentSetDF raises ValueError
-        with pytest.raises(ValueError):
-            result + agentset_pandas
-
     def test___contains__(
         self, fix_AgentsDF: AgentsDF, fix2_AgentSetPandas_with_unique_id: ExampleAgentSetPandas
     ):
@@ -995,3 +987,8 @@ def test_inactive_agents(self, fix_AgentsDF: AgentsDF):
                 == agents1._agentsets[1].select(mask1, negate=True).active_agents
             )
         )
+
+    def test_shift_indexes(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+        agents.shift_indexes(20)
+        assert agents._ids.to_list() == [20, 21, 22, 23, 24, 25, 26, 27]