Double header in Matplotlib Table

Question

I need to plot a table in matplotlib. The problem is some columns have one-level headers, some columns have double-level headers.

Here's what I need:

Here's simple example for one-level headers:

df = pd.DataFrame()
df['Animal'] = ['Cow', 'Bear']
df['Weight'] = [250, 450]
df['Favorite'] = ['Grass', 'Honey']
df['Least Favorite'] = ['Meat', 'Leaves']
df

fig = plt.figure(figsize=(9,2))
ax=plt.subplot(111)
ax.axis('off') 
table = ax.table(cellText=df.values, colColours=['grey']*df.shape[1], bbox=[0, 0, 1, 1], colLabels=df.columns)
plt.savefig('Table.jpg')

Last chunk of code produces next picture:

What changes do I need to make to have table I need?

Answer 1

Cell merge solution

You can merge the cells produced by ax.table , a la the cell merge function in an Excel spreadsheet. This allows for a completely automated solution in which you don't need to fiddle with any coordinates (save for the indices of the cell you want to merge):

import matplotlib.pyplot as plt
import pandas as pd

df = pd.DataFrame()
df['Animal'] = ['Cow', 'Bear']
df['Weight'] = [250, 450]
df['Favorite'] = ['Grass', 'Honey']
df['Least Favorite'] = ['Meat', 'Leaves']

fig = plt.figure(figsize=(9,2))
ax=fig.gca()
ax.axis('off')
r,c = df.shape

# ensure consistent background color
ax.table(cellColours=[['lightgray']] + [['none']], bbox=[0,0,1,1])

# plot the real table
table = ax.table(cellText=np.vstack([['', '', 'Food', ''], df.columns, df.values]), 
                 cellColours=[['none']*c]*(2 + r), bbox=[0, 0, 1, 1])

# need to draw here so the text positions are calculated
fig.canvas.draw()

# do the 3 cell merges needed
mergecells(table, (1,0), (0,0))
mergecells(table, (1,1), (0,1))
mergecells(table, (0,2), (0,3))

Output:

Here's the code for the mergecells function used above:

import matplotlib as mpl

def mergecells(table, ix0, ix1):
    ix0,ix1 = np.asarray(ix0), np.asarray(ix1)
    d = ix1 - ix0
    if not (0 in d and 1 in np.abs(d)):
        raise ValueError("ix0 and ix1 should be the indices of adjacent cells. ix0: %s, ix1: %s" % (ix0, ix1))

    if d[0]==-1:
        edges = ('BRL', 'TRL')
    elif d[0]==1:
        edges = ('TRL', 'BRL')
    elif d[1]==-1:
        edges = ('BTR', 'BTL')
    else:
        edges = ('BTL', 'BTR')

    # hide the merged edges
    for ix,e in zip((ix0, ix1), edges):
        table[ix[0], ix[1]].visible_edges = e

    txts = [table[ix[0], ix[1]].get_text() for ix in (ix0, ix1)]
    tpos = [np.array(t.get_position()) for t in txts]

    # center the text of the 0th cell between the two merged cells
    trans = (tpos[1] - tpos[0])/2
    if trans[0] > 0 and txts[0].get_ha() == 'right':
        # reduce the transform distance in order to center the text
        trans[0] /= 2
    elif trans[0] < 0 and txts[0].get_ha() == 'right':
        # increase the transform distance...
        trans[0] *= 2

    txts[0].set_transform(mpl.transforms.Affine2D().translate(*trans))

    # hide the text in the 1st cell
    txts[1].set_visible(False)

Answer 2

Yet another option would be to utilize matplotlib.gridspec.GridSpec to plot values and columns using a custom layout:

def format_axes(fig):
    for i, ax in enumerate(fig.axes):
        ax.tick_params(labelbottom=False, labelleft=False, labelright=False)
        ax.get_xaxis().set_ticks([])
        ax.get_yaxis().set_ticks([])


df = pd.DataFrame()
df['Animal'] = ['Cow', 'Bear']
df['Weight'] = [250, 450]
df['Favorite'] = ['Grass', 'Honey']
df['Least Favorite'] = ['Meat', 'Leaves']

fig = plt.figure(figsize=(9, 2))


gs = GridSpec(3, 4, figure=fig, wspace=0.0, hspace=0.0,height_ratios=[1, 1, 4])
# plot table header
ax1 = fig.add_subplot(gs[:-1, 0])
ax1.text(0.5, 0.5, df.columns[0], va="center", ha="center")
ax2 = fig.add_subplot(gs[:-1, 1])
ax2.text(0.5, 0.5, df.columns[1], va="center", ha="center")
ax3 = fig.add_subplot(gs[0, -2:])
ax3.text(0.5, 0.5, "Food", va="center", ha="center")
ax4 = fig.add_subplot(gs[1, -2])
ax4.text(0.5, 0.5, df.columns[2], va="center", ha="center")
ax5 = fig.add_subplot(gs[1, -1])
ax5.text(0.5, 0.5, df.columns[3], va="center", ha="center")
# plot table data
ax6 = fig.add_subplot(gs[-1, :])
table = ax6.table(cellText=df.values, cellLoc='center', bbox=[0, 0, 1, 1])

format_axes(fig)

plt.show()

Result

Answer 3

In addition to @tel's answer, I've made some changes to his code to solve my own problem - to merge more than 2 cells. Here's what I got:

def mergecells(table, cells):
    '''
    Merge N matplotlib.Table cells

    Parameters
    -----------
    table: matplotlib.Table
        the table
    cells: list[set]
        list of sets od the table coordinates
        - example: [(0,1), (0,0), (0,2)]

    Notes
    ------
    https://stackoverflow.com/a/53819765/12684122
    '''
    cells_array = [np.asarray(c) for c in cells]
    h = np.array([cells_array[i+1][0] - cells_array[i][0] for i in range(len(cells_array) - 1)])
    v = np.array([cells_array[i+1][1] - cells_array[i][1] for i in range(len(cells_array) - 1)])

    # if it's a horizontal merge, all values for `h` are 0
    if not np.any(h):
        # sort by horizontal coord
        cells = np.array(sorted(list(cells), key=lambda v: v[1]))
        edges = ['BTL'] + ['BT' for i in range(len(cells) - 2)] + ['BTR']
    elif not np.any(v):
        cells = np.array(sorted(list(cells), key=lambda h: h[0]))
        edges = ['TRL'] + ['RL' for i in range(len(cells) - 2)] + ['BRL']
    else:
        raise ValueError("Only horizontal and vertical merges allowed")

    for cell, e in zip(cells, edges):
        table[cell[0], cell[1]].visible_edges = e
        
    txts = [table[cell[0], cell[1]].get_text() for cell in cells]
    tpos = [np.array(t.get_position()) for t in txts]

    # transpose the text of the left cell
    trans = (tpos[-1] - tpos[0])/2
    # didn't had to check for ha because I only want ha='center'
    txts[0].set_transform(mpl.transforms.Affine2D().translate(*trans))
    for txt in txts[1:]:
        txt.set_visible(False)

Answer 4

I guess that the only way is to add the headers manually. You can control their exact position and size with the bbox argument. See my example below. You can get more details from this answer: https://stackoverflow.com/a/37440236/2912478

#!/usr/bin/env python

import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame()
df['Animal'] = ['Cow', 'Bear']
df['Weight'] = [250, 450]
df['Favorite'] = ['Grass', 'Honey']
df['Least Favorite'] = ['Meat', 'Leaves']
df

fig = plt.figure(figsize=(9,2))
ax=plt.subplot(111)
ax.axis('off') 

plt.table(cellText=[['Animal', 'Weight']],
                     loc='bottom',
                     bbox=[0, 0.6, 0.5, 0.3]
                     )

plt.table(cellText=[['Food']],
                     loc='bottom',
                     bbox=[0.5, 0.75, 0.5, 0.15]
                     )

plt.table(cellText=[['Favorite', 'Least favorite']],
                     loc='bottom',
                     bbox=[0.5, 0.6, 0.5, 0.15]
                     )

plt.table(cellText=df.values,
                     loc='bottom',
                     bbox=[0, 0, 1, 0.6]
                     )

plt.show()

Here is the output I get: