From fcb86c1e55f84fe9bbfc1e6e063d208d5270268f Mon Sep 17 00:00:00 2001 From: David Minton Date: Wed, 2 Nov 2022 17:06:16 -0400 Subject: [PATCH] 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 | 392 ++++++++++++++++++++++++++++++- 1 file changed, 391 insertions(+), 1 deletion(-) diff --git a/python/swiftest/swiftest/tool.py b/python/swiftest/swiftest/tool.py index 91f8c5a46..07e447ab4 100644 --- a/python/swiftest/swiftest/tool.py +++ b/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 \ No newline at end of file + 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) \ No newline at end of file