In [ ]:

    

from IPython.display import display
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from pyodesys.symbolic import SymbolicSys
%matplotlib inline


    

In [ ]:

    

def dydt(t, y, params=(), be=None):
    x0, y0, z0, x1, y1, z1, x2, y2, z2, px0, py0, pz0, px1, py1, pz1, px2, py2, pz2 = y
    m0, m1, m2 = params
    coord = [[x0, y0, z0], [x1, y1, z1], [x2, y2, z2]]
    mmntm = [[px0, py0, pz0], [px1, py1, pz1], [px2, py2, pz2]]

    def r(a, b):
        return be.sqrt(sum((c0 - c1)**2 for c0, c1 in zip(a, b)))
    r01 = r(coord[0], coord[1])
    r02 = r(coord[0], coord[2])
    r12 = r(coord[1], coord[2])
    f01 = m0*m1/r01
    f02 = m0*m2/r02
    f12 = m1*m2/r12
    def e(a, b, denom):
        return [(c1 - c0)/denom for c0, c1 in zip(a, b)]
    e01 = e(coord[0], coord[1], r01)
    e02 = e(coord[0], coord[2], r02)
    e12 = e(coord[1], coord[2], r12)
    f0 = [f01*c01 + f02*c02 for c01, c02 in zip(e01, e02)]
    f1 = [-f01*c01 + f12*c12 for c01, c12 in zip(e01, e12)]
    f2 = [-f02*c02 - f12*c12 for c02, c12 in zip(e02, e12)]
    return [mmntm[i][j]/params[i] for i in range(3) for j in range(3)] + f0 + f1 + f2


    

In [ ]:

    

odesys = SymbolicSys.from_callback(dydt, 18, 3)


    

In [ ]:

    

y0  = [0, 0, 0, 0, 0, 1, 0, 1, 0] + [0]*9
res = odesys.integrate(15, y0, [2, 3, 4], integrator='cvode', method='adams', nsteps=15000)
print(res.info['time_cpu'], res.info['time_wall'])


    

In [ ]:

    

_ = res.plot(title_info=1)


    

In [ ]:

    

fig1 = plt.figure()
ax = fig1.gca(projection='3d')
ax.plot(*res.yout[:, 0:3].T, alpha=1.0, color='r')
ax.plot(*res.yout[:, 3:6].T, alpha=1.0, color='g')
ax.plot(*res.yout[:, 6:9].T, alpha=1.0, color='b')


    

In [ ]:

    

import numpy as np
import time
from pyodesys.native import native_sys
nativesys = native_sys['cvode'].from_other(odesys)


    

In [ ]:

    

nlvls = 300
params_varied = np.tile([2,3,4], (nlvls, 1))
params_varied[:, 0] = np.linspace(2, 5, nlvls)
t0 = time.time()
results = nativesys.integrate(15, y0, params_varied, integrator='native', method='adams', nsteps=15000)
print(sum([r.info['time_cpu'] for r in results]), time.time() - t0)