Lagrange propagation and the state transition matrix#
In the gravitational two-body problem it is possible to compute the state of the system at an arbitrary time from an initial state \(\left( \boldsymbol{r}_0, \boldsymbol{v}_0 \right)\) via the so-called Lagrange coefficients \(F\), \(G\), \(F_t\) and \(G_t\):
Analytical expressions for these coefficients are available in terms of anomalies differences between the end state and the initial state. The details on these analytical expressions and their derivation can be found, for example, in [Bat87].
In this notebook, we will first show how to use the Lagrange propagator provided by heyoka.py. Because the propagator is implemented in terms of analytical formulae, we will then be able to differentiate it and effortlessly construct the state transtition matrix (i.e., the Jacobian of the propagator with respect to the initial state).
The Lagrange propagator#
Note
The Lagrange propagator presented here is currently limited to elliptic orbits.
heyoka.py’s Lagrange propagator is available as heyoka.model.lagrange_prop(). In order to instantiate it, we first need to introduce expressions representing:
the initial Cartesian position \(\boldsymbol{r}_0=\left( x_0, y_0, z_0 \right)\),
the initial Cartesian velocity \(\boldsymbol{v}_0=\left( v_{x0}, v_{y0}, v_{z0} \right)\),
the gravitational parameter of the system \(\mu\),
the propagation time \(t\).
import heyoka as hy
import numpy as np
x0, y0, z0 = hy.make_vars("x0", "y0", "z0")
vx0, vy0, vz0 = hy.make_vars("vx0", "vy0", "vz0")
mu, tm = hy.make_vars("mu", "t")
We can now proceed to building the propagator, which consists of expressions representing the propagated position and velocity vectors:
pos, vel = hy.model.lagrange_prop(pos0=[x0, y0, z0], vel0=[vx0, vy0, vz0], mu=mu, tm=tm)
# Concatenate position and velocity
# into a single state vector.
pos_vel = np.hstack([pos, vel])
Let us now build a compiled function for the evaluation of the state vector at time \(t\):
cf = hy.cfunc(
pos_vel,
# Specify the order in which the input
# variables are passed to the compiled
# function.
vars=[x0, y0, z0, vx0, vy0, vz0, mu, tm],
)
We can now run a quick test for our propagator. We set up a circular orbit with \(\boldsymbol{r}_0=\left( 1, 0, 0 \right)\), \(\boldsymbol{v}_0=\left( 0, 1, 0 \right)\) and \(\mu = 1\), and we ask for the state vector at \(t=\pi\) (i.e., half period):
cf(
[
# r0.
1.0,
0.0,
0.0,
# v0.
0.0,
1.0,
0.0,
# mu and t.
1.0,
np.pi,
]
)
array([-1.0000000e+00, 1.2246468e-16, 0.0000000e+00, -1.2246468e-16,
-1.0000000e+00, -0.0000000e+00])
Indeed, as expected, \(\boldsymbol{r}\left(\pi\right) = \left( -1, 0, 0 \right)\) and \(\boldsymbol{v}\left(\pi\right) = \left( 0, -1, 0 \right)\) (plus/minus epsilon).
Recall from the compiled functions tutorial that the propagator is fully vectorised, and that it takes advantage of SIMD instructions. For instance, on a modern x86 machine we can propagate four different trajectories at the cost of one single propagation by passing in a two-dimensional matrix of initial conditions like this:
cf(
[
# x0.
[1.0, 1.01, 1.02, 1.03],
# y0.
[0.0, 0.01, 0.02, 0.03],
# z0.
[0.0, 0.01, 0.02, 0.03],
# vx0.
[0.0, 0.01, 0.02, 0.03],
# vy0.
[1.0, 1.01, 1.02, 1.03],
# vz0.
[0.0, 0.01, 0.02, 0.03],
# mu.
[1.0, 1.01, 1.02, 1.03],
# t.
[np.pi, np.pi + 0.01, np.pi + 0.02, np.pi + 0.03],
]
)
array([[-1.00000000e+00, -1.03804031e+00, -1.05314825e+00,
-1.04768800e+00],
[ 1.22464680e-16, 1.88508928e-01, 3.79275881e-01,
5.68044550e-01],
[ 0.00000000e+00, -8.32873901e-03, -1.29590841e-02,
-1.35748147e-02],
[-1.22464680e-16, -1.56913323e-01, -2.92251543e-01,
-4.03032933e-01],
[-1.00000000e+00, -9.54125221e-01, -8.82265186e-01,
-7.93231703e-01],
[-0.00000000e+00, -1.08925347e-02, -2.25868602e-02,
-3.38565463e-02]])
Constructing the STM#
We can now proceed to the construction of an analytical expression for the state transition matrix (STM).
The STM is nothing but the Jacobian of the Lagrange propagator with respect to the initial conditions. In order to compute the derivatives, we employ the diff_tensors() function as explained here:
dt = hy.diff_tensors(pos_vel, diff_args=[x0, y0, z0, vx0, vy0, vz0], diff_order=1)
We can then extract the Jacobian from dt,
jac = dt.jacobian
and proceed to the creation of a compiled function for the numerical evaluation of the STM:
cf_stm = hy.cfunc(
jac.flatten(),
# Specify the order in which the input
# variables are passed to the compiled
# function.
vars=[x0, y0, z0, vx0, vy0, vz0, mu, tm],
)
Let us take a look at the STM for our test circular orbit:
cf_stm([1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0, np.pi]).reshape((6, -1))
array([[-3.00000000e+00, 2.44929360e-16, 0.00000000e+00,
3.67394040e-16, -4.00000000e+00, 0.00000000e+00],
[ 9.42477796e+00, 3.00000000e+00, 0.00000000e+00,
4.00000000e+00, 9.42477796e+00, 0.00000000e+00],
[ 0.00000000e+00, 0.00000000e+00, -1.00000000e+00,
0.00000000e+00, 0.00000000e+00, 1.22464680e-16],
[-9.42477796e+00, -2.00000000e+00, 0.00000000e+00,
-3.00000000e+00, -9.42477796e+00, 0.00000000e+00],
[ 2.00000000e+00, 3.67394040e-16, 0.00000000e+00,
4.89858720e-16, 3.00000000e+00, 0.00000000e+00],
[ 0.00000000e+00, 0.00000000e+00, -1.22464680e-16,
0.00000000e+00, 0.00000000e+00, -1.00000000e+00]])
We can notice that the top-left element of the STM is the value \(-3\). If we slightly perturb by \(10^{-5}\) the value of \(x_0\) and re-evaluate the state of the system at \(t=\pi\),
cf([1 + 1e-5, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0, np.pi])
array([-1.00003000e+00, 9.42473082e-05, 0.00000000e+00, -9.42435384e-05,
-9.99979996e-01, -0.00000000e+00])
we can indeed see that the the initial perturbation of \(10^{-5}\) has been amplified by a factor of \(3\) in the final state of \(x\) as predicted by the STM.
Like the Lagrange propagator, the compiled function for the STM is also fully vectorised and SIMD-enabled. Note also that it is also possible via diff_tensors() to compute not only the Jacobian, but also higher-order tensors of derivatives (e.g., such as the Hessians).