Wrong plot of the solution of the equation, solved by the RK4 algorithm in Python

50 Views Asked by At

I'm trying to make a 2D-model of rocket flying from one planet to another in the Solar System. Before simulating the spaceflight, I must make a model of the planets itself. So, for every planet I have to solve the system of differential equations (using such units of measure of velocity and time, that make G*M_Sun/A.U**2 = 1 to reduce rounding errors) via RK4 method:

x' = vx
vx'= -x/(x**2+y**2)**1.5
y' = vy
vy'= -y/(x**2+y**2)**1.5

Since I don't want to write similar expressions 4 times, I use 4-vector to make an equation U' = f(U, t), where U = {x, y, vx, vy}, f = {vx, vy , -x/(x^2+y^2)^1.5, -y/(x^2+y^2)^1.5}, with starting conditions U0 = {x_0, 0, 0, vy_0}

However, this code (rewritten a bit from my original to be "standalone", sorry if your eyes are bleeding...)

import numpy as np
from scipy import integrate
from matplotlib import pyplot as plt

a1 = 1
eps1 = 0.0167 #Earth parameters
r_01 = (1-eps1)*(a1) 
v_01 = 1*((1+eps1)/(1-eps1))**0.5

def runge_kutta(x, t, func, dt):
    k1 = func(x, t)
    k2 = func(x+k1/2, t+dt/2)
    k3 = func(x+k2/2, t+dt/2)
    k4 = func(x+k3, t+dt)
    return dt*(k1+2*k2+2*k3+k4)/6

def f(U, t): #U = [x, y, vx, vy]
    res = np.zeros(4)
    res[0] = U[2]
    res[1] = U[3]
    res[2] = -U[0]/((U[0]**2+U[1]**2)**1.5)
    res[3] = -U[1]/((U[0]**2+U[1]**2)**1.5)
    return res

U_0 = np.zeros(4)
U_0[0]+=r_01
U_0[3]+=v_01

x=[]
y=[]
t=0
dt=0.1
while t<=2*np.pi:
    x.append(U_0[0])
    y.append(U_0[1])
    res = runge_kutta(U_0, t, f, dt)
    print(res)
    U_0 += res
    t += dt 

x = np.array(x)
y = np.array(y)
fig1, ax1 = plt.subplots(layout="tight")
ax1.scatter(x, y)
plt.show()

doesn’t draw an ellipse, as I expected, but something like this:

Using the scipy.integrate.solve_ivp function to solve the equation makes (practically) correct elliptical orbit. So, the mistake is, I think, in my RK4 function, and definitely very stupid mistake since the function is very simple... but I can't find it. Asking for help!

1

There are 1 best solutions below

1
lastchance On

Your RK4 function is, indeed, at fault.

If the changes in X are k = dt*gradient then you should include dt in the original expression for k ( and remove it from the weighted sum at the end). Thus:

def runge_kutta(x, t, func, dt):
    k1 = dt * func(x, t)
    k2 = dt * func(x+k1/2, t+dt/2)
    k3 = dt * func(x+k2/2, t+dt/2)
    k4 = dt * func(x+k3, t+dt)
    return (k1+2*k2+2*k3+k4)/6

Then you get enter image description here