Add orb_params function

pyproject.toml

@@ -17,15 +17,14 @@ classifiers = [
     "Development Status :: 5 - Production/Stable",
     "Natural Language :: English",
     "Operating System :: POSIX",
-    "Operating System :: Microsoft :: Windows",
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Topic :: Communications",
     "Topic :: Software Development :: Libraries",
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
-dependencies = ['numpy', 'scipy']
+dependencies = ['numpy', 'scipy', 'pynbody']
 
 [project.urls]
 "Homepage" = "https://github.com/spencerw/keplercalc"

src/keplercalc.py

@@ -4,6 +4,48 @@
 def hello(name):
     return 'Hello ' + name + '!'
 
+def orb_params(snap, isHelio=False, mCentral=1.0):
+    """
+    Takes a Pynbody snapshot of particles and calculates gives them fields
+    corresponding to Kepler orbital elements. Assumes that the central star
+    is particle #0 and that the positions and velocities are in barycentric
+    coordinates.
+
+    Parameters
+    ----------
+    snap: SimArray
+        A snapshot of the particles
+    isHelio: boolean
+        Skip frame transformation if the snap is already in heliocentric coordinates
+    mCentral: float
+        Mass of central star in simulation units. Need to provide if no star particle in snap
+
+    Returns
+    -------
+    SimArray
+        A copy of 'snap' with additional fields 'a', 'ecc', 'inc', 'asc_node',
+        'omega' and 'M' added.
+    """
+
+    x = np.array(snap.d['pos'])
+    x_h = x[1:] - x[0]
+    v = np.array(snap.d['vel'])
+    v_h = v[1:] - v[0]
+    m1 = np.array(snap['mass'][0])
+    pl = snap[1:]
+    m2 = np.array(pl['mass'])
+
+    if isHelio:
+        x_h = x
+        v_h = v
+        pl = snap
+        m1 = mCentral
+
+    pl['a'], pl['e'], pl['inc'], pl['asc_node'], pl['omega'], pl['M'] \
+        = cart2kep(x_h[:,0], x_h[:,1], x_h[:,2], v_h[:,0], v_h[:,1], v_h[:,2], m1, m2)
+
+    return pl
+
 def cart2kep(X, Y, Z, vx, vy, vz, m1, m2):
     """
     Convert an array of cartesian positions and velocities

tests/test_keplercalc.py

@@ -1,6 +1,8 @@
 from unittest import TestCase
 import numpy as np
-from keplercalc import hello, kep2cart, cart2kep
+import pynbody as pb
+import os
+from keplercalc import hello, kep2cart, cart2kep, orb_params
 
 class TestKeplercalc(TestCase):
     def test_hello1(self):
@@ -46,4 +48,18 @@ def test_cart2kep2cart(self):
         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)
+        self.assertTrue(np.fabs(vz1 - vz) < tol)
+
+    def test_snap(self):
+        tol = 1e-8
+
+        path = os.path.dirname(os.path.abspath(__file__))
+        snap = pb.load(path + '/test.ic')
+        pl = orb_params(snap)
+
+        self.assertTrue(np.mean(pl['a']) - 2.96233536 < tol)
+        self.assertTrue(np.mean(pl['e']) - 0.06809352 < tol)
+        self.assertTrue(np.mean(pl['inc']) - 0.005003 < tol)
+        self.assertTrue(np.mean(pl['asc_node']) - 3.15458748 < tol)
+        self.assertTrue(np.mean(pl['omega']) - 3.13253456)
+        self.assertTrue(np.mean(pl['M']) - 3.13972053 < tol)