在Python中为3个身体问题实现Runge-Kutta 4方法

时间:2019-03-04 17:05:58

标签: python differential-equations gravity runge-kutta

我正在尝试在Python中实现RK4方法来解决太阳,地球和木星系统。这段代码适用于Halley的彗星,当我有4个ODE而不是8个要解决时,我尝试将方法扩展到8个ODE,但是现在我一直保持直线,不知道为什么。有人可以向我指出我正在犯的愚蠢错误,因为我对为什么它不起作用感到困惑。这是代码:

import math
import numpy as np
import matplotlib.pyplot as plt

G = 6.67e-11
M = 1.989e+30       #mass of the sun
m1 = 5.972e+24      #mass of planet 1
m2 = 1.898e+27      #mass of planet 2
AU = 1.496e+11
a1 = 1.0*AU         #distance from planet 1 to the sun
a2 = 5.2*AU         #distance from planet 2 to the sun

#A = -(G*M)/(a1**3)
#B = -(G*M)/(a2**3)
##C = G/(abs(a2 - a1)**3)
#
#print("A = ", A, "B = ", B)

x_i1 = a1               #initial values for planet 1 in x and y direction
y_i1 = 0.0
v_x1i = 0.0
v_y1i = 3e+4

x_i2 = a2               #initial values for planet 2 in x and y direction
y_i2 = 0.0
v_x2i = 0.0
v_y2i = 13.07e+3

t_upper = 365*24*3600   #run the program to simulate a year
t = 0.0
t_i = 0.0
N = 10000


def f1(r, t):       #function containing the relevant equations and arrays

    x1 = r[0]       #storing position and velocity components in x and y direction for planet 1
    y1 = r[2]
    v_x1 = r[1]
    v_y1 = r[3]

    x2 = r[4]       #storing position and vecloity components in x and y direction for planet 2
    y2 = r[6]
    v_x2 = r[5]
    v_y2 = r[7]

    C = G/(((x1 - x2)**2 + (y1 - y2)**2)**1.5)      #equations for 3 body problem, just to condense the maths
    D = G/(((x2 - x1)**2 + (y2 - y1)**2)**1.5)
    A = (G*M)/(a1**3)
    B = (G*M)/(a2**3)


    dvx1 = (A*x1) + ((C*m2)*(x1 - x2))      #Acceleration in x and y for each planet
    dvy1 = (A*y1) + ((C*m2)*(y1 - y2))    
    dvx2 = (B*x2) - ((D*m1)*(x2 - x1))
    dvy2 = (B*y2) - ((D*m1)*(y2 - y1))

    return np.array([v_x1, dvx1, v_y1, dvy1, v_x2, dvx2, v_y2, dvy2]) #return array storing all the differential equations we want to solve



r = np.array([x_i1, v_x1i, y_i1, v_y1i, x_i2, v_x2i, y_i2, v_y2i], float) #define array with initial values


x_pnts1 = [x_i1]                #array containing the position in x for planet 1                 
v_x_pnts1 = [v_x1i]             #array containing the velocity in x for planet 1
y_pnts1 = [y_i1]                #position in y for planet 1
v_y_pnts1 = [v_y1i]             #velocity in y for planet 1

x_pnts2 = [x_i2]                #same as above, but for planet 2                 
v_x_pnts2 = [v_x2i]
y_pnts2 = [y_i2]
v_y_pnts2 = [v_y2i]

t_values = [t_i]

h = t_upper/N

while t < t_upper:

    k1 = h*f1(r, t)                 #RK4 routine, using the function defined above
    k2 = h*f1(r + 0.5*k1, t+h/2)
    k3 = h*f1(r + 0.5*k2, t+h/2)
    k4 = h*f1(r + k3, t+h)

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


    x_pnts1.append(r[0])     #append the relevant arrays with the values from the RK4 routine
    v_x_pnts1.append(r[1])
    y_pnts1.append(r[2])
    v_y_pnts1.append(r[3])

    x_pnts2.append(r[4])     
    v_x_pnts2.append(r[5])
    y_pnts2.append(r[6])
    v_y_pnts2.append(r[7])

    t += h    

    t_values.append(t)


aphelion_e = np.max(x_pnts1)
perihelion_e = abs(np.min(x_pnts1))

print("Perihelion of earth = ", perihelion_e/AU, "AU")
print("Aphelion of earth = ", aphelion_e/AU, "AU")

plt.plot(x_pnts1, y_pnts1)
plt.plot(0, 0, "o", 15)
plt.show()

0 个答案:

没有答案