option to use quad_precision on kepler elements, also store precision and version field for snapshot

2 files changed, 28 insertions, 13 deletions

diff --git a/Commons/Loaders.py b/Commons/Loaders.py
index 1aa422e..3f5eb67 100644
--- a/Commons/Loaders.py
+++ b/Commons/Loaders.py
@@ -9,20 +9,22 @@ from Structs import Snapshot, Particle
 import numpy as np
 
 class Loader():
-    def __init__(self, nstep, fname, ellipses):
+    def __init__(self, nstep, fname, ellipses, precision):
         self.nstep = nstep
         self.fname = fname
         self.snapshot = Snapshot()
         self.snapshot.nstep = nstep
         self.snapshot.ellipses = ellipses
+        self.snapshot.precision = precision
 
 class GengaIC(Loader):
     pass
 
 class GengaOut(Loader):
-    def __init__(self, nstep=1, ellipses=False, run_name="gasrun"):
+    def __init__(self, nstep=1, ellipses=False, run_name="gasrun", \
+                       precision="double"):
         fname = "Out%s_%012d.dat" % (run_name, nstep)
-        Loader.__init__(self, nstep, fname, ellipses)
+        Loader.__init__(self, nstep, fname, ellipses, precision)
 
     def load(self):
         with open(self.fname, 'r') as f:
@@ -35,7 +37,7 @@ class GengaOut(Loader):
                 if first:
                     self.snapshot.tout = float(line[0])     # yr
                     first = False
-                particle = Particle()
+                particle = Particle(self.snapshot.precision)
                 particle.id = float(line[1])    # -
                 particle.m  = float(line[2])    # Msun
                 particle.r  = float(line[3])    # AU
@@ -52,9 +54,9 @@ class GengaOut(Loader):
             self.snapshot.particles = particles
 
 class SSAscii(Loader):
-    def __init__(self, nstep=1, ellipses=False):
+    def __init__(self, nstep=1, ellipses=False, precision="double"):
         fname = "Out.%012d.dat" % nstep
-        Loader.__init__(self, nstep, fname, ellipses)
+        Loader.__init__(self, nstep, fname, ellipses, precision)
 
     def load(self):
         with open(self.fname, 'r') as f:
@@ -67,7 +69,7 @@ class SSAscii(Loader):
                 if first:
                     self.snapshot.tout = float(line[0]) / 2. / np.pi    # TU -> yr
                     first = False
-                particle = Particle()
+                particle = Particle(self.snapshot.precision)
                 particle.id = float(line[1])    # -
                 particle.m  = float(line[2])    # Msun
                 particle.r  = float(line[3])    # AU
diff --git a/Commons/Structs.py b/Commons/Structs.py
index 5cae38b..ada3608 100644
--- a/Commons/Structs.py
+++ b/Commons/Structs.py
@@ -19,7 +19,7 @@ class Particle():
 
     NB: For G=M=1, 1 Year = 2 Pi (From Kepler's Third Law).
     """
-    def __init__(self):
+    def __init__(self, precision='double'):
         # General
         self.id = None
         # Cartesian
@@ -32,14 +32,25 @@ class Particle():
         self.m, self.r = None, None
         # Ellipse
         self.xell = None; self.yell = None; self.zell = None
+        # Set Precision
+        if precision == 'double' or precision == 'quad':
+            self.precision = precision
+        else:
+            raise Exception("Invalid Kepler Elements Precision Requested.")
+        # File Version
+        self.version = 2
 
     def cart2kep(self):
+        # Set Precision
+        if self.precision == 'quad':
+            x = np.array([self.x, self.y, self.z], dtype=np.float128)
+            v = np.array([self.vx, self.vy, self.vz], dtype=np.float128)
+        elif self.precision == 'double':
+            x = np.array([self.x, self.y, self.z], dtype=np.float64)
+            v = np.array([self.vx, self.vy, self.vz], dtype=np.float64)
+        # Compute Kepler Elements
         self.a, self.ecc, self.inc, self.Omega, self.omega, self.M0 = \
-            cart2kep(np.array([self.x, self.y, self.z], dtype=np.float128), \
-                     np.array([self.vx, self.vy, self.vz], \
-                              dtype=np.float128), \
-                     self.m, \
-                     central_mass=1.0)
+            cart2kep(x, v, self.m, central_mass=1.0)
     def kep2cart(self):
         x, v = \
             kep2cart(self.a, self.ecc, self.inc, \
@@ -71,3 +82,5 @@ class Snapshot():
         self.nstep = None
         self.nparticles = None
         self.ellipses = None
+        self.version = 2
+        self.precision = 'double'