benchmark: Split pandas agent into native and concise (#24)

docs/scripts/readme_plot.py

@@ -167,7 +167,7 @@ def give_money(self):
         )
 
 
-class MoneyAgentPandas(AgentSetPandas):
+class MoneyAgentPandasConcise(AgentSetPandas):
     def __init__(self, n: int, model: ModelDF) -> None:
         super().__init__(model)
         ## Adding the agents to the agent set
@@ -184,41 +184,59 @@ def step(self) -> None:
 
     def give_money(self):
         ## Active agents are changed to wealthy agents
-        # 1. Using a native expression
-        # self.select(self.agents['wealth'] > 0)
-        # 2. Using the __getitem__ method
+        # 1. Using the __getitem__ method
         # self.select(self["wealth"] > 0)
-        # 3. Using the fallback __getattr__ method
+        # 2. Using the fallback __getattr__ method
         self.select(self.wealth > 0)
 
         # Receiving agents are sampled (only native expressions currently supported)
         other_agents = self.agents.sample(n=len(self.active_agents), replace=True)
 
         # Wealth of wealthy is decreased by 1
-        # 1. Using a native expression
-        """b_mask = self.active_agents.index.isin(self.agents)
-        self.agents.loc[b_mask, "wealth"] -= 1"""
-        # 2. Using the __setitem__ method with self.active_agents mask
+        # 1. Using the __setitem__ method with self.active_agents mask
         # self[self.active_agents, "wealth"] -= 1
-        # 3. Using the __setitem__ method with "active" mask
+        # 2. Using the __setitem__ method with "active" mask
         self["active", "wealth"] -= 1
 
         # Compute the income of the other agents (only native expressions currently supported)
         new_wealth = other_agents.groupby("unique_id").count()
 
         # Add the income to the other agents
-        # 1. Using native expressions
-        """merged = pd.merge(
+        # 1. Using the set method
+        # self.set(attr_names="wealth", values=self["wealth"] + new_wealth["wealth"], mask=new_wealth)
+        # 2. Using the __setitem__ method
+        self[new_wealth, "wealth"] += new_wealth["wealth"]
+
+
+class MoneyAgentPandasNative(AgentSetPandas):
+    def __init__(self, n: int, model: ModelDF) -> None:
+        super().__init__(model)
+        ## Adding the agents to the agent set
+        self += pd.DataFrame({"unique_id": np.arange(n), "wealth": np.ones(n)})
+
+    def step(self) -> None:
+        # The give_money method is called
+        self.do("give_money")
+
+    def give_money(self):
+        self.select(self.agents["wealth"] > 0)
+
+        # Receiving agents are sampled (only native expressions currently supported)
+        other_agents = self.agents.sample(n=len(self.active_agents), replace=True)
+
+        # Wealth of wealthy is decreased by 1
+        b_mask = self.active_agents.index.isin(self.agents)
+        self.agents.loc[b_mask, "wealth"] -= 1
+
+        # Compute the income of the other agents (only native expressions currently supported)
+        new_wealth = other_agents.groupby("unique_id").count()
+
+        # Add the income to the other agents
+        merged = pd.merge(
             self.agents, new_wealth, on="unique_id", how="left", suffixes=("", "_new")
         )
         merged["wealth"] = merged["wealth"] + merged["wealth_new"].fillna(0)
-        self.agents = merged.drop(columns=["wealth_new"])"""
-
-        # 2. Using the set method
-        # self.set(attr_names="wealth", values=self["wealth"] + new_wealth["wealth"], mask=new_wealth)
-
-        # 3. Using the __setitem__ method
-        self[new_wealth, "wealth"] += new_wealth["wealth"]
+        self.agents = merged.drop(columns=["wealth_new"])
 
 
 class MoneyModelDF(ModelDF):
@@ -246,8 +264,13 @@ def mesa_frames_polars_native(n_agents: int) -> None:
     model.run_model(100)
 
 
-def mesa_frames_pandas(n_agents: int) -> None:
-    model = MoneyModelDF(n_agents, MoneyAgentPandas)
+def mesa_frames_pandas_concise(n_agents: int) -> None:
+    model = MoneyModelDF(n_agents, MoneyAgentPandasConcise)
+    model.run_model(100)
+
+
+def mesa_frames_pandas_native(n_agents: int) -> None:
+    model = MoneyModelDF(n_agents, MoneyAgentPandasNative)
     model.run_model(100)
 
 
@@ -258,15 +281,17 @@ def main():
         # "mesa",
         "mesa-frames (pl concise)",
         "mesa-frames (pl native)",
-        "mesa-frames (pandas)",
+        "mesa-frames (pd concise)",
+        "mesa-frames (pd native)",
     ]
     out = perfplot.bench(
         setup=lambda n: n,
         kernels=[
             # mesa_implementation,
             mesa_frames_polars_concise,
             mesa_frames_polars_native,
-            mesa_frames_pandas,
+            mesa_frames_pandas_concise,
+            mesa_frames_pandas_native,
         ],
         labels=labels,
         n_range=[k for k in range(100, 10000, 1000)],