Interp orbit using agama between snapshots

orbit_interp.py

def interp_orbit(ix, snap_list, npts=51, rtol=1e-5):
    ns = len(snap_list) - 1
    ts = np.array([si.t for si in snap_list.values()])

    t_list = np.empty((ns, npts))
    xv_list = np.empty((len(ix), ns, npts, 6))

    for i, si in enumerate(snap_list.values()):
        print(i, end=',', flush=True)

        pot = si.pot_t
        ic = si.xv[ix]

        if i < ns:
            res = agama.orbit(potential=pot, ic=ic, timestart=ts[i], time=ts[i + 1] - ts[i],
                              trajsize=npts, accuracy=rtol)
            for j, (t, xv) in enumerate(res):
                assert len(t) == npts
                xv_list[j, i] = xv
            t_list[i] = t

        if i > 0:
            res = agama.orbit(potential=pot, ic=ic, timestart=ts[i], time=ts[i - 1] - ts[i],
                              trajsize=npts, accuracy=rtol)
            for j, (t, xv) in enumerate(res):
                assert len(t) == npts
                fac = (t[::-1] - ts[i - 1]).reshape(-1, 1) / (ts[i] - ts[i - 1])
                xv_list[j, i - 1] = xv_list[j, i - 1] * (1 - fac) + xv[::-1] * fac

    t_list = t_list.reshape(-1)
    xv_list = xv_list.reshape(len(ix), -1, 6)
    return ts, t_list, xv_list