Merge pull request #89 from NDF-Poli-USP/fixing_docs

Fixing docs

spyro/domains/quadrature.py

@@ -73,7 +73,18 @@ def quadrature_rules(V):
 # Spectral method - Gauss-Lobatto-Legendre rule
 # 1D
 def gauss_lobatto_legendre_line_rule(degree):
-    """Returns GLL quad rule for a given degree in a line"""
+    """Returns GLL quad rule for a given degree in a line
+
+    Parameters
+    ----------
+    degree : int
+        degree of the polynomial
+    
+    Returns
+    -------
+    result : obj
+        quadrature rule
+    """
     fiat_make_rule = FIAT.quadrature.GaussLobattoLegendreQuadratureLineRule
     fiat_rule = fiat_make_rule(FIAT.ufc_simplex(1), degree + 1)
     finat_ps = finat.point_set.GaussLobattoLegendrePointSet
@@ -83,7 +94,20 @@ def gauss_lobatto_legendre_line_rule(degree):
 
 # 3D
 def gauss_lobatto_legendre_cube_rule(dimension, degree):
-    """Returns GLL quad rule for a given degree in a multidimensional space"""
+    """Returns GLL quad rule for a given degree in a multidimensional space
+
+    Parameters
+    ----------
+    dimension : int
+        dimension of the space
+    degree : int
+        degree of the polynomial
+    
+    Returns
+    -------
+    result : obj
+        quadrature rule
+    """
     make_tensor_rule = finat.quadrature.TensorProductQuadratureRule
     result = gauss_lobatto_legendre_line_rule(degree)
     for _ in range(1, dimension):

spyro/domains/space.py

@@ -4,7 +4,22 @@
 def FE_method(mesh, method, degree):
     """Define the finite element method:
     Space discretization - Continuous
-    or Discontinuous Galerkin methods"""
+    or Discontinuous Galerkin methods
+
+    Parameters
+    ----------
+    mesh : obj
+        Firedrake mesh
+    method : str
+        Finite element method
+    degree : int
+        Degree of the finite element method
+    
+    Returns
+    -------
+    element : obj
+        Firedrake finite element
+    """
     cell_geometry = mesh.ufl_cell()
     if method == "CG" or method == "spectral":
         # CG - Continuous Galerkin

spyro/io/io.py

@@ -12,7 +12,18 @@
 
 
 def ensemble_save(func):
-    """Decorator for read and write shots for ensemble parallelism"""
+    """Decorator for read and write shots for ensemble parallelism
+    
+    Parameters
+    ----------
+    func : function
+        Function to be decorated
+        
+    Returns
+    -------
+    wrapper : function
+        Decorated function
+    """
 
     def wrapper(*args, **kwargs):
         acq = args[0].get("acquisition")
@@ -179,7 +190,26 @@ def wrapper(*args, **kwargs):
 
 
 def write_function_to_grid(function, V, grid_spacing):
-    """Interpolate a Firedrake function to a structured grid"""
+    """Interpolate a Firedrake function to a structured grid
+    
+    Parameters
+    ----------
+    function : firedrake.Function
+        Function to interpolate
+    V : firedrake.FunctionSpace
+        Function space of function
+    grid_spacing : float
+        Spacing of grid points in metres
+
+    Returns
+    -------
+    x : numpy.ndarray
+        x coordinates of grid points
+    y : numpy.ndarray
+        y coordinates of grid points
+    z : numpy.ndarray
+        z coordinates of grid points
+    """
     # get DoF coordinates
     m = V.ufl_domain()
     W = fire.VectorFunctionSpace(m, V.ufl_element())
@@ -206,7 +236,20 @@ def write_function_to_grid(function, V, grid_spacing):
 
 
 def create_segy(velocity, filename):
-    """Write the velocity data into a segy file named filename"""
+    """Write the velocity data into a segy file named filename
+    
+    Parameters
+    ----------
+    velocity : firedrake.Function
+        Velocity in a firedrake function
+    filename : str
+        Name of the segy file to write
+
+    Returns
+    -------
+    None
+    
+    """
     spec = segyio.spec()
 
     velocity = np.flipud(velocity.T)
@@ -223,6 +266,7 @@ def create_segy(velocity, filename):
         for tr, il in enumerate(spec.ilines):
             f.trace[tr] = velocity[:, tr]
 
+    return None
 
 @ensemble_save
 def save_shots(model, comm, array, file_name=None):

spyro/pml/damping.py

@@ -21,7 +21,51 @@ def functions(
     y2=None,
     b_pml=None,
 ):
-    """Damping functions for the perfect matched layer for 2D and 3D"""
+    """Damping functions for the perfect matched layer for 2D and 3D
+    
+    Parameters
+    ----------
+    model : dict
+        Dictionary with the model parameters
+    V : obj
+        Firedrake function space
+    dimension : int
+        Dimension of the problem
+    x : obj
+        Firedrake spatial coordinate
+    x1 : float
+        x coordinate of the left boundary of the PML
+    x2 : float
+        x coordinate of the right boundary of the PML
+    a_pml : float
+        Width of the PML in the x direction
+    z : obj
+        Firedrake spatial coordinate
+    z1 : float
+        z coordinate of the bottom boundary of the PML
+    z2 : float
+        z coordinate of the top boundary of the PML
+    c_pml : float
+        Width of the PML in the z direction
+    y : obj, optional
+        Firedrake spatial coordinate, by default None
+    y1 : float, optional
+        y coordinate of the back boundary of the PML, by default None
+    y2 : float, optional
+        y coordinate of the front boundary of the PML, by default None
+    b_pml : float, optional
+        Width of the PML in the y direction, by default None
+
+    Returns
+    -------
+    sigma_x : obj
+        Firedrake function with the damping function in the x direction
+    sigma_z : obj
+        Firedrake function with the damping function in the z direction
+    sigma_y : obj
+        Firedrake function with the damping function in the y direction
+    
+    """
 
     damping_type = model["BCs"]["damping_type"]
     if damping_type == "polynomial":

spyro/receivers/Receivers.py

@@ -115,6 +115,17 @@ def build_maps(self):
         Is always automatticaly called when initializing the class,
         therefore should only be called again if a mesh related attribute
         changes.
+
+        Returns
+        -------
+        cellIDs: list
+            List of cell IDs for each receiver
+        cellVertices: list
+            List of vertices for each receiver
+        cellNodeMaps: list
+            List of node maps for each receiver
+        cell_tabulations: list
+            List of tabulations for each receiver
         """
 
         for rid in range(self.num_receivers):
@@ -143,10 +154,12 @@ def build_maps(self):
     def interpolate(self, field):
         """Interpolate the solution to the receiver coordinates for
         one simulation timestep.
+
         Parameters
         ----------
         field: array-like
             An array of the solution at a given timestep at all nodes
+        
         Returns
         -------
         solution_at_receivers: list
@@ -171,6 +184,7 @@ def apply_receivers_as_source(self, rhs_forcing, residual, IT):
             and timesteps
         IT: int
             Desired time step number to get residual value from
+        
         Returns
         -------
         rhs_forcing: object
@@ -288,6 +302,7 @@ def __new_at(self, udat, receiver_id):
         receiver_id: a list of integers
             A list of receiver ids, ranging from 0 to total receivers
             minus one.
+        
         Returns
         -------
         at: Function value at given receiver
@@ -328,6 +343,7 @@ def __func_receiver_locator_3D(self):
         the receiver.
         The matrix has the deegres of freedom of the nodes inside
         same element as the receiver.
+        
         """
         num_recv = self.num_receivers
 
@@ -360,6 +376,7 @@ def __func_node_locations_3D(self):
         """Function that returns a list which includes a numpy matrix
         where line n has the x and y values of the nth degree of freedom,
         and a numpy matrix of the vertex coordinates.
+        
         """
         x, y, z = SpatialCoordinate(self.mesh)
         ux = Function(self.space).interpolate(x)
@@ -496,7 +513,21 @@ def set_point_cloud(self, comm):
 
 
 def choosing_element(V, degree):
-    """Chooses UFL element based on desired function space"""
+    """Chooses UFL element based on desired function space
+    and degree of interpolation.
+    
+    Parameters
+    ----------
+    V : firedrake.FunctionSpace
+        Function space to be used.
+    degree : int
+        Degree of interpolation.
+
+    Returns
+    -------
+    element : UFL element
+        UFL element to be used in the interpolation.    
+    """
     cell_geometry = V.mesh().ufl_cell()
     if cell_geometry == quadrilateral:
         T = UFCQuadrilateral()

spyro/solvers/helpers.py

@@ -16,7 +16,25 @@
 
 def fill(usol_recv, is_local, nt, nr):
     """Fills usol_recv with -99999 value
-    when it isn't local to any core"""
+    when it isn't local to any core
+    
+    Parameters
+    ----------
+    usol_recv : list
+        List of numpy arrays
+    is_local : list
+        List of booleans indicating if the receiver is local to the core
+    nt : int
+        Number of timesteps
+    nr : int
+        Number of receivers
+    
+    Returns
+    -------
+    usol_recv : list
+        List of numpy arrays
+    
+    """
     usol_recv = np.asarray(usol_recv)
     for ti in range(nt):
         for rn in range(nr):
@@ -26,7 +44,22 @@ def fill(usol_recv, is_local, nt, nr):
 
 
 def create_output_file(name, comm, source_num):
-    """Saves shots in output file"""
+    """Saves shots in output file
+    
+    Parameters
+    ----------
+    name : str
+        Name of the output file
+    comm : object
+        MPI communicator
+    source_num : int
+        Source number
+
+    Returns
+    -------
+    outfile : object
+        Firedrake.File object
+    """
     if io.is_owner(comm, source_num):
         outfile = File(
             os.getcwd()
@@ -41,7 +74,16 @@ def create_output_file(name, comm, source_num):
 
 
 def display(comm, source_num):
-    """Displays current shot and ensemble in terminal"""
+    """Displays current shot and ensemble in terminal
+    
+    Parameters
+    ----------
+    comm : object
+        MPI communicator
+    source_num : int
+        Source number
+    
+    """
     if comm.comm.rank == 0:
         print(
             "Timestepping for shot #",
@@ -54,7 +96,15 @@ def display(comm, source_num):
 
 
 def display_progress(comm, t):
-    """Displays progress time"""
+    """Displays progress time
+    
+    Parameters
+    ----------
+    comm : object
+        MPI communicator
+    t : float
+        Current time
+    """
     if comm.ensemble_comm.rank == 0 and comm.comm.rank == 0:
         print(f"Simulation time is: {t:{10}.{4}} seconds", flush=True)
 
@@ -66,7 +116,22 @@ def parallel_print(string, comm):
 
 
 def receivers_local(mesh, dimension, receiver_locations):
-    """Locates receiveirs in cells"""
+    """Locates receivers in cells
+    
+    Parameters
+    ----------
+    mesh : object
+        Firedrake mesh object
+    dimension : int
+        Dimension of the mesh
+    receiver_locations : list
+        List of receiver locations
+    
+    Returns
+    -------
+    list
+        List of receiver locations in cells
+    """
     if dimension == 2:
         return [mesh.locate_cell([z, x], tolerance=0.01) for z, x in receiver_locations]
     elif dimension == 3: