simple adaptive Heun integrator

ahoyn.py

def heun(y, t, dt):
    d1 = dfun(y, t)
    d2 = dfun(y + dt*d1, t + dt)
    err = np.mean(np.abs((d1 - d2)))#/(1e-9+d2)))
    return y + dt/2*(d1 + d2), err

def solve_adapt1(y0, ts, tol):
    ys = [y0]
    max_dt = dt = ts[1] - ts[0]
    t0 = ts[0]
    y = y0
    t = t0
    i = 0
    hu = []
    nfe = 0
    for t1 in ts[1:]:
        avg_dt = 0
        inner_steps = 0
        while t <= t1:
            i += 1
            dt = max(1e-8, min(dt, max_dt))
            if dt < 1e-13:
                raise RuntimeError(f'excessive work! dt={dt:0.3g}')
            ny, err = heun(y, t, dt)
            nfe += 2
            # print(i, t, dt, err)

            if err > tol:
                dt = dt * (tol/err)
                # don't accept
                continue

            # accept
            y = ny
            t = t + dt
            avg_dt += dt
            inner_steps += 1

            if err < tol:
                dt = dt * tol/err

        hu.append(dt / inner_steps)
        ys.append(y.copy())
    return np.array(ys), np.array(hu), nfe

ts = np.r_[0.0:10.0:256j]

for tol in (100.0, 10.0, 1.0):#, 1e-3):
    ys, hu, nfe = solve_adapt1(y0, ts, tol)
    # yt, infodict = odeint(dfun, y0, ts, atol=tol, rtol=tol, full_output=True)
    plt.subplot(211); plt.semilogy(ts[:-1], hu,
                                   label=f'tol{tol:0.1g} nfe {nfe}'); plt.legend()
    plt.subplot(212); plt.semilogy(ts, ys[:,0], label=f'heun')

yt = odeint(dfun, y0, ts, atol=1e-13, rtol=1e-13)
plt.semilogy(ts, yt[:,0], label='LSODA')
plt.legend()