Update asr.py

asr.py

@@ -1,5 +1,6 @@
 import math # physics is applied mathematics - xkcd.com/435
 import time # timer
+from random import randint # random initial conditions
 #
 # n-body physics simulation
 # copyright 2022 ryland goldman, all rights reserved
@@ -41,17 +42,20 @@
 # version 0.4.2 - 8 nov 2022
 #   better memory management in render script
 #
+# version 0.5.0 - 15 nov 2022
+#   collisions! - they're elastic
 
 # constants, 1 now because units don't exist when there's nothing to compare them to (time/distance/speed/direction/etc. are relative)
-G = 6.7e-11  # gravitational constant
+G = 1  # gravitational constant
 k = 1 # coulumb constant
 E = 1e-100 # softening constant
 t = 1e-7 # time constant
+s = 1 # size constant
 
 # initial conditions
 particles = []
-iterations = 5e4 # iterations of simulation
-frequency = 100 # frequency of recording frames
+iterations = 1e5 # iterations of simulation
+frequency = 1e3 # frequency of recording frames
 
 # structure of a particle has position, velocity, mass, and charge
 class Particle:
@@ -68,12 +72,36 @@ def toArrStr(self):
         return "Particle("+str(self.x)+","+str(self.y)+","+str(self.z)+","+str(self.vx)+","+str(self.vy)+","+str(self.vz)+","+str(self.m)+","+str(self.q)+")"
     def __str__(self):
         print("Particle at (",str(self.x),",",str(self.y),",",str(self.z),") at |v|=(",str(self.vx),",",str(self.vy),",",str(self.vz),") - m=",str(self.m),", q=",str(self.q))
-
+
+def check_collision(p1, p2, r):
+    if r < s: # if the distance is greater than the particle size, a collision occurs
+        p1vx = p1.vx # temporary variables to hold distance
+        p1vy = p1.vy
+        p1vz = p1.vz
+        p2vx = p2.vx
+        p2vy = p2.vy
+        p2vz = p2.vz
+        p1.vx = (p1vx*(p1.m-p2.m)+2*p2.m*p2vx)/(p1.m+p2.m) # formulas for elastic collision
+        p1.vy = (p1vy*(p1.m-p2.m)+2*p2.m*p2vy)/(p1.m+p2.m)
+        p1.vz = (p1vz*(p1.m-p2.m)+2*p2.m*p2vz)/(p1.m+p2.m)
+        p2.vx = (p2vx*(p2.m-p1.m)+2*p1.m*p1vx)/(p1.m+p2.m)
+        p2.vy = (p2vy*(p2.m-p1.m)+2*p1.m*p1vy)/(p1.m+p2.m)
+        p2.vz = (p2vz*(p2.m-p1.m)+2*p1.m*p1vz)/(p1.m+p2.m)
+        # push again
+        while math.sqrt( (p1.x-p2.x)**2 + (p1.y-p2.y)**2 + (p1.z-p2.z)**2) < s:
+            p1.x = p1.x + p1.vx
+            p1.y = p1.y + p1.vy
+            p1.z = p1.z + p1.vz
+            p2.x = p2.x + p2.vx
+            p2.y = p2.y + p2.vy
+            p2.z = p2.z + p2.vz
+
 # calculates interactions between particles, define as a cuda kernel later for testing
 def particle_interaction(p1,p2):
     if p1.x == p2.x and p1.y == p2.y and p1.z == p2.z:
         return p1 # don't have two of the same particles interact
-    r = math.sqrt( (p1.x-p2.x)**2 + (p1.y-p2.y)**2 + (p1.z-p2.z)**2)    # 3d distance formula - sqrt( [x1-x2]^2 + [y1-y2] ^2 + [z1-z2]^2 )
+    r = math.sqrt( (p1.x-p2.x)**2 + (p1.y-p2.y)**2 + (p1.z-p2.z)**2) # 3d distance formula - sqrt( [x1-x2]^2 + [y1-y2] ^2 + [z1-z2]^2 )
+    check_collision(p1, p2, r)
     force = (G*p1.m*p2.m/((r+E)**2)) + (-k*p1.q*p2.q/((r+E)**2))        # gravitational force (Gmm/r^2) plus electromagnetic force (kqq/r^2)
     acc_p1 = - force/p1.m                                                 # convert to acceleration via newton's 2nd law
 
@@ -83,11 +111,17 @@ def particle_interaction(p1,p2):
     dz = p1.z - p2.z
 
     # calculate angles
-    alpha = math.asin(dy/r)
+    try:
+        alpha = math.asin(dy/r)
+    except ValueError:
+        if dy/r > 0:
+            alpha = math.pi/2
+        else:
+            alpha = -math.pi/2
     if dz == 0:
         beta = math.pi
     else:
-         beta = math.atan(dx/(dz+0.001))
+         beta = math.atan(dx/(dz))
 
     if dx < 0: alpha = -alpha - math.pi
 
@@ -100,17 +134,8 @@ def particle_interaction(p1,p2):
 
 def populate_universe():
     # again, units don't exist
-    particles.append(Particle(1,5,9,1*t,-2*t,-1*t,1e8,0))
-    particles.append(Particle(2,6,2,2*t,1*t,0,1e3,0))
-    particles.append(Particle(3,4,1,-1*t,-1*t,2*t,1e5,0))
-    particles.append(Particle(4,7,3,-2*t,2*t,2*t,1e8,0))
-    particles.append(Particle(5,3,8,1*t,2*t,-1*t,1e7,0))
-    particles.append(Particle(6,8,7,2*t,0,-1*t,1e8,0))
-    particles.append(Particle(7,2,4,-1*t,1*t,2*t,1e5,0))
-    particles.append(Particle(8,9,6,-2*t,1*t,0,1e9,0))
-    particles.append(Particle(9,1,5,1*t,2*t,0,1e8,0))
-    '''particles.append(Particle(0,0,0,0,0,0,2e30,0))
-    particles.append(Particle(1.5e11,0,0,0,0,1e7*t,6e24,0))'''
+    for n in range(20):
+        particles.append(Particle(randint(-100,100), randint(-100,100), randint(-100,100), randint(-10,10)*t, randint(-10,10)*t, randint(-10,10)*t, randint(1,100), randint(-10, 10)))
 
 n = 0 # counter n
 # main loop of program
@@ -149,11 +174,11 @@ def plot_particles(itr_num):
     return rstr
 
 populate_universe()
-s = time.time()
+starttime = time.time()
 main()
-e = time.time()
+endtime = time.time()
 
-print("Processing done in ",e-s," seconds")
+print("Processing done in ",endtime-starttime," seconds")
 
 print("Simulation completed. Preparing postprocessing...")
 fcmd = open("/Users/rylandgoldman/Desktop/Nbody.sh","w")