Python绘制散点图

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.

Answer 1

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()