Convert calcC2C3 into calcStumpff (add in first 6 Stumpff equations as opposed)

Author: moeyensj

thor/orbits/iod/gauss.py

@@ -4,7 +4,7 @@
 from ...constants import Constants as c
 from ...coordinates import equatorialToEclipticCartesian
 from ...coordinates import equatorialAngularToCartesian
-from ..propagate import calcC2C3
+from ..propagate import calcStumpff
 from ..propagate import calcChi
 from .gibbs import calcGibbs
 from .herrick_gibbs import calcHerrickGibbs
@@ -104,7 +104,7 @@ def _iterateGaussIOD(orbit, t21, t32, q1, q2, q3, rho1hat, rho2hat, rho3hat, mu=
             chi2 = chi**2
 
             psi = alpha * chi2
-            c2, c3 = calcC2C3(psi)
+            c0, c1, c2, c3, c4, c5 = calcStumpff(psi)
 
             f = 1 - chi**2 / r_mag * c2
             g = dt - 1 / sqrt_mu * chi**3 * c3

thor/orbits/propagate/stumpff.py

@@ -2,29 +2,31 @@
 from numba import jit
 
 __all__ = [
-    "calcC2C3"
+    "calcStumpff"
 ]
 
-@jit("UniTuple(f8, 2)(f8)", nopython=True)
-def calcC2C3(psi):
+@jit("UniTuple(f8, 6)(f8)", nopython=True)
+def calcStumpff(psi):
     """
-    Calculate the second and third Stumpff functions. 
+    Calculate the first 6 Stumpff functions for universal variable psi. 
     
     .. math::
+    
+        c_0(\Psi) = \begin{cases}
+            \cos{\sqrt{\Psi}} & \text{ if } \Psi > 0 \\ 
+            \cosh{\sqrt{-\Psi}} & \text{ if } \Psi < 0 \\ 
+            1 & \text{ if } \Psi= 0
+        \end{cases} 
         
-        c_2(\Psi) = \begin{cases}
-            \frac{1 - \cos{\sqrt{\Psi}}}{\Psi} & \text{ if } \Psi > 0 \\ 
-            \frac{1 - \cosh{\sqrt{-\Psi}}}{\Psi} & \text{ if } \Psi < 0 \\ 
-            \frac{1}{2} & \text{ if } \Psi= 0
+        c_1(\Psi) = \begin{cases}
+            \frac{\sin{\sqrt{\Psi}}{\sqrt{\Psi}} & \text{ if } \Psi > 0 \\ 
+            \frac{\sinh{\sqrt{-\Psi}}{\sqrt{-\Psi}} & \text{ if } \Psi < 0 \\ 
+            1 & \text{ if } \Psi= 0
         \end{cases}
         
-        c_3(\Psi) = \begin{cases}
-            \frac{\sqrt{\Psi} - \sin{\sqrt{\Psi}}}{\sqrt{\Psi^3}} & \text{ if } \Psi > 0 \\ 
-            \frac{\sinh{\sqrt{-\Psi}} - \sqrt{-\Psi}}{\sqrt{(-\Psi)^3}} & \text{ if } \Psi < 0 \\ 
-            \frac{1}{6} & \text{ if } \Psi= 0
-        \end{cases}
+        \Psi c_{n+2} = \frac{1}{k!} - c_n(\Psi)
         
-    For more details on theory see Chapter 2 in David A. Vallado's "Fundamentals of Astrodynamics
+    For more details on the universal variable formalism see Chapter 2 in David A. Vallado's "Fundamentals of Astrodynamics
     and Applications" or Chapter 3 in Howard Curtis' "Orbital Mechanics for Engineering Students".
     
     Parameters
@@ -34,17 +36,29 @@ def calcC2C3(psi):
         
     Returns
     -------
-    c2, c3 : float, float
-        Second and third Stumpff functions.
+    s0, s1, s2, s3, s4, s5 : 6 x float
+        First six Stumpff functions.
     """
     if psi > 0.0:
-        c2 = (1 - np.cos(np.sqrt(psi))) / psi
-        c3 = (np.sqrt(psi) - np.sin(np.sqrt(psi))) / np.sqrt(psi)**3
+        c0 = np.cos(np.sqrt(psi))
+        c1 = np.sin(np.sqrt(psi)) / np.sqrt(psi)
+        c2 = (1. - c0) / psi
+        c3 = (1. - c1) / psi
+        c4 = (1/2. - c2) / psi
+        c5 = (1/6. - c3) / psi
     elif psi < 0.0:
-        c2 = (np.cosh(np.sqrt(-psi)) - 1) / (-psi)
-        c3 = (np.sinh(np.sqrt(-psi)) - np.sqrt(-psi)) / np.sqrt(-psi)**3
+        c0 = np.cosh(np.sqrt(-psi))
+        c1 = np.sinh(np.sqrt(-psi)) / np.sqrt(-psi)
+        c2 = (1. - c0) / psi
+        c3 = (1. - c1) / psi
+        c4 = (1/2. - c2) / psi
+        c5 = (1/6. - c3) / psi
     else:
+        c0 = 1.
+        c1 = 1.
         c2 = 1/2.
         c3 = 1/6.
+        c4 = 1/24.
+        c5 = 1/120.
 
-    return c2, c3
+    return c0, c1, c2, c3, c4, c5

thor/orbits/propagate/universal.py

@@ -2,7 +2,7 @@
 from numba import jit
 
 from ...constants import Constants as c
-from .stumpff import calcC2C3
+from .stumpff import calcStumpff
 
 __all__ = [
     "calcChi",
@@ -54,7 +54,7 @@ def calcChi(orbit, dt, mu=MU, max_iter=10000, tol=1e-14):
     while np.abs(ratio) > tol:
         chi2 = chi**2
         psi = alpha * chi2
-        c2, c3 = calcC2C3(psi)
+        c0, c1, c2, c3, c4, c5 = calcStumpff(psi)
 
         # Newton-Raphson
         f = (r_mag * rv_mag / sqrt_mu * chi2 * c2 
@@ -121,7 +121,7 @@ def propagateUniversal(orbits, t0, t1, mu=MU, max_iter=10000, tol=1e-14):
 
             alpha = -v_mag**2 / mu + 2 / r_mag
             psi = alpha * chi2
-            c2, c3 = calcC2C3(psi)
+            c0, c1, c2, c3, c4, c5 = calcStumpff(psi)
 
             f = 1 - chi**2 / r_mag * c2
             g = dt - 1 / sqrt_mu * chi**3 * c3