orbiter/orbits/body.py

@@ -0,0 +1,57 @@
+from enum import Enum
+import numpy as np
+
+from ..units import *
+from .calc import format_sig_figs
+
+class Body:
+    """
+    Base class for any orbital body
+    """
+    def __init__(self, X: Position, V: Velocity, m: Mass, name: str = ""):
+        """
+        x (position) and v (velocity) 
+        """
+        self.X = X
+        self.V = V
+        self.A = Acceleration([0,0,0])
+        self.m = m
+        self.name = name
+
+    def save(self):
+        return (self.X, self.V, self.m)
+    
+    @classmethod
+    def load(cls, tup):
+        return cls(tup[0], tup[1], tup[2])
+
+    def step(self, step_size: float):
+        self.X += step_size*self.V
+        self.V += step_size*self.A
+        self.A = HighPrecisionVector([0,0,0])
+    
+    def __str__(self):
+        pos = 'm '.join([format_sig_figs(real_pos(x), 3) for x in self.X])
+        vel = 'm/s '.join([format_sig_figs(real_vel(v), 3) for v in self.V])
+        return str(f"{self.name}: X = {pos}, V = {vel}")
+    
+    def E(self):
+        return self.ke() + self.pe()
+
+    def pe(self):
+        return -self.m/self.dist_from_o()
+
+    def dist_from_o(self):
+        return sum([x**2 for x in self.X]).sqrt()
+
+
+    def ke(self):
+        return Decimal(0.5)*self.m*(self._speed()**2)
+
+    def _speed(self):
+        return sum([v**2 for v in self.V]).sqrt()
+
+    def speed(self):
+        return str(f"{format_sig_figs(real_vel(self._speed),5)}")
+    
+

orbiter/orbits/calc.py

@@ -0,0 +1,84 @@
+import numpy as np
+import sys
+from decimal import Decimal
+from time import time
+
+from ..units import HighPrecisionMatrix
+
+def calculate_distances(positions):
+    N = len(positions)
+    dists = HighPrecisionMatrix(N,N)
+    for i in range(N):
+        dists[i][i] = Decimal(0)
+        for j in range(i+1, N):
+            d = sum([x**2 for x in (positions[i]-positions[j])]).sqrt()
+            dists[i][j] = Decimal(d)
+            dists[j][i] = Decimal(d)
+    return dists
+
+def calculate_distances_np(positions):
+    positions = np.array(positions)
+    diffs = positions[:, np.newaxis] - positions[np.newaxis, :]
+    dists = np.linalg.norm(diffs, axis=-1)
+    np.fill_diagonal(dists, 0)
+    return dists
+
+def print_progress_bar(iteration, total, start_time, length=50):
+    """Prints a progress bar to the console."""
+    percent = (iteration / total) * 100
+    filled_length = int(length * iteration // total)
+    bar = '#' * filled_length + '-' * (length - filled_length)
+    sys.stdout.write(f'\r[{bar}] {percent:.2f}% {int(iteration/(time()-start_time))} steps/s')
+    sys.stdout.flush()
+
+
+def format_sig_figs(value, sig_figs):
+    """Format a number to a specified number of significant figures for printing."""
+    if value == 0:
+        return "0"
+    return f"{value:.{sig_figs-1}e}"
+
+
+
+def plot_points_terminal(vectors, stdscr, scale=500000, grid_width=30, grid_height=30):
+    """Plots multiple points in the terminal, scaled and centered at (0,0)."""
+    stdscr.clear()
+
+    if not vectors:
+        stdscr.addstr(0, 0, "No vectors provided.")
+        stdscr.refresh()
+        return
+
+    # Apply scaling
+    scaled_vectors = {(round(vec[0] / scale), round(vec[1] / scale)) for vec in vectors}
+
+    # Find min and max coordinates
+    min_x = min(vec[0] for vec in scaled_vectors)
+    max_x = max(vec[0] for vec in scaled_vectors)
+    min_y = min(vec[1] for vec in scaled_vectors)
+    max_y = max(vec[1] for vec in scaled_vectors)
+
+    # Center offsets to keep (0,0) in middle
+    center_x = (grid_width // 2) - min_x
+    center_y = (grid_height // 2) - min_y
+
+    # Adjust coordinates for plotting
+    adjusted_vectors = {(vec[0] + center_x, vec[1] + center_y) for vec in scaled_vectors}
+
+    # Ensure grid boundaries
+    max_terminal_y, max_terminal_x = stdscr.getmaxyx()
+    max_x = min(grid_width, max_terminal_x - 5)
+    max_y = min(grid_height, max_terminal_y - 5)
+
+    # Draw grid with points
+    for i in range(grid_height, -1, -1):
+        row = f"{i - center_y:2} | "
+        for j in range(grid_width + 1):
+            row += "● " if (j, i) in adjusted_vectors else ". "
+        stdscr.addstr(max_y - i, 0, row[:max_terminal_x - 1])  # Ensure no overflow
+
+    # Print X-axis labels
+    x_labels = "   " + " ".join(f"{j - center_x:2}" for j in range(max_x + 1))
+    stdscr.addstr(max_y + 1, 0, x_labels[:max_terminal_x - 1])  # Avoid out-of-bounds error
+
+    stdscr.refresh()

orbiter/orbits/simulator.py

@@ -0,0 +1,136 @@
+from pathlib import Path
+import numpy as np
+from decimal import InvalidOperation, DivisionByZero
+from time import time
+
+from .body import Body
+from .calc import *
+from ..units import *
+
+
+
+class Simulator:
+    """
+    Everything a simulator needs to run:
+        - a step size
+        - bodies with initial positions and momenta
+        - number of steps to take
+        - a reset mechanism
+        - a checkpoint mechanism
+        - how often to save?
+        - overwrite output file if exists? default false
+    if output file is a saved checkpoint file, then use the
+    from_checkpoint class method to create the class.
+
+    otherwise if the output file exists, class will not start.
+
+    output is as text:
+    first line of file is list of masses of bodies
+    each subsequent line is list of positions and velocities of each bodies:
+    body1x, body1y, body1vx, body1vy, body2x, body2y, body2vx etc
+
+    For some of these, it should come with sensible defaults with
+    the ability to change
+    """
+    def __init__(
+                self,
+                bodies: list[Body],
+                step_size: float,
+                steps_per_save: int,
+                output_file: str,
+                current_step: int = 0,
+                overwrite_output: bool = False
+    ):
+        self.output_file = Path(output_file)
+        if output_file.exists() and not overwrite_output:
+            raise FileExistsError(f"File {output_file} exists and overwrite flag not given.")
+        
+        self.bodies = bodies
+        self.step_size = step_size
+        self.steps_per_save = steps_per_save
+        self.current_step = current_step
+
+        if output_file.exists() and overwrite_output:
+            print(f"Warning! Overwriting file: {output_file}")
+
+        self._checkpoint()
+
+    
+    @classmethod
+    def from_checkpoint(cls, output_file: Path):
+        data = np.load("last_checkpoint.npz")
+        positions = data["positions"]
+        velocities = data["velocities"]
+        masses = data["masses"]
+        step_size = data["steps"][0]
+        current_step = data["steps"][1]
+        steps_per_save = data["steps"][2]
+
+        bodies = [
+            Body(val[0], val[1], val[2]) for val in zip(
+                positions, velocities, masses
+            )
+        ]
+        return cls(
+            bodies, 
+            step_size, 
+            steps_per_save,
+            output_file,
+            current_step,
+        )
+
+    
+    def _checkpoint(self):
+        """
+        Two things - save high precision last checkpoint for resuming
+        then save lower precision text for trajectories
+        """
+        body_X_np = np.array([
+            body.X for body in self.bodies
+        ])
+        body_V_np = np.array([
+            body.V for body in self.bodies
+        ])
+        body_m_np = np.array([
+            body.m for body in self.bodies
+        ])
+        stepsz_n_np = np.array([
+            self.step_size,
+            self.current_step,
+            self.steps_per_save
+        ])
+
+        np.savez("last_checkpoint.npz", 
+                positions=body_X_np,
+                velocities=body_V_np,
+                masses=body_m_np,
+                steps=stepsz_n_np)
+
+    def run(self, steps):
+        time_start = time()
+        for i in range(steps):
+            self.calculate_forces()
+            self.move_bodies()
+            self.current_step += 1
+            if (self.current_step % self.steps_per_save == 0):
+                print_progress_bar(i, steps, time_start)
+                #self._checkpoint()
+
+    def calculate_forces(self):
+        positions = [
+            body.X for body in self.bodies
+        ]
+        dists = calculate_distances_np(positions)
+        for i in range(len(self.bodies)):
+            for j in range(i, len(self.bodies)):
+                if i == j:
+                    continue
+                vec = self.bodies[i].X - self.bodies[j].X
+                f = vec/(dists[i][j]**3)
+                self.bodies[i].A += f*self.bodies[j].m
+                self.bodies[j].A += f*self.bodies[i].m
+
+
+    def move_bodies(self):
+        for body in self.bodies:
+            body.step(self.step_size)

orbiter/units.py

@@ -0,0 +1,74 @@
+import numpy as np
+from decimal import Decimal
+
+#Declare these as units to make code clearer
+class HighPrecisionVector(np.ndarray):
+    def __new__(cls, input_array, *args, **kwargs):
+        decimal_array = [Decimal(i) for i in input_array]
+        obj = np.asarray(decimal_array).view(cls)
+        return obj
+
+class HighPrecisionMatrix(np.ndarray):
+    def __new__(cls, dim1, dim2, *args, **kwargs):
+        decimal_array = [Decimal(0) for _ in range(dim1)]
+        decimal_matrix = [decimal_array for _ in range(dim2)]
+        obj = np.asarray(decimal_matrix).view(cls)
+        return obj
+
+class Position(HighPrecisionVector):
+    pass
+class Velocity(HighPrecisionVector):
+    pass
+class Acceleration(HighPrecisionVector):
+    pass
+
+Mass = Decimal
+
+####################
+# USEFUL CONSTANTS
+####################
+EARTH_MASS = Decimal(5972 * 10**21) #kg
+EARTH_RADIUS = Decimal(6378 * 10**3) #meters
+EARTH_ORBITAL_VELOCITY = Decimal(29780) # m/s
+AU = Decimal(149_597_870_700) #meters
+
+MOON_MASS = Decimal(734767309 * 10**14)
+MOON_ORBITAL_VELOCITY = Decimal(1022) #m/s relative to earth
+
+SUN_MASS = Decimal(1989 * 10**27) #kg
+SUN_RADIUS = Decimal(6957 * 10**5) #meters
+
+pi_approx = Decimal("3.14159265358979323846264338327950288419716939937510")
+
+
+#NORMALIZING CONSTANTS
+G = Decimal(6.67430e-11)
+r_0 = Decimal(EARTH_RADIUS) #1.496e11
+m_0 = Decimal(5.972e24)
+t_0 = np.sqrt((r_0**3) / (G*m_0))
+
+def norm_pos(pos):
+    return Decimal(pos) / Decimal(r_0)
+
+def real_pos(pos):
+    return pos * Decimal(r_0)
+
+def norm_mass(mass):
+    return Decimal(mass) / Decimal(m_0)
+
+def real_mass(mass):
+    return Decimal(mass) * Decimal(m_0)
+
+def norm_time(time):
+    return Decimal(time) / Decimal(t_0)
+
+def real_time(time):
+    return Decimal(time) * Decimal(t_0)
+
+def norm_vel(vel):
+    return vel / Decimal(r_0/t_0)
+
+def real_vel(vel):
+    return vel * Decimal(r_0/t_0)
+
+