Imitating the waterfall plots in Origin with Matplotlib

Question

I am attempting to create the waterfall plots made by Origin (see image below) with Python and Matplotlib.

or

The general scheme makes sense to me, you start with a 2D matrix as if you wanted to make a surface plot, and then you can follow any of the recipes shown in the StackOverflow question here. The idea is to plot each line of the matrix as an individual curve in 3D space.

This matplotlib method results in a plot like the one below:

The struggle I am having is that the sense of perspective that is clear in the Origin plot is lost in the matplotlib version. You can argue that this is partially due to the camera angle, but I think more importantly it comes from the closer lines appearing "in front" of the lines that are farther away.

My question is, how would you properly imitate the waterfall plot from Origin in Matplotlib with the perspective effect? I don't really understand what it is about the two plots that's so different, so even defining the exact problem is difficult.

Answer 1

Update: as you've now updated your question to make clearer what you're after, let me demonstrate three different ways to plot such data, which all have lots of pros and cons. The general gist (at least for me!) is that matplotlib is bad in 3D, especially when it comes to creating publishable figures (again, my personal opinion, your mileage may vary.)

What I did: I've used the original data behind the second image you've posted. In all cases, I used zorder and added polygon data (in 2D: fill_between(), in 3D: PolyCollection) to enhance the "3D effect", i.e. to enable "plotting in front of each other". The code below shows:

• plot_2D_a() uses color to indicate angle, hence keeping the original y-axis; though this technically can now only be used to read out the foremost line plot, it still gives the reader a "feeling" for the y scale.

• plot_2D_b() removes unnecessary spines/ticks and rather adds the angle as text labels; this comes closest to the second image you've posted

• plot_3D() uses mplot3d to make a "3D" plot; while this can now be rotated to analyze the data, it breaks (at least for me) when trying to zoom, yielding cut-off data and/or hidden axes.

In the end there are many ways to achieve a waterfall plot in matplotlib, and you have to decide yourself what you're after. Personally, I'd probably us plot_2D_a() most of the time, since it allows for easy rescaling in more or less "all 3 dimensions" while also keeping proper axes (+colorbar) that allow the reader to get all relevant information once you publish it somewhere as a static image.

---

Code:

import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.collections import PolyCollection
import numpy as np


def offset(myFig,myAx,n=1,yOff=60):
    dx, dy = 0., yOff/myFig.dpi 
    return myAx.transData + mpl.transforms.ScaledTranslation(dx,n*dy,myFig.dpi_scale_trans)

## taken from 
## http://www.gnuplotting.org/data/head_related_impulse_responses.txt
df=pd.read_csv('head_related_impulse_responses.txt',delimiter="\t",skiprows=range(2),header=None)
df=df.transpose()

def plot_2D_a():
    """ a 2D plot which uses color to indicate the angle"""
    fig,ax=plt.subplots(figsize=(5,6))
    sampling=2
    thetas=range(0,360)[::sampling]

    cmap = mpl.cm.get_cmap('viridis')
    norm = mpl.colors.Normalize(vmin=0,vmax=360)

    for idx,i in enumerate(thetas):
        z_ind=360-idx ## to ensure each plot is "behind" the previous plot
        trans=offset(fig,ax,idx,yOff=sampling)

        xs=df.loc[0]
        ys=df.loc[i+1]

        ## note that I am using both .plot() and .fill_between(.. edgecolor="None" ..) 
        #  in order to circumvent showing the "edges" of the fill_between 
        ax.plot(xs,ys,color=cmap(norm(i)),linewidth=1, transform=trans,zorder=z_ind)
        ## try alpha=0.05 below for some "light shading"
        ax.fill_between(xs,ys,-0.5,facecolor="w",alpha=1, edgecolor="None",transform=trans,zorder=z_ind)

    cbax = fig.add_axes([0.9, 0.15, 0.02, 0.7]) # x-position, y-position, x-width, y-height
    cb1 = mpl.colorbar.ColorbarBase(cbax, cmap=cmap, norm=norm, orientation='vertical')
    cb1.set_label('Angle')

    ## use some sensible viewing limits
    ax.set_xlim(-0.2,2.2)
    ax.set_ylim(-0.5,5)

    ax.set_xlabel('time [ms]')

def plot_2D_b():
    """ a 2D plot which removes the y-axis and replaces it with text labels to indicate angles """
    fig,ax=plt.subplots(figsize=(5,6))
    sampling=2
    thetas=range(0,360)[::sampling]

    for idx,i in enumerate(thetas):
        z_ind=360-idx ## to ensure each plot is "behind" the previous plot
        trans=offset(fig,ax,idx,yOff=sampling)

        xs=df.loc[0]
        ys=df.loc[i+1]

        ## note that I am using both .plot() and .fill_between(.. edgecolor="None" ..) 
        #  in order to circumvent showing the "edges" of the fill_between 
        ax.plot(xs,ys,color="k",linewidth=0.5, transform=trans,zorder=z_ind)
        ax.fill_between(xs,ys,-0.5,facecolor="w", edgecolor="None",transform=trans,zorder=z_ind)

        ## for every 10th line plot, add a text denoting the angle. 
        #  There is probably a better way to do this.
        if idx%10==0:
            textTrans=mpl.transforms.blended_transform_factory(ax.transAxes, trans)
            ax.text(-0.05,0,u'{0}º'.format(i),ha="center",va="center",transform=textTrans,clip_on=False)

    ## use some sensible viewing limits
    ax.set_xlim(df.loc[0].min(),df.loc[0].max())
    ax.set_ylim(-0.5,5)

    ## turn off the spines
    for side in ["top","right","left"]:
        ax.spines[side].set_visible(False)
    ## and turn off the y axis
    ax.set_yticks([])

    ax.set_xlabel('time [ms]')

#--------------------------------------------------------------------------------
def plot_3D():
    """ a 3D plot of the data, with differently scaled axes"""
    fig=plt.figure(figsize=(5,6))
    ax= fig.gca(projection='3d')

    """                                                                                                                                                    
    adjust the axes3d scaling, taken from https://stackoverflow.com/a/30419243/565489
    """
    # OUR ONE LINER ADDED HERE:                to scale the    x, y, z   axes
    ax.get_proj = lambda: np.dot(Axes3D.get_proj(ax), np.diag([1, 2, 1, 1]))

    sampling=2
    thetas=range(0,360)[::sampling]
    verts = []
    count = len(thetas)

    for idx,i in enumerate(thetas):
        z_ind=360-idx

        xs=df.loc[0].values
        ys=df.loc[i+1].values

        ## To have the polygons stretch to the bottom, 
        #  you either have to change the outermost ydata here, 
        #  or append one "x" pixel on each side and then run this.
        ys[0] = -0.5 
        ys[-1]= -0.5

        verts.append(list(zip(xs, ys)))        

    zs=thetas

    poly = PolyCollection(verts, facecolors = "w", edgecolors="k",linewidth=0.5 )
    ax.add_collection3d(poly, zs=zs, zdir='y')

    ax.set_ylim(0,360)
    ax.set_xlim(df.loc[0].min(),df.loc[0].max())
    ax.set_zlim(-0.5,1)

    ax.set_xlabel('time [ms]')

# plot_2D_a()
# plot_2D_b()
plot_3D()
plt.show()