improve plot.py documentation

Author: jaimesouza

simwave/plots/plot.py

@@ -6,7 +6,32 @@
 
 
 def plot_wavefield(wavefield, file_name="wavefield", colorbar=True,
-                   cmap="gray", extent=None, show=False, clim=[-5, 5],):
+                   cmap="gray", extent=None, show=False, clim=[-5, 5]):
+    """
+    Plot the wavefield.
+
+    Parameters
+    ----------
+    wavefield : ndarray
+        Wavefield data.
+    file_name : str, optional
+        Name of the image to be saved.
+        Default is wavefield.
+    colorbar : bool, optional
+        If True, show a colorbar. Default is True.
+    cmap : str, optional
+        The Colormap instance or registered colormap name
+        used to map scalar data to colors.
+        Default is gray.
+    extent : floats(left, right, bottom, top), optional
+        The bounding box in data coordinates that the image will fill.
+    show : bool, optional
+        If True, show the image on a pop up window.
+        Default is False.
+    clim : list
+        Set the color limits of the current image.
+        Default is (vmin=-5, vmax=5).
+    """
 
     # create the destination dir
     os.makedirs("plots", exist_ok=True)
@@ -44,8 +69,16 @@ def plot_shotrecord(rec, file_name="shotrecord", colorbar=True, show=False):
 
     Parameters
     ----------
-    rec :
+    rec : ndarray
         Receiver data with shape (time, points).
+    file_name : str, optional
+        Name of the image to be saved.
+        Default is wavelet.
+    colorbar : bool, optional
+        If True, show a colorbar. Default is True.
+    show : bool, optional
+        If True, show the image on a pop up window.
+        Default is False.
     """
     scale = np.max(rec) / 10.0
     plot = plt.imshow(rec, vmin=-scale, vmax=scale, cmap=cm.gray)
@@ -73,6 +106,28 @@ def plot_shotrecord(rec, file_name="shotrecord", colorbar=True, show=False):
 
 def plot_velocity_model(model, file_name="velocity_model", colorbar=True,
                         extent=None, cmap="jet", show=False):
+    """
+    Plot the velocity model.
+
+    Parameters
+    ----------
+    model : ndarray
+        Velocity model data.
+    file_name : str, optional
+        Name of the image to be saved.
+        Default is velocity_model.
+    colorbar : bool, optional
+        If True, show a colorbar. Default is True.
+    cmap : str, optional
+        The Colormap instance or registered colormap name
+        used to map scalar data to colors.
+        Default is jet.
+    extent : floats(left, right, bottom, top), optional
+        The bounding box in data coordinates that the image will fill.
+    show : bool, optional
+        If True, show the image on a pop up window.
+        Default is False.
+    """
 
     # create the destination dir
     os.makedirs("plots", exist_ok=True)
@@ -114,6 +169,12 @@ def plot_wavelet(time_values, wavelet_values,
         Discretized values of time in seconds.
     wavelet_values : list
         Pulse of the wavelet.
+    file_name : str, optional
+        Name of the image to be saved.
+        Default is wavelet.
+    show : bool, optional
+        If True, show the image on a pop up window.
+        Default is False.
     """
     plt.plot(time_values, wavelet_values)
     plt.xlabel("Time (s)")