wmaee.utils.plot

 1from typing import Optional, Tuple
 2import matplotlib.pyplot as plt
 3import numpy as np
 4
 5def scatter_hist(x: np.ndarray,
 6                 y: np.ndarray,
 7                 xlabel: Optional[str] = None,
 8                 ylabel: Optional[str] = None,
 9                 figsize: Tuple[float, float] = (6, 6),
10                 alpha: float = 0.1,
11                 alphax: float = 0.7,
12                 alphay: float = 0.7,
13                 marker: str = '.') -> Tuple[plt.Figure, plt.Axes]:
14    """
15    Create a scatter plot with histograms on the margins.
16
17    Parameters
18    ----------
19    x : np.ndarray
20        Data for the x-axis.
21    y : np.ndarray
22        Data for the y-axis.
23    xlabel : Optional[str], optional
24        Label for the x-axis. Default is None.
25    ylabel : Optional[str], optional
26        Label for the y-axis. Default is None.
27    figsize : Tuple[float, float], optional
28        Size of the figure (width, height). Default is (6, 6).
29    alpha : float, optional
30        Transparency of the scatter plot points. Default is 0.1.
31    alphax : float, optional
32        Transparency of the x-axis histogram bars. Default is 0.7.
33    alphay : float, optional
34        Transparency of the y-axis histogram bars. Default is 0.7.
35    marker : str, optional
36        Marker style for the scatter plot. Default is '.'.
37
38    Returns
39    -------
40    Tuple[plt.Figure, plt.Axes]
41        The created figure and the main scatter plot axes.
42    """
43
44    # Create a figure with constrained layout and specified size
45    fig = plt.figure(constrained_layout=True, figsize=figsize)
46    
47    # Create a main scatter plot axes
48    ax = fig.add_gridspec(top=0.75, right=0.75).subplots()
49
50    # Create inset axes for x-axis histogram
51    ax_histx = ax.inset_axes([0, 1.05, 1, 0.25], sharex=ax)
52    ax_histx.tick_params(axis='x', labelbottom=False)
53
54    # Create inset axes for y-axis histogram
55    ax_histy = ax.inset_axes([1.05, 0, 0.25, 1], sharey=ax)
56    ax_histy.tick_params(axis='y', labelleft=False)
57
58    # Plot scatter plot on the main axes with specified marker
59    ax.scatter(x, y, alpha=alpha, marker=marker)
60
61    # Plot histogram for x-axis on the inset axes
62    ax_histx.hist(x, bins=100, alpha=alphax)
63
64    # Plot histogram for y-axis on the inset axes with orientation set to horizontal
65    ax_histy.hist(y, bins=100, alpha=alphay, orientation='horizontal')
66
67    # Set labels for the x and y axes
68    ax.set_ylabel(ylabel)
69    ax.set_xlabel(xlabel)
70
71    # Return the created figure and axes
72    return fig, ax