1import json
 2import numpy as np
 3from scipy.optimize import root
 4from experiment import E, I, Iinv, loglog_slope
 5
 6z0 = np.r_[np.array([0.2, -0.1, 0.3]), np.array([0.7, 1.1, 1.5])]
 7
 8def grad_h(z):
 9    q, p = z[:3], z[3:]
10    return np.r_[0.12*np.sin(q), Iinv*p]
11
12def field(z):
13    g = grad_h(z)
14    coad = np.einsum('a,aij,j->i', z[3:], E, g[3:])
15    return np.r_[g[3:], -g[:3] + coad]
16
17def step(z, dt):
18    def r(x):
19        m = (z+x)/2
20        return x-z-dt*field(m)
21    sol = root(r, z+dt*field(z), method='hybr')
22    assert sol.success and np.max(np.abs(r(sol.x))) < 1e-9
23    return sol.x
24
25def rollout(dt, T=5.0):
26    z=z0.copy()
27    for _ in range(round(T/dt)): z=step(z,dt)
28    return z
29
30# Fine reference and successively refined midpoint trajectories.
31ref=rollout(0.0005)
32dts=np.array([0.08,0.04,0.02,0.01])
33errs=np.array([np.linalg.norm(rollout(float(dt))-ref) for dt in dts])
34slope=loglog_slope(dts, errs)
35out={'dts':dts.tolist(),'errors_vs_fine_reference':errs.tolist(),
36     'predicted_global_error_exponent':2.0,'observed_global_error_exponent':float(slope)}
37open('nonlinear_results.json','w').write(json.dumps(out,indent=2))
38print(json.dumps(out,indent=2))