Python- compress lower end of y-axis in contourf plot

Tags: , , , ,



The issue

I have a contourf plot I made with a pandas dataframe that plots some 2-dimensional value with time on the x-axis and vertical pressure level on the y-axis. The field, time, and pressure data I’m pulling is all from a netCDF file. I can plot it fine, but I’d like to scale the y-axis to better represent the real atmosphere. (The default scaling is linear, but the pressure levels in the file imply a different king of scaling.) Basically, it should look something like the plot below on the y-axis. It’s like a log scale, but compressing the bottom part of the axis instead of the top. (I don’t know the term for this… like a log scale but inverted?) It doesn’t need to be exact.

Example of a properly-scaled time vs. pressure plot

Working example (written in Jupyter notebook)

#modules
import numpy             as np
import pandas            as pd
import matplotlib.pyplot as plt
from   matplotlib        import ticker, colors

#data
time = np.arange(0,10)
lev  = np.array([900,800,650,400,100])
df = pd.DataFrame(np.arange(50).reshape(5,10),index=lev,columns=time)
df.index.name = 'Level'
print(df)
        0   1   2   3   4   5   6   7   8   9
Level                                        
900     0   1   2   3   4   5   6   7   8   9
800    10  11  12  13  14  15  16  17  18  19
650    20  21  22  23  24  25  26  27  28  29
400    30  31  32  33  34  35  36  37  38  39
100    40  41  42  43  44  45  46  47  48  49
#lists for plotting
levtick = np.arange(len(lev))
clevels = np.arange(0,55,5)

#Main plot
fig, ax = plt.subplots(figsize=(10, 5))
im      = ax.contourf(df,levels=clevels,cmap='RdBu_r')

#x-axis customization
plt.xticks(time)
ax.set_xticklabels(time)
ax.set_xlabel('Time')

#y-axis customization
plt.yticks(levtick)
ax.set_yticklabels(lev)
ax.set_ylabel('Pressure')

#title and colorbar
ax.set_title('Some mean time series')
cbar = plt.colorbar(im,values=clevels,pad=0.01)
tick_locator = ticker.MaxNLocator(nbins=11)
cbar.locator = tick_locator
cbar.update_ticks()

The Question

How can I scale the y-axis such that values near the bottom (900, 800) are compressed while values near the top (200) are expanded and given more plot space, like in the sample above my code? I tried using ax.set_yscale('function', functions=(forward, inverse)) but didn’t understand how it works. I also tried simply ax.set_yscale('log'), but log isn’t what I need.

Answer

You can use a custom scale transformation with ax.set_yscale('function', functions=(forward, inverse)) as you suggested. From the documentation:

forward and inverse are callables that return the scale transform and its inverse.

In this case, define in forward() the function you want, such as the inverse of the log function, or a more custom one for your need. Call this function before your y-axis customization.

def forward(x):
    return 2**x

def inverse(x):
    return np.log2(x)

ax.set_yscale('function', functions=(forward,inverse))


Source: stackoverflow