在Python Matplotlib中制作瀑布图

2022-06-04 10:11:53 浏览数 (1)

标签:Python,Matplotlib,瀑布图

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

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

1.创建标准的条形图。

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

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

代码语言:javascript复制
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。

代码语言:javascript复制
df['tot']= df['num'].cumsum()
df['tot1']= df['tot'].shift(1).fillna(0)

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

图2

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

代码语言:javascript复制
lower= df[['tot','tot1']].min(axis=1)
upper= df[['tot','tot1']].max(axis=1)

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

代码语言:javascript复制
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列来存储每个类别的适当颜色。

代码语言:javascript复制
df.loc[df['num']>= 0, 'color'] = 'green'
df.loc[df['num']< 0, 'color'] = 'red'

使用新颜色重新绘制条形图,如下所示:

代码语言:javascript复制
ax.bar(x=df['category'],height=upper,color= df['color'])

运行结果如下图4所示。

图4

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

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

代码语言:javascript复制
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,供有兴趣的朋友参考。

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

0 人点赞