Generate automatically unique_id

.vscode/settings.json

@@ -4,5 +4,5 @@
     ],
     "python.testing.unittestEnabled": false,
     "python.testing.pytestEnabled": true,
-    "python.testing.envFile": "${workspaceFolder}/.env.test",
+    "python.envFile": "${workspaceFolder}/.env.test"
 }

mesa_frames/abstract/space.py

@@ -436,29 +436,35 @@
         self, agents: IdsLike | AgentContainer | Collection[AgentContainer]
     ) -> Series:
         if isinstance(agents, Sized) and len(agents) == 0:
-            return self._srs_constructor([], name="agent_id")
+            return self._srs_constructor([], name="agent_id", dtype="uint64")
         if isinstance(agents, AgentSetDF):
             return self._srs_constructor(
-                self._df_index(agents.agents, "unique_id"), name="agent_id"
+                self._df_index(agents.agents, "unique_id"),
+                name="agent_id",
+                dtype="uint64",
             )
         elif isinstance(agents, AgentsDF):
-            return self._srs_constructor(agents._ids, name="agent_id")
+            return self._srs_constructor(agents._ids, name="agent_id", dtype="uint64")
         elif isinstance(agents, Collection) and (isinstance(agents[0], AgentContainer)):
             ids = []
             for a in agents:
                 if isinstance(a, AgentSetDF):
                     ids.append(
                         self._srs_constructor(
-                            self._df_index(a.agents, "unique_id"), name="agent_id"
+                            self._df_index(a.agents, "unique_id"),
+                            name="agent_id",
+                            dtype="uint64",
                         )
                     )
                 elif isinstance(a, AgentsDF):
-                    ids.append(self._srs_constructor(a._ids, name="agent_id"))
+                    ids.append(
+                        self._srs_constructor(a._ids, name="agent_id", dtype="uint64")
+                    )
             return self._df_concat(ids, ignore_index=True)
         elif isinstance(agents, int):
-            return self._srs_constructor([agents], name="agent_id")
+            return self._srs_constructor([agents], name="agent_id", dtype="uint64")
         else:  # IDsLike
-            return self._srs_constructor(agents, name="agent_id")
+            return self._srs_constructor(agents, name="agent_id", dtype="uint64")
 
     @abstractmethod
     def _place_or_move_agents(
@@ -1178,7 +1184,7 @@
         self._agents = self._df_constructor(
             columns=["agent_id"] + self._pos_col_names,
             index_cols="agent_id",
-            dtypes={col: int for col in ["agent_id"] + self._pos_col_names},
+            dtypes={"agent_id": "uint64"} | {col: int for col in self._pos_col_names},
         )
 
         cells_df_dtypes = {col: int for col in self._pos_col_names}

mesa_frames/concrete/agents.py

@@ -50,6 +50,7 @@ def step(self):
 from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
 from typing import Any, Literal, Self, cast, overload
 
+import numpy as np
 import polars as pl
 
 from mesa_frames.abstract.agents import AgentContainer, AgentSetDF
@@ -80,7 +81,7 @@ def __init__(self, model: mesa_frames.concrete.model.ModelDF) -> None:
         """
         self._model = model
         self._agentsets = []
-        self._ids = pl.Series(name="unique_id", dtype=pl.Int64)
+        self._ids = pl.Series(name="unique_id", dtype=pl.UInt64)
 
     def add(
         self, agents: AgentSetDF | Iterable[AgentSetDF], inplace: bool = True
@@ -139,7 +140,7 @@ def contains(
             else:  # IdsLike
                 agents = cast(IdsLike, agents)
 
-                return pl.Series(agents).is_in(self._ids)
+                return pl.Series(agents, dtype=pl.UInt64).is_in(self._ids)
 
     @overload
     def do(
@@ -225,19 +226,22 @@ def remove(
             # We have to get the index of the original AgentSetDF because the copy made AgentSetDFs with different hash
             ids = [self._agentsets.index(agentset) for agentset in iter(agents)]
             ids.sort(reverse=True)
-            removed_ids = pl.Series(dtype=pl.Int64)
+            removed_ids = pl.Series(dtype=pl.UInt64)
             for id in ids:
                 removed_ids = pl.concat(
-                    [removed_ids, pl.Series(obj._agentsets[id].index)]
+                    [
+                        removed_ids,
+                        pl.Series(obj._agentsets[id]["unique_id"], dtype=pl.UInt64),
+                    ]
                 )
                 obj._agentsets.pop(id)
 
         else:  # IDsLike
-            if isinstance(agents, int):
+            if isinstance(agents, (int, np.uint64)):
                 agents = [agents]
             elif isinstance(agents, DataFrame):
                 agents = agents["unique_id"]
-            removed_ids = pl.Series(agents)
+            removed_ids = pl.Series(agents, dtype=pl.UInt64)
             deleted = 0
 
             for agentset in obj._agentsets:
@@ -355,10 +359,10 @@ def _check_ids_presence(self, other: list[AgentSetDF]) -> pl.DataFrame:
         """
         presence_df = pl.DataFrame(
             data={"unique_id": self._ids, "present": True},
-            schema={"unique_id": pl.Int64, "present": pl.Boolean},
+            schema={"unique_id": pl.UInt64, "present": pl.Boolean},
         )
         for agentset in other:
-            new_ids = pl.Series(agentset.index)
+            new_ids = pl.Series(agentset.index, dtype=pl.UInt64)
             presence_df = pl.concat(
                 [
                     presence_df,

mesa_frames/concrete/agentset.py

@@ -68,7 +68,7 @@
 from mesa_frames.concrete.agents import AgentSetDF
 from mesa_frames.concrete.mixin import PolarsMixin
 from mesa_frames.concrete.model import ModelDF
-from mesa_frames.types_ import AgentPolarsMask, IdsLike, IntoExpr, PolarsIdsLike
+from mesa_frames.types_ import AgentPolarsMask, IntoExpr, PolarsIdsLike
 from mesa_frames.utils import copydoc
 
 
@@ -92,7 +92,8 @@
             The model that the agent set belongs to.
         """
         self._model = model
-        self._agents = pl.DataFrame(schema={"unique_id": pl.Int64})
+        # No definition of schema with unique_id, as it becomes hard to add new agents
+        self._agents = pl.DataFrame()
         self._mask = pl.repeat(True, len(self._agents), dtype=pl.Boolean, eager=True)
 
     def add(
@@ -115,33 +116,51 @@
             The updated AgentSetPolars.
         """
         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 isinstance(agents, AgentSetDF):
+            raise TypeError(
+                "AgentSetPolars.add() does not accept AgentSetDF objects. "
+                "Extract the DataFrame with agents.agents.drop('unique_id') first."
+            )
+        elif isinstance(agents, pl.DataFrame):
+            if "unique_id" in agents.columns:
+                raise ValueError("Dataframe should not have a unique_id column.")
             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:
+                raise ValueError("Dictionary should not have a unique_id key.")
             new_agents = pl.DataFrame(agents)
-        else:
-            if len(agents) != len(obj._agents.columns):
+        else:  # Sequence
+            if len(obj._agents) != 0:
+                # For non-empty AgentSet, check column count
+                expected_columns = len(obj._agents.columns) - 1  # Exclude unique_id
+                if len(agents) != expected_columns:
+                    raise ValueError(
+                        f"Length of data ({len(agents)}) must match the number of columns in the AgentSet (excluding unique_id): {expected_columns}"
+                    )
+                new_agents = pl.DataFrame(
+                    [list(agents)],
+                    schema=[col for col in obj._agents.schema if col != "unique_id"],
+                    orient="row",
+                )
+            else:
+                # For empty AgentSet, cannot infer schema from sequence
                 raise ValueError(
-                    "Length of data must match the number of columns in the AgentSet if being added as a Collection."
+                    "Cannot add a sequence to an empty AgentSet. Use a DataFrame or dict with column names."
                 )
-            new_agents = pl.DataFrame([agents], schema=obj._agents.schema)
 
-        if new_agents["unique_id"].dtype != pl.Int64:
-            raise TypeError("unique_id column must be of type int64.")
+        new_agents = new_agents.with_columns(
+            self._generate_unique_ids(len(new_agents)).alias("unique_id")
+        )
 
         # If self._mask is pl.Expr, then new mask is the same.
         # If self._mask is pl.Series[bool], then new mask has to be updated.
-
-        if isinstance(obj._mask, pl.Series):
+        originally_empty = len(obj._agents) == 0
+        if isinstance(obj._mask, pl.Series) and not originally_empty:
             original_active_indices = obj._agents.filter(obj._mask)["unique_id"]
 
         obj._agents = pl.concat([obj._agents, new_agents], how="diagonal_relaxed")
 
-        if isinstance(obj._mask, pl.Series):
+        if isinstance(obj._mask, pl.Series) and not originally_empty:
             obj._update_mask(original_active_indices, new_agents["unique_id"])
 
         return obj
@@ -158,8 +177,8 @@
     ) -> bool | pl.Series:
         if isinstance(agents, pl.Series):
             return agents.is_in(self._agents["unique_id"])
-        elif isinstance(agents, Collection):
-            return pl.Series(agents).is_in(self._agents["unique_id"])
+        elif isinstance(agents, Collection) and not isinstance(agents, str):
+            return pl.Series(agents, dtype=pl.UInt64).is_in(self._agents["unique_id"])
         else:
             return agents in self._agents["unique_id"]
 
@@ -185,7 +204,6 @@
         inplace: bool = True,
     ) -> Self:
         obj = self._get_obj(inplace)
-        b_mask = obj._get_bool_mask(mask)
         masked_df = obj._get_masked_df(mask)
 
         if not attr_names:
@@ -196,17 +214,12 @@
             masked_df: pl.DataFrame, attr_name: str, values: Any
         ) -> pl.DataFrame:
             if isinstance(values, pl.DataFrame):
-                return masked_df.with_columns(values.to_series().alias(attr_name))
-            elif isinstance(values, pl.Expr):
-                return masked_df.with_columns(values.alias(attr_name))
-            if isinstance(values, pl.Series):
-                return masked_df.with_columns(values.alias(attr_name))
+                values_series = values.to_series()
+            elif isinstance(values, (pl.Expr, pl.Series, Collection)):
+                values_series = pl.Series(values)
             else:
-                if isinstance(values, Collection):
-                    values = pl.Series(values)
-                else:
-                    values = pl.repeat(values, len(masked_df))
-                return masked_df.with_columns(values.alias(attr_name))
+                values_series = pl.repeat(values, len(masked_df))
+            return masked_df.with_columns(values_series.alias(attr_name))
 
         if isinstance(attr_names, str) and values is not None:
             masked_df = process_single_attr(masked_df, attr_names, values)
@@ -224,10 +237,19 @@
             raise ValueError(
                 "attr_names must be a string, a collection of string or a dictionary with columns as keys and values."
             )
+        unique_id_column = None
+        if "unique_id" not in obj._agents:
+            unique_id_column = self._generate_unique_ids(len(masked_df)).alias(
+                "unique_id"
+            )
+            obj._agents = obj._agents.with_columns(unique_id_column)
+            masked_df = masked_df.with_columns(unique_id_column)
+        b_mask = obj._get_bool_mask(mask)
         non_masked_df = obj._agents.filter(b_mask.not_())
         original_index = obj._agents.select("unique_id")
         obj._agents = pl.concat([non_masked_df, masked_df], how="diagonal_relaxed")
         obj._agents = original_index.join(obj._agents, on="unique_id", how="left")
+        obj._update_mask(original_index["unique_id"], unique_id_column)
         return obj
 
     def select(
@@ -368,9 +390,9 @@
         elif mask == "active":
             return self._mask
         elif isinstance(mask, Collection):
-            return bool_mask_from_series(pl.Series(mask))
+            return bool_mask_from_series(pl.Series(mask, dtype=pl.UInt64))
         else:
-            return bool_mask_from_series(pl.Series([mask]))
+            return bool_mask_from_series(pl.Series([mask], dtype=pl.UInt64))
 
     def _get_masked_df(
         self,
@@ -401,9 +423,9 @@
             return self._agents.filter(self._mask)
         else:
             if isinstance(mask, Collection):
-                mask_series = pl.Series(mask)
+                mask_series = pl.Series(mask, dtype=pl.UInt64)
             else:
-                mask_series = pl.Series([mask])
+                mask_series = pl.Series([mask], dtype=pl.UInt64)
             if not mask_series.is_in(self._agents["unique_id"]).all():
                 raise KeyError(
                     "Some 'unique_id' of mask are not present in DataFrame 'unique_id'."
@@ -447,6 +469,11 @@
         super().__getattr__(key)
         return self._agents[key]
 
+    def _generate_unique_ids(self, n: int) -> pl.Series:
+        return pl.Series(
+            self.random.integers(1, np.iinfo(np.uint64).max, size=n, dtype=np.uint64)
+        )
+
     @overload
     def __getitem__(
         self,

mesa_frames/concrete/mixin.py

@@ -58,7 +58,11 @@ class PolarsMixin(DataFrameMixin):
     """Polars-specific implementation of DataFrame operations."""
 
     # TODO: complete with other dtypes
-    _dtypes_mapping: dict[str, Any] = {"int64": pl.Int64, "bool": pl.Boolean}
+    _dtypes_mapping: dict[str, Any] = {
+        "int64": pl.Int64,
+        "bool": pl.Boolean,
+        "uint64": pl.UInt64,
+    }
 
     def _df_add(
         self,
@@ -599,7 +603,7 @@ def _df_with_columns(
             )
         ):
             # This means that data is a Sequence of Sequences (rows)
-            data = pl.DataFrame(data, new_columns)
+            data = pl.DataFrame(data, new_columns, orient="row")
             original_df = original_df.select(pl.exclude(data.columns))
             return original_df.hstack(data)
         if not isinstance(data, dict):
@@ -619,7 +623,7 @@ def _srs_constructor(
     ) -> pl.Series:
         if dtype is not None:
             dtype = self._dtypes_mapping[dtype]
-        return pl.Series(name=name, values=data, dtype=dtype)
+        return pl.Series(name=name, values=list(data), dtype=dtype)
 
     def _srs_contains(
         self,
@@ -628,7 +632,7 @@ def _srs_contains(
     ) -> pl.Series:
         if not isinstance(values, Collection):
             values = [values]
-        return pl.Series(values).is_in(srs)
+        return pl.Series(values).is_in(pl.Series(srs))
 
     def _srs_range(
         self,

mesa_frames/types_.py

@@ -53,6 +53,7 @@
 AgentPolarsMask = AgnosticAgentMask | pl.Expr | pl.Series | pl.DataFrame | Sequence[int]
 PolarsIdsLike = AgnosticIds | pl.Series | pl.DataFrame
 PolarsGridCapacity = list[pl.Expr]
+# Had to define this because beartype has issues with polars._typing.IntoExpr
 IntoExpr = (
     int
     | float