fastrk 0.0.5

FastRK, a generator of fast jit-compiled code for ODE propagation by ERK methods with adaptive step and events

FastRK

• developed as fast alternative for subset of scipy.integrate.ode methods (i.e. DOP853);
• is a python code generator for Ordinary Differential Equations (ODE) propagation;
• uses explicit embedded Runge-Kutta (ERK) methods with adaptive step technique;
• calculates events using event functions (like scipy.integrate.solve_ivp);
• is jit-compiled by numba;
  • compiled code cached on SSD/HDD to prevent unnecessary recompilation;
• reentry, i.e. can be used in multithreaded applications;
• OS-independent (to the same extent as numba);
• contains Butcher Tables for several ERK methods:
  • Dormand and Prince 6(5)8M;
  • Dormand and Prince 8(7)13M;
    • 1.5x - 4.5x faster than DOP853 from scipy.integrate.ode
  • Verner's 8(9)16;
• user-defined Butcher Tables also supported;
• generated code is open and user-modifiable;

Butcher Tables was adapted from TrackerComponentLibrary.

Installation

pip install fastrk

Fast example

import numpy as np
import matplotlib.pyplot as plt
from fastrk import BT8713M, RKCodeGen

# CRTBP ODE, https://github.com/BoberSA/fastrk/tree/master/examples/model_crtbp.py
from model_crtbp import crtbp

rk_module = RKCodeGen(BT8713M, autonomous=True).save_and_import()
rk_prop = rk_module.rk_prop

t0, t1 = 0., 3 * np.pi
# initial state for halo orbit
s0 = np.zeros(6)
s0[[0, 2, 4]] = 9.949942666080747733e-01, 4.732924802139452415e-03, -1.973768492871211949e-02
mc = np.array([3.001348389698916e-06])  # CRTBP constant
rtol, atol = 1e-12, 1e-12

# integrate CRTBP ODE from t0, s0 to t1
arr = rk_prop(crtbp, s0, t0, t1, np.inf, rtol, atol, mc)

plt.plot(arr[:, 1], arr[:, 2], 'r')
plt.axis('equal')
plt.show()

Performance measurement

import numpy as np
import pandas as pd
from scipy.integrate import ode
from numba import jit
from timeit import timeit
from fastrk import BT8713M, RKCodeGen, default_jitkwargs
from model_crtbp import crtbp
rk_prop = RKCodeGen(BT8713M, autonomous=False).save_and_import().rk_prop
#%%
# scipy.integrate.ode's fortran implementation of DOP853 works
# significantly faster with @jit-compiled function rather than @cfunc-compiled
# default_jitkwargs = {'nopython': True, 'nogil': True, 'fastmath': True, 'cache': True}
jit_crtbp = jit(**default_jitkwargs)(crtbp._pyfunc).compile('f8[:](f8, f8[:], f8[:])')

#%% integration parameters
rtol, atol = 1e-12, 1e-12
max_step = np.inf

#%% CTRBP constant
mc = np.array([3.001348389698916e-06])  # CRTBP constant

#%% initial states
s0 = [0.9919060293647325, 0., 0.0016194537148125807, 0., -0.010581643111837302, 0.]
s1 = [0.9966271059324971, 0., 0.0050402173579027045, 0., -0.024398902561093703, 0.]
s2 = [0.4999857344807682, -0.866005893551121, 0., 3.902066111769351e-05, 2.252789194673211e-05, 0.]
s3 = [0.4966615415563801, -0.8602481879501589, 0., 0.011597147217611577, 0.0066415463209149195, 0.]

#%% tests
#         name                 t          initial state
tests = {'halo (small)':      (4*np.pi,   np.array(s0)),
         'halo (big)':        (4*np.pi,   np.array(s1)),
         'stable L5 (small)': (100*np.pi, np.array(s2)),
         'stable L5 (big)':   (100*np.pi, np.array(s3)),
         }
#%%
# define DOP853 integrator with same parameters
dop853_prop = ode(jit_crtbp)
dop853_prop.set_integrator('DOP853', max_step=np.inf, rtol=rtol, atol=atol, nsteps=100000)
dop853_prop.set_f_params(mc)

# to retrieve all integrator steps a callback with side effect needed
lst = []
def solout(t, s):
    lst.append([t, *s])
dop853_prop.set_solout(solout)

#%% measure execution time

loops = 1000
res = []
for i, name in enumerate(tests):
    print(name)
    t, s = tests[name]
    lst = []
    dop853_prop.set_initial_value(s, 0.).integrate(t)
    steps0 = len(lst)
    steps1 = rk_prop(crtbp, s, 0., t, max_step, rtol, atol, mc).shape[0]
    r0 = timeit("dop853_prop.set_initial_value(s, 0.).integrate(t)",
                number=loops, globals=globals())
    r1 = timeit("rk_prop(crtbp, s, 0., t, max_step, rtol, atol, mc)",
                number=loops, globals=globals())
    res.append([t, steps0, r0, steps1, r1])

#%% print results

columns = pd.MultiIndex.from_tuples([('integration', 'time'),
                                     ('dop853', 'steps'),
                                     ('dop853', 'time'),
                                     ('fastrk', 'steps'),
                                     ('fastrk', 'time')], names=['', ''])
df = pd.DataFrame(res, columns=columns, index=tests.keys())
df['speedup'] = df[('dop853', 'time')] / df[('fastrk', 'time')]
print(df)

Output for AMD Ryzen 7 4700U @ 4GHz:

                  integration dop853            fastrk             speedup
                         time  steps       time  steps      time          
halo (small)        12.566371     96   1.422070    235  0.964003  1.475171
halo (big)          12.566371    145   2.212505    151  0.611676  3.617117
stable L5 (small)  314.159265    452   5.898220    422  1.333944  4.421639
stable L5 (big)    314.159265    916  12.021953    856  2.701954  4.449355

Detailed examples

Example 0: Propagate spacecraft motion in Circular Restricted Three Body Problem

Example 1: Calculate events in Circular Restricted Three Body Problem

Required modules:

• model_crtbp.py

Parallel (OpenMP) example

Required modules:

• model_crtbp.py
• map_vy_funcs.py

Output for AMD Ryzen 7 4700U @ 4GHz:

    compiling sequential...
    compiling parallel...
    states count: 2500
    sequential calculation started (Ms)
    map calculation time: 46.31 s
    parallel calculation started (Mp)
    map calculation time: 22.19 s
    speedup x 2.09
    ||Ms - Mp|| = 0.0

Core Developer

Stanislav Bober, MIEM NRU HSE, IKI RAS