python - 如何使用 Plotly 制作简单的多级桑基图？

Question

我有一个像这样的 DataFrame，我试图用桑基图来描述它:

import pandas as pd

pd.DataFrame({
    'animal': ['dog', 'cat', 'cat', 'dog', 'cat'],
    'sex': ['male', 'female', 'female', 'male', 'male'],
    'status': ['wild', 'domesticated', 'domesticated', 'wild', 'domesticated'],
    'count': [8, 10, 11, 14, 6]
})

    animal  sex     status          count
0   dog     male    wild            8
1   cat     female  domesticated    10
2   cat     female  domesticated    11
3   dog     male    wild            14
4   cat     male    domesticated    6

我正在尝试按照 documentation 中的步骤操作但我无法让它发挥作用 - 我不明白什么分支在哪里。示例代码如下:

import plotly.graph_objects as go

fig = go.Figure(data=[go.Sankey(
    node = dict(
      pad = 15,
      thickness = 20,
      line = dict(color = "black", width = 0.5),
      label = ["A1", "A2", "B1", "B2", "C1", "C2"],
      color = "blue"
    ),
    link = dict(
      source = [0, 1, 0, 2, 3, 3], 
      target = [2, 3, 3, 4, 4, 5],
      value = [8, 4, 2, 8, 4, 2]
  ))])

fig.update_layout(title_text="Basic Sankey Diagram", font_size=10)
fig.show()

这就是我想要实现的目标:

Answer 1

您可以通过以下方式使用 Plotly 创建 Sankey 图:

import pandas as pd
import plotly.graph_objects as go

label_list = ['cat', 'dog', 'domesticated', 'female', 'male', 'wild']
# cat: 0, dog: 1, domesticated: 2, female: 3, male: 4, wild: 5
source = [0, 0, 1, 3, 4, 4]
target = [3, 4, 4, 2, 2, 5]
count = [21, 6, 22, 21, 6, 22]

fig = go.Figure(data=[go.Sankey(
    node = {"label": label_list},
    link = {"source": source, "target": target, "value": count}
    )])
fig.show()

它是如何工作的:列表 source、target 和 count 的长度均为 6，桑基图有 6 个箭头。 source和target的元素是label_list的索引。所以source的第一个元素是0，意思是“猫”。 target的第一个元素是3，意思是“女性”。 count 的第一个元素是 21。因此，图中的第一个箭头从 cat 到 Female，大小为 21。相应地，列表 source、target 和 count 的第二个元素定义第二个箭头，依此类推。

---

您可能想要创建一个更大的桑基图，如本例所示。手动定义源、目标和计数列表变得非常乏味。因此，这里有一段代码，可以根据您的格式的数据帧创建这些列表。

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'animal': ['dog', 'cat', 'cat', 'dog', 'cat'],
    'sex': ['male', 'female', 'female', 'male', 'male'],
    'status': ['wild', 'domesticated', 'domesticated', 'wild', 'domesticated'],
    'count': [8, 10, 11, 14, 6]
})

categories = ['animal', 'sex', 'status']

newDf = pd.DataFrame()
for i in range(len(categories)-1):
    tempDf = df[[categories[i],categories[i+1],'count']]
    tempDf.columns = ['source','target','count']
    newDf = pd.concat([newDf,tempDf])    
newDf = newDf.groupby(['source','target']).agg({'count':'sum'}).reset_index()

label_list = list(np.unique(df[categories].values))
source = newDf['source'].apply(lambda x: label_list.index(x))
target = newDf['target'].apply(lambda x: label_list.index(x))
count = newDf['count']