From 29fab077b6e65dfcf1060dd763d2eee02c59fc42 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Levente=20Temesv=C3=A1ri-Nagy?=
 <147416790+leventetn@users.noreply.github.com>
Date: Sat, 26 Sep 2026 10:53:48 +0200
Subject: [PATCH 01/10] refactor/reconstruct_corrs_df_symmetric_from_long_df
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

proteopy/utils/copf.py → new proteopy/utils/_matrix_wrangling.py with
reconstruct_symmetric_matrix_from_long(df, var_a_col, var_b_col,
value_col, *, diagonal=1.0) — generic naming/docstring, diagonal param,
and a raise ValueError (with np.isclose) replacing the bare assert. Not
exported anywhere (private).
Updated the 3 production call sites in tl/copf.py and the test imports;
removed the stray test that lived in tests/utils/helpers.py.
New tests/utils/test_matrix_wrangling.py — 9 short one-concern tests
(basic reconstruction, symmetry, mirroring, diagonal, int-vs-str
selectors, generic labels, single pair, missing-pair raise, conflict
raise)
---
 proteopy/utils/_matrix_wrangling.py  |  98 ++++++
 proteopy/utils/copf.py               |  65 ----
 tests/tl/test_copro.py               | 464 ++++++++++++++-------------
 tests/utils/helpers.py               | 103 +++---
 tests/utils/test_matrix_wrangling.py | 143 +++++++++
 5 files changed, 524 insertions(+), 349 deletions(-)
 create mode 100644 proteopy/utils/_matrix_wrangling.py
 delete mode 100644 proteopy/utils/copf.py
 create mode 100644 tests/utils/test_matrix_wrangling.py

diff --git a/proteopy/utils/_matrix_wrangling.py b/proteopy/utils/_matrix_wrangling.py
new file mode 100644
index 0000000..d733eb7
--- /dev/null
+++ b/proteopy/utils/_matrix_wrangling.py
@@ -0,0 +1,98 @@
+import numpy as np
+import pandas as pd
+
+
+def reconstruct_symmetric_matrix_from_long(
+    df,
+    var_a_col=0,
+    var_b_col=1,
+    value_col=2,
+    *,
+    diagonal=1.0,
+):
+    """Reconstruct a square symmetric matrix from a long DataFrame.
+
+    Build a full symmetric matrix from a long-format DataFrame that
+    lists pairs of labels and an associated value. Each ``(a, b)`` pair
+    fills both ``M[a, b]`` and ``M[b, a]``. Labels are the sorted union
+    of the two label columns and become both the index and the columns
+    of the result.
+
+    Parameters
+    ----------
+    df : pandas.DataFrame
+        Long-format table with a column for the first label, a column
+        for the second label, and a column with the pair value.
+    var_a_col : str | int
+        Name or positional index of the first label column.
+    var_b_col : str | int
+        Name or positional index of the second label column.
+    value_col : str | int
+        Name or positional index of the value column.
+    diagonal : float
+        Value written on the matrix diagonal (``1.0`` by default, the
+        correlation convention).
+
+    Returns
+    -------
+    pandas.DataFrame
+        Square symmetric matrix with the sorted labels as both index
+        and columns.
+
+    Raises
+    ------
+    ValueError
+        If a label pair is absent from both triangles, or if a pair is
+        present in both triangles with conflicting values.
+    """
+    if isinstance(var_a_col, int):
+        var_a_col = df.columns[var_a_col]
+
+    if isinstance(var_b_col, int):
+        var_b_col = df.columns[var_b_col]
+
+    if isinstance(value_col, int):
+        value_col = df.columns[value_col]
+
+    labels = set(df[var_a_col]).union(set(df[var_b_col]))
+    labels = sorted(labels)
+    n = len(labels)
+
+    label_to_idx = {label: i for i, label in enumerate(labels)}
+
+    # -- Initialise with NaN and a fixed diagonal
+    matrix = np.full((n, n), np.nan)
+    np.fill_diagonal(matrix, diagonal)
+
+    # -- Fill in the known values
+    for _, row in df.iterrows():
+        i = label_to_idx[row[var_a_col]]
+        j = label_to_idx[row[var_b_col]]
+
+        matrix[i, j] = row[value_col]
+
+    # -- Mirror across the diagonal, validating any overlap
+    idx_to_label = {i: label for label, i in label_to_idx.items()}
+    for i in range(n):
+        for j in range(i + 1, n):
+
+            upper_nan = np.isnan(matrix[i, j])
+            lower_nan = np.isnan(matrix[j, i])
+
+            if upper_nan and not lower_nan:
+                matrix[i, j] = matrix[j, i]
+            elif lower_nan and not upper_nan:
+                matrix[j, i] = matrix[i, j]
+            elif upper_nan and lower_nan:
+                raise ValueError(
+                    "No value found for the combination of labels: "
+                    f"{idx_to_label[i]} and {idx_to_label[j]}."
+                )
+            elif not np.isclose(matrix[i, j], matrix[j, i]):
+                raise ValueError(
+                    "Conflicting values for the combination of labels "
+                    f"{idx_to_label[i]} and {idx_to_label[j]}: "
+                    f"{matrix[i, j]} != {matrix[j, i]}."
+                )
+
+    return pd.DataFrame(matrix, index=labels, columns=labels)
diff --git a/proteopy/utils/copf.py b/proteopy/utils/copf.py
deleted file mode 100644
index 7f5a577..0000000
--- a/proteopy/utils/copf.py
+++ /dev/null
@@ -1,65 +0,0 @@
-import numpy as np
-import pandas as pd
-
-
-def reconstruct_corrs_df_symmetric_from_long_df(df, var_a_col=0, var_b_col=1, corr_col=2):
-    '''Reconstruct correlation dataframe in symmetrical matrix format.
-
-    Reconstruct a full correlation matrix from a long DataFrame containing asymmetric correlation data.
-    
-    Args:
-        df (pd.DataFrame): DataFrame with columns for peptide A, peptide B, and their correlation value
-        var_a_col (str | int): Name of column containing first peptide identifier
-        var_b_col (str | int): Name of column containing second peptide identifier
-        corr_col (str | int): Name of column containing correlation values
-        
-    Returns:
-        pd.DataFrame: Fully symmetric correlation matrix as a pd.DataFrame with peptide labels as columns and rows.
-    '''
-    if isinstance(var_a_col, int):
-        var_a_col = df.columns[var_a_col]
-
-    if isinstance(var_b_col, int):
-        var_b_col = df.columns[var_b_col]
-
-    if isinstance(corr_col, int):
-        corr_col = df.columns[corr_col]
-
-    all_peptides = set(df[var_a_col]).union(set(df[var_b_col]))
-    all_peptides = sorted(list(all_peptides))
-    n = len(all_peptides)
-    
-    pep_to_idx = {pep: i for i, pep in enumerate(all_peptides)}
-    
-    # Init
-    corr_matrix = np.full((n, n), np.nan)
-    np.fill_diagonal(corr_matrix, 1.0)
-    
-    # Fill in the known correlation values
-    for _, row in df.iterrows():
-        i = pep_to_idx[row[var_a_col]]
-        j = pep_to_idx[row[var_b_col]]
-
-        corr_matrix[i, j] = row[corr_col]
-    
-    # Fill in the symmetric values where possible
-    for i in range(n):
-        for j in range(i+1, n):
-
-            if np.isnan(corr_matrix[i, j]) and not np.isnan(corr_matrix[j, i]):
-                corr_matrix[i, j] = corr_matrix[j, i]
-            elif np.isnan(corr_matrix[j, i]) and not np.isnan(corr_matrix[i, j]):
-                corr_matrix[j, i] = corr_matrix[i, j]
-            elif np.isnan(corr_matrix[j, i]) and np.isnan(corr_matrix[i, j]):
-                rev = {i: pep for pep, i in pep_to_idx.items()}
-                raise ValueError((
-                    f'Logical bug. For combination of peptides: {rev[i]} and '
-                    f'{rev[j]} there was no value found.'
-                    ))
-            elif not np.isnan(corr_matrix[j, i]) and not np.isnan(corr_matrix[i, j]):
-                assert corr_matrix[i,j] == corr_matrix[j,i]
-    
-
-    corr_df = pd.DataFrame(corr_matrix, index=all_peptides, columns=all_peptides)
-    
-    return corr_df
diff --git a/tests/tl/test_copro.py b/tests/tl/test_copro.py
index 8d80969..fba1f9f 100644
--- a/tests/tl/test_copro.py
+++ b/tests/tl/test_copro.py
@@ -13,24 +13,27 @@
     peptide_dendograms_by_correlation_,
     peptide_clusters_from_dendograms_,
     proteoform_scores_,
-    )
+)
 
 from proteopy.utils.data_structures import ListDict
+from proteopy.utils._matrix_wrangling import (
+    reconstruct_symmetric_matrix_from_long,
+)
 
 from tests.utils.helpers import (
     transform_dendogram_r2py,
     remap_dendogram_leaf_order,
-    reconstruct_corrs_df_symmetric_from_long_df,
     check_dendogram_equality,
-    )
+)
 
 TEST_DIR = Path(__file__).parent.parent
 DATA_DIR = TEST_DIR / "data"
 
 NOISE = 1e6
 
+
 def compare_clusters_dsVlist(ds, ref):
-    '''
+    """
     Check if peptide cluster annotations defined in a pd.DataSeries are the
     same as a refrence list of clusters annotations.
 
@@ -38,7 +41,7 @@ def compare_clusters_dsVlist(ds, ref):
         ds (pd.DataSeries): values are the categorical annotations (clusters) and
             the indices represent the (peptide) labels.
         ref (list): reference cluster annotations.
-    '''
+    """
     groups = ds.groupby(ds).groups
     clusters_ds = [v.tolist() for _, v in groups.items()]
     clusters_ds = [tuple(sorted(c)) for c in clusters_ds]
@@ -46,136 +49,151 @@ def compare_clusters_dsVlist(ds, ref):
     ref = [tuple(sorted(c)) for c in ref]
     ref_log = set(ref)
 
-    counter=0
+    counter = 0
     for c in clusters_ds:
-        counter+=1
+        counter += 1
         assert c in ref
         ref_log.remove(c)
 
     assert len(ref_log) == 0
 
+
 @pytest.fixture
 def traces_preproc():
-    '''
+    """
     Get COPF mouse tissue pre-processed traces df.
-    '''
-    traces_path = DATA_DIR / 'mouse_tissue/traces_pre-processed_rcopf.tsv'
-    traces = pd.read_csv(traces_path, sep='\t', header=0)
-    traces = traces.rename(columns={'id': 'peptide_id'})
+    """
+    traces_path = DATA_DIR / "mouse_tissue/traces_pre-processed_rcopf.tsv"
+    traces = pd.read_csv(traces_path, sep="\t", header=0)
+    traces = traces.rename(columns={"id": "peptide_id"})
 
     return traces
 
+
 @pytest.fixture
 def traces_preproc_anns():
-    '''
+    """
     Get COPF mouse tissue pre-processed traces annotations df.
-    '''
+    """
     anns_path = (
-            DATA_DIR / 'mouse_tissue/traces_pre-processed_trace-annotations_rcopf.tsv'
-            )
-    anns = pd.read_csv(anns_path, sep='\t', header=0)
-    anns = anns.rename(columns={'id': 'peptide_id'})
+        DATA_DIR
+        / "mouse_tissue/traces_pre-processed_trace-annotations_rcopf.tsv"
+    )
+    anns = pd.read_csv(anns_path, sep="\t", header=0)
+    anns = anns.rename(columns={"id": "peptide_id"})
 
     return anns
 
 
 @pytest.fixture
 def traces_preproc_ext(traces_preproc, traces_preproc_anns):
-    '''
+    """
     Extend pre-processed traces with annotations.
-    '''
-    anns_select = traces_preproc_anns[['peptide_id', 'protein_id']]
-    traces_ext = traces_preproc.merge(anns_select, on='peptide_id')
-    traces_ext = pd.melt(traces_ext, id_vars=('protein_id', 'peptide_id'))
-    traces_ext = traces_ext.rename(columns={'value': 'intensity', 'variable': 'sample'})
+    """
+    anns_select = traces_preproc_anns[["peptide_id", "protein_id"]]
+    traces_ext = traces_preproc.merge(anns_select, on="peptide_id")
+    traces_ext = pd.melt(traces_ext, id_vars=("protein_id", "peptide_id"))
+    traces_ext = traces_ext.rename(
+        columns={"value": "intensity", "variable": "sample"}
+    )
 
     return traces_ext
 
 
 @pytest.fixture
 def fraction_annotation():
-    frac_ann_path = DATA_DIR / \
-        'mouse_tissue/fraction_annotation.tsv'
-    frac_ann = pd.read_csv(frac_ann_path, sep='\t', header=0)
-    frac_ann = frac_ann.copy(deep=True).set_index('filename')
+    frac_ann_path = DATA_DIR / "mouse_tissue/fraction_annotation.tsv"
+    frac_ann = pd.read_csv(frac_ann_path, sep="\t", header=0)
+    frac_ann = frac_ann.copy(deep=True).set_index("filename")
 
     return frac_ann
 
 
 @pytest.fixture
 def traces_corrs():
-    '''
+    """
     Get COPF mouse tissue correlations df.
-    '''
-    traces_corrs_path = DATA_DIR / 'mouse_tissue/traces_correlations_rcopf.tsv'
-    col_names = ['pepA', 'pepB', 'PCC', 'protein_id']
-    df = pd.read_csv(traces_corrs_path, sep='\t', names=col_names)
+    """
+    traces_corrs_path = DATA_DIR / "mouse_tissue/traces_correlations_rcopf.tsv"
+    col_names = ["pepA", "pepB", "PCC", "protein_id"]
+    df = pd.read_csv(traces_corrs_path, sep="\t", names=col_names)
 
     return df
 
 
 @pytest.fixture
 def traces_corrs_ref(traces_corrs):
-    '''
-    Filter COPF mouse tissue correlations df for 
+    """
+    Filter COPF mouse tissue correlations df for
     unique (non-symmetrical) correlation values.
-    '''
-    corrs_ref = traces_corrs.set_index('protein_id')
-    corrs_ref = corrs_ref[corrs_ref['PCC'] != 1]
+    """
+    corrs_ref = traces_corrs.set_index("protein_id")
+    corrs_ref = corrs_ref[corrs_ref["PCC"] != 1]
 
-    sort_peps_ab = lambda row: tuple(sorted([row['pepA'], row['pepB']]))
+    sort_peps_ab = lambda row: tuple(sorted([row["pepA"], row["pepB"]]))
 
-    corrs_ref['sorted_pair'] = corrs_ref.apply(sort_peps_ab, axis=1)
-    corrs_ref = corrs_ref.drop_duplicates(subset=['sorted_pair'])
-    corrs_ref = corrs_ref.drop(columns=['sorted_pair'])
-    corrs_ref = corrs_ref.sort_values(['pepA', 'pepB']).sort_index()
+    corrs_ref["sorted_pair"] = corrs_ref.apply(sort_peps_ab, axis=1)
+    corrs_ref = corrs_ref.drop_duplicates(subset=["sorted_pair"])
+    corrs_ref = corrs_ref.drop(columns=["sorted_pair"])
+    corrs_ref = corrs_ref.sort_values(["pepA", "pepB"]).sort_index()
 
     return corrs_ref
 
 
-def test_pairwise_peptide_correlations_vs_rcopf(traces_preproc_ext, traces_corrs_ref):
-    '''
+def test_pairwise_peptide_correlations_vs_rcopf(
+    traces_preproc_ext, traces_corrs_ref
+):
+    """
     Test pairwise_peptide_correlations() application for equality to rCOPF correlations df.
     Uses COPF mouse tissue dataset as reference results.
-    '''
+    """
 
     # Apply pairwise_peptide_correlations on the entire mouse tissue df
-    pep_corrs = lambda x: pairwise_peptide_correlations_(x,
-                                                        sample_column='sample',
-                                                        peptide_column='peptide_id',
-                                                        value_column='intensity')
-
+    pep_corrs = lambda x: pairwise_peptide_correlations_(
+        x,
+        sample_column="sample",
+        peptide_column="peptide_id",
+        value_column="intensity",
+    )
 
-    corrs = traces_preproc_ext.groupby('protein_id').apply(pep_corrs, include_groups=False)
+    corrs = traces_preproc_ext.groupby("protein_id").apply(
+        pep_corrs, include_groups=False
+    )
     corrs = corrs.droplevel(1, axis=0)
-    corrs = corrs.sort_values(['pepA', 'pepB']).sort_index()
+    corrs = corrs.sort_values(["pepA", "pepB"]).sort_index()
 
     # Compare to rCOPF reference output
-    pep_cols = ['pepA', 'pepB']
-    assert corrs[pep_cols].equals(traces_corrs_ref[pep_cols]) # Both prev. sorted
-
-    abs_tolerance = 1e-14 # loaded reference corrs precision 1e-15
-    assert corrs['PCC'].values == approx(traces_corrs_ref['PCC'].values, abs=abs_tolerance)
+    pep_cols = ["pepA", "pepB"]
+    assert corrs[pep_cols].equals(
+        traces_corrs_ref[pep_cols]
+    )  # Both prev. sorted
+
+    abs_tolerance = 1e-14  # loaded reference corrs precision 1e-15
+    assert corrs["PCC"].values == approx(
+        traces_corrs_ref["PCC"].values, abs=abs_tolerance
+    )
 
 
 @pytest.fixture
 def prot_dends():
 
-    clusts_ref_path = DATA_DIR / 'mouse_tissue/traces_cluster-dendograms_rcopf.json'
+    clusts_ref_path = (
+        DATA_DIR / "mouse_tissue/traces_cluster-dendograms_rcopf.json"
+    )
 
-    with open(clusts_ref_path, 'r') as f:
+    with open(clusts_ref_path) as f:
         dends_R = json.load(f)
 
     # Reformat to match python sklearn dendograms
     dends = {}
 
     for prot_id in dends_R:
-    
+
         dend = dends_R[prot_id]
         dend = transform_dendogram_r2py(dend)
 
-        if isinstance(dend['heights'], float):
-            dend['heights'] = [dend['heights']]
+        if isinstance(dend["heights"], float):
+            dend["heights"] = [dend["heights"]]
 
         dends[prot_id] = dend
 
@@ -187,18 +205,22 @@ def test_peptide_dendograms_by_correlation_vs_rcopf(traces_corrs, prot_dends):
     # Construct map: {protein: dendogram}
     dends = {}
 
-    for protein_id, df in traces_corrs.groupby('protein_id'):
+    for protein_id, df in traces_corrs.groupby("protein_id"):
 
-        corr_df_sym = reconstruct_corrs_df_symmetric_from_long_df(df, var_a_col='pepA', var_b_col='pepB', corr_col='PCC')
+        corr_df_sym = reconstruct_symmetric_matrix_from_long(
+            df, var_a_col="pepA", var_b_col="pepB", value_col="PCC"
+        )
         corr_dists = 1 - corr_df_sym
         dends[protein_id] = peptide_dendograms_by_correlation_(corr_dists)
 
     dends_ref = copy.deepcopy(prot_dends)
 
-    # Remap 
+    # Remap
     for prot_id, dend in dends_ref.items():
 
-        dend_corrected = remap_dendogram_leaf_order(dend, ref_labels=dends[prot_id]['labels'])
+        dend_corrected = remap_dendogram_leaf_order(
+            dend, ref_labels=dends[prot_id]["labels"]
+        )
         dends_ref[prot_id] = dend_corrected
 
     # Equal dendogram dict structure
@@ -206,14 +228,14 @@ def test_peptide_dendograms_by_correlation_vs_rcopf(traces_corrs, prot_dends):
     assert len(dends_ref.keys()) == len(dends.keys())
 
     for prot_id in dends_ref.keys():
-        abs_tolerance = 1e-4 # loaded reference heights precision = 1e-4
-        check_dendogram_equality(dends[prot_id],
-                                 dends_ref[prot_id],
-                                 abs_tolerance=abs_tolerance)
+        abs_tolerance = 1e-4  # loaded reference heights precision = 1e-4
+        check_dendogram_equality(
+            dends[prot_id], dends_ref[prot_id], abs_tolerance=abs_tolerance
+        )
 
 
 def test_peptide_clusters_from_dendograms_():
-    '''Test protein-level peptide_clusters_from_dendograms_() on a single peptide group.'''
+    """Test protein-level peptide_clusters_from_dendograms_() on a single peptide group."""
     # Using dendogram-based toy data
     #
     #                                     (11)
@@ -234,61 +256,54 @@ def test_peptide_clusters_from_dendograms_():
     #   cluster numbers may be different order, which is accounted for in test comparisons.
 
     dendogram = {
-        'type': 'sklearn_agglomerative_clustering',
-        'labels': ['pepA', 'pepB', 'pepC', 'pepD', 'pepE', 'pepF'],
-        'merge': [[0,1], [2,3], [4,5], [6,7], [8,9]],
-        'heights': [0.1, 0.2, 0.4, 0.8, 0.9]
-        }
+        "type": "sklearn_agglomerative_clustering",
+        "labels": ["pepA", "pepB", "pepC", "pepD", "pepE", "pepF"],
+        "merge": [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]],
+        "heights": [0.1, 0.2, 0.4, 0.8, 0.9],
+    }
 
     # Config 1
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 1,
-        min_peptides_per_cluster=1)
-    expected_clusters = [['pepA', 'pepB', 'pepC', 'pepD', 'pepE', 'pepF']]
+        dendogram, n_clusters=1, min_peptides_per_cluster=1
+    )
+    expected_clusters = [["pepA", "pepB", "pepC", "pepD", "pepE", "pepF"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
 
     # Config 2
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 1,
-        min_peptides_per_cluster=2)
-    expected_clusters = [['pepA', 'pepB', 'pepC', 'pepD', 'pepE', 'pepF']]
+        dendogram, n_clusters=1, min_peptides_per_cluster=2
+    )
+    expected_clusters = [["pepA", "pepB", "pepC", "pepD", "pepE", "pepF"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
 
     # Config 3
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 2,
-        min_peptides_per_cluster=1)
-    expected_clusters = [['pepA', 'pepB', 'pepC', 'pepD'], ['pepE', 'pepF']]
+        dendogram, n_clusters=2, min_peptides_per_cluster=1
+    )
+    expected_clusters = [["pepA", "pepB", "pepC", "pepD"], ["pepE", "pepF"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
 
     # Config 4
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 2,
-        min_peptides_per_cluster=2)
-    expected_clusters = [['pepA', 'pepB', 'pepC', 'pepD'], ['pepE', 'pepF']]
+        dendogram, n_clusters=2, min_peptides_per_cluster=2
+    )
+    expected_clusters = [["pepA", "pepB", "pepC", "pepD"], ["pepE", "pepF"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
 
     # Config 5
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 3,
-        min_peptides_per_cluster=1)
-    expected_clusters = [['pepA', 'pepB'], ['pepC', 'pepD'], ['pepE', 'pepF']]
+        dendogram, n_clusters=3, min_peptides_per_cluster=1
+    )
+    expected_clusters = [["pepA", "pepB"], ["pepC", "pepD"], ["pepE", "pepF"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
 
     # Config 6
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 3,
-        min_peptides_per_cluster=2)
-    expected_clusters = [['pepA', 'pepB'], ['pepC', 'pepD'], ['pepE', 'pepF']]
+        dendogram, n_clusters=3, min_peptides_per_cluster=2
+    )
+    expected_clusters = [["pepA", "pepB"], ["pepC", "pepD"], ["pepE", "pepF"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
 
-
     # Using dendogram-based toy data
     #
     #                                     (8)
@@ -307,65 +322,57 @@ def test_peptide_clusters_from_dendograms_():
     #        2           1         0           0         x   n_clust=3, min_pep=2
 
     dendogram = {
-        'type': 'sklearn_agglomerative_clustering',
-        'labels': ['pepA', 'pepB', 'pepC', 'pepD', 'pepE'],
-        'merge': [[0,1], [2,3], [5,6], [4,7]],
-        'heights': [0.1, 0.2, 0.4, 0.8]
-        }
+        "type": "sklearn_agglomerative_clustering",
+        "labels": ["pepA", "pepB", "pepC", "pepD", "pepE"],
+        "merge": [[0, 1], [2, 3], [5, 6], [4, 7]],
+        "heights": [0.1, 0.2, 0.4, 0.8],
+    }
 
     # Config 1
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 1,
-        min_peptides_per_cluster=1)
-    expected_clusters = [['pepA', 'pepB', 'pepC', 'pepD', 'pepE']]
+        dendogram, n_clusters=1, min_peptides_per_cluster=1
+    )
+    expected_clusters = [["pepA", "pepB", "pepC", "pepD", "pepE"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
 
     # Config 2
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 1,
-        min_peptides_per_cluster=2)
-    expected_clusters = [['pepA', 'pepB', 'pepC', 'pepD', 'pepE']]
+        dendogram, n_clusters=1, min_peptides_per_cluster=2
+    )
+    expected_clusters = [["pepA", "pepB", "pepC", "pepD", "pepE"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
 
     # Config 3
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 2,
-        min_peptides_per_cluster=1)
-    expected_clusters = [['pepA', 'pepB', 'pepC', 'pepD'], ['pepE']]
+        dendogram, n_clusters=2, min_peptides_per_cluster=1
+    )
+    expected_clusters = [["pepA", "pepB", "pepC", "pepD"], ["pepE"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
 
     # Config 4
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 2,
-        min_peptides_per_cluster=2)
-    expected_clusters = [['pepA', 'pepB'], ['pepC', 'pepD'], ['pepE']]
+        dendogram, n_clusters=2, min_peptides_per_cluster=2
+    )
+    expected_clusters = [["pepA", "pepB"], ["pepC", "pepD"], ["pepE"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
-    assert clusters['pepE'] == NOISE
+    assert clusters["pepE"] == NOISE
 
     # Config 5
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 3,
-        min_peptides_per_cluster=1)
-    expected_clusters = [['pepA', 'pepB'], ['pepC', 'pepD'], ['pepE']]
+        dendogram, n_clusters=3, min_peptides_per_cluster=1
+    )
+    expected_clusters = [["pepA", "pepB"], ["pepC", "pepD"], ["pepE"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
 
     # Config 6
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 3,
-        min_peptides_per_cluster=2)
-    expected_clusters = [['pepA', 'pepB', 'pepC', 'pepD', 'pepE']]
+        dendogram, n_clusters=3, min_peptides_per_cluster=2
+    )
+    expected_clusters = [["pepA", "pepB", "pepC", "pepD", "pepE"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
     assert clusters.nunique() == 1
     assert clusters.iloc[0] == NOISE
 
-
-    
     # Using correlation based toy data
     #
     #                                     (8)
@@ -383,28 +390,33 @@ def test_peptide_clusters_from_dendograms_():
     #        2           2         1           1         0   n_clust=3, min_pep=1
     #        2           1         0           0         x   n_clust=3, min_pep=2
 
-    corrs = pd.DataFrame({
-        'pepA': [0, 1, 3, 3, 4],
-        'pepB': [1, 0, 3, 3, 4],
-        'pepC': [3, 3, 0, 2, 4],
-        'pepD': [3, 3, 2, 0, 4],
-        'pepE': [4, 4, 4, 4, 0],
-        }, index=['pepA', 'pepB', 'pepC', 'pepD', 'pepE'])
+    corrs = pd.DataFrame(
+        {
+            "pepA": [0, 1, 3, 3, 4],
+            "pepB": [1, 0, 3, 3, 4],
+            "pepC": [3, 3, 0, 2, 4],
+            "pepD": [3, 3, 2, 0, 4],
+            "pepE": [4, 4, 4, 4, 0],
+        },
+        index=["pepA", "pepB", "pepC", "pepD", "pepE"],
+    )
 
-    model = AgglomerativeClustering(n_clusters=None,
-                                    metric='precomputed',
-                                    linkage='average',
-                                    distance_threshold=0,
-                                    compute_distances=True)
+    model = AgglomerativeClustering(
+        n_clusters=None,
+        metric="precomputed",
+        linkage="average",
+        distance_threshold=0,
+        compute_distances=True,
+    )
 
     model.fit(corrs)
 
     # pylint: disable=no-member
     dendogram = {
-        'type': 'sklearn_agglomerative_clustering', 
-        'labels': model.feature_names_in_.tolist(),
-        'heights': model.distances_.tolist(),
-        'merge': model.children_.tolist()
+        "type": "sklearn_agglomerative_clustering",
+        "labels": model.feature_names_in_.tolist(),
+        "heights": model.distances_.tolist(),
+        "merge": model.children_.tolist(),
     }
     # pylint: enable=no-member
 
@@ -412,94 +424,89 @@ def test_peptide_clusters_from_dendograms_():
 
     # Config 1
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 1,
-        min_peptides_per_cluster=1)
-    expected_clusters = [['pepA', 'pepB', 'pepC', 'pepD', 'pepE']]
+        dendogram, n_clusters=1, min_peptides_per_cluster=1
+    )
+    expected_clusters = [["pepA", "pepB", "pepC", "pepD", "pepE"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
 
     # Config 2
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 1,
-        min_peptides_per_cluster=2)
-    expected_clusters = [['pepA', 'pepB', 'pepC', 'pepD', 'pepE']]
+        dendogram, n_clusters=1, min_peptides_per_cluster=2
+    )
+    expected_clusters = [["pepA", "pepB", "pepC", "pepD", "pepE"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
 
     # Config 3
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 2,
-        min_peptides_per_cluster=1)
-    expected_clusters = [['pepA', 'pepB', 'pepC', 'pepD'], ['pepE']]
+        dendogram, n_clusters=2, min_peptides_per_cluster=1
+    )
+    expected_clusters = [["pepA", "pepB", "pepC", "pepD"], ["pepE"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
 
     # Config 4
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 2,
-        min_peptides_per_cluster=2)
-    expected_clusters = [['pepA', 'pepB'], ['pepC', 'pepD'], ['pepE']]
+        dendogram, n_clusters=2, min_peptides_per_cluster=2
+    )
+    expected_clusters = [["pepA", "pepB"], ["pepC", "pepD"], ["pepE"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
-    assert clusters['pepE'] == NOISE
+    assert clusters["pepE"] == NOISE
 
     # Config 5
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 3,
-        min_peptides_per_cluster=1)
-    expected_clusters = [['pepA', 'pepB'], ['pepC', 'pepD'], ['pepE']]
+        dendogram, n_clusters=3, min_peptides_per_cluster=1
+    )
+    expected_clusters = [["pepA", "pepB"], ["pepC", "pepD"], ["pepE"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
 
     # Config 6
     clusters = peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters = 3,
-        min_peptides_per_cluster=2)
-    expected_clusters = [['pepA', 'pepB', 'pepC', 'pepD', 'pepE']]
+        dendogram, n_clusters=3, min_peptides_per_cluster=2
+    )
+    expected_clusters = [["pepA", "pepB", "pepC", "pepD", "pepE"]]
     compare_clusters_dsVlist(clusters, expected_clusters)
     assert clusters.nunique() == 1
     assert clusters.iloc[0] == NOISE
 
+
 @pytest.fixture
 def prot_clust_ann():
-    '''Get protein-level peptide cluster annotations from COPF mouse tissue df.'''
+    """Get protein-level peptide cluster annotations from COPF mouse tissue df."""
 
     clusts_ann_path = DATA_DIR / (
-        'mouse_tissue/traces_annotation_cluster-assignment_rcopf.tsv'
-        )
-    clusts_ann = pd.read_csv(clusts_ann_path, sep='\t', header=0)
+        "mouse_tissue/traces_annotation_cluster-assignment_rcopf.tsv"
+    )
+    clusts_ann = pd.read_csv(clusts_ann_path, sep="\t", header=0)
 
     return clusts_ann
 
-def test_peptide_clusters_from_dendograms_vs_rcopf_(prot_dends, prot_clust_ann):
-    '''
+
+def test_peptide_clusters_from_dendograms_vs_rcopf_(
+    prot_dends, prot_clust_ann
+):
+    """
     Test protein-level peptide_clusters_from_dendograms_()
     on an rCOPF-derived reference dataset.
-    '''
+    """
     cluster_ann_ref_df = copy.deepcopy(prot_clust_ann)
 
     # Compute protein-level cluster annotations
     cluster_ann = {}
     for prot, dend in prot_dends.items():
         dend_upd = copy.deepcopy(dend)
-        dend_upd['type'] = 'sklearn_agglomerative_clustering'
+        dend_upd["type"] = "sklearn_agglomerative_clustering"
         clusters = peptide_clusters_from_dendograms_(
-                dend_upd,
-                n_clusters=2,
-                min_peptides_per_cluster=2,
-                noise=NOISE
-            )
+            dend_upd, n_clusters=2, min_peptides_per_cluster=2, noise=NOISE
+        )
         cluster_ann[prot] = clusters
-    
+
     # Format reference cluster annotations
-    cluster_ann_ref_df = cluster_ann_ref_df.set_index('id')
-    prot_ids = cluster_ann_ref_df['protein_id'].unique()
+    cluster_ann_ref_df = cluster_ann_ref_df.set_index("id")
+    prot_ids = cluster_ann_ref_df["protein_id"].unique()
 
     cluster_ann_ref = {}
     for p in prot_ids:
-        clusters = cluster_ann_ref_df[cluster_ann_ref_df['protein_id'] == p]
-        clusters_map = clusters['cluster'].to_dict()
+        clusters = cluster_ann_ref_df[cluster_ann_ref_df["protein_id"] == p]
+        clusters_map = clusters["cluster"].to_dict()
 
         clusters_map_inv = ListDict()
         for pep, clust in clusters_map.items():
@@ -521,66 +528,67 @@ def test_peptide_clusters_from_dendograms_vs_rcopf_(prot_dends, prot_clust_ann):
 
 @pytest.fixture
 def trace_annotation_proteoform_scores():
-    '''Get protein-level proteoform annotations from COPF mouse tissue dataset.'''
+    """Get protein-level proteoform annotations from COPF mouse tissue dataset."""
 
     annotation_path = DATA_DIR / (
-        'mouse_tissue/trace_annotation_proteoform-scores_rcopf.tsv'
-        )
-    clusts_ann = pd.read_csv(annotation_path, sep='\t', header=0)
+        "mouse_tissue/trace_annotation_proteoform-scores_rcopf.tsv"
+    )
+    clusts_ann = pd.read_csv(annotation_path, sep="\t", header=0)
 
     return clusts_ann
 
+
 def test_proteoform_scores_vs_rcopf_(
     traces_corrs,
     prot_clust_ann,
     fraction_annotation,
     trace_annotation_proteoform_scores,
     summary_func=np.mean,
-    ):
+):
 
     n_fractions = len(fraction_annotation)
 
     columns = [
-        'protein_id',
-        'proteoform_score',
-        'proteoform_score_z',
-        'proteoform_score_dz',
-        'proteoform_score_pval',
-        ]
+        "protein_id",
+        "proteoform_score",
+        "proteoform_score_z",
+        "proteoform_score_dz",
+        "proteoform_score_pval",
+    ]
 
     proteoform_scores_list = []
 
-    for prot, corrs in traces_corrs.groupby('protein_id'):
+    for prot, corrs in traces_corrs.groupby("protein_id"):
 
-        corrs_mat = reconstruct_corrs_df_symmetric_from_long_df(
-            corrs,
-            var_a_col='pepA',
-            var_b_col='pepB',
-            corr_col='PCC')
+        corrs_mat = reconstruct_symmetric_matrix_from_long(
+            corrs, var_a_col="pepA", var_b_col="pepB", value_col="PCC"
+        )
 
-        clusters = prot_clust_ann[prot_clust_ann['protein_id'] == prot]
-        clusters = clusters.set_index('id')['cluster']
+        clusters = prot_clust_ann[prot_clust_ann["protein_id"] == prot]
+        clusters = clusters.set_index("id")["cluster"]
         clusters[clusters == 100] = NOISE
 
         scores = proteoform_scores_(
-            corrs_mat,
-            clusters,
-            n_fractions,
-            summary_func=np.mean)
+            corrs_mat, clusters, n_fractions, summary_func=np.mean
+        )
 
-        scores_entry = {column:value for column, value in zip(columns[1:5], scores)}
-        scores_entry['protein_id'] = prot
+        scores_entry = {
+            column: value for column, value in zip(columns[1:5], scores)
+        }
+        scores_entry["protein_id"] = prot
         scores_entry = pd.DataFrame([scores_entry])
         proteoform_scores_list.append(scores_entry)
 
     proteoform_scores = pd.concat(proteoform_scores_list, ignore_index=True)
 
-    proteoform_scores = proteoform_scores.loc[:, columns].set_index('protein_id')
+    proteoform_scores = proteoform_scores.loc[:, columns].set_index(
+        "protein_id"
+    )
 
     # Format reference proteoforms
     proteoform_scores_ref = trace_annotation_proteoform_scores[columns]
     proteoform_scores_ref = proteoform_scores_ref.drop_duplicates()
-    proteoform_scores_ref = proteoform_scores_ref.set_index('protein_id')
+    proteoform_scores_ref = proteoform_scores_ref.set_index("protein_id")
 
     # Compare
     pfs_idx = set(proteoform_scores.index)
@@ -588,13 +596,23 @@ def test_proteoform_scores_vs_rcopf_(
 
     assert len(pfs_idx.symmetric_difference(pfs_ref_idx)) == 0
     assert len(pfs_idx) == len(pfs_ref_idx)
-    assert len(set(proteoform_scores.columns).symmetric_difference(set(proteoform_scores_ref.columns))) == 0
+    assert (
+        len(
+            set(proteoform_scores.columns).symmetric_difference(
+                set(proteoform_scores_ref.columns)
+            )
+        )
+        == 0
+    )
 
-    proteoform_scores = proteoform_scores.loc[proteoform_scores_ref.index,proteoform_scores_ref.columns]
+    proteoform_scores = proteoform_scores.loc[
+        proteoform_scores_ref.index, proteoform_scores_ref.columns
+    ]
 
     assert np.allclose(
         proteoform_scores,
         proteoform_scores_ref,
         rtol=0,
         atol=1e-12,
-        equal_nan=True)
+        equal_nan=True,
+    )
diff --git a/tests/utils/helpers.py b/tests/utils/helpers.py
index 3bbba95..51761b0 100644
--- a/tests/utils/helpers.py
+++ b/tests/utils/helpers.py
@@ -1,34 +1,7 @@
 import numpy as np
-import pandas as pd
 import copy
 from pytest import approx
 
-from proteopy.utils.copf import reconstruct_corrs_df_symmetric_from_long_df
-
-def test_reconstruct_corrs_df_symmetric_from_long_df():
-
-    # labels: a-c
-    #
-    # [[ x, 0.1, 0.5],        [[ 1  , 0.1, 0.5],
-    #  [ x, 1  , x  ],   ==>   [ 0.1, 1  , 0.4],
-    #  [ x  0.4, 1  ]]         [ 0.5, 0.4, 1  ]]
-
-    df = pd.DataFrame({
-        'colA': ['a', 'a', 'b', 'c', 'c'],
-        'colB': ['b', 'c', 'b', 'b', 'c'],
-        'value':[0.1, 0.5, 1, 0.4, 1]
-        })
-
-    df_expected = pd.DataFrame({
-        'a': [1, 0.1, 0.5],
-        'b': [0.1, 1, 0.4],
-        'c':[0.5, 0.4, 1]
-        }, index=['a', 'b', 'c'])
-
-    df_reconstructed = reconstruct_corrs_df_symmetric_from_long_df(df, 'colA', 'colB', 2)
-    assert np.isclose(df_reconstructed, df_expected, atol=1e-4).all().all()
-    assert all(df_reconstructed.index == df_reconstructed.columns)
-
 
 def transform_dendogram_merge_arr_r2py(merge_arr: list):
 
@@ -37,35 +10,35 @@ def transform_dendogram_merge_arr_r2py(merge_arr: list):
 
     merge_new = np.where(merge_new > 0, merge_new + n_samples, merge_new)
     merge_new = np.abs(merge_new)
-    merge_new = merge_new - 1 # 1-based -> 0-based order
+    merge_new = merge_new - 1  # 1-based -> 0-based order
 
     return merge_new
 
 
 def transform_dendogram_r2py(dendogram: dict):
-    '''
+    """
     Parameters:
     -----------
     dendogram: dict
         Dictionary with the following structure: {..., merge = [[int, int], ...]}
-    '''
+    """
 
-    if not 'merge' in dendogram.keys():
-        raise ValueError('Dendogram slot missing!')
+    if not "merge" in dendogram.keys():
+        raise ValueError("Dendogram slot missing!")
 
-    merge = dendogram['merge']
+    merge = dendogram["merge"]
     merge_new = transform_dendogram_merge_arr_r2py(merge)
 
     dendogram_new = copy.deepcopy(dendogram)
-    dendogram_new['merge'] = merge_new.tolist()
+    dendogram_new["merge"] = merge_new.tolist()
 
     return dendogram_new
 
 
 def remap_dendogram_leaf_order(dendogram: dict, ref_labels: list):
-    '''
+    """
     Remap nodes in dendogram['merge'] using a reference label order.
-    
+
     Parameters:
     -----------
     - dendogram: dict
@@ -73,31 +46,33 @@ def remap_dendogram_leaf_order(dendogram: dict, ref_labels: list):
         - merge: np.ndarray of shape (n_samples-1, 2)
         - heights: list of length n_samples
     - ref_annotation: list of labels in desired new leaf order
-    
+
     Returns:
     --------
     - dendogram with updated node indices remapped to match ref_annotation order
-    '''
-    orig_labels = dendogram['labels']
+    """
+    orig_labels = dendogram["labels"]
     n_samples = len(orig_labels)
 
-    assert set(orig_labels) == set(ref_labels), f'orig_labels: {orig_labels},\nref_labels: {ref_labels}'
+    assert set(orig_labels) == set(
+        ref_labels
+    ), f"orig_labels: {orig_labels},\nref_labels: {ref_labels}"
     assert len(orig_labels) == len(ref_labels)
-    assert len(orig_labels) == len(dendogram['merge']) + 1
+    assert len(orig_labels) == len(dendogram["merge"]) + 1
+
+    merge_arr = np.array(dendogram["merge"])
 
-    merge_arr = np.array(dendogram['merge'])
-    
     # Mapping from original index to ref index
     orig_label_to_index = {label: i for i, label in enumerate(orig_labels)}
     ref_label_to_index = {label: i for i, label in enumerate(ref_labels)}
-    
+
     # Create remapping array
     leaf_map = np.zeros(n_samples, dtype=int)
     for label in orig_labels:
         orig_idx = orig_label_to_index[label]
         ref_idx = ref_label_to_index[label]
         leaf_map[orig_idx] = ref_idx
-    
+
     # Now remap only values < n_leaves
     merge_remapped = merge_arr.copy()
 
@@ -108,37 +83,43 @@ def remap_dendogram_leaf_order(dendogram: dict, ref_labels: list):
 
     # Replace old merge
     dendogram_remapped = copy.deepcopy(dendogram)
-    dendogram_remapped['labels'] = ref_labels
-    dendogram_remapped['merge'] = merge_remapped.tolist()
-    
+    dendogram_remapped["labels"] = ref_labels
+    dendogram_remapped["merge"] = merge_remapped.tolist()
+
     return dendogram_remapped
 
 
-def check_dendogram_equality(dend, dend_ref, rel_tolerance=None, abs_tolerance=None):
-    '''
+def check_dendogram_equality(
+    dend, dend_ref, rel_tolerance=None, abs_tolerance=None
+):
+    """
     Note:
         To choose the tolerances view API: pytest.approx
-    '''
+    """
 
-    keys = ('labels', 'merge', 'heights')
+    keys = ("labels", "merge", "heights")
 
     # Correct dict keys
-    assert all([key in dend_ref.keys() for key in keys]), f'dend.keys: {list(dend.keys())}\nkeys:{keys}'
-    assert all([key in dend.keys() for key in keys]), f'dend.keys: {list(dend.keys())}\nkeys:{keys}'
+    assert all(
+        [key in dend_ref.keys() for key in keys]
+    ), f"dend.keys: {list(dend.keys())}\nkeys:{keys}"
+    assert all(
+        [key in dend.keys() for key in keys]
+    ), f"dend.keys: {list(dend.keys())}\nkeys:{keys}"
 
     # Equal labels
-    labels_ref = dend_ref['labels']
-    labels = dend['labels']
+    labels_ref = dend_ref["labels"]
+    labels = dend["labels"]
     assert labels_ref == labels
 
     # Equal merge arrays
-    merge_arr_ref = dend_ref['merge']
-    merge_arr = dend['merge']
+    merge_arr_ref = dend_ref["merge"]
+    merge_arr = dend["merge"]
 
     for i, (pair_ref, pair) in enumerate(zip(merge_arr, merge_arr_ref)):
-        assert pair_ref == pair or pair_ref == pair[::-1], f'{i}'
+        assert pair_ref == pair or pair_ref == pair[::-1], f"{i}"
 
     # Equal heights
-    heights_ref = dend_ref['heights']
-    heights = dend['heights']
+    heights_ref = dend_ref["heights"]
+    heights = dend["heights"]
     assert heights == approx(heights_ref, rel=rel_tolerance, abs=abs_tolerance)
diff --git a/tests/utils/test_matrix_wrangling.py b/tests/utils/test_matrix_wrangling.py
new file mode 100644
index 0000000..3732cee
--- /dev/null
+++ b/tests/utils/test_matrix_wrangling.py
@@ -0,0 +1,143 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from proteopy.utils._matrix_wrangling import (
+    reconstruct_symmetric_matrix_from_long,
+)
+
+
+def test_basic_reconstruction():
+    # labels a-c
+    #
+    # [[ x, 0.1, 0.5],        [[ 1  , 0.1, 0.5],
+    #  [ x, 1  , x  ],   ==>   [ 0.1, 1  , 0.4],
+    #  [ x  0.4, 1  ]]         [ 0.5, 0.4, 1  ]]
+    df = pd.DataFrame(
+        {
+            "colA": ["a", "a", "b", "c", "c"],
+            "colB": ["b", "c", "b", "b", "c"],
+            "value": [0.1, 0.5, 1, 0.4, 1],
+        }
+    )
+    expected = pd.DataFrame(
+        {
+            "a": [1.0, 0.1, 0.5],
+            "b": [0.1, 1.0, 0.4],
+            "c": [0.5, 0.4, 1.0],
+        },
+        index=["a", "b", "c"],
+    )
+
+    result = reconstruct_symmetric_matrix_from_long(df, "colA", "colB", 2)
+    assert np.isclose(result, expected, atol=1e-4).all().all()
+
+
+def test_result_is_symmetric_with_sorted_labels():
+    df = pd.DataFrame(
+        {
+            "a": ["c", "c", "a"],
+            "b": ["a", "b", "b"],
+            "v": [0.5, 0.4, 0.1],
+        }
+    )
+    result = reconstruct_symmetric_matrix_from_long(df, "a", "b", "v")
+
+    assert list(result.index) == ["a", "b", "c"]
+    assert list(result.columns) == ["a", "b", "c"]
+    assert np.allclose(result.values, result.values.T)
+
+
+def test_lower_triangle_mirrored_from_upper():
+    # Only upper-triangle pairs supplied.
+    df = pd.DataFrame(
+        {
+            "a": ["a", "a", "b"],
+            "b": ["b", "c", "c"],
+            "v": [0.2, 0.3, 0.4],
+        }
+    )
+    result = reconstruct_symmetric_matrix_from_long(df, "a", "b", "v")
+
+    assert result.loc["b", "a"] == 0.2
+    assert result.loc["c", "a"] == 0.3
+    assert result.loc["c", "b"] == 0.4
+
+
+def test_diagonal_value():
+    df = pd.DataFrame({"a": ["x"], "b": ["y"], "v": [0.7]})
+
+    default = reconstruct_symmetric_matrix_from_long(df, "a", "b", "v")
+    assert list(np.diag(default.values)) == [1.0, 1.0]
+
+    custom = reconstruct_symmetric_matrix_from_long(
+        df, "a", "b", "v", diagonal=0.0
+    )
+    assert list(np.diag(custom.values)) == [0.0, 0.0]
+
+
+def test_column_selectors_int_vs_str_equivalent():
+    df = pd.DataFrame(
+        {
+            "a": ["a", "a", "b"],
+            "b": ["b", "c", "c"],
+            "v": [0.2, 0.3, 0.4],
+        }
+    )
+    by_name = reconstruct_symmetric_matrix_from_long(df, "a", "b", "v")
+    by_index = reconstruct_symmetric_matrix_from_long(df, 0, 1, 2)
+
+    assert np.allclose(by_name.values, by_index.values)
+    assert list(by_name.index) == list(by_index.index)
+
+
+def test_generic_labels_and_values():
+    # Non-"peptide" labels and a non-correlation value column.
+    df = pd.DataFrame(
+        {
+            "src": ["CityA", "CityA", "CityB"],
+            "dst": ["CityB", "CityC", "CityC"],
+            "distance": [10.0, 25.0, 7.0],
+        }
+    )
+    result = reconstruct_symmetric_matrix_from_long(
+        df, "src", "dst", "distance", diagonal=0.0
+    )
+    assert result.loc["CityA", "CityB"] == 10.0
+    assert result.loc["CityC", "CityA"] == 25.0
+    assert result.loc["CityA", "CityA"] == 0.0
+
+
+def test_single_pair_minimal_input():
+    df = pd.DataFrame({"a": ["p"], "b": ["q"], "v": [0.9]})
+    result = reconstruct_symmetric_matrix_from_long(df, "a", "b", "v")
+
+    assert result.shape == (2, 2)
+    assert result.loc["p", "q"] == 0.9
+    assert result.loc["q", "p"] == 0.9
+
+
+def test_missing_pair_raises():
+    # Labels a, b, c present but the (b, c) pair has no value anywhere.
+    df = pd.DataFrame(
+        {
+            "a": ["a", "a"],
+            "b": ["b", "c"],
+            "v": [0.1, 0.2],
+        }
+    )
+    with pytest.raises(ValueError, match="No value found"):
+        reconstruct_symmetric_matrix_from_long(df, "a", "b", "v")
+
+
+def test_conflicting_values_raise():
+    # (a, b) supplied as 0.1 and (b, a) supplied as 0.9.
+    df = pd.DataFrame(
+        {
+            "a": ["a", "b"],
+            "b": ["b", "a"],
+            "v": [0.1, 0.9],
+        }
+    )
+    with pytest.raises(ValueError, match="Conflicting values"):
+        reconstruct_symmetric_matrix_from_long(df, "a", "b", "v")

From 849960e41213d3af555e12760c804e5bd62d82c3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Levente=20Temesv=C3=A1ri-Nagy?=
 <147416790+leventetn@users.noreply.github.com>
Date: Sat, 26 Sep 2026 10:55:54 +0200
Subject: [PATCH 02/10] feature/pairwise_peptide_correlations_heatmap

pr.pl.pairwise_peptide_correlations_heatmap(adata, protein, ...) in
pl/copf.py, registered in pl/init.py. Single protein; reads
adata.uns["pairwise_peptide_correlations"], reconstructs the per-protein
matrix, draws a sns.clustermap with symmetric row/col clustering
(method, unified linkage shadow, cluster toggle) and one-or-more
margin_color annotation strips (default cluster_id). Tick labels from
var["peptide_id"]; 1e6 NOISE left as an ordinary category; full numpydoc
docstring.
---
 HISTORY.md              |  11 ++
 proteopy/pl/__init__.py |   5 +-
 proteopy/pl/copf.py     | 360 +++++++++++++++++++++++++++++++++++--
 proteopy/tl/copf.py     | 387 +++++++++++++++++++++-------------------
 4 files changed, 566 insertions(+), 197 deletions(-)

diff --git a/HISTORY.md b/HISTORY.md
index b139c35..4141158 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -12,6 +12,10 @@ and this project adheres to
 
 **Plotting** (`pr.pl`)
 
+- `pairwise_peptide_correlations_heatmap()`: per-protein clustered
+  peptide-correlation heatmap with proteoform (`cluster_id` /
+  `proteoform_id`) annotation strips, symmetric row/column linkage
+  (`method`, `linkage`, `cluster`) and multiple `margin_color` columns
 - `peptide_intensities()`, `proteoform_intensities()`: new `facet_by`
   parameter splitting the samples (`.obs`) across a grid of subplots
 
@@ -50,6 +54,13 @@ and this project adheres to
 
 ### Changed
 
+**Internal**
+
+- Moved and renamed the private matrix helper
+  `utils.copf.reconstruct_corrs_df_symmetric_from_long_df` to
+  `utils._matrix_wrangling.reconstruct_symmetric_matrix_from_long`
+  (generic labels/values, `raise` instead of `assert` on conflicts)
+
 **Preprocessing** (`pr.pp`)
 
 - `normalize_median()`:
diff --git a/proteopy/pl/__init__.py b/proteopy/pl/__init__.py
index 8f0fba2..f72c78d 100644
--- a/proteopy/pl/__init__.py
+++ b/proteopy/pl/__init__.py
@@ -23,7 +23,10 @@
     hclustv_profiles_heatmap,
 )
 
-from .copf import proteoform_scores
+from .copf import (
+    proteoform_scores,
+    pairwise_peptide_correlations_heatmap,
+)
 from .stat_tests import (
     volcano,
     differential_abundance_box,
diff --git a/proteopy/pl/copf.py b/proteopy/pl/copf.py
index 4d2e3b7..0fa1394 100644
--- a/proteopy/pl/copf.py
+++ b/proteopy/pl/copf.py
@@ -5,14 +5,23 @@
 
 import numpy as np
 import pandas as pd
+import matplotlib as mpl
 import matplotlib.pyplot as plt
 from matplotlib.ticker import LogLocator
+from matplotlib.patches import Patch
 import seaborn as sns
 import anndata as ad
 from matplotlib.axes import Axes
+from scipy.spatial.distance import squareform
+from scipy.cluster.hierarchy import linkage as scipy_linkage
 from adjustText import adjust_text
 
 from proteopy.utils.anndata import check_proteodata
+from proteopy.utils.matplotlib import _resolve_color_scheme
+from proteopy.utils._matrix_wrangling import (
+    reconstruct_symmetric_matrix_from_long,
+)
+
 
 def proteoform_scores(
     adata: ad.AnnData,
@@ -259,18 +268,15 @@ def _validate_threshold(
         if protein_id_key is not None:
             if protein_id_key not in adata.var.columns:
                 raise ValueError(
-                    f"Column '{protein_id_key}' not found "
-                    "in `adata.var`."
+                    f"Column '{protein_id_key}' not found " "in `adata.var`."
                 )
             # Validate 1-to-1 mapping.
             mapping_df = adata.var[
                 ["protein_id", protein_id_key]
             ].drop_duplicates()
-            dup_proteins = (
-                mapping_df
-                .groupby("protein_id")[protein_id_key]
-                .nunique()
-            )
+            dup_proteins = mapping_df.groupby("protein_id")[
+                protein_id_key
+            ].nunique()
             bad = dup_proteins[dup_proteins > 1]
             if not bad.empty:
                 raise ValueError(
@@ -296,7 +302,7 @@ def _validate_threshold(
             # values — resolve them to protein_ids.
             known_labels = set(mapping_df[protein_id_key])
             resolved_pids = set()
-            unknown = (set(highlight_prots) - known_labels)
+            unknown = set(highlight_prots) - known_labels
             if unknown:
                 raise ValueError(
                     "The following values from "
@@ -304,9 +310,7 @@ def _validate_threshold(
                     f"`adata.var['{protein_id_key}']`: "
                     f"{sorted(unknown)}"
                 )
-            highlight_pids = {
-                label_to_pid[v] for v in highlight_prots
-            }
+            highlight_pids = {label_to_pid[v] for v in highlight_prots}
         else:
             pid_to_label = None
             known_ids = set(adata.var["protein_id"])
@@ -377,11 +381,339 @@ def _validate_threshold(
 
     if save is not None:
         if not isinstance(save, (str, Path)):
-            raise TypeError(
-                "`save` must be a path-like object or None."
-            )
+            raise TypeError("`save` must be a path-like object or None.")
         _fig.savefig(save, dpi=300, bbox_inches="tight")
     if show:
         plt.show()
 
     return _ax
+
+
+def pairwise_peptide_correlations_heatmap(
+    adata: ad.AnnData,
+    protein: str,
+    *,
+    corr_key: str = "pairwise_peptide_correlations",
+    margin_color: str | list[str] = "cluster_id",
+    method: str = "average",
+    cluster: bool = True,
+    linkage=None,
+    color_scheme=None,
+    cmap: str = "coolwarm",
+    xticklabels: bool = True,
+    yticklabels: bool = True,
+    figsize: tuple[float, float] = (8.0, 8.0),
+    show: bool = True,
+    ax: bool = False,
+    print_stats: bool = False,
+    save: str | Path | None = None,
+) -> Axes | None:
+    """Clustered peptide-correlation heatmap for a single protein.
+
+    Draw the peptide x peptide Pearson-correlation matrix of one protein
+    as a clustered heatmap, so that peptides of the same proteoform (COPF
+    cluster) group into visible blocks. One or more categorical
+    annotation strips (``margin_color``) are drawn alongside the heatmap,
+    by default the per-protein ``cluster_id`` assignment.
+
+    The correlations are read from
+    ``adata.uns[corr_key]`` (produced by
+    :func:`proteopy.tl.pairwise_peptide_correlations`) and reconstructed
+    into a full symmetric matrix for the selected protein. The clustering
+    is symmetric: the same tree is used for rows and columns.
+
+    Parameters
+    ----------
+    adata : AnnData
+        :class:`~anndata.AnnData` carrying COPF annotations.
+    protein : str
+        ``protein_id`` of the protein to plot.
+    corr_key : str
+        Key in ``adata.uns`` holding the long-form pairwise correlations
+        (columns ``pepA``, ``pepB``, ``PCC``; indexed by ``protein_id``).
+    margin_color : str | list[str]
+        Column(s) in ``adata.var`` used for the annotation strip(s). A
+        string is treated as a single-item list; a list draws one strip
+        per entry, in the given order. Proteoform labels are obtained by
+        passing ``"proteoform_id"``.
+    method : str
+        Linkage method handed to :func:`scipy.cluster.hierarchy.linkage`
+        when the tree is recomputed from ``1 - correlation``.
+    cluster : bool
+        Cluster and draw dendrograms on both axes. When ``False`` the
+        matrix is shown in its input order without dendrograms.
+    linkage : numpy.ndarray | None
+        Precomputed SciPy linkage matrix. When provided it is used for
+        both axes, shadowing the default ``1 - correlation`` computation.
+    color_scheme : Any
+        Color palette specification understood by
+        :func:`proteopy.utils.matplotlib._resolve_color_scheme`.
+    cmap : str
+        Continuous colormap for the heatmap body.
+    xticklabels, yticklabels : bool
+        Whether to show peptide tick labels on each axis.
+    figsize : tuple[float, float]
+        Matplotlib figure size in inches.
+    show : bool
+        Display the figure with :func:`matplotlib.pyplot.show`.
+    ax : bool
+        Return the heatmap :class:`matplotlib.axes.Axes` when ``True``.
+    print_stats : bool
+        Print correlation summary statistics before drawing.
+    save : str | Path | None
+        File path to save the figure. ``None`` skips saving.
+
+    Returns
+    -------
+    Axes or None
+        Heatmap axes when ``ax`` is ``True``; otherwise ``None``.
+
+    Raises
+    ------
+    ValueError
+        If ``protein`` is unknown, ``corr_key`` is missing, the protein
+        has no pairwise correlations, or clustering is requested on a
+        matrix containing missing values.
+    KeyError
+        If a ``margin_color`` column is not present in ``adata.var``.
+
+    Examples
+    --------
+    >>> import proteopy as pr
+    >>> adata = pr.read.long(...)
+    >>> pr.tl.pairwise_peptide_correlations(adata)
+    >>> pr.tl.peptide_dendograms_by_correlation(
+    ...     adata,
+    ...     method='agglomerative-hierarchical-clustering',
+    ... )
+    >>> pr.tl.peptide_clusters_from_dendograms(
+    ...     adata,
+    ...     n_clusters=2,
+    ...     min_peptides_per_cluster=2,
+    ... )
+    >>> pr.pl.pairwise_peptide_correlations_heatmap(adata, protein="P12345")
+
+    Annotate by proteoform and draw two strips:
+
+    >>> pr.pl.pairwise_peptide_correlations_heatmap(
+    ...     adata,
+    ...     protein="P12345",
+    ...     margin_color=["cluster_id", "proteoform_id"],
+    ... )
+    """
+    check_proteodata(adata)
+
+    # -- Validate arguments up front
+    if not isinstance(method, str):
+        raise TypeError("`method` must be a string.")
+    if not isinstance(cluster, bool):
+        raise TypeError("`cluster` must be a bool.")
+
+    if isinstance(margin_color, str):
+        margin_cols = [margin_color]
+    elif isinstance(margin_color, list) and all(
+        isinstance(c, str) for c in margin_color
+    ):
+        margin_cols = list(margin_color)
+    else:
+        raise TypeError("`margin_color` must be a string or list of strings.")
+
+    if protein not in set(adata.var["protein_id"]):
+        raise ValueError(
+            f"protein '{protein}' not found in adata.var['protein_id']."
+        )
+
+    if corr_key not in adata.uns:
+        raise ValueError(
+            f"'{corr_key}' not found in adata.uns; run "
+            "pr.tl.pairwise_peptide_correlations() first."
+        )
+
+    missing_cols = [c for c in margin_cols if c not in adata.var.columns]
+    if missing_cols:
+        raise KeyError(
+            "margin_color column(s) not found in adata.var: "
+            f"{', '.join(missing_cols)}."
+        )
+
+    corrs = adata.uns[corr_key]
+    if protein not in corrs.index:
+        raise ValueError(
+            f"protein '{protein}' has no pairwise correlations in "
+            f"adata.uns['{corr_key}'] (a protein needs at least two "
+            "peptides)."
+        )
+
+    # -- Reconstruct the symmetric peptide x peptide correlation matrix
+    corr_long = corrs.loc[[protein]]
+    corr_df = reconstruct_symmetric_matrix_from_long(
+        corr_long,
+        var_a_col="pepA",
+        var_b_col="pepB",
+        value_col="PCC",
+    )
+
+    peptides = list(corr_df.index)
+
+    # -- Labels come from the peptide_id column, never the index
+    pep_ids = adata.var.loc[peptides, "peptide_id"].astype(str).tolist()
+
+    # -- Build annotation strip(s) and their legend handles
+    annot_colors = pd.DataFrame(index=peptides)
+    legend_handles: list[Patch] = []
+    for col in margin_cols:
+        groups = adata.var.loc[peptides, col]
+        if isinstance(groups.dtype, pd.CategoricalDtype):
+            cats = [c for c in groups.cat.categories if c in set(groups)]
+        else:
+            cats = sorted(map(str, groups.dropna().unique()))
+
+        resolved = _resolve_color_scheme(color_scheme, cats)
+        if resolved is None:
+            resolved = (
+                sns.color_palette(n_colors=len(cats)) if len(cats) else []
+            )
+        palette = {str(cat): color for cat, color in zip(cats, resolved)}
+
+        color_series = groups.astype("string").map(palette)
+        if color_series.isna().any() and groups.notna().any():
+            bad = color_series.isna() & groups.notna()
+            missing_cats = sorted(groups[bad].astype(str).unique())
+            raise ValueError(
+                f"No color provided for categories in '{col}': "
+                f"{', '.join(missing_cats)}."
+            )
+
+        prefix = f"{col}: " if len(margin_cols) > 1 else ""
+        legend_handles.extend(
+            Patch(
+                facecolor=palette[str(cat)],
+                edgecolor="none",
+                label=f"{prefix}{cat}",
+            )
+            for cat in cats
+        )
+
+        if groups.isna().any():
+            na_color = mpl.colors.to_rgba("lightgray")
+            color_series = color_series.astype(object)
+            color_series[groups.isna()] = na_color
+            legend_handles.append(
+                Patch(
+                    facecolor=na_color, edgecolor="none", label=f"{prefix}NA"
+                )
+            )
+
+        annot_colors[col] = color_series.to_numpy()
+
+    annot_colors.index = pep_ids
+    corr_df.index = pep_ids
+    corr_df.columns = pep_ids
+
+    # -- Color center at the off-diagonal mean (as sample_correlation_matrix)
+    A = corr_df.to_numpy(dtype=float)
+    n = A.shape[0]
+    if n > 1:
+        offdiag = A[~np.eye(n, dtype=bool)]
+        center_val = np.nanmean(offdiag)
+    else:
+        offdiag = A.ravel()
+        center_val = float(np.nanmean(A))
+
+    # -- Resolve the row/column linkage (symmetric)
+    row_linkage = None
+    col_linkage = None
+    do_cluster = cluster
+    if cluster:
+        if linkage is not None:
+            row_linkage = col_linkage = linkage
+        else:
+            if np.isnan(A).any():
+                raise ValueError(
+                    "Correlation matrix for protein "
+                    f"'{protein}' contains missing values; clustering "
+                    "is not possible. Pass cluster=False to skip "
+                    "clustering."
+                )
+            dist = 1.0 - A
+            np.fill_diagonal(dist, 0.0)
+            dist = np.clip(dist, 0, 2)
+            Z = scipy_linkage(squareform(dist, checks=False), method=method)
+            row_linkage = col_linkage = Z
+
+    # -- Optional statistics printout
+    if print_stats and n > 1:
+        summary = pd.DataFrame(
+            {
+                "min": [np.nanmin(offdiag)],
+                "max": [np.nanmax(offdiag)],
+                "mean": [np.nanmean(offdiag)],
+                "median": [np.nanmedian(offdiag)],
+                "std": [np.nanstd(offdiag)],
+            }
+        )
+        print(f"Peptide correlation summary (off-diagonal, {protein}):")
+        print(summary.to_string(index=False))
+        print()
+
+        clusters = adata.var.loc[peptides, "cluster_id"]
+        rows = []
+        for cid, grp in clusters.groupby(clusters, observed=True):
+            idx = [peptides.index(p) for p in grp.index]
+            block = A[np.ix_(idx, idx)]
+            if len(idx) > 1:
+                within = block[~np.eye(len(idx), dtype=bool)]
+                mean_within = np.nanmean(within)
+            else:
+                mean_within = np.nan
+            rows.append({"cluster_id": cid, "mean_within": mean_within})
+        print(f"\nPer cluster_id (within-cluster mean correlation):")
+        print(pd.DataFrame(rows).to_string(index=False))
+        print()
+
+    # -- Draw the clustered heatmap
+    clustermap_kwargs = dict(
+        row_colors=annot_colors,
+        col_colors=annot_colors,
+        cmap=cmap,
+        center=center_val,
+        figsize=figsize,
+        xticklabels=xticklabels,
+        yticklabels=yticklabels,
+        cbar_kws={"label": "PCC"},
+    )
+    if do_cluster:
+        clustermap_kwargs["row_linkage"] = row_linkage
+        clustermap_kwargs["col_linkage"] = col_linkage
+    else:
+        clustermap_kwargs["row_cluster"] = False
+        clustermap_kwargs["col_cluster"] = False
+
+    g = sns.clustermap(corr_df, **clustermap_kwargs)
+
+    if legend_handles:
+        g.ax_heatmap.legend(
+            handles=legend_handles,
+            title=", ".join(margin_cols),
+            bbox_to_anchor=(1.05, 1),
+            loc="upper left",
+            borderaxespad=0.0,
+            frameon=False,
+        )
+
+    g.ax_heatmap.set_xlabel("Peptides")
+    g.ax_heatmap.set_ylabel("Peptides")
+    g.figure.suptitle(protein)
+
+    plt.tight_layout()
+
+    if save is not None:
+        if not isinstance(save, (str, Path)):
+            raise TypeError("`save` must be a path-like object or None.")
+        g.savefig(save, dpi=300, bbox_inches="tight")
+    if show:
+        plt.show()
+
+    if ax:
+        return g.ax_heatmap
+    return None
diff --git a/proteopy/tl/copf.py b/proteopy/tl/copf.py
index 923230b..cd13038 100644
--- a/proteopy/tl/copf.py
+++ b/proteopy/tl/copf.py
@@ -8,7 +8,9 @@
 from scipy.stats import norm
 from statsmodels.stats.multitest import multipletests
 
-from proteopy.utils.copf import reconstruct_corrs_df_symmetric_from_long_df
+from proteopy.utils._matrix_wrangling import (
+    reconstruct_symmetric_matrix_from_long,
+)
 from proteopy.utils.data_structures import BinaryClusterTree
 
 NOISE = 1e6
@@ -19,8 +21,8 @@ def pairwise_peptide_correlations_(
     sample_column="filename",
     peptide_column="peptide_id",
     value_column="intensity",
-    ):
-    '''
+):
+    """
     Calculate pairwise peptide correlations.
     Only outputs unique (non-symmetrical) correlations.
 
@@ -34,13 +36,15 @@ def pairwise_peptide_correlations_(
     - result (pandas.DataFrame): A DataFrame containing the pairwise peptide
         correlations. Columns: 'pepA', 'pepB', 'PCC' (Pearson correlation coefficient).
         Only outputs unique (non-symmetrical) correlations (AB, not AB, B-A, AA, BB).
-    '''
+    """
 
     # TODO: modify df input to be obs x vars. Here we have redundant steps with
     # AnnDataTrces pairwise_peptide_correlations()
     df = df[[sample_column, peptide_column, value_column]]
 
-    pivot_df = df.pivot_table(index=sample_column, columns=peptide_column, values=value_column)
+    pivot_df = df.pivot_table(
+        index=sample_column, columns=peptide_column, values=value_column
+    )
     columns = pivot_df.columns.tolist()
 
     corr_dict = {}
@@ -49,13 +53,17 @@ def pairwise_peptide_correlations_(
 
         pivot_col_a = pivot_df.loc[:, col_a]
         pivot_col_b = pivot_df.loc[:, col_b]
-        corr_dict[col_a + '_' + col_b] = stats.pearsonr(pivot_col_a, pivot_col_b)
+        corr_dict[col_a + "_" + col_b] = stats.pearsonr(
+            pivot_col_a, pivot_col_b
+        )
 
-    corr_df = pd.DataFrame.from_dict(corr_dict, orient='index')
-    corr_df.columns = ['PCC', 'p-value']
-    corr_df['peptide_pair'] = corr_df.index
-    corr_df[['pepA', 'pepB']] = corr_df['peptide_pair'].str.split('_', expand=True)
-    corr_df = corr_df[["pepA","pepB","PCC"]]
+    corr_df = pd.DataFrame.from_dict(corr_dict, orient="index")
+    corr_df.columns = ["PCC", "p-value"]
+    corr_df["peptide_pair"] = corr_df.index
+    corr_df[["pepA", "pepB"]] = corr_df["peptide_pair"].str.split(
+        "_", expand=True
+    )
+    corr_df = corr_df[["pepA", "pepB", "PCC"]]
     corr_df = corr_df.reset_index(drop=True)
 
     return corr_df
@@ -63,19 +71,19 @@ def pairwise_peptide_correlations_(
 
 def pairwise_peptide_correlations(
     adata,
-    protein_id='protein_id',
+    protein_id="protein_id",
     inplace=True,
     copy=False,
-    batch_key: str | None = None,   # per-batch if provided → always pooled
-    min_contrib_batches: int = 1,   # pooling threshold
-    min_wsum: float = 0.0,          # pooling threshold on sum(n_b-3)
-    ):
+    batch_key: str | None = None,  # per-batch if provided → always pooled
+    min_contrib_batches: int = 1,  # pooling threshold
+    min_wsum: float = 0.0,  # pooling threshold on sum(n_b-3)
+):
 
     if inplace and copy:
-        raise ValueError('Arguments raise and copy are mutually exclusive')
+        raise ValueError("Arguments raise and copy are mutually exclusive")
 
     if protein_id not in adata.var.columns:
-        raise ValueError(f'protein_id: {protein_id} not in .var.columns')
+        raise ValueError(f"protein_id: {protein_id} not in .var.columns")
 
     STORE_KEY = "pairwise_peptide_correlations"
     PER_BATCH_STORE_KEY = "pairwise_peptide_correlations_by_batch"
@@ -93,34 +101,38 @@ def _finalize(out, per_batch=None):
                 adata.uns[PER_BATCH_STORE_KEY] = per_batch
             return
         return out
-    
+
     def compute_corrs(df):
         corrs = pairwise_peptide_correlations_(
             df,
-            sample_column='obs_id',
-            peptide_column='var_id',
-            value_column='intensity')
+            sample_column="obs_id",
+            peptide_column="var_id",
+            value_column="intensity",
+        )
 
         return corrs
 
-    anns = adata.var[['protein_id']].reset_index()
+    anns = adata.var[["protein_id"]].reset_index()
     traces_df = adata.to_df().T.reset_index()
-    traces_df = traces_df.merge(anns, on='index')
-    traces_df = traces_df.rename(columns={'index': 'var_id'})
+    traces_df = traces_df.merge(anns, on="index")
+    traces_df = traces_df.rename(columns={"index": "var_id"})
 
     # TODO: remove unnecessary step of melting which gets unmelted
     #   in protein-level function
 
     traces_df = pd.melt(
         traces_df,
-        id_vars=['protein_id', 'var_id'],
-        var_name='obs_id',
-        value_name='intensity')
+        id_vars=["protein_id", "var_id"],
+        var_name="obs_id",
+        value_name="intensity",
+    )
 
     if batch_key is None:
-        corrs = traces_df.groupby('protein_id', observed=True).apply(compute_corrs, include_groups=False)
+        corrs = traces_df.groupby("protein_id", observed=True).apply(
+            compute_corrs, include_groups=False
+        )
         corrs = corrs.droplevel(1, axis=0)
-        corrs = corrs.sort_values(['pepA', 'pepB']).sort_index()
+        corrs = corrs.sort_values(["pepA", "pepB"]).sort_index()
         return _finalize(corrs)
 
     if batch_key not in adata.obs.columns:
@@ -129,43 +141,49 @@ def compute_corrs(df):
     batches = (
         adata.obs[[batch_key]]
         .reset_index()
-        .rename(columns={'index': 'obs_id', batch_key: 'batch_id'})
+        .rename(columns={"index": "obs_id", batch_key: "batch_id"})
     )
-    long = traces_df.merge(batches, on='obs_id', how='left')
+    long = traces_df.merge(batches, on="obs_id", how="left")
 
     batch_sizes = adata.obs[batch_key].value_counts().to_dict()
     batch_weights = {b: max(n - 3.0, 0.0) for b, n in batch_sizes.items()}
 
-    per_batch = (
-        long
-        .groupby([protein_id, 'batch_id'], observed=True)
-        .apply(compute_corrs, include_groups=False)
+    per_batch = long.groupby([protein_id, "batch_id"], observed=True).apply(
+        compute_corrs, include_groups=False
     )
 
     if per_batch.empty:
-        per_batch_df = pd.DataFrame(columns=['pepA', 'pepB', 'PCC'])
-        per_batch_df.index = pd.MultiIndex.from_tuples([], names=[protein_id, 'batch_id'])
+        per_batch_df = pd.DataFrame(columns=["pepA", "pepB", "PCC"])
+        per_batch_df.index = pd.MultiIndex.from_tuples(
+            [], names=[protein_id, "batch_id"]
+        )
     else:
         per_batch_df = (
-            per_batch
-            .reset_index(level=2, drop=True)
-            .sort_values(['pepA', 'pepB'])
+            per_batch.reset_index(level=2, drop=True)
+            .sort_values(["pepA", "pepB"])
             .sort_index()
         )
 
     # Fisher pooling across batches
     rows = []
-    for prot, gprot in per_batch_df.reset_index().groupby(protein_id, observed=True, sort=False):
-        for (pa, pb), gp in gprot.groupby(['pepA', 'pepB'], observed=True, sort=False):
-            r = gp['PCC'].to_numpy(dtype=float)
-            bids = gp['batch_id'].to_numpy()
+    for prot, gprot in per_batch_df.reset_index().groupby(
+        protein_id, observed=True, sort=False
+    ):
+        for (pa, pb), gp in gprot.groupby(
+            ["pepA", "pepB"], observed=True, sort=False
+        ):
+            r = gp["PCC"].to_numpy(dtype=float)
+            bids = gp["batch_id"].to_numpy()
             r = np.clip(r, -0.999999, 0.999999)
             z = np.arctanh(r)
-            w = np.array([batch_weights.get(b, 0.0) for b in bids], dtype=float)
+            w = np.array(
+                [batch_weights.get(b, 0.0) for b in bids], dtype=float
+            )
             mask = w > 0
             if not np.any(mask):
                 continue
-            w = w[mask]; z = z[mask]
+            w = w[mask]
+            z = z[mask]
             wsum = float(w.sum())
             if (mask.sum() >= min_contrib_batches) and (wsum >= min_wsum):
                 # Fixed-effects mean (zbar_fe) and weighted between-batch variance (var_z_between)
@@ -180,13 +198,18 @@ def compute_corrs(df):
 
     if rows:
         pooled_df = (
-            pd.DataFrame(rows, columns=[protein_id, 'pepA', 'pepB', 'PCC', 'var_z_between'])
+            pd.DataFrame(
+                rows,
+                columns=[protein_id, "pepA", "pepB", "PCC", "var_z_between"],
+            )
             .set_index(protein_id)
-            .sort_values(['pepA', 'pepB'])
+            .sort_values(["pepA", "pepB"])
             .sort_index()
         )
     else:
-        pooled_df = pd.DataFrame(columns=['pepA', 'pepB', 'PCC', 'var_z_between'])
+        pooled_df = pd.DataFrame(
+            columns=["pepA", "pepB", "PCC", "var_z_between"]
+        )
         pooled_df.index.name = protein_id
 
     return _finalize(pooled_df, per_batch=per_batch_df)
@@ -194,9 +217,9 @@ def compute_corrs(df):
 
 def peptide_dendograms_by_correlation_(
     df,
-    method: str = 'agglomerative-hierarchical-clustering',
-    ):
-    '''
+    method: str = "agglomerative-hierarchical-clustering",
+):
+    """
     Perform peptide clustering grouped by protein annotation.
 
 
@@ -220,24 +243,26 @@ def peptide_dendograms_by_correlation_(
                      The two ids included for every step represent the index of the peptide in 'labels'.
             - heights: The height of each merging step in 'merge'.
                        The idx of the height corresponds to the index of the step in 'merge'.
-    '''
+    """
 
     assert all(df.index == df.columns)
 
-    model = AgglomerativeClustering(n_clusters=None,
-                                    metric='precomputed',
-                                    linkage='average',
-                                    distance_threshold=0,
-                                    compute_distances=True)
+    model = AgglomerativeClustering(
+        n_clusters=None,
+        metric="precomputed",
+        linkage="average",
+        distance_threshold=0,
+        compute_distances=True,
+    )
 
     model.fit(df)
 
     # pylint: disable=no-member
     dendogram = {
-        'type': 'sklearn_agglomerative_clustering',
-        'labels': model.feature_names_in_.tolist(),
-        'heights': model.distances_.tolist(),
-        'merge': model.children_.tolist()
+        "type": "sklearn_agglomerative_clustering",
+        "labels": model.feature_names_in_.tolist(),
+        "heights": model.distances_.tolist(),
+        "merge": model.children_.tolist(),
     }
     # pylint: enable=no-member
 
@@ -246,59 +271,55 @@ def peptide_dendograms_by_correlation_(
 
 def peptide_dendograms_by_correlation(
     adata,
-    method='agglomerative-hierarchical-clustering',
+    method="agglomerative-hierarchical-clustering",
     inplace=True,
     copy=False,
-    ):
+):
 
     if inplace and copy:
-        raise ValueError('Arguments raise and copy are mutually exclusive')
+        raise ValueError("Arguments raise and copy are mutually exclusive")
 
+    if "pairwise_peptide_correlations" not in adata.uns:
+        raise ValueError(f"pairwise_peptide_correlations not in .uns")
 
-    if 'pairwise_peptide_correlations' not in adata.uns:
-        raise ValueError(f'pairwise_peptide_correlations not in .uns')
-
-
-    corrs = adata.uns['pairwise_peptide_correlations'].copy()
+    corrs = adata.uns["pairwise_peptide_correlations"].copy()
 
     dends = {}
 
-    for protein_id, df in corrs.groupby('protein_id', observed=True):
+    for protein_id, df in corrs.groupby("protein_id", observed=True):
 
-        corr_sym = reconstruct_corrs_df_symmetric_from_long_df(
-            df,
-            var_a_col='pepA',
-            var_b_col='pepB',
-            corr_col='PCC')
+        corr_sym = reconstruct_symmetric_matrix_from_long(
+            df, var_a_col="pepA", var_b_col="pepB", value_col="PCC"
+        )
 
         corr_dists = 1 - corr_sym
 
         dends[protein_id] = peptide_dendograms_by_correlation_(
-            corr_dists,
-            method= 'agglomerative-hierarchical-clustering')
+            corr_dists, method="agglomerative-hierarchical-clustering"
+        )
 
     if inplace:
-        adata.uns['dendograms'] = dends
+        adata.uns["dendograms"] = dends
 
     elif copy:
         adata_new = adata.copy()
-        adata_new.uns['dendograms'] = dends
+        adata_new.uns["dendograms"] = dends
         return adata_new
-    
+
     else:
         return dends
 
 
 def peptide_clusters_from_dendograms_(
-        dendogram,
-        n_clusters=2,
-        min_peptides_per_cluster=2,
-        noise=1e6,
-        ):
-    '''
-    Cut clusters from cluster_peptides into N clusters with more than 1 peptide. 
-    '''
-    n_peptides = len(dendogram['labels'])
+    dendogram,
+    n_clusters=2,
+    min_peptides_per_cluster=2,
+    noise=1e6,
+):
+    """
+    Cut clusters from cluster_peptides into N clusters with more than 1 peptide.
+    """
+    n_peptides = len(dendogram["labels"])
     n_real_clusters = 0
     k = n_clusters
     cluster_tree = BinaryClusterTree(constructor=dendogram)
@@ -319,7 +340,7 @@ def peptide_clusters_from_dendograms_(
 
     # Rename cluster_ids to systematic format
     max_cluster = clusters.max()
-    cats = clusters.astype('category').cat.categories
+    cats = clusters.astype("category").cat.categories
     n_clusters = len(cats)
 
     if max_cluster != n_clusters:
@@ -339,64 +360,58 @@ def peptide_clusters_from_dendograms(
     noise=NOISE,
     inplace=True,
     copy=False,
-    ):
+):
 
     if inplace and copy:
-        raise ValueError('Arguments raise and copy are mutually exclusive')
+        raise ValueError("Arguments raise and copy are mutually exclusive")
 
-    if 'dendograms' not in adata.uns:
-        raise ValueError(f'dendograms not in .uns')
+    if "dendograms" not in adata.uns:
+        raise ValueError(f"dendograms not in .uns")
 
     var = adata.var.copy()
-    var['cluster_id'] = np.nan
+    var["cluster_id"] = np.nan
 
     clusters_ann = {}
 
-    dends = adata.uns['dendograms']
+    dends = adata.uns["dendograms"]
     for prot, dend in dends.items():
         dend_upd = copym.deepcopy(dend)
-        dend_upd['type'] = 'sklearn_agglomerative_clustering'
+        dend_upd["type"] = "sklearn_agglomerative_clustering"
 
         clusters = peptide_clusters_from_dendograms_(
-            dend_upd,
-            n_clusters=2,
-            min_peptides_per_cluster=2,
-            noise=noise)
+            dend_upd, n_clusters=2, min_peptides_per_cluster=2, noise=noise
+        )
 
-        mask = (var['protein_id'] == prot) & (var.index.isin(clusters.index))
-        var.loc[mask, 'cluster_id'] = clusters.reindex(var.index[mask])
+        mask = (var["protein_id"] == prot) & (var.index.isin(clusters.index))
+        var.loc[mask, "cluster_id"] = clusters.reindex(var.index[mask])
 
         clusters_ann[prot] = clusters
 
-    assert not any((var['cluster_id'] == -1).tolist())
+    assert not any((var["cluster_id"] == -1).tolist())
 
-    var['proteoform_id'] = (
-            var['protein_id'].astype(str) +
-            '_' + 
-            var['cluster_id'].astype(int).astype(str)
-            )
+    var["proteoform_id"] = (
+        var["protein_id"].astype(str)
+        + "_"
+        + var["cluster_id"].astype(int).astype(str)
+    )
 
     if inplace:
-        adata.uns['clusters'] = clusters_ann
+        adata.uns["clusters"] = clusters_ann
         adata.var = var
 
     elif copy:
         adata_new = adata.copy()
-        adata_new.uns['clusters'] = clusters_ann
+        adata_new.uns["clusters"] = clusters_ann
         return adata_new
-    
+
     else:
         return clusters_ann
 
 
 def proteoform_scores_(
-    corrs,
-    clusters,
-    n_fractions,
-    summary_func=np.mean,
-    noise=NOISE
-    ):
-    '''
+    corrs, clusters, n_fractions, summary_func=np.mean, noise=NOISE
+):
+    """
     Calculates a score for proteoforms based on the difference of within
     cluster distances and between cluster distances.
 
@@ -410,7 +425,7 @@ def proteoform_scores_(
         n_fractions (int): Number of samples.
         summary_func (Callable): Summary function to apply to intra- and inter-
             cluster correlation coefficients.
-    '''
+    """
 
     def replace_upper_triangle(df, replacement, k=0):
         arr = df.to_numpy().astype(float)
@@ -422,15 +437,15 @@ def replace_upper_triangle(df, replacement, k=0):
         return new_df
 
     if isinstance(clusters, pd.DataFrame):
-        clusters = clusters['cluster']
+        clusters = clusters["cluster"]
 
     if np.issubdtype(clusters.dtype, np.floating):
         clusters = clusters.astype(int)
 
     assert any(corrs.index == corrs.columns)
-    assert all([i in clusters.index for i in corrs.index]), \
-        f'clusters.index = {clusters.index}' \
-        f'\ncorrs_index = {corrs.index}'
+    assert all([i in clusters.index for i in corrs.index]), (
+        f"clusters.index = {clusters.index}" f"\ncorrs_index = {corrs.index}"
+    )
 
     if (clusters == noise).all().all():
         return np.array([0, np.nan, np.nan, np.nan])
@@ -440,7 +455,9 @@ def replace_upper_triangle(df, replacement, k=0):
 
     if len(cluster_ids) > 2:
 
-        raise ValueError('Functionality with n_clusters > 2 not implemented yet.')
+        raise ValueError(
+            "Functionality with n_clusters > 2 not implemented yet."
+        )
 
         mat = corrs.copy(deep=True)
         stat_v = []
@@ -452,12 +469,16 @@ def replace_upper_triangle(df, replacement, k=0):
             clust_ids_ord = clust1_ids + clust2_ids
             mat_inv = corrs.loc[clust_ids_ord, clust_ids_ord]
 
-            cross = mat_inv.loc[clust1_ids, clust2_ids] # QUESTION: why no diagonal removal as below?
+            cross = mat_inv.loc[
+                clust1_ids, clust2_ids
+            ]  # QUESTION: why no diagonal removal as below?
             values = cross.to_numpy().flatten()
             values = values[~np.isnan(values)]
             stat_across = np.apply_along_axis(summary_func, 0, cross)
 
-            rows, cols = np.triu_indices_from(mat_inv, k=0)  # k=1 excludes diagonal
+            rows, cols = np.triu_indices_from(
+                mat_inv, k=0
+            )  # k=1 excludes diagonal
             mat_inv.to_numpy()[rows, cols] = np.nan
 
             within_c1 = mat_inv.loc[clust1_ids, clust1_ids]
@@ -478,7 +499,9 @@ def replace_upper_triangle(df, replacement, k=0):
             z_stat_across = np.atanh(stat_across)
             z_diff_stat = z_stat_within - z_stat_across
 
-            dz = z_diff_stat / (np.sqrt((1 / (n_fractions-3)) + (1 / (n_fractions-3))))
+            dz = z_diff_stat / (
+                np.sqrt((1 / (n_fractions - 3)) + (1 / (n_fractions - 3)))
+            )
             pval = 2 * (1 - norm.cdf(np.abs(dz)))
 
             stat_v.append([diff_stat, z_diff_stat, dz, pval])
@@ -523,7 +546,9 @@ def replace_upper_triangle(df, replacement, k=0):
 
         # T-test: intra-cluster peptide correlations are significantly different
         #   from cross-cluster peptide correlations
-        dz = z_diff_stat / np.sqrt((1 / (n_fractions-3)) + (1 / (n_fractions-3)))
+        dz = z_diff_stat / np.sqrt(
+            (1 / (n_fractions - 3)) + (1 / (n_fractions - 3))
+        )
         pval = 2 * (1 - norm.cdf(np.abs(dz)))
 
         return np.array([diff_stat, z_diff_stat, dz, pval])
@@ -537,27 +562,26 @@ def proteoform_scores(
     noise=NOISE,
     inplace=True,
     copy=False,
-    ):
+):
 
     if inplace and copy:
-        raise ValueError('Arguments raise and copy are mutually exclusive')
+        raise ValueError("Arguments raise and copy are mutually exclusive")
 
+    if "pairwise_peptide_correlations" not in adata.uns:
+        raise ValueError(f"pairwise_peptide_correlations not in .uns")
 
-    if 'pairwise_peptide_correlations' not in adata.uns:
-        raise ValueError(f'pairwise_peptide_correlations not in .uns')
-
-    if 'dendograms' not in adata.uns:
-        raise ValueError(f'dendograms not in .uns')
+    if "dendograms" not in adata.uns:
+        raise ValueError(f"dendograms not in .uns")
 
     columns = [
-        'protein_id',
-        'proteoform_score',
-        'proteoform_score_z',
-        'proteoform_score_dz',
-        'proteoform_score_pval',
-        ]
-
-    corrs = adata.uns['pairwise_peptide_correlations'].copy().reset_index()
+        "protein_id",
+        "proteoform_score",
+        "proteoform_score_z",
+        "proteoform_score_dz",
+        "proteoform_score_pval",
+    ]
+
+    corrs = adata.uns["pairwise_peptide_correlations"].copy().reset_index()
     # pylint: disable=access-member-before-definition
     var = adata.var
     # pylint: enable=access-member-before-definition
@@ -565,24 +589,22 @@ def proteoform_scores(
 
     proteoform_scores_list = []
 
-    for prot, corrs_prot in corrs.groupby('protein_id', observed=True):
+    for prot, corrs_prot in corrs.groupby("protein_id", observed=True):
 
-        corrs_mat = reconstruct_corrs_df_symmetric_from_long_df(
-            corrs_prot,
-            var_a_col='pepA',
-            var_b_col='pepB',
-            corr_col='PCC')
+        corrs_mat = reconstruct_symmetric_matrix_from_long(
+            corrs_prot, var_a_col="pepA", var_b_col="pepB", value_col="PCC"
+        )
 
-        clusters = var.loc[var['protein_id'] == prot, 'cluster_id']
+        clusters = var.loc[var["protein_id"] == prot, "cluster_id"]
 
         scores = proteoform_scores_(
-            corrs_mat,
-            clusters,
-            n_fractions,
-            summary_func=np.mean)
+            corrs_mat, clusters, n_fractions, summary_func=np.mean
+        )
 
-        scores_entry = {column:value for column, value in zip(columns[1:5], scores)}
-        scores_entry['protein_id'] = prot
+        scores_entry = {
+            column: value for column, value in zip(columns[1:5], scores)
+        }
+        scores_entry["protein_id"] = prot
         scores_entry = pd.DataFrame([scores_entry])
         proteoform_scores_list.append(scores_entry)
 
@@ -591,53 +613,54 @@ def proteoform_scores(
 
     # Perform multiple-testing correction
 
-    mask_nonan = proteoform_scores['proteoform_score_pval'].notna()
-    pvals = proteoform_scores.loc[mask_nonan, 'proteoform_score_pval']
+    mask_nonan = proteoform_scores["proteoform_score_pval"].notna()
+    pvals = proteoform_scores.loc[mask_nonan, "proteoform_score_pval"]
 
     bh_alpha = min_pval_adj if min_pval_adj is not None else 0.05
     _, corrected_pvals, _, _ = multipletests(
         pvals,
         alpha=bh_alpha,
-        method='fdr_bh',
+        method="fdr_bh",
     )
 
-    proteoform_scores['proteoform_score_pval_adj'] = np.nan
-    proteoform_scores['is_proteoform'] = np.nan
+    proteoform_scores["proteoform_score_pval_adj"] = np.nan
+    proteoform_scores["is_proteoform"] = np.nan
 
-    proteoform_scores.loc[
-        pvals.index, 'proteoform_score_pval_adj'
-    ] = corrected_pvals
+    proteoform_scores.loc[pvals.index, "proteoform_score_pval_adj"] = (
+        corrected_pvals
+    )
 
     if min_pval_adj is not None or min_score is not None:
         is_pf = pd.Series(True, index=pvals.index)
         if min_pval_adj is not None:
             is_pf &= corrected_pvals <= min_pval_adj
         if min_score is not None:
-            scores = proteoform_scores.loc[pvals.index, 'proteoform_score']
+            scores = proteoform_scores.loc[pvals.index, "proteoform_score"]
             is_pf &= scores >= min_score
-        proteoform_scores.loc[
-            pvals.index, 'is_proteoform'
-        ] = is_pf.astype(int).values
+        proteoform_scores.loc[pvals.index, "is_proteoform"] = is_pf.astype(
+            int
+        ).values
 
     # --- drop existing score columns before merge (safe for re-runs) ---
     score_cols = [
-        'proteoform_score',
-        'proteoform_score_z',
-        'proteoform_score_dz',
-        'proteoform_score_pval',
-        'proteoform_score_pval_adj',
-        'is_proteoform',
+        "proteoform_score",
+        "proteoform_score_z",
+        "proteoform_score_dz",
+        "proteoform_score_pval",
+        "proteoform_score_pval_adj",
+        "is_proteoform",
     ]
     var = var.drop(columns=[c for c in score_cols if c in var.columns])
     # Add all new scores to .var
     var_upd = pd.merge(
         var,
         proteoform_scores,
-        on='protein_id',
-        how='left',
-        validate='many_to_one')
+        on="protein_id",
+        how="left",
+        validate="many_to_one",
+    )
 
-    var_upd = var_upd.set_index('peptide_id', drop=False)
+    var_upd = var_upd.set_index("peptide_id", drop=False)
     var_upd.index.name = None
 
     assert (var.index == var_upd.index).all()
@@ -649,6 +672,6 @@ def proteoform_scores(
         adata_new = adata.copy()
         adata_new.var = var_upd
         return adata_new
-    
+
     else:
         return proteoform_scores

From 6f3424d4a8a45e7faff13e1fc9aa9e1fb3f2af54 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Levente=20Temesv=C3=A1ri-Nagy?=
 <147416790+leventetn@users.noreply.github.com>
Date: Sat, 26 Sep 2026 10:57:22 +0200
Subject: [PATCH 03/10] tests/pairwise_peptide_correlations_heatmap
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

New tests/pl/test_pairwise_peptide_correlations_heatmap.py — 16 tests
(signature lock, return semantics, symmetry, linkage pass-through,
cluster=False, str-vs-list strips, peptide_id labels, validation,
save/non-mutation).
---
 ...t_pairwise_peptide_correlations_heatmap.py | 305 ++++++++++++++++++
 1 file changed, 305 insertions(+)
 create mode 100644 tests/pl/test_pairwise_peptide_correlations_heatmap.py

diff --git a/tests/pl/test_pairwise_peptide_correlations_heatmap.py b/tests/pl/test_pairwise_peptide_correlations_heatmap.py
new file mode 100644
index 0000000..c6737cf
--- /dev/null
+++ b/tests/pl/test_pairwise_peptide_correlations_heatmap.py
@@ -0,0 +1,305 @@
+"""Tests for pr.pl.pairwise_peptide_correlations_heatmap.
+
+Covers: signature contract, return-value semantics, symmetric
+clustering and linkage pass-through, cluster toggle, annotation
+strips, ID-column labelling, argument validation, and lifecycle
+(save / non-mutation).
+"""
+
+import inspect
+
+import numpy as np
+import pandas as pd
+import anndata as ad
+import matplotlib.pyplot as plt
+from matplotlib.axes import Axes
+from scipy.spatial.distance import squareform
+from scipy.cluster.hierarchy import linkage as scipy_linkage
+import pytest
+
+import proteopy as pr
+
+plt.switch_backend("Agg")
+
+
+@pytest.fixture(autouse=True)
+def close_test_figures():
+    before = set(plt.get_fignums())
+    yield
+    for num in set(plt.get_fignums()) - before:
+        plt.close(num)
+
+
+def _function():
+    return getattr(pr.pl, "pairwise_peptide_correlations_heatmap")
+
+
+def _copf_adata(seed=0):
+    """Peptide-level AnnData with the COPF chain already run.
+
+    P1: 5 peptides forming two proteoforms; P2: 3 peptides.
+    """
+    rng = np.random.default_rng(seed)
+    prot_map = {f"pep{i}": "P1" for i in range(5)}
+    prot_map.update({f"pep{i}": "P2" for i in range(5, 8)})
+    peps = list(prot_map)
+
+    n_obs = 12
+    pat_a = rng.normal(size=n_obs)
+    pat_b = rng.normal(size=n_obs)
+    pat_c = rng.normal(size=n_obs)
+    patterns = {
+        "pep0": pat_a,
+        "pep1": pat_a,
+        "pep2": pat_a,
+        "pep3": pat_b,
+        "pep4": pat_b,
+        "pep5": pat_c,
+        "pep6": pat_c,
+        "pep7": pat_c,
+    }
+    cols = {
+        p: np.abs(patterns[p] + rng.normal(scale=0.1, size=n_obs)) + 1.0
+        for p in peps
+    }
+    X = np.vstack([cols[p] for p in peps]).T
+
+    obs = pd.DataFrame(
+        {"sample_id": [f"s{i}" for i in range(n_obs)]},
+        index=[f"s{i}" for i in range(n_obs)],
+    )
+    var = pd.DataFrame(
+        {"peptide_id": peps, "protein_id": [prot_map[p] for p in peps]},
+        index=peps,
+    )
+    adata = ad.AnnData(X=X, obs=obs, var=var)
+
+    pr.tl.pairwise_peptide_correlations(adata)
+    pr.tl.peptide_dendograms_by_correlation(adata)
+    pr.tl.peptide_clusters_from_dendograms(
+        adata, n_clusters=2, min_peptides_per_cluster=2
+    )
+    return adata
+
+
+def _protein_peptides(adata, protein):
+    var = adata.var
+    return sorted(var.index[var["protein_id"] == protein].tolist())
+
+
+def _spy_clustermap(monkeypatch):
+    """Capture kwargs passed to sns.clustermap; call through."""
+    import proteopy.pl.copf as copf_mod
+
+    captured = {}
+    original = copf_mod.sns.clustermap
+
+    def spy(data, **kwargs):
+        captured["data"] = data
+        captured["kwargs"] = kwargs
+        return original(data, **kwargs)
+
+    monkeypatch.setattr(copf_mod.sns, "clustermap", spy)
+    return captured
+
+
+# -- Signature contract --------------------------------------------------
+
+
+def test_signature_locked():
+    sig = inspect.signature(_function())
+    params = list(sig.parameters.values())
+    names = [p.name for p in params]
+    assert names == [
+        "adata",
+        "protein",
+        "corr_key",
+        "margin_color",
+        "method",
+        "cluster",
+        "linkage",
+        "color_scheme",
+        "cmap",
+        "xticklabels",
+        "yticklabels",
+        "figsize",
+        "show",
+        "ax",
+        "print_stats",
+        "save",
+    ]
+    # protein is required (no default)
+    assert sig.parameters["protein"].default is inspect.Parameter.empty
+    defaults = sig.parameters
+    assert defaults["corr_key"].default == "pairwise_peptide_correlations"
+    assert defaults["margin_color"].default == "cluster_id"
+    assert defaults["method"].default == "average"
+    assert defaults["cluster"].default is True
+    assert defaults["linkage"].default is None
+    assert defaults["cmap"].default == "coolwarm"
+    assert defaults["show"].default is True
+    assert defaults["ax"].default is False
+    assert defaults["print_stats"].default is False
+    assert defaults["save"].default is None
+
+
+# -- Return-value semantics ----------------------------------------------
+
+
+def test_returns_axes_when_ax_true():
+    adata = _copf_adata()
+    out = _function()(adata, protein="P1", show=False, ax=True)
+    assert isinstance(out, Axes)
+
+
+def test_returns_none_when_ax_false():
+    adata = _copf_adata()
+    out = _function()(adata, protein="P1", show=False, ax=False)
+    assert out is None
+
+
+# -- Clustering: symmetry, linkage pass-through, toggle -------------------
+
+
+def test_default_clustering_is_symmetric():
+    adata = _copf_adata()
+    axm = _function()(adata, protein="P1", show=False, ax=True)
+    x = [t.get_text() for t in axm.get_xticklabels()]
+    y = [t.get_text() for t in axm.get_yticklabels()]
+    assert x == y
+
+
+def test_cluster_false_keeps_input_order():
+    adata = _copf_adata()
+    axm = _function()(adata, protein="P1", cluster=False, show=False, ax=True)
+    y = [t.get_text() for t in axm.get_yticklabels()]
+    assert y == _protein_peptides(adata, "P1")
+
+
+def test_supplied_linkage_shadows_both_axes(monkeypatch):
+    adata = _copf_adata()
+    captured = _spy_clustermap(monkeypatch)
+
+    # Build a linkage over P1's peptides.
+    corr = adata.uns["pairwise_peptide_correlations"].loc[["P1"]]
+    from proteopy.utils._matrix_wrangling import (
+        reconstruct_symmetric_matrix_from_long,
+    )
+
+    mat = reconstruct_symmetric_matrix_from_long(
+        corr, var_a_col="pepA", var_b_col="pepB", value_col="PCC"
+    ).to_numpy(dtype=float)
+    dist = 1.0 - mat
+    np.fill_diagonal(dist, 0.0)
+    Z = scipy_linkage(
+        squareform(np.clip(dist, 0, 2), checks=False), method="complete"
+    )
+
+    _function()(adata, protein="P1", linkage=Z, show=False, ax=True)
+    kwargs = captured["kwargs"]
+    assert kwargs["row_linkage"] is Z
+    assert kwargs["col_linkage"] is Z
+
+
+# -- Annotation strips ----------------------------------------------------
+
+
+def test_margin_color_string_single_strip(monkeypatch):
+    adata = _copf_adata()
+    captured = _spy_clustermap(monkeypatch)
+    _function()(
+        adata, protein="P1", margin_color="cluster_id", show=False, ax=False
+    )
+    row_colors = captured["kwargs"]["row_colors"]
+    assert list(row_colors.columns) == ["cluster_id"]
+
+
+def test_margin_color_list_multiple_strips(monkeypatch):
+    adata = _copf_adata()
+    captured = _spy_clustermap(monkeypatch)
+    _function()(
+        adata,
+        protein="P1",
+        margin_color=["cluster_id", "proteoform_id"],
+        show=False,
+        ax=False,
+    )
+    row_colors = captured["kwargs"]["row_colors"]
+    assert list(row_colors.columns) == ["cluster_id", "proteoform_id"]
+    # Same strips mirrored to columns (symmetric).
+    assert captured["kwargs"]["col_colors"] is row_colors
+
+
+def test_annotation_distinct_colors_match_categories(monkeypatch):
+    adata = _copf_adata()
+    captured = _spy_clustermap(monkeypatch)
+    _function()(
+        adata, protein="P1", margin_color="cluster_id", show=False, ax=False
+    )
+    strip = captured["kwargs"]["row_colors"]["cluster_id"]
+    n_colors = len({tuple(np.atleast_1d(c)) for c in strip})
+    n_cats = adata.var.loc[
+        _protein_peptides(adata, "P1"), "cluster_id"
+    ].nunique()
+    assert n_colors == n_cats
+
+
+# -- Labelling from ID columns -------------------------------------------
+
+
+def test_tick_labels_from_peptide_id():
+    adata = _copf_adata()
+    axm = _function()(adata, protein="P1", show=False, ax=True)
+    y = [t.get_text() for t in axm.get_yticklabels()]
+    assert sorted(y) == _protein_peptides(adata, "P1")
+
+
+# -- Validation -----------------------------------------------------------
+
+
+def test_unknown_protein_raises():
+    adata = _copf_adata()
+    with pytest.raises(ValueError, match="not found"):
+        _function()(adata, protein="NOPE", show=False)
+
+
+def test_missing_corr_key_raises():
+    adata = _copf_adata()
+    del adata.uns["pairwise_peptide_correlations"]
+    with pytest.raises(ValueError, match="pairwise_peptide_correlations"):
+        _function()(adata, protein="P1", show=False)
+
+
+def test_missing_margin_color_column_raises():
+    adata = _copf_adata()
+    with pytest.raises(KeyError, match="not found in adata.var"):
+        _function()(
+            adata, protein="P1", margin_color="does_not_exist", show=False
+        )
+
+
+def test_bad_margin_color_type_raises():
+    adata = _copf_adata()
+    with pytest.raises(TypeError, match="margin_color"):
+        _function()(adata, protein="P1", margin_color=123, show=False)
+
+
+# -- Lifecycle ------------------------------------------------------------
+
+
+def test_save_writes_png(tmp_path):
+    adata = _copf_adata()
+    out = tmp_path / "heatmap.png"
+    _function()(adata, protein="P1", show=False, save=out)
+    assert out.exists() and out.stat().st_size > 0
+
+
+def test_input_not_mutated():
+    adata = _copf_adata()
+    var_before = adata.var.copy()
+    uns_before = adata.uns["pairwise_peptide_correlations"].copy()
+    _function()(adata, protein="P1", show=False)
+    pd.testing.assert_frame_equal(adata.var, var_before)
+    pd.testing.assert_frame_equal(
+        adata.uns["pairwise_peptide_correlations"], uns_before
+    )

From f56302a111cd225c0e9a9b3c316e3976451f197b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Levente=20Temesv=C3=A1ri-Nagy?=
 <147416790+leventetn@users.noreply.github.com>
Date: Sat, 26 Sep 2026 12:01:24 +0200
Subject: [PATCH 04/10] fixed/test warnings for karayel_2020 and
 summarize_peptides_by_neighbourhood_union

the pandas test now sets observed=False, and the peptide test now
represents multiple protein mappings without duplicate variable names
---
 tests/datasets/test_karayel_2020.py           | 44 +++++++++----------
 ...mmarize_peptides_by_neighbourhood_union.py | 10 ++---
 2 files changed, 25 insertions(+), 29 deletions(-)

diff --git a/tests/datasets/test_karayel_2020.py b/tests/datasets/test_karayel_2020.py
index b7ac407..54c6355 100644
--- a/tests/datasets/test_karayel_2020.py
+++ b/tests/datasets/test_karayel_2020.py
@@ -1,4 +1,5 @@
 """Tests for proteopy.datasets.karayel_2020."""
+
 import hashlib
 
 import anndata as ad
@@ -13,28 +14,30 @@
 _EXPECTED_SHAPE = (20, 7758)
 
 _EXPECTED_X_HASH = (
-    "3f40838356b56b8f230bdb02bde8d16d"
-    "c574fcc41d106bdce02ebe666f4e02db"
+    "3f40838356b56b8f230bdb02bde8d16d" "c574fcc41d106bdce02ebe666f4e02db"
 )
 _EXPECTED_OBS_NAMES_HASH = (
-    "fef7fd91a6e93d20b719f61c63098865"
-    "bbd3f886dabfa46786cde09e520c0abe"
+    "fef7fd91a6e93d20b719f61c63098865" "bbd3f886dabfa46786cde09e520c0abe"
 )
 _EXPECTED_VAR_NAMES_HASH = (
-    "17d3bd09174bad3544738f30ed2867c4"
-    "bd431feccdf36515e5fae415110fc456"
+    "17d3bd09174bad3544738f30ed2867c4" "bd431feccdf36515e5fae415110fc456"
 )
 
 _EXPECTED_OBS_COLUMNS = ["sample_id", "cell_type", "replicate"]
 _EXPECTED_VAR_COLUMNS = ["protein_id", "gene_id"]
 _EXPECTED_CELL_TYPES = [
-    "LBaso", "Ortho", "Poly", "ProE&EBaso", "Progenitor",
+    "LBaso",
+    "Ortho",
+    "Poly",
+    "ProE&EBaso",
+    "Progenitor",
 ]
 _EXPECTED_REPLICATES = ["rep1", "rep2", "rep3", "rep4"]
 
 
 # -- Fixtures --------------------------------------------------------
 
+
 @pytest.fixture(scope="module")
 def adata():
     """Load karayel_2020 dataset once for all tests."""
@@ -43,6 +46,7 @@ def adata():
 
 # -- Helpers ---------------------------------------------------------
 
+
 def _sha256(data: bytes) -> str:
     return hashlib.sha256(data).hexdigest()
 
@@ -59,6 +63,7 @@ def _encode_index(index) -> bytes:
 
 # -- Content tests ---------------------------------------------------
 
+
 class TestKarayel2020:
     """Verify structure and content of the karayel_2020 dataset."""
 
@@ -75,32 +80,23 @@ def test_var_columns(self, adata):
         assert adata.var.columns.tolist() == _EXPECTED_VAR_COLUMNS
 
     def test_cell_types(self, adata):
-        assert (
-            sorted(adata.obs["cell_type"].unique())
-            == _EXPECTED_CELL_TYPES
-        )
+        assert sorted(adata.obs["cell_type"].unique()) == _EXPECTED_CELL_TYPES
 
     def test_replicates(self, adata):
-        assert (
-            sorted(adata.obs["replicate"].unique())
-            == _EXPECTED_REPLICATES
-        )
+        assert sorted(adata.obs["replicate"].unique()) == _EXPECTED_REPLICATES
 
     def test_four_replicates_per_cell_type(self, adata):
-        counts = adata.obs.groupby("cell_type").size()
+        counts = adata.obs.groupby(
+            "cell_type",
+            observed=False,
+        ).size()
         assert (counts == 4).all()
 
     def test_obs_names_match_sample_id(self, adata):
-        assert (
-            list(adata.obs_names)
-            == list(adata.obs["sample_id"])
-        )
+        assert list(adata.obs_names) == list(adata.obs["sample_id"])
 
     def test_var_names_match_protein_id(self, adata):
-        assert (
-            list(adata.var_names)
-            == list(adata.var["protein_id"])
-        )
+        assert list(adata.var_names) == list(adata.var["protein_id"])
 
     def test_x_dtype(self, adata):
         assert adata.X.dtype == np.float64
diff --git a/tests/pp/test_summarize_peptides_by_neighbourhood_union.py b/tests/pp/test_summarize_peptides_by_neighbourhood_union.py
index a591158..f16e647 100644
--- a/tests/pp/test_summarize_peptides_by_neighbourhood_union.py
+++ b/tests/pp/test_summarize_peptides_by_neighbourhood_union.py
@@ -1879,14 +1879,14 @@ def test_multi_mapped_peptide_raises(self):
         proteodata either."""
         var = pd.DataFrame(
             {
-                "peptide_id": ["ACDEF", "ACDEF"],
-                "protein_id": ["P1", "P2"],
+                "peptide_id": ["ACDEF"],
+                "protein_id": ["P1;P2"],
             },
-            index=["ACDEF", "ACDEF"],
+            index=["ACDEF"],
         )
         obs = pd.DataFrame({"sample_id": ["s1"]}, index=["s1"])
-        adata = AnnData(X=np.array([[1.0, 2.0]]), obs=obs, var=var)
-        with pytest.raises(ValueError):
+        adata = AnnData(X=np.array([[1.0]]), obs=obs, var=var)
+        with pytest.raises(ValueError, match="mapping to exactly one protein"):
             summarize(adata, _FASTA, inplace=False)
 
     def test_obs_is_preserved(self):

From 6c67611886213c5999938c90576b0098954610e6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Levente=20Temesv=C3=A1ri-Nagy?=
 <147416790+leventetn@users.noreply.github.com>
Date: Mon, 28 Sep 2026 11:58:59 +0200
Subject: [PATCH 05/10] fixed/pairwise_peptide_correlations_heatmap
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Silently wrong output:
A table with several correlations per peptide pair (e.g. …_by_batch)
plotted whichever batch came last. It now raises.
A precomputed linkage of the wrong size drew a wrong tree. It now raises
with the expected shape.

Crashes:
Missing values in an annotation column: drawn light gray, with an "NA"
legend entry.
Colour dicts: color_scheme dicts keyed by the column's real values (e.g.
{0.0: …}) now work.
NaN correlations (missing intensities, constant peptides): cells are
drawn hatched with a legend entry, and are treated as r = 0 for
clustering, with a warning. NaN no longer blocks the default
cluster=True. The helper gained allow_missing=False, which the heatmap
sets to True.
Peptides filtered out after the correlations were computed: left out
with a warning.
print_stats without a cluster_id column: it now prints stats per
annotation column, in legend order.
pdf/svg backends: the plot now works when saved as vector graphics.

Layout and readability:
One legend per annotation column, placed to the right of all labels. The
figure grows if needed, and plt.tight_layout() was removed because it
left a gap under the annotation strip.
Several annotation strips never share a colour.
The title sits on the dendrogram.
Peptides that have no correlations at all are shown as empty rows
instead of disappearing.
Tick labels default to "auto" so labels on large proteins don't overlap.

Defaults, validation and AGENTS.md compliance:
margin_color defaults to "proteoform_id".
With cluster=False, a categorical peptide_id order is respected.
Clear errors for an empty or repeated margin_color and for missing COPF
columns (with a hint to run pr.tl.peptide_clusters_from_dendograms).
An invalid save type is rejected before drawing.
A warning when there are fewer than 3 samples.
The docstring now has Raises, Warns and Notes sections (data scale, NaN
behaviour of the upstream correlation).
Tests: heatmap tests went from 16 to 47 and helper tests from 9 to 10.
Each new test that targets a fix fails on the earlier version.
---
 proteopy/pl/copf.py                           | 398 +++++++++++++----
 proteopy/utils/_matrix_wrangling.py           |   9 +-
 ...t_pairwise_peptide_correlations_heatmap.py | 419 +++++++++++++++++-
 tests/utils/test_matrix_wrangling.py          |  16 +
 4 files changed, 753 insertions(+), 89 deletions(-)

diff --git a/proteopy/pl/copf.py b/proteopy/pl/copf.py
index 0fa1394..12a5202 100644
--- a/proteopy/pl/copf.py
+++ b/proteopy/pl/copf.py
@@ -22,6 +22,9 @@
     reconstruct_symmetric_matrix_from_long,
 )
 
+NAN_HATCH = "////"
+NAN_HATCH_COLOR = "#9e9e9e"
+
 
 def proteoform_scores(
     adata: ad.AnnData,
@@ -394,14 +397,14 @@ def pairwise_peptide_correlations_heatmap(
     protein: str,
     *,
     corr_key: str = "pairwise_peptide_correlations",
-    margin_color: str | list[str] = "cluster_id",
+    margin_color: str | list[str] = "proteoform_id",
     method: str = "average",
     cluster: bool = True,
     linkage=None,
     color_scheme=None,
     cmap: str = "coolwarm",
-    xticklabels: bool = True,
-    yticklabels: bool = True,
+    xticklabels: bool | str = "auto",
+    yticklabels: bool | str = "auto",
     figsize: tuple[float, float] = (8.0, 8.0),
     show: bool = True,
     ax: bool = False,
@@ -434,24 +437,40 @@ def pairwise_peptide_correlations_heatmap(
     margin_color : str | list[str]
         Column(s) in ``adata.var`` used for the annotation strip(s). A
         string is treated as a single-item list; a list draws one strip
-        per entry, in the given order. Proteoform labels are obtained by
-        passing ``"proteoform_id"``.
+        per entry, in the given order. The default shows the COPF
+        proteoform labels (``<protein_id>_<cluster_id>``) written by
+        :func:`proteopy.tl.peptide_clusters_from_dendograms`. Colours
+        are assigned across all columns together, so strips never share
+        a colour, and each column gets its own legend, placed right of
+        all heatmap labels. Legend entries follow the column's category
+        order if it is categorical, otherwise the lexicographic order of
+        the ``str``-coerced values; store the column as an ordered
+        :class:`pandas.Categorical` to change it. Missing values are
+        drawn in light gray.
     method : str
         Linkage method handed to :func:`scipy.cluster.hierarchy.linkage`
         when the tree is recomputed from ``1 - correlation``.
     cluster : bool
-        Cluster and draw dendrograms on both axes. When ``False`` the
-        matrix is shown in its input order without dendrograms.
+        Cluster and draw dendrograms on both axes. When ``False``,
+        peptides follow the category order of ``adata.var["peptide_id"]``
+        if it is categorical, otherwise lexicographic order. ``NaN``
+        correlations (e.g. from constant peptides) are drawn hatched and
+        treated as ``r = 0`` when building the tree.
     linkage : numpy.ndarray | None
-        Precomputed SciPy linkage matrix. When provided it is used for
-        both axes, shadowing the default ``1 - correlation`` computation.
+        Precomputed SciPy linkage matrix of shape ``(n_peptides - 1, 4)``.
+        When provided it is used for both axes, shadowing the default
+        ``1 - correlation`` computation. Leaf indices refer to the
+        peptides in the default order described under ``cluster``.
     color_scheme : Any
         Color palette specification understood by
         :func:`proteopy.utils.matplotlib._resolve_color_scheme`.
     cmap : str
         Continuous colormap for the heatmap body.
-    xticklabels, yticklabels : bool
-        Whether to show peptide tick labels on each axis.
+    xticklabels, yticklabels : bool | str
+        Peptide tick labels on each axis. ``"auto"`` labels as many
+        peptides as fit without overlapping; ``True`` labels every
+        peptide (enlarge ``figsize`` for large proteins); ``False``
+        hides the labels.
     figsize : tuple[float, float]
         Matplotlib figure size in inches.
     show : bool
@@ -472,10 +491,39 @@ def pairwise_peptide_correlations_heatmap(
     ------
     ValueError
         If ``protein`` is unknown, ``corr_key`` is missing, the protein
-        has no pairwise correlations, or clustering is requested on a
-        matrix containing missing values.
+        has no pairwise correlations, the table holds several
+        correlations per peptide pair (e.g. per-batch results),
+        ``margin_color`` is empty or repeats a column, ``linkage`` does
+        not match the number of peptides, or clustering is requested on
+        a matrix containing missing values.
     KeyError
         If a ``margin_color`` column is not present in ``adata.var``.
+    TypeError
+        If ``method``, ``cluster``, ``margin_color`` or ``save`` has the
+        wrong type.
+
+    Warns
+    -----
+    UserWarning
+        If peptides in ``adata.uns[corr_key]`` are no longer in
+        ``adata.var`` (filtered after the correlations were computed;
+        they are left out), if peptides of ``protein`` have no
+        correlations (shown as empty rows and columns), if NaN
+        correlations are treated as ``r = 0`` for clustering, or if
+        ``adata`` has fewer than 3 samples.
+
+    Notes
+    -----
+    The heatmap only displays correlations; it does not transform
+    intensities. They are computed by
+    :func:`proteopy.tl.pairwise_peptide_correlations` on ``adata.X`` as
+    stored, so their scale is the scale of ``.X``: Pearson correlations
+    of raw and log-transformed intensities differ, and a single extreme
+    sample can dominate raw-scale correlations. Log-transform (and
+    aggregate precursors to peptides) before computing correlations.
+    ``scipy.stats.pearsonr`` is not NaN-aware, so a pair is ``NaN`` as
+    soon as either peptide is missing in any sample, and there is no
+    minimum-overlap cutoff.
 
     Examples
     --------
@@ -517,6 +565,18 @@ def pairwise_peptide_correlations_heatmap(
         margin_cols = list(margin_color)
     else:
         raise TypeError("`margin_color` must be a string or list of strings.")
+    if not margin_cols:
+        raise ValueError("`margin_color` must name at least one column.")
+    duplicated_cols = sorted(
+        {c for c in margin_cols if margin_cols.count(c) > 1}
+    )
+    if duplicated_cols:
+        raise ValueError(
+            "Duplicate columns in `margin_color`: "
+            f"{', '.join(duplicated_cols)}."
+        )
+    if save is not None and not isinstance(save, (str, Path)):
+        raise TypeError("`save` must be a path-like object or None.")
 
     if protein not in set(adata.var["protein_id"]):
         raise ValueError(
@@ -531,9 +591,24 @@ def pairwise_peptide_correlations_heatmap(
 
     missing_cols = [c for c in margin_cols if c not in adata.var.columns]
     if missing_cols:
+        hint = ""
+        if {"cluster_id", "proteoform_id"} & set(missing_cols):
+            hint = (
+                " These columns are written by "
+                "pr.tl.peptide_clusters_from_dendograms(); run it first or "
+                "pass another .var column."
+            )
         raise KeyError(
             "margin_color column(s) not found in adata.var: "
-            f"{', '.join(missing_cols)}."
+            f"{', '.join(missing_cols)}.{hint}"
+        )
+
+    if adata.n_obs < 3:
+        warnings.warn(
+            f"adata has {adata.n_obs} sample(s); Pearson correlations "
+            "from fewer than 3 samples are always +/-1 or undefined.",
+            UserWarning,
+            stacklevel=2,
         )
 
     corrs = adata.uns[corr_key]
@@ -546,79 +621,147 @@ def pairwise_peptide_correlations_heatmap(
 
     # -- Reconstruct the symmetric peptide x peptide correlation matrix
     corr_long = corrs.loc[[protein]]
+    pairs = pd.DataFrame(
+        np.sort(corr_long[["pepA", "pepB"]].to_numpy(dtype=str), axis=1)
+    )
+    if pairs.duplicated().any():
+        raise ValueError(
+            f"adata.uns['{corr_key}'] holds several correlations per "
+            f"peptide pair for protein '{protein}' (e.g. per-batch "
+            "results); pass a table with one correlation per pair, such "
+            "as the pooled 'pairwise_peptide_correlations'."
+        )
     corr_df = reconstruct_symmetric_matrix_from_long(
         corr_long,
         var_a_col="pepA",
         var_b_col="pepB",
         value_col="PCC",
+        allow_missing=True,
     )
 
+    # Peptides filtered out after pr.tl.pairwise_peptide_correlations
+    # still sit in .uns; their correlations with the rest stay valid.
+    stale = [p for p in corr_df.index if p not in adata.var.index]
+    if stale:
+        warnings.warn(
+            f"{len(stale)} peptide(s) of protein '{protein}' in "
+            f"adata.uns['{corr_key}'] are no longer in adata.var and are "
+            "not shown.",
+            UserWarning,
+            stacklevel=2,
+        )
+        kept = [p for p in corr_df.index if p not in set(stale)]
+        if len(kept) < 2:
+            raise ValueError(
+                f"Fewer than two peptides of protein '{protein}' remain "
+                "in adata.var; recompute "
+                "pr.tl.pairwise_peptide_correlations()."
+            )
+        corr_df = corr_df.loc[kept, kept]
+
+    protein_peptides = adata.var.index[adata.var["protein_id"] == protein]
+    without_corr = [p for p in protein_peptides if p not in corr_df.index]
+    if without_corr:
+        warnings.warn(
+            f"{len(without_corr)} peptide(s) of protein '{protein}' have "
+            f"no correlations in adata.uns['{corr_key}'] (e.g. all "
+            "intensities missing); they are shown as empty rows and "
+            "columns.",
+            UserWarning,
+            stacklevel=2,
+        )
+        full = list(corr_df.index) + without_corr
+        corr_df = corr_df.reindex(index=full, columns=full)
+
     peptides = list(corr_df.index)
 
     # -- Labels come from the peptide_id column, never the index
     pep_ids = adata.var.loc[peptides, "peptide_id"].astype(str).tolist()
 
-    # -- Build annotation strip(s) and their legend handles
-    annot_colors = pd.DataFrame(index=peptides)
-    legend_handles: list[Patch] = []
+    # -- Categories per annotation column (category order, else
+    # lexicographic)
+    col_groups, col_cats = {}, {}
     for col in margin_cols:
         groups = adata.var.loc[peptides, col]
         if isinstance(groups.dtype, pd.CategoricalDtype):
             cats = [c for c in groups.cat.categories if c in set(groups)]
         else:
-            cats = sorted(map(str, groups.dropna().unique()))
-
-        resolved = _resolve_color_scheme(color_scheme, cats)
+            cats = sorted(groups.dropna().unique(), key=str)
+        col_groups[col], col_cats[col] = groups, cats
+
+    # -- Colours are resolved across all columns at once, so several
+    # strips never reuse the same colour
+    all_cats = [cat for col in margin_cols for cat in col_cats[col]]
+    cycle = plt.rcParams["axes.prop_cycle"].by_key().get("color", [])
+    if color_scheme is None and len(all_cats) > len(cycle):
+        resolved = sns.color_palette("husl", len(all_cats))
+    else:
+        resolved = _resolve_color_scheme(color_scheme, all_cats)
         if resolved is None:
-            resolved = (
-                sns.color_palette(n_colors=len(cats)) if len(cats) else []
-            )
-        palette = {str(cat): color for cat, color in zip(cats, resolved)}
+            resolved = sns.color_palette(n_colors=len(all_cats))
+
+    annot_colors = pd.DataFrame(index=peptides)
+    legend_groups: dict[str, list[Patch]] = {}
+    offset = 0
+    for col in margin_cols:
+        groups, cats = col_groups[col], col_cats[col]
+        colors = resolved[offset : offset + len(cats)]
+        offset += len(cats)
+        palette = {str(cat): color for cat, color in zip(cats, colors)}
 
         color_series = groups.astype("string").map(palette)
-        if color_series.isna().any() and groups.notna().any():
-            bad = color_series.isna() & groups.notna()
+        bad = color_series.isna() & groups.notna()
+        if bad.any():
             missing_cats = sorted(groups[bad].astype(str).unique())
             raise ValueError(
                 f"No color provided for categories in '{col}': "
                 f"{', '.join(missing_cats)}."
             )
 
-        prefix = f"{col}: " if len(margin_cols) > 1 else ""
-        legend_handles.extend(
-            Patch(
-                facecolor=palette[str(cat)],
-                edgecolor="none",
-                label=f"{prefix}{cat}",
-            )
+        handles = [
+            Patch(facecolor=palette[str(cat)], edgecolor="none", label=cat)
             for cat in cats
-        )
+        ]
 
         if groups.isna().any():
             na_color = mpl.colors.to_rgba("lightgray")
-            color_series = color_series.astype(object)
-            color_series[groups.isna()] = na_color
-            legend_handles.append(
-                Patch(
-                    facecolor=na_color, edgecolor="none", label=f"{prefix}NA"
-                )
+            # Element-wise: a masked assignment would broadcast the tuple
+            color_series = pd.Series(
+                [
+                    na_color if is_na else color
+                    for color, is_na in zip(color_series, groups.isna())
+                ],
+                index=color_series.index,
+                dtype=object,
+            )
+            handles.append(
+                Patch(facecolor=na_color, edgecolor="none", label="NA")
             )
 
         annot_colors[col] = color_series.to_numpy()
+        legend_groups[col] = handles
 
     annot_colors.index = pep_ids
     corr_df.index = pep_ids
     corr_df.columns = pep_ids
 
+    # -- Default order (used when not clustering): peptide_id categories,
+    # else lexicographic
+    pep_col = adata.var.loc[peptides, "peptide_id"]
+    if isinstance(pep_col.dtype, pd.CategoricalDtype):
+        present = set(pep_ids)
+        order = [str(c) for c in pep_col.cat.categories if str(c) in present]
+    else:
+        order = sorted(pep_ids)
+    corr_df = corr_df.loc[order, order]
+    annot_colors = annot_colors.loc[order]
+
     # -- Color center at the off-diagonal mean (as sample_correlation_matrix)
     A = corr_df.to_numpy(dtype=float)
     n = A.shape[0]
-    if n > 1:
-        offdiag = A[~np.eye(n, dtype=bool)]
-        center_val = np.nanmean(offdiag)
-    else:
-        offdiag = A.ravel()
-        center_val = float(np.nanmean(A))
+    offdiag = A[~np.eye(n, dtype=bool)]
+    finite = offdiag[~np.isnan(offdiag)]
+    center_val = float(finite.mean()) if finite.size else 0.0
 
     # -- Resolve the row/column linkage (symmetric)
     row_linkage = None
@@ -626,16 +769,26 @@ def pairwise_peptide_correlations_heatmap(
     do_cluster = cluster
     if cluster:
         if linkage is not None:
+            linkage = np.asarray(linkage, dtype=float)
+            if linkage.shape != (n - 1, 4):
+                raise ValueError(
+                    f"`linkage` must have shape ({n - 1}, 4) for the {n} "
+                    f"peptides of protein '{protein}'; got "
+                    f"{linkage.shape}."
+                )
             row_linkage = col_linkage = linkage
         else:
-            if np.isnan(A).any():
-                raise ValueError(
-                    "Correlation matrix for protein "
-                    f"'{protein}' contains missing values; clustering "
-                    "is not possible. Pass cluster=False to skip "
-                    "clustering."
+            n_nan = int(np.isnan(A[np.triu_indices(n, k=1)]).sum())
+            if n_nan:
+                warnings.warn(
+                    f"{n_nan} correlation(s) of protein '{protein}' are "
+                    "NaN (e.g. constant peptides or missing intensities); "
+                    "they are drawn hatched and treated as r = 0 when "
+                    "clustering.",
+                    UserWarning,
+                    stacklevel=2,
                 )
-            dist = 1.0 - A
+            dist = 1.0 - np.nan_to_num(A, nan=0.0)
             np.fill_diagonal(dist, 0.0)
             dist = np.clip(dist, 0, 2)
             Z = scipy_linkage(squareform(dist, checks=False), method=method)
@@ -643,33 +796,33 @@ def pairwise_peptide_correlations_heatmap(
 
     # -- Optional statistics printout
     if print_stats and n > 1:
+        values = finite if finite.size else np.array([np.nan])
         summary = pd.DataFrame(
             {
-                "min": [np.nanmin(offdiag)],
-                "max": [np.nanmax(offdiag)],
-                "mean": [np.nanmean(offdiag)],
-                "median": [np.nanmedian(offdiag)],
-                "std": [np.nanstd(offdiag)],
+                "min": [values.min()],
+                "max": [values.max()],
+                "mean": [values.mean()],
+                "median": [np.median(values)],
+                "std": [values.std()],
             }
         )
         print(f"Peptide correlation summary (off-diagonal, {protein}):")
         print(summary.to_string(index=False))
         print()
 
-        clusters = adata.var.loc[peptides, "cluster_id"]
-        rows = []
-        for cid, grp in clusters.groupby(clusters, observed=True):
-            idx = [peptides.index(p) for p in grp.index]
-            block = A[np.ix_(idx, idx)]
-            if len(idx) > 1:
-                within = block[~np.eye(len(idx), dtype=bool)]
-                mean_within = np.nanmean(within)
-            else:
-                mean_within = np.nan
-            rows.append({"cluster_id": cid, "mean_within": mean_within})
-        print(f"\nPer cluster_id (within-cluster mean correlation):")
-        print(pd.DataFrame(rows).to_string(index=False))
-        print()
+        for col in margin_cols:
+            groups = pd.Series(col_groups[col].to_numpy(), index=pep_ids)
+            rows = []
+            for gid in col_cats[col]:
+                members = groups.index[groups == gid]
+                block = corr_df.loc[members, members].to_numpy()
+                within = block[~np.eye(len(members), dtype=bool)]
+                within = within[~np.isnan(within)]
+                mean_within = within.mean() if within.size else np.nan
+                rows.append({col: gid, "mean_within": mean_within})
+            print(f"Per {col} (within-group mean correlation):")
+            print(pd.DataFrame(rows).to_string(index=False))
+            print()
 
     # -- Draw the clustered heatmap
     clustermap_kwargs = dict(
@@ -691,25 +844,27 @@ def pairwise_peptide_correlations_heatmap(
 
     g = sns.clustermap(corr_df, **clustermap_kwargs)
 
-    if legend_handles:
-        g.ax_heatmap.legend(
-            handles=legend_handles,
-            title=", ".join(margin_cols),
-            bbox_to_anchor=(1.05, 1),
-            loc="upper left",
-            borderaxespad=0.0,
-            frameon=False,
-        )
+    # NaN cells are masked by seaborn; a hatched background keeps them
+    # distinct from every colour on the correlation scale
+    if np.isnan(A).any():
+        g.ax_heatmap.patch.set_facecolor("white")
+        g.ax_heatmap.patch.set_edgecolor(NAN_HATCH_COLOR)
+        g.ax_heatmap.patch.set_hatch(NAN_HATCH)
+        legend_groups["correlation"] = [
+            Patch(
+                facecolor="white",
+                edgecolor=NAN_HATCH_COLOR,
+                hatch=NAN_HATCH,
+                label="NaN",
+            )
+        ]
 
     g.ax_heatmap.set_xlabel("Peptides")
     g.ax_heatmap.set_ylabel("Peptides")
-    g.figure.suptitle(protein)
-
-    plt.tight_layout()
+    g.ax_col_dendrogram.set_title(protein)
+    _place_margin_legends(g, legend_groups)
 
     if save is not None:
-        if not isinstance(save, (str, Path)):
-            raise TypeError("`save` must be a path-like object or None.")
         g.savefig(save, dpi=300, bbox_inches="tight")
     if show:
         plt.show()
@@ -717,3 +872,74 @@ def pairwise_peptide_correlations_heatmap(
     if ax:
         return g.ax_heatmap
     return None
+
+
+def _place_margin_legends(g, legend_groups):
+    """Stack one legend per annotation column right of all heatmap text.
+
+    The figure is widened when the legends would extend past its right
+    edge; every axes keeps its size in inches, so nothing overlaps and
+    nothing is cut off on screen or when saved.
+    """
+    fig = g.figure
+    if hasattr(fig.canvas, "get_renderer"):
+        renderer = fig.canvas.get_renderer()
+    else:
+        # Vector backends (pdf, svg) have no canvas renderer
+        renderer = fig._get_renderer()
+    to_fig = fig.transFigure.inverted()
+
+    text_axes = [
+        ax for ax in (g.ax_heatmap, g.ax_col_colors) if ax is not None
+    ]
+    right = max(
+        ax.get_tightbbox(renderer).transformed(to_fig).x1 for ax in text_axes
+    )
+    x = right + 0.02
+    y = g.ax_col_dendrogram.get_position().y1
+
+    legends = []
+    legends_right, legends_bottom = 1.0, 0.0
+    for title, handles in legend_groups.items():
+        legend = fig.legend(
+            handles=handles,
+            title=title,
+            loc="upper left",
+            bbox_to_anchor=(x, y),
+            borderaxespad=0.0,
+            frameon=False,
+        )
+        box = legend.get_window_extent(renderer).transformed(to_fig)
+        legends.append((legend, y))
+        legends_right = max(legends_right, box.x1)
+        legends_bottom = min(legends_bottom, box.y0)
+        y -= box.height + 0.02
+
+    if legends_right <= 1.0 and legends_bottom >= 0.0:
+        return
+    width, height = fig.get_size_inches()
+    new_width = width * max(legends_right + 0.01, 1.0)
+    extra_bottom = height * max(-legends_bottom + 0.01, 0.0)
+    new_height = height + extra_bottom
+    fig.set_size_inches(new_width, new_height)
+
+    def new_x(v):
+        return v * width / new_width
+
+    def new_y(v):
+        return (v * height + extra_bottom) / new_height
+
+    for ax in fig.axes:
+        pos = ax.get_position()
+        ax.set_position(
+            [
+                new_x(pos.x0),
+                new_y(pos.y0),
+                pos.width * width / new_width,
+                pos.height * height / new_height,
+            ]
+        )
+    for legend, top in legends:
+        legend.set_bbox_to_anchor(
+            (new_x(x), new_y(top)), transform=fig.transFigure
+        )
diff --git a/proteopy/utils/_matrix_wrangling.py b/proteopy/utils/_matrix_wrangling.py
index d733eb7..fb0e87d 100644
--- a/proteopy/utils/_matrix_wrangling.py
+++ b/proteopy/utils/_matrix_wrangling.py
@@ -9,6 +9,7 @@ def reconstruct_symmetric_matrix_from_long(
     value_col=2,
     *,
     diagonal=1.0,
+    allow_missing=False,
 ):
     """Reconstruct a square symmetric matrix from a long DataFrame.
 
@@ -32,6 +33,9 @@ def reconstruct_symmetric_matrix_from_long(
     diagonal : float
         Value written on the matrix diagonal (``1.0`` by default, the
         correlation convention).
+    allow_missing : bool
+        If ``True``, a label pair without a value in either triangle
+        (absent or ``NaN``) is left as ``NaN`` instead of raising.
 
     Returns
     -------
@@ -42,7 +46,8 @@ def reconstruct_symmetric_matrix_from_long(
     Raises
     ------
     ValueError
-        If a label pair is absent from both triangles, or if a pair is
+        If a label pair is absent from both triangles (and
+        ``allow_missing`` is ``False``), or if a pair is
         present in both triangles with conflicting values.
     """
     if isinstance(var_a_col, int):
@@ -84,6 +89,8 @@ def reconstruct_symmetric_matrix_from_long(
             elif lower_nan and not upper_nan:
                 matrix[j, i] = matrix[i, j]
             elif upper_nan and lower_nan:
+                if allow_missing:
+                    continue
                 raise ValueError(
                     "No value found for the combination of labels: "
                     f"{idx_to_label[i]} and {idx_to_label[j]}."
diff --git a/tests/pl/test_pairwise_peptide_correlations_heatmap.py b/tests/pl/test_pairwise_peptide_correlations_heatmap.py
index c6737cf..bc983d3 100644
--- a/tests/pl/test_pairwise_peptide_correlations_heatmap.py
+++ b/tests/pl/test_pairwise_peptide_correlations_heatmap.py
@@ -7,6 +7,7 @@
 """
 
 import inspect
+import warnings
 
 import numpy as np
 import pandas as pd
@@ -87,6 +88,28 @@ def _protein_peptides(adata, protein):
     return sorted(var.index[var["protein_id"] == protein].tolist())
 
 
+def tick_labels(ax, axis="y"):
+    ticks = ax.get_yticklabels() if axis == "y" else ax.get_xticklabels()
+    return [t.get_text() for t in ticks]
+
+
+def _rendered(ax):
+    """Tick labels and the matrix exactly as drawn (NaN for masked)."""
+    rows, cols = tick_labels(ax), tick_labels(ax, "x")
+    values = ax.collections[0].get_array().astype(float)
+    return (
+        rows,
+        cols,
+        np.ma.filled(values, np.nan).reshape(len(rows), len(cols)),
+    )
+
+
+def _recorrelate(adata):
+    """Recompute .uns correlations after editing adata.X."""
+    pr.tl.pairwise_peptide_correlations(adata)
+    return adata
+
+
 def _spy_clustermap(monkeypatch):
     """Capture kwargs passed to sns.clustermap; call through."""
     import proteopy.pl.copf as copf_mod
@@ -132,11 +155,13 @@ def test_signature_locked():
     assert sig.parameters["protein"].default is inspect.Parameter.empty
     defaults = sig.parameters
     assert defaults["corr_key"].default == "pairwise_peptide_correlations"
-    assert defaults["margin_color"].default == "cluster_id"
+    assert defaults["margin_color"].default == "proteoform_id"
     assert defaults["method"].default == "average"
     assert defaults["cluster"].default is True
     assert defaults["linkage"].default is None
     assert defaults["cmap"].default == "coolwarm"
+    assert defaults["xticklabels"].default == "auto"
+    assert defaults["yticklabels"].default == "auto"
     assert defaults["show"].default is True
     assert defaults["ax"].default is False
     assert defaults["print_stats"].default is False
@@ -169,13 +194,24 @@ def test_default_clustering_is_symmetric():
     assert x == y
 
 
-def test_cluster_false_keeps_input_order():
+def test_cluster_false_orders_plain_ids_lexicographically():
     adata = _copf_adata()
     axm = _function()(adata, protein="P1", cluster=False, show=False, ax=True)
     y = [t.get_text() for t in axm.get_yticklabels()]
     assert y == _protein_peptides(adata, "P1")
 
 
+def test_cluster_false_follows_peptide_id_categories():
+    adata = _copf_adata()
+    categories = list(reversed(adata.var["peptide_id"].tolist()))
+    adata.var["peptide_id"] = pd.Categorical(
+        adata.var["peptide_id"], categories=categories, ordered=True
+    )
+    axm = _function()(adata, protein="P1", cluster=False, show=False, ax=True)
+    y = [t.get_text() for t in axm.get_yticklabels()]
+    assert y == ["pep4", "pep3", "pep2", "pep1", "pep0"]
+
+
 def test_supplied_linkage_shadows_both_axes(monkeypatch):
     adata = _copf_adata()
     captured = _spy_clustermap(monkeypatch)
@@ -303,3 +339,382 @@ def test_input_not_mutated():
     pd.testing.assert_frame_equal(
         adata.uns["pairwise_peptide_correlations"], uns_before
     )
+
+
+# -- Robustness: NaN correlations, NaN annotations, colour schemes -------
+
+
+def _with_nan_correlation(adata):
+    corrs = adata.uns["pairwise_peptide_correlations"].copy()
+    first = np.flatnonzero(corrs.index == "P1")[0]
+    corrs.iloc[first, corrs.columns.get_loc("PCC")] = np.nan
+    adata.uns["pairwise_peptide_correlations"] = corrs
+    return adata
+
+
+def test_nan_correlations_are_clustered_as_zero_with_warning():
+    adata = _with_nan_correlation(_copf_adata())
+    with pytest.warns(UserWarning, match="treated as r = 0"):
+        axm = _function()(adata, protein="P1", show=False, ax=True)
+    assert tick_labels(axm, "x") == tick_labels(axm)
+
+
+def test_nan_correlations_cluster_false_draws():
+    adata = _with_nan_correlation(_copf_adata())
+    out = _function()(adata, protein="P1", cluster=False, show=False, ax=True)
+    assert isinstance(out, Axes)
+
+
+def test_margin_color_with_missing_values(monkeypatch):
+    adata = _copf_adata()
+    adata.var["region"] = pd.Series(
+        ["N-term", None, "C-term", None, "N-term", None, None, None],
+        index=adata.var_names,
+        dtype=object,
+    )
+    captured = _spy_clustermap(monkeypatch)
+    _function()(adata, protein="P1", margin_color="region", show=False)
+    strip = captured["kwargs"]["row_colors"]["region"]
+    gray = plt.matplotlib.colors.to_rgba("lightgray")
+    assert sum(tuple(c) == gray for c in strip) == 2
+
+
+def test_color_scheme_dict_keyed_by_raw_values(monkeypatch):
+    adata = _copf_adata()
+    values = adata.var.loc[_protein_peptides(adata, "P1"), "cluster_id"]
+    scheme = {v: "black" for v in values.unique()}
+    captured = _spy_clustermap(monkeypatch)
+    _function()(
+        adata,
+        protein="P1",
+        margin_color="cluster_id",
+        color_scheme=scheme,
+        show=False,
+    )
+    strip = captured["kwargs"]["row_colors"]["cluster_id"]
+    assert set(strip) == {"black"}
+
+
+def test_print_stats_per_margin_color(capsys):
+    adata = _copf_adata()
+    adata.var = adata.var.drop(columns=["cluster_id"])
+    _function()(
+        adata,
+        protein="P1",
+        margin_color="proteoform_id",
+        print_stats=True,
+        show=False,
+    )
+    out = capsys.readouterr().out
+    assert "Peptide correlation summary" in out
+    assert "Per proteoform_id" in out
+
+
+# -- Margin annotations: default, colours, legends ----------------------
+
+
+def test_default_margin_is_proteoform_id(monkeypatch):
+    adata = _copf_adata()
+    captured = _spy_clustermap(monkeypatch)
+    _function()(adata, protein="P1", show=False)
+    assert list(captured["kwargs"]["row_colors"].columns) == ["proteoform_id"]
+
+
+def test_multiple_margins_use_disjoint_colours(monkeypatch):
+    adata = _copf_adata()
+    captured = _spy_clustermap(monkeypatch)
+    _function()(
+        adata,
+        protein="P1",
+        margin_color=["cluster_id", "proteoform_id", "protein_id"],
+        show=False,
+    )
+    strips = captured["kwargs"]["row_colors"]
+    colours = {
+        col: {tuple(np.atleast_1d(c)) for c in strips[col]}
+        for col in strips.columns
+    }
+    assert not colours["cluster_id"] & colours["proteoform_id"]
+    assert not colours["cluster_id"] & colours["protein_id"]
+    assert not colours["proteoform_id"] & colours["protein_id"]
+
+
+def test_one_legend_per_margin_column():
+    adata = _copf_adata()
+    margins = ["cluster_id", "proteoform_id"]
+    axm = _function()(
+        adata, protein="P1", margin_color=margins, show=False, ax=True
+    )
+    titles = [leg.get_title().get_text() for leg in axm.figure.legends]
+    assert titles == margins
+
+
+def test_legends_clear_of_all_axes_text_and_inside_figure():
+    adata = _copf_adata()
+    axm = _function()(
+        adata,
+        protein="P1",
+        margin_color=["cluster_id", "proteoform_id", "protein_id"],
+        show=False,
+        ax=True,
+    )
+    fig = axm.figure
+    renderer = fig.canvas.get_renderer()
+    text_right = max(ax.get_tightbbox(renderer).x1 for ax in fig.axes)
+    boxes = [leg.get_window_extent(renderer) for leg in fig.legends]
+    assert all(box.x0 >= text_right for box in boxes)
+    assert all(box.x1 <= fig.bbox.x1 for box in boxes)
+    for upper, lower in zip(boxes, boxes[1:]):
+        assert lower.y1 <= upper.y0
+
+
+# -- Edge cases: input tables, arguments, backends, layout -------------
+
+
+def test_several_correlations_per_pair_raise():
+    adata = _copf_adata()
+    corrs = adata.uns["pairwise_peptide_correlations"]
+    duplicate = corrs.loc[["P1"]].iloc[[0]].assign(PCC=-0.5)
+    adata.uns["pairwise_peptide_correlations"] = pd.concat([corrs, duplicate])
+    with pytest.raises(ValueError, match="several correlations per"):
+        _function()(adata, protein="P1", show=False)
+
+
+def test_peptides_filtered_after_correlations_warn_and_are_dropped():
+    adata = _copf_adata()
+    sub = adata[:, adata.var_names != "pep0"].copy()
+    with pytest.warns(UserWarning, match="no longer in adata.var"):
+        axm = _function()(sub, protein="P1", show=False, ax=True)
+    assert "pep0" not in tick_labels(axm)
+    assert len(tick_labels(axm)) == 4
+
+
+def test_fewer_than_two_remaining_peptides_raise():
+    adata = _copf_adata()
+    dropped = {"pep1", "pep2", "pep3", "pep4"}
+    sub = adata[:, [p for p in adata.var_names if p not in dropped]].copy()
+    with pytest.warns(UserWarning):
+        with pytest.raises(ValueError, match="Fewer than two"):
+            _function()(sub, protein="P1", show=False)
+
+
+def test_empty_margin_list_raises():
+    with pytest.raises(ValueError, match="at least one column"):
+        _function()(_copf_adata(), protein="P1", margin_color=[], show=False)
+
+
+def test_duplicate_margin_columns_raise():
+    with pytest.raises(ValueError, match="Duplicate columns"):
+        _function()(
+            _copf_adata(),
+            protein="P1",
+            margin_color=["cluster_id", "cluster_id"],
+            show=False,
+        )
+
+
+def test_bad_save_type_raises_before_drawing():
+    adata = _copf_adata()
+    before = set(plt.get_fignums())
+    with pytest.raises(TypeError, match="save"):
+        _function()(adata, protein="P1", save=123, show=False)
+    assert set(plt.get_fignums()) == before
+
+
+def test_linkage_of_wrong_size_raises():
+    adata = _copf_adata()
+    Z = scipy_linkage(np.random.default_rng(0).random((3, 2)))
+    with pytest.raises(ValueError, match=r"shape \(4, 4\)"):
+        _function()(adata, protein="P1", linkage=Z, show=False)
+
+
+def test_print_stats_groups_follow_legend_order(capsys):
+    adata = _copf_adata()
+    adata.var["grp"] = [2, 10, 2, 10, 2, 2, 10, 2]
+    axm = _function()(
+        adata,
+        protein="P1",
+        margin_color="grp",
+        print_stats=True,
+        show=False,
+        ax=True,
+    )
+    legend = [t.get_text() for t in axm.figure.legends[0].get_texts()]
+    table = capsys.readouterr().out.split("Per grp")[1].splitlines()[2:]
+    stats = [line.split()[0] for line in table if line.strip()]
+    assert legend == ["10", "2"]
+    assert stats == legend
+
+
+def test_all_nan_correlations_draw_without_runtime_warnings():
+    adata = _copf_adata()
+    corrs = adata.uns["pairwise_peptide_correlations"].copy()
+    corrs.loc["P1", "PCC"] = np.nan
+    adata.uns["pairwise_peptide_correlations"] = corrs
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        _function()(
+            adata,
+            protein="P1",
+            margin_color="protein_id",
+            cluster=False,
+            print_stats=True,
+            show=False,
+        )
+    assert not [w for w in caught if issubclass(w.category, RuntimeWarning)]
+
+
+def test_missing_copf_columns_point_to_clustering_step():
+    adata = _copf_adata()
+    adata.var = adata.var.drop(columns=["cluster_id", "proteoform_id"])
+    with pytest.raises(KeyError, match="peptide_clusters_from_dendograms"):
+        _function()(adata, protein="P1", show=False)
+
+
+@pytest.mark.parametrize("backend", ["pdf", "svg"])
+def test_vector_backends(backend, tmp_path):
+    adata = _copf_adata()
+    plt.switch_backend(backend)
+    try:
+        _function()(
+            adata, protein="P1", show=False, save=tmp_path / f"h.{backend}"
+        )
+    finally:
+        plt.switch_backend("Agg")
+    assert (tmp_path / f"h.{backend}").stat().st_size > 0
+
+
+# The deliberately tiny figure is too small for seaborn's own layout
+@pytest.mark.filterwarnings("ignore:Tight layout not applied:UserWarning")
+def test_tall_legends_stay_inside_figure():
+    adata = _copf_adata()
+    adata.var["per_peptide"] = adata.var["peptide_id"].astype(str)
+    axm = _function()(
+        adata,
+        protein="P1",
+        margin_color=["per_peptide", "proteoform_id", "protein_id"],
+        figsize=(2.5, 1.2),
+        show=False,
+        ax=True,
+    )
+    fig = axm.figure
+    renderer = fig.canvas.get_renderer()
+    for legend in fig.legends:
+        box = legend.get_window_extent(renderer)
+        assert box.y0 >= 0
+        assert box.x1 <= fig.bbox.x1
+
+
+# -- Checklist scenarios: missing data, degenerate inputs, rendering -----
+
+
+def test_all_na_peptide_shown_as_empty_row_and_column():
+    adata = _copf_adata()
+    adata.X[:, 0] = np.nan  # pep0 never measured
+    _recorrelate(adata)
+    with pytest.warns(UserWarning) as record:
+        axm = _function()(adata, protein="P1", show=False, ax=True)
+    messages = " | ".join(str(w.message) for w in record)
+    assert "no correlations" in messages
+    assert "treated as r = 0" in messages
+    rows, cols, values = _rendered(axm)
+    assert "pep0" in rows
+    assert np.isnan(values[rows.index("pep0")]).all()
+    assert np.isnan(values[:, cols.index("pep0")]).all()
+
+
+def test_constant_peptide_draws_with_default_clustering():
+    adata = _copf_adata()
+    adata.X[:, 1] = 5.0  # zero variance -> NaN PCC
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")  # scipy ConstantInputWarning
+        _recorrelate(adata)
+    with pytest.warns(UserWarning, match="NaN"):
+        axm = _function()(adata, protein="P1", show=False, ax=True)
+    rows, _, values = _rendered(axm)
+    assert np.isnan(values[rows.index("pep1")]).sum() == len(rows) - 1
+
+
+def test_nan_cells_hatched_and_in_legend():
+    adata = _with_nan_correlation(_copf_adata())
+    with pytest.warns(UserWarning):
+        axm = _function()(adata, protein="P1", show=False, ax=True)
+    assert axm.patch.get_hatch()
+    titles = [leg.get_title().get_text() for leg in axm.figure.legends]
+    assert "correlation" in titles
+
+
+def test_complete_matrix_has_no_nan_legend():
+    axm = _function()(_copf_adata(), protein="P1", show=False, ax=True)
+    titles = [leg.get_title().get_text() for leg in axm.figure.legends]
+    assert "correlation" not in titles
+    assert not axm.patch.get_hatch()
+
+
+def test_fewer_than_three_samples_warn():
+    adata = _recorrelate(_copf_adata()[:2].copy())
+    with pytest.warns(UserWarning, match="fewer than 3 samples"):
+        _function()(adata, protein="P1", show=False)
+
+
+def test_two_peptide_protein_renders_2x2():
+    adata = _copf_adata()
+    adata = _recorrelate(adata[:, ["pep0", "pep1", "pep5", "pep6"]].copy())
+    axm = _function()(adata, protein="P1", show=False, ax=True)
+    rows, cols, values = _rendered(axm)
+    assert values.shape == (2, 2)
+    assert np.allclose(np.diag(values), 1.0)
+    assert np.allclose(values, values.T)
+
+
+def test_identical_peptides_render_r_of_one():
+    adata = _copf_adata()
+    adata.X[:, 1] = adata.X[:, 0]
+    _recorrelate(adata)
+    rows, cols, values = _rendered(
+        _function()(adata, protein="P1", show=False, ax=True)
+    )
+    assert values[rows.index("pep0"), cols.index("pep1")] == pytest.approx(1)
+
+
+@pytest.mark.parametrize("cluster", [True, False])
+def test_rendered_cells_match_correlations(cluster):
+    adata = _copf_adata()
+    long = adata.uns["pairwise_peptide_correlations"].loc["P1"]
+    truth = {}
+    for a, b, r in long[["pepA", "pepB", "PCC"]].itertuples(index=False):
+        truth[(a, b)] = truth[(b, a)] = r
+    axm = _function()(
+        adata, protein="P1", cluster=cluster, show=False, ax=True
+    )
+    rows, cols, values = _rendered(axm)
+    assert rows == cols
+    assert np.allclose(np.diag(values), 1.0)
+    for i, a in enumerate(rows):
+        for j, b in enumerate(cols):
+            if i != j:
+                assert values[i, j] == pytest.approx(truth[(a, b)])
+
+
+def test_default_labels_do_not_overlap_for_many_peptides():
+    rng = np.random.default_rng(3)
+    n_obs, n_pep = 12, 45
+    peptides = [f"PEPTIDE{i:02d}K" for i in range(n_pep)]
+    obs_names = [f"s{i}" for i in range(n_obs)]
+    adata = ad.AnnData(
+        X=rng.normal(20, 1, (n_obs, n_pep)),
+        obs=pd.DataFrame({"sample_id": obs_names}, index=obs_names),
+        var=pd.DataFrame(
+            {"peptide_id": peptides, "protein_id": "P9"}, index=peptides
+        ),
+    )
+    _recorrelate(adata)
+    axm = _function()(
+        adata, protein="P9", margin_color="protein_id", show=False, ax=True
+    )
+    renderer = axm.figure.canvas.get_renderer()
+    boxes = sorted(
+        (t.get_window_extent(renderer) for t in axm.get_yticklabels()),
+        key=lambda box: box.y0,
+    )
+    assert all(lower.y1 <= upper.y0 for lower, upper in zip(boxes, boxes[1:]))
diff --git a/tests/utils/test_matrix_wrangling.py b/tests/utils/test_matrix_wrangling.py
index 3732cee..1e668da 100644
--- a/tests/utils/test_matrix_wrangling.py
+++ b/tests/utils/test_matrix_wrangling.py
@@ -141,3 +141,19 @@ def test_conflicting_values_raise():
     )
     with pytest.raises(ValueError, match="Conflicting values"):
         reconstruct_symmetric_matrix_from_long(df, "a", "b", "v")
+
+
+def test_allow_missing_leaves_nan():
+    df = pd.DataFrame(
+        {
+            "a": ["a", "a", "b"],
+            "b": ["b", "c", "c"],
+            "v": [0.1, 0.2, np.nan],
+        }
+    )
+    result = reconstruct_symmetric_matrix_from_long(
+        df, "a", "b", "v", allow_missing=True
+    )
+    assert np.isnan(result.loc["b", "c"])
+    assert np.isnan(result.loc["c", "b"])
+    assert result.loc["a", "b"] == 0.1

From 7fe7ec015c2fcce924b9804a05b5f0f454099ba6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Levente=20Temesv=C3=A1ri-Nagy?=
 <147416790+leventetn@users.noreply.github.com>
Date: Wed, 7 Oct 2026 16:29:40 +0200
Subject: [PATCH 06/10] fix: Address peptide correlation heatmap review

  - Remove ax and print_stats; always return None.
  - Support shared and per-strip annotation palettes, resolving both
    through _resolve_color_scheme.
  - Preserve shared flat color lists and validate per-strip palette
counts.
  - Omit peptides absent from the correlation table without warning.
  - Extract input validation into _validate_heatmap_inputs().
  - Shorten docstrings and the changelog; correct the default annotation
    description to proteoform_id.
  - Replace placeholder examples with a runnable synthetic AnnData
example.
  - Update API and rendering tests; cover both representations of
peptides
    mapping to multiple proteins.

  Validation: 61 heatmap tests, 129 summarization tests, and 13 doctest
  statements passed.
---
 HISTORY.md                                    |   5 +-
 proteopy/pl/copf.py                           | 359 +++++++++---------
 ...t_pairwise_peptide_correlations_heatmap.py | 324 +++++++++++-----
 ...mmarize_peptides_by_neighbourhood_union.py |  18 +-
 4 files changed, 428 insertions(+), 278 deletions(-)

diff --git a/HISTORY.md b/HISTORY.md
index 4141158..47e9a18 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -12,10 +12,7 @@ and this project adheres to
 
 **Plotting** (`pr.pl`)
 
-- `pairwise_peptide_correlations_heatmap()`: per-protein clustered
-  peptide-correlation heatmap with proteoform (`cluster_id` /
-  `proteoform_id`) annotation strips, symmetric row/column linkage
-  (`method`, `linkage`, `cluster`) and multiple `margin_color` columns
+- `pairwise_peptide_correlations_heatmap()`: new feature
 - `peptide_intensities()`, `proteoform_intensities()`: new `facet_by`
   parameter splitting the samples (`.obs`) across a grid of subplots
 
diff --git a/proteopy/pl/copf.py b/proteopy/pl/copf.py
index 12a5202..c91e50f 100644
--- a/proteopy/pl/copf.py
+++ b/proteopy/pl/copf.py
@@ -407,63 +407,57 @@ def pairwise_peptide_correlations_heatmap(
     yticklabels: bool | str = "auto",
     figsize: tuple[float, float] = (8.0, 8.0),
     show: bool = True,
-    ax: bool = False,
-    print_stats: bool = False,
     save: str | Path | None = None,
-) -> Axes | None:
-    """Clustered peptide-correlation heatmap for a single protein.
+) -> None:
+    """Peptide-correlation heatmap for a single protein.
 
-    Draw the peptide x peptide Pearson-correlation matrix of one protein
-    as a clustered heatmap, so that peptides of the same proteoform (COPF
-    cluster) group into visible blocks. One or more categorical
-    annotation strips (``margin_color``) are drawn alongside the heatmap,
-    by default the per-protein ``cluster_id`` assignment.
-
-    The correlations are read from
-    ``adata.uns[corr_key]`` (produced by
-    :func:`proteopy.tl.pairwise_peptide_correlations`) and reconstructed
-    into a full symmetric matrix for the selected protein. The clustering
-    is symmetric: the same tree is used for rows and columns.
+    Draw a protein's peptide Pearson correlations as a heatmap,
+    optionally with hierarchical clustering. Correlations are read
+    from ``adata.uns[corr_key]``, computed by
+    :func:`proteopy.tl.pairwise_peptide_correlations`. By default,
+    margin annotations show COPF proteoform assignments from
+    ``adata.var["proteoform_id"]``.
 
     Parameters
     ----------
     adata : AnnData
-        :class:`~anndata.AnnData` carrying COPF annotations.
+        Peptide-level :class:`~anndata.AnnData` with pairwise peptide
+        correlations in ``.uns[corr_key]`` and margin annotations
+        in ``.var``.
     protein : str
         ``protein_id`` of the protein to plot.
     corr_key : str
         Key in ``adata.uns`` holding the long-form pairwise correlations
-        (columns ``pepA``, ``pepB``, ``PCC``; indexed by ``protein_id``).
+        (columns ``pepA``, ``pepB``, ``PCC``; indexed by
+        ``protein_id``).
     margin_color : str | list[str]
-        Column(s) in ``adata.var`` used for the annotation strip(s). A
-        string is treated as a single-item list; a list draws one strip
-        per entry, in the given order. The default shows the COPF
-        proteoform labels (``<protein_id>_<cluster_id>``) written by
-        :func:`proteopy.tl.peptide_clusters_from_dendograms`. Colours
-        are assigned across all columns together, so strips never share
-        a colour, and each column gets its own legend, placed right of
-        all heatmap labels. Legend entries follow the column's category
-        order if it is categorical, otherwise the lexicographic order of
-        the ``str``-coerced values; store the column as an ordered
-        :class:`pandas.Categorical` to change it. Missing values are
-        drawn in light gray.
+        Columns in ``adata.var`` shown as annotation strips, in the
+        given order. The default uses ``proteoform_id`` written by
+        :func:`proteopy.tl.peptide_clusters_from_dendograms`.
+        Legends follow category order, otherwise lexicographic order;
+        use an ordered :class:`pandas.Categorical` to customize it.
+        Missing annotations appear in light gray.
     method : str
         Linkage method handed to :func:`scipy.cluster.hierarchy.linkage`
         when the tree is recomputed from ``1 - correlation``.
     cluster : bool
-        Cluster and draw dendrograms on both axes. When ``False``,
-        peptides follow the category order of ``adata.var["peptide_id"]``
-        if it is categorical, otherwise lexicographic order. ``NaN``
-        correlations (e.g. from constant peptides) are drawn hatched and
-        treated as ``r = 0`` when building the tree.
+        Cluster peptides and draw dendrograms on both axes. When
+        ``False``, order peptides lexicographically unless
+        ``adata.var["peptide_id"]`` is categorical, in which case
+        its category order is used.
     linkage : numpy.ndarray | None
-        Precomputed SciPy linkage matrix of shape ``(n_peptides - 1, 4)``.
+        Precomputed SciPy linkage matrix of shape
+        ``(n_peptides - 1, 4)``.
         When provided it is used for both axes, shadowing the default
         ``1 - correlation`` computation. Leaf indices refer to the
         peptides in the default order described under ``cluster``.
     color_scheme : Any
-        Color palette specification understood by
-        :func:`proteopy.utils.matplotlib._resolve_color_scheme`.
+        Palette for margin annotations: a colormap name or object,
+        color, color list, category-to-color dict, or callable. A single
+        palette is shared across strips. For multiple strips, a list
+        of palettes must match ``margin_color`` in length and order;
+        nest color lists to supply one per strip. ``None`` assigns
+        colors that do not overlap between strips.
     cmap : str
         Continuous colormap for the heatmap body.
     xticklabels, yticklabels : bool | str
@@ -475,17 +469,13 @@ def pairwise_peptide_correlations_heatmap(
         Matplotlib figure size in inches.
     show : bool
         Display the figure with :func:`matplotlib.pyplot.show`.
-    ax : bool
-        Return the heatmap :class:`matplotlib.axes.Axes` when ``True``.
-    print_stats : bool
-        Print correlation summary statistics before drawing.
     save : str | Path | None
         File path to save the figure. ``None`` skips saving.
 
     Returns
     -------
-    Axes or None
-        Heatmap axes when ``ax`` is ``True``; otherwise ``None``.
+    None
+        Creates its own figure with :func:`seaborn.clustermap`.
 
     Raises
     ------
@@ -494,8 +484,9 @@ def pairwise_peptide_correlations_heatmap(
         has no pairwise correlations, the table holds several
         correlations per peptide pair (e.g. per-batch results),
         ``margin_color`` is empty or repeats a column, ``linkage`` does
-        not match the number of peptides, or clustering is requested on
-        a matrix containing missing values.
+        not match the number of peptides, a per-strip ``color_scheme``
+        list has the wrong length, or clustering is requested on a
+        matrix containing missing values.
     KeyError
         If a ``margin_color`` column is not present in ``adata.var``.
     TypeError
@@ -507,101 +498,61 @@ def pairwise_peptide_correlations_heatmap(
     UserWarning
         If peptides in ``adata.uns[corr_key]`` are no longer in
         ``adata.var`` (filtered after the correlations were computed;
-        they are left out), if peptides of ``protein`` have no
-        correlations (shown as empty rows and columns), if NaN
-        correlations are treated as ``r = 0`` for clustering, or if
+        they are left out), if NaN correlations are treated as
+        ``r = 0`` for clustering, or if
         ``adata`` has fewer than 3 samples.
 
-    Notes
-    -----
-    The heatmap only displays correlations; it does not transform
-    intensities. They are computed by
-    :func:`proteopy.tl.pairwise_peptide_correlations` on ``adata.X`` as
-    stored, so their scale is the scale of ``.X``: Pearson correlations
-    of raw and log-transformed intensities differ, and a single extreme
-    sample can dominate raw-scale correlations. Log-transform (and
-    aggregate precursors to peptides) before computing correlations.
-    ``scipy.stats.pearsonr`` is not NaN-aware, so a pair is ``NaN`` as
-    soon as either peptide is missing in any sample, and there is no
-    minimum-overlap cutoff.
-
     Examples
     --------
+    >>> import anndata as ad
+    >>> import numpy as np
+    >>> import pandas as pd
     >>> import proteopy as pr
-    >>> adata = pr.read.long(...)
-    >>> pr.tl.pairwise_peptide_correlations(adata)
-    >>> pr.tl.peptide_dendograms_by_correlation(
-    ...     adata,
-    ...     method='agglomerative-hierarchical-clustering',
+    >>> samples = ["s1", "s2", "s3", "s4"]
+    >>> peptides = ["pep1", "pep2", "pep3", "pep4"]
+    >>> adata = ad.AnnData(
+    ...     X=np.array(
+    ...         [[1, 2, 4, 5], [2, 3, 3, 4],
+    ...          [3, 4, 2, 3], [4, 5, 1, 2]],
+    ...         dtype=float,
+    ...     ),
+    ...     obs=pd.DataFrame({"sample_id": samples}, index=samples),
+    ...     var=pd.DataFrame(
+    ...         {"peptide_id": peptides, "protein_id": "P1"},
+    ...         index=peptides,
+    ...     ),
     ... )
+    >>> pr.tl.pairwise_peptide_correlations(adata)
+    >>> pr.tl.peptide_dendograms_by_correlation(adata)
     >>> pr.tl.peptide_clusters_from_dendograms(
     ...     adata,
     ...     n_clusters=2,
     ...     min_peptides_per_cluster=2,
     ... )
-    >>> pr.pl.pairwise_peptide_correlations_heatmap(adata, protein="P12345")
+    >>> pr.pl.pairwise_peptide_correlations_heatmap(
+    ...     adata, protein="P1", show=False,
+    ... )
 
     Annotate by proteoform and draw two strips:
 
     >>> pr.pl.pairwise_peptide_correlations_heatmap(
     ...     adata,
-    ...     protein="P12345",
+    ...     protein="P1",
     ...     margin_color=["cluster_id", "proteoform_id"],
+    ...     show=False,
     ... )
     """
     check_proteodata(adata)
 
-    # -- Validate arguments up front
-    if not isinstance(method, str):
-        raise TypeError("`method` must be a string.")
-    if not isinstance(cluster, bool):
-        raise TypeError("`cluster` must be a bool.")
-
-    if isinstance(margin_color, str):
-        margin_cols = [margin_color]
-    elif isinstance(margin_color, list) and all(
-        isinstance(c, str) for c in margin_color
-    ):
-        margin_cols = list(margin_color)
-    else:
-        raise TypeError("`margin_color` must be a string or list of strings.")
-    if not margin_cols:
-        raise ValueError("`margin_color` must name at least one column.")
-    duplicated_cols = sorted(
-        {c for c in margin_cols if margin_cols.count(c) > 1}
+    margin_cols = _validate_heatmap_inputs(
+        adata=adata,
+        protein=protein,
+        corr_key=corr_key,
+        margin_color=margin_color,
+        method=method,
+        cluster=cluster,
+        save=save,
     )
-    if duplicated_cols:
-        raise ValueError(
-            "Duplicate columns in `margin_color`: "
-            f"{', '.join(duplicated_cols)}."
-        )
-    if save is not None and not isinstance(save, (str, Path)):
-        raise TypeError("`save` must be a path-like object or None.")
-
-    if protein not in set(adata.var["protein_id"]):
-        raise ValueError(
-            f"protein '{protein}' not found in adata.var['protein_id']."
-        )
-
-    if corr_key not in adata.uns:
-        raise ValueError(
-            f"'{corr_key}' not found in adata.uns; run "
-            "pr.tl.pairwise_peptide_correlations() first."
-        )
-
-    missing_cols = [c for c in margin_cols if c not in adata.var.columns]
-    if missing_cols:
-        hint = ""
-        if {"cluster_id", "proteoform_id"} & set(missing_cols):
-            hint = (
-                " These columns are written by "
-                "pr.tl.peptide_clusters_from_dendograms(); run it first or "
-                "pass another .var column."
-            )
-        raise KeyError(
-            "margin_color column(s) not found in adata.var: "
-            f"{', '.join(missing_cols)}.{hint}"
-        )
 
     if adata.n_obs < 3:
         warnings.warn(
@@ -659,20 +610,6 @@ def pairwise_peptide_correlations_heatmap(
             )
         corr_df = corr_df.loc[kept, kept]
 
-    protein_peptides = adata.var.index[adata.var["protein_id"] == protein]
-    without_corr = [p for p in protein_peptides if p not in corr_df.index]
-    if without_corr:
-        warnings.warn(
-            f"{len(without_corr)} peptide(s) of protein '{protein}' have "
-            f"no correlations in adata.uns['{corr_key}'] (e.g. all "
-            "intensities missing); they are shown as empty rows and "
-            "columns.",
-            UserWarning,
-            stacklevel=2,
-        )
-        full = list(corr_df.index) + without_corr
-        corr_df = corr_df.reindex(index=full, columns=full)
-
     peptides = list(corr_df.index)
 
     # -- Labels come from the peptide_id column, never the index
@@ -689,24 +626,13 @@ def pairwise_peptide_correlations_heatmap(
             cats = sorted(groups.dropna().unique(), key=str)
         col_groups[col], col_cats[col] = groups, cats
 
-    # -- Colours are resolved across all columns at once, so several
-    # strips never reuse the same colour
-    all_cats = [cat for col in margin_cols for cat in col_cats[col]]
-    cycle = plt.rcParams["axes.prop_cycle"].by_key().get("color", [])
-    if color_scheme is None and len(all_cats) > len(cycle):
-        resolved = sns.color_palette("husl", len(all_cats))
-    else:
-        resolved = _resolve_color_scheme(color_scheme, all_cats)
-        if resolved is None:
-            resolved = sns.color_palette(n_colors=len(all_cats))
-
+    # -- Resolve shared or per-strip palettes
+    margin_palettes = _resolve_margin_palettes(color_scheme, col_cats)
     annot_colors = pd.DataFrame(index=peptides)
-    legend_groups: dict[str, list[Patch]] = {}
-    offset = 0
+    legend_groups = {}
     for col in margin_cols:
         groups, cats = col_groups[col], col_cats[col]
-        colors = resolved[offset : offset + len(cats)]
-        offset += len(cats)
+        colors = margin_palettes[col]
         palette = {str(cat): color for cat, color in zip(cats, colors)}
 
         color_series = groups.astype("string").map(palette)
@@ -794,36 +720,6 @@ def pairwise_peptide_correlations_heatmap(
             Z = scipy_linkage(squareform(dist, checks=False), method=method)
             row_linkage = col_linkage = Z
 
-    # -- Optional statistics printout
-    if print_stats and n > 1:
-        values = finite if finite.size else np.array([np.nan])
-        summary = pd.DataFrame(
-            {
-                "min": [values.min()],
-                "max": [values.max()],
-                "mean": [values.mean()],
-                "median": [np.median(values)],
-                "std": [values.std()],
-            }
-        )
-        print(f"Peptide correlation summary (off-diagonal, {protein}):")
-        print(summary.to_string(index=False))
-        print()
-
-        for col in margin_cols:
-            groups = pd.Series(col_groups[col].to_numpy(), index=pep_ids)
-            rows = []
-            for gid in col_cats[col]:
-                members = groups.index[groups == gid]
-                block = corr_df.loc[members, members].to_numpy()
-                within = block[~np.eye(len(members), dtype=bool)]
-                within = within[~np.isnan(within)]
-                mean_within = within.mean() if within.size else np.nan
-                rows.append({col: gid, "mean_within": mean_within})
-            print(f"Per {col} (within-group mean correlation):")
-            print(pd.DataFrame(rows).to_string(index=False))
-            print()
-
     # -- Draw the clustered heatmap
     clustermap_kwargs = dict(
         row_colors=annot_colors,
@@ -869,17 +765,114 @@ def pairwise_peptide_correlations_heatmap(
     if show:
         plt.show()
 
-    if ax:
-        return g.ax_heatmap
     return None
 
 
+def _validate_heatmap_inputs(
+    adata: ad.AnnData,
+    protein: str,
+    *,
+    corr_key: str,
+    margin_color: str | list[str],
+    method: str,
+    cluster: bool,
+    save: str | Path | None,
+) -> list[str]:
+    """Check heatmap arguments and annotations; return margin columns."""
+    if not isinstance(method, str):
+        raise TypeError("`method` must be a string.")
+    if not isinstance(cluster, bool):
+        raise TypeError("`cluster` must be a bool.")
+
+    if isinstance(margin_color, str):
+        margin_cols = [margin_color]
+    elif isinstance(margin_color, list) and all(
+        isinstance(c, str) for c in margin_color
+    ):
+        margin_cols = list(margin_color)
+    else:
+        raise TypeError("`margin_color` must be a string or list of strings.")
+    if not margin_cols:
+        raise ValueError("`margin_color` must name at least one column.")
+    duplicated_cols = sorted(
+        {c for c in margin_cols if margin_cols.count(c) > 1}
+    )
+    if duplicated_cols:
+        raise ValueError(
+            "Duplicate columns in `margin_color`: "
+            f"{', '.join(duplicated_cols)}."
+        )
+    if save is not None and not isinstance(save, (str, Path)):
+        raise TypeError("`save` must be a path-like object or None.")
+
+    if protein not in set(adata.var["protein_id"]):
+        raise ValueError(
+            f"protein '{protein}' not found in adata.var['protein_id']."
+        )
+
+    if corr_key not in adata.uns:
+        raise ValueError(
+            f"'{corr_key}' not found in adata.uns; run "
+            "pr.tl.pairwise_peptide_correlations() first."
+        )
+
+    missing_cols = [c for c in margin_cols if c not in adata.var.columns]
+    if missing_cols:
+        hint = ""
+        if {"cluster_id", "proteoform_id"} & set(missing_cols):
+            hint = (
+                " These columns are written by "
+                "pr.tl.peptide_clusters_from_dendograms(); run it first or "
+                "pass another .var column."
+            )
+        raise KeyError(
+            "margin_color column(s) not found in adata.var: "
+            f"{', '.join(missing_cols)}.{hint}"
+        )
+
+    return margin_cols
+
+
+def _resolve_margin_palettes(color_scheme, col_cats):
+    """Resolve annotation colors, sharing a palette unless given several."""
+    per_strip = (
+        len(col_cats) > 1
+        and isinstance(color_scheme, list)
+        and bool(color_scheme)
+        and not all(mpl.colors.is_color_like(c) for c in color_scheme)
+    )
+    if per_strip:
+        if len(color_scheme) != len(col_cats):
+            raise ValueError(
+                "A per-strip `color_scheme` list must have the same "
+                "length as `margin_color`."
+            )
+        return {
+            col: _resolve_color_scheme(scheme, cats) or []
+            for (col, cats), scheme in zip(col_cats.items(), color_scheme)
+        }
+
+    all_cats = [cat for cats in col_cats.values() for cat in cats]
+    cycle = plt.rcParams["axes.prop_cycle"].by_key().get("color", [])
+    if color_scheme is None and len(all_cats) > len(cycle):
+        color_scheme = sns.color_palette("husl", len(all_cats))
+    resolved = _resolve_color_scheme(color_scheme, all_cats)
+    if resolved is None:
+        resolved = sns.color_palette(n_colors=len(all_cats))
+
+    palettes = {}
+    offset = 0
+    for col, cats in col_cats.items():
+        palettes[col] = resolved[offset : offset + len(cats)]
+        offset += len(cats)
+    return palettes
+
+
 def _place_margin_legends(g, legend_groups):
-    """Stack one legend per annotation column right of all heatmap text.
+    """Place annotation legends to the right of the heatmap and its labels.
 
-    The figure is widened when the legends would extend past its right
-    edge; every axes keeps its size in inches, so nothing overlaps and
-    nothing is cut off on screen or when saved.
+    Stack legends vertically and expand the figure as needed to keep
+    them inside its bounds without shrinking the existing axes.
     """
     fig = g.figure
     if hasattr(fig.canvas, "get_renderer"):
diff --git a/tests/pl/test_pairwise_peptide_correlations_heatmap.py b/tests/pl/test_pairwise_peptide_correlations_heatmap.py
index bc983d3..4a87b56 100644
--- a/tests/pl/test_pairwise_peptide_correlations_heatmap.py
+++ b/tests/pl/test_pairwise_peptide_correlations_heatmap.py
@@ -13,7 +13,6 @@
 import pandas as pd
 import anndata as ad
 import matplotlib.pyplot as plt
-from matplotlib.axes import Axes
 from scipy.spatial.distance import squareform
 from scipy.cluster.hierarchy import linkage as scipy_linkage
 import pytest
@@ -35,6 +34,16 @@ def _function():
     return getattr(pr.pl, "pairwise_peptide_correlations_heatmap")
 
 
+def _plot_heatmap(*args, **kwargs):
+    """Draw the plot and find its heatmap axes in the generated figure."""
+    assert _function()(*args, **kwargs) is None
+    return next(
+        ax
+        for ax in plt.gcf().axes
+        if ax.get_xlabel() == "Peptides" and ax.get_ylabel() == "Peptides"
+    )
+
+
 def _copf_adata(seed=0):
     """Peptide-level AnnData with the COPF chain already run.
 
@@ -147,8 +156,6 @@ def test_signature_locked():
         "yticklabels",
         "figsize",
         "show",
-        "ax",
-        "print_stats",
         "save",
     ]
     # protein is required (no default)
@@ -163,24 +170,20 @@ def test_signature_locked():
     assert defaults["xticklabels"].default == "auto"
     assert defaults["yticklabels"].default == "auto"
     assert defaults["show"].default is True
-    assert defaults["ax"].default is False
-    assert defaults["print_stats"].default is False
     assert defaults["save"].default is None
 
 
 # -- Return-value semantics ----------------------------------------------
 
 
-def test_returns_axes_when_ax_true():
-    adata = _copf_adata()
-    out = _function()(adata, protein="P1", show=False, ax=True)
-    assert isinstance(out, Axes)
-
-
-def test_returns_none_when_ax_false():
+@pytest.mark.parametrize("show", [True, False])
+@pytest.mark.parametrize("cluster", [True, False])
+def test_always_returns_none(show, cluster, monkeypatch):
     adata = _copf_adata()
-    out = _function()(adata, protein="P1", show=False, ax=False)
-    assert out is None
+    shown = []
+    monkeypatch.setattr(plt, "show", lambda: shown.append(True))
+    assert _function()(adata, protein="P1", show=show, cluster=cluster) is None
+    assert shown == ([True] if show else [])
 
 
 # -- Clustering: symmetry, linkage pass-through, toggle -------------------
@@ -188,7 +191,7 @@ def test_returns_none_when_ax_false():
 
 def test_default_clustering_is_symmetric():
     adata = _copf_adata()
-    axm = _function()(adata, protein="P1", show=False, ax=True)
+    axm = _plot_heatmap(adata, protein="P1", show=False)
     x = [t.get_text() for t in axm.get_xticklabels()]
     y = [t.get_text() for t in axm.get_yticklabels()]
     assert x == y
@@ -196,7 +199,7 @@ def test_default_clustering_is_symmetric():
 
 def test_cluster_false_orders_plain_ids_lexicographically():
     adata = _copf_adata()
-    axm = _function()(adata, protein="P1", cluster=False, show=False, ax=True)
+    axm = _plot_heatmap(adata, protein="P1", cluster=False, show=False)
     y = [t.get_text() for t in axm.get_yticklabels()]
     assert y == _protein_peptides(adata, "P1")
 
@@ -207,7 +210,7 @@ def test_cluster_false_follows_peptide_id_categories():
     adata.var["peptide_id"] = pd.Categorical(
         adata.var["peptide_id"], categories=categories, ordered=True
     )
-    axm = _function()(adata, protein="P1", cluster=False, show=False, ax=True)
+    axm = _plot_heatmap(adata, protein="P1", cluster=False, show=False)
     y = [t.get_text() for t in axm.get_yticklabels()]
     assert y == ["pep4", "pep3", "pep2", "pep1", "pep0"]
 
@@ -231,7 +234,7 @@ def test_supplied_linkage_shadows_both_axes(monkeypatch):
         squareform(np.clip(dist, 0, 2), checks=False), method="complete"
     )
 
-    _function()(adata, protein="P1", linkage=Z, show=False, ax=True)
+    _function()(adata, protein="P1", linkage=Z, show=False)
     kwargs = captured["kwargs"]
     assert kwargs["row_linkage"] is Z
     assert kwargs["col_linkage"] is Z
@@ -243,9 +246,7 @@ def test_supplied_linkage_shadows_both_axes(monkeypatch):
 def test_margin_color_string_single_strip(monkeypatch):
     adata = _copf_adata()
     captured = _spy_clustermap(monkeypatch)
-    _function()(
-        adata, protein="P1", margin_color="cluster_id", show=False, ax=False
-    )
+    _function()(adata, protein="P1", margin_color="cluster_id", show=False)
     row_colors = captured["kwargs"]["row_colors"]
     assert list(row_colors.columns) == ["cluster_id"]
 
@@ -258,7 +259,6 @@ def test_margin_color_list_multiple_strips(monkeypatch):
         protein="P1",
         margin_color=["cluster_id", "proteoform_id"],
         show=False,
-        ax=False,
     )
     row_colors = captured["kwargs"]["row_colors"]
     assert list(row_colors.columns) == ["cluster_id", "proteoform_id"]
@@ -269,9 +269,7 @@ def test_margin_color_list_multiple_strips(monkeypatch):
 def test_annotation_distinct_colors_match_categories(monkeypatch):
     adata = _copf_adata()
     captured = _spy_clustermap(monkeypatch)
-    _function()(
-        adata, protein="P1", margin_color="cluster_id", show=False, ax=False
-    )
+    _function()(adata, protein="P1", margin_color="cluster_id", show=False)
     strip = captured["kwargs"]["row_colors"]["cluster_id"]
     n_colors = len({tuple(np.atleast_1d(c)) for c in strip})
     n_cats = adata.var.loc[
@@ -285,7 +283,7 @@ def test_annotation_distinct_colors_match_categories(monkeypatch):
 
 def test_tick_labels_from_peptide_id():
     adata = _copf_adata()
-    axm = _function()(adata, protein="P1", show=False, ax=True)
+    axm = _plot_heatmap(adata, protein="P1", show=False)
     y = [t.get_text() for t in axm.get_yticklabels()]
     assert sorted(y) == _protein_peptides(adata, "P1")
 
@@ -326,7 +324,7 @@ def test_bad_margin_color_type_raises():
 def test_save_writes_png(tmp_path):
     adata = _copf_adata()
     out = tmp_path / "heatmap.png"
-    _function()(adata, protein="P1", show=False, save=out)
+    assert _function()(adata, protein="P1", show=False, save=out) is None
     assert out.exists() and out.stat().st_size > 0
 
 
@@ -355,16 +353,10 @@ def _with_nan_correlation(adata):
 def test_nan_correlations_are_clustered_as_zero_with_warning():
     adata = _with_nan_correlation(_copf_adata())
     with pytest.warns(UserWarning, match="treated as r = 0"):
-        axm = _function()(adata, protein="P1", show=False, ax=True)
+        axm = _plot_heatmap(adata, protein="P1", show=False)
     assert tick_labels(axm, "x") == tick_labels(axm)
 
 
-def test_nan_correlations_cluster_false_draws():
-    adata = _with_nan_correlation(_copf_adata())
-    out = _function()(adata, protein="P1", cluster=False, show=False, ax=True)
-    assert isinstance(out, Axes)
-
-
 def test_margin_color_with_missing_values(monkeypatch):
     adata = _copf_adata()
     adata.var["region"] = pd.Series(
@@ -395,19 +387,174 @@ def test_color_scheme_dict_keyed_by_raw_values(monkeypatch):
     assert set(strip) == {"black"}
 
 
-def test_print_stats_per_margin_color(capsys):
+def _color_scheme_adata():
     adata = _copf_adata()
-    adata.var = adata.var.drop(columns=["cluster_id"])
+    adata.var["region"] = ["b", "a", "b", "a", "b", "a", "b", "a"]
+    adata.var["phase"] = [2, 10, 2, 10, 2, 10, 2, 10]
+    return adata
+
+
+def _spy_color_resolver(monkeypatch):
+    import proteopy.pl.copf as copf_mod
+
+    calls = []
+    original = copf_mod._resolve_color_scheme
+
+    def spy(scheme, labels):
+        calls.append((scheme, list(labels)))
+        return original(scheme, labels)
+
+    monkeypatch.setattr(copf_mod, "_resolve_color_scheme", spy)
+    return calls
+
+
+def _assert_strip_colors(strip, groups, palette):
+    actual = [plt.matplotlib.colors.to_rgba(color) for color in strip]
+    expected = [
+        plt.matplotlib.colors.to_rgba(palette[group]) for group in groups
+    ]
+    np.testing.assert_allclose(actual, expected)
+
+
+@pytest.mark.parametrize(
+    "schemes, expected",
+    [
+        (
+            ["Blues", "Reds"],
+            [
+                {
+                    "a": plt.get_cmap("Blues")(0.0),
+                    "b": plt.get_cmap("Blues")(1.0),
+                },
+                {10: plt.get_cmap("Reds")(0.0), 2: plt.get_cmap("Reds")(1.0)},
+            ],
+        ),
+        (
+            [["red", "blue"], ["green", "orange"]],
+            [{"a": "red", "b": "blue"}, {10: "green", 2: "orange"}],
+        ),
+        (
+            [{"a": "blue", "b": "red"}, {2: "green", 10: "orange"}],
+            [{"a": "blue", "b": "red"}, {10: "orange", 2: "green"}],
+        ),
+        (
+            [plt.get_cmap("Blues"), plt.get_cmap("Reds")],
+            [
+                {
+                    "a": plt.get_cmap("Blues")(0.0),
+                    "b": plt.get_cmap("Blues")(1.0),
+                },
+                {10: plt.get_cmap("Reds")(0.0), 2: plt.get_cmap("Reds")(1.0)},
+            ],
+        ),
+        (
+            [lambda position: "red", lambda position: "blue"],
+            [{"a": "red", "b": "red"}, {10: "blue", 2: "blue"}],
+        ),
+    ],
+    ids=["named", "color_lists", "dicts", "colormaps", "callables"],
+)
+def test_per_strip_palettes_resolve_separately(schemes, expected, monkeypatch):
+    adata = _color_scheme_adata()
+    captured = _spy_clustermap(monkeypatch)
+    calls = _spy_color_resolver(monkeypatch)
+    margins = ["region", "phase"]
+    _function()(
+        adata,
+        protein="P1",
+        margin_color=margins,
+        color_scheme=schemes,
+        show=False,
+    )
+    assert len(calls) == 2
+    assert calls[0][0] is schemes[0]
+    assert calls[1][0] is schemes[1]
+    assert calls[0][1] == ["a", "b"]
+    assert calls[1][1] == [10, 2]
+    strips = captured["kwargs"]["row_colors"]
+    assert list(strips.columns) == margins
+    assert captured["kwargs"]["col_colors"] is strips
+    for col, palette in zip(margins, expected):
+        _assert_strip_colors(
+            strips[col], adata.var.loc[strips.index, col], palette
+        )
+
+
+@pytest.mark.parametrize(
+    "scheme, expected",
+    [
+        (
+            "viridis",
+            [plt.get_cmap("viridis")(p) for p in [0.0, 1 / 3, 2 / 3, 1.0]],
+        ),
+        (
+            ["red", "blue", "green", "orange"],
+            ["red", "blue", "green", "orange"],
+        ),
+        (
+            {"a": "red", "b": "blue", 10: "green", 2: "orange"},
+            ["red", "blue", "green", "orange"],
+        ),
+    ],
+    ids=["named", "flat_color_list", "dict"],
+)
+def test_shared_palette_resolves_all_strips_together(
+    scheme, expected, monkeypatch
+):
+    adata = _color_scheme_adata()
+    captured = _spy_clustermap(monkeypatch)
+    calls = _spy_color_resolver(monkeypatch)
     _function()(
         adata,
         protein="P1",
-        margin_color="proteoform_id",
-        print_stats=True,
+        margin_color=["region", "phase"],
+        color_scheme=scheme,
         show=False,
     )
-    out = capsys.readouterr().out
-    assert "Peptide correlation summary" in out
-    assert "Per proteoform_id" in out
+    assert len(calls) == 1
+    assert calls[0][0] is scheme
+    assert calls[0][1] == ["a", "b", 10, 2]
+    strips = captured["kwargs"]["row_colors"]
+    palettes = {
+        "region": {"a": expected[0], "b": expected[1]},
+        "phase": {10: expected[2], 2: expected[3]},
+    }
+    for col, palette in palettes.items():
+        _assert_strip_colors(
+            strips[col], adata.var.loc[strips.index, col], palette
+        )
+
+
+@pytest.mark.parametrize("margin", ["region", ["region"]])
+def test_single_strip_accepts_flat_color_list(margin, monkeypatch):
+    adata = _color_scheme_adata()
+    captured = _spy_clustermap(monkeypatch)
+    _function()(
+        adata,
+        protein="P1",
+        margin_color=margin,
+        color_scheme=["red", "blue"],
+        show=False,
+    )
+    strip = captured["kwargs"]["row_colors"]["region"]
+    _assert_strip_colors(
+        strip, adata.var.loc[strip.index, "region"], {"a": "red", "b": "blue"}
+    )
+
+
+@pytest.mark.parametrize("schemes", [["Blues"], ["Blues", "Reds", "Greens"]])
+def test_per_strip_palette_count_must_match_margins(schemes):
+    adata = _color_scheme_adata()
+    before = set(plt.get_fignums())
+    with pytest.raises(ValueError, match="same length as `margin_color`"):
+        _function()(
+            adata,
+            protein="P1",
+            margin_color=["region", "phase"],
+            color_scheme=schemes,
+            show=False,
+        )
+    assert set(plt.get_fignums()) == before
 
 
 # -- Margin annotations: default, colours, legends ----------------------
@@ -442,21 +589,18 @@ def test_multiple_margins_use_disjoint_colours(monkeypatch):
 def test_one_legend_per_margin_column():
     adata = _copf_adata()
     margins = ["cluster_id", "proteoform_id"]
-    axm = _function()(
-        adata, protein="P1", margin_color=margins, show=False, ax=True
-    )
+    axm = _plot_heatmap(adata, protein="P1", margin_color=margins, show=False)
     titles = [leg.get_title().get_text() for leg in axm.figure.legends]
     assert titles == margins
 
 
 def test_legends_clear_of_all_axes_text_and_inside_figure():
     adata = _copf_adata()
-    axm = _function()(
+    axm = _plot_heatmap(
         adata,
         protein="P1",
         margin_color=["cluster_id", "proteoform_id", "protein_id"],
         show=False,
-        ax=True,
     )
     fig = axm.figure
     renderer = fig.canvas.get_renderer()
@@ -484,7 +628,7 @@ def test_peptides_filtered_after_correlations_warn_and_are_dropped():
     adata = _copf_adata()
     sub = adata[:, adata.var_names != "pep0"].copy()
     with pytest.warns(UserWarning, match="no longer in adata.var"):
-        axm = _function()(sub, protein="P1", show=False, ax=True)
+        axm = _plot_heatmap(sub, protein="P1", show=False)
     assert "pep0" not in tick_labels(axm)
     assert len(tick_labels(axm)) == 4
 
@@ -528,24 +672,6 @@ def test_linkage_of_wrong_size_raises():
         _function()(adata, protein="P1", linkage=Z, show=False)
 
 
-def test_print_stats_groups_follow_legend_order(capsys):
-    adata = _copf_adata()
-    adata.var["grp"] = [2, 10, 2, 10, 2, 2, 10, 2]
-    axm = _function()(
-        adata,
-        protein="P1",
-        margin_color="grp",
-        print_stats=True,
-        show=False,
-        ax=True,
-    )
-    legend = [t.get_text() for t in axm.figure.legends[0].get_texts()]
-    table = capsys.readouterr().out.split("Per grp")[1].splitlines()[2:]
-    stats = [line.split()[0] for line in table if line.strip()]
-    assert legend == ["10", "2"]
-    assert stats == legend
-
-
 def test_all_nan_correlations_draw_without_runtime_warnings():
     adata = _copf_adata()
     corrs = adata.uns["pairwise_peptide_correlations"].copy()
@@ -558,7 +684,6 @@ def test_all_nan_correlations_draw_without_runtime_warnings():
             protein="P1",
             margin_color="protein_id",
             cluster=False,
-            print_stats=True,
             show=False,
         )
     assert not [w for w in caught if issubclass(w.category, RuntimeWarning)]
@@ -589,13 +714,12 @@ def test_vector_backends(backend, tmp_path):
 def test_tall_legends_stay_inside_figure():
     adata = _copf_adata()
     adata.var["per_peptide"] = adata.var["peptide_id"].astype(str)
-    axm = _function()(
+    axm = _plot_heatmap(
         adata,
         protein="P1",
         margin_color=["per_peptide", "proteoform_id", "protein_id"],
         figsize=(2.5, 1.2),
         show=False,
-        ax=True,
     )
     fig = axm.figure
     renderer = fig.canvas.get_renderer()
@@ -608,29 +732,49 @@ def test_tall_legends_stay_inside_figure():
 # -- Checklist scenarios: missing data, degenerate inputs, rendering -----
 
 
-def test_all_na_peptide_shown_as_empty_row_and_column():
+def test_all_na_peptide_is_omitted_without_warning():
     adata = _copf_adata()
-    adata.X[:, 0] = np.nan  # pep0 never measured
+    X = np.asarray(adata.X).copy()
+    X[:, 0] = np.nan  # pep0 never measured
+    adata.X = X
     _recorrelate(adata)
-    with pytest.warns(UserWarning) as record:
-        axm = _function()(adata, protein="P1", show=False, ax=True)
-    messages = " | ".join(str(w.message) for w in record)
-    assert "no correlations" in messages
-    assert "treated as r = 0" in messages
+    with warnings.catch_warnings(record=True) as record:
+        warnings.simplefilter("always")
+        axm = _plot_heatmap(adata, protein="P1", show=False)
+    assert not record
+    rows, cols, values = _rendered(axm)
+    assert "pep0" not in rows
+    assert "pep0" not in cols
+    assert values.shape == (4, 4)
+    assert not np.isnan(values).any()
+
+
+def test_peptide_removed_from_correlations_stays_omitted():
+    adata = _copf_adata()
+    corrs = adata.uns["pairwise_peptide_correlations"]
+    keep = (corrs["pepA"] != "pep0") & (corrs["pepB"] != "pep0")
+    adata.uns["pairwise_peptide_correlations"] = corrs.loc[keep].copy()
+    assert "pep0" in adata.var["peptide_id"].values
+    with warnings.catch_warnings(record=True) as record:
+        warnings.simplefilter("always")
+        axm = _plot_heatmap(adata, protein="P1", show=False)
+    assert not record
     rows, cols, values = _rendered(axm)
-    assert "pep0" in rows
-    assert np.isnan(values[rows.index("pep0")]).all()
-    assert np.isnan(values[:, cols.index("pep0")]).all()
+    assert "pep0" not in rows
+    assert "pep0" not in cols
+    assert values.shape == (4, 4)
 
 
 def test_constant_peptide_draws_with_default_clustering():
     adata = _copf_adata()
-    adata.X[:, 1] = 5.0  # zero variance -> NaN PCC
+    X = np.asarray(adata.X).copy()
+    X[:, 1] = 5.0  # zero variance -> NaN PCC
+    adata.X = X
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")  # scipy ConstantInputWarning
         _recorrelate(adata)
     with pytest.warns(UserWarning, match="NaN"):
-        axm = _function()(adata, protein="P1", show=False, ax=True)
+        axm = _plot_heatmap(adata, protein="P1", show=False)
     rows, _, values = _rendered(axm)
     assert np.isnan(values[rows.index("pep1")]).sum() == len(rows) - 1
 
@@ -638,14 +782,14 @@ def test_constant_peptide_draws_with_default_clustering():
 def test_nan_cells_hatched_and_in_legend():
     adata = _with_nan_correlation(_copf_adata())
     with pytest.warns(UserWarning):
-        axm = _function()(adata, protein="P1", show=False, ax=True)
+        axm = _plot_heatmap(adata, protein="P1", show=False)
     assert axm.patch.get_hatch()
     titles = [leg.get_title().get_text() for leg in axm.figure.legends]
     assert "correlation" in titles
 
 
 def test_complete_matrix_has_no_nan_legend():
-    axm = _function()(_copf_adata(), protein="P1", show=False, ax=True)
+    axm = _plot_heatmap(_copf_adata(), protein="P1", show=False)
     titles = [leg.get_title().get_text() for leg in axm.figure.legends]
     assert "correlation" not in titles
     assert not axm.patch.get_hatch()
@@ -660,7 +804,7 @@ def test_fewer_than_three_samples_warn():
 def test_two_peptide_protein_renders_2x2():
     adata = _copf_adata()
     adata = _recorrelate(adata[:, ["pep0", "pep1", "pep5", "pep6"]].copy())
-    axm = _function()(adata, protein="P1", show=False, ax=True)
+    axm = _plot_heatmap(adata, protein="P1", show=False)
     rows, cols, values = _rendered(axm)
     assert values.shape == (2, 2)
     assert np.allclose(np.diag(values), 1.0)
@@ -669,10 +813,12 @@ def test_two_peptide_protein_renders_2x2():
 
 def test_identical_peptides_render_r_of_one():
     adata = _copf_adata()
-    adata.X[:, 1] = adata.X[:, 0]
+    X = np.asarray(adata.X).copy()
+    X[:, 1] = X[:, 0]
+    adata.X = X
     _recorrelate(adata)
     rows, cols, values = _rendered(
-        _function()(adata, protein="P1", show=False, ax=True)
+        _plot_heatmap(adata, protein="P1", show=False)
     )
     assert values[rows.index("pep0"), cols.index("pep1")] == pytest.approx(1)
 
@@ -684,9 +830,7 @@ def test_rendered_cells_match_correlations(cluster):
     truth = {}
     for a, b, r in long[["pepA", "pepB", "PCC"]].itertuples(index=False):
         truth[(a, b)] = truth[(b, a)] = r
-    axm = _function()(
-        adata, protein="P1", cluster=cluster, show=False, ax=True
-    )
+    axm = _plot_heatmap(adata, protein="P1", cluster=cluster, show=False)
     rows, cols, values = _rendered(axm)
     assert rows == cols
     assert np.allclose(np.diag(values), 1.0)
@@ -709,8 +853,8 @@ def test_default_labels_do_not_overlap_for_many_peptides():
         ),
     )
     _recorrelate(adata)
-    axm = _function()(
-        adata, protein="P9", margin_color="protein_id", show=False, ax=True
+    axm = _plot_heatmap(
+        adata, protein="P9", margin_color="protein_id", show=False
     )
     renderer = axm.figure.canvas.get_renderer()
     boxes = sorted(
diff --git a/tests/pp/test_summarize_peptides_by_neighbourhood_union.py b/tests/pp/test_summarize_peptides_by_neighbourhood_union.py
index f16e647..fa01858 100644
--- a/tests/pp/test_summarize_peptides_by_neighbourhood_union.py
+++ b/tests/pp/test_summarize_peptides_by_neighbourhood_union.py
@@ -50,7 +50,6 @@
 )
 from proteopy.utils.anndata import check_proteodata
 
-
 # -- Synthetic reference protein --------------------------------------
 #
 # Twenty distinct residues, so every substring occurs exactly once and
@@ -1889,6 +1888,23 @@ def test_multi_mapped_peptide_raises(self):
         with pytest.raises(ValueError, match="mapping to exactly one protein"):
             summarize(adata, _FASTA, inplace=False)
 
+    def test_multi_mapped_peptide_with_duplicate_names_raises(self):
+        """Duplicate peptide rows fail index uniqueness validation first."""
+        var = pd.DataFrame(
+            {
+                "peptide_id": ["ACDEF", "ACDEF"],
+                "protein_id": ["P1", "P2"],
+            },
+            index=["ACDEF", "ACDEF"],
+        )
+        obs = pd.DataFrame({"sample_id": ["s1"]}, index=["s1"])
+        with pytest.warns(UserWarning, match="Variable names are not unique"):
+            adata = AnnData(X=np.array([[1.0, 2.0]]), obs=obs, var=var)
+        with pytest.raises(
+            ValueError, match="Duplicate names detected in var"
+        ):
+            summarize(adata, _FASTA, inplace=False)
+
     def test_obs_is_preserved(self):
         adata = make_adata(
             ["ACDEF", "EFGHI"],

From 072b0a72c99865f4797563b52498ec092cbcccce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Levente=20Temesv=C3=A1ri-Nagy?=
 <147416790+leventetn@users.noreply.github.com>
Date: Thu, 8 Oct 2026 00:00:50 +0200
Subject: [PATCH 07/10] fix: seaborn mathplotlib deprecated warning silenced

---
 .../pl/test_pairwise_peptide_correlations_heatmap.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tests/pl/test_pairwise_peptide_correlations_heatmap.py b/tests/pl/test_pairwise_peptide_correlations_heatmap.py
index 4a87b56..f6b4800 100644
--- a/tests/pl/test_pairwise_peptide_correlations_heatmap.py
+++ b/tests/pl/test_pairwise_peptide_correlations_heatmap.py
@@ -740,6 +740,12 @@ def test_all_na_peptide_is_omitted_without_warning():
     _recorrelate(adata)
     with warnings.catch_warnings(record=True) as record:
         warnings.simplefilter("always")
+        warnings.filterwarnings(
+            "ignore",
+            message=r"The set_bad function will be deprecated\b",
+            category=PendingDeprecationWarning,
+            module=r"seaborn\.matrix$",
+        )
         axm = _plot_heatmap(adata, protein="P1", show=False)
     assert not record
     rows, cols, values = _rendered(axm)
@@ -757,6 +763,12 @@ def test_peptide_removed_from_correlations_stays_omitted():
     assert "pep0" in adata.var["peptide_id"].values
     with warnings.catch_warnings(record=True) as record:
         warnings.simplefilter("always")
+        warnings.filterwarnings(
+            "ignore",
+            message=r"The set_bad function will be deprecated\b",
+            category=PendingDeprecationWarning,
+            module=r"seaborn\.matrix$",
+        )
         axm = _plot_heatmap(adata, protein="P1", show=False)
     assert not record
     rows, cols, values = _rendered(axm)

From 753cc05749471b872399a2311ab84f3f6fa4d480 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Levente=20Temesv=C3=A1ri-Nagy?=
 <147416790+leventetn@users.noreply.github.com>
Date: Thu, 8 Oct 2026 10:37:40 +0200
Subject: [PATCH 08/10] fix: margin annotation

Updated copf.py:
- Added half-cell padding between the matrix and annotation bars.
- Kept annotation names only on the row bars.
---
 proteopy/pl/copf.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/proteopy/pl/copf.py b/proteopy/pl/copf.py
index c91e50f..b127002 100644
--- a/proteopy/pl/copf.py
+++ b/proteopy/pl/copf.py
@@ -739,6 +739,8 @@ def pairwise_peptide_correlations_heatmap(
         clustermap_kwargs["col_cluster"] = False
 
     g = sns.clustermap(corr_df, **clustermap_kwargs)
+    _space_heatmap_margins(g, n)
+    g.ax_row_colors.set_xticks([])
 
     # NaN cells are masked by seaborn; a hatched background keeps them
     # distinct from every colour on the correlation scale
@@ -768,6 +770,25 @@ def pairwise_peptide_correlations_heatmap(
     return None
 
 
+def _space_heatmap_margins(g, n_peptides: int) -> None:
+    """Reserve half a cell between the heatmap and annotation strips."""
+    pos = g.ax_heatmap.get_position()
+    row_pos = g.ax_row_colors.get_position()
+    col_pos = g.ax_col_colors.get_position()
+    width = (pos.x1 - row_pos.x1) / (1 + 0.5 / n_peptides)
+    height = (col_pos.y0 - pos.y0) / (1 + 0.5 / n_peptides)
+    left = pos.x1 - width
+    g.ax_heatmap.set_position([left, pos.y0, width, height])
+
+    # Keep annotation cells and dendrogram leaves aligned with peptides.
+    for ax in (g.ax_col_colors, g.ax_col_dendrogram):
+        annot_pos = ax.get_position()
+        ax.set_position([left, annot_pos.y0, width, annot_pos.height])
+    for ax in (g.ax_row_colors, g.ax_row_dendrogram):
+        annot_pos = ax.get_position()
+        ax.set_position([annot_pos.x0, pos.y0, annot_pos.width, height])
+
+
 def _validate_heatmap_inputs(
     adata: ad.AnnData,
     protein: str,

From 3db9057ef4c7ad4141f7e23ca5be75919e987580 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Levente=20Temesv=C3=A1ri-Nagy?=
 <147416790+leventetn@users.noreply.github.com>
Date: Fri, 9 Oct 2026 09:24:59 +0200
Subject: [PATCH 09/10] added pairwise_peptide_correlations_heatmap functionto
 sphinx api

---
 docs/sphinx/source/api/pl.rst | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/sphinx/source/api/pl.rst b/docs/sphinx/source/api/pl.rst
index 34a7bd3..2d45192 100644
--- a/docs/sphinx/source/api/pl.rst
+++ b/docs/sphinx/source/api/pl.rst
@@ -52,6 +52,7 @@ control, exploratory analysis, and statistical results.
    :nosignatures:
 
    proteopy.pl.proteoform_scores
+   proteopy.pl.pairwise_peptide_correlations_heatmap
    proteopy.pl.peptides_on_sequence
    proteopy.pl.peptides_on_prot_sequence
 

From 64fc27808a7fef368d5a170509e3739167ca685a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Levente=20Temesv=C3=A1ri-Nagy?=
 <147416790+leventetn@users.noreply.github.com>
Date: Fri, 9 Oct 2026 10:27:25 +0200
Subject: [PATCH 10/10] fix: added example for
 pairwise_peptide_correlations_heatmap

---
 proteopy/pl/copf.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/proteopy/pl/copf.py b/proteopy/pl/copf.py
index b127002..3592c4a 100644
--- a/proteopy/pl/copf.py
+++ b/proteopy/pl/copf.py
@@ -533,12 +533,20 @@ def pairwise_peptide_correlations_heatmap(
     ...     adata, protein="P1", show=False,
     ... )
 
-    Annotate by proteoform and draw two strips:
-
+    Compare proteoforms with illustrative peptide membrane positions.
+    The categorical annotation controls the membrane-position legend
+    order; each strip uses its own palette:
+
+    >>> adata.var["membrano_pos"] = pd.Categorical(
+    ...     ["trans", "cytosol", "extra", "trans"],
+    ...     categories=["trans", "cytosol", "extra"],
+    ...     ordered=True,
+    ... )
     >>> pr.pl.pairwise_peptide_correlations_heatmap(
     ...     adata,
     ...     protein="P1",
-    ...     margin_color=["cluster_id", "proteoform_id"],
+    ...     margin_color=["proteoform_id", "membrano_pos"],
+    ...     color_scheme=["Set2", "Dark2"],
     ...     show=False,
     ... )
     """