Commit 2d2b58c8 by Olaf

### Calculate potential energy

parent 3f614fce
 ... ... @@ -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!