From e1b55dbad2c1c32dc4fa594451df80c35972d3ea Mon Sep 17 00:00:00 2001
From: Michael Bradley <michaelmbrad@gmail.com>
Date: Sat, 7 Oct 2023 21:10:37 -0400
Subject: [PATCH] Add slightly better n_body_matrix

---
 data.py    |  4 ++--
 physics.py | 21 ++++++++++++++++++++-
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/data.py b/data.py
index 478c542..29df928 100644
--- a/data.py
+++ b/data.py
@@ -3,7 +3,7 @@ import matplotlib.cm as cm
 import matplotlib.pyplot as plt
 import numpy as np
 
-from physics import n_body
+from physics import n_body_matrix
 
 
 def parse_csv(filename: str):
@@ -54,7 +54,7 @@ class Animator:
         return self.scat,
 
     def update(self, *_args, **_kwargs):
-        n_body(self.pos, self.vel, self.mass)
+        n_body_matrix(self.pos, self.vel, self.mass)
         self.scat.set_offsets(self.pos)
         return self.scat,
 
diff --git a/physics.py b/physics.py
index 6b3fb0c..ac0b2a5 100644
--- a/physics.py
+++ b/physics.py
@@ -13,5 +13,24 @@ def rotations(a: np.ndarray):
 def n_body(pos: np.ndarray, vel: np.ndarray, mass: np.ndarray):
     for (o_pos, o_mass) in zip(rotations(pos), rotations(mass)):
         dist = o_pos - pos
-        vel += G * (dist / np.linalg.norm(dist, axis=1)[:, np.newaxis]) * o_mass / np.sum(dist ** 2, axis=1)[:, np.newaxis]
+        vel += G * dist * o_mass / (np.linalg.norm(dist, axis=1) ** 3)[:, np.newaxis]
+    pos += vel
+
+
+def n_body_matrix(pos: np.ndarray, vel: np.ndarray, mass: np.ndarray):
+    dist = np.zeros((len(pos) - 1, len(pos), 2))
+    rot_mass = np.zeros((len(mass) - 1, len(mass), 1))
+
+    pos2 = np.concatenate((pos, pos))
+    mass2 = np.concatenate((mass, mass))
+
+    for i in range(1, len(pos)):
+        dist[i - 1] = pos2[i: i + len(pos)] - pos
+        rot_mass[i - 1] = mass2[i: i + len(mass)]
+
+    vel += G * np.sum(
+        dist * rot_mass / (np.linalg.norm(dist, axis=2) ** 3)[:, :, np.newaxis],
+        axis=0
+    )
+
     pos += vel