feat: Adding agent_id

Author: AgrimGupta04

mesa_frames/concrete/datacollector.py

@@ -172,28 +172,57 @@ def _collect_model_reporters(self, current_model_step: int, batch_id: int):
 
     def _collect_agent_reporters(self, current_model_step: int, batch_id: int):
         """
-        Collect agent-level data using the agent_reporters.
+        Collect agent-level data using the agent_reporters, including unique agent IDs
 
         Constructs a LazyFrame with one column per reporter and
-        includes `step` and `seed` metadata. Appends it to internal storage.
+        includes 
+        - agent_id : unique identifier for each agent
+        - step, seed and batch columns for context.
+        - Columns for all requested agent reporters.
         """
-        agent_data_dict = {}
+        all_agent_frames = []
+
         for col_name, reporter in self._agent_reporters.items():
             if isinstance(reporter, str):
-                for k, v in self._model.sets[reporter].items():
-                    agent_data_dict[col_name + "_" + str(k.__class__.__name__)] = v
+                agent_set = self._model.sets[reporter]
+
+                if hasattr(agent_set, "df"):
+                    df = agent_set.df.select(["id", col_name]).rename({"id": "agent_id"})
+                elif hasattr(agent_set, "to_polars"):
+                    df = agent_set.to_polars().select(["id", col_name]).rename({"id": "agent_id"})
+                else:
+                    records = []
+                    for agent in agent_set.values():
+                        if hasattr(agent, "unique_id"):
+                            agent_id = agent.unique_id
+                        elif hasattr(agent, "id"):
+                            agent_id = agent.id
+                        else:
+                            agent_id = None
+                        records.append({"agent_id": agent_id, col_name: getattr(agent, col_name, None)})
+                    df = pl.DataFrame(records)
             else:
-                agent_data_dict[col_name] = reporter(self._model)
-        agent_lazy_frame = pl.LazyFrame(agent_data_dict)
-        agent_lazy_frame = agent_lazy_frame.with_columns(
-            [
+                df = reporter(self._model)
+                if not isinstance(df, pl.DataFrame):
+                    raise TypeError(f"Agent reporter {col_name} must return a Polars DataFrame")
+                
+            df = df.with_columns([
                 pl.lit(current_model_step).alias("step"),
                 pl.lit(str(self.seed)).alias("seed"),
                 pl.lit(batch_id).alias("batch"),
-            ]
-        )
-        self._frames.append(("agent", current_model_step, batch_id, agent_lazy_frame))
-
+            ])
+            all_agent_frames.append(df)
+
+        if all_agent_frames:
+            merged_df = all_agent_frames[0]
+            for next_df in all_agent_frames[1:]:
+                if "agent_id" not in next_df.columns:
+                    continue
+                merged_df = merged_df.join(next_df, on=["agent_id", "step", "seed", "batch"], how="outer")
+
+            agent_lazy_frame = merged_df.lazy()
+            self._frames.append(("agent", current_model_step, batch_id, agent_lazy_frame))
+                
     @property
     def data(self) -> dict[str, pl.DataFrame]:
         """