From 317bbbe90a155be88231687135a78de5e51130b1 Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Wed, 30 Oct 2024 19:59:19 +0100
Subject: [PATCH 1/8] expand ax.scatter kwargs that can be used

---
 .../basic/boltzmann_wealth_model/app.py       | 13 +++----
 mesa/visualization/components/matplotlib.py   | 37 ++++++++++++++-----
 2 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/mesa/examples/basic/boltzmann_wealth_model/app.py b/mesa/examples/basic/boltzmann_wealth_model/app.py
index 2ab6d06bf73..14b7f9ca297 100644
--- a/mesa/examples/basic/boltzmann_wealth_model/app.py
+++ b/mesa/examples/basic/boltzmann_wealth_model/app.py
@@ -7,12 +7,8 @@
 
 
 def agent_portrayal(agent):
-    size = 10
-    color = "tab:red"
-    if agent.wealth > 0:
-        size = 50
-        color = "tab:blue"
-    return {"size": size, "color": color}
+    color = agent.wealth
+    return {"color": color}
 
 
 model_params = {
@@ -28,6 +24,9 @@ def agent_portrayal(agent):
     "height": 10,
 }
 
+def post_process(ax):
+    ax.get_figure().colorbar(ax.collections[0], label="wealth", ax=ax)
+
 # Create initial model instance
 model1 = BoltzmannWealthModel(50, 10, 10)
 
@@ -36,7 +35,7 @@ def agent_portrayal(agent):
 # Under the hood these are just classes that receive the model instance.
 # You can also author your own visualization elements, which can also be functions
 # that receive the model instance and return a valid solara component.
-SpaceGraph = make_space_component(agent_portrayal)
+SpaceGraph = make_space_component(agent_portrayal, cmap='viridis', vmin=0, vmax=10, post_process=post_process)
 GiniPlot = make_plot_measure("Gini")
 
 # Create the SolaraViz page. This will automatically create a server and display the
diff --git a/mesa/visualization/components/matplotlib.py b/mesa/visualization/components/matplotlib.py
index 7e9982a7387..898886cbddc 100644
--- a/mesa/visualization/components/matplotlib.py
+++ b/mesa/visualization/components/matplotlib.py
@@ -300,6 +300,7 @@ def draw_orthogonal_grid(
     agent_portrayal: Callable,
     ax: Axes | None = None,
     draw_grid: bool = True,
+    **kwargs
 ):
     """Visualize a orthogonal grid.
 
@@ -308,6 +309,7 @@ def draw_orthogonal_grid(
         agent_portrayal: a callable that is called with the agent and returns a dict
         ax: a Matplotlib Axes instance. If none is provided a new figure and ax will be created using plt.subplots
         draw_grid: whether to draw the grid
+        kwargs: additional keyword arguments passed to ax.scatter
 
     Returns:
         Returns the Axes object with the plot drawn onto it.
@@ -324,7 +326,7 @@ def draw_orthogonal_grid(
     arguments = collect_agent_data(space, agent_portrayal, size=s_default)
 
     # plot the agents
-    _scatter(ax, arguments)
+    _scatter(ax, arguments, **kwargs)
 
     # further styling
     ax.set_xlim(-0.5, space.width - 0.5)
@@ -345,6 +347,7 @@ def draw_hex_grid(
     agent_portrayal: Callable,
     ax: Axes | None = None,
     draw_grid: bool = True,
+    **kwargs
 ):
     """Visualize a hex grid.
 
@@ -353,6 +356,7 @@ def draw_hex_grid(
         agent_portrayal: a callable that is called with the agent and returns a dict
         ax: a Matplotlib Axes instance. If none is provided a new figure and ax will be created using plt.subplots
         draw_grid: whether to draw the grid
+        kwargs: additional keyword arguments passed to ax.scatter
 
     Returns:
         Returns the Axes object with the plot drawn onto it.
@@ -385,7 +389,7 @@ def draw_hex_grid(
     arguments["loc"] = loc
 
     # plot the agents
-    _scatter(ax, arguments)
+    _scatter(ax, arguments, **kwargs)
 
     # further styling and adding of grid
     ax.set_xlim(-1, space.width + 0.5)
@@ -434,6 +438,7 @@ def draw_network(
     draw_grid: bool = True,
     layout_alg=nx.spring_layout,
     layout_kwargs=None,
+    **kwargs
 ):
     """Visualize a network space.
 
@@ -444,6 +449,7 @@ def draw_network(
         draw_grid: whether to draw the grid
         layout_alg: a networkx layout algorithm or other callable with the same behavior
         layout_kwargs: a dictionary of keyword arguments for the layout algorithm
+        kwargs: additional keyword arguments passed to ax.scatter
 
     Returns:
         Returns the Axes object with the plot drawn onto it.
@@ -479,7 +485,7 @@ def draw_network(
     arguments["loc"] = pos[arguments["loc"]]
 
     # plot the agents
-    _scatter(ax, arguments)
+    _scatter(ax, arguments, **kwargs)
 
     # further styling
     ax.set_axis_off()
@@ -497,7 +503,8 @@ def draw_network(
 
 
 def draw_continuous_space(
-    space: ContinuousSpace, agent_portrayal: Callable, ax: Axes | None = None
+    space: ContinuousSpace, agent_portrayal: Callable, ax: Axes | None = None,
+    **kwargs
 ):
     """Visualize a continuous space.
 
@@ -505,6 +512,7 @@ def draw_continuous_space(
         space: the space to visualize
         agent_portrayal: a callable that is called with the agent and returns a dict
         ax: a Matplotlib Axes instance. If none is provided a new figure and ax will be created using plt.subplots
+        kwargs: additional keyword arguments passed to ax.scatter
 
     Returns:
         Returns the Axes object with the plot drawn onto it.
@@ -527,7 +535,7 @@ def draw_continuous_space(
     arguments = collect_agent_data(space, agent_portrayal, size=s_default)
 
     # plot the agents
-    _scatter(ax, arguments)
+    _scatter(ax, arguments, **kwargs)
 
     # further visual styling
     border_style = "solid" if not space.torus else (0, (5, 10))
@@ -543,7 +551,8 @@ def draw_continuous_space(
 
 
 def draw_voroinoi_grid(
-    space: VoronoiGrid, agent_portrayal: Callable, ax: Axes | None = None
+    space: VoronoiGrid, agent_portrayal: Callable, ax: Axes | None = None,
+    **kwargs
 ):
     """Visualize a voronoi grid.
 
@@ -551,6 +560,7 @@ def draw_voroinoi_grid(
         space: the space to visualize
         agent_portrayal: a callable that is called with the agent and returns a dict
         ax: a Matplotlib Axes instance. If none is provided a new figure and ax will be created using plt.subplots
+        kwargs: additional keyword arguments passed to ax.scatter
 
     Returns:
         Returns the Axes object with the plot drawn onto it.
@@ -580,7 +590,7 @@ def draw_voroinoi_grid(
     ax.set_xlim(x_min - x_padding, x_max + x_padding)
     ax.set_ylim(y_min - y_padding, y_max + y_padding)
 
-    _scatter(ax, arguments)
+    _scatter(ax, arguments, **kwargs)
 
     for cell in space.all_cells:
         polygon = cell.properties["polygon"]
@@ -595,8 +605,16 @@ def draw_voroinoi_grid(
     return ax
 
 
-def _scatter(ax: Axes, arguments):
-    """Helper function for plotting the agents."""
+def _scatter(ax: Axes, arguments, **kwargs):
+    """Helper function for plotting the agents.
+
+    Args
+        ax: a Matplotlib Axes instance
+        arguments: the agents specific arguments for platting
+        additional keyword arguments for ax.scatter
+
+
+    """
     loc = arguments.pop("loc")
 
     x = loc[:, 0]
@@ -615,6 +633,7 @@ def _scatter(ax: Axes, arguments):
                 marker=mark,
                 zorder=z_order,
                 **{k: v[logical] for k, v in arguments.items()},
+                **kwargs,
             )
 
 

From fdd45c1176cfd643721b0b0d3e74775e8db5449f Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 30 Oct 2024 19:04:00 +0000
Subject: [PATCH 2/8] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 mesa/examples/basic/boltzmann_wealth_model/app.py |  6 +++++-
 mesa/visualization/components/matplotlib.py       | 14 ++++++--------
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/mesa/examples/basic/boltzmann_wealth_model/app.py b/mesa/examples/basic/boltzmann_wealth_model/app.py
index 14b7f9ca297..e092e2a1187 100644
--- a/mesa/examples/basic/boltzmann_wealth_model/app.py
+++ b/mesa/examples/basic/boltzmann_wealth_model/app.py
@@ -24,9 +24,11 @@ def agent_portrayal(agent):
     "height": 10,
 }
 
+
 def post_process(ax):
     ax.get_figure().colorbar(ax.collections[0], label="wealth", ax=ax)
 
+
 # Create initial model instance
 model1 = BoltzmannWealthModel(50, 10, 10)
 
@@ -35,7 +37,9 @@ def post_process(ax):
 # Under the hood these are just classes that receive the model instance.
 # You can also author your own visualization elements, which can also be functions
 # that receive the model instance and return a valid solara component.
-SpaceGraph = make_space_component(agent_portrayal, cmap='viridis', vmin=0, vmax=10, post_process=post_process)
+SpaceGraph = make_space_component(
+    agent_portrayal, cmap="viridis", vmin=0, vmax=10, post_process=post_process
+)
 GiniPlot = make_plot_measure("Gini")
 
 # Create the SolaraViz page. This will automatically create a server and display the
diff --git a/mesa/visualization/components/matplotlib.py b/mesa/visualization/components/matplotlib.py
index 898886cbddc..8d4d04d5d0a 100644
--- a/mesa/visualization/components/matplotlib.py
+++ b/mesa/visualization/components/matplotlib.py
@@ -300,7 +300,7 @@ def draw_orthogonal_grid(
     agent_portrayal: Callable,
     ax: Axes | None = None,
     draw_grid: bool = True,
-    **kwargs
+    **kwargs,
 ):
     """Visualize a orthogonal grid.
 
@@ -347,7 +347,7 @@ def draw_hex_grid(
     agent_portrayal: Callable,
     ax: Axes | None = None,
     draw_grid: bool = True,
-    **kwargs
+    **kwargs,
 ):
     """Visualize a hex grid.
 
@@ -438,7 +438,7 @@ def draw_network(
     draw_grid: bool = True,
     layout_alg=nx.spring_layout,
     layout_kwargs=None,
-    **kwargs
+    **kwargs,
 ):
     """Visualize a network space.
 
@@ -503,8 +503,7 @@ def draw_network(
 
 
 def draw_continuous_space(
-    space: ContinuousSpace, agent_portrayal: Callable, ax: Axes | None = None,
-    **kwargs
+    space: ContinuousSpace, agent_portrayal: Callable, ax: Axes | None = None, **kwargs
 ):
     """Visualize a continuous space.
 
@@ -551,8 +550,7 @@ def draw_continuous_space(
 
 
 def draw_voroinoi_grid(
-    space: VoronoiGrid, agent_portrayal: Callable, ax: Axes | None = None,
-    **kwargs
+    space: VoronoiGrid, agent_portrayal: Callable, ax: Axes | None = None, **kwargs
 ):
     """Visualize a voronoi grid.
 
@@ -608,7 +606,7 @@ def draw_voroinoi_grid(
 def _scatter(ax: Axes, arguments, **kwargs):
     """Helper function for plotting the agents.
 
-    Args
+    Args:
         ax: a Matplotlib Axes instance
         arguments: the agents specific arguments for platting
         additional keyword arguments for ax.scatter

From f02a8bee0b8dbc6377c5eda509ae96d4675f17dd Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Wed, 30 Oct 2024 20:09:03 +0100
Subject: [PATCH 3/8] updated docstring

---
 mesa/visualization/components/matplotlib.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/mesa/visualization/components/matplotlib.py b/mesa/visualization/components/matplotlib.py
index 8d4d04d5d0a..c4b02b5c900 100644
--- a/mesa/visualization/components/matplotlib.py
+++ b/mesa/visualization/components/matplotlib.py
@@ -609,8 +609,7 @@ def _scatter(ax: Axes, arguments, **kwargs):
     Args:
         ax: a Matplotlib Axes instance
         arguments: the agents specific arguments for platting
-        additional keyword arguments for ax.scatter
-
+        kwargs: additional keyword arguments for ax.scatter
 
     """
     loc = arguments.pop("loc")

From db03ebfc739c416a8d4fdb81e895b40548348aa8 Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Wed, 30 Oct 2024 20:35:42 +0100
Subject: [PATCH 4/8] small explanation as requested

---
 mesa/examples/basic/boltzmann_wealth_model/app.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mesa/examples/basic/boltzmann_wealth_model/app.py b/mesa/examples/basic/boltzmann_wealth_model/app.py
index e092e2a1187..58d05164c8a 100644
--- a/mesa/examples/basic/boltzmann_wealth_model/app.py
+++ b/mesa/examples/basic/boltzmann_wealth_model/app.py
@@ -7,7 +7,7 @@
 
 
 def agent_portrayal(agent):
-    color = agent.wealth
+    color = agent.wealth  # we are using a colormap to translate wealth to color
     return {"color": color}
 
 

From e760c514dc918004b21083816f74febc2275ecd2 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 31 Oct 2024 07:29:09 +0000
Subject: [PATCH 5/8] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 mesa/examples/basic/boltzmann_wealth_model/app.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/mesa/examples/basic/boltzmann_wealth_model/app.py b/mesa/examples/basic/boltzmann_wealth_model/app.py
index 8de27a7919b..79619721c60 100644
--- a/mesa/examples/basic/boltzmann_wealth_model/app.py
+++ b/mesa/examples/basic/boltzmann_wealth_model/app.py
@@ -1,7 +1,6 @@
 from mesa.examples.basic.boltzmann_wealth_model.model import BoltzmannWealthModel
 from mesa.visualization import (
     SolaraViz,
-    make_plot_component,
     make_space_component,
 )
 

From 617c0cc721a50bd269f35e3dc87b973022d875ff Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Thu, 31 Oct 2024 08:34:34 +0100
Subject: [PATCH 6/8] merge fix

---
 mesa/examples/basic/boltzmann_wealth_model/app.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mesa/examples/basic/boltzmann_wealth_model/app.py b/mesa/examples/basic/boltzmann_wealth_model/app.py
index 79619721c60..1c934c7aa84 100644
--- a/mesa/examples/basic/boltzmann_wealth_model/app.py
+++ b/mesa/examples/basic/boltzmann_wealth_model/app.py
@@ -2,6 +2,7 @@
 from mesa.visualization import (
     SolaraViz,
     make_space_component,
+    make_plot_component
 )
 
 
@@ -40,7 +41,7 @@ def post_process(ax):
 SpaceGraph = make_space_component(
     agent_portrayal, cmap="viridis", vmin=0, vmax=10, post_process=post_process
 )
-GiniPlot = make_plot_measure("Gini")
+GiniPlot = make_plot_component("Gini")
 
 # Create the SolaraViz page. This will automatically create a server and display the
 # visualization elements in a web browser.

From 6b0a38da7e36c4cbd38535e0b8df8581ddd0434c Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Thu, 31 Oct 2024 08:35:43 +0100
Subject: [PATCH 7/8] spaces fixes

---
 docs/overview.md | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/docs/overview.md b/docs/overview.md
index 92f7be06907..38f0ef9d2d0 100644
--- a/docs/overview.md
+++ b/docs/overview.md
@@ -177,20 +177,20 @@ def agent_portrayal(agent):
 
 model_params = {
     "N": {
-       "type": "SliderInt",
-       "value": 50,
-       "label": "Number of agents:",
-       "min": 10,
-       "max": 100,
-       "step": 1,
+        "type": "SliderInt",
+        "value": 50,
+        "label": "Number of agents:",
+        "min": 10,
+        "max": 100,
+        "step": 1,
    }
 }
 
 page = SolaraViz(
     MyModel,
     [
-       make_space_component(agent_portrayal),
-       make_plot_component("mean_age")
+        make_space_component(agent_portrayal),
+        make_plot_component("mean_age")
     ],
     model_params=model_params
 )

From 2e7e2027e310accd7e90620bcd231e839bb3998b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 31 Oct 2024 07:35:52 +0000
Subject: [PATCH 8/8] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 mesa/examples/basic/boltzmann_wealth_model/app.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/mesa/examples/basic/boltzmann_wealth_model/app.py b/mesa/examples/basic/boltzmann_wealth_model/app.py
index 1c934c7aa84..ff329ab9667 100644
--- a/mesa/examples/basic/boltzmann_wealth_model/app.py
+++ b/mesa/examples/basic/boltzmann_wealth_model/app.py
@@ -1,9 +1,5 @@
 from mesa.examples.basic.boltzmann_wealth_model.model import BoltzmannWealthModel
-from mesa.visualization import (
-    SolaraViz,
-    make_space_component,
-    make_plot_component
-)
+from mesa.visualization import SolaraViz, make_plot_component, make_space_component
 
 
 def agent_portrayal(agent):