#!/usr/bin/env python3
from math import *G = 6.67e-11def negate(f):
# Avoids "minus zero" floating point result
return f if f == 0.0 else - fclass V3(object):
# Poor man's 3-vector implementation, just the properties we need for this program
def __init__(self, x, y, z):
self.x = x
self.y = y
self.z = z
def __str__(self):
# For convenient print output
return "({}, {}, {})".format(self.x, self.y, self.z)
def __neg__(self):
# For unary - operator
return V3(negate(self.x), negate(self.y), negate(self.z))
def __add__(self, other):
# For + operator for two vectors
return V3(self.x + other.x, self.y + other.y, self.z + other.z)
def __sub__(self, other):
# For - operator for two vectors
return V3(self.x - other.x, self.y - other.y, self.z - other.z)
def __mul__(self, other):
# For * operator for vector and scalar
if isinstance(other, (int, float)):
return V3(other * self.x, other * self.y, other * self.z)
raise TypeError("Cannot multiply V3 by {}".format(type(other)))
__rmul__ = __mul__ # allows * to work regardless of order of operands
def radius(self):
# For convenience in calculations
return sqrt(self.x**2 + self.y**2 + self.z**2)
def unit(self):
# Unit vector pointing in the same direction as this vector
r = self.radius()
return V3(self.x / r, self.y / r, self.z / r)
def is_zero(self):
# For convenience in calculations
return (self.x == 0.0) and (self.y == 0.0) and (self.z == 0.0)class Run(object):
# Simulation run with given initial conditions
headings = "r_pos v_pos r_neg v_neg r_diff v_diff".split()
def __init__(self, r_pos, v_pos, r_neg, v_neg):
# Initialize time step arrays with initial conditions
self.r_pos = [r_pos]
self.v_pos = [v_pos]
self.r_neg = [r_neg]
self.v_neg = [v_neg]
def increment(self, vec, dvec, dt):
# Convenience operation to add new value to time step array
vec.append(vec[-1] + dt * dvec)
def output_line(self, *fields):
# Convenience for line output
return " ".join(str(f) for f in fields)
def execute(self, steps=10000, dt=0.001, show_output=True, full_output=False):
# Run simulation for given number of time steps
for _ in range(steps):
r_vec = self.r_pos[-1] - self.r_neg[-1]
if r_vec.is_zero():
break # stop iterating if particles collide
u_vec = r_vec.unit()
r = r_vec.radius()
a = (G / r**2) * u_vec
self.increment(self.r_pos, self.v_pos[-1], dt)
self.increment(self.r_neg, self.v_neg[-1], dt)
self.increment(self.v_pos, a, dt)
self.increment(self.v_neg, a, dt)
if not show_output:
return
lines = []
lines.append(self.output_line(*self.headings))
lines.append("")
itemcount = min(steps + 1, len(self.r_pos), len(self.r_neg), len(self.v_pos), len(self.v_neg))
indexes = range(itemcount) if full_output else (0, -1)
for i in indexes:
r_pos = self.r_pos[i]
r_neg = self.r_neg[i]
v_pos = self.v_pos[i]
v_neg = self.v_neg[i]
r_diff = r_neg - r_pos
v_diff = v_neg - v_pos
lines.append(self.output_line(r_pos, v_pos, r_neg, v_neg, r_diff, v_diff))
return linesZERO = V3(0.0, 0.0, 0.0)
INCX = V3(0.000001, 0.0, 0.0)
INCY = V3(0.0, 0.000001, 0.0)initial_datasets = [
# r_pos, v_pos, r_neg, v_neg
(ZERO, ZERO, INCX, ZERO), # base case, no relative velocity
(ZERO, ZERO, INCX, INCX), # moving apart
(ZERO, ZERO, INCX, - INCX), # moving together
(ZERO, ZERO, INCX, INCY), # moving transverse
]if __name__ == '__main__':
for dataset in initial_datasets:
R = Run(*dataset)
for line in R.execute():
print(line)
print("")