diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/orbits/__init__.py b/src/orbits/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/orbits/body.py b/src/orbits/body.py new file mode 100644 index 0000000..c2b5236 --- /dev/null +++ b/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 + + diff --git a/src/orbits/simulator.py b/src/orbits/simulator.py new file mode 100644 index 0000000..a75249f --- /dev/null +++ b/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) + diff --git a/src/units.py b/src/units.py new file mode 100644 index 0000000..02bddb9 --- /dev/null +++ b/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