a simulator than can... step!

orbiter/orbits/calc.py

@@ -0,0 +1,4 @@
+from scipy.spatial import distance_matrix
+
+def calculate_distances(positions):
+    return distance_matrix(positions, positions)

orbiter/orbits/simulator.py

@@ -1,8 +1,11 @@
 from pathlib import Path
 import numpy as np
+from scipy.spatial import distance_matrix
 
 from .body import Body
 
+G = 1
+
 class Simulator:
     """
     Everything a simulator needs to run:
@@ -31,14 +34,14 @@ class Simulator:
                 bodies: list[Body],
                 step_size: float,
                 steps_per_save: int,
-                output_file: Path,
+                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.output_file = output_file
         self.bodies = bodies
         self.step_size = step_size
         self.steps_per_save = steps_per_save
@@ -47,7 +50,6 @@ class Simulator:
         if output_file.exists() and overwrite_output:
             print(f"Warning! Overwriting file: {output_file}")
 
-        #self._save_body_masses_to_file()
         self._checkpoint()
 
     
@@ -101,3 +103,30 @@ class Simulator:
                 masses=body_m_np,
                 steps=stepsz_n_np)
 
+    def run(self, steps):
+        self.calculate_forces()
+        self.move_bodies()
+        if (steps % self.steps_per_save == 0):
+            self._checkpoint()
+
+    def calculate_forces(self):
+        positions = [
+            body.X for body in self.bodies
+        ]
+        dists = distance_matrix(positions, positions)
+        for i in range(len(self.bodies)):
+            force_sum = np.array([0.0,0.0,0.0])
+            for j in range(len(self.bodies)):
+                vec = self.bodies[i].X - self.bodies[j].X
+                print(vec)
+                norm = np.linalg.norm(vec)
+                vec_norm = vec/norm
+                if (i != j):
+                    f = G*self.bodies[i].m*self.bodies[j].m/(dists[i][j]**2)
+                    print(vec_norm, f)
+                    force_sum += vec_norm*f
+            self.bodies[i].A = force_sum/self.bodies[i].m
+
+    def move_bodies(self):
+        for body in self.bodies:
+            body.step(self.step_size)