smftools 0.2.3__py3-none-any.whl → 0.2.5__py3-none-any.whl

smftools/plotting/general_plotting.py

@@ -1,22 +1,87 @@
 from __future__ import annotations
 
-import numpy as np
-import seaborn as sns
-import matplotlib.pyplot as plt
-import scipy.cluster.hierarchy as sch
-import matplotlib.gridspec as gridspec
-import os
 import math
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
+
+import matplotlib.gridspec as gridspec
+import matplotlib.pyplot as plt
+import numpy as np
 import pandas as pd
+import scipy.cluster.hierarchy as sch
+import seaborn as sns
+
+
+def _fixed_tick_positions(n_positions: int, n_ticks: int) -> np.ndarray:
+    """
+    Return indices for ~n_ticks evenly spaced labels across [0, n_positions-1].
+    Always includes 0 and n_positions-1 when possible.
+    """
+    n_ticks = int(max(2, n_ticks))
+    if n_positions <= n_ticks:
+        return np.arange(n_positions)
+
+    # linspace gives fixed count
+    pos = np.linspace(0, n_positions - 1, n_ticks)
+    return np.unique(np.round(pos).astype(int))
+
+
+def _select_labels(subset, sites: np.ndarray, reference: str, index_col_suffix: str | None):
+    """
+    Select tick labels for the heatmap axis.
+
+    Parameters
+    ----------
+    subset : AnnData view
+        The per-bin subset of the AnnData.
+    sites : np.ndarray[int]
+        Indices of the subset.var positions to annotate.
+    reference : str
+        Reference name (e.g., '6B6_top').
+    index_col_suffix : None or str
+        If None → use subset.var_names
+        Else     → use subset.var[f"{reference}_{index_col_suffix}"]
+
+    Returns
+    -------
+    np.ndarray[str]
+        The labels to use for tick positions.
+    """
+    if sites.size == 0:
+        return np.array([])
+
+    # Default behavior: use var_names
+    if index_col_suffix is None:
+        return subset.var_names[sites].astype(str)
+
+    # Otherwise: use a computed column adata.var[f"{reference}_{suffix}"]
+    colname = f"{reference}_{index_col_suffix}"
+
+    if colname not in subset.var:
+        raise KeyError(
+            f"index_col_suffix='{index_col_suffix}' requires var column '{colname}', "
+            f"but it is not present in adata.var."
+        )
+
+    labels = subset.var[colname].astype(str).values
+    return labels[sites]
 
-from typing import Optional, Mapping, Sequence, Any, Dict, List
-from pathlib import Path
 
 def normalized_mean(matrix: np.ndarray) -> np.ndarray:
+    """Compute normalized column means for a matrix.
+
+    Args:
+        matrix: Input matrix.
+
+    Returns:
+        1D array of normalized means.
+    """
     mean = np.nanmean(matrix, axis=0)
     denom = (mean.max() - mean.min()) + 1e-9
     return (mean - mean.min()) / denom
 
+
 def methylation_fraction(matrix: np.ndarray) -> np.ndarray:
     """
     Fraction methylated per column.
@@ -31,14 +96,20 @@ def methylation_fraction(matrix: np.ndarray) -> np.ndarray:
     valid = valid_mask.sum(axis=0)
 
     return np.divide(
-        methylated, valid,
-        out=np.zeros_like(methylated, dtype=float),
-        where=valid != 0
+        methylated, valid, out=np.zeros_like(methylated, dtype=float), where=valid != 0
     )
 
+
 def clean_barplot(ax, mean_values, title):
+    """Format a barplot with consistent axes and labels.
+
+    Args:
+        ax: Matplotlib axes.
+        mean_values: Values to plot.
+        title: Plot title.
+    """
     x = np.arange(len(mean_values))
-    ax.bar(x, mean_values, color="gray", width=1.0, align='edge')
+    ax.bar(x, mean_values, color="gray", width=1.0, align="edge")
     ax.set_xlim(0, len(mean_values))
     ax.set_ylim(0, 1)
     ax.set_yticks([0.0, 0.5, 1.0])
@@ -47,9 +118,10 @@ def clean_barplot(ax, mean_values, title):
 
     # Hide all spines except left
     for spine_name, spine in ax.spines.items():
-        spine.set_visible(spine_name == 'left')
+        spine.set_visible(spine_name == "left")
+
+    ax.tick_params(axis="x", which="both", bottom=False, top=False, labelbottom=False)
 
-    ax.tick_params(axis='x', which='both', bottom=False, top=False, labelbottom=False)
 
 # def combined_hmm_raw_clustermap(
 #     adata,
@@ -92,7 +164,7 @@ def clean_barplot(ax, mean_values, title):
 #                     (adata.obs['read_length'] >= min_length) &
 #                     (adata.obs['mapped_length_to_reference_length_ratio'] > min_mapped_length_to_reference_length_ratio)
 #                 ]
-                
+
 #                 mask = subset.var[f"{ref}_valid_fraction"].astype(float) > float(min_position_valid_fraction)
 #                 subset = subset[:, mask]
 
@@ -204,7 +276,7 @@ def clean_barplot(ax, mean_values, title):
 #                         clean_barplot(axes_bar[1], mean_gpc, f"GpC Accessibility Signal")
 #                         clean_barplot(axes_bar[2], mean_cpg, f"CpG Accessibility Signal")
 #                         clean_barplot(axes_bar[3], mean_any_c, f"Any C Accessibility Signal")
-                        
+
 #                         hmm_labels = subset.var_names.astype(int)
 #                         hmm_label_spacing = 150
 #                         sns.heatmap(hmm_matrix, cmap=cmap_hmm, ax=axes_heat[0], xticklabels=hmm_labels[::hmm_label_spacing], yticklabels=False, cbar=False)
@@ -258,7 +330,7 @@ def clean_barplot(ax, mean_values, title):
 #                             "bin_boundaries": bin_boundaries,
 #                             "percentages": percentages
 #                         })
-                        
+
 #                         #adata.uns['clustermap_results'] = results
 
 #             except Exception as e:
@@ -271,83 +343,131 @@ def combined_hmm_raw_clustermap(
     adata,
     sample_col: str = "Sample_Names",
     reference_col: str = "Reference_strand",
-
     hmm_feature_layer: str = "hmm_combined",
-
     layer_gpc: str = "nan0_0minus1",
     layer_cpg: str = "nan0_0minus1",
-    layer_any_c: str = "nan0_0minus1",
+    layer_c: str = "nan0_0minus1",
     layer_a: str = "nan0_0minus1",
-
     cmap_hmm: str = "tab10",
     cmap_gpc: str = "coolwarm",
     cmap_cpg: str = "viridis",
-    cmap_any_c: str = "coolwarm",
+    cmap_c: str = "coolwarm",
     cmap_a: str = "coolwarm",
-
     min_quality: int = 20,
     min_length: int = 200,
     min_mapped_length_to_reference_length_ratio: float = 0.8,
     min_position_valid_fraction: float = 0.5,
-
+    demux_types: Sequence[str] = ("single", "double", "already"),
+    sample_mapping: Optional[Mapping[str, str]] = None,
     save_path: str | Path | None = None,
     normalize_hmm: bool = False,
-
     sort_by: str = "gpc",
     bins: Optional[Dict[str, Any]] = None,
-
     deaminase: bool = False,
     min_signal: float = 0.0,
-
     # ---- fixed tick label controls (counts, not spacing)
     n_xticks_hmm: int = 10,
     n_xticks_any_c: int = 8,
     n_xticks_gpc: int = 8,
     n_xticks_cpg: int = 8,
     n_xticks_a: int = 8,
+    index_col_suffix: str | None = None,
 ):
     """
     Makes a multi-panel clustermap per (sample, reference):
-      HMM panel (always) + optional raw panels for any_C, GpC, CpG, and A sites.
+      HMM panel (always) + optional raw panels for C, GpC, CpG, and A sites.
 
     Panels are added only if the corresponding site mask exists AND has >0 sites.
 
     sort_by options:
-      'gpc', 'cpg', 'any_c', 'any_a', 'gpc_cpg', 'none', or 'obs:<col>'
+      'gpc', 'cpg', 'c', 'a', 'gpc_cpg', 'none', 'hmm', or 'obs:<col>'
     """
+
     def pick_xticks(labels: np.ndarray, n_ticks: int):
+        """Pick tick indices/labels from an array."""
         if labels.size == 0:
             return [], []
         idx = np.linspace(0, len(labels) - 1, n_ticks).round().astype(int)
         idx = np.unique(idx)
         return idx.tolist(), labels[idx].tolist()
-    
+
+    # Helper: build a True mask if filter is inactive or column missing
+    def _mask_or_true(series_name: str, predicate):
+        """Return a mask from predicate or an all-True mask."""
+        if series_name not in adata.obs:
+            return pd.Series(True, index=adata.obs.index)
+        s = adata.obs[series_name]
+        try:
+            return predicate(s)
+        except Exception:
+            # Fallback: all True if bad dtype / predicate failure
+            return pd.Series(True, index=adata.obs.index)
+
     results = []
     signal_type = "deamination" if deaminase else "methylation"
 
     for ref in adata.obs[reference_col].cat.categories:
         for sample in adata.obs[sample_col].cat.categories:
+            # Optionally remap sample label for display
+            display_sample = sample_mapping.get(sample, sample) if sample_mapping else sample
+            # Row-level masks (obs)
+            qmask = _mask_or_true(
+                "read_quality",
+                (lambda s: s >= float(min_quality))
+                if (min_quality is not None)
+                else (lambda s: pd.Series(True, index=s.index)),
+            )
+            lm_mask = _mask_or_true(
+                "mapped_length",
+                (lambda s: s >= float(min_length))
+                if (min_length is not None)
+                else (lambda s: pd.Series(True, index=s.index)),
+            )
+            lrr_mask = _mask_or_true(
+                "mapped_length_to_reference_length_ratio",
+                (lambda s: s >= float(min_mapped_length_to_reference_length_ratio))
+                if (min_mapped_length_to_reference_length_ratio is not None)
+                else (lambda s: pd.Series(True, index=s.index)),
+            )
+
+            demux_mask = _mask_or_true(
+                "demux_type",
+                (lambda s: s.astype("string").isin(list(demux_types)))
+                if (demux_types is not None)
+                else (lambda s: pd.Series(True, index=s.index)),
+            )
+
+            ref_mask = adata.obs[reference_col] == ref
+            sample_mask = adata.obs[sample_col] == sample
+
+            row_mask = ref_mask & sample_mask & qmask & lm_mask & lrr_mask & demux_mask
+
+            if not bool(row_mask.any()):
+                print(
+                    f"No reads for {display_sample} - {ref} after read quality and length filtering"
+                )
+                continue
 
             try:
                 # ---- subset reads ----
-                subset = adata[
-                    (adata.obs[reference_col] == ref) &
-                    (adata.obs[sample_col] == sample) &
-                    (adata.obs["read_quality"] >= min_quality) &
-                    (adata.obs["read_length"] >= min_length) &
-                    (
-                        adata.obs["mapped_length_to_reference_length_ratio"]
-                        > min_mapped_length_to_reference_length_ratio
-                    )
-                ]
-
-                # ---- valid fraction filter ----
-                vf_key = f"{ref}_valid_fraction"
-                if vf_key in subset.var:
-                    mask = subset.var[vf_key].astype(float) > float(min_position_valid_fraction)
-                    subset = subset[:, mask]
+                subset = adata[row_mask, :].copy()
+
+                # Column-level mask (var)
+                if min_position_valid_fraction is not None:
+                    valid_key = f"{ref}_valid_fraction"
+                    if valid_key in subset.var:
+                        v = pd.to_numeric(subset.var[valid_key], errors="coerce").to_numpy()
+                        col_mask = np.asarray(v > float(min_position_valid_fraction), dtype=bool)
+                        if col_mask.any():
+                            subset = subset[:, col_mask].copy()
+                        else:
+                            print(
+                                f"No positions left after valid_fraction filter for {display_sample} - {ref}"
+                            )
+                            continue
 
                 if subset.shape[0] == 0:
+                    print(f"No reads left after filtering for {display_sample} - {ref}")
                     continue
 
                 # ---- bins ----
@@ -358,6 +478,7 @@ def combined_hmm_raw_clustermap(
 
                 # ---- site masks (robust) ----
                 def _sites(*keys):
+                    """Return indices for the first matching site key."""
                     for k in keys:
                         if k in subset.var:
                             return np.where(subset.var[k].values)[0]
@@ -368,13 +489,14 @@ def combined_hmm_raw_clustermap(
                 any_c_sites = _sites(f"{ref}_any_C_site", f"{ref}_C_site")
                 any_a_sites = _sites(f"{ref}_A_site", f"{ref}_any_A_site")
 
-                def _labels(sites):
-                    return subset.var_names[sites].astype(int) if sites.size else np.array([])
-
-                gpc_labels = _labels(gpc_sites)
-                cpg_labels = _labels(cpg_sites)
-                any_c_labels = _labels(any_c_sites)
-                any_a_labels = _labels(any_a_sites)
+                # ---- labels via _select_labels ----
+                # HMM uses *all* columns
+                hmm_sites = np.arange(subset.n_vars, dtype=int)
+                hmm_labels = _select_labels(subset, hmm_sites, ref, index_col_suffix)
+                gpc_labels = _select_labels(subset, gpc_sites, ref, index_col_suffix)
+                cpg_labels = _select_labels(subset, cpg_sites, ref, index_col_suffix)
+                any_c_labels = _select_labels(subset, any_c_sites, ref, index_col_suffix)
+                any_a_labels = _select_labels(subset, any_a_sites, ref, index_col_suffix)
 
                 # storage
                 stacked_hmm = []
@@ -411,11 +533,11 @@ def combined_hmm_raw_clustermap(
                         linkage = sch.linkage(sb[:, cpg_sites].layers[layer_cpg], method="ward")
                         order = sch.leaves_list(linkage)
 
-                    elif sort_by == "any_c" and any_c_sites.size:
-                        linkage = sch.linkage(sb[:, any_c_sites].layers[layer_any_c], method="ward")
+                    elif sort_by == "c" and any_c_sites.size:
+                        linkage = sch.linkage(sb[:, any_c_sites].layers[layer_c], method="ward")
                         order = sch.leaves_list(linkage)
 
-                    elif sort_by == "any_a" and any_a_sites.size:
+                    elif sort_by == "a" and any_a_sites.size:
                         linkage = sch.linkage(sb[:, any_a_sites].layers[layer_a], method="ward")
                         order = sch.leaves_list(linkage)
 
@@ -423,6 +545,12 @@ def combined_hmm_raw_clustermap(
                         linkage = sch.linkage(sb.layers[layer_gpc], method="ward")
                         order = sch.leaves_list(linkage)
 
+                    elif sort_by == "hmm" and hmm_sites.size:
+                        linkage = sch.linkage(
+                            sb[:, hmm_sites].layers[hmm_feature_layer], method="ward"
+                        )
+                        order = sch.leaves_list(linkage)
+
                     else:
                         order = np.arange(n)
 
@@ -431,7 +559,7 @@ def combined_hmm_raw_clustermap(
                     # ---- collect matrices ----
                     stacked_hmm.append(sb.layers[hmm_feature_layer])
                     if any_c_sites.size:
-                        stacked_any_c.append(sb[:, any_c_sites].layers[layer_any_c])
+                        stacked_any_c.append(sb[:, any_c_sites].layers[layer_c])
                     if gpc_sites.size:
                         stacked_gpc.append(sb[:, gpc_sites].layers[layer_gpc])
                     if cpg_sites.size:
@@ -446,46 +574,62 @@ def combined_hmm_raw_clustermap(
 
                 # ---------------- stack ----------------
                 hmm_matrix = np.vstack(stacked_hmm)
-                mean_hmm = normalized_mean(hmm_matrix) if normalize_hmm else np.nanmean(hmm_matrix, axis=0)
+                mean_hmm = (
+                    normalized_mean(hmm_matrix) if normalize_hmm else np.nanmean(hmm_matrix, axis=0)
+                )
 
                 panels = [
-                    ("HMM", hmm_matrix, subset.var_names.astype(int), cmap_hmm, mean_hmm, n_xticks_hmm),
+                    (
+                        f"HMM - {hmm_feature_layer}",
+                        hmm_matrix,
+                        hmm_labels,
+                        cmap_hmm,
+                        mean_hmm,
+                        n_xticks_hmm,
+                    ),
                 ]
 
                 if stacked_any_c:
                     m = np.vstack(stacked_any_c)
-                    panels.append(("any_C", m, any_c_labels, cmap_any_c, methylation_fraction(m), n_xticks_any_c))
+                    panels.append(
+                        ("C", m, any_c_labels, cmap_c, methylation_fraction(m), n_xticks_any_c)
+                    )
 
                 if stacked_gpc:
                     m = np.vstack(stacked_gpc)
-                    panels.append(("GpC", m, gpc_labels, cmap_gpc, methylation_fraction(m), n_xticks_gpc))
+                    panels.append(
+                        ("GpC", m, gpc_labels, cmap_gpc, methylation_fraction(m), n_xticks_gpc)
+                    )
 
                 if stacked_cpg:
                     m = np.vstack(stacked_cpg)
-                    panels.append(("CpG", m, cpg_labels, cmap_cpg, methylation_fraction(m), n_xticks_cpg))
+                    panels.append(
+                        ("CpG", m, cpg_labels, cmap_cpg, methylation_fraction(m), n_xticks_cpg)
+                    )
 
                 if stacked_any_a:
                     m = np.vstack(stacked_any_a)
-                    panels.append(("A", m, any_a_labels, cmap_a, methylation_fraction(m), n_xticks_a))
+                    panels.append(
+                        ("A", m, any_a_labels, cmap_a, methylation_fraction(m), n_xticks_a)
+                    )
 
                 # ---------------- plotting ----------------
                 n_panels = len(panels)
                 fig = plt.figure(figsize=(4.5 * n_panels, 10))
                 gs = gridspec.GridSpec(2, n_panels, height_ratios=[1, 6], hspace=0.01)
-                fig.suptitle(f"{sample} — {ref} — {total_reads} reads ({signal_type})",
-                             fontsize=14, y=0.98)
+                fig.suptitle(
+                    f"{sample} — {ref} — {total_reads} reads ({signal_type})", fontsize=14, y=0.98
+                )
 
                 axes_heat = [fig.add_subplot(gs[1, i]) for i in range(n_panels)]
                 axes_bar = [fig.add_subplot(gs[0, i], sharex=axes_heat[i]) for i in range(n_panels)]
 
                 for i, (name, matrix, labels, cmap, mean_vec, n_ticks) in enumerate(panels):
-
                     # ---- your clean barplot ----
                     clean_barplot(axes_bar[i], mean_vec, name)
 
                     # ---- heatmap ----
-                    sns.heatmap(matrix, cmap=cmap, ax=axes_heat[i],
-                                yticklabels=False, cbar=False)
+                    sns.heatmap(matrix, cmap=cmap, ax=axes_heat[i], yticklabels=False, cbar=False)
 
                     # ---- xticks ----
                     xtick_pos, xtick_labels = pick_xticks(np.asarray(labels), n_ticks)
@@ -509,6 +653,7 @@ def combined_hmm_raw_clustermap(
 
             except Exception:
                 import traceback
+
                 traceback.print_exc()
                 continue
 
@@ -628,7 +773,7 @@ def combined_hmm_raw_clustermap(
 #                             order = np.arange(num_reads)
 #                         elif sort_by == "any_a":
 #                             linkage = sch.linkage(subset_bin.layers[layer_a], method="ward")
-#                             order = sch.leaves_list(linkage)                            
+#                             order = sch.leaves_list(linkage)
 #                         else:
 #                             raise ValueError(f"Unsupported sort_by option: {sort_by}")
 
@@ -657,13 +802,13 @@ def combined_hmm_raw_clustermap(
 #                             order = np.arange(num_reads)
 #                         elif sort_by == "any_a":
 #                             linkage = sch.linkage(subset_bin.layers[layer_a], method="ward")
-#                             order = sch.leaves_list(linkage)                            
+#                             order = sch.leaves_list(linkage)
 #                         else:
 #                             raise ValueError(f"Unsupported sort_by option: {sort_by}")
-                        
+
 #                         stacked_any_a.append(subset_bin[order][:, any_a_sites].layers[layer_a])
-                        
-                        
+
+
 #                     row_labels.extend([bin_label] * num_reads)
 #                     bin_labels.append(f"{bin_label}: {num_reads} reads ({percent_reads:.1f}%)")
 #                     last_idx += num_reads
@@ -686,7 +831,7 @@ def combined_hmm_raw_clustermap(
 #                     if any_a_matrix.size > 0:
 #                         mean_any_a = methylation_fraction(any_a_matrix)
 #                         gs_dim += 1
-                    
+
 
 #                 fig = plt.figure(figsize=(18, 12))
 #                 gs = gridspec.GridSpec(2, gs_dim, height_ratios=[1, 6], hspace=0.01)
@@ -718,8 +863,8 @@ def combined_hmm_raw_clustermap(
 #                         sns.heatmap(cpg_matrix, cmap=cmap_cpg, ax=axes_heat[2], xticklabels=cpg_labels, yticklabels=False, cbar=False)
 #                         axes_heat[current_ax].set_xticklabels(cpg_labels, rotation=90, fontsize=10)
 #                         for boundary in bin_boundaries[:-1]:
-#                             axes_heat[current_ax].axhline(y=boundary, color="black", linewidth=2)  
-#                         current_ax +=1  
+#                             axes_heat[current_ax].axhline(y=boundary, color="black", linewidth=2)
+#                         current_ax +=1
 
 #                         results.append({
 #                             "sample": sample,
@@ -731,7 +876,7 @@ def combined_hmm_raw_clustermap(
 #                             "bin_labels": bin_labels,
 #                             "bin_boundaries": bin_boundaries,
 #                             "percentages": percentages
-#                         })                    
+#                         })
 
 #                 if stacked_any_a:
 #                     if any_a_matrix.size > 0:
@@ -751,7 +896,7 @@ def combined_hmm_raw_clustermap(
 #                             "bin_labels": bin_labels,
 #                             "bin_boundaries": bin_boundaries,
 #                             "percentages": percentages
-#                         }) 
+#                         })
 
 #                 plt.tight_layout()
 
@@ -769,7 +914,7 @@ def combined_hmm_raw_clustermap(
 #                 print(f"Summary for {sample} - {ref}:")
 #                 for bin_label, percent in percentages.items():
 #                     print(f"  - {bin_label}: {percent:.1f}%")
-                
+
 #                 adata.uns['clustermap_results'] = results
 
 #             except Exception as e:
@@ -777,52 +922,41 @@ def combined_hmm_raw_clustermap(
 #                 traceback.print_exc()
 #                 continue
 
-def _fixed_tick_positions(n_positions: int, n_ticks: int) -> np.ndarray:
-    """
-    Return indices for ~n_ticks evenly spaced labels across [0, n_positions-1].
-    Always includes 0 and n_positions-1 when possible.
-    """
-    n_ticks = int(max(2, n_ticks))
-    if n_positions <= n_ticks:
-        return np.arange(n_positions)
-
-    # linspace gives fixed count
-    pos = np.linspace(0, n_positions - 1, n_ticks)
-    return np.unique(np.round(pos).astype(int))
 
 def combined_raw_clustermap(
     adata,
     sample_col: str = "Sample_Names",
     reference_col: str = "Reference_strand",
     mod_target_bases: Sequence[str] = ("GpC", "CpG"),
-    layer_any_c: str = "nan0_0minus1",
+    layer_c: str = "nan0_0minus1",
     layer_gpc: str = "nan0_0minus1",
     layer_cpg: str = "nan0_0minus1",
     layer_a: str = "nan0_0minus1",
-    cmap_any_c: str = "coolwarm",
+    cmap_c: str = "coolwarm",
     cmap_gpc: str = "coolwarm",
     cmap_cpg: str = "viridis",
     cmap_a: str = "coolwarm",
-    min_quality: float = 20,
-    min_length: int = 200,
-    min_mapped_length_to_reference_length_ratio: float = 0.8,
-    min_position_valid_fraction: float = 0.5,
+    min_quality: float | None = 20,
+    min_length: int | None = 200,
+    min_mapped_length_to_reference_length_ratio: float | None = 0,
+    min_position_valid_fraction: float | None = 0,
+    demux_types: Sequence[str] = ("single", "double", "already"),
     sample_mapping: Optional[Mapping[str, str]] = None,
     save_path: str | Path | None = None,
-    sort_by: str = "gpc",  # 'gpc','cpg','any_c','gpc_cpg','any_a','none','obs:<col>'
+    sort_by: str = "gpc",  # 'gpc','cpg','c','gpc_cpg','a','none','obs:<col>'
     bins: Optional[Dict[str, Any]] = None,
     deaminase: bool = False,
     min_signal: float = 0,
-    # NEW tick controls
     n_xticks_any_c: int = 10,
     n_xticks_gpc: int = 10,
     n_xticks_cpg: int = 10,
     n_xticks_any_a: int = 10,
     xtick_rotation: int = 90,
     xtick_fontsize: int = 9,
+    index_col_suffix: str | None = None,
 ):
     """
-    Plot stacked heatmaps + per-position mean barplots for any_C, GpC, CpG, and optional A.
+    Plot stacked heatmaps + per-position mean barplots for C, GpC, CpG, and optional A.
 
     Key fixes vs old version:
       - order computed ONCE per bin, applied to all matrices
@@ -838,6 +972,18 @@ def combined_raw_clustermap(
         One entry per (sample, ref) plot with matrices + bin metadata.
     """
 
+    # Helper: build a True mask if filter is inactive or column missing
+    def _mask_or_true(series_name: str, predicate):
+        """Return a mask from predicate or an all-True mask."""
+        if series_name not in adata.obs:
+            return pd.Series(True, index=adata.obs.index)
+        s = adata.obs[series_name]
+        try:
+            return predicate(s)
+        except Exception:
+            # Fallback: all True if bad dtype / predicate failure
+            return pd.Series(True, index=adata.obs.index)
+
     results: List[Dict[str, Any]] = []
     save_path = Path(save_path) if save_path is not None else None
     if save_path is not None:
@@ -856,24 +1002,63 @@ def combined_raw_clustermap(
 
     for ref in adata.obs[reference_col].cat.categories:
         for sample in adata.obs[sample_col].cat.categories:
-
             # Optionally remap sample label for display
             display_sample = sample_mapping.get(sample, sample) if sample_mapping else sample
 
-            try:
-                subset = adata[
-                    (adata.obs[reference_col] == ref) &
-                    (adata.obs[sample_col] == sample) &
-                    (adata.obs["read_quality"] >= min_quality) &
-                    (adata.obs["mapped_length"] >= min_length) &
-                    (adata.obs["mapped_length_to_reference_length_ratio"] >= min_mapped_length_to_reference_length_ratio)
-                ]
+            # Row-level masks (obs)
+            qmask = _mask_or_true(
+                "read_quality",
+                (lambda s: s >= float(min_quality))
+                if (min_quality is not None)
+                else (lambda s: pd.Series(True, index=s.index)),
+            )
+            lm_mask = _mask_or_true(
+                "mapped_length",
+                (lambda s: s >= float(min_length))
+                if (min_length is not None)
+                else (lambda s: pd.Series(True, index=s.index)),
+            )
+            lrr_mask = _mask_or_true(
+                "mapped_length_to_reference_length_ratio",
+                (lambda s: s >= float(min_mapped_length_to_reference_length_ratio))
+                if (min_mapped_length_to_reference_length_ratio is not None)
+                else (lambda s: pd.Series(True, index=s.index)),
+            )
+
+            demux_mask = _mask_or_true(
+                "demux_type",
+                (lambda s: s.astype("string").isin(list(demux_types)))
+                if (demux_types is not None)
+                else (lambda s: pd.Series(True, index=s.index)),
+            )
+
+            ref_mask = adata.obs[reference_col] == ref
+            sample_mask = adata.obs[sample_col] == sample
+
+            row_mask = ref_mask & sample_mask & qmask & lm_mask & lrr_mask & demux_mask
+
+            if not bool(row_mask.any()):
+                print(
+                    f"No reads for {display_sample} - {ref} after read quality and length filtering"
+                )
+                continue
 
-                # position-level mask
-                valid_key = f"{ref}_valid_fraction"
-                if valid_key in subset.var:
-                    mask = subset.var[valid_key].astype(float).values > float(min_position_valid_fraction)
-                    subset = subset[:, mask]
+            try:
+                subset = adata[row_mask, :].copy()
+
+                # Column-level mask (var)
+                if min_position_valid_fraction is not None:
+                    valid_key = f"{ref}_valid_fraction"
+                    if valid_key in subset.var:
+                        v = pd.to_numeric(subset.var[valid_key], errors="coerce").to_numpy()
+                        col_mask = np.asarray(v > float(min_position_valid_fraction), dtype=bool)
+                        if col_mask.any():
+                            subset = subset[:, col_mask].copy()
+                        else:
+                            print(
+                                f"No positions left after valid_fraction filter for {display_sample} - {ref}"
+                            )
+                            continue
 
                 if subset.shape[0] == 0:
                     print(f"No reads left after filtering for {display_sample} - {ref}")
@@ -893,19 +1078,19 @@ def combined_raw_clustermap(
 
                 if include_any_c:
                     any_c_sites = np.where(subset.var.get(f"{ref}_C_site", False).values)[0]
-                    gpc_sites   = np.where(subset.var.get(f"{ref}_GpC_site", False).values)[0]
-                    cpg_sites   = np.where(subset.var.get(f"{ref}_CpG_site", False).values)[0]
+                    gpc_sites = np.where(subset.var.get(f"{ref}_GpC_site", False).values)[0]
+                    cpg_sites = np.where(subset.var.get(f"{ref}_CpG_site", False).values)[0]
 
                     num_any_c, num_gpc, num_cpg = len(any_c_sites), len(gpc_sites), len(cpg_sites)
 
-                    any_c_labels = subset.var_names[any_c_sites].astype(str)
-                    gpc_labels   = subset.var_names[gpc_sites].astype(str)
-                    cpg_labels   = subset.var_names[cpg_sites].astype(str)
+                    any_c_labels = _select_labels(subset, any_c_sites, ref, index_col_suffix)
+                    gpc_labels = _select_labels(subset, gpc_sites, ref, index_col_suffix)
+                    cpg_labels = _select_labels(subset, cpg_sites, ref, index_col_suffix)
 
                 if include_any_a:
                     any_a_sites = np.where(subset.var.get(f"{ref}_A_site", False).values)[0]
                     num_any_a = len(any_a_sites)
-                    any_a_labels = subset.var_names[any_a_sites].astype(str)
+                    any_a_labels = _select_labels(subset, any_a_sites, ref, index_col_suffix)
 
                 stacked_any_c, stacked_gpc, stacked_cpg, stacked_any_a = [], [], [], []
                 row_labels, bin_labels, bin_boundaries = [], [], []
@@ -932,23 +1117,31 @@ def combined_raw_clustermap(
                         order = np.argsort(subset_bin.obs[colname].values)
 
                     elif sort_by == "gpc" and num_gpc > 0:
-                        linkage = sch.linkage(subset_bin[:, gpc_sites].layers[layer_gpc], method="ward")
+                        linkage = sch.linkage(
+                            subset_bin[:, gpc_sites].layers[layer_gpc], method="ward"
+                        )
                         order = sch.leaves_list(linkage)
 
                     elif sort_by == "cpg" and num_cpg > 0:
-                        linkage = sch.linkage(subset_bin[:, cpg_sites].layers[layer_cpg], method="ward")
+                        linkage = sch.linkage(
+                            subset_bin[:, cpg_sites].layers[layer_cpg], method="ward"
+                        )
                         order = sch.leaves_list(linkage)
 
-                    elif sort_by == "any_c" and num_any_c > 0:
-                        linkage = sch.linkage(subset_bin[:, any_c_sites].layers[layer_any_c], method="ward")
+                    elif sort_by == "c" and num_any_c > 0:
+                        linkage = sch.linkage(
+                            subset_bin[:, any_c_sites].layers[layer_c], method="ward"
+                        )
                         order = sch.leaves_list(linkage)
 
                     elif sort_by == "gpc_cpg":
                         linkage = sch.linkage(subset_bin.layers[layer_gpc], method="ward")
                         order = sch.leaves_list(linkage)
 
-                    elif sort_by == "any_a" and num_any_a > 0:
-                        linkage = sch.linkage(subset_bin[:, any_a_sites].layers[layer_a], method="ward")
+                    elif sort_by == "a" and num_any_a > 0:
+                        linkage = sch.linkage(
+                            subset_bin[:, any_a_sites].layers[layer_a], method="ward"
+                        )
                         order = sch.leaves_list(linkage)
 
                     elif sort_by == "none":
@@ -961,7 +1154,7 @@ def combined_raw_clustermap(
 
                     # stack consistently
                     if include_any_c and num_any_c > 0:
-                        stacked_any_c.append(subset_bin[:, any_c_sites].layers[layer_any_c])
+                        stacked_any_c.append(subset_bin[:, any_c_sites].layers[layer_c])
                     if include_any_c and num_gpc > 0:
                         stacked_gpc.append(subset_bin[:, gpc_sites].layers[layer_gpc])
                     if include_any_c and num_cpg > 0:
@@ -981,57 +1174,65 @@ def combined_raw_clustermap(
 
                 if include_any_c and stacked_any_c:
                     any_c_matrix = np.vstack(stacked_any_c)
-                    gpc_matrix   = np.vstack(stacked_gpc) if stacked_gpc else np.empty((0, 0))
-                    cpg_matrix   = np.vstack(stacked_cpg) if stacked_cpg else np.empty((0, 0))
+                    gpc_matrix = np.vstack(stacked_gpc) if stacked_gpc else np.empty((0, 0))
+                    cpg_matrix = np.vstack(stacked_cpg) if stacked_cpg else np.empty((0, 0))
 
                     mean_any_c = methylation_fraction(any_c_matrix) if any_c_matrix.size else None
-                    mean_gpc   = methylation_fraction(gpc_matrix) if gpc_matrix.size else None
-                    mean_cpg   = methylation_fraction(cpg_matrix) if cpg_matrix.size else None
+                    mean_gpc = methylation_fraction(gpc_matrix) if gpc_matrix.size else None
+                    mean_cpg = methylation_fraction(cpg_matrix) if cpg_matrix.size else None
 
                     if any_c_matrix.size:
-                        blocks.append(dict(
-                            name="any_c",
-                            matrix=any_c_matrix,
-                            mean=mean_any_c,
-                            labels=any_c_labels,
-                            cmap=cmap_any_c,
-                            n_xticks=n_xticks_any_c,
-                            title="any C site Modification Signal"
-                        ))
+                        blocks.append(
+                            dict(
+                                name="c",
+                                matrix=any_c_matrix,
+                                mean=mean_any_c,
+                                labels=any_c_labels,
+                                cmap=cmap_c,
+                                n_xticks=n_xticks_any_c,
+                                title="any C site Modification Signal",
+                            )
+                        )
                     if gpc_matrix.size:
-                        blocks.append(dict(
-                            name="gpc",
-                            matrix=gpc_matrix,
-                            mean=mean_gpc,
-                            labels=gpc_labels,
-                            cmap=cmap_gpc,
-                            n_xticks=n_xticks_gpc,
-                            title="GpC Modification Signal"
-                        ))
+                        blocks.append(
+                            dict(
+                                name="gpc",
+                                matrix=gpc_matrix,
+                                mean=mean_gpc,
+                                labels=gpc_labels,
+                                cmap=cmap_gpc,
+                                n_xticks=n_xticks_gpc,
+                                title="GpC Modification Signal",
+                            )
+                        )
                     if cpg_matrix.size:
-                        blocks.append(dict(
-                            name="cpg",
-                            matrix=cpg_matrix,
-                            mean=mean_cpg,
-                            labels=cpg_labels,
-                            cmap=cmap_cpg,
-                            n_xticks=n_xticks_cpg,
-                            title="CpG Modification Signal"
-                        ))
+                        blocks.append(
+                            dict(
+                                name="cpg",
+                                matrix=cpg_matrix,
+                                mean=mean_cpg,
+                                labels=cpg_labels,
+                                cmap=cmap_cpg,
+                                n_xticks=n_xticks_cpg,
+                                title="CpG Modification Signal",
+                            )
+                        )
 
                 if include_any_a and stacked_any_a:
                     any_a_matrix = np.vstack(stacked_any_a)
                     mean_any_a = methylation_fraction(any_a_matrix) if any_a_matrix.size else None
                     if any_a_matrix.size:
-                        blocks.append(dict(
-                            name="any_a",
-                            matrix=any_a_matrix,
-                            mean=mean_any_a,
-                            labels=any_a_labels,
-                            cmap=cmap_a,
-                            n_xticks=n_xticks_any_a,
-                            title="any A site Modification Signal"
-                        ))
+                        blocks.append(
+                            dict(
+                                name="a",
+                                matrix=any_a_matrix,
+                                mean=mean_any_a,
+                                labels=any_a_labels,
+                                cmap=cmap_a,
+                                n_xticks=n_xticks_any_a,
+                                title="any A site Modification Signal",
+                            )
+                        )
 
                 if not blocks:
                     print(f"No matrices to plot for {display_sample} - {ref}")
@@ -1043,7 +1244,7 @@ def combined_raw_clustermap(
                 fig.suptitle(f"{display_sample} - {ref} - {total_reads} reads", fontsize=14, y=0.97)
 
                 axes_heat = [fig.add_subplot(gs[1, i]) for i in range(gs_dim)]
-                axes_bar  = [fig.add_subplot(gs[0, i], sharex=axes_heat[i]) for i in range(gs_dim)]
+                axes_bar = [fig.add_subplot(gs[0, i], sharex=axes_heat[i]) for i in range(gs_dim)]
 
                 # ----------------------------
                 # plot blocks
@@ -1059,20 +1260,14 @@ def combined_raw_clustermap(
 
                     # heatmap
                     sns.heatmap(
-                        mat,
-                        cmap=blk["cmap"],
-                        ax=axes_heat[i],
-                        yticklabels=False,
-                        cbar=False
+                        mat, cmap=blk["cmap"], ax=axes_heat[i], yticklabels=False, cbar=False
                     )
 
                     # fixed tick labels
                     tick_pos = _fixed_tick_positions(len(labels), n_xticks)
                     axes_heat[i].set_xticks(tick_pos)
                     axes_heat[i].set_xticklabels(
-                        labels[tick_pos],
-                        rotation=xtick_rotation,
-                        fontsize=xtick_fontsize
+                        labels[tick_pos], rotation=xtick_rotation, fontsize=xtick_fontsize
                     )
 
                     # bin separators
@@ -1085,7 +1280,12 @@ def combined_raw_clustermap(
 
                 # save or show
                 if save_path is not None:
-                    safe_name = f"{ref}__{display_sample}".replace("=", "").replace("__", "_").replace(",", "_").replace(" ", "_")
+                    safe_name = (
+                        f"{ref}__{display_sample}".replace("=", "")
+                        .replace("__", "_")
+                        .replace(",", "_")
+                        .replace(" ", "_")
+                    )
                     out_file = save_path / f"{safe_name}.png"
                     fig.savefig(out_file, dpi=300)
                     plt.close(fig)
@@ -1111,20 +1311,15 @@ def combined_raw_clustermap(
                 for bin_label, percent in percentages.items():
                     print(f"  - {bin_label}: {percent:.1f}%")
 
-            except Exception as e:
+            except Exception:
                 import traceback
+
                 traceback.print_exc()
                 continue
 
-    # store once at the end (HDF5 safe)
-    # matrices won't be HDF5-safe; store only metadata + maybe hit counts
-    # adata.uns["clustermap_results"] = [
-    #     {k: v for k, v in r.items() if not k.endswith("_matrix")}
-    #     for r in results
-    # ]
-
     return results
 
+
 def plot_hmm_layers_rolling_by_sample_ref(
     adata,
     layers: Optional[Sequence[str]] = None,
@@ -1141,7 +1336,7 @@ def plot_hmm_layers_rolling_by_sample_ref(
     output_dir: Optional[str] = None,
     save: bool = True,
     show_raw: bool = False,
-    cmap: str = "tab10",
+    cmap: str = "tab20",
     use_var_coords: bool = True,
 ):
     """
@@ -1191,7 +1386,9 @@ def plot_hmm_layers_rolling_by_sample_ref(
 
     # --- basic checks / defaults ---
     if sample_col not in adata.obs.columns or ref_col not in adata.obs.columns:
-        raise ValueError(f"sample_col '{sample_col}' and ref_col '{ref_col}' must exist in adata.obs")
+        raise ValueError(
+            f"sample_col '{sample_col}' and ref_col '{ref_col}' must exist in adata.obs"
+        )
 
     # canonicalize samples / refs
     if samples is None:
@@ -1214,7 +1411,9 @@ def plot_hmm_layers_rolling_by_sample_ref(
     if layers is None:
         layers = list(adata.layers.keys())
         if len(layers) == 0:
-            raise ValueError("No adata.layers found. Please pass `layers=[...]` of the HMM layers to plot.")
+            raise ValueError(
+                "No adata.layers found. Please pass `layers=[...]` of the HMM layers to plot."
+            )
     layers = list(layers)
 
     # x coordinates (positions)
@@ -1253,19 +1452,29 @@ def plot_hmm_layers_rolling_by_sample_ref(
 
         fig_w = figsize_per_cell[0] * ncols
         fig_h = figsize_per_cell[1] * nrows
-        fig, axes = plt.subplots(nrows=nrows, ncols=ncols,
-                                 figsize=(fig_w, fig_h), dpi=dpi,
-                                 squeeze=False)
+        fig, axes = plt.subplots(
+            nrows=nrows, ncols=ncols, figsize=(fig_w, fig_h), dpi=dpi, squeeze=False
+        )
 
         for r_idx, sample_name in enumerate(chunk):
             for c_idx, ref_name in enumerate(refs_all):
                 ax = axes[r_idx][c_idx]
 
                 # subset adata
-                mask = (adata.obs[sample_col].values == sample_name) & (adata.obs[ref_col].values == ref_name)
+                mask = (adata.obs[sample_col].values == sample_name) & (
+                    adata.obs[ref_col].values == ref_name
+                )
                 sub = adata[mask]
                 if sub.n_obs == 0:
-                    ax.text(0.5, 0.5, "No reads", ha="center", va="center", transform=ax.transAxes, color="gray")
+                    ax.text(
+                        0.5,
+                        0.5,
+                        "No reads",
+                        ha="center",
+                        va="center",
+                        transform=ax.transAxes,
+                        color="gray",
+                    )
                     ax.set_xticks([])
                     ax.set_yticks([])
                     if r_idx == 0:
@@ -1315,7 +1524,11 @@ def plot_hmm_layers_rolling_by_sample_ref(
                         smoothed = col_mean
                     else:
                         ser = pd.Series(col_mean)
-                        smoothed = ser.rolling(window=window, min_periods=min_periods, center=center).mean().to_numpy()
+                        smoothed = (
+                            ser.rolling(window=window, min_periods=min_periods, center=center)
+                            .mean()
+                            .to_numpy()
+                        )
 
                     # x axis: x_coords (trim/pad to match length)
                     L = len(col_mean)
@@ -1325,7 +1538,15 @@ def plot_hmm_layers_rolling_by_sample_ref(
                     if show_raw:
                         ax.plot(x, col_mean[:L], linewidth=0.7, alpha=0.25, zorder=1)
 
-                    ax.plot(x, smoothed[:L], label=layer, color=colors[li], linewidth=1.2, alpha=0.95, zorder=2)
+                    ax.plot(
+                        x,
+                        smoothed[:L],
+                        label=layer,
+                        color=colors[li],
+                        linewidth=1.2,
+                        alpha=0.95,
+                        zorder=2,
+                    )
                     plotted_any = True
 
                 # labels / titles
@@ -1343,11 +1564,15 @@ def plot_hmm_layers_rolling_by_sample_ref(
 
                 ax.grid(True, alpha=0.2)
 
-        fig.suptitle(f"Rolling mean of layer positional means (window={window}) — page {page+1}/{total_pages}", fontsize=11, y=0.995)
+        fig.suptitle(
+            f"Rolling mean of layer positional means (window={window}) — page {page + 1}/{total_pages}",
+            fontsize=11,
+            y=0.995,
+        )
         fig.tight_layout(rect=[0, 0, 1, 0.97])
 
         if save:
-            fname = os.path.join(outdir, f"hmm_layers_rolling_page{page+1}.png")
+            fname = os.path.join(outdir, f"hmm_layers_rolling_page{page + 1}.png")
             plt.savefig(fname, bbox_inches="tight", dpi=dpi)
             saved_files.append(fname)
         else: