Scale heatmaps of different dimensions to make cells equal in size

Question

I'm using Python package Seaborn to plot matrices as heatmaps to PDF document, each heatmap on a separate page. Heatmaps have the same column count (same within one run, but not constant in general) and possibly different row counts, so the PDF pages get different sizes. I would like heatmap cells to be squares of equal size throughout the whole PDF document. Ideally, I would specify cell size explicitly upon drawing heatmaps, but it seems there is no such option. I tried resizing the figure proportionally to row and column counts, but even for constant figure_width and figure_height = cell_size * row_count it doesn't result in cells of fixed cell_size, while figure_width = cell_size * column_count and figure_height = cell_size * row_count additionally creates an interplay between width and height of the heatmap (since cells are square, i.e. aspect ratio is 1:1) and the resulting scale becomes even more obscure.

Is there any solution to specify cell size in seaborn.heatmap (or matplotlib colorbar, colormesh etc.) excplicitly, or force it to certain value otherwise?

My code when figure_width is constant (20in) and figure_height = cell_size * row_count (cell_size is 0.2in), producing last heatmap (with 5-cell height) with cell size different from the rest:

from matplotlib.backends.backend_pdf import PdfPages
import seaborn
import numpy as np

with PdfPages('test.pdf') as pdf_pages:
    column_count = 150
    for row_count in [100, 30, 10, 5]:
        data = np.random.rand(row_count, column_count)
        fig = plt.figure(figsize=(20, 0.2 * row_count))
        seaborn.heatmap(data, square=True, xticklabels=range(column_count),
                             yticklabels=range(row_count), cbar=None)
        plt.tight_layout()
        pdf_pages.savefig(fig)

My code when figure_width = cell_size * column_count and figure_height = cell_size * row_count (cell_size is 0.2in), producing all the heatmaps with different cell sizes:

with PdfPages('test.pdf') as pdf_pages:
    column_count = 150
    for row_count in [100, 30, 10, 5]:
        data = np.random.rand(row_count, column_count)
        fig = plt.figure(figsize=(0.2 * column_count, 0.2 * row_count))
        seaborn.heatmap(data, square=True, xticklabels=range(column_count),
                             yticklabels=range(row_count), cbar=None)
        plt.tight_layout()
        pdf_pages.savefig(fig)

Ideally, solution should be parametric with respect to column_count, since it remains constant only within one PDF document, while changing throughout the runs.

Answer 1

You will need to calculate some parameters to obtain the desired result.
You may start with the figure width. Then define some margins around the heatmap and the number of rows (which should be constant).
From those numbers you can calculate the size in inches of once cell (pixel).
Then for each plot you can calculate the figure height needed to host a certain number of cells. This is done inside the loop. Also each figure needs to have the relative margins (calculated from the absolute ones) being set.

At the end this would give (ignoring the pdfpages, which seem irrelevant here):

import matplotlib.pyplot as plt
import seaborn
import numpy as np

####  Free input parameters ####
cols = [100, 30, 10, 5]
# start by specifying figure width and margins
figwidth= 20. #inch
marg_top = 0.7
marg_bottom = 0.7
marg_left = 0.7
marg_right = 0.7

# number of cells along width
cells_in_row = 150

####  Automatic calculation ####
# determine cell size in inch
cellsize = (figwidth-marg_left-marg_right)/float(cells_in_row)
# now loop over cols:
for cells_in_column in cols:
    # calculate figure height in inches
    figheight = cellsize*cells_in_column+marg_top+marg_bottom
    # set figure size
    fig = plt.figure(figsize=(figwidth, figheight))
    # adjust margins (relative numbers) according to absolute values
    fig.subplots_adjust(bottom =marg_bottom/figheight ,top=1.-marg_top/figheight,
                        left=marg_left/figwidth, right=1.-marg_right/figwidth)

    data = np.random.rand(cells_in_column, cells_in_row)
    seaborn.heatmap(data, square=True, xticklabels=range(cells_in_row),
                         yticklabels=range(cells_in_column), cbar=None)
    fig.savefig("fig{}.png".format(cells_in_column))

Of course you could also start with a given cellsize and calculate the figure width from it:

# start by specifying cell size and margins
cellsize = 0.08 #inch
marg_top = 0.7
marg_bottom = 0.7
marg_left = 0.7
marg_right = 0.7

# number of cells along width
cells_in_row = 150
# determine figure width
figwidth = cellsize* cells_in_row + marg_left+marg_right