社区首页 >专栏 >在Python Matplotlib中制作瀑布图

在Python Matplotlib中制作瀑布图

fanjy

发布于 2022-06-04 02:11:53

2.9K00

代码可运行

文章被收录于专栏：完美Excel完美Excel

运行总次数：0

代码可运行

标签：Python，Matplotlib，瀑布图

我们将用Python制作瀑布图，特别是使用matplotlib库。瀑布图显示了运行总数以及增减，这对于属性分析来说是很好的选择。

Matplotlib没有像“waterfall_chart()”这样的神奇函数，使我们能够用一行代码就绘制瀑布图。然而，可以使用一点小小的技巧在Python中自定义自己的瀑布图。

1.创建标准的条形图。

2.创建另一个条形图并将其放在第一个条形图的顶部，然后将新条形图的颜色设置为与背景色相同的颜色，以隐藏第一个条形图的底部。

实际上，因为我们看不到第二组条形图，所以我们可以使用它们来“隐藏”另一组条形图。

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
df= pd.DataFrame({'category':['Sales','Service','Expenses','Taxes','Interest'],
                   'num':[100,10,-20,-30,60]})

代码运行结果如下图1所示。

图1

任务现在变成创建两个条形图，其中一个应该记录运行总数，另一个只是运行总数的变化，我们稍后就会看到。

可以使用cumsum()方法计算一个运行总数，然后将其下移1行。这两个新的列tot和tot1为我们提供了每个瀑布条的起点和终点。例如，在第2行Expenses（费用）中，起点是110，终点是90。

df['tot']= df['num'].cumsum()
df['tot1']= df['tot'].shift(1).fillna(0)

代码运行结果如下图2所示。

图2

由于起点和终点可以位于两个新列中的任意一列（取决于值的符号），因此我们可以再创建两列来捕获upper点和lower点：

lower= df[['tot','tot1']].min(axis=1)
upper= df[['tot','tot1']].max(axis=1)

我们使用upper点绘制第一组条形图。注意，这些条形的颜色与背景颜色不同。然后，我们使用lower点绘制第二组条形图，并将颜色设置为与背景颜色相同，默认情况下为白色。

fig,ax= plt.subplots()
ax.bar(x=df['category'],height=upper,)
ax.bar(x=df['category'],height=lower,color='white')

现在我们得到一个类似下面图3所示的图表。基本上，由于与背景颜色相同，高度为“lower点”的条形图是不可见的。

图3

现在，我们有了一个基本的瀑布图，再给它添加一些颜色。这里使用绿色表示增加，红色表示减少。

数据在num列中随时可用，让我们创建一个新的color列来存储每个类别的适当颜色。

df.loc[df['num']>= 0, 'color'] = 'green'
df.loc[df['num']< 0, 'color'] = 'red'

使用新颜色重新绘制条形图，如下所示：

ax.bar(x=df['category'],height=upper,color= df['color'])

运行结果如下图4所示。

图4

瀑布图显示了每个类别对总数的贡献，因此可在每个条形的中间添加标签信息。也可以添加“连接符”，将上一个条形的起点和终点连接到下一个条形。

下面将完整的瀑布图代码转换为一个方便的Python函数，以便以后可以重用它。该函数接受三个参数：包含数据的数据框架、要放置为x轴的数据列的名称以及要用作y轴的数据列的名称。

def waterfall(df, x, y):
    # 计算运行总数
    df['tot'] = df[y].cumsum()
    df['tot1']=df['tot'].shift(1).fillna(0)
    # 条形图的lower点和upper点
    lower = df[['tot','tot1']].min(axis=1)
    upper = df[['tot','tot1']].max(axis=1)
    # 为标签位置的中间点
    mid = (lower + upper)/2
    # 正数显示绿色, 负数显示红色
    df.loc[df[y] >= 0, 'color'] = 'green'
    df.loc[df[y] < 0, 'color'] = 'red'
    # 计算连接点
    connect= df['tot1'].repeat(3).shift(-1)
    connect[1::3] = np.nan
    fig,ax = plt.subplots()
    # 绘制具有颜色的第一个条形
    bars = ax.bar(x=df[x],height=upper, color=df['color'])
    # 绘制第二个条形 - 不可见
    plt.bar(x=df[x],height=lower,color='white')
    # 绘制连接线
    plt.plot(connect.index,connect.values, 'k')
    # 绘制条形标签
    for i, v in enumerate(upper):
        plt.text(i-.15, mid[i],f"{df[y][i]:,.0f}")
waterfall(df,'category','num')