API implementation of AgentSet and Agents

1) Agents of the same type are stored in a pd.DataFrame of the AgentSetDF class. This avoids having many missing values (which would occupy memory) if Agents or different type were stored in the same Dataframe. All agents of the model are stored in the AgentsDF class as a list of the AgentSetDF, which virtually supports the same operations as AgentSet.
2) Went with encapsulation, avoiding extensions and subclassing because they didn't always work well with storing additional attributes and it wasn't easy to extend the created classes by subclassing (as it's often done with base mesa.Agent).
3) All operations are inplace (no immutability): if we wanted to keep immutability, there couldn't be method chaining with encapsulation (since methods would have to return a pd.Dataframe) If we were to return an AgentSet. I think it also aligns well with base mesa API.
4) The "select" operations was changed to selecting "active_agents" (by a latent mask) which are the ones effectively used for the "do" operations (since usually you don't want to remove the other agents from the df but simply use a method on a part of agents of the DF). If you wanted to remove the agents you could simply use the "discard" or "remove" methods.

Author: adamamer20

docs/scripts/readme_plot.py

@@ -1,9 +1,13 @@
+from hmac import new
+from random import seed
+from typing import TYPE_CHECKING
+
 import mesa
 import numpy as np
+import pandas as pd
 import perfplot
 
-from mesa_frames.agent import AgentDF
-from mesa_frames.model import ModelDF
+from mesa_frames import AgentSetDF, ModelDF
 
 
 # Mesa implementation
@@ -57,6 +61,14 @@ def run_model(self, n_steps) -> None:
             self.step()
 
 
+"""def compute_gini(model):
+    agent_wealths = model.agents.get("wealth")
+    x = sorted(agent_wealths)
+    N = model.num_agents
+    B = sum(xi * (N - i) for i, xi in enumerate(x)) / (N * sum(x))
+    return 1 + (1 / N) - 2 * B"""
+
+
 # Mesa Frames implementation
 def mesa_frames_implementation(n_agents: int) -> None:
     model = MoneyModelDF(n_agents)
@@ -67,45 +79,48 @@ class MoneyModelDF(ModelDF):
     def __init__(self, N):
         super().__init__()
         self.num_agents = N
-        self.create_agents(N, {MoneyAgentDF: 1})
+        self.agents = self.agents.add(MoneyAgentsDF(N, model=self))
+
+    def step(self):
+        self.agents = self.agents.do("step")
 
-    def step(self, merged_mro=True):
-        self.agents = self.agents.sample(frac=1)
-        self.update_agents_masks()
-        super().step(merged_mro)
+    def run_model(self, n):
+        for _ in range(n):
+            self.step()
 
 
-class MoneyAgentDF(AgentDF):
-    dtypes: dict[str, str] = {"wealth": "int64"}
+class MoneyAgentsDF(AgentSetDF):
+    def __init__(self, n: int, model: MoneyModelDF):
+        super().__init__(model=model)
+        self.add(n, data={"wealth": np.ones(n)})
 
-    @classmethod
-    def __init__(cls):
-        super().__init__()
-        cls.model.agents.loc[cls.mask, "wealth"] = 1
+    def step(self):
+        wealthy_agents = self.agents["wealth"] > 0
+        self.select(wealthy_agents).do("give_money")
 
-    @classmethod
-    def step(cls):
-        wealthy_agents = cls.model.agents.loc[cls.mask, "wealth"] > 0
-        if wealthy_agents.any():
-            other_agents = cls.model.agents.index.isin(
-                cls.model.agents.sample(n=wealthy_agents.sum()).index
-            )
-            cls.model.agents.loc[wealthy_agents, "wealth"] -= 1
-            cls.model.agents.loc[other_agents, "wealth"] += 1
+    def give_money(self):
+        other_agents = self.agents.sample(len(self.active_agents), replace=True)
+        new_wealth = (
+            other_agents.index.value_counts()
+            .reindex(self.active_agents.index)
+            .fillna(-1)
+        )
+        self.set_attribute("wealth", self.get_attribute("wealth") + new_wealth)
 
 
 def main():
+    mesa_frames_implementation(100)
     out = perfplot.bench(
         setup=lambda n: n,
         kernels=[mesa_implementation, mesa_frames_implementation],
         labels=["mesa", "mesa-frames"],
-        n_range=[k for k in range(10, 10000, 100)],
+        n_range=[k for k in range(100, 1000, 100)],
         xlabel="Number of agents",
         equality_check=None,
         title="100 steps of the Boltzmann Wealth model",
     )
     out.show()
-    out.save("docs/images/readme_plot.png")
+    # out.save("docs/images/readme_plot.png")
 
 
 if __name__ == "__main__":

mesa_frames/__init__.py

@@ -1,3 +1,2 @@
-from .agent import AgentSetPandas, AgentDF
-from .datacollection import DataCollectorDF
+from .agent import AgentsDF, AgentSetDF
 from .model import ModelDF