from mpl_toolkits.mplot3d import Axes3D
from matplotlib.colors import LogNorm
import matplotlib.pyplot as plt
import numpy as np
fig = plt.figure()
ax = Axes3D(fig, azim=-128, elev=43)
s = .05
X = np.arange(-2, 2.+s, s)
Y = np.arange(-1, 3.+s, s)
X, Y = np.meshgrid(X, Y)
Z = (1.-X)**2 + 100.*(Y-X*X)**2
# ax.plot_surface(X, Y, Z, rstride=1, cstride=1, norm = LogNorm(),
# cmap="viridis")
# Without using `` linewidth=0, edgecolor='none' '', the code may produce a
# graph with wide black edges, which will make the surface look much darker
# than the one illustrated in the figure above.
ax.plot_surface(X, Y, Z, rstride=1, cstride=1, norm=LogNorm(),
 linewidth=0, edgecolor='none', cmap="viridis")
# Set the axis limits so that they are the same as in the figure above.
ax.set_xlim([-2, 2.0]) 
ax.set_ylim([-1, 3.0]) 
ax.set_zlim([0, 2500]) 
plt.xlabel("x")
plt.ylabel("y")
plt.savefig("Rosenbrock function.svg", bbox_inches="tight")
plt.show()

clearall;
closeall;
hfig=figure(1);
s=0.05;
X=[-2:s:2+s];
Y=[-1:s:3+s];
[X,Y]=meshgrid(X,Y);
Z=(1-X).^2+100*(Y-X.^2).^2;
% Use log scale of Z for a batter usage of the color spectrum 
minZ=min(Z(:));
maxZ=max(Z(:));
C=minZ+(maxZ-minZ).*log(1+Z-minZ)./log(1+maxZ-minZ);
% See
% http://stackoverflow.com/questions/5073865/how-to-color-surface-with-stronger-contrast
% for how to color a surface with a even stronger contrast. 
% The method is as follows:
%C = Z;
%[~, index] = sort(C(:));
%C(index) = 1 : numel(index);
colormap(jet);
surf(X,Y,Z,C,'EdgeColor','none','LineStyle','none');

axis([-2,2,-1,3,0,2500]);
xlabel('x','fontsize',18);
ylabel('y','fontsize',18);
zlabel('f','fontsize',18);

% Note that the `-dsvg' option is only supported for Simulink systems
print(hfig,'-dsvg','rosenbrock');
% To produce eps and pdf, use the following code. Notice that `epstopdf' may not work on Windows. 
print(hfig,'-depsc','rosenbrock');
system('epstopdf rosenbrock.eps');