在Python Matplotlib中制作瀑布图

excelperfect

标签：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')

 

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

图5

注：本文学习整理自pythoninoffice.com，供有兴趣的朋友参考。

欢迎在下面留言，完善本文内容，让更多的人学到更完美的知识。

欢迎到知识星球：完美Excel社群，进行技术交流和提问，获取更多电子资料，并通过社群加入专门的微信讨论群，更方便交流。