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 the seminal book by Richard Battin “An introduction to the mathematics and methods of astrodynamics” (section 4.3). See here for another derivation.
In this notebook, we will first show how to implement a Lagrange propagator using heyoka.py’s expression system. Because the propagator will be 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 propagator presented here is implemented in terms of eccentric anomalies, and thus it is limited to elliptic orbits.
We begin by introducing the symbolic variables corresponding to the inputs of the propagator:
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")
# Package initial position and velocity into
# arrays for later use.
pos_0 = np.array([x0, y0, z0])
vel_0 = np.array([vx0, vy0, vz0])
Next, we compute the semi-major axis \(a\) from the specific orbital energy \(\epsilon\):
v02 = vx0**2+vy0**2+vz0**2
r0 = hy.sqrt(x0**2+y0**2+z0**2)
eps = v02 * 0.5 - mu/r0
a = -mu/(2.*eps)
Now we compute the quantities \(\sigma_0=\frac{\boldsymbol{r}_0 \cdot \boldsymbol{v}_0}{\sqrt{\mu}}\), \(s_0=\frac{\sigma_0}{\sqrt{a}}\) and \(c_0=1-\frac{r_0}{a}\):
sigma0 = np.dot(pos_0, vel_0) / hy.sqrt(mu)
s0 = sigma0 / hy.sqrt(a)
c0 = 1. - r0 / a
We can now calculate the difference in mean anomaly \(\Delta M\) from the mean motion \(n\) and the propagation time \(t\),
n = hy.sqrt(mu / (a * a * a))
DM = n*tm
and then proceed to convert it to a difference in eccentric anomaly \(\Delta E\) via the kepDE() function:
DE = hy.kepDE(s0, c0, DM)
# Compute cos(DE) and sin(DE).
cDE = hy.cos(DE)
sDE = hy.sin(DE)
We can now calculate \(r(t)\),
r = a+(r0-a)*cDE+sigma0*hy.sqrt(a)*sDE
and the Lagrange coefficients \(F\), \(G\), \(F_t\) and \(G_t\):
F = 1. - a/r0*(1. - cDE)
G = a*sigma0/hy.sqrt(mu)*(1. - cDE) + r0*hy.sqrt(a/mu)*sDE
Ft = -hy.sqrt(mu*a)/(r*r0)*sDE
Gt = 1-a/r*(1.-cDE)
Finally, we can calculate the position and velocity vectors at time \(t\):
pos = F*pos_0 + G*vel_0
vel = Ft*pos_0 + Gt*vel_0
# Concatenate position and velocity
# into a single state vector.
pos_vel = np.hstack([pos, vel])
We can now proceed to create 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])
[2024-12-17 22:51:34.819] [heyoka] [info] heyoka logger initialised
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.,
# v0.
0., 1., 0.,
# mu and t.
1., 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., 1.01, 1.02, 1.03],
# y0.
[0., 0.01, 0.02, 0.03],
# z0.
[0., 0.01, 0.02, 0.03],
# vx0.
[0., 0.01, 0.02, 0.03],
# vy0.
[1., 1.01, 1.02, 1.03],
# vz0.
[0., 0.01, 0.02, 0.03],
# mu.
[1., 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., 1., 0., 1., 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., 1., 0., 1., 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).