wawi 0.0.16__py3-none-any.whl → 0.0.18__py3-none-any.whl

wawi/plot.py

@@ -8,7 +8,42 @@ from scipy.interpolate import RectBivariateSpline, interp1d
 
 
 def plot_ads(ad_dict, v, terms='stiffness', num=None, test_v=dict(), test_ad=dict(), zasso_type=False, ranges=None):
-    # v: v or K
+    """
+    Plot aerodynamic derivative (AD) curves for multiple terms and test data.
+
+    Parameters
+    ----------
+    ad_dict : dict
+        Dictionary mapping term names (e.g., 'P4', 'H6', etc.) to functions that compute aerodynamic derivative values for a given `v`.
+    v : array-like
+        Array of reduced velocity (or reduced frequency) at which to evaluate the aerodynamic derivative functions.
+    terms : str or list of list of str, optional
+        Specifies which terms to plot. If 'stiffness' or 'damping', uses predefined groupings. Otherwise, expects a nested list of term names.
+    num : int or None, optional
+        Figure number for matplotlib. If None, a new figure is created.
+    test_v : dict, optional
+        Dictionary mapping term names to arrays of test reduced velocity values for overlaying test data.
+    test_ad : dict, optional
+        Dictionary mapping term names to arrays of test aerodynamic derivative data corresponding to `test_v`.
+    zasso_type : bool, optional
+        If True, applies special formatting and scaling for Zasso-type plots.
+    ranges : dict or None, optional
+        Dictionary mapping term names to (min, max) tuples, specifying valid ranges for plotting. Values outside the range are clipped.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        The matplotlib Figure object containing the plots.
+    ax : numpy.ndarray of matplotlib.axes.Axes
+        Array of Axes objects for each subplot.
+
+    Notes
+    -----
+    - The function arranges subplots in a grid according to the structure of `terms`.
+    - Each subplot shows the fitted aerodynamic derivative curve and, if provided, test data points.
+    - Axis labels and scaling are adjusted depending on `zasso_type` and the term type.
+    """
+
     if terms == 'stiffness':
         terms = [['P4', 'P6', 'P3'], ['H6', 'H4', 'H3'], ['A6', 'A4', 'A3']]
     elif terms == 'damping':
@@ -68,6 +103,22 @@ def plot_ads(ad_dict, v, terms='stiffness', num=None, test_v=dict(), test_ad=dic
     return fig, ax
 
 def save_plot(pl, path, w=None, h=None):
+    '''
+    Saves pyvista plotter screenshot to file.
+
+    Parameters
+    ----------
+    pl : pyvista.Plotter
+        PyVista plotter object.
+    path : str
+        Path to save the screenshot.
+    w : int, optional
+        Width of the screenshot. If None, uses the width of the plotter window.
+    h : int, optional
+        Height of the screenshot. If None, uses the height of the plotter window.
+    
+
+    '''
     ws = pl.window_size
     if w is not None and h is None:
         w = int(np.round(w))
@@ -83,22 +134,74 @@ def save_plot(pl, path, w=None, h=None):
 
     pl.screenshot(path, window_size=[w,h], return_img=False)
 
-def plot_dir_and_crests(theta0, Tp, arrow_length=100, origin=np.array([0,0]), 
-                        ax=None, n_repeats=2, crest_length=1000,
-                        alpha_crests=0.2, arrow_options={}):
-    arr_opts = {'head_width': 4, 'width':2, 'edgecolor':'none'}
-    arr_opts.update(**arrow_options)
+def plot_dir_and_crests(theta0, Tp, U=0.0, thetaU=0.0, arrow_length=100, origin=np.array([0,0]), add_text=True, 
+                        ax=None, n_repeats=2, crest_length=1000, arrow_options={}, crest_options={}):
     
+    """
+    Plot wave direction arrow and wave crests on a matplotlib axis.
+
+    Parameters
+    ----------
+    theta0 : float
+        Wave direction in degrees (0 degrees is along the x-axis).
+    Tp : float
+        Peak wave period in seconds.
+    U : float, optional
+        Uniform current velocity [m/s]. Default is 0.0.
+    thetaU : float, optional
+        Angle of current in degrees. Default is 0.0.
+    arrow_length : float, optional
+        Length of the direction arrow (default is 100).
+    origin : np.ndarray, optional
+        2D coordinates of the arrow origin (default is np.array([0, 0])).
+    ax : matplotlib.axes.Axes, optional
+        Axis to plot on. If None, uses current axis (default is None).
+    n_repeats : int, optional
+        Number of wave crests to plot (default is 2).
+    crest_length : float, optional
+        Length of each wave crest line (default is 1000).
+    arrow_options : dict, optional
+        Additional keyword arguments for the arrow (default is {}).
+    crest_options : dict, optional
+        Additional keyword arguments for the crest lines (default is {})
+    add_text : bool, optional
+        Whether or not to add text annotation, default is True.
+    
+
+    Returns
+    -------
+    ax : matplotlib.axes.Axes
+        The axis with the plotted wave direction and crests.
+
+    Notes
+    -----
+    - Requires `matplotlib.pyplot` as `plt` and `numpy` as `np` to be imported.
+    - The function also requires a `get_kappa` function to compute the wavenumber.
+
+    Docstring generated by GitHub Copilot.
+    """
+    arrow_options = {'head_width': 4, 'width':2, 'edgecolor':'none'} | arrow_options
+    crest_options = {'alpha': 0.2, 'color': 'black'} | crest_options
+
     if ax is None:
         ax = plt.gca()
         
     # Plot wave angle and crests
     v = np.array([np.cos(theta0*np.pi/180), np.sin(theta0*np.pi/180)])
     v_norm = np.array([-np.sin(theta0*np.pi/180), np.cos(theta0*np.pi/180)])
-    wave_length = 2*np.pi/get_kappa(2*np.pi/Tp, U=0.0)
 
-    plt.arrow(origin[0],origin[1], arrow_length*v[0], arrow_length*v[1], **arr_opts)
-    plt.text(origin[0], origin[1], f'$\\theta_0$ = {theta0}$^o$\n $T_p$={Tp} s\n $\\lambda=${wave_length:.0f} m')
+    theta_rel_U = (thetaU - theta0)*180/np.pi
+    wave_length = 2*np.pi/get_kappa(2*np.pi/Tp, U=U, theta_rel_U=theta_rel_U)
+    
+    plt.arrow(origin[0],origin[1], arrow_length*v[0], arrow_length*v[1], **arrow_options)
+
+    if add_text:
+        plot_string = f'$\\theta_0$ = {theta0}$^o$\n $T_p$={Tp} s\n $\\lambda=${wave_length:.0f} m'
+
+        if U != 0.0:
+            plot_string = plot_string + f'\nU = {U:.1f}m/s @ {thetaU:.1f} deg'
+    
+        plt.text(origin[0], origin[1], plot_string)
 
     dv = v*wave_length
     for n in range(n_repeats):
@@ -106,12 +209,48 @@ def plot_dir_and_crests(theta0, Tp, arrow_length=100, origin=np.array([0,0]),
         p2 = origin+v_norm*crest_length/2
         pts = np.vstack([p1,p2])
         
-        ax.plot(pts[:,0], pts[:,1], alpha=alpha_crests, color='black', zorder=0)
+        ax.plot(pts[:,0], pts[:,1], zorder=0, **crest_options)
         origin = origin + dv
 
+    ax.axis('equal')
     return ax
 
 def rotate_image_about_pivot(Z, x, y, angle, x0=0, y0=0):
+    """
+    Rotate an image array about a specified pivot point.
+    This function rotates a 2D image array `Z` by a given angle (in degrees) about a pivot point (`x0`, `y0`), 
+    where the image axes are defined by the coordinate arrays `x` and `y`. The rotation is performed in the 
+    coordinate space defined by `x` and `y`, not necessarily pixel indices.
+
+    Parameters
+    ----------
+    Z : ndarray
+        2D array representing the image to be rotated.
+    x : array_like
+        1D array of x-coordinates corresponding to the columns of `Z`.
+    y : array_like
+        1D array of y-coordinates corresponding to the rows of `Z`.
+    angle : float
+        Rotation angle in degrees. Positive values correspond to counter-clockwise rotation.
+    x0 : float, optional
+        X-coordinate of the pivot point about which to rotate. Default is 0.
+    y0 : float, optional
+        Y-coordinate of the pivot point about which to rotate. Default is 0.
+
+    Returns
+    -------
+    rotated_Z : ndarray
+        The rotated image array, with the same shape as `Z`.
+
+    Notes
+    -----
+    - The function uses interpolation to map between coordinate space and pixel indices.
+    - The rotation is performed about the specified pivot point (`x0`, `y0`) in the coordinate system defined by `x` and `y`.
+    - Requires `numpy`, `scipy.ndimage.shift`, `scipy.ndimage.rotate`, and `scipy.interpolate.interp1d`.
+
+    Docstring generated by GitHub Copilot.
+    """
+
     xc = np.mean(x)
     yc = np.mean(y)
     
@@ -130,6 +269,45 @@ def rotate_image_about_pivot(Z, x, y, angle, x0=0, y0=0):
     return shift(rotate(shift(Z, ds[::-1]), angle), -ds_rot[::-1])
 
 def combine_eta(eta_fine, eta_course, x_fine, y_fine, x_course, y_course, x=None, y=None):
+    """
+    Combine two 2D fields (fine and coarse) into a single field, using the fine field where available.
+
+    Parameters
+    ----------
+    eta_fine : ndarray
+        2D array of the fine-resolution field values.
+    eta_course : ndarray
+        2D array of the coarse-resolution field values.
+    x_fine : ndarray
+        1D array of x-coordinates for the fine field.
+    y_fine : ndarray
+        1D array of y-coordinates for the fine field.
+    x_course : ndarray
+        1D array of x-coordinates for the coarse field.
+    y_course : ndarray
+        1D array of y-coordinates for the coarse field.
+    x : ndarray, optional
+        1D array of x-coordinates for the output grid. If None, generated from coarse grid.
+    y : ndarray, optional
+        1D array of y-coordinates for the output grid. If None, generated from coarse grid.
+
+    Returns
+    -------
+    eta_combined : ndarray
+        2D array of the combined field, using fine field values where available and coarse elsewhere.
+    x : ndarray
+        1D array of x-coordinates for the combined field.
+    y : ndarray
+        1D array of y-coordinates for the combined field.
+
+    Notes
+    -----
+    The function interpolates both input fields onto a common grid. The fine field overwrites the coarse field
+    in the region where it is defined.
+
+    Docstring generated by GitHub Copilot.
+    """
+
     dx_fine = x_fine[1]-x_fine[0]
     dy_fine = y_fine[1]-y_fine[0]
     
@@ -156,7 +334,47 @@ def combine_eta(eta_fine, eta_course, x_fine, y_fine, x_course, y_course, x=None
 def animate_surface(eta, x, y, t, filename=None, fps=None, 
                     speed_ratio=1.0, figsize=None, writer='ffmpeg', 
                     ax=None, surface=None):
-    
+    """
+    Animates a time-evolving eta (sea surface).
+
+    Parameters
+    ----------
+    eta : ndarray
+        3D array of surface values with shape (nx, ny, nt), where nt is the number of time steps.
+    x : ndarray
+        1D or 2D array representing the x-coordinates of the surface grid.
+    y : ndarray
+        1D or 2D array representing the y-coordinates of the surface grid.
+    t : ndarray
+        1D array of time points corresponding to the third dimension of `eta`.
+    filename : str or None, optional
+        If provided, the animation is saved to this file. If None, the animation is displayed interactively.
+    fps : float or None, optional
+        Frames per second for the animation. If None, it is computed from the time step and `speed_ratio`.
+    speed_ratio : float, optional
+        Ratio to speed up or slow down the animation relative to real time. Default is 1.0.
+    figsize : tuple or None, optional
+        Size of the figure in inches. Passed to `plt.subplots` if a new figure is created.
+    writer : str, optional
+        Writer to use for saving the animation (e.g., 'ffmpeg'). Default is 'ffmpeg'.
+    ax : matplotlib.axes.Axes or None, optional
+        Axes object to plot on. If None, a new figure and axes are created.
+    surface : matplotlib.image.AxesImage or None, optional
+        Existing surface image to update. If None, a new surface is created.
+
+    Returns
+    -------
+    None
+        The function either displays the animation or saves it to a file.
+
+    Notes
+    -----
+    - Requires matplotlib and numpy.
+    - The function uses `matplotlib.animation.FuncAnimation` for animation.
+    - If `filename` is provided, the animation is saved and not displayed interactively.
+
+    Docstring generated by GitHub Copilot.
+    """
         
     if surface is None:
         if ax is None:
@@ -195,6 +413,31 @@ def animate_surface(eta, x, y, t, filename=None, fps=None,
         plt.show()
 
 def xy_to_latlon(latlon0, coors, rot=0):
+    """
+    Convert local Cartesian (x, y) coordinates to latitude and longitude using geodesic calculations.
+
+    Parameters
+    ----------
+    latlon0 : array-like, shape (2,)
+        The reference latitude and longitude (in degrees) as a 2-element array or list [lat, lon].
+    coors : array-like, shape (N, 2)
+        Array of local Cartesian coordinates (x, y) to be converted, where each row is a point.
+    rot : float, optional
+        Rotation angle in degrees to be subtracted from the computed azimuth, default is 0.
+
+    Returns
+    -------
+    latlon : ndarray, shape (N, 2)
+        Array of latitude and longitude pairs (in degrees) corresponding to the input coordinates.
+
+    Notes
+    -----
+    Uses Cartopy's geodesic calculations to convert local (x, y) displacements from a reference point
+    (latlon0) to geographic coordinates, accounting for Earth's curvature.
+
+    Docstring generated by GitHub Copilot.
+    """
+
     import cartopy.crs as ccrs, cartopy.geodesic as cgds
     dist = np.linalg.norm(coors, axis=1)
     azi = np.arctan2(coors[:,0], coors[:,1])*180/np.pi - rot # atan2(x,y) to get azimuth (relative to N-vector)
@@ -208,6 +451,56 @@ def plot_surface_in_map(eta, x, y, eta_geo0, extent,
                         wms_url='https://openwms.statkart.no/skwms1/wms.terrengmodell?request=GetCapabilities&service=WMS', 
                         wms_layers=['relieff'], ax=None,
                         cm='Blues_r', colorbar=True, figsize=None, eta_rot=0, labels=False):
+    """
+    Plot a 2D surface (e.g., elevation or other field) on a map with optional scatter points and WMS background.
+
+    Parameters
+    ----------
+    eta : ndarray or None
+        2D array representing the surface to plot (e.g., elevation values). If None, only scatter and WMS are shown.
+    x : ndarray
+        1D or 2D array of x-coordinates corresponding to `eta`.
+    y : ndarray
+        1D or 2D array of y-coordinates corresponding to `eta`.
+    eta_geo0 : array-like
+        Reference geographic coordinates (e.g., origin for coordinate transformation).
+    extent : array-like
+        Map extent in the format [xmin, xmax, ymin, ymax] (in longitude and latitude).
+    eta_scatter : ndarray or None, optional
+        Array of points to scatter on the map, shape (N, 2). Default is None.
+    wms_url : str, optional
+        URL to the WMS service for background map. Default is Statkart's terrain model.
+    wms_layers : list of str, optional
+        List of WMS layer names to display. Default is ['relieff'].
+    ax : matplotlib.axes.Axes or None, optional
+        Existing axes to plot on. If None, a new axes with Mercator projection is created.
+    cm : str or Colormap, optional
+        Colormap for the surface plot. Default is 'Blues_r'.
+    colorbar : bool, optional
+        If True, display a colorbar for the surface. Default is True.
+    figsize : tuple or None, optional
+        Figure size. Not used if `ax` is provided. Default is None.
+    eta_rot : float, optional
+        Rotation angle (degrees) to apply to `eta` and scatter points. Default is 0.
+    labels : bool, optional
+        If True, draw gridlines with labels. Default is False.
+
+    Returns
+    -------
+    ax : matplotlib.axes.Axes
+        The axes with the plotted map.
+    scatter : matplotlib.collections.PathCollection or None
+        The scatter plot object, or None if `eta_scatter` is None.
+    surface : matplotlib.image.AxesImage or None
+        The surface plot object, or None if `eta` is None.
+
+    Notes
+    -----
+    - Requires cartopy and matplotlib.
+    - Assumes existence of `xy_to_latlon` and `rotate` helper functions.
+
+    Docstring generated by GitHub Copilot.
+    """
     
     import cartopy.crs as ccrs, cartopy.geodesic as cgds
     
@@ -270,6 +563,42 @@ def plot_surface_in_map(eta, x, y, eta_geo0, extent,
 def plot_surface(eta, x, y, ax=None, 
                  cm='Blues_r', colorbar=True, 
                  labels=True, figsize=None, interpolation='none'): 
+    """
+    Plot a 2D surface using imshow.
+    
+    Parameters
+    ----------
+    eta : ndarray
+        2D array representing the surface to plot (e.g., elevation values).
+    x : ndarray
+        1D or 2D array of x-coordinates corresponding to `eta`.
+    y : ndarray
+        1D or 2D array of y-coordinates corresponding to `eta`.
+    ax : matplotlib.axes.Axes or None, optional
+        Existing axes to plot on. If None, a new axes is created.
+    cm : str or Colormap, optional
+        Colormap for the surface plot. Default is 'Blues_r'.
+    colorbar : bool, optional
+        If True, display a colorbar for the surface. Default is True.
+    labels : bool, optional
+        If True, draw gridlines with labels. Default is True.
+    figsize : tuple or None, optional
+        Figure size. Not used if `ax` is provided. Default is None.
+    interpolation : str, optional
+        Interpolation method for the surface plot. Default is 'none'.
+
+    Returns
+    ------- 
+    surface : matplotlib.image.AxesImage
+        The surface plot object.
+    ax : matplotlib.axes.Axes
+        The axes with the plotted surface.
+
+    Notes
+    -----
+    Docstring generated by GitHub Copilot.
+
+"""
 
     if ax is None:
         fig, ax = plt.subplots(figsize=figsize)
@@ -295,7 +624,6 @@ def plot_surface(eta, x, y, ax=None,
 
 
 def set_axes_equal(ax: plt.Axes):
-    import matplotlib.pyplot as plt
     """Set 3D plot axes to equal scale.
 
     Make axes of 3D plot have equal scale so that spheres appear as
@@ -375,103 +703,38 @@ def plot_transformation_mats(x,y,z,T,figno=None, ax=None, scaling='auto'):
     equal_3d(ax)
     return ax,h
 
-
-def plot_elements(element_matrix, node_matrix, chosen_nodes_ix=[], disp=None, node_labels=False, element_labels=False, plot_nodes=True, plot_elements=True, ax=None, fig=None, element_settings={}, node_settings={}, node_label_settings={}, chosen_node_settings={}, disp_settings={}, element_label_settings={}, three_d=True):
-    e_dict = {'color': 'LimeGreen', 'alpha': 1}
-    e_dict.update(**element_settings)
-
-    n_dict = {'color':'Black', 'linestyle':'', 'marker':'.', 'markersize':4, 'alpha':0.8}
-    n_dict.update(**node_settings)
-       
-    n_chosen_dict = {'color':'GreenYellow', 'linestyle':'', 'marker':'o', 'markersize':8, 'alpha':1, 'markeredgecolor':'dimgray'}
-    n_chosen_dict.update(**chosen_node_settings)
-    
-    disp_dict = {'color':'IndianRed', 'alpha':1}
-    disp_dict.update(**disp_settings)
-
-    l_nodes_dict = {'color':'Black', 'fontsize': 8, 'fontweight':'normal'}
-    l_nodes_dict.update(**node_label_settings)
-    
-    l_elements_dict = {'color':'LimeGreen', 'fontsize': 8, 'fontweight':'bold', 'style':'italic'}
-    l_elements_dict.update(**element_label_settings)
-    
-    if ax is None and fig is None:
-        fig = plt.figure()
-        
-    if ax == None and three_d:
-        ax = fig.gca(projection='3d')
-    elif ax == None:
-        ax = fig.gca()
-    elif three_d:
-        1
-        # ax.set(projection='3d')  #mangler funksjonalitet her...
-    
-    element_handles = [None]*len(element_matrix[:,0])
-
-    if plot_elements:
-        for element_ix, __ in enumerate(element_matrix[:,0]):
-            node1 = element_matrix[element_ix, 1]
-            node2 = element_matrix[element_ix, 2]
-            nodeix1 = np.where(node_matrix[:,0]==node1)[0]
-            nodeix2 = np.where(node_matrix[:,0]==node2)[0]
-            x01 = node_matrix[nodeix1,1:4]
-            x02 = node_matrix[nodeix2,1:4]
-            x0 = np.vstack([x01,x02])
-
-            if three_d:
-                element_handles[element_ix] = ax.plot(xs=x0[:,0], ys=x0[:,1], zs=x0[:,2], **e_dict)
-            else:
-                element_handles[element_ix] = ax.plot(x0[:,0], x0[:,1], **e_dict)
-  
-            if element_labels:
-                xmean = np.mean(x0, axis=0)
-                if three_d:
-                    ax.text(xmean[0],xmean[1],xmean[2],'%i' % element_matrix[element_ix,0], **l_elements_dict)
-                else:
-                    ax.text(xmean[0],xmean[1],s='%i' % element_matrix[element_ix,0], **l_elements_dict)
-
-            if disp is not None:
-                disp_node1 = disp[nodeix1[0]*6:(nodeix1[0]*6+6)]
-                disp_node2 = disp[nodeix2[0]*6:(nodeix2[0]*6+6)]
-                x1 = x01+disp_node1[0:3]
-                x2 = x02+disp_node2[0:3]
-                x = np.vstack([x1,x2])
-                
-                if three_d:
-                    ax.plot(xs=x[:,0], ys=x[:,1], zs=x[:,2], **disp_dict)
-                else:
-                    ax.plot(x[:,0], x[:,1], **disp_dict)
-
-    if plot_nodes:
-        if three_d:
-            ax.plot(xs=node_matrix[:, 1], ys=node_matrix[:, 2], zs=node_matrix[:, 3], **n_dict)
-        else:
-           ax.plot(node_matrix[:, 1], node_matrix[:, 2], **n_dict)
-        
-        if chosen_nodes_ix != []:
-            if three_d:
-                ax.plot(xs=node_matrix[chosen_nodes_ix, 1], ys=node_matrix[chosen_nodes_ix, 2], zs=node_matrix[chosen_nodes_ix, 3], **n_chosen_dict)
-            else:
-               ax.plot(node_matrix[chosen_nodes_ix, 1], node_matrix[chosen_nodes_ix, 2], **n_chosen_dict)
-        
-    if node_labels:
-        if three_d:
-            for node_ix in range(0, np.shape(node_matrix)[0]):
-                ax.text(node_matrix[node_ix, 1], node_matrix[node_ix, 2], node_matrix[node_ix, 3], '%i' % node_matrix[node_ix, 0], **l_nodes_dict)
-        else:
-            for node_ix in range(0, np.shape(node_matrix)[0]):
-                ax.text(node_matrix[node_ix, 1], node_matrix[node_ix, 2], '%i' % node_matrix[node_ix, 0], **l_nodes_dict)
+def plot_2d(S2d, x1, x2, ax=None, levels=80, discrete=False, **kwargs):
+    """
+    Plot a 2D array as either a filled contour plot or a pseudocolor mesh.
+
+    Parameters
+    ----------
+    S2d : array_like
+        2D array of values to plot.
+    x1 : array_like
+        1D array representing the x-coordinates.
+    x2 : array_like
+        1D array representing the y-coordinates.
+    ax : matplotlib.axes.Axes, optional
+        The axes on which to plot. If None, uses the current axes.
+    levels : int, optional
+        Number of contour levels to use for filled contour plot. Default is 80.
+    discrete : bool, optional
+        If True, use `pcolormesh` for a discrete colormap. If False, use `contourf` for a filled contour plot.
+    **kwargs
+        Additional keyword arguments passed to the plotting function (`contourf` or `pcolormesh`).
+
+    Returns
+    -------
+    contour : QuadMesh or QuadContourSet
+        The resulting plot object from `pcolormesh` or `contourf`.
+
+    Notes
+    -----
+    Docstring generated by GitHub Copilot.
+    """
         
-    if three_d:
-        equal_3d(ax)
-    else:
-        ax.set_aspect('equal', adjustable='box')
-    
-    ax.grid('off')
-    return ax, element_handles
-
 
-def plot_2d(S2d, x1, x2, ax=None, levels=80, discrete=False, **kwargs):
     if ax is None:
         ax = plt.gca()
         
@@ -484,6 +747,39 @@ def plot_2d(S2d, x1, x2, ax=None, levels=80, discrete=False, **kwargs):
 
 
 def plot_S2d(S, omega, theta, D=None, omega_range=None, theta_range=None):
+    """
+    Plot a 2D spectral density (S2d) as a contour plot, with optional marginal plots (S and D).
+
+    Parameters
+    ----------
+    S : array_like
+        1D array of spectral density values as a function of omega.
+    omega : array_like
+        1D array of frequency values (rad/s).
+    theta : array_like
+        1D array of direction values (rad).
+    D : array_like, optional
+        1D array of directional spreading function values as a function of theta.
+        If provided, marginal plots of S(omega) and D(theta) are shown.
+    omega_range : list or tuple, optional
+        Range [min, max] for the omega axis. If None, uses [0, max(omega)].
+    theta_range : list or tuple, optional
+        Range [min, max] for the theta axis. If None, uses [min(theta), max(theta)].
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        The matplotlib Figure object containing the plot.
+
+    Notes
+    -----
+    - If `D` is provided, the function plots the 2D spectral density as a contour plot,
+      with marginal line plots for S(omega) and D(theta).
+    - If `D` is not provided, only the contour plot is shown.
+    - The function handles NaN values in the spectral density by setting them to zero.
+
+    Docstring generated by GitHub Copilot.
+    """
 
     if theta_range is None:
         theta_range = [np.min(theta), np.max(theta)]

wawi/prob.py

@@ -1,6 +1,34 @@
 import numpy as np
 
 def gumbel_log(umax):
+    """
+    Compute the Gumbel reduced variate for a given array of maxima.
+
+    Parameters
+    ----------
+    umax : array_like
+        Array of maxima values.
+
+    Returns
+    -------
+    umax_ordered : ndarray
+        The input maxima sorted in descending order.
+    loglogF : ndarray
+        The Gumbel reduced variate corresponding to each sorted maxima.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> umax = np.array([2.3, 3.1, 1.8, 2.9])
+    >>> umax_ordered, loglogF = gumbel_log(umax)
+
+    Notes
+    -------
+    This function sorts the input array of maxima in descending order and computes the
+    Gumbel reduced variate (log-log of the empirical cumulative distribution function)
+    for each value.
+    """
+    
     umax_ordered = np.sort(umax)[::-1]
     N_stat = len(umax)
     F = 1-np.arange(1, N_stat+1)/(N_stat+1)