How to arrange graphs and texts like titles and axis lables appealingly with matplotlib with multiple subplots?

Question

I want to create a before-and-after plot of the three axis of my measurement system. This is close to what I want. However:

• How do I have the "before" and "after" titles span subplot 1+2 and 4+5, respectively? (the tabs dont work as expected)
• Like "before" and "after" should be above a column, i would like to have the "x-Axis", "y-Axis" etc infront of the row of graphs. How do I do that?
• How do I join subplot 1+2 and 4+5 together, so that they touch? wspace=.0 doesnt work as expected.
• How do I reduce the width in the middle, where subplot 3 would be, so that the other sides can take up more space?
• How do I add some hspace between the fig.suptitle and the graphs?
• How can I make my code more elegant? There is a lot of repetition.

from matplotlib.pyplot import figure

def plot_before_and_after(data_before, data_after, title):
    shape = data_before.shape
    sensor_num = shape[0]
    n_start = 20  # number picked at random
    N = 2 ** 12  # power of two is good
    n_stop = n_start + N
    p_stop = n_start + 40  # one periode @50Hz at the sampling rate
    x_long = range(n_start, n_stop)
    x_short = range(n_start, p_stop)
    
    cmap = plt.get_cmap('jet_r')
    axis_list = ['x', 'y', 'z']
    fig = figure(num=None, figsize=(10, 10), dpi=80, facecolor='w', edgecolor='k')
    
    fig.suptitle(title + "\nbefore \t\t\t\t\tafter")
    plots = []
    for axis_cnt in range(0, 3):
        ax0 = plt.subplot(3, 5, axis_cnt * 5 + 1,
                                              title="before, {}-Axis".format(axis_list[axis_cnt]))
        for sensor_cnt in range(0, sensor_num):
            color = cmap(float(sensor_cnt) / sensor_num)
            plt.plot(x_long, data_before[sensor_cnt, n_start:n_stop, axis_cnt], color=color,
                     label="sensor" + str(sensor_cnt))
    
        ax1 = plt.subplot(3, 5, axis_cnt * 5 + 2,
                                              title="before, {}-Axis".format(axis_list[axis_cnt]),
                                              sharey=ax0)
        for sensor_cnt in range(0, sensor_num):
            color = cmap(float(sensor_cnt) / sensor_num)
            plt.plot(x_short, data_before[sensor_cnt, n_start:p_stop, axis_cnt], color=color,
                     label="sensor" + str(sensor_cnt))
        plt.setp(ax1.get_yticklabels(), visible=False)
    
        ax3 = plt.subplot(3, 5, axis_cnt * 5 + 4,
                                              title="after, {}-Axis".format(axis_list[axis_cnt]))
        for sensor_cnt in range(0, sensor_num):
            color = cmap(float(sensor_cnt) / sensor_num)
            plt.plot(x_long, data_after[sensor_cnt, n_start:n_stop, axis_cnt], color=color,
                     label="sensor" + str(sensor_cnt))
    
        ax4 = plt.subplot(3, 5, axis_cnt * 5 + 5,
                                              title="after, {}-Axis".format(axis_list[axis_cnt]),sharey=ax3)
        for sensor_cnt in range(0, sensor_num):
            color = cmap(float(sensor_cnt) / sensor_num)
            plt.plot(x_short, data_after[sensor_cnt, n_start:p_stop, axis_cnt], color=color,
                     label="sensor" + str(sensor_cnt))
        plt.setp(ax4.get_yticklabels(), visible=False)
    plt.subplots_adjust(wspace=.0)
    plt.legend()
    plt.show()

Answer 1

Here's a preliminary answer that may help you. I used Matplotlib's GridSpec (see here for useful information) and the add_subplot method, both of which seem to be convenient in these cases. The gridspec is what allows us to create the two groups of subplots independently formatted; the add_subplot generates the individual axes.

Code

import matplotlib.pyplot as plt

nrow, ncol = 3, 2             # Number of rows and cols of gridspecs     
lborder = [0.1, 0.6]          # Left border coordinates of gridspecs
rborder = [0.45, 0.95]        # Right border coordinates of gridspecs
tborder = 0.92                # Top border coordinate of gridspecs
gtitles = ['Before', 'After']
txt_axis = ['X-axis', 'Y-axis', 'Z-axis']

fig = plt.figure(figsize=(10, 10), dpi=80, facecolor='w', edgecolor='k')

for i in range(2):
    gs = fig.add_gridspec(nrows=nrow, ncols=ncol, left=lborder[i], 
                          right=rborder[i], top=tborder, hspace=0.45, wspace=0)
    for j in range(nrow):
        ax0 = fig.add_subplot(gs[j,0])
        ax0.plot([1,2,3])
        plt.text(-0.4, 0.5, txt_axis[j],
                 horizontalalignment='center',verticalalignment='center',
                 transform = ax0.transAxes, fontsize = 12)
        if j == 0:
            fig.text(1, 1.1, gtitles[i], fontsize=12, horizontalalignment = 
                     'center', transform = ax0.transAxes)
        for k in range(1, ncol):
            ax1 = fig.add_subplot(gs[j,k], sharey = ax0)
            plt.setp(ax1.get_yticklabels(), visible=False)
            ax1.plot([1,2,3])
            
fig.suptitle('Figure title', fontsize = 14)

As for your questions:

• I created the 'Before' and 'After' titles using text, based on this answer).
• Same thing for the "-axis" text. Note that it will probably overlap with any axes label you write on the vertical axis. Also note that now we have to shift the left gridspec slightly to the right (via the leftargument of add_gridspec).
• wspace can be introduced in add_gridspec too. I don't know why it doesn't work in your code.
• For the space in between the 2 gridspecs, use the left and right arguments in the add_gridspec function.
• The space between the main title and the subplots can be achieved via the top argument in add_gridspec.
• Your inner loops seem very similar, perhaps you could define a function and save some lines of code. In my case, I tried to encase everything in a loop.

Hope it helps.