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    for k in range(K):
21        rows_k = np.where(np.argmax(Q_cols_sorted, axis=1) == k)[0]
22        # Sort these rows by their proportion in col k
23        rows_k_sorted = rows_k[np.argsort(Q_cols_sorted[rows_k, k])[::-1]]
24        row_groups.append(rows_k_sorted)
25        boundary_list.append(boundary_list[-1] + len(rows_k_sorted))
26    
27    # Combine them into one final row order
28    row_order = np.concatenate(row_groups)
29    Q_mat_sorted = Q_cols_sorted[row_order, :]
30    
31    return Q_mat_sorted, row_order, boundary_list, col_order
32
33def plot_admixture(ax, Q_mat_sorted, boundary_list, col_order=None, colors=None):
34    """
35    Plot a structure-style bar chart of Q_mat_sorted in the given Axes ax.
36    If colors is not None, it should be a list or array of length K.
37    If col_order is not None, colors are reordered according to col_order.
38    """
39    n_samples, K = Q_mat_sorted.shape
40    
41    # If we have a specific color list and a col_order, reorder the colors to match the columns
42    if (colors is not None) and (col_order is not None):
43        # re-map the user-provided colors to match the new column order
44        colors = [colors[idx] for idx in col_order]
45    
46    # cumulative sum across columns for stacked fill
47    Q_cum = np.cumsum(Q_mat_sorted, axis=1)
48    x_vals = np.arange(n_samples)
49    
50    step_mode = "post"
51    ax.step(x_vals, Q_cum, linewidth=0.0, where=step_mode)
52    
53    # fill-between for a stacked bar effect
54    for j in range(K):
55        c = colors[j] if (colors is not None) else None
56        if j == 0:
57            ax.fill_between(x_vals, 0, Q_cum[:, j], step=step_mode, color=c)
58        else:
59            ax.fill_between(x_vals, Q_cum[:, j - 1], Q_cum[:, j], step=step_mode, color=c)
60    
61    # Vertical lines for group boundaries
62    for boundary in boundary_list:
63        ax.axvline(boundary, color='black', ls='--', lw=1.0)
64    
65    ax.set_xlim(0, n_samples - 1)
66    ax.set_ylim(0, 1)
67    ax.set_xlabel("Samples")
68    ax.set_ylabel("Ancestry Proportion")