Sunburst/treemap path

packages/python/plotly/plotly/express/_chart_types.py

@@ -1252,6 +1252,7 @@ def sunburst(
     names=None,
     values=None,
     parents=None,
+    path=None,
     ids=None,
     color=None,
     color_continuous_scale=None,
@@ -1278,6 +1279,13 @@ def sunburst(
         layout_patch = {"sunburstcolorway": color_discrete_sequence}
     else:
         layout_patch = {}
+    if path is not None and (ids is not None or parents is not None):
+        raise ValueError(
+            "Either `path` should be provided, or `ids` and `parents`."
+            "These parameters are mutually exclusive and cannot be passed together."
+        )
+    if path is not None and branchvalues is None:
+        branchvalues = "total"
     return make_figure(
         args=locals(),
         constructor=go.Sunburst,
@@ -1295,6 +1303,7 @@ def treemap(
     values=None,
     parents=None,
     ids=None,
+    path=None,
     color=None,
     color_continuous_scale=None,
     range_color=None,
@@ -1320,6 +1329,12 @@ def treemap(
         layout_patch = {"treemapcolorway": color_discrete_sequence}
     else:
         layout_patch = {}
+    if path is not None and (ids is not None or parents is not None):
+        raise ValueError(
+            "Either `path` should be provided, or `ids` and `parents`."
+            "These parameters are mutually exclusive and cannot be passed together."
+        )
+
     return make_figure(
         args=locals(),
         constructor=go.Treemap,

packages/python/plotly/plotly/express/_core.py

@@ -890,6 +890,9 @@ def build_dataframe(args, attrables, array_attrables):
         else:
             df_output[df_input.columns] = df_input[df_input.columns]
 
+    # This should be improved + tested - HACK
+    if "path" in args and args["path"] is not None:
+        df_output[args["path"]] = df_input[args["path"]]
     # Loop over possible arguments
     for field_name in attrables:
         # Massaging variables
@@ -1007,6 +1010,66 @@ def build_dataframe(args, attrables, array_attrables):
     return args
 
 
+def process_dataframe_hierarchy(args):
+    """
+    Build dataframe for sunburst or treemap when the path argument is provided.
+    """
+    df = args["data_frame"]
+    path = args["path"]
+    # Other columns (for color, hover_data, custom_data etc.)
+    cols = list(set(df.columns).difference(path))
+    df_all_trees = pd.DataFrame(columns=["labels", "parent", "id"] + cols)
+    # Set column type here (useful for continuous vs discrete colorscale)
+    for col in cols:
+        df_all_trees[col] = df_all_trees[col].astype(df[col].dtype)
+    for i, level in enumerate(path):
+        df_tree = pd.DataFrame(columns=df_all_trees.columns)
+        dfg = df.groupby(path[i:]).sum(numerical_only=True)
+        dfg = dfg.reset_index()
+        df_tree["labels"] = dfg[level].copy().astype(str)
+        df_tree["parent"] = ""
+        df_tree["id"] = dfg[level].copy().astype(str)
+        if i < len(path) - 1:
+            j = i + 1
+            while j < len(path):
+                df_tree["parent"] += dfg[path[j]].copy().astype(str)
+                df_tree["id"] += dfg[path[j]].copy().astype(str)
+                j += 1
+        else:
+            df_tree["parent"] = "total"
+
+        if i == 0 and cols:
+            df_tree[cols] = dfg[cols]
+        elif cols:
+            for col in cols:
+                df_tree[col] = "n/a"
+        if args["values"]:
+            df_tree[args["values"]] = dfg[args["values"]]
+        df_all_trees = df_all_trees.append(df_tree, ignore_index=True)
+
+    # Root node
+    total_dict = {
+        "labels": "total",
+        "id": "total",
+        "parent": "",
+    }
+    for col in cols:
+        if not col == args["values"]:
+            total_dict[col] = "n/a"
+        if col == args["values"]:
+            total_dict[col] = df[col].sum()
+    total = pd.Series(total_dict)
+
+    df_all_trees = df_all_trees.append(total, ignore_index=True)
+    # Now modify arguments
+    args["data_frame"] = df_all_trees
+    args["path"] = None
+    args["ids"] = "id"
+    args["names"] = "labels"
+    args["parents"] = "parent"
+    return args
+
+
 def infer_config(args, constructor, trace_patch):
     # Declare all supported attributes, across all plot types
     attrables = (
@@ -1017,7 +1080,7 @@ def infer_config(args, constructor, trace_patch):
         + ["error_y", "error_y_minus", "error_z", "error_z_minus"]
         + ["lat", "lon", "locations", "animation_group"]
     )
-    array_attrables = ["dimensions", "custom_data", "hover_data"]
+    array_attrables = ["dimensions", "custom_data", "hover_data", "path"]
     group_attrables = ["animation_frame", "facet_row", "facet_col", "line_group"]
     all_attrables = attrables + group_attrables + ["color"]
     group_attrs = ["symbol", "line_dash"]
@@ -1026,6 +1089,8 @@ def infer_config(args, constructor, trace_patch):
             all_attrables += [group_attr]
 
     args = build_dataframe(args, all_attrables, array_attrables)
+    if constructor in [go.Treemap, go.Sunburst] and args["path"] is not None:
+        args = process_dataframe_hierarchy(args)
 
     attrs = [k for k in attrables if k in args]
     grouped_attrs = []

packages/python/plotly/plotly/express/_doc.py

@@ -82,6 +82,7 @@
         colref_desc,
         "Values from this column or array_like are used to set ids of sectors",
     ],
+    path=[colref_type, colref_desc],
     lat=[
         colref_type,
         colref_desc,

packages/python/plotly/plotly/tests/test_core/test_px/test_px_functions.py

@@ -2,6 +2,7 @@
 import plotly.graph_objects as go
 from numpy.testing import assert_array_equal
 import numpy as np
+import pandas as pd
 
 
 def _compare_figures(go_trace, px_fig):
@@ -111,6 +112,42 @@ def test_sunburst_treemap_colorscales():
         assert list(fig.layout[colorway]) == color_seq
 
 
+def test_sunburst_treemap_with_path():
+    vendors = ["A", "B", "C", "D", "E", "F", "G", "H"]
+    sectors = [
+        "Tech",
+        "Tech",
+        "Finance",
+        "Finance",
+        "Tech",
+        "Tech",
+        "Finance",
+        "Finance",
+    ]
+    regions = ["North", "North", "North", "North", "South", "South", "South", "South"]
+    values = [1, 3, 2, 4, 2, 2, 1, 4]
+    df = pd.DataFrame(
+        dict(vendors=vendors, sectors=sectors, regions=regions, values=values)
+    )
+    # No values
+    fig = px.sunburst(df, path=["vendors", "sectors", "regions"])
+    assert fig.data[0].branchvalues == "total"
+    # Values passed
+    fig = px.sunburst(df, path=["vendors", "sectors", "regions"], values="values")
+    assert fig.data[0].branchvalues == "total"
+    assert fig.data[0].values[-1] == np.sum(values)
+    # Values passed
+    fig = px.sunburst(df, path=["vendors", "sectors", "regions"], values="values")
+    assert fig.data[0].branchvalues == "total"
+    assert fig.data[0].values[-1] == np.sum(values)
+    # Continuous colorscale
+    fig = px.sunburst(
+        df, path=["vendors", "sectors", "regions"], values="values", color="values"
+    )
+    assert "coloraxis" in fig.data[0].marker
+    assert np.all(np.array(fig.data[0].marker.colors) == np.array(fig.data[0].values))
+
+
 def test_pie_funnelarea_colorscale():
     labels = ["A", "B", "C", "D"]
     values = [3, 2, 1, 4]