diff --git a/docs/images/readme_plot_2.png b/docs/images/readme_plot_2.png deleted file mode 100644 index db95895..0000000 Binary files a/docs/images/readme_plot_2.png and /dev/null differ diff --git a/docs/scripts/readme_plot.py b/docs/scripts/readme_plot.py index 91ea578..a2de783 100644 --- a/docs/scripts/readme_plot.py +++ b/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,7 +281,8 @@ 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, @@ -266,7 +290,8 @@ def main(): # 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)],