How to automatize imshow plots in python

Question

I want to automatize an imshow degrading figure with python3. I would like to give a data frame and this to be plot no matter how many columns are given.

I tried this:

vmin = 3.5
vmax = 6
fig, axes = plt.subplots(len(list(df.columns)),1)

for i,j in zip(list(df.columns),range(1,len(list(df.columns))+1)):

    df = df.sort_values([i], ascending = False) 
    y = df[i].tolist()

    gradient = [y,y]

    plt.imshow(gradient, aspect='auto', cmap=plt.get_cmap('hot_r'), vmin=vmin, vmax=vmax)

    axes = plt.subplot(len(list(df.columns)),1,j)


sm = plt.cm.ScalarMappable(cmap=plt.get_cmap('hot_r'),norm=plt.Normalize(vmin,vmax))
sm._A = []
plt.colorbar(sm,ax=axes)

plt.show()

My problem is that the first set of data (first column of the df) is never showed. Also the map is not where I want it to be. This is exactly what I get:

But this is what I want:

Answer 1

You shouldn't use plt.subplot if you already have created your subplots via plt.subplots .

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

f = lambda x, s: x*np.exp(-x**2/s)/2
df = pd.DataFrame({"A" : f(np.linspace(0,50,600),70)+3.5,
                   "B" : f(np.linspace(0,50,600),110)+3.5,
                   "C" : f(np.linspace(0,50,600),150)+3.5,})

vmin = 3.5
vmax = 6

fig, axes = plt.subplots(len(list(df.columns)),1)

for col, ax in zip(df.columns,axes.flat):

    df = df.sort_values([col], ascending = False) 
    y = df[col].values

    gradient = [y,y]

    im = ax.imshow(gradient, aspect='auto', 
                   cmap=plt.get_cmap('hot_r'), vmin=vmin, vmax=vmax)

# Since all images have the same vmin/vmax, we can take any of them for the colorbar
fig.colorbar(im, ax=axes)

plt.show()