We Keep Coding

Plot an energy potential with matplotlib

I want to plot the gravitational energy potential to highlight its extremums (the Lagrangian points around two celestial bodies).
Here is the function that returns the potential for each set of coordinates x and y:
def gravitational_potential(M,m,R,x,y):
G = 6.674*10**(-11)
omega2 = G*(M+m)/(R**3)
r = np.sqrt(x**2+y**2)
r2 = R*m/(M+m)
r1 = R-r2
phi = -G*(M/abs(r-r1)+m/abs(r-r2))-1/2*omega2*(x**2+y**2)
return phi
I want to use meshgrid and plot_surface to plot the surface and the contour of the potential but it doesn't work.
What am I doing wrong ?
PS: I managed to plot the potential with WolframAlpha so I know the math works.
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
from matplotlib import cm
import numpy as np
def gravitational_potential(M,m,R,x,y):
G = 6.674*10**(-11)
omega2 = G*(M+m)/(R**3)
r = np.sqrt(x**2+y**2)
r2 = R*m/(M+m)
r1 = R-r2
phi = -G*(M/abs(r-r1)+m/abs(r-r2))-1/2*omega2*(x**2+y**2)
return phi
fig = plt.figure()
ax = fig.gca(projection='3d')
X, Y = np.meshgrid(np.arange(-20, 20, 0.5), np.arange(-20, 20, 0.5))
M = 10
m = 1
R = 10
Z = gravitational_potential(M,m,R,X,Y)
ax.plot_surface(X, Y, Z, rstride=8, cstride=8, alpha=0.9)
cset = ax.contour(X, Y, Z, zdir='z', offset=-40, cmap=cm.coolwarm)
cset = ax.contour(X, Y, Z, zdir='x', offset=-20, cmap=cm.coolwarm)
cset = ax.contour(X, Y, Z, zdir='y', offset=20, cmap=cm.coolwarm)
ax.set_xlabel('X')
ax.set_xlim(-20, 20)
ax.set_ylabel('Y')
ax.set_ylim(-20, 20)
ax.set_zlabel('Z')
ax.set_zlim(-40, 40)
plt.show()
When I execute it I get the following:
runfile('C:/Users/python/Google Drive/lagrangepoint_maths/potential/gravitational_potential.py', wdir='C:/Users/python/Google Drive/lagrangepoint_maths/potential')
C:/Users/python/Google Drive/lagrangepoint_maths/potential/gravitational_potential.py:13: RuntimeWarning: divide by zero encountered in divide
phi = -G*(M/abs(r-r1)+m/abs(r-r2))-1/2*omega2*(x**2+y**2)
This is not really what I want. There is something wrong with Z. I want something like that:
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
from matplotlib import cm
fig = plt.figure()
ax = fig.gca(projection='3d')
X, Y, Z = axes3d.get_test_data(0.05)
ax.plot_surface(X, Y, Z, rstride=8, cstride=8, alpha=0.9)
cset = ax.contour(X, Y, Z, zdir='z', offset=-100, cmap=cm.coolwarm)
cset = ax.contour(X, Y, Z, zdir='x', offset=-40, cmap=cm.coolwarm)
cset = ax.contour(X, Y, Z, zdir='y', offset=40, cmap=cm.coolwarm)
ax.set_xlabel('X')
ax.set_xlim(-40, 40)
ax.set_ylabel('Y')
ax.set_ylim(-40, 40)
ax.set_zlabel('Z')
ax.set_zlim(-100, 100)
plt.show()

All of this are things one may just debug one by one:
Integer division in python 2 results in 0 if the nominator is smaller than the denominator. You may from __future__ import division or correct your code to divide by floats.
If you want to show numbers between -2 x 10^-8 and +2 x 10^-8 it is not useful to set the z_limits to -40 to 40.
If you want to show small features in the plot, you should not set the plotting resolution coarsely to rstride=8, cstride=8.
In total you would arrive at something like this:

Plotting streamlines in a Matplotlib 3Dplot?

I'm trying to plot streamlines on a plane in a 3D plot using Matplotlib. For the streamlines, I'd like to use the function streamplot(), because of its simplicity. Here's a MWE that I modified from the gallery (to attempt the streamplot() call):
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
from matplotlib import cm
X, Y, Z = axes3d.get_test_data(0.05)
fig = plt.figure()
ax = fig.gca()
plt.streamplot(X, Y, X, Y)
plt.show()
fig = plt.figure()
ax = fig.gca(projection='3d')
cset = ax.contour(X, Y, Z, zdir='z', offset=-10, cmap=cm.coolwarm)
ax.set_zlim(-100, 100)
plt.show()
fig = plt.figure()
ax = fig.gca(projection='3d')
cset = ax.streamplot(X, Y, X, Y, zdir='z', offset=-10, cmap=cm.coolwarm)
ax.set_zlim(-100, 100)
plt.show()
The first and second examples work as intended, however, the third example gives me this error:
TypeError: streamplot() got an unexpected keyword argument 'zdir'
which hints to the possibility of this not being implemented. If I check Axes3D, the function streamplot() I get:
In [28]: Axes3D.streamplot
Out[28]: <function matplotlib.axes._axes.Axes.streamplot>
From which we get that streamplot() is the 2D version, and hence can't be directly used in 3D plots.
Is there a way to circumvent this and get a streamlines plane in a 3D plot?

Set 3d plot linewidth in matplotlib 2.x

I'm using a newer version of matplotlib and the argument that sets the linewidth was removed. They seem to have changed it so I set it in Collections object, but I can't find a way of doing this.
I tried their example with a different linewidth:
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import numpy as np
fig = plt.figure()
ax = fig.gca(projection='3d')
# Make data.
X = np.arange(-5, 5, 0.25)
Y = np.arange(-5, 5, 0.25)
X, Y = np.meshgrid(X, Y)
R = np.sqrt(X**2 + Y**2)
Z = np.sin(R)
# Plot the surface.
surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm,
linewidth=10, antialiased=False)
# Customize the z axis.
ax.set_zlim(-1.01, 1.01)
ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
# Add a color bar which maps values to colors.
fig.colorbar(surf, shrink=0.5, aspect=5)
plt.show()
But as the figure shows, it doesn't add lines to my surface.
What is the new method for setting linewidths?
Thanks!

The linewidth can of course only take effect if there is actually a line to be shown. So one would need to specify the color of the lines to show in order to see them.
surf = ax.plot_surface(X, Y, Z, cmap="RdYlGn", linewidth=2, edgecolor="limegreen")

Minimum path on potential energy surface using matplotlib

I used Matplotlib to plot a 3D potential energy surface.
But, I also want to show a reaction path with the lowest potential barrier between the minima of a potential energy surface by arrows.
Can any one help me with this: this is the my code for the 3D surface:
#!/usr/bin/python
INPUT='input.txt'
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.mlab import griddata
from matplotlib import cm
from pylab import *
from matplotlib.ticker import LinearLocator, FormatStrFormatter
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
data = np.genfromtxt(INPUT)
x = data[:,0]
y = data[:,1]
z = data[:,2]
xi = np.linspace(min(x), max(x))
yi = np.linspace(min(y), max(y))
X, Y = np.meshgrid(xi, yi)
Z = griddata(x, y, z, xi, yi, interp='linear')
ax.contourf(X, Y, Z, 30, zdir='z', offset=-7, linewidth=0.2, cmap=cm.jet, antialiased=False, shade=True)
C = ax.contour(X, Y, Z, 30, colors='k', linewidths=0.2)
surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.jet, linewidth=0, antialiased=False, shade=False)
fig.colorbar(surf)
ax.clabel(C, inline=1, fontsize=10)
ax.set_xlabel(r"$d$$_{3}$ ($\AA$)")
ax.set_ylabel(r"$d$$_{4}$ ($\AA$)")
ax.set_zlabel(r"$\Delta$$E$$_{e}$ (kcal/mol)")
plt.show()
I want a graph like the one below:
3D surface
Source: http://pubs.acs.org/doi/abs/10.1021/ja801727k

Matplotlib: imshow in 3d plot

In the plot below, taken from matplotlib's gallery, contourf is used to create a 2d plot beneath the 3d one. My question is, is it possible to use imshow to do the same thing? I would like the colors in the 2d plot to be smoother.
Making the 2d plot seems to be possible because contourf accepts a zdir argument, while I've looked and imshow doesn't. That suggests that it isn't possible, but why not? pcolor would also get the job done, is it possible with that?

Just specify the levels= option for the contourf, e.g.
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt,numpy as np
plt.clf()
fig = plt.figure(1)
ax = fig.gca(projection='3d')
X, Y, Z = axes3d.get_test_data(0.05)
ax.plot_surface(X, Y, Z, rstride=8, cstride=8, alpha=0.3)
cset = ax.contourf(X, Y, Z, zdir='z', offset=-100,
levels=np.linspace(-100,100,1200),cmap=plt.cm.jet)
cset = ax.contourf(X, Y, Z, zdir='x', offset=-40, cmap=plt.cm.jet)
cset = ax.contourf(X, Y, Z, zdir='y', offset=40, cmap=plt.cm.jet)
ax.set_xlabel('X')
ax.set_xlim(-40, 40)
ax.set_ylabel('Y')
ax.set_ylim(-40, 40)
ax.set_zlabel('Z')
ax.set_zlim(-100, 100)
plt.show()

A little longer code then sega_sai's answer but faster and to my experience much better for more complex surfaces.
Use plot_surface to plot a flat surface where you want it and facecolors to color it with the values you want
You might need to make your data smoother with scipy's zoom
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt,numpy as np
plt.clf()
fig = plt.figure(1)
ax = fig.gca(projection='3d')
X, Y, Z = axes3d.get_test_data(0.05)
ax.plot_surface(X, Y, Z, rstride=8, cstride=8, alpha=0.3)
cset = ax.contourf(X, Y, Z, zdir='x', offset=-40, cmap=plt.cm.jet)
cset = ax.contourf(X, Y, Z, zdir='y', offset=40, cmap=plt.cm.jet)
### strating here:
# normalize Z to [0..1]
Z=Z-Z.min()
Z=Z/Z.max()
#use zoom to make your data smoother
from scipy.ndimage.interpolation import zoom
#make data 5 times smoother
X=zoom(X,5)
Y=zoom(Y,5)
Z=zoom(Z,5)
#draw a surface at -100, using the facecolors command to color it with the values of Z
cset = ax.plot_surface(X, Y, np.zeros_like(Z)-100,facecolors=plt.cm.jet(Z),shade=False)
ax.set_xlabel('X')
ax.set_xlim(-40, 40)
ax.set_ylabel('Y')
ax.set_ylim(-40, 40)
ax.set_zlabel('Z')
ax.set_zlim(-100, 100)
plt.show()
This also makes it a little harder to create a color bar, in order to that:
cb = plt.cm.ScalarMappable(cmap=plt.cm.jet)
cb.set_array(Z)
plt.colorbar(cb)
plt.show()