Commit d489d565 authored by Max Capelle's avatar Max Capelle
parents 047ac500 16cf5e80
......@@ -165,10 +165,10 @@ def force_tree(theta, tree, pos):
force: np.array([N,D])
all net forces on all bodies
"""
force = force_node(pos, tree, theta, np.array(range(pos.shape[0])), pos.shape)
force, U = force_node(pos, tree, theta, np.array(range(pos.shape[0])), pos.shape)
# return
return force
return force, U
def force_node(pos, node, theta, nbody, Mshape):
......@@ -196,35 +196,38 @@ def force_node(pos, node, theta, nbody, Mshape):
forces on all bodies
"""
force = np.zeros(Mshape)
U = np.zeros(Mshape[0])
if len(node) == 3:
if nbody.shape[0] == 1 and node[-1] != nbody:
force[nbody] = force_cal(pos, node[0], node[1])
force[nbody], U[nbody] = force_cal(pos, node[0], node[1])
elif nbody.shape[0] != 1:
if nbody[nbody == node[-1]].shape[0] == 0:
force[nbody] = force_cal(pos, node[0], node[1])
force[nbody], U[nbody] = force_cal(pos, node[0], node[1])
else:
nbody_for = np.delete(nbody, np.array(range(pos.shape[0]))[nbody == node[-1]])
pos_for = np.delete(pos, np.array(range(pos.shape[0]))[nbody == node[-1]], 0)
force[nbody_for] = force_cal(pos_for, node[0], node[1])
force[nbody_for], U[nbody_for] = force_cal(pos_for, node[0], node[1])
elif len(node) == 8:
d = np.linalg.norm(pos - np.asarray(node[4]), axis = 1)
nbody_calc = nbody[node[-1]/d < theta]
if nbody_calc.shape[0] != 0:
force[nbody_calc] = force_cal(pos[node[-1]/d < theta], node[4], node[5])
force[nbody_calc], U[nbody_calc] = force_cal(pos[node[-1]/d < theta], node[4], node[5])
pos_cont = pos[node[-1]/d >= theta]
nbody_cont = nbody[node[-1]/d >= theta]
for i in range(len(node) - 4):
force += force_node(pos_cont, node[i], theta, nbody_cont, Mshape)
force1, U1 = force_node(pos_cont, node[i], theta, nbody_cont, Mshape)
force += force1
U += U1
# return
return force
return force, U
def force_cal(pos1 , pos2, m):
......@@ -248,9 +251,10 @@ def force_cal(pos1 , pos2, m):
dx = np.asarray(pos2) - pos1
dr = np.linalg.norm(dx, axis = 1)
F = ((G*m/(dr**3))*(dx.T)).T
U = G*m/dr
# return
return F
return F, U
def force_barneshut(pos, mass, theta):
"""
......@@ -271,7 +275,7 @@ def force_barneshut(pos, mass, theta):
all net forces on all bodies
"""
tree = create_tree(pos, mass)
force = force_tree(theta, tree, pos)
force, U = force_tree(theta, tree, pos)
# return
return force
\ No newline at end of file
return force, -U
\ No newline at end of file
......@@ -3,6 +3,7 @@ import warnings
from predef import *
from barnes_hut import *
import math
import matplotlib.pyplot as plt
def extract_data(bodies, D):
......@@ -105,20 +106,20 @@ def force_calc(mass, positions, D):
force exerted on each body
"""
N = mass.shape[0]
m = (mass * np.ones((N,N))).T
m = (mass * np.ones((N,N))).T
x = positions*np.ones([N,N,D])
dx = np.transpose(x, (1,0,2)) - x
dr_mag = np.linalg.norm(dx, axis = 2)
dr_mag += np.eye(N)
U = -np.sum(G*m*(1 - np.eye(N))/dr_mag, axis = 1)
F_mag = G*m*(1 - np.eye(N))/(dr_mag**3)
F = np.transpose(dx, (2,0,1))*F_mag
F = np.sum(F, axis = 1) # All forces on all particles ([D,N] matrix)
F = F.transpose() # Transpose matrix to obtain desired [N,D] matrix
# return
return F
return F, U
def velocity(v0, f0, f1, h):
......@@ -204,11 +205,12 @@ def force_asteroids(mass, asteroids, positions, D):
dr = np.linalg.norm(dx, axis = 2)
F_mag = G*m/(dr**3)
U = -np.sum(G*m/dr, axis = 1)
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
return F, U
def force_approx(mass, positions, D, n):
"""
......@@ -218,11 +220,12 @@ def force_approx(mass, positions, D, n):
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)
U = np.zeros(mass.shape[0])
Force[:n], U[:n] = force_calc(mass[:n], positions[:n], D)
Force[n:], U[n:] = force_asteroids(mass[:n], positions[n:], positions[:n], D)
# return
return Force
return Force, U
def dynamics(bodies, D, h, t_max, barnes_hut, theta, approximation, j, startsave):
......@@ -266,58 +269,64 @@ def dynamics(bodies, D, h, t_max, barnes_hut, theta, approximation, j, startsave
positions, velo, mass, n = extract_data(bodies, D)
if barnes_hut and theta != 0 and not approximation:
force0 = force_barneshut(positions, mass, theta)
force0, U = force_barneshut(positions, mass, theta)
elif approximation and not barnes_hut:
force0 = force_approx(mass, positions, D, n)
force0, U = force_approx(mass, positions, D, n)
else:
force0 = force_calc(mass, positions, D)
force0, U = force_calc(mass, positions, D)
all_pos = np.zeros((math.ceil((T-startsave+1)/j), positions.shape[0], positions.shape[1]))
E_kin = []
E_kin = np.zeros((math.ceil((T-startsave+1)/j), positions.shape[0]))
E_pot = np.zeros((math.ceil((T-startsave+1)/j), positions.shape[0]))
times = []
index = -1
if startsave == 0:
index += 1
times.append(0)
E_kin.append(0.5*mass*(np.linalg.norm(velo, axis=1)**2))
all_pos[0] = positions
E_kin[index] = 0.5*(np.linalg.norm(velo, axis=1)**2)
E_pot[index] = U
all_pos[index] = positions
if barnes_hut and theta != 0 and not approximation:
for i in range(1,T+1):
positions = position(positions, velo, force0, h)
force1 = force_barneshut(positions, mass, theta)
force1, U = force_barneshut(positions, mass, theta)
velo = velocity(velo, force0, force1, h)
force0 = force1
if i%j == 0 and i>=startsave:
index += 1
all_pos[index] = positions
E_kin.append(0.5*mass*(np.linalg.norm(velo, axis=1)**2))
E_kin[index] = 0.5*(np.linalg.norm(velo, axis=1)**2)
E_pot[index] = U
times.append(h*i)
elif approximation and not barnes_hut:
for i in range(1,T+1):
positions = position(positions, velo, force0, h)
force1 = force_approx(mass, positions, D, n)
force1, U = force_approx(mass, positions, D, n)
velo = velocity(velo, force0, force1, h)
force0 = force1
if i%j == 0 and i>=startsave:
index += 1
all_pos[index] = positions
E_kin.append(0.5*mass*(np.linalg.norm(velo, axis=1)**2))
E_kin[index] = 0.5*(np.linalg.norm(velo, axis=1)**2)
E_pot[index] = U
times.append(h*i)
else:
for i in range(1,T+1):
positions = position(positions, velo, force0, h)
force1 = force_calc(mass, positions, D)
force1, U = force_calc(mass, positions, D)
velo = velocity(velo, force0, force1, h)
force0 = force1
if i%j == 0 and i>=startsave:
index += 1
all_pos[index] = positions
E_kin.append(0.5*mass*(np.linalg.norm(velo, axis=1)**2))
E_kin[index] = 0.5*(np.linalg.norm(velo, axis=1)**2)
E_pot[index] = U
times.append(h*i)
all_pos_corr = np.transpose(np.transpose(all_pos, (1,0,2)) - np.transpose(all_pos, (1,0,2))[0], (1,0,2))
# return
return all_pos_corr, velo, times, E_kin
\ No newline at end of file
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