snputils.visualization.admixture

  1import numpy as np
  2import matplotlib.pyplot as plt
  3import matplotlib
  4
  5def reorder_admixture(Q_mat):
  6    """
  7    Reorder Q_mat rows so that rows are grouped by each sample's dominant ancestry,
  8    and columns are sorted by descending average ancestry proportion.
  9    """
 10    n_samples, K = Q_mat.shape
 11    
 12    # Reorder columns by descending average proportion
 13    col_means = Q_mat.mean(axis=0)
 14    col_order = np.argsort(col_means)[::-1]   # largest first
 15    Q_cols_sorted = Q_mat[:, col_order]
 16    
 17    # Group samples by whichever column is argmax
 18    row_groups = []
 19    boundary_list = [0]
 20    argmax_all = np.argmax(Q_cols_sorted, axis=1)
 21    for k in range(K):
 22        rows_k = np.where(argmax_all == k)[0]
 23        # Sort these rows by their proportion in col k
 24        rows_k_sorted = rows_k[np.argsort(Q_cols_sorted[rows_k, k])[::-1]]
 25        row_groups.append(rows_k_sorted)
 26        boundary_list.append(boundary_list[-1] + len(rows_k_sorted))
 27    
 28    # Combine them into one final row order
 29    row_order = np.concatenate(row_groups) if row_groups else np.arange(n_samples)
 30    Q_mat_sorted = Q_cols_sorted[row_order, :]
 31    
 32    return Q_mat_sorted, row_order, boundary_list, col_order
 33
 34def plot_admixture(ax, Q_mat_sorted, boundary_list, col_order=None, colors=None, show_boundaries=True, show_axes_labels=True, show_ticks=True, set_limits=True, minimal=False):
 35    """
 36    Plot a structure-style bar chart of Q_mat_sorted in the given Axes ax.
 37    If colors is not None, it should be a list or array of length K.
 38    If col_order is not None, colors are reordered according to col_order.
 39
 40    Optional controls:
 41    - show_boundaries (bool): draw vertical lines at group boundaries. Default True.
 42    - show_axes_labels (bool): set X/Y axis labels. Default True.
 43    - show_ticks (bool): show axis ticks. Default True.
 44    - set_limits (bool): set xlim and ylim to [0, n_samples-1] and [0,1]. Default True.
 45    - minimal (bool): if True, overrides to disable boundaries, labels, ticks, limits and hides spines.
 46    """
 47    n_samples, K = Q_mat_sorted.shape
 48
 49    # Minimal overrides
 50    if minimal:
 51        show_boundaries = False
 52        show_axes_labels = False
 53        show_ticks = False
 54        set_limits = False
 55
 56    # If we have a specific color list and a col_order, reorder the colors to match the columns
 57    if (colors is not None) and (col_order is not None):
 58        colors = [colors[idx] for idx in col_order]
 59
 60    # cumulative sum across columns for stacked fill
 61    Q_cum = np.cumsum(Q_mat_sorted, axis=1)
 62    # Use step='post' with padded x/y so the last bar renders fully and no thin band appears
 63    x_edges = np.arange(n_samples + 1)
 64    Q_pad = np.vstack([Q_cum, Q_cum[-1]])
 65
 66    # fill-between for a stacked bar effect
 67    for j in range(K):
 68        c = colors[j] if (colors is not None) else None
 69        lower = Q_pad[:, j - 1] if j > 0 else np.zeros(n_samples + 1)
 70        upper = Q_pad[:, j]
 71        ax.fill_between(
 72            x_edges,
 73            lower,
 74            upper,
 75            step="post",
 76            color=c,
 77            linewidth=0,
 78            edgecolor='none',
 79        )
 80
 81    # Vertical lines for group boundaries
 82    if show_boundaries:
 83        for boundary in boundary_list:
 84            ax.axvline(boundary, color='black', ls='--', lw=1.0)
 85
 86    if set_limits:
 87        ax.set_xlim(0, n_samples)
 88        ax.set_ylim(0, 1)
 89
 90    if show_axes_labels:
 91        ax.set_xlabel("Samples")
 92        ax.set_ylabel("Ancestry Proportion")
 93
 94    if not show_ticks:
 95        ax.set_xticks([])
 96        ax.set_yticks([])
 97
 98    if minimal:
 99        # Hide all spines for a clean, minimal appearance
100        for spine in ax.spines.values():
101            spine.set_visible(False)