From e2d9a8bafec4831a3264894ad9c9fae7758973a8 Mon Sep 17 00:00:00 2001
From: rht <rhtbot@protonmail.com>
Date: Sat, 6 Jan 2024 19:38:47 -0500
Subject: [PATCH] jupyter_viz: Implement multiline plot

Fixes #1821.

This provides a declarative interface to plotting multiple lines in 1
chart.

Example usage for virus on network
```python
page = JupyterViz(
    VirusOnNetwork,
    model_params,
    measures=[
        # This is subject to change, as it could have been
        # {"Infected": {"color": "tab:red"}}, but let's keep it simple.
        {"Infected": "tab:red", "Susceptible": "tab:green", "Resistant": "tab:gray"},
        # Alternatively, if you don't want to manually specify colors:
        # ["Infected", "Susceptible", "Resistant"],
        make_text(get_resistant_susceptible_ratio),
    ],
    name="Virus Model",
    agent_portrayal=agent_portrayal,
```
---
 mesa/experimental/jupyter_viz.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/mesa/experimental/jupyter_viz.py b/mesa/experimental/jupyter_viz.py
index 63165eb6ca4..f4ac1271054 100644
--- a/mesa/experimental/jupyter_viz.py
+++ b/mesa/experimental/jupyter_viz.py
@@ -406,8 +406,17 @@ def make_plot(model, measure):
     fig = Figure()
     ax = fig.subplots()
     df = model.datacollector.get_model_vars_dataframe()
-    ax.plot(df.loc[:, measure])
-    ax.set_ylabel(measure)
+    if isinstance(measure, str):
+        ax.plot(df.loc[:, measure])
+        ax.set_ylabel(measure)
+    elif isinstance(measure, dict):
+        for m, color in measure.items():
+            ax.plot(df.loc[:, m], label=m, color=color)
+        fig.legend()
+    elif isinstance(measure, (list, tuple)):
+        for m in measure:
+            ax.plot(df.loc[:, m], label=m)
+        fig.legend()
     # Set integer x axis
     ax.xaxis.set_major_locator(MaxNLocator(integer=True))
     solara.FigureMatplotlib(fig)