matplotlib: plot interpolated data with a small set of colors

Viewed 78

Here is a plot of 2D data:

import scipy.interpolate
import numpy as np
x = np.tile(np.append(np.linspace(0,85,18), np.linspace(85,0,18)), 9)
y = np.repeat(np.linspace(0,85,18), 18) 
z = np.array([0.05036811, 0.06607374, 0.08139464, 0.0956418 , 0.10787732,
       0.11657655, 0.12201094, 0.12443012, 0.12548632, 0.11411325,
       0.12529082, 0.12252976, 0.12296086, 0.11887216, 0.11128453,
       0.10016285, 0.08659644, 0.07139233, 0.08428328, 0.0981452 ,
       0.11158404, 0.12046799, 0.12423737, 0.12547053, 0.12633549,
       0.12688213, 0.11537648, 0.12672348, 0.12594674, 0.12489706,
       0.12354561, 0.11875938, 0.10829915, 0.09419385, 0.07766486,
       0.06038621, 0.07106876, 0.08942542, 0.10561504, 0.11797838,
       0.12304032, 0.12518906, 0.12642471, 0.12728667, 0.12777341,
       0.11630886, 0.12813333, 0.12769195, 0.12678413, 0.12593546,
       0.12446095, 0.1212733 , 0.11107825, 0.09576687, 0.10478517,
       0.11746357, 0.12376542, 0.12545398, 0.12674177, 0.12769897,
       0.12845215, 0.12907674, 0.11716154, 0.1289449 , 0.12828296,
       0.12731298, 0.12625227, 0.12393535, 0.12311751, 0.11453633,
       0.09909011, 0.08059009, 0.08776442, 0.10639528, 0.11985718,
       0.12445017, 0.125364  , 0.12697913, 0.12817744, 0.12899654,
       0.1296274 , 0.11748387, 0.1296289 , 0.12927725, 0.12866469,
       0.12768468, 0.12655931, 0.12502639, 0.12237487, 0.11177202,
       0.11643469, 0.12337241, 0.12543418, 0.12697763, 0.12823584,
       0.12913138, 0.12976049, 0.13008331, 0.11780469, 0.13003795,
       0.1296289 , 0.1286692 , 0.12772002, 0.12614174, 0.12489882,
       0.12243226, 0.11138503, 0.09316923, 0.09630299, 0.11423958,
       0.12320298, 0.12537829, 0.12662147, 0.12734732, 0.128984  ,
       0.12944268, 0.13000837, 0.11779341, 0.13005624, 0.12987904,
       0.12955321, 0.12845691, 0.12730797, 0.12617959, 0.12413762,
       0.11881527, 0.11971632, 0.1244251 , 0.12625002, 0.12716185,
       0.12858724, 0.12953065, 0.13019485, 0.13030062, 0.1180636 ,
       0.13012016, 0.12981137, 0.12900305, 0.1271438 , 0.12705909,
       0.12556803, 0.12339597, 0.11510454, 0.0979304 , 0.09679149,
       0.1144243 , 0.12337166, 0.12506549, 0.12646858, 0.12671872,
       0.12907048, 0.12966525, 0.13017279, 0.1180463 , 0.12971537,
       0.13024698, 0.12944017, 0.12855541, 0.12735033, 0.12641795,
       0.12444014, 0.11885838, 0.11649133, 0.12365464, 0.12567354,
       0.12713403, 0.12830903, 0.12922487, 0.13004722, 0.13008607,
       0.11795131, 0.13005449, 0.12940909, 0.1286692 , 0.12692098,
       0.12629413, 0.12512364, 0.12294933, 0.11188782, 0.09370009,
       0.08859805, 0.10712565, 0.12041937, 0.12441132, 0.12558056,
       0.1265771 , 0.12834988, 0.12911359, 0.12956399, 0.11764503,
       0.12985423, 0.12941335, 0.12883537, 0.12787818, 0.1267122 ,
       0.12531814, 0.12274657, 0.11220312, 0.10501977, 0.11832201,
       0.12393736, 0.12547078, 0.12735509, 0.12775461, 0.12866093,
       0.12929781, 0.11706455, 0.12899929, 0.12827344, 0.1277303 ,
       0.1265756 , 0.12475996, 0.12330599, 0.11498323, 0.09990394,
       0.08120792, 0.07206606, 0.0905934 , 0.10643238, 0.11908471,
       0.12351754, 0.12546627, 0.1259099 , 0.12727388, 0.12796114,
       0.11629809, 0.12844864, 0.12781853, 0.12714556, 0.1261069 ,
       0.12449328, 0.12122392, 0.11140358, 0.09569093, 0.08407925,
       0.09956633, 0.11271694, 0.12126302, 0.12414389, 0.12542516,
       0.12636832, 0.12697587, 0.11476191, 0.1265568 , 0.1258129 ,
       0.12407671, 0.12347092, 0.11950553, 0.10953856, 0.0950751 ,
       0.07892734, 0.06170332, 0.05015381, 0.06627701, 0.08156583,
       0.09561674, 0.10810916, 0.11713773, 0.12168986, 0.12396017,
       0.12482788, 0.113526  , 0.12514144, 0.12426043, 0.1230942 ,
       0.11985016, 0.1120372 , 0.1003571 , 0.08672451, 0.07088604,
       0.05653386, 0.07055369, 0.08260624, 0.09597916, 0.10580327,
       0.11302447, 0.11736256, 0.11762372, 0.10776979, 0.11869822,
       0.11601411, 0.1109048 , 0.1031921 , 0.09314843, 0.08104826,
       0.066176  , 0.05259229, 0.03810226, 0.02648483, 0.03889203,
       0.05155138, 0.06425685, 0.07563518, 0.08548961, 0.09354795,
       0.09955981, 0.10351091, 0.09563303, 0.10492126, 0.10177397,
       0.09629798, 0.08844341, 0.07845388, 0.06749337, 0.05530898,
       0.04241153, 0.02875563, 0.03946199, 0.05035257, 0.06024861,
       0.06917767, 0.0765059 , 0.08184855, 0.0847219 , 0.0768177 ,
       0.08371633, 0.07972989, 0.07399749, 0.06678105, 0.05774195,
       0.04757397, 0.03656985, 0.02585146, 0.01638251])
z *= 100.
xi, yi = np.mgrid[x.min():x.max():1000j, y.min():y.max():1000j]
rbf = scipy.interpolate.Rbf(x, y, z) 
zi = rbf(xi, yi) 
im = plt.imshow(zi.T, origin='lower', extent=[x.min(), x.max(), y.min(), y.max()], cmap='gray')

enter image description here

Now I want to plot it with a small number of levels, from 0.85*max(z) to max(z), 5% per level. Here is a picture made with gnuplot:

enter image description here

How can I do it?

Here is a part of gnuplot script:

set size ratio -1
set pm3d at b interp 25,25
set view map
stats '18.08.2021_14-53-57.txt' u 3
set cbrange [STATS_max*0.8*100:STATS_max*100]
set palette gray maxcolors 4
splot '18.08.2021_14-53-57.txt' u ($1/10.+1.):($2/10.+1.):($3*100) w pm3d
4 Answers

The contourf function does this kind of thresholding automatically and additionally does some interpolation for you to get smooth contours. It also has a levels argument which cleanly handles your requirements for the color splits.

axz = zi.max()
levels = np.arange(0.85*maxz, maxz, .5)

plt.contourf(xi, yi, zi, levels=levels, cmap='gray')
plt.colorbar()
plt.show()

resulting plot

You need to bin your data, you can use pandas.cut for that.

Replace the last line in your code with:

import pandas as pd
bins = list(np.linspace(zi.min(), zi.max(), 10))
zi_cut = pd.cut(zi.flatten(), bins=[0]+bins, labels=bins).astype(float).reshape(zi.shape)

im = plt.imshow(zi_cut.T, origin='lower', extent=[x.min(), x.max(), y.min(), y.max()], cmap='gray')

output:

binned data

Of course you can pick whichever bins you want.

You can also use numpy.digitize, but this gives you the indexes of the bins:

plt.imshow(np.digitize(zi, bins=bins).T, origin='lower', extent=[x.min(), x.max(), y.min(), y.max()], cmap='gray')

numpy digitize

The key is using the vmin/vmax parameters in plt.imshow() like this, and plt.colorbar() with ticks brings up the colorbar:

vmin = 0.85*zi.max()
vmax = zi.max()
ticks = np.linspace(vmin, vmax, 6)

im = plt.imshow(
    zi.T, origin='lower', extent=[x.min(), x.max(), y.min(), y.max()], cmap='gray', 
    vmin=vmin, vmax=vmax)
plt.colorbar(ticks=ticks)

which gives this:

image screenshot

If you do want your data binned before plotting and therefore approximate your data in that way, you can use the code in @mozway answer.

You can also just define a custom colormap that suits your needs. For instance with ListedColormap you can easily define your own:

from matplotlib.colors import ListedColormap
color_steps = [0, 0.33, 0.66, 1]
colors = np.array(color_steps).reshape(-1, 1) @ np.ones((1, 3)) # needs to be an N x 3 or N x 4 array
cmap = ListedColormap(colors, 'stepped')
plt.imshow(zi.T, origin='lower',  cmap=cmap)
plt.colorbar()
Related