Add transformation functions and test

.gitignore

@@ -1,3 +1,5 @@
+runtest.sh
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

pyproject.toml

@@ -25,7 +25,7 @@ classifiers = [
     "Topic :: Software Development :: Libraries",
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
-dependencies = []
+dependencies = ['numpy', 'scipy']
 
 [project.urls]
 "Homepage" = "https://github.com/spencerw/keplercalc"

src/keplercalc.py

@@ -1,4 +1,266 @@
-class Keplercalc:
-    @classmethod
-    def hello(self, name):
-        return 'Hello ' + name + '!'
+import numpy as np
+from mathhelpers import cross, dot, nr, PQW
+from scipy import optimize
+
+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.)
+    magv = np.sqrt(vx ** 2. + vy ** 2. + vz ** 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)
+    l = M
+    g = omega
+    h = Omega
+
+    return l, 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)
+    l = M
+    g = omega
+    h = Omega
+
+    Lam = L
+    P = L - G
+    Q = G - H
+    lam = l + g + h
+    p = -g - h
+    q = -h
+    return Lam, P, Q, lam, p, q

src/mathhelpers.py

@@ -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

tests/test_keplercalc.py

@@ -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)