import rebound
import numpy as np
def test_case():
    sim = rebound.Simulation()
    sim.integrator = 'whfast'
    sim.add(m=1.) # add the Sun
    sim.add(m=3.e-6,e=0.99, a=1.) # add Earth
    sim.move_to_com()
    sim.dt = 0.2
    return sim

"""
Explanation: Advanced settings for WHFast:  Extra speed, accuracy, and additional forces
There are several performance enhancements one can make to WHFast.  However, each one has pitfalls that an inexperienced user can unwittingly fall into.  We therefore chose safe default settings that make the integrator difficult to misuse.  This makes the default WHFast substantially slower and less accurate than it can be.  Here we describe how to alter the integrator settings to improve WHFast's performance.
TL;DR
As long as 

you don't add, remove or otherwise modify particles between timesteps
you get your outputs by passing a list of output times ahead of time and access the particles pointer between calls to sim.integrate() (see, e.g., the Visualization section of WHFast.ipynb)

you can set sim.ri_whfast.safe_mode = 0 to get a substantial performance boost.  Under the same stipulations, you can set sim.ri_whfast.corrector = 11 to get much higher accuracy, at a nearly negligible loss of performance (as long as there are many timesteps between outputs).
If you want to modify particles, or if the code breaks with these advanced settings, read below for details, and check out the Common mistake with WHFast section at the bottom of WHFast.ipynb.
The Wisdom-Holman algorithm
In order to understand and apply the various integrator flags, we need to first understand the Wisdom-Holman scheme (see, e.g., Wisdom & Holman 1991, or Rein & Tamayo 2015 for more details).
The Wisdom-Holman algorithm consists of alternating Keplerian steps that evolve particles on their two-body Keplerian orbits around the star with interaction steps that apply impulses to the particles' velocities from the interactions between bodies.  The basic algorithm for a single timestep $dt$ is a Leapfrog Drift-Kick-Drift scheme with an interaction kick over the full $dt$ sandwiched between half timesteps of Keplerian drift:
$H_{Kepler}(dt/2)\:H_{Interaction}(dt)\:H_{Kepler}(dt/2)$
Timesteps like the one above are then concatenated over the full integration:
$H_{Kepler}(dt/2)\:H_{Interaction}(dt)\:H_{Kepler}(dt/2)$ $H_{Kepler}(dt/2)\:H_{Interaction}(dt)\:H_{Kepler}(dt/2)$ ... $H_{Kepler}(dt/2)\:H_{Interaction}(dt)\:H_{Kepler}(dt/2)$
Combining Kepler steps and synchronizing
It turns out that Kepler steps take longer than interaction steps as long as you don't have many planets, so an obvious and important performance boost would be to combine adjacent Kepler half-steps into full ones, i.e.:
$H_{Kepler}(dt/2)\:H_{Interaction}(dt)\:H_{Kepler}(dt)\:H_{Interaction}(dt)\:H_{Kepler}(dt) ... \:H_{Interaction}(dt)\:H_{Kepler}(dt/2)$
The issue is that if you were to, say, output the state of the particles as the simulation progressed, the positions would not correspond to anything real, since the beginning (or end) of one of the full $H_{Kepler}(dt)$ steps corresponds to some intermediate step in an abstract sequence of calculations for a given timestep.  In order to get the particles' actual positions, we would have to calculate to the end the timestep we want the output for by splitting a full Kepler step back into two half-steps, e.g.,
$H_{Kepler}(dt/2)\:H_{Interaction}(dt)\:H_{Kepler}(dt)\:H_{Interaction}(dt)\:H_{Kepler}(dt/2) \text{PRINT OUTPUT} H_{Kepler}(dt/2) H_{Interaction}(dt)\:H_{Kepler}(dt)$...
We call this step of reinserting half-Kepler steps to obtain the physical state of the particles synchronizing.  This must be done whenever the actual states of the particles are required, e.g., before every output, or if one wanted to use the particles' states to compute additional changes to the particle orbits between timesteps.  It is also necessary to synchronize each timestep whenever the MEGNO chaos indicator is being computed.
Conversions between Jacobi and Inertial Coordinates
It turns out that the most convenient coordinate system to work in for performing the Kepler steps is often Jacobi coordinates (see, e.g., 9.5.4 of Murray & Dermott).  WHFast therefore works in Jacobi coordinates by default, converting to inertial coordinates when it needs to (e.g. for output, and for doing the direct gravity calculation in the interaction step, which is most easily done in inertial coordinates).
One feature of WHFast is that it works in whatever inertial coordinate system you choose for your initial conditions.  This means that whatever happens behind the scenes, the user always gets the particles' inertial coordinates at the front end.  At the beginning of every timestep, WHFast therefore has to somehow obtain the Jacobi coordinates.  The straightforward thing would be to convert from the inertial coordinates to Jacobi coordinates every timestep, but these conversions slow things down, and they represent extra operations that grow the round-off error.
WHFast therefore stores the Jacobi coordinates internally throughout the time it is running, and only recalculates Jacobi coordinates from the inertial ones if told to do so.  Since Jacobi coordinates reference particles to the center of mass of all the particles with indices lower than their own (typically all the particles interior to them), the main reason you would have to recalculate Jacobi coordinates is if between timesteps you choose to somehow change the particles' positions or velocities (give them kicks in addition to their mutual gravity), or change the particles' masses.  
Overriding the defaults
Let's begin by importing rebound, and defining a simple function to reset rebound and initialize a new simulation with a test case,
End of explanation
"""


sim = test_case()
sim.ri_whfast.safe_mode = 0

"""
Explanation: By default WHFast synchronizes and recalculates the Jacobi coordinates from the inertial ones every timestep.  This guarantees that the user always gets physical particle states for output, and ensures reliable output if the user decides to, e.g., grow the particles' masses between timesteps.  
Now that you understand the pitfalls, if you want to boost WHFast's performance, you simply set
End of explanation
"""


import time
Porb = 2*np.pi # orbital period for Earth, using units of G = 1, solar masses, AU and yr/2pi

sim = test_case()
print("safe_mode = {0}".format(sim.ri_whfast.safe_mode))
start_time = time.time()
sim.integrate(1.e5*Porb)
sim.status()
print("Safe integration took {0} seconds".format(time.time() - start_time))

sim = test_case()
sim.ri_whfast.safe_mode = 0
start_time = time.time()
sim.integrate(1.e5*Porb)
sim.status()
print("Manual integration took {0} seconds".format(time.time() - start_time))

"""
Explanation: Now it becomes the user's responsibility to appropriately synchronize and recalculate jacobi coordinates when needed.  You can tell WHFast to recalculate Jacobi coordinates for a given timestep (say after you change a particle's mass) with the sim.ri_whfast.recalculate_coordinates_this_timestep flag.  After it recalculates Jacobi coordinates, WHFast will reset this flag to zero, so you just set it each time you mess with the particles.
End of explanation
"""


sim.init_megno()

"""
Explanation: In our test case with a single planet, there is effectively no interaction step, and by combining Kepler steps we get almost the full factor of 2 speedup we expect.  Because Kepler steps are expensive (by virtue of having to solve the transcendental Kepler equation), this will always be an important performance boost for few-planet cases.
Note that one case where REBOUND needs to synchronize every timestep is if you're using the MEGNO chaos indicator.  So if you call
End of explanation
"""


sim = test_case()
sim.ri_whfast.safe_mode = 0
def integrate_mod(sim, t_final):
    while sim.t < t_final:
        sim.step()
        sim.particles[1].m += 1.e-10
        sim.ri_whfast.recalculate_coordinates_this_timestep = 1
    sim.integrator_synchronize()

"""
Explanation: REBOUND will synchronize every timestep even if you set sim.ri_whfast.safe_mode = 0 and never explicitly call sim.integrator_synchronize().
Modifying particles/forces
Again, if performance is a factor in your simulations, you would not want to write a custom stepper in python that modifies the particles, since this will be very slow.  You could either write a modified C version of reb_integrate in src/librebound.c (the flags are defined in librebound.h, and have the same name as the python ones, just without sim. in front), or you can use the REBOUNDXF library, which takes care of this for you and supports many typically used modifications.  We again illustrate a simple scheme with python code:
End of explanation
"""


sim = test_case()
sim.ri_whfast.safe_mode = 0
def integrate_mod(sim, t_final):
    while sim.t < t_final:
        sim.step()
        sim.particles[1].vx += 1.e-10*sim.dt
        sim.ri_whfast.recalculate_coordinates_this_timestep = 1
    sim.integrator_synchronize()

"""
Explanation: Here, because we grow the mass of the planet every timestep, we have to recalculate Jacobi coordinates every timestep (since they depend on the masses of the particles).  We therefore manually set the flag to recalculate them the next timestep every time we make a change.  Here we would actually get the same result if we just left sim.ri_whfast.safe_mode = 1, since when recalculating Jacobi coordinates, WHFast automatically has to synchronize in order to get real positions and velocities for the planets.  In this case WHFast is therefore synchronizing and recalculating Jacobi coordinates every timestep.
But imagine now that instead of growing the mass, we continually add an impulse to vx:
End of explanation
"""


sim.ri_whfast.corrector = 11

"""
Explanation: This would not give accurate results, because the sim.particles[1].vx we access after sim.step() isn't a physical velocity (it's missing a half-Kepler step).  It's basically at an intermediate point in the calculation.  In order to make this work, one would call sim.integrator_synchronize() between sim.step() and accessing sim.particles[1].vx, to ensure the velocity is physical.
Symplectic correctors
Symplectic correctors make the Wisdom-Holman scheme higher order (without symplectic correctors it's second order).  The great thing about them is that they only need to get applied when you synchronize.  So if you just need to synchronize to output, and there are many timesteps between outputs, they represent a very small performance loss for a huge boost in accuracy (compare for example the green line (11th order corrector) to the red line (no corrector) in Fig. 4 of Rein & Tamayo 2015--beyond the right of the plot, where the round-off errors dominate, the two lines would rise in unison).  We have implemented symplectic correctors up to order 11.  You can set the order with (must be an odd number), e.g.,
End of explanation
"""


sim.ri_whfast.coordinates = 'jacobi' #default
sim.ri_whfast.coordinates = 'democraticheliocentric' 
sim.ri_whfast.coordinates = 'whds' 

"""
Explanation: By default, WHFast does not use correctors, i.e., sim.integrator_whfast_corrector = 0.  This is because the default is also to synchronize every timestep.  An Nth order corrector does N-1 Kepler steps of various sizes, so an 11th order corrector done every timestep would increase the number of Kepler steps by an order of magnitude, making WHFast unacceptably slow.  So keep in mind that if you're doing modifications that require recalculating jacobi coordinates or synchronizing every timestep, you should turn off symplectic correctors (the default) unless you really need the accuracy.
Changing the internal coordinate system
WHFast by default uses Jacobi coordinates internally. This works well for planetary systems which are stable and orbits are not crossing. However, in some cases a different coordinate system might perform better. WHFast also support so-called democratic heliocentric coordinates and the so called WHDS coordinates. For more information on these coordinates systems see Hernandez and Dehnen (2016). To select a different coordinate system, use the following syntax:
End of explanation
"""


sim = test_case()
sim.dt = 1000.
sim.integrate(1000.)

"""
Explanation: Note that symplectic corrector are only compatible with Jacobi coordinates because both democratic heliocentric and WHDS include a so called jump step.
Warning messages
If you choose a timestep that is larger than the smallest dynamical timescale and WHFast has difficulties to solve the Kepler problem, you will receive a warning message.
End of explanation
"""


