Add sankey chart

mckinsey · Jul 15, 2024 · 1da246d · 1da246d
1 parent 0a7ba67
commit 1da246d
Show file tree

Hide file tree

Showing 3 changed files with 189 additions and 26 deletions.
diff --git a/vizro-core/examples/_dev/app.py b/vizro-core/examples/_dev/app.py
@@ -1,36 +1,63 @@
 """Dev app to try things out."""
 
-import json
-from urllib.request import urlopen
+from typing import List
 
 import pandas as pd
+import plotly.graph_objects as go
 import vizro.models as vm
-import vizro.plotly.express as px
 from vizro import Vizro
+from vizro.models.types import capture
 
-with urlopen(
-    "https://raw.githubusercontent.com/plotly/datasets/master/geojson-counties-fips.json"
-) as response:
-    counties = json.load(response)
-
-fips_unemp = pd.read_csv(
-    "https://raw.githubusercontent.com/plotly/datasets/master/fips-unemp-16.csv",
-    dtype={"fips": str},
+sankey_data = pd.DataFrame(
+    {
+        "Origin": [0, 1, 2, 1, 2, 4, 0],  # indices inside labels
+        "Destination": [1, 2, 3, 4, 5, 5, 6],  # indices inside labels
+        "Value": [10, 4, 8, 6, 4, 8, 8],
+    }
 )
 
 
+@capture("graph")
+def sankey(
+    data_frame: pd.DataFrame,
+    source: str,
+    target: str,
+    value: str,
+    labels: List[str],
+):
+    fig = go.Figure(
+        data=[
+            go.Sankey(
+                node=dict(
+                    pad=16,
+                    thickness=16,
+                    label=labels,
+                ),
+                link=dict(
+                    source=data_frame[source],
+                    target=data_frame[target],
+                    value=data_frame[value],
+                    label=labels,
+                    color="rgba(205, 209, 228, 0.4)",
+                ),
+            )
+        ]
+    )
+    fig.update_layout(barmode="relative")
+    return fig
+
+
 page = vm.Page(
-    title="Choropleth",
+    title="Sankey",
     components=[
         vm.Graph(
-            figure=px.choropleth(
-                fips_unemp,
-                geojson=counties,
-                locations="fips",
-                color="unemp",
-                range_color=(0, 12),
-                scope="usa",
-            )
+            figure=sankey(
+                data_frame=sankey_data,
+                labels=["A1", "A2", "B1", "B2", "C1", "C2", "D1"],
+                source="Origin",
+                target="Destination",
+                value="Value",
+            ),
         ),
     ],
 )

diff --git a/vizro-core/examples/chart-gallery/app.py b/vizro-core/examples/chart-gallery/app.py
@@ -46,6 +46,7 @@
         butterfly_page,
         distribution_butterfly,
         choropleth,
+        sankey,
     ],
     navigation=vm.Navigation(
         nav_selector=vm.NavBar(
@@ -88,7 +89,7 @@
                 ),
                 vm.NavLink(
                     label="Flow",
-                    pages={"Flow": ["Line"]},  # TODO: Replace with Sankey
+                    pages={"Flow": ["Sankey"]},
                     icon="Stacked Line Chart",
                 ),
                 vm.NavLink(

diff --git a/vizro-core/examples/chart-gallery/utils/_pages.py b/vizro-core/examples/chart-gallery/utils/_pages.py
@@ -35,6 +35,14 @@
     dtype={"fips": str},
 )
 
+sankey_data = pd.DataFrame(
+    {
+        "Origin": [0, 1, 2, 1, 2, 4, 0],  # indices inside labels
+        "Destination": [1, 2, 3, 4, 5, 5, 6],  # indices inside labels
+        "Value": [10, 4, 8, 6, 4, 8, 8],
+    }
+)
+
 
 vm.Page.add_type("components", CodeClipboard)
 vm.Page.add_type("components", FlexContainer)
@@ -64,6 +72,30 @@ def butterfly(data_frame: pd.DataFrame, x1: str, x2: str, y: str):
     return fig
 
 
+@capture("graph")
+def sankey(data_frame: pd.DataFrame, source: str, target: str, value: str):
+    fig = go.Figure(
+        data=[
+            go.Sankey(
+                node=dict(
+                    pad=16,
+                    thickness=16,
+                    #     label=labels,
+                ),
+                link=dict(
+                    source=data_frame[source],
+                    target=data_frame[target],
+                    value=data_frame[value],
+                    #    label=labels,
+                    color="rgba(205, 209, 228, 0.4)",
+                ),
+            )
+        ]
+    )
+    fig.update_layout(barmode="relative")
+    return fig
+
+
 # All functions ending with `_factory` are there to reuse the existing content of a page. Given the restriction that
 # one page can only be mapped to one navigation group, we have to create a new page with a new ID.
 def line_factory(id: str, title: str):
@@ -764,23 +796,23 @@ def butterfly(data_frame: pd.DataFrame, x1: str, x2: str, y: str):
                 ```python
                 import json
                 from urllib.request import urlopen
-                
+
                 import pandas as pd
                 import vizro.models as vm
                 import vizro.plotly.express as px
                 from vizro import Vizro
-                
+
                 with urlopen(
                     "https://raw.githubusercontent.com/plotly/datasets/master/geojson-counties-fips.json"
                 ) as response:
                     counties = json.load(response)
-                
+
                 fips_unemp = pd.read_csv(
                     "https://raw.githubusercontent.com/plotly/datasets/master/fips-unemp-16.csv",
                     dtype={"fips": str},
                 )
-                
-                
+
+
                 page = vm.Page(
                     title="Choropleth",
                     components=[
@@ -804,3 +836,106 @@ def butterfly(data_frame: pd.DataFrame, x1: str, x2: str, y: str):
         ),
     ],
 )
+
+
+sankey = vm.Page(
+    title="Sankey",
+    layout=vm.Layout(grid=PAGE_GRID),
+    components=[
+        vm.Card(
+            text="""
+            #### What is a sankey chart?
+
+            A Sankey chart is a type of flow diagram that illustrates how resources or values move between different
+            stages or entities. The width of the arrows in the chart is proportional to the quantity of the flow,
+            making it easy to see where the largest movements occur.
+
+            &nbsp;
+
+            #### When to use it?
+
+            Use a Sankey chart when you want to visualize the flow of resources, energy, money, or other values from
+            one point to another. It is particularly useful for showing distributions and transfers within a system,
+            such as energy usage, cost breakdowns, or material flows.
+
+            Be mindful that Sankey charts can become cluttered if there are too many nodes or flows.
+            To maintain clarity, focus on highlighting the most significant flows and keep the chart as simple as
+            possible.
+        """
+        ),
+        vm.Graph(
+            figure=sankey(
+                data_frame=sankey_data,
+                labels=["A1", "A2", "B1", "B2", "C1", "C2", "D1"],
+                source="Origin",
+                target="Destination",
+                value="Value",
+            ),
+        ),
+        CodeClipboard(
+            text="""
+                ```python
+                import pandas as pd
+                import plotly.graph_objects as go
+                import vizro.models as vm
+                from vizro import Vizro
+                from vizro.models.types import capture
+                from typing import List
+
+                sankey_data = pd.DataFrame(
+                    {
+                        "Origin": [0, 1, 2, 1, 2, 4, 0],  # indices inside labels
+                        "Destination": [1, 2, 3, 4, 5, 5, 6],  # indices inside labels
+                        "Value": [10, 4, 8, 6, 4, 8, 8],
+                    }
+                )
+
+
+                @capture("graph")
+                def sankey(
+                    data_frame: pd.DataFrame,
+                    source: str,
+                    target: str,
+                    value: str,
+                    labels: List[str],
+                ):
+                    fig = go.Figure(
+                        data=[
+                            go.Sankey(
+                                node=dict(pad=16, thickness=16, label=labels,),
+                                link=dict(
+                                    source=data_frame[source],
+                                    target=data_frame[target],
+                                    value=data_frame[value],
+                                    label=labels,
+                                    color="rgba(205, 209, 228, 0.4)",
+                                ),
+                            )
+                        ]
+                    )
+                    fig.update_layout(barmode="relative")
+                    return fig
+
+
+                page = vm.Page(
+                    title="Sankey",
+                    components=[
+                        vm.Graph(
+                            figure=sankey(
+                                data_frame=sankey_data,
+                                labels=["A1", "A2", "B1", "B2", "C1", "C2", "D1"],
+                                source="Origin",
+                                target="Destination",
+                                value="Value",
+                            ),
+                        ),
+                    ],
+                )
+
+                dashboard = vm.Dashboard(pages=[page])
+                Vizro().build(dashboard).run()
+                ```
+                """
+        ),
+    ],
+)