From a7a6c4b1b1435d23c7815784faba28f580c34fd3 Mon Sep 17 00:00:00 2001
From: Ankit Kumar <ankit.kumar@tum.de>
Date: Fri, 8 Dec 2023 06:16:07 +0100
Subject: [PATCH 01/11] introduce altair grid

---
 mesa/experimental/altair_grid.py | 69 ++++++++++++++++++++++++++++++++
 mesa/experimental/jupyter_viz.py | 16 ++++++++
 2 files changed, 85 insertions(+)
 create mode 100644 mesa/experimental/altair_grid.py

diff --git a/mesa/experimental/altair_grid.py b/mesa/experimental/altair_grid.py
new file mode 100644
index 00000000000..851aae0120c
--- /dev/null
+++ b/mesa/experimental/altair_grid.py
@@ -0,0 +1,69 @@
+import json
+from typing import Callable
+
+import altair as alt
+import solara
+
+import mesa
+
+
+def get_agent_data_from_coord_iter(data):
+    for agent, (x, y) in data:
+        if agent:
+            print(agent[0], x, y)
+            agent_data = json.loads(
+                json.dumps(agent[0].__dict__, skipkeys=True, default=str)
+            )
+            agent_data["x"] = x
+            agent_data["y"] = y
+            agent_data.pop("model", None)
+            agent_data.pop("pos", None)
+            yield agent_data
+
+
+def create_grid(
+    color: str | None = None,
+    on_click: Callable[[mesa.Model, mesa.space.Coordinate], None] | None = None,
+) -> Callable[[mesa.Model], solara.component]:
+    return lambda model: Grid(model, color, on_click)
+
+
+def Grid(model, color=None, on_click=None):
+    if color is None:
+        color = "unique_id:N"
+
+    if color[-2] != ":":
+        color = color + ":N"
+
+    print(model.grid.coord_iter())
+
+    data = solara.reactive(
+        list(get_agent_data_from_coord_iter(model.grid.coord_iter()))
+    )
+    print("data-- ", data)
+
+    def update_data():
+        data.value = list(get_agent_data_from_coord_iter(model.grid.coord_iter()))
+
+    def click_handler(datum):
+        if datum is None:
+            return
+        on_click(model, datum["x"], datum["y"])
+        update_data()
+
+    default_tooltip = [f"{key}:N" for key in data.value[0]]
+    chart = (
+        alt.Chart(alt.Data(values=data.value))
+        .mark_rect()
+        .encode(
+            x=alt.X("x:N", scale=alt.Scale(domain=list(range(model.grid.width)))),
+            y=alt.Y(
+                "y:N",
+                scale=alt.Scale(domain=list(range(model.grid.height - 1, -1, -1))),
+            ),
+            color=color,
+            tooltip=default_tooltip,
+        )
+        .properties(width=600, height=600)
+    )
+    return solara.FigureAltair(chart, on_click=click_handler)
diff --git a/mesa/experimental/jupyter_viz.py b/mesa/experimental/jupyter_viz.py
index c34a3d9fef3..201d9143d6e 100644
--- a/mesa/experimental/jupyter_viz.py
+++ b/mesa/experimental/jupyter_viz.py
@@ -11,6 +11,7 @@
 from solara.alias import rv
 
 import mesa
+from mesa.experimental.altair_grid import create_grid
 
 # Avoid interactive backend
 plt.switch_backend("agg")
@@ -75,6 +76,12 @@ def ColorCard(color, layout_type):
                     SpaceMatplotlib(
                         model, agent_portrayal, dependencies=[current_step.value]
                     )
+                elif space_drawer == "altair":
+                    # draw with the altair implementation
+                    SpaceAltair(
+                        model, agent_portrayal, dependencies=[current_step.value]
+                    )
+
                 elif space_drawer:
                     # if specified, draw agent space with an alternate renderer
                     space_drawer(model, agent_portrayal)
@@ -109,6 +116,9 @@ def render_in_jupyter():
                 SpaceMatplotlib(
                     model, agent_portrayal, dependencies=[current_step.value]
                 )
+            elif space_drawer == "altair":
+                # draw with the default implementation
+                SpaceAltair(model, agent_portrayal, dependencies=[current_step.value])
             elif space_drawer:
                 # if specified, draw agent space with an alternate renderer
                 space_drawer(model, agent_portrayal)
@@ -334,6 +344,12 @@ def SpaceMatplotlib(model, agent_portrayal, dependencies: Optional[List[any]] =
     solara.FigureMatplotlib(space_fig, format="png", dependencies=dependencies)
 
 
+@solara.component
+def SpaceAltair(model, agent_portrayal, dependencies: Optional[List[any]] = None):
+    grid = create_grid(color="wealth")
+    grid(model)
+
+
 def _draw_grid(space, space_ax, agent_portrayal):
     def portray(g):
         x = []

From 6236e64f92ff480717cfe96ce0a54b83bec3201b Mon Sep 17 00:00:00 2001
From: Ankit Kumar <ankit.kumar@tum.de>
Date: Fri, 15 Dec 2023 11:19:33 +0100
Subject: [PATCH 02/11] add statistics tab

---
 mesa/experimental/altair_grid.py | 2 --
 mesa/experimental/jupyter_viz.py | 8 +++++++-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/mesa/experimental/altair_grid.py b/mesa/experimental/altair_grid.py
index 851aae0120c..f753d7a1da3 100644
--- a/mesa/experimental/altair_grid.py
+++ b/mesa/experimental/altair_grid.py
@@ -10,7 +10,6 @@
 def get_agent_data_from_coord_iter(data):
     for agent, (x, y) in data:
         if agent:
-            print(agent[0], x, y)
             agent_data = json.loads(
                 json.dumps(agent[0].__dict__, skipkeys=True, default=str)
             )
@@ -40,7 +39,6 @@ def Grid(model, color=None, on_click=None):
     data = solara.reactive(
         list(get_agent_data_from_coord_iter(model.grid.coord_iter()))
     )
-    print("data-- ", data)
 
     def update_data():
         data.value = list(get_agent_data_from_coord_iter(model.grid.coord_iter()))
diff --git a/mesa/experimental/jupyter_viz.py b/mesa/experimental/jupyter_viz.py
index 201d9143d6e..8396796caae 100644
--- a/mesa/experimental/jupyter_viz.py
+++ b/mesa/experimental/jupyter_viz.py
@@ -149,6 +149,12 @@ def render_in_browser():
                 ModelController(model, play_interval, current_step, reset_counter)
             with solara.Card("Progress", margin=1, elevation=2):
                 solara.Markdown(md_text=f"####Step - {current_step}")
+            with solara.Card("Analytics", margin=1, elevation=2):
+                df = model.datacollector.get_model_vars_dataframe()
+                for col in list(df.columns):
+                    solara.Markdown(
+                        md_text=f"####Avg. {col} - {df.loc[:, f'{col}'].mean()}"
+                    )
 
         items = [
             ColorCard(color="white", layout_type=layout_types[i])
@@ -440,7 +446,7 @@ def get_initial_grid_layout(layout_types):
     grid_lay = []
     y_coord = 0
     for ii in range(len(layout_types)):
-        template_layout = {"h": 10, "i": 0, "moved": False, "w": 6, "y": 0, "x": 0}
+        template_layout = {"h": 20, "i": 0, "moved": False, "w": 6, "y": 0, "x": 0}
         if ii == 0:
             grid_lay.append(template_layout)
         else:

From 7707d5b0ab3c41d85a8d4249ec6e5a79cbfd0266 Mon Sep 17 00:00:00 2001
From: Ankit Kumar <ankit.kumar@tum.de>
Date: Fri, 15 Dec 2023 12:03:41 +0100
Subject: [PATCH 03/11] add altair to depencency

---
 pyproject.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyproject.toml b/pyproject.toml
index ca20d2d88ed..8fa9504405a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -47,6 +47,7 @@ dependencies = [
   "pandas",
   "solara",
   "tqdm",
+  "altair"
 ]
 dynamic = ["version"]
 

From 300e48cc17f9ebfecba1e6f397f7f1aa3becfdeb Mon Sep 17 00:00:00 2001
From: Ankit Kumar <ankit.kumar@tum.de>
Date: Tue, 26 Dec 2023 08:52:31 +0100
Subject: [PATCH 04/11] update unittest

---
 mesa/experimental/jupyter_viz.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/mesa/experimental/jupyter_viz.py b/mesa/experimental/jupyter_viz.py
index e4100233662..3617890fda0 100644
--- a/mesa/experimental/jupyter_viz.py
+++ b/mesa/experimental/jupyter_viz.py
@@ -133,7 +133,7 @@ def render_in_jupyter():
                 else:
                     make_plot(model, measure)
 
-    def render_in_browser():
+    def render_in_browser(statistics=False):
         # if space drawer is disabled, do not include it
         layout_types = [{"Space": "default"}] if space_drawer else []
 
@@ -150,11 +150,12 @@ def render_in_browser():
             with solara.Card("Progress", margin=1, elevation=2):
                 solara.Markdown(md_text=f"####Step - {current_step}")
             with solara.Card("Analytics", margin=1, elevation=2):
-                df = model.datacollector.get_model_vars_dataframe()
-                for col in list(df.columns):
-                    solara.Markdown(
-                        md_text=f"####Avg. {col} - {df.loc[:, f'{col}'].mean()}"
-                    )
+                if statistics:
+                    df = model.datacollector.get_model_vars_dataframe()
+                    for col in list(df.columns):
+                        solara.Markdown(
+                            md_text=f"####Avg. {col} - {df.loc[:, f'{col}'].mean()}"
+                        )
 
         items = [
             ColorCard(color="white", layout_type=layout_types[i])

From 9b1ba7f5bb1b27c9242c143f1052ee1ec3a0bec5 Mon Sep 17 00:00:00 2001
From: Ankit Kumar <ankit.kumar@tum.de>
Date: Tue, 26 Dec 2023 13:21:35 +0100
Subject: [PATCH 05/11] add altrair to unittest

---
 tests/test_jupyter_viz.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/tests/test_jupyter_viz.py b/tests/test_jupyter_viz.py
index dd125ffcadf..ca807c5f261 100644
--- a/tests/test_jupyter_viz.py
+++ b/tests/test_jupyter_viz.py
@@ -100,6 +100,14 @@ def test_call_space_drawer(self, mock_space_matplotlib):
                 agent_portrayal=agent_portrayal,
             )
         )
+        solara.render(
+            JupyterViz(
+                model_class=mock_model_class,
+                model_params={},
+                agent_portrayal=agent_portrayal,
+                space_drawer="altair",
+            )
+        )
         # should call default method with class instance and agent portrayal
         mock_space_matplotlib.assert_called_with(
             mock_model_class.return_value, agent_portrayal, dependencies=dependencies
@@ -132,3 +140,8 @@ def test_call_space_drawer(self, mock_space_matplotlib):
         altspace_drawer.assert_called_with(
             mock_model_class.return_value, agent_portrayal
         )
+
+
+if __name__ == "__main__":
+    tjv = TestJupyterViz()
+    tjv.test_call_space_drawer()

From 9dbb49449541b138e3db29602d2a4b7710c3e374 Mon Sep 17 00:00:00 2001
From: Ankit Kumar <ankit.kumar@tum.de>
Date: Tue, 26 Dec 2023 13:36:16 +0100
Subject: [PATCH 06/11] add fix for python 3.9/8

---
 mesa/experimental/altair_grid.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mesa/experimental/altair_grid.py b/mesa/experimental/altair_grid.py
index f753d7a1da3..565854a720d 100644
--- a/mesa/experimental/altair_grid.py
+++ b/mesa/experimental/altair_grid.py
@@ -1,5 +1,5 @@
 import json
-from typing import Callable
+from typing import Callable, Optional
 
 import altair as alt
 import solara
@@ -21,8 +21,8 @@ def get_agent_data_from_coord_iter(data):
 
 
 def create_grid(
-    color: str | None = None,
-    on_click: Callable[[mesa.Model, mesa.space.Coordinate], None] | None = None,
+    color: Optional[str] = None,
+    on_click: Optional[Callable[[mesa.Model, mesa.space.Coordinate], None]] = None,
 ) -> Callable[[mesa.Model], solara.component]:
     return lambda model: Grid(model, color, on_click)
 

From 9e0c00e433c4fe79140a5eb9b39f0ae1d7ed1135 Mon Sep 17 00:00:00 2001
From: Ankit Kumar <ankit.kumar@tum.de>
Date: Wed, 17 Jan 2024 23:24:55 +0100
Subject: [PATCH 07/11] updates

---
 mesa/experimental/altair_grid.py | 16 +++++++++-------
 mesa/experimental/jupyter_viz.py |  2 +-
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/mesa/experimental/altair_grid.py b/mesa/experimental/altair_grid.py
index 565854a720d..28e145ccb78 100644
--- a/mesa/experimental/altair_grid.py
+++ b/mesa/experimental/altair_grid.py
@@ -10,11 +10,8 @@
 def get_agent_data_from_coord_iter(data):
     for agent, (x, y) in data:
         if agent:
-            agent_data = json.loads(
-                json.dumps(agent[0].__dict__, skipkeys=True, default=str)
-            )
-            agent_data["x"] = x
-            agent_data["y"] = y
+            agent_data = agent[0].__dict__.copy()
+            agent_data.update({"x": x, "y": y})
             agent_data.pop("model", None)
             agent_data.pop("pos", None)
             yield agent_data
@@ -24,7 +21,10 @@ def create_grid(
     color: Optional[str] = None,
     on_click: Optional[Callable[[mesa.Model, mesa.space.Coordinate], None]] = None,
 ) -> Callable[[mesa.Model], solara.component]:
-    return lambda model: Grid(model, color, on_click)
+    def create_grid_function(model: mesa.Model) -> solara.component:
+        return solara.component.Grid(model, color, on_click)
+
+    return create_grid_function
 
 
 def Grid(model, color=None, on_click=None):
@@ -49,7 +49,9 @@ def click_handler(datum):
         on_click(model, datum["x"], datum["y"])
         update_data()
 
-    default_tooltip = [f"{key}:N" for key in data.value[0]]
+    default_tooltip = [
+        f"{key}:N" for key in data.value[0]
+    ]  # add all agent attributes to tooltip
     chart = (
         alt.Chart(alt.Data(values=data.value))
         .mark_rect()
diff --git a/mesa/experimental/jupyter_viz.py b/mesa/experimental/jupyter_viz.py
index 3617890fda0..700ff7718ad 100644
--- a/mesa/experimental/jupyter_viz.py
+++ b/mesa/experimental/jupyter_viz.py
@@ -24,7 +24,7 @@ def JupyterViz(
     measures=None,
     name="Mesa Model",
     agent_portrayal=None,
-    space_drawer="default",
+    space_drawer="altair",
     play_interval=150,
 ):
     """Initialize a component to visualize a model.

From a0ce9d98529542643b6260c74086c7e4fc533f39 Mon Sep 17 00:00:00 2001
From: Ankit Kumar <ankit.kumar@tum.de>
Date: Fri, 26 Jan 2024 09:45:11 +0100
Subject: [PATCH 08/11] add docstring

---
 mesa/experimental/altair_grid.py | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/mesa/experimental/altair_grid.py b/mesa/experimental/altair_grid.py
index 28e145ccb78..20274ea44ec 100644
--- a/mesa/experimental/altair_grid.py
+++ b/mesa/experimental/altair_grid.py
@@ -1,4 +1,3 @@
-import json
 from typing import Callable, Optional
 
 import altair as alt
@@ -8,6 +7,15 @@
 
 
 def get_agent_data_from_coord_iter(data):
+    """
+    Extracts agent data from a sequence of tuples containing agent objects and their coordinates.
+
+    Parameters:
+    - data (iterable): A sequence of tuples where each tuple contains an agent object and its coordinates.
+
+    Yields:
+    - dict: A dictionary containing agent data with updated coordinates. The dictionary excludes 'model' and 'pos' attributes.
+    """
     for agent, (x, y) in data:
         if agent:
             agent_data = agent[0].__dict__.copy()
@@ -21,6 +29,18 @@ def create_grid(
     color: Optional[str] = None,
     on_click: Optional[Callable[[mesa.Model, mesa.space.Coordinate], None]] = None,
 ) -> Callable[[mesa.Model], solara.component]:
+    """
+    Factory function for creating a grid component for a Mesa model.
+
+    Parameters:
+    - color (Optional[str]): Color of the grid lines. Defaults to None.
+    - on_click (Optional[Callable[[mesa.Model, mesa.space.Coordinate], None]]):
+    Function to be called when a grid cell is clicked. Defaults to None.
+
+    Returns:
+    - Callable[[mesa.Model], solara.component]: A function that creates a grid component for the given model.
+    """
+
     def create_grid_function(model: mesa.Model) -> solara.component:
         return solara.component.Grid(model, color, on_click)
 
@@ -28,6 +48,16 @@ def create_grid_function(model: mesa.Model) -> solara.component:
 
 
 def Grid(model, color=None, on_click=None):
+    """
+    Handles click events on grid cells.
+
+    Parameters:
+    - datum (dict): Data associated with the clicked cell.
+
+    Notes:
+    - Invokes the provided `on_click` function with the model and cell coordinates.
+    - Updates the data displayed on the grid.
+    """
     if color is None:
         color = "unique_id:N"
 

From 5182198cb75bc33fcaa0c5e4168ae666cbc9f623 Mon Sep 17 00:00:00 2001
From: Ankit Kumar <ankit.kumar@tum.de>
Date: Fri, 26 Jan 2024 10:37:47 +0100
Subject: [PATCH 09/11] fix grid

---
 mesa/experimental/altair_grid.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mesa/experimental/altair_grid.py b/mesa/experimental/altair_grid.py
index 20274ea44ec..eaa423aab37 100644
--- a/mesa/experimental/altair_grid.py
+++ b/mesa/experimental/altair_grid.py
@@ -42,7 +42,7 @@ def create_grid(
     """
 
     def create_grid_function(model: mesa.Model) -> solara.component:
-        return solara.component.Grid(model, color, on_click)
+        return Grid(model, color, on_click)
 
     return create_grid_function
 

From 3fe18384feed5df5fae007c47a02d90c3eb0a164 Mon Sep 17 00:00:00 2001
From: Ankit Kumar <ankit.kumar@tum.de>
Date: Thu, 8 Feb 2024 21:48:48 +0100
Subject: [PATCH 10/11] fix tests

---
 mesa/experimental/jupyter_viz.py |  2 +-
 tests/test_jupyter_viz.py        | 12 ++++--------
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/mesa/experimental/jupyter_viz.py b/mesa/experimental/jupyter_viz.py
index fe4ac1cd1c4..3ea858c70e2 100644
--- a/mesa/experimental/jupyter_viz.py
+++ b/mesa/experimental/jupyter_viz.py
@@ -21,7 +21,7 @@ def JupyterViz(
     measures=None,
     name=None,
     agent_portrayal=None,
-    space_drawer="altair",
+    space_drawer="default",
     play_interval=150,
 ):
     """Initialize a component to visualize a model.
diff --git a/tests/test_jupyter_viz.py b/tests/test_jupyter_viz.py
index 490ca6e64ea..50eeef734e8 100644
--- a/tests/test_jupyter_viz.py
+++ b/tests/test_jupyter_viz.py
@@ -92,6 +92,7 @@ def test_call_space_drawer(self, mock_space_matplotlib):
         }
         current_step = 0
         dependencies = [current_step]
+
         # initialize with space drawer unspecified (use default)
         # component must be rendered for code to run
         solara.render(
@@ -99,16 +100,11 @@ def test_call_space_drawer(self, mock_space_matplotlib):
                 model_class=mock_model_class,
                 model_params={},
                 agent_portrayal=agent_portrayal,
+                space_drawer="default",
             )
         )
-        solara.render(
-            JupyterViz(
-                model_class=mock_model_class,
-                model_params={},
-                agent_portrayal=agent_portrayal,
-                space_drawer="altair",
-            )
-        )
+
+      
         # should call default method with class instance and agent portrayal
         mock_space_matplotlib.assert_called_with(
             mock_model_class.return_value, agent_portrayal, dependencies=dependencies

From 487dca329e6363fd49c3722ff244638931bff50b Mon Sep 17 00:00:00 2001
From: Ankit Kumar <ankit.kumar@tum.de>
Date: Thu, 8 Feb 2024 21:51:49 +0100
Subject: [PATCH 11/11] apply ruff

---
 tests/test_jupyter_viz.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/test_jupyter_viz.py b/tests/test_jupyter_viz.py
index 67d51461fdc..ae0b324c0fa 100644
--- a/tests/test_jupyter_viz.py
+++ b/tests/test_jupyter_viz.py
@@ -104,7 +104,6 @@ def test_call_space_drawer(self, mock_space_matplotlib):
             )
         )
 
-      
         # should call default method with class instance and agent portrayal
         mock_space_matplotlib.assert_called_with(
             mock_model_class.return_value, agent_portrayal, dependencies=dependencies