distinction between abstract and concrete implementation, test implem…

…entations, update script

docs/scripts/readme_plot.py

@@ -1,16 +1,15 @@
-from hmac import new
-from random import seed
-from typing import TYPE_CHECKING
-
+import matplotlib.pyplot as plt
 import mesa
 import numpy as np
 import pandas as pd
 import perfplot
+import polars as pl
+import seaborn as sns
 
-from mesa_frames import AgentSetDF, ModelDF
+from mesa_frames import AgentSetPandas, AgentSetPolars, ModelDF
 
 
-# Mesa implementation
+### ---------- Mesa implementation ---------- ###
 def mesa_implementation(n_agents: int) -> None:
     model = MoneyModel(n_agents)
     model.run_model(100)
@@ -39,7 +38,7 @@ class MoneyModel(mesa.Model):
     """A model with some number of agents."""
 
     def __init__(self, N):
-        super().__init__
+        super().__init__()
         self.num_agents = N
         # Create scheduler and assign it to the model
         self.schedule = mesa.time.RandomActivation(self)
@@ -69,58 +68,178 @@ def run_model(self, n_steps) -> None:
     return 1 + (1 / N) - 2 * B"""
 
 
-# Mesa Frames implementation
-def mesa_frames_implementation(n_agents: int) -> None:
-    model = MoneyModelDF(n_agents)
-    model.run_model(100)
+### ---------- Mesa-frames implementation ---------- ###
+
+
+class MoneyAgentPolars(AgentSetPolars):
+    def __init__(self, n: int, model: ModelDF):
+        super().__init__(model)
+        ## Adding the agents to the agent set
+        # 1. Changing the agents attribute directly (not reccomended, if other agents were added before, they will be lost)
+        """self.agents = pl.DataFrame(
+            {"unique_id": pl.arange(n, eager=True), "wealth": pl.ones(n, eager=True)}
+        )"""
+        # 2. Adding the dataframe with add
+        """self.add(
+            pl.DataFrame(
+                {
+                    "unique_id": pl.arange(n, eager=True),
+                    "wealth": pl.ones(n, eager=True),
+                }
+            )
+        )"""
+        # 3. Adding the dataframe with __iadd__
+        self += pl.DataFrame(
+            {"unique_id": pl.arange(n, eager=True), "wealth": pl.ones(n, eager=True)}
+        )
+
+    def step(self) -> None:
+        # The give_money method is called
+        # self.give_money()
+        self.do("give_money")
+
+    def give_money(self):
+        ## Active agents are changed to wealthy agents
+        # 1. Using a native expression
+        # self.select(pl.col("wealth") > 0)
+        # 2. Using the __getitem__ method
+        # self.select(self["wealth"] > 0)
+        # 3. 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), with_replacement=True
+        )
+
+        # Wealth of wealthy is decreased by 1
+        # 1. Using a native expression
+        """self.agents = self.agents.with_columns(
+            wealth=pl.when(pl.col("unique_id").is_in(self.active_agents["unique_id"]))
+            .then(pl.col("wealth") - 1)
+            .otherwise(pl.col("wealth"))
+        )"""
+        # 2. Using the __setitem__ method with self.active_agents mask
+        # self[self.active_agents, "wealth"] -= 1
+        # 3. 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.group_by("unique_id").len()
+
+        # Add the income to the other agents
+        # 1. Using native expressions
+        """self.agents = self.agents.with_columns(
+            pl.when(pl.col("unique_id").is_in(new_wealth["unique_id"]))
+            .then(pl.col("wealth") + new_wealth["wealth"])
+            .otherwise(pl.col("wealth"))
+        )"""
+
+        # 2. Using the set method
+        """self.set(
+            attr_names="wealth",
+            values=pl.col("wealth") + new_wealth["len"],
+            mask=new_wealth,
+        )"""
+
+        # 3. Using the __setitem__ method
+        self[new_wealth, "wealth"] += new_wealth["len"]
+
+
+class MoneyAgentPandas(AgentSetPandas):
+    def __init__(self, n: int, model: ModelDF) -> None:
+        super().__init__(model)
+        ## Adding the agents to the agent set
+        # 1. Changing the agents attribute directly (not reccomended, if other agents were added before, they will be lost)
+        # self.agents = pd.DataFrame({"unique_id": np.arange(n), "wealth": np.ones(n)})
+        # 2. Adding the dataframe with add
+        # self.add(pd.DataFrame({"unique_id": np.arange(n), "wealth": np.ones(n)}))
+        # 3. Adding the dataframe with __iadd__
+        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):
+        ## Active agents are changed to wealthy agents
+        # 1. Using a native expression
+        # self.select(self.agents['wealth'] > 0)
+        # 2. Using the __getitem__ method
+        # self.select(self["wealth"] > 0)
+        # 3. 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
+        # self[self.active_agents, "wealth"] -= 1
+        # 3. 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(
+            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"]
 
 
 class MoneyModelDF(ModelDF):
-    def __init__(self, N):
+    def __init__(self, N: int, agents_cls):
         super().__init__()
-        self.num_agents = N
-        self.agents = self.agents.add(MoneyAgentsDF(N, model=self))
+        self.n_agents = N
+        self.agents += agents_cls(N, self)
 
     def step(self):
-        self.agents = self.agents.do("step")
+        # Executes the step method for every agentset in self.agents
+        self.agents.do("step")
 
     def run_model(self, n):
         for _ in range(n):
             self.step()
 
 
-class MoneyAgentsDF(AgentSetDF):
-    def __init__(self, n: int, model: MoneyModelDF):
-        super().__init__(model=model)
-        self.add(n, data={"wealth": np.ones(n)})
+def mesa_frames_polars(n_agents: int) -> None:
+    model = MoneyModelDF(n_agents, MoneyAgentPolars)
+    model.run_model(100)
 
-    def step(self):
-        wealthy_agents = self.agents["wealth"] > 0
-        self.select(wealthy_agents).do("give_money")
 
-    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 mesa_frames_pandas(n_agents: int) -> None:
+    model = MoneyModelDF(n_agents, MoneyAgentPandas)
+    model.run_model(100)
 
 
 def main():
-    mesa_frames_implementation(100)
+
+    sns.set_theme(style="whitegrid")
+
     out = perfplot.bench(
         setup=lambda n: n,
-        kernels=[mesa_implementation, mesa_frames_implementation],
-        labels=["mesa", "mesa-frames"],
-        n_range=[k for k in range(100, 1000, 100)],
+        kernels=[mesa_implementation, mesa_frames_polars, mesa_frames_pandas],
+        labels=["mesa", "mesa-frames (polars)", "mesa-frames (pandas)"],
+        n_range=[k for k in range(100, 10000, 1000)],
         xlabel="Number of agents",
         equality_check=None,
-        title="100 steps of the Boltzmann Wealth model",
+        title="100 steps of the Boltzmann Wealth model: mesa vs mesa-frames (polars) vs mesa-frames (pandas)",
     )
-    out.show()
-    # out.save("docs/images/readme_plot.png")
+    plt.ylabel("Execution time (s)")
+    out.save("docs/images/readme_plot_1.png")
 
 
 if __name__ == "__main__":

mesa_frames/__init__.py

@@ -1,2 +1,4 @@
-from mesa_frames.agent import AgentSetPandas, AgentSetPolars, AgentsPandas, AgentsPolars
-from mesa_frames.model import ModelDF
+from mesa_frames.concrete.agents import AgentsDF
+from mesa_frames.concrete.agentset_pandas import AgentSetPandas
+from mesa_frames.concrete.agentset_polars import AgentSetPolars
+from mesa_frames.concrete.model import ModelDF