1# celestial_mechanics.py 2# celestial mechanics utilities for exoplanet ephemerides 3# 4# intellectual property: 5# Copyright 2009 David W. Hogg. All rights reserved. 6# Licensed under a 3-clause BSD style license - see LICENSE 7# 8# comments: 9# - Written for clarity, not speed. The code is intended to be human- 10# readable. 11# 12# bugs: 13# - Need to make Fourier expansion functions. 14# - What to do if e is close to 1.0 in eccentric_anomaly 15# 16from __future__ import print_function 17 18from math import pi 19import unittest 20import sys 21 22import numpy 23import numpy as np 24from numpy import * 25from scipy.special import jn 26import matplotlib.pyplot as plt 27 28from astrometry.util.starutil_numpy import * 29 30ihat = array([1.,0.,0.]) 31jhat = array([0.,1.,0.]) 32khat = array([0.,0.,1.]) 33default_tolerance = 1e-15 # (radians) don't set to zero or less 34default_maximum_iteration = 1000 # should never hit this limit! 35default_order = 32 36default_K = 1.0 37 38(Equinox, Solstice, EclipticPole) = ecliptic_basis() 39 40c_au_per_yr = 63239.6717 # google says 41 42#GM_sun = 2.9591310798672560E-04 #AU^3/d^2 43# Google says mass of the sun * G = 39.4775743 (au^3) / (yr^2) 44GM_sun = 39.4775743 # AU^3 / yr^2 45 46def norm1d(x): 47 assert(len(x.shape) == 1) 48 return np.sqrt(np.sum(x**2)) 49 50def deg2rad(x): 51 return x * pi/180. 52 #return radians(x) 53 54def orbital_elements_to_ss_xyz(E, observer=None, light_travel=True): 55 (a,e,i,Omega,pomega,M,GM) = E 56 # ugh, it's hard to be units-agnostic. 57 # we just assert here so we have to think about this! 58 # This means: 59 # distances in AU 60 # angles in radians 61 # times in years 62 assert(GM == GM_sun) 63 if light_travel: 64 assert(observer is not None) 65 # orbital angular velocity [radians/yr] 66 meanfrequency = np.sqrt(GM / a**3) 67 # Correct for light-time delay. 68 # dM: [radians] 69 dM = 0. 70 lastdM = dM 71 dx = None 72 for ii in range(100): 73 (x,v) = phase_space_coordinates_from_orbital_elements( 74 a,e,i,Omega,pomega,M-dM,GM) 75 if not light_travel: 76 if observer is not None: 77 dx = x - observer 78 break 79 dx = x - observer 80 # light-travel distance [AU] 81 r = norm1d(dx) 82 # light-travel time [yr] 83 travel = r / c_au_per_yr 84 #print 'light-travel time:', travel, 'yr, or', travel*365.25, 'd' 85 # light travel in angular periods [radians] 86 dM = travel * meanfrequency 87 if abs(lastdM - dM) < 1e-12: 88 break 89 lastdM = dM 90 if ii == 99: 91 print('Warning: orbital_elements_to_ss_xyz: niters', ii) 92 return x,dx 93 94def orbital_elements_to_xyz(E, observer, light_travel=True, normalize=True): 95 (x,dx) = orbital_elements_to_ss_xyz(E, observer, light_travel) 96 if normalize: 97 dx /= norm1d(dx) 98 edx = dx[0] * Equinox + dx[1] * Solstice + dx[2] * EclipticPole 99 return edx 100 101# E = (a,e,i,Omega,pomega,M, GM) 102# observer = 3-vector 103# light_travel: correct for light travel time? 104# Returns RA,Dec in degrees. 105def orbital_elements_to_radec(E, observer, light_travel=True): 106 xyz = orbital_elements_to_xyz(E, observer, light_travel) 107 return xyztoradec(xyz) 108 109# convert orbital elements into vectors in the plane of the orbit. 110def orbital_vectors_from_orbital_elements(i, Omega, pomega): 111 ascendingnodevector = np.cos(Omega) * ihat + np.sin(Omega) * jhat 112 tmpydir= np.cross(khat, ascendingnodevector) 113 zhat= np.cos(i) * khat - np.sin(i) * tmpydir 114 tmpydir= np.cross(zhat, ascendingnodevector) 115 xhat= np.cos(pomega) * ascendingnodevector + np.sin(pomega) * tmpydir 116 yhat = np.cross(zhat, xhat) 117 return (xhat, yhat, zhat) 118 119def position_from_orbital_vectors(xhat, yhat, a, e, M): 120 E = eccentric_anomaly_from_mean_anomaly(M, e) 121 cosE = np.cos(E) 122 sinE = np.sin(E) 123 b = a*np.sqrt(1. - e**2) 124 x = a * (cosE - e) * xhat + b * sinE * yhat 125 return x 126 127# convert orbital elements to phase-space coordinates 128# a - semi-major axis (length units) 129# e - eccentricity 130# i - inclination (rad) 131# Omega - longitude of ascending node (rad) 132# pomega - argument of periapsis (rad) 133# M - mean anomaly (rad) 134# GM - Newton's constant times central mass (length units cubed over time units squared) 135# return - (x,v) 136# position, velocity (length units, length units per time unit) 137def phase_space_coordinates_from_orbital_elements(a, e, i, Omega, pomega, M, GM): 138 (xhat, yhat, zhat) = orbital_vectors_from_orbital_elements(i, Omega, pomega) 139 # [radians/yr] 140 dMdt = np.sqrt(GM / a**3) 141 E = eccentric_anomaly_from_mean_anomaly(M, e) 142 cosE = np.cos(E) 143 sinE = np.sin(E) 144 # [radians/yr] 145 dEdt = 1.0 / (1.0 - e * cosE) * dMdt 146 # [AU] 147 b = a*np.sqrt(1. - e**2) 148 x = a * (cosE - e) * xhat + b * sinE * yhat 149 # [AU/yr] 150 v = -a * sinE * dEdt * xhat + b * cosE * dEdt * yhat 151 return (x, v) 152 153class UnboundOrbitError(ValueError): 154 pass 155 156def potential_energy_from_position(x, GM): 157 return -1. * GM / norm1d(x) 158 159def energy_from_phase_space_coordinates(x, v, GM): 160 return 0.5 * np.dot(v, v) + potential_energy_from_position(x, GM) 161 162# convert phase-space coordinates to orbital elements 163# x - position (3-vector, length units) 164# v - velocity (3-vector, length units per time unit) 165# GM - Newton's constant times central mass (length units cubed over time units squared) 166# return - (a, e, i, Omega, pomega, M) 167# - see "phase_space_coordinates" for definitions 168def orbital_elements_from_phase_space_coordinates(x, v, GM): 169 energy = energy_from_phase_space_coordinates(x, v, GM) 170 if energy > 0: 171 raise UnboundOrbitError('orbital_elements_from_phase_space_coordinates: Unbound orbit') 172 173 angmom = np.cross(x, v) 174 zhat = angmom / norm1d(angmom) 175 evec = np.cross(v, angmom) / GM - x / norm1d(x) 176 e = norm1d(evec) 177 if e == 0: 178 # by convention: 179 xhat = np.cross(jhat, zhat) 180 xhat /= norm1d(xhat) 181 else: 182 xhat = evec / e 183 yhat = np.cross(zhat, xhat) 184 a = -0.5 * GM / energy 185 i = np.arccos(angmom[2] / norm1d(angmom)) 186 if i == 0: 187 Omega = 0.0 188 else: 189 Omega = np.arctan2(angmom[1], angmom[0]) + 0.5 * pi 190 if Omega < 0: 191 Omega += 2.*pi 192 if i < 0: 193 i *= -1. 194 Omega += pi 195 cosOmega = cos(Omega) 196 sinOmega = sin(Omega) 197 if e == 0: 198 pomega = 0. - Omega 199 else: 200 pomega = np.arccos(min(1.0, (evec[0] * cosOmega + evec[1] * sinOmega) / e)) 201 horriblescalar = ( sinOmega * evec[2] * angmom[0] 202 - cosOmega * evec[2] * angmom[1] 203 + cosOmega * evec[1] * angmom[2] 204 - sinOmega * evec[0] * angmom[2]) 205 if horriblescalar < 0.: 206 pomega = 2.0 * pi - pomega 207 if pomega < 0.0: 208 pomega += 2.0 * pi 209 if pomega > 2.0 * pi: 210 pomega -= 2.0 * pi 211 f = np.arctan2(np.dot(yhat, x), np.dot(xhat, x)) 212 M = mean_anomaly_from_true_anomaly(f, e) 213 if M < 0: 214 M += 2.*pi 215 return (a, e, i, Omega, pomega, M) 216 217# convert eccentric anomaly to mean anomaly 218# E - eccentric anomaly (radians) 219# e - eccentricity 220# return - mean anomaly (radians) 221def mean_anomaly_from_eccentric_anomaly(E, e): 222 return (E - e * np.sin(E)) 223 224def mean_anomaly_from_true_anomaly(f, e): 225 return mean_anomaly_from_eccentric_anomaly(eccentric_anomaly_from_true_anomaly(f, e), e) 226 227# convert mean anomaly to eccentric anomaly 228# M - [array of] mean anomaly (radians) 229# e - eccentricity 230# [tolerance - read the source] 231# [maximum_iteration - read the source] 232# return - eccentric anomaly (radians) 233def eccentric_anomaly_from_mean_anomaly(M, e, tolerance = default_tolerance, 234 maximum_iteration = default_maximum_iteration, verbose=False): 235 E = M + e * np.sin(M) 236 iteration = 0 237 deltaM = 100.0 238 while (iteration < maximum_iteration) and (abs(deltaM) > tolerance): 239 deltaM = (M - mean_anomaly_from_eccentric_anomaly(E, e)) 240 E = E + deltaM / (1. - e * cos(E)) 241 iteration += 1 242 if verbose: print('eccentric anomaly iterations:',iteration) 243 return E 244 245def eccentric_anomaly_from_true_anomaly(f, e): 246 E = np.arccos((np.cos(f) + e) / (1.0 + e * np.cos(f))) 247 E *= (np.sign(np.sin(f)) * np.sign(np.sin(E))) 248 return E 249 250# convert eccentric anomaly to true anomaly 251# E - eccentric anomaly (radians) 252# e - eccentricity 253# return - true anomaly (radians) 254def true_anomaly_from_eccentric_anomaly(E, e): 255 f = np.arccos((np.cos(E) - e) / (1.0 - e * np.cos(E))) 256 f *= (np.sign(np.sin(f)) * np.sign(np.sin(E))) 257 return f 258 259# compute radial velocity 260# K - radial velocity amplitude 261# f - true anomaly (radians) 262# e - eccentricity 263# pomega - eccentric longitude (radians) 264# return - radial velocity (same units as K) 265def radial_velocity(K, f, e, pomega): 266 return K * (np.sin(f + pomega) + e * np.sin(pomega)) 267 268# compute radial velocity 269# K - radial velocity amplitude 270# M - mean anomaly (radians) 271# e - eccentricity 272# pomega - eccentric longitude (radians) 273# return - radial velocity (same units as K) 274def radial_velocity_from_M(K, M, e, pomega): 275 E = M + e*np.sin(M) 276 term1 = np.cos(pomega) * np.sqrt(1 - e**2) * np.sin(E) / (1 - e*np.cos(E)) 277 term2 = np.sin(pomega) * (np.cos(E) - e) / (1 - e*np.cos(E)) 278 term3 = e*np.sin(pomega) 279 return K * (term1 + term2 + term3) 280 281# compute radial velocity using a truncated Fourier series 282# K - radial velocity amplitude 283# M - mean anomaly (radians) APW: you may want to change this input 284# e - eccentricity 285# pomega - eccentric longitude (radians) 286# phi - phase 287# [order - read the source] 288# return - radial velocity (same units as K) 289def radial_velocity_fourier_series(K, M, e, pomega, phi, order=default_order): 290 vr = 0.0 291 for n in arange(0, order+1, 1): 292 vr += K*(fourier_coeff_A(n, pomega, phi, e) * np.cos(n*(M-phi)) \ 293 + fourier_coeff_B(n, pomega, phi, e) * np.sin(n*(M-phi))) 294 return vr 295 296# the following is based on the naming convention in Itay's notes on Fourier analysis 297# - fourier_coeff_A and fourier_coeff_B are the actual coefficients in the series 298# - aprime and bprime are just used to simplify the code, and break it up to make it more readable 299 300# n - order of the coefficient 301# e - eccentricity 302def aprime(n,e): 303 return np.sqrt(1. - e**2)*( ((np.sqrt(1. - e**2) - 1)/e)*jn(n,n*e) + jn(n-1,n*e)) 304 305def bprime(n,e): 306 return np.sqrt(1. - e**2)*( ((np.sqrt(1. - e**2) - 1)/e)*jn(n,n*e) - jn(n-1,n*e)) 307 308def fourier_coeff_A(n, pomega, phi, e): 309 return 0.5 * (aprime(n,e) * np.sin(pomega + n * phi) + bprime(n,e) * np.sin(pomega - n * phi)) 310 311def fourier_coeff_B(n, pomega, phi, e): 312 return 0.5 * (aprime(n,e) * np.cos(pomega + n * phi) - bprime(n,e) * np.cos(pomega - n * phi)) 313 314# APW: adjust function call as necessary 315# return - amplitudes as a list of tuples (An,Bn) 316def radial_velocity_fourier_amplitudes(K, phi, e, pomega, order=default_order): 317 amplitudes = [] 318 for n in range(order): 319 amplitudes.append((fourier_coeff_A(n, pomega, phi, e), fourier_coeff_B(n, pomega, phi, e))) 320 return amplitudes 321 322def eccentricity_from_fourier_amplitudes(amplitudes): 323 K = np.sqrt(amplitude[0]**2 + amplitude[1]**2) 324 phi = np.arctan(amplitude[1] / amplitude[0]) # WRONG? 325 e = 0 # WRONG 326 pomega = 0 # WRONG 327 return (K, phi, e, pomega) 328 329 330 331# some functional testing 332if __name__ == '__main__': 333 from test_celestial_mechanics import * 334 #suite = unittest.TestLoader().loadTestsFromTestCase(TestOrbitalElements) 335 336 suite = unittest.TestSuite() 337 #suite.addTest(TestOrbitalElements('testEdgeCases')) 338 suite.addTest(TestOrbitalElements('testAgainstJPL_2')) 339 340 unittest.TextTestRunner(verbosity=2).run(suite) 341 import sys 342 sys.exit(0) 343 344 # -- test Earth ephemeris against JPL at Holmes' closest approach 345 # -- test Holmes against JPL ---------''------------ 346 # -- test direction2radec() 347 348 try: 349 arg1 = sys.argv[1] 350 except IndexError: 351 arg1 = 0 352 353 for e in [0.01, 0.1, 0.9, 0.99, 0.999]: 354 print('eccentricity:', e) 355 M = arange(-3.16,-3.14,0.001) # easy 356 #M = arange(-10., 10., 2.1) # easy 357 #M = arange(-0.01,0.01,0.001) # hard 358 print('mean anomaly input:', M) 359 E = eccentric_anomaly_from_mean_anomaly(M, e, verbose=True) 360 print('eccentric anomaly output:', E) 361 f = true_anomaly_from_eccentric_anomaly(E, e) 362 print('true anomaly output:', f) 363 M2 = mean_anomaly_from_eccentric_anomaly(E, e) 364 print('round-trip error:', M2 - M) 365 366 if arg1 == "plot": 367 # This code will do the plotting: 368 range_min = 0.0 369 range_max = 20.0 370 step_size = 0.01 371 phase = 0.0 372 pomegas = [0.0, pi/5., 3.*pi/5., 5.*pi/5., 7.*pi/5., 9.*pi/5.] 373 pomegas_str = ["pomega = $0.0$", "pomega = $\pi/5$", "pomega = $3\pi/5$", "pomega = $\pi$", \ 374 "pomega = $7\pi/5$", "pomega = $9\pi/5$"] 375 eccens = [0.01, 0.1, 0.5, 0.9, 0.99] 376 orders = [2, 4, 8, 16, 32] 377 M = arange(range_min,range_max,step_size) 378 379 for n in orders: 380 for e in eccens: 381 i = 1 382 for pomega in pomegas: 383 plt.clf() 384 plt.suptitle(pomegas_str[i-1] + ", e = %.2f, n = %i" % (e,n)) 385 plt.subplot(211) 386 plt.plot(M, radial_velocity_fourier_series(default_K, M, e, pomega, phase, order=n), 'r') 387 plt.plot(M, radial_velocity_from_M(default_K, M, e, pomega), 'k--') 388 plt.axis([range_min,range_max,-(default_K+1.),default_K+1.]) 389 plt.subplot(212) 390 plt.plot(M, radial_velocity_from_M(default_K, M, e, pomega) \ 391 - radial_velocity_fourier_series(default_K, M, e, pomega, phase, order=n), 'k') 392 plt.savefig("celestial_mechanics_plots/pomega_%i_e_%.2f_n_%i.png" % (i,e,n)) 393 i += 1 394