kirkwood.py

import matplotlib.pylab as plt
import numpy as np, heyoka as hey
from copy import deepcopy
from multiprocessing.pool import ThreadPool
pi = np.pi

nproc = 20
ninst = 50

n_ast = 25
G = 0.01720209895**2 * 365**2
msun = 1.00000597682
masses = np.array([msun,  1/1047.355, 1/3501.6] + [0.] * n_ast)
sys = hey.make_nbody_sys(3 + n_ast, masses=masses, Gconst=G)

ta_orig = hey.taylor_adaptive(sys, np.zeros(6*(n_ast + 3)), high_accuracy=True, tol=1e-18, compact_mode=True)

tas = [deepcopy(ta_orig) for _ in range(ninst)]

svs = [ta.state.reshape((n_ast + 3), 6) for ta in tas]

for sv in svs:
    sv[0] = [-4.06428567034226e-3, -6.08813756435987e-3, -1.66162304225834e-6, +6.69048890636161e-6 * 365,
         -6.33922479583593e-6 * 365, -3.13202145590767e-9 * 365]
    sv[1] = [+3.40546614227466e+0, +3.62978190075864e+0, +3.42386261766577e-2, -5.59797969310664e-3 * 365,
         +5.51815399480116e-3 * 365, -2.66711392865591e-6 * 365]
    sv[2] = [+6.60801554403466e+0, +6.38084674585064e+0, -1.36145963724542e-1, -4.17354020307064e-3 * 365,
         +3.99723751748116e-3 * 365, +1.67206320571441e-5 * 365]

    for i in range(3, n_ast + 3):
        sv[i] = sv[0] + hey.random_elliptic_state(msun * G,
                                              [(2.49, 2.51),(1e-9,.1),(1e-9, .1), (0, 2*pi), (0, 2*pi), (0, 2*pi)])

    com_x = np.sum(sv[:, 0] * masses) / np.sum(masses)
    com_y = np.sum(sv[:, 1] * masses) / np.sum(masses)
    com_z = np.sum(sv[:, 2] * masses) / np.sum(masses)

    com_vx = np.sum(sv[:, 3] * masses) / np.sum(masses)
    com_vy = np.sum(sv[:, 4] * masses) / np.sum(masses)
    com_vz = np.sum(sv[:, 5] * masses) / np.sum(masses)
    
    sv[:, 0:3] -= [com_x, com_y, com_z]
    sv[:, 3:6] -= [com_vx, com_vy, com_vz]

def runner(i):
    sv = svs[i]
    ta = tas[i]

    def data_saver(j):
        oe = np.array([hey.cartesian_to_oe(msun * G, _ - sv[0]) for _ in sv[3:]])
        np.savetxt('kirkwood_a_{}_{}.txt'.format(i,j), oe[:, 0])
        np.savetxt('kirkwood_e_{}_{}.txt'.format(i,j), oe[:, 1])
        np.savetxt('kirkwood_i_{}_{}.txt'.format(i,j), oe[:, 2])
    
    for j in range(50):
        data_saver(j)
        ta.propagate_for(100000)

with ThreadPool(processes=nproc) as pool:
    pool.map(runner, range(ninst))