update docstrings

barnes_hut.py

@@ -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)

movements.py

@@ -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")