For a complete introduction to variational equations, please read the paper by Rein and Tamayo (2016).
Variational equations can be used to calculate derivatives in an $N$-body simulation. More specifically, given a set of initial conditions $\alpha_i$ and a set of variables at the end of the simulation $v_k$, we can calculate all first order derivatives $$\frac{\partial v_k}{\partial \alpha_i}$$ as well as all second order derivates $$\frac{\partial^2 v_k}{\partial \alpha_i\partial \alpha_j}$$
For this tutorial, we work with a two planet system.
We first chose the semi-major axis $a$ of the outer planet as an initial condition (this is our $\alpha_i$). At the end of the simulation we output the velocity of the star in the $x$ direction (this is our $v_k$).
To do that, let us first import REBOUND and numpy.
In [1]:
import rebound
import numpy as np
The following function takes $a$ as a parameter, then integrates the two planet system and returns the velocity of the star at the end of the simulation.
In [2]:
def calculate_vx(a):
sim = rebound.Simulation()
sim.add(m=1.) # star
sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet
sim.add(primary=sim.particles[0],m=1e-3, a=a) # outer planet
sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits
return sim.particles[0].vx # return star's velocity in the x direction
In [3]:
calculate_vx(a=1.5) # initial semi-major axis of the outer planet is 1.5
Out[3]:
If we run the simulation again, with a different initial $a$, we get a different velocity:
In [4]:
calculate_vx(a=1.51) # initial semi-major axis of the outer planet is 1.51
Out[4]:
We could now run many different simulations to map out the parameter space. This is a very simple examlpe of a typical use case: the fitting of a radial velocity datapoint.
However, we can be smarter than simple running an almost identical simulation over and over again by using variational equations. These will allow us to calculate the derivate of the stellar velocity at the end of the simulation. We can take derivative with respect to any of the initial conditions, i.e. a particles's mass, semi-major axis, x-coordinate, etc. Here, we want to take the derivative with respect to the semi-major axis of the outer planet. The following function does exactly that:
In [5]:
def calculate_vx_derivative(a):
sim = rebound.Simulation()
sim.add(m=1.) # star
sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet
sim.add(primary=sim.particles[0],m=1e-3, a=a) # outer planet
v1 = sim.add_variation() # add a set of variational particles
v1.vary(2,"a") # initialize the variational particles
sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits
return sim.particles[0].vx, v1.particles[0].vx # return star's velocity and its derivative
Note the two new functions. sim.add_variation() adds a set of variational particles to the simulation. All variational particles are by default initialized to zero. We use the vary() function to initialize them to a variation that we are interested in. Here, we initialize the variational particles corresponding to a change in the semi-major axis, $a$, of the particle with index 2 (the outer planet).
In [6]:
calculate_vx_derivative(a=1.5)
Out[6]:
We can use the derivative to construct a Taylor series expansion of the velocity around $a_0=1.5$: $$v(a) \approx v(a_0) + (a-a_0) \frac{\partial v}{\partial a}$$
In [7]:
a0=1.5
va0, dva0 = calculate_vx_derivative(a=a0)
def v(a):
return va0 + (a-a0)*dva0
print(v(1.51))
Compare this value with the explicitly calculate one above. They are almost the same! But we can do even better, by using second order variational equations to calculate second order derivatives.
In [8]:
def calculate_vx_derivative_2ndorder(a):
sim = rebound.Simulation()
sim.add(m=1.) # star
sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet
sim.add(primary=sim.particles[0],m=1e-3, a=a) # outer planet
v1 = sim.add_variation()
v1.vary(2,"a")
# The following lines add and initialize second order variational particles
v2 = sim.add_variation(order=2, first_order=v1)
v2.vary(2,"a")
sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits
# return star's velocity and its first and second derivatives
return sim.particles[0].vx, v1.particles[0].vx, v2.particles[0].vx
Using a Taylor series expansion to second order gives a better estimate of v(1.51).
In [9]:
a0=1.5
va0, dva0, ddva0 = calculate_vx_derivative_2ndorder(a=a0)
def v(a):
return va0 + (a-a0)*dva0 + 0.5*(a-a0)**2*ddva0
print(v(1.51))
Now that we know how to calculate first and second order derivates of positions and velocities of particles, we can simply use the chain rule to calculate more complicated derivates. For example, instead of the velocity $v_x$, you might be interested in the quanity $w\equiv(v_x - c)^2$ where $c$ is a constant. This is something that typically appears in a $\chi^2$ fit. The chain rule gives us: $$ \frac{\partial w}{\partial a} = 2 \cdot (v_x-c)\cdot \frac{\partial v_x}{\partial a}$$ The variational equations provide the $\frac{\partial v_x}{\partial a}$ part, the ordinary particles provide $v_x$.
In [10]:
def calculate_w_derivative(a):
sim = rebound.Simulation()
sim.add(m=1.) # star
sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet
sim.add(primary=sim.particles[0],m=1e-3, a=a) # outer planet
v1 = sim.add_variation() # add a set of variational particles
v1.vary(2,"a") # initialize the variational particles
sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits
c = 1.02 # some constant
w = (sim.particles[0].vx-c)**2
dwda = 2.*v1.particles[0].vx * (sim.particles[0].vx-c)
return w, dwda # return w and its derivative
calculate_w_derivative(1.5)
Out[10]:
Similarly, you can also use the chain rule to vary initial conditions of particles in a way that is not supported by REBOUND by default. For example, suppose you want to work in some fancy coordinate system, using $h\equiv e\sin(\omega)$ and $k\equiv e \cos(\omega)$ variables instead of $e$ and $\omega$. You might want to do that because $h$ and $k$ variables are often better behaved near $e\sim0$. In that case the chain rule gives us: $$\frac{\partial p(e(h, k), \omega(h, k))}{\partial h} = \frac{\partial p}{\partial e}\frac{\partial e}{\partial h} + \frac{\partial p}{\partial \omega}\frac{\partial \omega}{\partial h}$$ where $p$ is any of the particles initial coordinates. In our case the derivates of $e$ and $\omega$ with respect to $h$ are: $$\frac{\partial \omega}{\partial h} = -\frac{k}{e^2}\quad\text{and}\quad \frac{\partial e}{\partial h} = \frac{h}{e}$$
With REBOUND, you can easily implement this. The following function calculates the derivate of the star's velocity with respect to the outer planet's $h$ variable.
In [12]:
def calculate_vx_derivative_h():
h, k = 0.1, 0.2
e = float(np.sqrt(h**2+k**2))
omega = np.arctan2(k,h)
sim = rebound.Simulation()
sim.add(m=1.) # star
sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet
sim.add(primary=sim.particles[0],m=1e-3, a=1.5, e=e, omega=omega) # outer planet
v1 = sim.add_variation()
dpde = rebound.Particle(simulation=sim, particle=sim.particles[2], variation="e")
dpdomega = rebound.Particle(simulation=sim, particle=sim.particles[2], m=1e-3, a=1.5, e=e, omega=omega, variation="omega")
v1.particles[2] = h/e * dpde - k/(e*e) * dpdomega
sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits
# return star's velocity and its first derivatives
return sim.particles[0].vx, v1.particles[0].vx
calculate_vx_derivative_h()
Out[12]:
Note that in the above function, there are expressions such as h/e * dpde. h/e is just a number, but dpde is actually a particle structure. REBOUND multiplies each cartesian component of that particle with the number h/e. Similarly, the particles are subtracted componentwise when using the - operator.
We can use the v1.particles[i] = ... syntax to directly set a variational particle's initial conditions.
In [ ]: