Added xv2el and el2xv conversion functions to the swiftest.tools package (_one works on single sets of data, _vec does the conversion row-wise on an array of data)

python/swiftest/swiftest/tool.py

@@ -100,4 +100,394 @@ def follow_swift(ds, ifol=None, nskp=None):
     except IOError:
         print(f"Error writing to follow.out")
 
-    return fol
+    return fol
+
+def danby(M, ecc, accuracy=1e-14):
+    """
+    Danby's method to solve Kepler's equation.
+
+    Parameters
+    ----------
+    M : float
+        the mean anomaly in radians
+    ecc : float
+        the eccentricity
+    accuracy : float
+        the relative accuracy to obtain a solution. Default is 1e-12.
+
+    Returns
+    ----------
+    E : float
+        the eccentric anomaly in radians
+
+    References
+    __________
+    Danby, J.M.A. 1988. Fundamentals of celestial mechanics. Richmond, Va., U.S.A., Willmann-Bell, 1988. 2nd ed.
+    Murray, C.D., Dermott, S.F., 1999. Solar system dynamics, New York, New York. ed, Cambridge University Press.
+
+    """
+
+    def kepler_root(E, ecc, M, deriv):
+        """
+        The Kepler equation root function.
+
+        Parameters
+        ----------
+        E : float
+            the eccentric anomaly in radians
+        ecc : float
+            the eccentricity
+        M : float
+            the mean anomaly in radians
+        deriv : int between 0 and 3
+            The order of the derivative to compute
+
+        Returns
+        ----------
+        deriv = 0: E - e * np.sin(E) - M
+        deriv = 1: 1 - e * np.cos(E)
+        deriv = 2: e * np.sin(E)
+        deriv = 3: e * np.cos(E)
+
+        Note: The function will return 0 when E is correct for a given e and M
+        """
+
+        if deriv == 0:
+            return E - ecc * np.sin(E) - M
+        elif deriv == 1:
+            return 1.0 - ecc * np.cos(E)
+        elif deriv == 2:
+            return ecc * np.sin(E)
+        elif deriv == 3:
+            return ecc * np.cos(E)
+        else:
+            print(f"deriv={deriv} is undefined")
+            return None
+
+    def delta_ij(E, ecc, M, j):
+        """
+        Danby's intermediate delta functions. This function is recursive for j>1.
+
+        Parameters
+        ----------
+        E : float
+            the eccentric anomaly in radians
+        ecc : float
+            the eccentricity
+        M : float
+            the mean anomaly in radians
+        j : int between 1 and 3
+            The order of the delta function to compute
+
+        Returns
+        ----------
+        delta_ij value used in Danby's iterative Kepler equation solver
+
+        """
+        if j == 1:
+            return -kepler_root(E, ecc, M, 0) / kepler_root(E, ecc, M, 1)
+        elif j == 2:
+            return -kepler_root(E, ecc, M, 0) / (kepler_root(E, ecc, M, 1)
+                                                 - delta_ij(E, ecc, M, 1) * kepler_root(E, ecc, M, 2) / 2.0)
+        elif j == 3:
+            return -kepler_root(E, ecc, M, 0) / (kepler_root(E, ecc, M, 1)
+                                                 + delta_ij(E, ecc, M, 1) * kepler_root(E, ecc, M, 2) / 2.0
+                                                 + delta_ij(E, ecc, M, 2) ** 2 * kepler_root(E, ecc, M,
+                                                                                             3) / 6.0)
+        else:
+            print(f"j = {j} is not a valid input to the delta_ij function")
+
+    # If this is a circular orbit, just return the mean anomaly as the eccentric anomaly
+    if ecc < np.finfo(np.float64).tiny:  # Prevent floating point exception for 0 eccentricity orbits
+        return M
+
+    # Initial guess
+    k = 0.85
+    E = M + np.sign(np.sin(M)) * k * ecc
+    MAXLOOPS = 50  # Maximum number of times to iterate before we give up
+    for i in range(MAXLOOPS):
+        Enew = E + delta_ij(E, ecc, M, 3)
+        if np.abs((Enew - E) / E) < accuracy:
+            return Enew
+        E = Enew
+
+    raise RuntimeError("The danby function did not converge on a solution.")
+
+
+
+def el2xv_one(mu, a, ecc, inc, Omega, omega, M):
+    """
+    Compute osculating orbital elements from relative Cartesian position and velocity
+    All angular measures are returned in radians
+        If inclination < TINY, longitude of the ascending node is arbitrarily set to 0
+
+        If eccentricity < sqrt(TINY), argument of pericenter is arbitrarily set to 0
+
+          ALGORITHM:  See Fitzpatrick "Principles of Cel. Mech."
+
+    Adapted from Martin Duncan's el2xv.f
+
+    Parameters
+    ----------
+    mu : float
+        Central body gravitational constant
+    a : float
+        semimajor axis
+    ecc : float
+        eccentricity
+    inc : float
+        inclination (degrees)
+    Omega : float
+        longitude of ascending node (degrees)
+    omega : float
+        argument of periapsis (degrees)
+    M : flaot
+        Mean anomaly (degrees)
+
+    Returns
+    ----------
+    rvec, vvec  : tuple of float vectors
+
+    rvec : (3) float vector
+        Cartesian position vector
+    vvec : (3) float vector
+        Cartesian velocity vector
+
+    """
+
+    if ecc < 0.0:
+        print("Error in el2xv! Eccentricity cannot be negative. Setting it to 0")
+        e = 0.0
+        iorbit_type = "ellipse"
+    else:
+        e = ecc
+        em1 = e - 1.
+        if np.abs(em1) < 2 * np.finfo(float).eps:
+            iorbit_type = "parabola"
+        elif e > 1.0:
+            iorbit_type = "hyperbola"
+        else:
+            iorbit_type = "ellipse"
+
+    def scget(angle):
+        """
+        Efficiently compute the sine and cosine of an input angle
+        Input angle must be in radians
+
+        Adapted from David E. Kaufmann's Swifter routine: orbel_scget.f90
+        Adapted from Hal Levison's Swift routine orbel_scget.f
+
+        Parameters
+        ----------
+        angle : input angle
+
+        Returns
+        -------
+        sx : sin of angle
+        cx : cos of angle
+
+        """
+        TWOPI = 2 * np.pi
+        nper = int(angle / TWOPI)
+        x = angle - nper * TWOPI
+        if x < 0.0:
+            x += TWOPI
+
+        sx = np.sin(x)
+        cx = np.sqrt(1.0 - sx ** 2)
+        if x > np.pi / 2.0 and x < 3 * np.pi / 2.0:
+            cx = -cx
+        return sx, cx
+
+    sip, cip = scget(np.deg2rad(omega))
+    so, co = scget(np.deg2rad(Omega))
+    si, ci = scget(np.deg2rad(inc))
+
+    d11 = cip * co - sip * so * ci
+    d12 = cip * so + sip * co * ci
+    d13 = sip * si
+    d21 = -sip * co - cip * so * ci
+    d22 = -sip * so + cip * co * ci
+    d23 = cip * si
+
+    # Get the other quantities depending on orbit type
+    if iorbit_type == "ellipse":
+        E = danby(np.deg2rad(M),e)
+        scap, ccap = scget(E)
+        sqe = np.sqrt(1. - e**2)
+        sqgma = np.sqrt(mu * a)
+        xfac1 = a * (ccap - e)
+        xfac2 = a * sqe * scap
+        ri = 1. / (a * (1. - e * ccap))
+        vfac1 = -ri * sqgma * scap
+        vfac2 = ri * sqgma * sqe * ccap
+    else:
+        print(f"Orbit type: {iorbit_type} not yet implemented")
+        xfac1 = 0.0
+        xfac2 = 0.0
+        vfac1 = 0.0
+        vfac2 = 0.0
+
+    rvec = np.array([d11 * xfac1 + d21 * xfac2,
+                     d12 * xfac1 + d22 * xfac2,
+                     d13 * xfac1 + d23 * xfac2])
+    vvec = np.array([d11 * vfac1 + d21 * vfac2,
+                     d12 * vfac1 + d22 * vfac2,
+                     d13 * vfac1 + d23 * vfac2])
+
+    return rvec, vvec
+
+
+def el2xv_vec(mu, a, ecc, inc, Omega, omega, M):
+    """
+
+    Vectorized call to el2xv_one
+    Parameters
+    ----------
+    mu : float
+        Central body gravitational constant
+    a : (n) float array
+        semimajor axis
+    ecc : (n) float array
+        eccentricity
+    inc : (n) float array
+        inclination (degrees)
+    Omega : (n) float array
+        longitude of ascending node (degrees)
+    omega : (n) float array
+        argument of periapsis (degrees)
+    M : (n) flaot array
+        Mean anomaly (degrees)
+
+    Returns
+    ----------
+    rvec, vvec  : tuple of float vectors
+
+    rvec : (n,3) float rray
+        Cartesian position vector
+    vvec : (n,3) float array
+        Cartesian velocity vector
+    """
+    vecfunc = np.vectorize(el2xv_one, signature='(),(),(),(),(),(),()->(3),(3)')
+    return vecfunc(mu, a, ecc, inc, Omega, omega, M)
+
+def xv2el_one(mu,rvec,vvec):
+    """
+    Converts from cartesian position and velocity values to orbital elements
+
+    Parameters
+    ----------
+    mu : float
+        Central body gravitational constant
+    rvec : (3) float array
+        Cartesian position vector
+    vvec : (3) float array
+        Cartesian velocity vector
+
+    Returns
+    ----------
+    a, ecc, inc, Omega, omega, M, varpi, f, lam : tuple of floats
+
+    a : float
+        semimajor axis
+    ecc : float
+        eccentricity
+    inc : float
+        inclination (degrees)
+    Omega : float
+        longitude of ascending node (degrees)
+    omega : float
+        argument of periapsis (degrees)
+    M : float
+        mean anomaly (degrees)
+    varpi : flaot
+        longitude of periapsis (degrees)
+    f : float
+        true anomaly (degrees)
+    lam : float
+        mean longitude (degrees)
+
+    """
+
+    rmag = np.linalg.norm(rvec)
+    vmag2 = np.vdot(vvec,vvec)
+    h = np.cross(rvec,vvec)
+    hmag = np.linalg.norm(h)
+
+    rdot = np.sign(np.vdot(rvec,vvec)) * np.sqrt(vmag2 - (hmag / rmag)**2)
+
+    a = 1.0/(2.0 / rmag - vmag2/mu)
+    ecc = np.sqrt(1 - hmag**2 / (mu * a))
+    inc = np.arccos(h[2]/hmag)
+
+    goodinc = np.abs(inc) > np.finfo(np.float64).tiny
+    sO = np.where(goodinc,  np.sign(h[2]) * h[0] / (hmag * np.sin(inc)),0.0)
+    cO = np.where(goodinc, -np.sign(h[2]) * h[1] / (hmag * np.sin(inc)),1.0)
+
+    Omega = np.arctan2(sO, cO)
+
+    sof = np.where(goodinc,rvec[2] / (rmag * np.sin(inc)), rvec[1]/rmag)
+    cof = np.where(goodinc,(rvec[0] / rmag + np.sin(Omega) * sof * np.cos(inc)) / np.cos(Omega), rvec[0]/rmag)
+
+    of = np.arctan2(sof,cof)
+
+    goodecc = ecc > np.finfo(np.float64).tiny
+    sf = np.where(goodecc,a * (1.0 - ecc**2)  * rdot / (hmag * ecc), 0.0)
+    cf = np.where(goodecc,(1.0 / ecc) * (a * (1.0 - ecc**2) / rmag - 1.0), 1.0 )
+
+    f = np.arctan2(sf, cf)
+
+    omega = of - f
+
+    varpi = Omega + omega
+
+    # Compute eccentric anomaly & mean anomaly in order to get mean longitude
+    E = np.where(ecc > np.finfo(np.float64).tiny, np.arccos(-(rmag - a) / (a * ecc)), 0)
+    if np.sign(np.vdot(rvec, vvec)) < 0.0:
+        E = 2 * np.pi - E
+
+    M = E - ecc * np.sin(E)
+    lam = M + varpi
+
+    return a, ecc, np.rad2deg(inc), np.rad2deg(Omega), np.rad2deg(omega), np.rad2deg(M), np.rad2deg(varpi), np.rad2deg(f), np.rad2deg(lam)
+
+def xv2el_vec(mu, rvec, vvec):
+    """
+    Vectorized call to xv2el_one.
+
+    Parameters
+    ----------
+    mu : float
+        Central body gravitational constant
+    rvec : (n,3) float array
+        Cartesian position vector
+    vvec : (n,3) float array
+        Cartesian velocity vector
+
+    Returns
+    ----------
+    a, ecc, inc, Omega, omega, M, varpi, f, lam : tuple of float arrays
+
+    a : (n) float array
+        semimajor axis
+    ecc : (n) float array
+        eccentricity
+    inc : (n) float array
+        inclination (degrees)
+    Omega : (n) float array
+        longitude of ascending node (degrees)
+    omega : (n) float array
+        argument of periapsis (degrees)
+    M : (n) float array
+        mean anomaly (degrees)
+    varpi : (n) flaot array
+        longitude of periapsis (degrees)
+    f : (n) float array
+        true anomaly (degrees)
+    lam : (n) float array
+        mean longitude (degrees)
+
+    """
+
+    vecfunc = np.vectorize(xv2el_one, signature='(),(3),(3)->(),(),(),(),(),(),(),(),()')
+    return vecfunc(mu, rvec, vvec)