efficiency.py

"""Build a small error-versus-work plot from existing integration calls.

Collects measured errors, saves the full cost table, and plots one cost
component. The nested ``record()`` helper stores results; it is not an
integrator.

Usage
-----
python -m examples.efficiency

Model
-----
Two-body orbit, eccentricity 0.6, inclination 0.4 rad, nine output times over
half a period. The reference is analytic Kepler propagation; IAS15 is plotted
as a measured curve, not as truth. Error is the maximum absolute barycentric
position error over bodies and times.

Notes
-----
Resolutions 80, 160 and 320 set both dt = T/resolution and
rtol = 1e-6 * (80/resolution)**4, so this is neither a pure convergence test
nor an equal-accuracy benchmark. Forced output can cap adaptive steps. The
force axis omits jerk, Kepler and linear-algebra work, which remain in the
CSV. Overwrites results/efficiency.png and results/efficiency_costs.csv in
the current directory; opens no window. Requires NumPy, Matplotlib, REBOUND.
"""
import csv
from pathlib import Path
import numpy as np
from astrosimbench import integrate,METHODS
from astrosimbench.systems import eccentric_two_body,exact_two_body_states
from astrosimbench.references import rebound_reference
from astrosimbench.plotting import efficiency_plot


def main():
    r,v,m,T=eccentric_two_body(eccentricity=.6,inclination=.4)
    times=np.linspace(0,.5*T,9)
    exact,_=exact_two_body_states(r,v,m,times)
    curves={};rows=[]
    def record(out,setting):
        error=float(np.max(np.linalg.norm(out.positions-exact,axis=-1)))
        rows.append(dict(method=out.method,setting=setting,error=error,**out.cost.to_dict()))
        return out.cost.force_evaluations,error
    for name,spec in METHODS.items():
        points=[]
        for resolution in (80,160,320):
            tol=1e-6*(80/resolution)**4
            out=integrate(name,r,v,m,t_span=(0,times[-1]),dt=T/resolution,
                          rtol=tol,atol=tol*1e-3,sample_times=times)
            points.append(record(out,f'dt=T/{resolution};rtol={tol}'))
        curves[spec.name]=np.array(points).T
    points=[]
    for epsilon in (1e-5,1e-7,1e-9):
        out=rebound_reference(r,v,m,times,epsilon=epsilon)
        points.append(record(out,f'epsilon={epsilon}'))
    fig,ax=efficiency_plot(curves,ias15_curve=np.array(points).T,
        cost_label='Force-field evaluations (jerk / Kepler / linear algebra costs excluded)',
        error_label='Maximum position error against analytic Kepler solution')
    ax.set_title('Two-body illustration: WH/IAS15 may reach the numerical floor')
    fig.tight_layout()
    folder=Path('results');folder.mkdir(exist_ok=True)
    fig.savefig(folder/'efficiency.png',dpi=160)
    with (folder/'efficiency_costs.csv').open('w',newline='') as f:
        writer=csv.DictWriter(f,fieldnames=list(rows[0]))
        writer.writeheader();writer.writerows(rows)
    print('Saved results/efficiency.png and results/efficiency_costs.csv')

if __name__=='__main__': main()