Problem in creating leapfrog algorithm for 3-body problem using Python

问题

I'm trying to write a code for 3-body problem with leapfrog algorithm. I'm using "Moving Stars Around" by Piet Hut & Jun Makino as a guide.

The codes in the guide are written in C, but I'm trying to follow the exact workflow using Python as a start before experimenting with it.

The following is my attempt to follow the code from section 5.1.

import numpy as np

N = 3       #number of bodies
m = 1       #mass
dt = 0.01   #timestep
t_end = 10  #duration

r = [] 
v = []
a = [[0, 0, 0] for i in range(N)]

for i in range(N):
    phi = i * 2 * np.pi/3
    r.append([np.cos(phi), np.sin(phi), 0])

for i in range(N):
    for j in range(i+1, N):
        rji = []
        for k in range(3):
            rji.append(r[j][k] - r[i][k])
        r2 = 0
        for k in range(3):
            r2 += rji[k]**2
        r3 = r2 * np.sqrt(r2)
        for k in range(3):
            a[i][k] += m * rji[k] / r3
            a[j][k] -= m * rji[k] / r3


v_abs = np.sqrt(-a[0][0])
for i in range(N):
    phi = i * 2 * np.pi/3
    v.append([-v_abs * np.sin(phi),
              v_abs * np.cos(phi), 0])


ekin = 0
epot = 0
for i in range(N):

    for j in range(i+1, N):
        rji = [0, 0, 0]
        for k in range(3):
            rji[k] = r[j][k] - r[i][k]
        r2 = 0
        for k in range(3):
            r2 += rji[k]**2
        d = np.sqrt(r2)
        epot -= m**2 / d

    for k in range(3):
        ekin += 0.5 * m * v[i][k]**2

e_in = ekin + epot
print('Initial total energy E_in = ', e_in)


dt_out = 0.01
t_out = dt_out

for t in np.arange(0, t_end, dt):

    for i in range(N):
        for k in range(3):
            v[i][k] += a[i][k] * dt/2
        for k in range(3):
            r[i][k] += v[i][k] * dt

    for i in range(N):
        for k in range(3):
            a[i][k] = 0

    for i in range(N):
        for j in range(i+1, N):
            rji = []
        for k in range(3):
            rji.append(r[j][k] - r[i][k])
        r2 = 0
        for k in range(3):
            r2 += rji[k]**2
        r3 = r2 * np.sqrt(r2)
        for k in range(3):
            a[i][k] += m * rji[k] / r3
            a[j][k] -= m * rji[k] / r3

    for i in range(N):
        for k in range(3):
            v[i][k] += a[i][k] * dt/2
        '''
        if t >= t_out:
            for i in range(N):
                print(r[i][k], ' ')
            for k in range(N):
                print(v[i][k], ' ')
        '''
        t_out += dt_out


epot = 0
ekin = 0
for i in range(N):

    for j in range(i+1, N):
        rji = [0, 0, 0]
        for k in range(3):
            rji[k] = r[j][k] - r[i][k]
        r2 = 0
        for k in range(3):
            r2 += rji[k]**2
        d = np.sqrt(r2)
        epot -= m**2 / d

    for k in range(3):
        ekin += 0.5 * m * v[i][k]**2

e_out = ekin + epot
print('Final total energy E_out = ', e_out)
print('absolute energy error: E_out - E_in = ', e_out - e_in)
print('relative energy error: (E_out - E_in)/E_in = ', (e_out - e_in)/e_in)

I have defined the timestep dt = 0.01 and duration t_end = 10 as opposed to prompting an input. In section 5.4 the results should be:

|gravity> g++ -o leapfrog2 leapfrog2.C
|gravity> leapfrog2 > leapfrog2_0.01_10.out
Please provide a value for the time step
0.01
and for the duration of the run
10
Initial total energy E_in = -0.866025
Final total energy E_out = -0.866025
absolute energy error: E_out - E_in = 2.72254e-10
relative energy error: (E_out - E_in) / E_in = -3.14372e-10

along with a circular plot. However, the results from my code diverges:

Initial total energy E_in =  -0.8660254037844386
Final total energy E_out =  -0.39922101519288833
absolute energy error: E_out - E_in =  0.46680438859155027
relative energy error: (E_out - E_in)/E_in =  -0.5390192788244604

and of course, after I plot my results, they don't go in circle.

I'm wondering if there is a mistake that I made in translating the code. Any help would be appreciated!

回答1:

welcome to stack overflow!

First of all, the bug is a classic python issue: a part of your code is indented incorrectly. Specifically:

for i in range(N):
    for j in range(i+1, N):
        rji = []
    for k in range(3):
        rji.append(r[j][k] - r[i][k])
    r2 = 0
    for k in range(3):
        r2 += rji[k]**2
    r3 = r2 * np.sqrt(r2)
    for k in range(3):
        a[i][k] += m * rji[k] / r3
        a[j][k] -= m * rji[k] / r3

should rather be:

for i in range(N):
    for j in range(i+1, N):
        rji = []
        for k in range(3):
            rji.append(r[j][k] - r[i][k])
        r2 = 0
        for k in range(3):
            r2 += rji[k]**2
        r3 = r2 * np.sqrt(r2)
        for k in range(3):
            a[i][k] += m * rji[k] / r3
            a[j][k] -= m * rji[k] / r3

Let me give you an advise: if you are trying to learn python while following the book, try to code up a version that does its job (like the one we are discussing here) and then make an effort to make it more idiomatic. By using numpy you might remove the majority (if not all) of the loops over the spatial dimensions (at the very least!).

来源：https://stackoverflow.com/questions/53056901/problem-in-creating-leapfrog-algorithm-for-3-body-problem-using-python