Commit 0a56cdbb authored by Olaf's avatar Olaf

Add second approximation

parent fff77f8e
......@@ -27,9 +27,11 @@ def extract_data(bodies, D):
masses = []
positions = []
velocities = []
n = 0
for i in range(len(bodies)):
if not isinstance(bodies[i], Asteroids):
masses.append(bodies[i].mass)
n += 1
if bodies[i].geo == 'cylindrical':
r = bodies[i].pos[0]
......@@ -80,7 +82,7 @@ def extract_data(bodies, D):
velocities = np.asarray(velocities).reshape(N, D)
# return
return positions, velocities, masses
return positions, velocities, masses, n
def force_calc(mass, positions, D):
......@@ -170,7 +172,59 @@ def position(x0, v0, f, h):
return x1
def dynamics(bodies, D, h, t_max, barnes_hut, theta):
def force_asteroids(mass, asteroids, positions, D):
"""
Calculates all forces exerted on each body using "brute force" algorithm.
Parameters:
-----------
mass: np.array([n,1])
masses of all planets and the sun
asteroids: np.array([N,D])
positions of all asteroids
positions: np.array([n,D])
positions of all planets and the sun
D: float
number of spatial dimensions
Returns:
--------
F: np.array([N,D])
force exerted on each body
"""
n = mass.shape[0]
N = asteroids.shape[0]
m = mass * np.ones((N,n))
x = asteroids*np.ones([n,N,D])
x = np.transpose(x, (1,0,2)) # np.array([N,n,D])
dx = positions - x
dr = np.linalg.norm(dx, axis = 2)
F_mag = G*m/(dr**3)
F = np.transpose(dx, (2,0,1))*F_mag
F = np.sum(F, axis = 2) # All forces on all particles ([D,N] matrix)
F = F.transpose() # Transpose matrix to obtain desired [N,D] matrix
# return
return F
def force_approx(mass, positions, D, n):
"""
mass: all masses
positions: all positions
D: dimensions
n: all non-asteroid bodies
"""
Force = np.zeros([mass.shape[0],D])
Force[:n] = force_calc(mass[:n], positions[:n], D)
Force[n:] = force_asteroids(mass[:n], positions[n:], positions[:n], D)
# return
return Force
def dynamics(bodies, D, h, t_max, barnes_hut, theta, approximation):
"""
Calculates the dynamics (trajectories) of all bodies
......@@ -198,8 +252,11 @@ def dynamics(bodies, D, h, t_max, barnes_hut, theta):
times: list
all times of the simulated points of the trajectories
"""
if barnes_hut and approximation:
raise ValueError("Can't combine Barnes-Hut with our approximation. Set at least one to false")
T = round(t_max/h)
positions, velo, mass = extract_data(bodies, D)
positions, velo, mass, n = extract_data(bodies, D)
if barnes_hut and theta != 0:
force0 = force_barneshut(positions, mass, theta)
else:
......@@ -207,27 +264,34 @@ def dynamics(bodies, D, h, t_max, barnes_hut, theta):
all_pos = np.zeros((T+1, positions.shape[0], positions.shape[1]))
times = [0]
E_kin = []
all_pos[0] = positions
if barnes_hut and theta != 0:
if barnes_hut and theta != 0 and not approximation:
for i in range(1,T+1):
all_pos[i] = position(all_pos[i-1], velo, force0, h)
force1 = force_barneshut(all_pos[i], mass, theta)
velo = velocity(velo, force0, force1, h)
force0 = force1
E_kin = 0.5*mass*(np.linalg.norm(velo, axis=1)**2)
E_kin.append(0.5*mass*(np.linalg.norm(velo, axis=1)**2))
times.append(h*i)
elif approximation and not barnes_hut:
for i in range(1,T+1):
all_pos[i] = position(all_pos[i-1], velo, force0, h)
force1 = force_approx(mass, all_pos[i], D, n)
velo = velocity(velo, force0, force1, h)
force0 = force1
E_kin.append(0.5*mass*(np.linalg.norm(velo, axis=1)**2))
times.append(i*h)
else:
for i in range(1,T+1):
all_pos[i] = position(all_pos[i-1], velo, force0, h)
force1 = force_calc(mass, all_pos[i], D)
velo = velocity(velo, force0, force1, h)
force0 = force1
E_kin = 0.5*mass*(np.linalg.norm(velo, axis=1)**2)
E_kin.append(0.5*mass*(np.linalg.norm(velo, axis=1)**2))
times.append(h*i)
# return
return all_pos, velo, times, E_kin
\ No newline at end of file
......@@ -414,6 +414,11 @@ class Settings:
self.theta = sets['Theta']
except:
self.theta = 0.5
try:
self.approx = sets['approx']
except:
self.approx = False
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment