snputils.visualization.scatter_plot

  1import numpy as np
  2import pandas as pd
  3import matplotlib.pyplot as plt
  4from matplotlib import cm
  5from typing import Optional
  6from adjustText import adjust_text
  7
  8
  9def scatter(
 10    dimredobj: np.ndarray,
 11    labels_file: str,
 12    abbreviation_inside_dots: bool = True,
 13    arrows_for_titles: bool = False,
 14    dots: bool = True,
 15    legend: bool = True,
 16    color_palette=None,
 17    show: bool = True,
 18    save_path: Optional[str] = None
 19) -> None:
 20    """
 21    Plot a scatter plot with centroids for each group, with options for labeling and display styles.
 22    
 23    Args:
 24        dimredobj (np.ndarray): 
 25            Reduced dimensionality data; expected to have `(n_haplotypes, 2)` shape.
 26        labels_file (str): 
 27            Path to a TSV file with columns 'indID' and 'label', providing labels for coloring and annotating points.
 28        abbreviation_inside_dots (bool): 
 29            If True, displays abbreviated labels (first 3 characters) inside the centroid markers.
 30        arrows_for_titles (bool): 
 31            If True, adds arrows pointing to centroids with group labels displayed near the centroids.
 32        legend (bool): 
 33            If True, includes a legend indicating each group label.
 34        color_palette (optional): 
 35            Color map or list of colors to use for unique labels. Defaults to 'tab10' if None.
 36        show (bool, optional):
 37            Whether to display the plot. Defaults to False.
 38        save_path (str, optional):
 39            Path to save the plot image. If None, the plot is not saved.
 40        
 41    Returns:
 42        None
 43    """
 44    # Load labels from TSV
 45    labels_df = pd.read_csv(labels_file, sep='\t')
 46
 47    # Filter labels based on the indIDs in dimredobj
 48    sample_ids = dimredobj.samples_
 49    filtered_labels_df = labels_df[labels_df['indID'].isin(sample_ids)]
 50
 51    # Define unique colors for each group label, either from color_palette or defaulting to 'tab10'
 52    unique_labels = filtered_labels_df['label'].unique()
 53    colors = color_palette if color_palette else cm.get_cmap('tab10', len(unique_labels))
 54
 55    # Initialize the plot
 56    fig, ax = plt.subplots(figsize=(10, 8))
 57
 58    # Calculate the overall center of the plot (used for positioning arrows)
 59    plot_center = dimredobj.X_new_.mean(axis=0)
 60
 61    # Dictionary to hold centroid positions for each label
 62    centroids = {}
 63
 64    # Plot data points and centroids by label
 65    for i, label in enumerate(unique_labels):
 66        # Get sample IDs corresponding to the current label
 67        sample_ids_for_label = filtered_labels_df[filtered_labels_df['label'] == label]['indID']
 68        
 69        # Filter points based on sample IDs
 70        points = dimredobj.X_new_[np.isin(dimredobj.samples_, sample_ids_for_label)]
 71
 72        if dots:
 73            # Plot individual points for the current group
 74            ax.scatter(points[:, 0], points[:, 1], s=30, color=colors(i), alpha=0.6, label=label)
 75        else:
 76            # TODO: solve bug
 77            for point in points:
 78                print(point[0], point[1])
 79                ax.text(point[0], point[1], label[:2].upper(), ha='center', va='center', color=colors(i), fontsize=8, weight='bold')
 80
 81        # Calculate and mark the centroid for the current group
 82        centroid = points.mean(axis=0)
 83        centroids[label] = centroid  # Store centroid for later use
 84        
 85        # Plot the centroid as a larger dot
 86        ax.scatter(*centroid, color=colors(i), s=300)
 87
 88        # Optionally add an abbreviation inside the centroid dot
 89        if abbreviation_inside_dots:
 90            ax.text(centroid[0], centroid[1], label[:2].upper(), ha='center', va='center', color='white', fontsize=8, weight='bold')
 91
 92    # Adding arrows and labels with `adjust_text` for no overlap
 93    texts = []
 94    for label, centroid in centroids.items():
 95        # Determine the direction of the arrow based on centroid position
 96        offset_x = 0.07 if centroid[0] >= plot_center[0] else -0.07
 97        offset_y = 0.07 if centroid[1] >= plot_center[1] else -0.07
 98        
 99        if arrows_for_titles:
100            ax.annotate('', xy=centroid, 
101                        xytext=(centroid[0] + offset_x, centroid[1] + offset_y),
102                        arrowprops=dict(facecolor=colors(unique_labels.tolist().index(label)), 
103                                        shrink=0.05, width=1.5, headwidth=8))
104        
105            # Label text for centroid
106            texts.append(ax.text(centroid[0] + offset_x, centroid[1] + offset_y, label, 
107                                color=colors(unique_labels.tolist().index(label)), 
108                                fontsize=12, weight='bold'))
109
110        # Adjust text to prevent overlap
111        adjust_text(texts, arrowprops=dict(arrowstyle="->", color='gray', lw=0.6))
112
113    # Configure additional plot settings
114    ax.set_xlabel("Component 1")
115    ax.set_ylabel("Component 2")
116    if legend:
117        ax.legend(loc='upper right')
118    
119    # Save plot if save_path is provided
120    if save_path:
121        plt.savefig(save_path)
122
123    # Display plot if show is True
124    if show:
125        plt.show()
126    else:
127        plt.close()