Normalizing colors in matplotlib

I am trying to plot a surface using matplotlib using the code below:

from matplotlib import cm
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import axes3d, Axes3D
import pylab as p

vima=0.5

fig = plt.figure()
ax = fig.gca(projection='3d')
X = np.arange(0, 16.67, vima)
Y = np.arange(0, 12.5, vima)
X, Y = np.meshgrid(X, Y)

Z = np.sqrt(((1.2*Y+0.6*X)**2+(0.2*Y+1.6*X)**2)/(0.64*Y**2+0.36*X**2))

surf = ax.plot_surface(X, Y, Z,rstride=1, cstride=1, alpha=1,cmap=cm.jet,  linewidth=0)
fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()

If you run it you will see a blue surface, but I want to use the whole color range of jet... I know there is a class "matplotlib开发者_Python百科.colors.Normalize", but I don't know how to use it. Could you please add the necessary code in order to do it?

---

I realise that the poster's issue has already been resolved, but the question of normalizing the colors was never dealt with. Since I've figured out how I thought I'd just drop this here for anyone else who might need it.

First you create a norm and pass that to the plotting function, I've tried to add this to the OP's code.

from matplotlib import cm
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import axes3d, Axes3D
import pylab as p
import matplotlib

vima=0.5

fig = plt.figure()
ax = fig.gca(projection='3d')
X = np.arange(0, 16.67, vima)
Y = np.arange(0, 12.5, vima)
X, Y = np.meshgrid(X, Y)

Z = np.sqrt(((1.2*Y+0.6*X)**2+(0.2*Y+1.6*X)**2)/(0.64*Y**2+0.36*X**2))
Z = np.nan_to_num(Z)

# Make the norm
norm = matplotlib.colors.Normalize(vmin = np.min(Z), vmax = np.max(Z), clip = False)

# Plot with the norm
surf = ax.plot_surface(X, Y, Z,rstride=1, cstride=1, norm=norm, alpha=1,cmap=cm.jet,     linewidth=0)
fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()

The norm works the same way for the "imshow" command.

---

As JoshAdel noted in a comment (credit belongs to him), it appears that the surface plot is improperly ranging the colormap when a NaN is in the Z array. A simple work-around is to simply convert the NaN's to zero or very large or very small numbers so that the colormap can be normalized to the z-axis range.

from matplotlib import cm
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import axes3d, Axes3D
import pylab as p

vima=0.5

fig = plt.figure()
ax = fig.gca(projection='3d')
X = np.arange(0, 16.67, vima)
Y = np.arange(0, 12.5, vima)
X, Y = np.meshgrid(X, Y)

Z = np.sqrt(((1.2*Y+0.6*X)**2+(0.2*Y+1.6*X)**2)/(0.64*Y**2+0.36*X**2))
Z = np.nan_to_num(Z) # added this line

surf = ax.plot_surface(X, Y, Z,rstride=1, cstride=1, alpha=1,cmap=cm.jet,  linewidth=0)
fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()

---

Replying to an old question, I know, but the answers posted were at least in my case somewhat unsatisfactory. For those still stumbling here, I give a solution that worked for me.

Firstly, I did not want use zeros to replace NaNs, as for me they represent points with missing or undefined data. I'd rather not have anything plotted at these points. Secondly, the whole z range of my data was way above zero, so dotting the plot with zeros would result in an ugly and badly scaled plot.

Solution given by leifdenby was quite close, so +1 for that (though as pointed out, the explicit normalisation does not add to the earlier solution). I just dropped the NaN-to-zero replacement, and used the functions nanmin and nanmax instead of min and max in the color scale normalisation. These functions give the min and max of an array but simply ignore all NaNs. The code now reads:

# Added colors to the matplotlib import list
from matplotlib import cm, colors
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import axes3d, Axes3D
import pylab as p


vima=0.5

fig = plt.figure()
ax = fig.gca(projection='3d')
X = np.arange(0, 16.67, vima)
Y = np.arange(0, 12.5, vima)
X, Y = np.meshgrid(X, Y)

Z = np.sqrt(((1.2*Y+0.6*X)**2+(0.2*Y+1.6*X)**2)/(0.64*Y**2+0.36*X**2))

# MAIN IDEA: Added normalisation using nanmin and nanmax functions
norm = colors.Normalize(vmin = np.nanmin(Z), 
                        vmax = np.nanmax(Z))

# Added the norm=norm parameter
surf = ax.plot_surface(X, Y, Z,rstride=1, cstride=1, alpha=1, norm=norm, cmap=cm.jet,  linewidth=0)
fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()

Running this, I get a correctly scaled plot, with the (0, 0) datapoint missing. This is also the behaviour that I find most preferable, as the limit (x, y) to (0, 0) does not seem to exist for the function in question.

This has been my first contribution to StackOverflow, I hope it was a good one (wink).