aboutsummaryrefslogtreecommitdiffstats
path: root/Commons/Structs.py
blob: ada36080c2fa5daef472d03550bb9f9121baa9d0 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
from kepler_helpers import cart2kep, kep2cart, compute_ellipse
import numpy as np

class Particle():
    """
    Semi-Major Axis               - a           - [AU]
    Eccentricity                  - ecc         - []
    Inclination                   - inc         - [rad]
    Longitude of Ascending Node   - Omega       - [rad]
    Argument of Perigee           - omega       - [rad]
    Orbital Phase (Mean Anomaly)  - M0          - [rad]
    XYZ-Position                  - x, y, z     - [AU]
    XYZ-Velocity                  - vx, vy, vz  - [AU/yr]
                                                - [AU/day/0.01720209895]
    Particle Mass                 - m           - [Msun]
    Particle ID                   - id          - []

    To convert velocities to km/s, do *au2km/24.0/3600.0*0.0172020989.

    NB: For G=M=1, 1 Year = 2 Pi (From Kepler's Third Law).
    """
    def __init__(self, precision='double'):
        # General
        self.id = None
        # 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
        # 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(x, v, 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
        self.ellipses = None
        self.version = 2
        self.precision = 'double'