import matplotlib.pyplot as plt
import numpy as np

from matplotlib import cm

a = 1;   # velka polos
G = 1;
M1 = 10;   # hlavna zlozka polozena do centra sustavy [x, y] = [0, 0]
M2 = 1;   # vedlajsia zlozka v pozicii  [x, y] = [1, 0]

# Pre polohu taziska plati: M1*a1 = M2*a2, z coho po odvodeni dostaneme:
a1 = a*M2/(M1 + M2)
X = np.arange(-2, 2, 0.01)
Y = np.arange(-2, 2, 0.01)
X, Y = np.meshgrid(X, Y)

# Pre potencial v sustave dvoch telies plati v rovine obeznej drahy:
# fi(x, y) = -G*M1*r1 - G*M2/r2 - G*(M1 + M2)/2/a^3*ro**2
# kde ro je vzdialenost od taziska.
fi = (-G*M1/np.sqrt(X**2 + Y**2) - G*M2/np.sqrt((X - 1)**2 + Y**2)
      - G*(M1 + M2)/2*a**3*((X - a1)**2 + Y**2))

fig = plt.figure()
ax = fig.gca(projection='3d')

ax.plot_surface(X, Y, fi, cmap=cm.coolwarm, vmin=-50, vmax=0, rstride=5,
                cstride=5)
ax.set_zlim(-60, 0)