Python plotting a scatter plot

Question

Here is my code:

    def twomasses(M1,M2,x1,x2,p,h,n): 
    global gamma
    global m1
    global m2 
    gamma = 1
    m1 = M1
    m2 = M2
    x0_1 = [1, 2]
    x0_2 = [4, 5]
    p = 3
    v1 = [0, p/m1]
    v2 = [0, -p/m2]

    def F(x1, x2):
        Fa = ((gamma*m1*m2)/(la.norm((x2 - x1),2) ** 3))*(x2 - x1)
        return Fa

    def a1(f, m1):
        a1 = f/m1
        return a1

    def a2(f, m2):
        a2 = f/m2
        return a2

    def ruku_step(F, y, h): #first ruku step

        k1 = F(y)
        k2 = F(y + (h/2)*k1)
        k3 = F(y + (h/2)*k2)
        k4 = F(y + h*k3)

        y = y + (h/6)*(k1 + 2*k2 + 2*k3 + k4)

        return y

    f = lambda y: np.array([y[2],y[3],a1(F(y[0],y[1]),m1),a2(F(y[0],y[1]),m2)])
    y = list()
    y.append(np.array([x0_1,x0_2, v1, v2]))
    for i in range(0,n):
        y.append(ruku_step(f, np.array(y[i]), h))
    return y

y = twomasses(1,2,-1,2,5,.1, 50)


maxy = np.max([e[0:2,1] for e in y])
maxx = np.max([e[0:2,0] for e in y])
minx = np.min([e[0:2,0] for e in y])
miny = np.min([e[0:2,1] for e in y])
fig, ax = plt.subplots()
def animate(t):
    plt.clf()
    plt.scatter(y[t][0:2,0],y[t][0:2,1])
anim = FuncAnimation(fig, animate, interval=100, frames=100) 
plt.show()

I want to animate the graph so that you can see the movement of the masses. I tried following How to animate a scatter plot? but it is quite complex and wouldnt run for me. This will refresh the graph each time new points are introduced, but I want them all within one graph.

@undergrad Don't edit your question with the solution that you have found. You should post it as an answer instead — DavidG, Mar 29 '18 at 21:11
my answer is not the answer you posted. post yourself an answer and choose that one. (i can't delete mine now) — bobrobbob, Mar 29 '18 at 21:27

score 0 · Answer 1 · answered Mar 29 '18 at 21:00

lots of problems here: bad indent, linspace feeds floats and some parts of your code seem useless. but hey, it moves

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation
from numpy import linalg as la

def twomasses(M1,M2,x1,x2,p,h,n):
    global gamma
    global m1
    global m2
    gamma = 1
    m1 = M1
    m2 = M2
    x0_1 = [1, 2]
    x0_2 = [4, 5]
    p = 3
    v1 = [0, p/m1]
    v2 = [0, -p/m2]

    def F(x1, x2):
        Fa = ((gamma*m1*m2)/(la.norm((x2 - x1),2) ** 3))*(x2 - x1)
        return Fa

    def a1(f, m1):
        a1 = f/m1
        return a1

    def a2(f, m2):
        a2 = f/m2
        return a2

    def ruku_step(F, y, h): #first ruku step

        k1 = F(y)
        k2 = F(y + (h/2)*k1)
        k3 = F(y + (h/2)*k2)
        k4 = F(y + h*k3)

        y = y + (h/6)*(k1 + 2*k2 + 2*k3 + k4)

        return y

    f = lambda y: np.array([y[2],y[3],a1(F(y[0],y[1]),m1),a2(F(y[0],y[1]),m2)])
    y = list()
    y.append(np.array([x0_1,x0_2, v1, v2]))
    for i in range(0,n):
        y.append(ruku_step(f, np.array(y[i]), h))
    return y

y = twomasses(1,2,-1,2,5,.1, 50)
#~ print(y)

fig, ax = plt.subplots()
def animate(t):
    xdata = y[t][0:2,0]
    ydata = y[t][0:2,1]
    #~ def update(frame):
        #~ xdata.append(frame)
        #~ ydata.append(frame)
    ln.set_data(xdata, ydata)
    return ln,
ln, = plt.plot([], [], 'bs', animated=True)
maxy = np.max([e[0:2,1] for e in y])
maxx = np.max([e[0:2,0] for e in y])
minx = np.min([e[0:2,0] for e in y])
miny = np.min([e[0:2,1] for e in y])
def init():
    ax.set_xlim(minx-1, maxx+1)
    ax.set_ylim(miny-1, maxy+1)
    return ln,
ani =FuncAnimation(fig, animate, frames=np.arange(len(y)),
    init_func=init, blit=True)
plt.show()

Python plotting a scatter plot

1 Answers1