gravity sim

import matplotlib.pyplot as plt
import matplotlib.animation as anim
import numpy as np
from numpy.random import normal as nr
import math

class Planet(object):
    def __init__(self,x,y,vx,vy,m):
        self.x = x
        self.y = y
        self.vx = vx
        self.vy = vy
        self.m = m

fig, ax = plt.subplots()
plist = []
planets, = ax.plot([], [], 'o', ms=10, alpha=0.75)

au = 1.5e11

earth_x = 1 * au
earth_y = 0.
earth_vx = 0.
earth_vy = 29783.
earth_m = 5.74e24
sun_m = 2.0e30




plist.append(Planet(earth_x,earth_y,earth_vx,earth_vy,earth_m))
plist.append(Planet(0.,0.,0.,0.,sun_m))

def force(x1,y1,x2,y2,m1,m2):
    g = 6.67e-11

    dx = x2 - x1
    dy = y2 - y1

    dd = dx**2 + dy**2
    dist = np.sqrt(dd)

    f = g * m1 * m2 * (1./dist**2)

    a = math.atan2(dy, dx)

    fx = f * np.cos(a)
    fy = f * np.sin(a)

    return fx, fy

def move(x,y,vx,vy,fx,fy,m):

    dt = 3600.

    vx += (fx/m)
    vy += (fy/m)
    x += (vx * dt)
    y += (vy * dt)

    return x, y, vx, vy

def init():

    planets.set_data([],[])

    ax.set_xlim(-1e12,1e12)
    ax.set_ylim(-1e12,1e12)

    return planets,

def animate(i):
    xdata, ydata = [],[]

    for p in plist:
        t_fx = t_fy = 0.0

        for o in plist:

            if o != p:
                fx, fy = force(p.x,p.y,o.x,o.y,p.m,o.m)

                t_fx += fx
                t_fy += fy

        p.x, p.y, p.vx, p.vy = move(p.x,p.y,p.vx,p.vy,t_fx,t_fy,p.m)

        xdata.append(p.x)
        ydata.append(p.y)

    planets.set_data(xdata,ydata)

    return planets,

animation = anim.FuncAnimation(fig,animate,init_func=init, frames=10000, interval=1, blit=True)
plt.show()