first commit

2 files changed, 127 insertions, 0 deletions

diff --git a/Commons/Loaders.py b/Commons/Loaders.py
new file mode 100644
index 0000000..27a762e
--- /dev/null
+++ b/Commons/Loaders.py
@@ -0,0 +1,76 @@
+from Structs import Snapshot, Particle
+
+class Loader():
+    def __init__(self, nstep, fname):
+        self.nstep = nstep
+        self.fname = fname
+        self.snapshot = Snapshot()
+        self.snapshot.nstep = nstep
+
+class GengaIC(Loader):
+    pass
+
+class GengaOut(Loader):
+    def __init__(self, nstep=1):
+        fname = "Outgasrun_%012d.dat" % nstep
+        Loader.__init__(self, nstep, fname)
+
+    def load(self):
+        with open(self.fname, 'r') as f:
+            lines = f.readlines()
+            first = True
+            particles = []
+            self.snapshot.nparticles = len(lines)
+            for line in lines:
+                line = line.strip().split(" ")
+                if first:
+                    self.snapshot.tout = float(line[0])     # yr
+                    first = False
+                particle = Particle()
+                particle.id = float(line[1])                # -
+                particle.m  = float(line[2])                # Msun
+                particle.r  = float(line[3])                # 
+                particle.x  = float(line[4])
+                particle.y  = float(line[5])
+                particle.z  = float(line[6])
+                particle.vx = float(line[7])
+                particle.vy = float(line[8])
+                particle.vz = float(line[9])
+                particle.cart2kep()
+                particle.compute_ellipse()
+                particles.append(particle)
+            self.snapshot.particles = particles
+
+class SSAscii(Loader):
+    def __init__(self, nstep=1):
+        fname = "Out.%010d.dat" % nstep
+        Loader.__init__(self, nstep, fname)
+
+    def load(self):
+        with open(self.fname, 'r') as f:
+            lines = f.readlines()
+            first = True
+            particles = []
+            self.snapshot.nparticles = len(lines)
+            for line in lines:
+                line = line.strip().split(" ")
+                if first:
+                    self.snapshot.tout = float(line[0]) / 2. / pi       # TU -> yr
+                    first = False
+                particle = Particle()
+                particle.id = float(line[1])
+                particle.m  = float(line[2])
+                particle.r  = float(line[3])
+                particle.x  = float(line[4])
+                particle.y  = float(line[5])
+                particle.z  = float(line[6])
+                particle.vx = float(line[7])
+                particle.vy = float(line[8])
+                particle.vz = float(line[9])
+                particle.cart2kep()
+                particle.compute_ellipse()
+                particles.append(particle)
+            self.snapshot.particles = particles
+
+class SSBinary(Loader):
+    pass
diff --git a/Commons/Structs.py b/Commons/Structs.py
new file mode 100644
index 0000000..3bccc5d
--- /dev/null
+++ b/Commons/Structs.py
@@ -0,0 +1,51 @@
+from kepler_helpers import cart2kep, kep2cart, compute_ellipse
+import numpy as np
+
+class Particle():
+    def __init__(self):
+        # Cartesian
+        self.x, self.y, self.z = None, None, None
+        self.vx, self.vy, self.vz = None, None, None
+        # Keplerian
+        self.a = None; self.ecc = None; self.inc = None
+        self.Omega = None; self.omega = None; self.M0 = None
+        # Other
+        self.m, self.r = None, None
+        # Ellipse
+        self.xell = None; self.yell = None; self.zell = None
+
+    def cart2kep(self):
+        self.a, self.ecc, self.inc, self.Omega, self.omega, self.M0 = \
+            cart2kep(np.array([self.x, self.y, self.z]), \
+                     np.array([self.vx, self.vy, self.vz]), self.m, \
+                     central_mass=1.0)
+    def kep2cart(self):
+        x, v = \
+            kep2cart(self.a, self.ecc, self.inc, \
+                     self.Omega, self.omega, self.M0, self.m, \
+                     central_mass=1.0)
+        self.x = x[0]; self.y = x[1]; self.z = x[2]
+        self.vx = v[0]; self.vy = v[1]; self.vz = v[2]
+
+    def compute_ellipse(self):
+        self.xell, self.yell, self.zell = compute_ellipse(self.a, \
+                                                          self.ecc, \
+                                                          self.inc, \
+                                                          self.Omega, \
+                                                          self.omega)
+
+class Planet(Particle):
+    pass
+
+class Jupiter(Planet):
+    pass
+
+class Saturn(Planet):
+    pass
+
+class Snapshot():
+    def __init__(self):
+        self.particles = []
+        self.tout = None
+        self.nstep = None
+        self.nparticles = None