getting the structure set up

src/orbits/body.py

@@ -0,0 +1,30 @@
+from enum import Enum
+import numpy as np
+
+from ..units import Position, Velocity, Mass, Acceleration, Force
+
+class Body:
+    """
+    Base class for any orbital body
+    """
+    def __init__(self, X: Position, V: Velocity, m: Mass):
+        """
+        x (position) and v (velocity) 
+        """
+        self.X = X
+        self.V = V
+        self.A = Acceleration([0,0,0])
+        self.m = m
+
+    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
+    
+

src/orbits/simulator.py

@@ -0,0 +1,82 @@
+from pathlib import Path
+
+from body import Body
+from ..units import Position, Velocity, Mass, Acceleration, Force
+
+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: Path
+                overwrite_output: bool = False
+    ):
+        if output_file.exists() and not overwrite_output:
+            raise FileExistsError(f"File {output_file} exists and overwrite flag not given.")
+        
+        self.output_file = output_file
+        self.bodies = bodies
+        self.step_size = step_size
+        self.steps_to_take = steps_to_take
+        self.steps_per_save = steps_per_save
+
+        if output_file.exists() and overwrite_output:
+            print(f"Warning! Overwriting file: {output_file}")
+
+        _save_body_masses_to_file()
+        _checkpoint()
+
+    
+        self.current_step = 0
+
+    @classmethod
+    def from_checkpoint(cls, input_file: Path):
+
+
+    def _save_body_masses_to_file(self):
+        masses_str = ' '.join([
+            str(body.m) for body in bodies
+        ]) + '\n'
+        #if saving masses, we are always starting a new file,
+        #this will overwrite file.
+        self.output_file.write_text(masses_str)
+    
+    def _checkpoint(self):
+        """
+        Two things - save high precision last checkpoint for resuming
+        then save lower precision text for trajectories
+        """
+        body_X_V_np = np.array([
+            [body.X, body.Y] for body in self.bodies
+        ])
+        body_ints_np = np.array([
+            body.m for body in self.bodies
+        ])
+        body_m_np.append(self.step_size, self.current_step)
+        np.savez("last_checkpoint.npz", 
+                array=body_X_V_np,
+                ints = body_m_np)
+

src/units.py

@@ -0,0 +1,16 @@
+import numpy as np
+
+class Position(np.array):
+    pass
+
+class Velocity(np.array):
+    pass
+
+class Acceleration(np.array):
+    pass
+
+class Force(np.array):
+    pass
+
+class Mass(int):
+    pass