-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add transformation functions and test
- Loading branch information
Showing
6 changed files
with
358 additions
and
9 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,3 +1,5 @@ | ||
runtest.sh | ||
|
||
# Byte-compiled / optimized / DLL files | ||
__pycache__/ | ||
*.py[cod] | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,4 +1,264 @@ | ||
class Keplercalc: | ||
@classmethod | ||
def hello(self, name): | ||
return 'Hello ' + name + '!' | ||
import numpy as np | ||
from mathhelpers import cross, dot, nr, PQW | ||
|
||
def hello(name): | ||
return 'Hello ' + name + '!' | ||
|
||
def cart2kep(X, Y, Z, vx, vy, vz, m1, m2): | ||
""" | ||
Convert an array of cartesian positions and velocities | ||
(in heliocentric coordinates) to an array of kepler orbital | ||
elements. Units are such that G=1. This function is | ||
fully vectorized. | ||
Parameters | ||
---------- | ||
X: numpy array of floats | ||
X position of body | ||
Y: numpy array of floats | ||
Y position of body | ||
Z: numpy array of floats | ||
Z position of body | ||
vx: numpy array of floats | ||
x velocity of body | ||
vy: numpy array of floats | ||
y velocity of body | ||
vz: numpy array of floats | ||
z velocity of body | ||
m1: float | ||
mass of central body | ||
m2: numpy array of floats | ||
mass of orbiting body | ||
Returns | ||
------- | ||
numpy array of floats | ||
a, e, inc, asc_node, omega, M - Semimajor axis, eccentricity | ||
inclination, longitude of ascending node, longitude of | ||
perihelion and mean anomaly of orbiting body | ||
""" | ||
|
||
mu = m1 + m2 | ||
magr = np.sqrt(X ** 2. + Y ** 2. + Z ** 2.) | ||
|
||
hx, hy, hz = cross(X, Y, Z, vx, vy, vz) | ||
magh = np.sqrt(hx ** 2. + hy ** 2. + hz ** 2.) | ||
tmpx, tmpy, tmpz = cross(vx, vy, vz, hx, hy, hz) | ||
evecx = tmpx / mu - X / magr | ||
evecy = tmpy / mu - Y / magr | ||
evecz = tmpz / mu - Z / magr | ||
|
||
e = np.sqrt(evecx ** 2. + evecy ** 2. + evecz ** 2.) | ||
|
||
a = dot(hx, hy, hz, hx, hy, hz) / (mu * (1. - e ** 2.)) | ||
|
||
ivec = [1., 0., 0.] | ||
jvec = [0., 1., 0.] | ||
kvec = [0., 0., 1.] | ||
|
||
inc = np.arccos(dot(kvec[0], kvec[1], kvec[2], hx, hy, hz) / magh) | ||
|
||
nx, ny, nz = cross(kvec[0], kvec[1], kvec[2], hx, hy, hz) | ||
nmag = np.sqrt(nx ** 2. + ny ** 2. + nz ** 2.) | ||
asc_node = np.where(inc == 0., 0., np.arccos(dot(ivec[0], ivec[1], ivec[2], nx, ny, nz) / nmag)) | ||
asc_node[dot(nx, ny, nz, jvec[0], jvec[1], jvec[2]) < 0.] = 2. * np.pi - asc_node[ | ||
dot(nx, ny, nz, jvec[0], jvec[1], jvec[2]) < 0.] | ||
|
||
omega = np.where(inc == 0., np.arctan2(evecy / e, evecx / e), | ||
np.arccos(dot(nx, ny, nz, evecx, evecy, evecz) / (nmag * e))) | ||
omega[dot(evecx, evecy, evecz, kvec[0], kvec[1], kvec[2]) < 0.] = 2. * np.pi - omega[ | ||
dot(evecx, evecy, evecz, kvec[0], kvec[1], kvec[2]) < 0.] | ||
|
||
theta = np.arccos(dot(evecx, evecy, evecz, X, Y, Z) / (e * magr)) | ||
theta = np.where(dot(X, Y, Z, vx, vy, vz) < 0., 2 * np.pi - theta, theta) | ||
|
||
E = np.arccos((e + np.cos(theta)) / (1 + e * np.cos(theta))) | ||
E = np.where(np.logical_and(theta > np.pi, theta < 2 * np.pi), 2 * np.pi - E, theta) | ||
M = E - e * np.sin(E) | ||
|
||
return a, e, inc, asc_node % (2 * np.pi), omega, M | ||
|
||
def kep2cart(sma, ecc, inc, Omega, omega, M, mass, m_central): | ||
""" | ||
Convert a single set of kepler orbital elements into cartesian | ||
positions and velocities. Units are such that G=1. | ||
Parameters | ||
---------- | ||
sma: float | ||
semi-major axis body | ||
ecc: float | ||
eccentricity of body | ||
inc: float | ||
inclination of body | ||
Omega: float | ||
longitude of ascending node of body | ||
omega: float | ||
longitude of perihelion of body | ||
M: float | ||
Mean anomaly of body | ||
mass: float | ||
mass of body | ||
m_central: float | ||
mass of central body | ||
Returns | ||
------- | ||
float | ||
X, Y, Z, vx, vy, vz - Cartesian positions and velocities | ||
of the bodies | ||
""" | ||
|
||
if inc == 0.: | ||
Omega = 0. | ||
if ecc == 0.: | ||
omega = 0. | ||
E = nr(M, ecc) | ||
|
||
X = sma * (np.cos(E) - ecc) | ||
Y = sma * np.sqrt(1. - ecc**2.) * np.sin(E) | ||
mu = m_central + mass | ||
n = np.sqrt(mu / sma**3.) | ||
Edot = n / (1. - ecc * np.cos(E)) | ||
Vx = - sma * np.sin(E) * Edot | ||
Vy = sma * np.sqrt(1. - ecc**2.) * Edot * np.cos(E) | ||
|
||
Px, Py, Pz, Qx, Qy, Qz = PQW(Omega, omega, inc) | ||
|
||
# Rotate Positions | ||
x = X * Px + Y * Qx | ||
y = X * Py + Y * Qy | ||
z = X * Pz + Y * Qz | ||
|
||
# Rotate Velocities | ||
vx = Vx * Px + Vy * Qx | ||
vy = Vx * Py + Vy * Qy | ||
vz = Vx * Pz + Vy * Qz | ||
|
||
return x, y, z, vx, vy, vz | ||
|
||
def kep2poinc(a, e, i, omega, Omega, M , m1, m2): | ||
""" | ||
Convert kepler orbital elements into Poincaire variables. | ||
Can accept either single values or lists of coordinates | ||
Parameters | ||
---------- | ||
a: float | ||
semi-major axis body | ||
ecc: float | ||
eccentricity of body | ||
inc: float | ||
inclination of body | ||
Omega: float | ||
longitude of ascending node of body | ||
omega: float | ||
longitude of perihelion of body | ||
M: float | ||
Mean anomaly of body | ||
mass: float | ||
mass of body | ||
m_central: float | ||
mass of central body | ||
Returns | ||
------- | ||
float | ||
lam, gam, z, Lam, Gam, Z - Poincaire coordinates | ||
""" | ||
|
||
mu = m1 + m2 | ||
mustar = m1*m2/(m1 + m2) | ||
|
||
lam = M + omega + Omega | ||
gam = -omega - Omega | ||
z = -Omega | ||
Lam = mustar*np.sqrt(mu*a) | ||
Gam = mustar*np.sqrt(mu*a)*(1 - np.sqrt(1 - e**2)) | ||
Z = mustar*np.sqrt(mu*a*np.sqrt(1 - e**2))*(1 - np.cos(i)) | ||
|
||
return lam, gam, z, Lam, Gam, Z | ||
|
||
def kep2del(a, e, i, omega, Omega, M , m1, m2): | ||
""" | ||
Convert kepler orbital elements into Delunay variables. | ||
Can accept either single values or lists of coordinates | ||
Parameters | ||
---------- | ||
a: float | ||
semi-major axis body | ||
ecc: float | ||
eccentricity of body | ||
inc: float | ||
inclination of body | ||
Omega: float | ||
longitude of ascending node of body | ||
omega: float | ||
longitude of perihelion of body | ||
M: float | ||
Mean anomaly of body | ||
mass: float | ||
mass of body | ||
m_central: float | ||
mass of central body | ||
Returns | ||
------- | ||
float | ||
l, g, h, L, G, H - Delunay coordinates | ||
""" | ||
|
||
L = np.sqrt((m1 + m2)*a) | ||
G = L*np.sqrt(1 - e**2) | ||
H = G*np.cos(i) | ||
lv = M | ||
g = omega | ||
h = Omega | ||
|
||
return lv, g, h, L, G, H | ||
|
||
def kep2mdel(a, e, i, omega, Omega, M, m1, m2): | ||
""" | ||
Convert kepler orbital elements into modified Delunay variables. | ||
Can accept either single values or lists of coordinates | ||
Parameters | ||
---------- | ||
a: float | ||
semi-major axis body | ||
ecc: float | ||
eccentricity of body | ||
inc: float | ||
inclination of body | ||
Omega: float | ||
longitude of ascending node of body | ||
omega: float | ||
longitude of perihelion of body | ||
M: float | ||
Mean anomaly of body | ||
mass: float | ||
mass of body | ||
m_central: float | ||
mass of central body | ||
Returns | ||
------- | ||
float | ||
Lam, P, Q, lam, p, q - Delunay coordinates | ||
""" | ||
|
||
L = np.sqrt((m1 + m2)*a) | ||
G = L*np.sqrt(1 - e**2) | ||
H = G*np.cos(i) | ||
lv = M | ||
g = omega | ||
h = Omega | ||
|
||
Lam = L | ||
P = L - G | ||
Q = G - H | ||
lam = lv + g + h | ||
p = -g - h | ||
q = -h | ||
return Lam, P, Q, lam, p, q |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,45 @@ | ||
import numpy as np | ||
from scipy import optimize | ||
|
||
# Vectorized cross and dot product functions | ||
def cross(x1,y1,z1,x2,y2,z2): | ||
xc = y1 * z2 - z1 * y2 | ||
yc = z1 * x2 - x1 * z2 | ||
zc = x1 * y2 - y1 * x2 | ||
return xc, yc, zc | ||
|
||
def dot(x1,y1,z1,x2,y2,z2): | ||
return x1*x2+y1*y2+z1*z2 | ||
|
||
# Solve for true anomaly using newton-raphson iteration | ||
# Works with a single value or a list of mean anomalies | ||
def nr(M, ecc): | ||
def kep(E, M, ecc): | ||
return E - (ecc*np.sin(E)) - M | ||
|
||
if isinstance(M, list): | ||
return optimize.newton(kep, np.ones(len(M)), args=(M, ecc)) | ||
else: | ||
return optimize.newton(kep, 1.0, args=(M, ecc)) | ||
|
||
def PQW(Omega, omega, inc): | ||
""" | ||
Rotation Matrix Components (Orbit Frame => Inertial Frame) | ||
http://biomathman.com/pair/KeplerElements.pdf (End) | ||
http://astro.geo.tu-dresden.de/~klioner/celmech.pdf (Eqn. 2.30) | ||
NB: Shapiro Notation Tricky; Multiplies x = R.T * X, Dresden x = R * X | ||
""" | ||
|
||
Px = np.cos(omega) * np.cos(Omega) - \ | ||
np.sin(omega) * np.cos(inc) * np.sin(Omega) | ||
Py = np.cos(omega) * np.sin(Omega) + \ | ||
np.sin(omega) * np.cos(inc) * np.cos(Omega) | ||
Pz = np.sin(omega) * np.sin(inc) | ||
|
||
Qx = - np.sin(omega) * np.cos(Omega) - \ | ||
np.cos(omega) * np.cos(inc) * np.sin(Omega) | ||
Qy = - np.sin(omega) * np.sin(Omega) + \ | ||
np.cos(omega) * np.cos(inc) * np.cos(Omega) | ||
Qz = np.sin(inc) * np.cos(omega) | ||
|
||
return Px, Py, Pz, Qx, Qy, Qz |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,7 +1,49 @@ | ||
from unittest import TestCase | ||
from keplercalc import Keplercalc | ||
import numpy as np | ||
from keplercalc import hello, kep2cart, cart2kep | ||
|
||
class TestKeplercalc(TestCase): | ||
def test_hello1(self): | ||
output = Keplercalc.hello('Spencer') | ||
output = hello('Spencer') | ||
self.assertEqual(output, 'Hello Spencer!') | ||
|
||
def test_cart2kep2cart(self): | ||
tol = 1e-10 | ||
|
||
# Earth orbiting Sun, slightly inclined so angles are defined | ||
m1, m2 = 1, 1e-20 | ||
|
||
a = 1 | ||
e = 0.05 | ||
inc = 0.1 | ||
asc_node = np.pi | ||
omega = np.pi | ||
M = np.pi | ||
|
||
X, Y, Z, vx, vy, vz = kep2cart(a, e, inc, asc_node, omega, M, m1, m2) | ||
|
||
self.assertTrue(np.fabs(X - -1.05) < tol) | ||
self.assertTrue(np.fabs(Y - 3.782338790704024e-16) < tol) | ||
self.assertTrue(np.fabs(Z - -2.5048146051777413e-17) < tol) | ||
self.assertTrue(np.fabs(vx - -3.490253699036788e-16) < tol) | ||
self.assertTrue(np.fabs(vy - -0.9464377445249709) < tol) | ||
self.assertTrue(np.fabs(vz - 0.09496052074620637) < tol) | ||
|
||
a, e, inc, asc_node, omega, M = cart2kep(X, Y, Z, vx, vy, vz, m1, m2) | ||
|
||
self.assertTrue(np.fabs(a - 1) < tol) | ||
self.assertTrue(np.fabs(e - 0.05) < tol) | ||
self.assertTrue(np.fabs(inc - 0.1) < tol) | ||
self.assertTrue(np.fabs(asc_node - np.pi) < tol) | ||
self.assertTrue(np.fabs(omega - np.pi) < tol) | ||
self.assertTrue(np.fabs(M - np.pi) < tol) | ||
|
||
# Now try converting back to cartesian | ||
X1, Y1, Z1, vx1, vy1, vz1 = kep2cart(a, e, inc, asc_node, omega, M, m1, m2) | ||
|
||
self.assertTrue(np.fabs(X1 - X) < tol) | ||
self.assertTrue(np.fabs(Y1 - Y) < tol) | ||
self.assertTrue(np.fabs(Z1 - Z) < tol) | ||
self.assertTrue(np.fabs(vx1 - vx) < tol) | ||
self.assertTrue(np.fabs(vy1 - vy) < tol) | ||
self.assertTrue(np.fabs(vz1 - vz) < tol) |