Replace Curtis's iterator with the state transition matrix iterator (retrieves position to within 100 meters and velocity to within 10 cm/s)

Author: moeyensj

thor/orbits/iod/gauss.py

@@ -1,5 +1,6 @@
 import numpy as np
-from scipy import roots
+from numpy import roots
+from numba import jit
 
 from ...constants import Constants as c
 from ...coordinates import equatorialToEclipticCartesian
@@ -16,6 +17,8 @@
     "_calcLambdas",
     "_calcRhos",
     "_calcFG",
+    "_calcM",
+    "_calcStateTransitionMatrix",
     "_iterateGaussIOD",
     "calcGauss",
     "gaussIOD"
@@ -24,130 +27,207 @@
 MU = c.G * c.M_SUN
 C = c.C
 
-#@jit(["f8(f8[::1], f8[::1], f8[::1])"], nopython=True)
-def _calcV(rho1hat, rho2hat, rho3hat):
+def _calcV(rho1_hat, rho2_hat, rho3_hat):
     # Vector triple product that gives the area of 
     # the "volume of the parallelepiped" or according to 
     # to Milani et al. 2008: 3x volume of the pyramid with vertices q, r1, r2, r3.
     # Note that vector triple product rules apply here.
-    return np.dot(np.cross(rho1hat, rho2hat), rho3hat)
+    return np.dot(np.cross(rho1_hat, rho2_hat), rho3_hat)
 
-#@jit(["f8(f8[::1], f8[::1], f8[::1], f8[::1], f8[::1], f8, f8, f8)"], nopython=True)
-def _calcA(q1, q2, q3, rho1hat, rho3hat, t31, t32, t21):
+def _calcA(q1, q2, q3, rho1_hat, rho3_hat, t31, t32, t21):
     # Equation 21 from Milani et al. 2008
-    return np.linalg.norm(q2)**3 * np.dot(np.cross(rho1hat, rho3hat), (t32 * q1 - t31 * q2 + t21 * q3))
+    return np.linalg.norm(q2)**3 * np.dot(np.cross(rho1_hat, rho3_hat), (t32 * q1 - t31 * q2 + t21 * q3))
 
-#@jit(["f8(f8[::1], f8[::1], f8[::1], f8[::1], f8, f8, f8)"], nopython=True)
-def _calcB(q1, q3, rho1hat, rho3hat, t31, t32, t21):
+def _calcB(q1, q3, rho1_hat, rho3_hat, t31, t32, t21, mu=MU):
     # Equation 19 from Milani et al. 2008
-    return MU / 6 * t32 * t21 * np.dot(np.cross(rho1hat, rho3hat), ((t31 + t32) * q1 + (t31 + t21) * q3))
+    return mu / 6 * t32 * t21 * np.dot(np.cross(rho1_hat, rho3_hat), ((t31 + t32) * q1 + (t31 + t21) * q3))
 
-#@jit(["UniTuple(f8, 2)(f8, f8, f8, f8)"], nopython=True)
-def _calcLambdas(r2_mag, t31, t32, t21):
+def _calcLambdas(r2_mag, t31, t32, t21, mu=MU):
     # Equations 16 and 17 from Milani et al. 2008
-    lambda1 = t32 / t31 * (1 + MU / (6 * r2_mag**3) * (t31**2 - t32**2))
-    lambda3 = t21 / t31 * (1 + MU / (6 * r2_mag**3) * (t31**2 - t21**2))
+    lambda1 = t32 / t31 * (1 + mu / (6 * r2_mag**3) * (t31**2 - t32**2))
+    lambda3 = t21 / t31 * (1 + mu / (6 * r2_mag**3) * (t31**2 - t21**2))
     return lambda1, lambda3
 
-#@jit(["UniTuple(f8[::1], 3)(f8, f8, f8[::1], f8[::1], f8[::1], f8[::1], f8[::1], f8[::1], f8)"], nopython=True)
-def _calcRhos(lambda1, lambda3, q1, q2, q3, rho1hat, rho2hat, rho3hat, V):
+def _calcRhos(lambda1, lambda3, q1, q2, q3, rho1_hat, rho2_hat, rho3_hat, V):
     # This can be derived by taking a series of scalar products of the coplanarity condition equation
     # with cross products of unit vectors in the direction of the observer, in particular, see Chapter 9 in
     # Milani's book on the theory of orbit determination 
     numerator = -lambda1 * q1 + q2 - lambda3 * q3
-    rho1_mag = np.dot(numerator, np.cross(rho2hat, rho3hat)) / (lambda1 * V)
-    rho2_mag = np.dot(numerator, np.cross(rho1hat, rho3hat)) / V
-    rho3_mag = np.dot(numerator, np.cross(rho1hat, rho2hat)) / (lambda3 * V)
-    return np.dot(rho1_mag, rho1hat), np.dot(rho2_mag, rho2hat), np.dot(rho3_mag, rho3hat)
+    rho1_mag = np.dot(numerator, np.cross(rho2_hat, rho3_hat)) / (lambda1 * V)
+    rho2_mag = np.dot(numerator, np.cross(rho1_hat, rho3_hat)) / V
+    rho3_mag = np.dot(numerator, np.cross(rho1_hat, rho2_hat)) / (lambda3 * V)
+    rho1 = np.dot(rho1_mag, rho1_hat)
+    rho2 = np.dot(rho2_mag, rho2_hat)
+    rho3 = np.dot(rho3_mag, rho3_hat)
+    return rho1, rho2, rho3
 
-def _calcFG(r2_mag, t32, t21):
+def _calcFG(r2_mag, t32, t21, mu=MU):
     # Calculate the Lagrange coefficients (Gauss f and g series)
-    f1 = 1 - (1 / 2) * MU / r2_mag**3 * (-t21)**2 
-    f3 = 1 - (1 / 2) * MU / r2_mag**3 * t32**2
-    g1 = -t21 - (1 / 6) * MU / r2_mag**3 * (-t21)**2
-    g3 = t32 - (1 / 6) * MU / r2_mag**3 * t32**2
+    f1 = 1 - (1 / 2) * mu / r2_mag**3 * (-t21)**2 
+    f3 = 1 - (1 / 2) * mu / r2_mag**3 * t32**2
+    g1 = -t21 - (1 / 6) * mu / r2_mag**3 * (-t21)**2
+    g3 = t32 - (1 / 6) * mu / r2_mag**3 * t32**2
     return f1, g1, f3, g3
 
-def _iterateGaussIOD(orbit, t21, t32, q1, q2, q3, rho1hat, rho2hat, rho3hat, mu=MU, max_iter=10, tol=1e-15):
+def _calcM(r0_mag, r_mag, f, g, f_dot, g_dot, c0, c1, c2, c3, c4, c5, alpha, chi, mu=MU):
+    # Universal variables will differ between different texts and works in the literature.
+    # c0, c1, c2, c3, c4, c5 are expected to be 
+    w = chi / np.sqrt(mu)
+    alpha_alt = - mu * alpha
+    U0 = (1 - alpha_alt * chi**2) * c0
+    U1 = (chi - alpha_alt * chi**3) * c1 / np.sqrt(mu)
+    U2 = chi**2 * c2 / mu
+    U3 = chi**3 * c3 / mu**(3/2)
+    U4 = chi**4 * c4 / mu**(2)
+    U5 = chi**5 * c5 / mu**(5/2)
+    
+    F = f_dot
+    G = g_dot
+    
+    # Equations 18 and 19 in Sheppard 1985
+    U = (U2 * U3 + w * U4 - 3 * U5) / 3
+    W = g * U2 + 3 * mu * U
+    
+    # Calculate elements of the M matrix
+    m11 = (U0 / (r_mag * r0_mag) + 1 / r0_mag**2 + 1 / r_mag**2) * F - (mu**2 * W) / (r_mag * r0_mag)**3
+    m12 = F * U1 / r_mag + (G - 1) / r_mag**2
+    m13 = (G - 1) * U1 / r_mag - (mu * W) / r_mag**3
+    m21 = -F * U1 / r0_mag - (f - 1) / r0_mag**2
+    m22 = -F * U2
+    m23 = -(G - 1) * U2
+    m31 = (f - 1) * U1 / r0_mag - (mu * W) / r0_mag**3
+    m32 = (f - 1) * U2
+    m33 = g * U2 - W
+    
+    # Combine elements into matrix
+    M = np.array([
+        [m11, m12, m13],
+        [m21, m22, m23],
+        [m31, m32, m33],
+    ])
+    
+    return M
+
+def _calcStateTransitionMatrix(M, r0, v0, f, g, f_dot, g_dot, r, v):
+    I = np.identity(3)
+    state_0 = np.vstack((r0, v0))
+    state_1 = np.vstack((r, v))
+    
+    phi11 = f * I + state_1.T @ M[1:, :2] @ state_0 
+    phi12 = g * I + state_1.T @ M[1:, 1:] @ state_0
+    phi21 = f_dot * I + state_1.T @ M[:2, :2] @ state_0
+    phi22 = g_dot * I + state_1.T @ M[:2, 1:] @ state_0
+    
+    phi = np.block([
+        [phi11, phi12],
+        [phi21, phi22]
+    ])
+    return phi
+
+def _iterateGaussIOD(orbit, t21, t32, q1, q2, q3, rho1, rho2, rho3, mu=MU, max_iter=10, tol=1e-15):
     # Iterate over the polynomial solution from Gauss using the universal anomaly 
     # formalism until the solution converges or the maximum number of iterations is reached
 
     # Calculate variables that won't change per iteration
     sqrt_mu = np.sqrt(mu)
-    V = _calcV(rho1hat, rho2hat, rho3hat)
 
-    # Set up temporary variables that change iteration to iteration
-    orbit_iter = orbit
-    rho1_ratio = 1e10
-    rho2_ratio = 1e10
-    rho3_ratio = 1e10
-    rho1_mag = -1e10
-    rho2_mag = -1e10
-    rho3_mag = -1e10
+    # Calculate magntiude and unit rho vectors
+    rho1_mag = np.linalg.norm(rho1)
+    rho2_mag = np.linalg.norm(rho2)
+    rho3_mag = np.linalg.norm(rho3)
+    rho1_hat = rho1 / rho1_mag
+    rho2_hat = rho2 / rho2_mag
+    rho3_hat = rho3 / rho3_mag
 
+    orbit_iter = orbit
     i = 0
-    while ((np.abs(rho1_ratio) > tol) 
-            & (np.abs(rho2_ratio) > tol) 
-            & (np.abs(rho3_ratio) > tol)):
-        
+    for i in range(max_iter):
+        # Grab orbit position and velocity vectors 
+        # These should belong to the state of the object at the time of the second
+        # observation after applying Gauss's method the first time
         r = orbit_iter[:3]
         v = orbit_iter[3:]
         v_mag = np.linalg.norm(v)
         r_mag = np.linalg.norm(r)
-        
+
+        # Calculate the inverse semi-major axis
+        # Note: the definition of alpha will change between different works in the literature.
+        #   Here alpha is defined as 1 / a where a is the semi-major axis of the orbit
+        alpha = -v_mag**2 / mu + 2 / r_mag
+
         # Calculate the universal anomaly for both the first and third observations
-        # then calculate the Lagrange coefficients
+        # then calculate the Lagrange coefficients and the state for each observation.
+        # Use those to calculate the state transition matrix
         for j, dt in enumerate([-t21, t32]):
-            alpha = -v_mag**2 / mu + 2 / r_mag
-            
-            chi = calcChi(orbit_iter, dt, mu=mu, max_iter=max_iter, tol=tol)
+            # Calculate the universal anomaly 
+            # Universal anomaly here is defined in such a way that it satisfies the following
+            # differential equation:
+            #   d\chi / dt = \sqrt{mu} / r
+            chi = calcChi(orbit_iter, dt, mu=mu, max_iter=10, tol=tol)
             chi2 = chi**2
 
-            psi = alpha * chi2
+            # Calculate the values of the Stumpff functions
+            psi = alpha * chi2 
             c0, c1, c2, c3, c4, c5 = calcStumpff(psi)
 
-            f = 1 - chi**2 / r_mag * c2
+            # Calculate the Lagrange coefficients 
+            # and the corresponding state vector
+            f = 1 - chi2 / r_mag * c2
             g = dt - 1 / sqrt_mu * chi**3 * c3
 
+            r_new = f * r + g * v
+            r_new_mag = np.linalg.norm(r_new)
+
+            f_dot = sqrt_mu / (r_mag * r_new_mag) * (alpha * chi**3 * c3 - chi)
+            g_dot = 1 - chi2 / r_new_mag * c2
+
+            v_new = f_dot * r + g_dot * v
+            
+            # Calculate M matrix and use it to calculate the state transition matrix
+            M = _calcM(r_mag, r_new_mag, f, g, f_dot, g_dot, c0, c1, c2, c3, c4, c5, alpha, chi, mu=mu)
+            STM = _calcStateTransitionMatrix(M, r, v, f, g, f_dot, g_dot, r_new, v_new)
+            
             if j == 0:
-                g1 = g
-                f1 = f
+                STM1 = STM
+                v1 = v_new
+                r1 = r_new
             else:
-                g3 = g
-                f3 = f
-    
-        # Calculate the coplanarity coefficients
-        lambda1 = g3 / (f1 * g3 - f3 * g1) 
-        lambda3 = -g1 / (f1 * g3 - f3 * g1)   
-        
-        # Calculate new topocentric observer to target vectors
-        rho1_temp, rho2_temp, rho3_temp = _calcRhos(lambda1, lambda3, 
-                                                    q1, q2, 
-                                                    q3, rho1hat, 
-                                                    rho2hat, rho3hat, 
-                                                    V)
-        
-        # Calculate heliocentric position vector
-        r1 = q1 + rho1_temp
-        r2 = q2 + rho2_temp
-        r3 = q3 + rho3_temp
-        
-        # Update topocentric observer to target magnitudes ratios
-        rho1_ratio = rho1_mag / np.linalg.norm(rho1_temp)
-        rho2_ratio = rho2_mag / np.linalg.norm(rho2_temp)
-        rho3_ratio = rho3_mag / np.linalg.norm(rho3_temp)
-        
-        # Update running topocentric observer to target magnitudes variables
-        rho1_mag = np.linalg.norm(rho1_temp)
-        rho2_mag = np.linalg.norm(rho2_temp)
-        rho3_mag = np.linalg.norm(rho3_temp)
+                STM3 = STM
+                v3 = v_new
+                r3 = r_new
+
+        # Create phi error vector: as the estimate of the orbit 
+        # improves the elements in this vector should approach 0.
+        phi = np.hstack((
+            r1 - q1 - rho1_mag * rho1_hat, 
+            r - q2 - rho2_mag * rho2_hat, 
+            r3 - q3 - rho3_mag * rho3_hat))
+        if np.linalg.norm(phi) == 0:
+            break
 
-        # Calculate the velocity at the second observation
-        v2 = 1 / (f1 * g3 - f3 * g1) * (-f3 * r1 + f1 * r3)
+        dphi = np.zeros((9, 9), dtype=float)
+        dphi[0:3, 0:3] = STM1[0:3, 0:3]   # dr1/dr2
+        dphi[3:6, 0:3] = np.identity(3)   # dr2/dr2
+        dphi[6:9, 0:3] = STM3[0:3, 0:3]   # dr3/dr2
 
-        # Update the guess of the orbit
-        orbit_iter = np.concatenate((r2, v2))
-        
+        dphi[0:3, 3:6] = STM1[0:3, 3:6]   # dr1/dv2
+        dphi[3:6, 3:6] = np.zeros((3, 3)) # dr2/dv2
+        dphi[6:9, 3:6] = STM3[0:3, 3:6]   # dr3/dv2
+
+        dphi[0:3,6] = -v1 / C - rho1_hat
+
+        dphi[0:3,7] = v1 / C
+        dphi[3:6,7] = -rho2_hat
+        dphi[6:9,7] = v3 / C
+
+        dphi[6:9,8] = -v3 / C - rho3_hat
+
+        delta = np.linalg.solve(dphi, phi)
+        orbit_iter -= delta[0:6]
+        rho1_mag -= delta[6]
+        rho2_mag -= delta[7]
+        rho3_mag -= delta[8]
+                
         i += 1
         if i >= max_iter:
             break
@@ -156,7 +236,7 @@ def _iterateGaussIOD(orbit, t21, t32, q1, q2, q3, rho1hat, rho2hat, rho3hat, mu=
 
 def calcGauss(r1, r2, r3, t1, t2, t3):
     """
-    Calculates the velocity vector at the location of the second position vector (r2) with Gauss'
+    Calculates the velocity vector at the location of the second position vector (r2) with Gauss's
     original method.
     
     .. math::        
@@ -240,18 +320,18 @@ def gaussIOD(coords_eq_ang, t, coords_obs, velocity_method="gibbs", iterate=True
         Up to three preliminary orbits (as cartesian state vectors).
     """
     rho = equatorialToEclipticCartesian(equatorialAngularToCartesian(np.radians(coords_eq_ang)))
-    rho1hat = rho[0, :]
-    rho2hat = rho[1, :]
-    rho3hat = rho[2, :]
+    rho1_hat = rho[0, :]
+    rho2_hat = rho[1, :]
+    rho3_hat = rho[2, :]
     q1 = coords_obs[0,:]
     q2 = coords_obs[1,:]
     q3 = coords_obs[2,:]
     q2_mag = np.linalg.norm(q2)
 
     # Make sure rhohats are unit vectors
-    rho1hat = rho1hat / np.linalg.norm(rho1hat)
-    rho2hat = rho2hat / np.linalg.norm(rho2hat)
-    rho3hat = rho3hat / np.linalg.norm(rho3hat)
+    rho1_hat = rho1_hat / np.linalg.norm(rho1_hat)
+    rho2_hat = rho2_hat / np.linalg.norm(rho2_hat)
+    rho3_hat = rho3_hat / np.linalg.norm(rho3_hat)
 
     t1 = t[0]
     t2 = t[1]
@@ -260,10 +340,10 @@ def gaussIOD(coords_eq_ang, t, coords_obs, velocity_method="gibbs", iterate=True
     t21 = t2 - t1
     t32 = t3 - t2
 
-    A = _calcA(q1, q2, q3, rho1hat, rho3hat, t31, t32, t21)
-    B = _calcB(q1, q3, rho1hat, rho3hat, t31, t32, t21)
-    V = _calcV(rho1hat, rho2hat, rho3hat)
-    coseps2 = np.dot(q2, rho2hat) / q2_mag
+    A = _calcA(q1, q2, q3, rho1_hat, rho3_hat, t31, t32, t21)
+    B = _calcB(q1, q3, rho1_hat, rho3_hat, t31, t32, t21)
+    V = _calcV(rho1_hat, rho2_hat, rho3_hat)
+    coseps2 = np.dot(q2, rho2_hat) / q2_mag
     C0 = V * t31 * q2_mag**4 / B
     h0 = - A / B
 
@@ -292,11 +372,11 @@ def gaussIOD(coords_eq_ang, t, coords_obs, velocity_method="gibbs", iterate=True
     orbits = []
     for r2_mag in r2_mags:
         lambda1, lambda3 = _calcLambdas(r2_mag, t31, t32, t21)
-        rho1, rho2, rho3 = _calcRhos(lambda1, lambda3, q1, q2, q3, rho1hat, rho2hat, rho3hat, V)
+        rho1, rho2, rho3 = _calcRhos(lambda1, lambda3, q1, q2, q3, rho1_hat, rho2_hat, rho3_hat, V)
 
         # Test if we get the same rho2 as using equation 22 in Milani et al. 2008
         rho2_mag = (h0 - q2_mag**3 / r2_mag**3) * q2_mag / C0
-        np.testing.assert_almost_equal(np.dot(rho2_mag, rho2hat), rho2)
+        np.testing.assert_almost_equal(np.dot(rho2_mag, rho2_hat), rho2)
 
         r1 = q1 + rho1
         r2 = q2 + rho2
@@ -315,10 +395,10 @@ def gaussIOD(coords_eq_ang, t, coords_obs, velocity_method="gibbs", iterate=True
         if iterate == True:
             orbit = _iterateGaussIOD(orbit, t21, t32, 
                                      q1, q2, q3, 
-                                     rho1hat, rho2hat, rho3hat, 
-                                     mu=MU, max_iter=max_iter, tol=tol)
+                                     rho1, rho2, rho3,
+                                     mu=mu, max_iter=max_iter, tol=tol)
 
-        if np.linalg.norm(v2) >= C:
+        if np.linalg.norm(orbit[3:]) >= C:
             print("Velocity is greater than speed of light!")
         orbits.append(orbit)