Commit f2892109 authored by Olaf's avatar Olaf

update docstrings

parent f1892ea6
......@@ -164,6 +164,8 @@ def force_tree(theta, tree, pos):
--------
force: np.array([N,D])
all net forces on all bodies
U: np.array([N,])
Potential energy
"""
force, U = force_node(pos, tree, theta, np.array(range(pos.shape[0])), pos.shape)
......@@ -194,6 +196,8 @@ def force_node(pos, node, theta, nbody, Mshape):
--------
force: np.array([N,D])
forces on all bodies
U: np.array([N,])
Potential energy
"""
force = np.zeros(Mshape)
U = np.zeros(Mshape[0])
......@@ -247,6 +251,8 @@ def force_cal(pos1 , pos2, m):
--------
F: np.array([1,D])
force on the body at pos1
U: np.array([N,])
Potential energy
"""
dx = np.asarray(pos2) - pos1
dr = np.linalg.norm(dx, axis = 1)
......@@ -273,6 +279,8 @@ def force_barneshut(pos, mass, theta):
--------
force: np.array([N,D])
all net forces on all bodies
U: np.array([N,])
Potential energy
"""
tree = create_tree(pos, mass)
force, U = force_tree(theta, tree, pos)
......
......@@ -25,6 +25,8 @@ def extract_data(bodies, D):
velocities of all bodies
masses: np.array([N,1])
masses of all bodies
n: float
number of non-asteroids
"""
masses = []
positions = []
......@@ -91,12 +93,30 @@ def extract_file_data(file):
"""
Extracts data from a file to continue calculations.
Parameters
----------
file : string
Path of the pickle file to read data from
Returns:
pos: all positions
velo: all velocities
mass: all masses
n: number of non-asteroid bodies
t: final time of last calculation
--------
pos: np.array([N,D])
all final postions from previous calculations
velo: np.array([N,D])
all final velocities from previous calculations
mass: np.array([N,])
all masses of all bodies
n: float
number of non-asteroids
t: float
final time of previous calculations
barnes_hut: bool
use the Barnes-Hut algorithm or not
theta: float
ratio in Barnes-Hut algorithm
approximation: bool
use approximation or not
"""
# Check whether file extension is given
if file[-4:] != ".pkl":
......@@ -161,6 +181,8 @@ def force_calc(mass, positions, D):
--------
F: np.array([N,D])
force exerted on each body
U: np.array([N,])
Potential energy
"""
N = mass.shape[0]
m = (mass * np.ones((N,N))).T
......@@ -250,6 +272,8 @@ def force_asteroids(mass, asteroids, positions, D):
--------
F: np.array([N,D])
force exerted on each body
U: np.array([N,])
Potential energy
"""
n = mass.shape[0]
N = asteroids.shape[0]
......@@ -271,10 +295,25 @@ def force_asteroids(mass, asteroids, positions, D):
def force_approx(mass, positions, D, n):
"""
mass: all masses
positions: all positions
D: dimensions
n: all non-asteroid bodies
Combines the forces on asteroids and on non-asteroids.
Parameters:
-----------
mass: np.array([N,])
all masses
positions: np.array([N,D])
all positions
D: float
dimensions
n: float
number of non-asteroid bodies
Returns:
--------
Force: np.array([N,D])
force exerted on each body
U: np.array([N,])
Potential energy
"""
Force = np.zeros([mass.shape[0],D])
U = np.zeros(mass.shape[0])
......@@ -316,12 +355,16 @@ def dynamics(bodies, D, h, t_max, barnes_hut, theta, approximation, j, startsave
Returns:
--------
all_pos: np.array([T+1,N,D])
all positions of all bodies at all timesteps
all_pos_corr: np.array([T+1,N,D])
all positions of all bodies at all timesteps, corrected for the position of the sun
velo: np.array([N,D])
all velocities of all bodies at all timesteps
times: list
all times of the simulated points of the trajectories
E_kin: np.array([T,N])
kinetic energy
E_pot: np.array([T,N])
potential energy
"""
if barnes_hut and approximation:
raise ValueError("Can't combine Barnes-Hut with our approximation. Set at least one to 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