smftools 0.1.3__py3-none-any.whl → 0.1.7__py3-none-any.whl

smftools/plotting/classifiers.py

@@ -0,0 +1,355 @@
+
+import numpy as np
+import matplotlib.pyplot as plt
+import torch
+import os
+
+def plot_model_performance(metrics, save_path=None):
+    import matplotlib.pyplot as plt
+    import os
+    for ref in metrics.keys():
+        plt.figure(figsize=(12, 5))
+
+        # ROC Curve
+        plt.subplot(1, 2, 1)
+        for model_name, vals in metrics[ref].items():
+            model_type = model_name.split('_')[0]
+            data_type = model_name.split(f"{model_type}_")[1]
+            plt.plot(vals['fpr'], vals['tpr'], label=f"{model_type.upper()} - AUC: {vals['auc']:.4f}")
+        plt.xlabel('False Positive Rate')
+        plt.ylabel('True Positive Rate')
+        plt.title(f'{data_type} ROC Curve ({ref})')
+        plt.legend()
+
+        # PR Curve
+        plt.subplot(1, 2, 2)
+        for model_name, vals in metrics[ref].items():
+            model_type = model_name.split('_')[0]
+            data_type = model_name.split(f"{model_type}_")[1]
+            plt.plot(vals['recall'], vals['precision'], label=f"{model_type.upper()} - F1: {vals['f1']:.4f}")
+        plt.xlabel('Recall')
+        plt.ylabel('Precision')
+        plt.title(f'{data_type} Precision-Recall Curve ({ref})')
+        plt.legend()
+
+        plt.tight_layout()
+
+        if save_path:
+            save_name = f"{ref}"
+            os.makedirs(save_path, exist_ok=True)
+            safe_name = save_name.replace("=", "").replace("__", "_").replace(",", "_")
+            out_file = os.path.join(save_path, f"{safe_name}.png")
+            plt.savefig(out_file, dpi=300)
+            print(f"📁 Saved: {out_file}")
+        plt.show()
+        
+        # Confusion Matrices
+        for model_name, vals in metrics[ref].items():
+            print(f"Confusion Matrix for {ref} - {model_name.upper()}:")
+            print(vals['confusion_matrix'])
+            print()
+
+def plot_feature_importances_or_saliency(
+    models,
+    positions,
+    tensors,
+    site_config,
+    adata=None,
+    layer_name=None,
+    save_path=None,
+    shaded_regions=None
+):
+    import torch
+    import numpy as np
+    import matplotlib.pyplot as plt
+    import os
+
+    # Select device for NN models
+    device = (
+        torch.device('cuda') if torch.cuda.is_available() else
+        torch.device('mps') if torch.backends.mps.is_available() else
+        torch.device('cpu')
+    )
+
+    for ref, model_dict in models.items():
+        if layer_name:
+            suffix = layer_name
+        else:
+            suffix = "_".join(site_config[ref]) if ref in site_config else "full"
+
+        if ref not in positions or suffix not in positions[ref]:
+            print(f"Positions not found for {ref} with suffix {suffix}. Skipping {ref}.")
+            continue
+
+        coords_index = positions[ref][suffix]
+        coords = coords_index.astype(int)
+
+        # Classify positions using adata.var columns
+        cpg_sites = set()
+        gpc_sites = set()
+        other_sites = set()
+
+        if adata is None:
+            print("⚠️ AnnData object is required to classify site types. Skipping site type markers.")
+        else:
+            gpc_col = f"{ref}_GpC_site"
+            cpg_col = f"{ref}_CpG_site"
+            for idx_str in coords_index:
+                try:
+                    gpc = adata.var.at[idx_str, gpc_col] if gpc_col in adata.var.columns else False
+                    cpg = adata.var.at[idx_str, cpg_col] if cpg_col in adata.var.columns else False
+                    coord_int = int(idx_str)
+                    if gpc and not cpg:
+                        gpc_sites.add(coord_int)
+                    elif cpg and not gpc:
+                        cpg_sites.add(coord_int)
+                    else:
+                        other_sites.add(coord_int)
+                except KeyError:
+                    print(f"⚠️ Index '{idx_str}' not found in adata.var. Skipping.")
+                    continue
+
+        for model_key, model in model_dict.items():
+            if not model_key.endswith(suffix):
+                continue
+
+            if model_key.startswith("rf"):
+                if hasattr(model, "feature_importances_"):
+                    importances = model.feature_importances_
+                else:
+                    print(f"Random Forest model {model_key} has no feature_importances_. Skipping.")
+                    continue
+                plot_title = f"RF Feature Importances for {ref} ({model_key})"
+                y_label = "Feature Importance"
+            else:
+                if tensors is None or ref not in tensors or suffix not in tensors[ref]:
+                    print(f"No input data provided for NN saliency for {model_key}. Skipping.")
+                    continue
+                input_tensor = tensors[ref][suffix]
+                model.eval()
+                input_tensor = input_tensor.to(device)
+                input_tensor.requires_grad_()
+
+                with torch.enable_grad():
+                    logits = model(input_tensor)
+                    score = logits[:, 1].sum()
+                    score.backward()
+                    saliency = input_tensor.grad.abs().mean(dim=0).cpu().numpy()
+                importances = saliency
+                plot_title = f"Feature Saliency for {ref} ({model_key})"
+                y_label = "Feature Saliency"
+
+            sorted_idx = np.argsort(coords)
+            positions_sorted = coords[sorted_idx]
+            importances_sorted = np.array(importances)[sorted_idx]
+
+            plt.figure(figsize=(12, 4))
+            for pos, imp in zip(positions_sorted, importances_sorted):
+                if pos in cpg_sites:
+                    plt.plot(pos, imp, marker='*', color='black', markersize=10, linestyle='None',
+                             label='CpG site' if 'CpG site' not in plt.gca().get_legend_handles_labels()[1] else "")
+                elif pos in gpc_sites:
+                    plt.plot(pos, imp, marker='o', color='blue', markersize=6, linestyle='None',
+                             label='GpC site' if 'GpC site' not in plt.gca().get_legend_handles_labels()[1] else "")
+                else:
+                    plt.plot(pos, imp, marker='.', color='gray', linestyle='None',
+                             label='Other' if 'Other' not in plt.gca().get_legend_handles_labels()[1] else "")
+
+            plt.plot(positions_sorted, importances_sorted, linestyle='-', alpha=0.5, color='black')
+
+            if shaded_regions:
+                for (start, end) in shaded_regions:
+                    plt.axvspan(start, end, color='gray', alpha=0.3)
+
+            plt.xlabel("Genomic Position")
+            plt.ylabel(y_label)
+            plt.title(plot_title)
+            plt.grid(True)
+            plt.legend()
+            plt.tight_layout()
+
+            if save_path:
+                os.makedirs(save_path, exist_ok=True)
+                safe_name = plot_title.replace("=", "").replace("__", "_").replace(",", "_").replace(" ", "_")
+                out_file = os.path.join(save_path, f"{safe_name}.png")
+                plt.savefig(out_file, dpi=300)
+                print(f"📁 Saved: {out_file}")
+
+            plt.show()
+
+def plot_model_curves_from_adata(
+    adata, 
+    label_col='activity_status', 
+    model_names = ["cnn", "mlp", "rf"], 
+    suffix='GpC_site_CpG_site',
+    omit_training=True, 
+    save_path=None, 
+    ylim_roc=(0.0, 1.05), 
+    ylim_pr=(0.0, 1.05)):
+
+    from sklearn.metrics import precision_recall_curve, roc_curve, auc
+    import matplotlib.pyplot as plt
+    import seaborn as sns    
+
+    if omit_training:
+        subset = adata[adata.obs['used_for_training'].astype(bool) == False]
+
+    label = subset.obs[label_col].map({'Active': 1, 'Silent': 0}).values  
+
+    positive_ratio = np.sum(label.astype(int)) / len(label)
+
+    plt.figure(figsize=(12, 5))
+
+    # ROC curve
+    plt.subplot(1, 2, 1)
+    for model in model_names:
+        prob_col = f"{model}_active_prob_{suffix}"
+        if prob_col in subset.obs.columns:
+            probs = subset.obs[prob_col].astype(float).values
+            fpr, tpr, _ = roc_curve(label, probs)
+            roc_auc = auc(fpr, tpr)
+            plt.plot(fpr, tpr, label=f"{model.upper()} (AUC={roc_auc:.4f})")
+
+    plt.plot([0, 1], [0, 1], 'k--', alpha=0.5)
+    plt.xlabel("False Positive Rate")
+    plt.ylabel("True Positive Rate")
+    plt.title("ROC Curve")
+    plt.ylim(*ylim_roc)
+    plt.legend()
+
+    # PR curve
+    plt.subplot(1, 2, 2)
+    for model in model_names:
+        prob_col = f"{model}_active_prob_{suffix}"
+        if prob_col in subset.obs.columns:
+            probs = subset.obs[prob_col].astype(float).values
+            precision, recall, _ = precision_recall_curve(label, probs)
+            pr_auc = auc(recall, precision)
+            plt.plot(recall, precision, label=f"{model.upper()} (AUC={pr_auc:.4f})")
+
+    plt.xlabel("Recall")
+    plt.ylabel("Precision")
+    plt.ylim(*ylim_pr)
+    plt.axhline(y=positive_ratio, linestyle='--', color='gray', label='Random Baseline')
+    plt.title("Precision-Recall Curve")
+    plt.legend()
+
+    plt.tight_layout()
+    if save_path:
+        save_name = f"ROC_PR_curves"
+        os.makedirs(save_path, exist_ok=True)
+        safe_name = save_name.replace("=", "").replace("__", "_").replace(",", "_")
+        out_file = os.path.join(save_path, f"{safe_name}.png")
+        plt.savefig(out_file, dpi=300)
+        print(f"📁 Saved: {out_file}")
+    plt.show()
+
+def plot_model_curves_from_adata_with_frequency_grid(
+    adata,
+    label_col='activity_status',
+    model_names=["cnn", "mlp", "rf"],
+    suffix='GpC_site_CpG_site',
+    omit_training=True,
+    save_path=None,
+    ylim_roc=(0.0, 1.05),
+    ylim_pr=(0.0, 1.05),
+    pos_sample_count=500,
+    pos_freq_list=[0.01, 0.05, 0.1],
+    show_f1_iso_curves=False,
+    f1_levels=None):
+    import numpy as np
+    import matplotlib.pyplot as plt
+    import seaborn as sns
+    import os
+    from sklearn.metrics import precision_recall_curve, roc_curve, auc
+    
+    if f1_levels is None:
+        f1_levels = np.linspace(0.2, 0.9, 8)
+        
+    if omit_training:
+        subset = adata[adata.obs['used_for_training'].astype(bool) == False]
+    else:
+        subset = adata
+
+    label = subset.obs[label_col].map({'Active': 1, 'Silent': 0}).values
+    subset = subset.copy()
+    subset.obs["__label__"] = label
+
+    pos_indices = np.where(label == 1)[0]
+    neg_indices = np.where(label == 0)[0]
+
+    n_rows = len(pos_freq_list)
+    fig, axes = plt.subplots(n_rows, 2, figsize=(12, 5 * n_rows))
+    fig.suptitle(f'{suffix} Performance metrics')
+
+    for row_idx, pos_freq in enumerate(pos_freq_list):
+        desired_total = int(pos_sample_count / pos_freq)
+        neg_sample_count = desired_total - pos_sample_count
+
+        if pos_sample_count > len(pos_indices) or neg_sample_count > len(neg_indices):
+            print(f"⚠️ Skipping frequency {pos_freq:.3f}: not enough samples.")
+            continue
+
+        sampled_pos = np.random.choice(pos_indices, size=pos_sample_count, replace=False)
+        sampled_neg = np.random.choice(neg_indices, size=neg_sample_count, replace=False)
+        sampled_indices = np.concatenate([sampled_pos, sampled_neg])
+
+        data_sampled = subset[sampled_indices]
+        y_true = data_sampled.obs["__label__"].values
+
+        ax_roc = axes[row_idx, 0] if n_rows > 1 else axes[0]
+        ax_pr = axes[row_idx, 1] if n_rows > 1 else axes[1]
+
+        # ROC Curve
+        for model in model_names:
+            prob_col = f"{model}_active_prob_{suffix}"
+            if prob_col in data_sampled.obs.columns:
+                probs = data_sampled.obs[prob_col].astype(float).values
+                fpr, tpr, _ = roc_curve(y_true, probs)
+                roc_auc = auc(fpr, tpr)
+                ax_roc.plot(fpr, tpr, label=f"{model.upper()} (AUC={roc_auc:.4f})")
+        ax_roc.plot([0, 1], [0, 1], 'k--', alpha=0.5)
+        ax_roc.set_xlabel("False Positive Rate")
+        ax_roc.set_ylabel("True Positive Rate")
+        ax_roc.set_ylim(*ylim_roc)
+        ax_roc.set_title(f"ROC Curve (Pos Freq: {pos_freq:.2%})")
+        ax_roc.legend()
+        ax_roc.spines['top'].set_visible(False)
+        ax_roc.spines['right'].set_visible(False)
+
+        # PR Curve
+        for model in model_names:
+            prob_col = f"{model}_active_prob_{suffix}"
+            if prob_col in data_sampled.obs.columns:
+                probs = data_sampled.obs[prob_col].astype(float).values
+                precision, recall, _ = precision_recall_curve(y_true, probs)
+                pr_auc = auc(recall, precision)
+                ax_pr.plot(recall, precision, label=f"{model.upper()} (AUC={pr_auc:.4f})")
+        ax_pr.axhline(y=pos_freq, linestyle='--', color='gray', label='Random Baseline')
+
+        if show_f1_iso_curves:
+            recall_vals = np.linspace(0.01, 1, 500)
+            for f1 in f1_levels:
+                precision_vals = (f1 * recall_vals) / (2 * recall_vals - f1)
+                precision_vals[precision_vals < 0] = np.nan  # Avoid plotting invalid values
+                ax_pr.plot(recall_vals, precision_vals, color='gray', linestyle=':', linewidth=1, alpha=0.6)
+                x_val = 0.9
+                y_val = (f1 * x_val) / (2 * x_val - f1)
+                if 0 < y_val < 1:
+                    ax_pr.text(x_val, y_val, f"F1={f1:.1f}", fontsize=8, color='gray')
+
+        ax_pr.set_xlabel("Recall")
+        ax_pr.set_ylabel("Precision")
+        ax_pr.set_ylim(*ylim_pr)
+        ax_pr.set_title(f"PR Curve (Pos Freq: {pos_freq:.2%})")
+        ax_pr.legend()
+        ax_pr.spines['top'].set_visible(False)
+        ax_pr.spines['right'].set_visible(False)
+
+    plt.tight_layout(rect=[0, 0, 1, 0.97])
+    if save_path:
+        os.makedirs(save_path, exist_ok=True)
+        out_file = os.path.join(save_path, "ROC_PR_grid.png")
+        plt.savefig(out_file, dpi=300)
+        print(f"📁 Saved: {out_file}")
+    plt.show()

smftools/plotting/general_plotting.py

@@ -0,0 +1,205 @@
+import numpy as np
+import seaborn as sns
+import matplotlib.pyplot as plt
+
+def clean_barplot(ax, mean_values, title):
+    x = np.arange(len(mean_values))
+    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])
+    ax.set_ylabel("Mean")
+    ax.set_title(title, fontsize=12, pad=2)
+
+    # Hide all spines except left
+    for spine_name, spine in ax.spines.items():
+        spine.set_visible(spine_name == 'left')
+
+    ax.tick_params(axis='x', which='both', bottom=False, top=False, labelbottom=False)
+
+
+def combined_hmm_raw_clustermap(
+    adata,
+    sample_col='Sample_Names',
+    hmm_feature_layer="hmm_combined",
+    layer_gpc="nan0_0minus1",
+    layer_cpg="nan0_0minus1",
+    cmap_hmm="tab10",
+    cmap_gpc="coolwarm",
+    cmap_cpg="viridis",
+    min_quality=20,
+    min_length=2700,
+    sample_mapping=None,
+    save_path=None,
+    normalize_hmm=False,
+    sort_by="gpc",  # options: 'gpc', 'cpg', 'gpc_cpg', 'none', or 'obs:<column>'
+    bins=None
+):
+    import scipy.cluster.hierarchy as sch
+    import pandas as pd
+    import numpy as np
+    import seaborn as sns
+    import matplotlib.pyplot as plt
+    import matplotlib.gridspec as gridspec
+    import os
+
+    results = []
+
+    for ref in adata.obs["Reference_strand"].cat.categories:
+        for sample in adata.obs[sample_col].cat.categories:
+            try:
+                subset = adata[
+                    (adata.obs['Reference_strand'] == ref) &
+                    (adata.obs[sample_col] == sample) &
+                    (adata.obs['query_read_quality'] >= min_quality) &
+                    (adata.obs['read_length'] >= min_length) &
+                    (adata.obs['Raw_methylation_signal'] >= 20)
+                ]
+                
+                subset = subset[:, subset.var[f'position_in_{ref}'] == True]
+
+                if subset.shape[0] == 0:
+                    print(f"  ❌ No reads left after filtering for {sample} - {ref}")
+                    continue
+
+                if bins:
+                    pass
+                else:
+                    bins = {"All": (subset.obs['Reference_strand'] != None)}
+
+                # Get column positions (not var_names!) of site masks
+                gpc_sites = np.where(subset.var[f"{ref}_GpC_site"].values)[0]
+                cpg_sites = np.where(subset.var[f"{ref}_CpG_site"].values)[0]
+
+                # Use var_names for x-axis tick labels
+                gpc_labels = subset.var_names[gpc_sites].astype(int)
+                cpg_labels = subset.var_names[cpg_sites].astype(int)
+
+                stacked_hmm_feature, stacked_gpc, stacked_cpg = [], [], []
+                row_labels, bin_labels = [], []
+                bin_boundaries = []
+
+                total_reads = subset.shape[0]
+                percentages = {}
+                last_idx = 0
+
+                for bin_label, bin_filter in bins.items():
+                    subset_bin = subset[bin_filter].copy()
+                    num_reads = subset_bin.shape[0]
+                    percent_reads = (num_reads / total_reads) * 100 if total_reads > 0 else 0
+                    percentages[bin_label] = percent_reads
+
+                    if num_reads > 0:
+                        # Determine sorting order
+                        if sort_by.startswith("obs:"):
+                            colname = sort_by.split("obs:")[1]
+                            order = np.argsort(subset_bin.obs[colname].values)
+                        elif sort_by == "gpc":
+                            linkage = sch.linkage(subset_bin[:, gpc_sites].layers[layer_gpc], method="ward")
+                            order = sch.leaves_list(linkage)
+                        elif sort_by == "cpg":
+                            linkage = sch.linkage(subset_bin[:, cpg_sites].layers[layer_cpg], 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 == "none":
+                            order = np.arange(num_reads)
+                        else:
+                            raise ValueError(f"Unsupported sort_by option: {sort_by}")
+
+                        stacked_hmm_feature.append(subset_bin[order].layers[hmm_feature_layer])
+                        stacked_gpc.append(subset_bin[order][:, gpc_sites].layers[layer_gpc])
+                        stacked_cpg.append(subset_bin[order][:, cpg_sites].layers[layer_cpg])
+
+                        row_labels.extend([bin_label] * num_reads)
+                        bin_labels.append(f"{bin_label}: {num_reads} reads ({percent_reads:.1f}%)")
+                        last_idx += num_reads
+                        bin_boundaries.append(last_idx)
+
+                if stacked_hmm_feature:
+                    hmm_matrix = np.vstack(stacked_hmm_feature)
+                    gpc_matrix = np.vstack(stacked_gpc)
+                    cpg_matrix = np.vstack(stacked_cpg)
+
+                    def normalized_mean(matrix):
+                        mean = np.nanmean(matrix, axis=0)
+                        normalized = (mean - mean.min()) / (mean.max() - mean.min() + 1e-9)
+                        return normalized
+
+                    def methylation_fraction(matrix):
+                        methylated = (matrix == 1).sum(axis=0)
+                        valid = (matrix != 0).sum(axis=0)
+                        return np.divide(methylated, valid, out=np.zeros_like(methylated, dtype=float), where=valid != 0)
+
+                    if normalize_hmm:
+                        mean_hmm = normalized_mean(hmm_matrix)
+                    else:
+                        mean_hmm = np.nanmean(hmm_matrix, axis=0)
+                    mean_gpc = methylation_fraction(gpc_matrix)
+                    mean_cpg = methylation_fraction(cpg_matrix)
+
+                    fig = plt.figure(figsize=(18, 12))
+                    gs = gridspec.GridSpec(2, 3, height_ratios=[1, 6], hspace=0.01)
+                    fig.suptitle(f"{sample} - {ref}", fontsize=14, y=0.95)
+
+                    axes_heat = [fig.add_subplot(gs[1, i]) for i in range(3)]
+                    axes_bar = [fig.add_subplot(gs[0, i], sharex=axes_heat[i]) for i in range(3)]
+
+                    clean_barplot(axes_bar[0], mean_hmm, f"{hmm_feature_layer} HMM Features")
+                    clean_barplot(axes_bar[1], mean_gpc, f"GpC Methylation")
+                    clean_barplot(axes_bar[2], mean_cpg, f"CpG Methylation")
+                    
+                    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)
+                    axes_heat[0].set_xticks(range(0, len(hmm_labels), hmm_label_spacing))
+                    axes_heat[0].set_xticklabels(hmm_labels[::hmm_label_spacing], rotation=90, fontsize=10)
+                    for boundary in bin_boundaries[:-1]:
+                        axes_heat[0].axhline(y=boundary, color="black", linewidth=2)
+
+                    sns.heatmap(gpc_matrix, cmap=cmap_gpc, ax=axes_heat[1], xticklabels=gpc_labels[::5], yticklabels=False, cbar=False)
+                    axes_heat[1].set_xticks(range(0, len(gpc_labels), 5))
+                    axes_heat[1].set_xticklabels(gpc_labels[::5], rotation=90, fontsize=10)
+                    for boundary in bin_boundaries[:-1]:
+                        axes_heat[1].axhline(y=boundary, color="black", linewidth=2)
+
+                    sns.heatmap(cpg_matrix, cmap=cmap_cpg, ax=axes_heat[2], xticklabels=cpg_labels, yticklabels=False, cbar=False)
+                    axes_heat[2].set_xticklabels(cpg_labels, rotation=90, fontsize=10)
+                    for boundary in bin_boundaries[:-1]:
+                        axes_heat[2].axhline(y=boundary, color="black", linewidth=2)
+
+                    plt.tight_layout()
+
+                    if save_path:
+                        save_name = f"{ref} — {sample}"
+                        os.makedirs(save_path, exist_ok=True)
+                        safe_name = save_name.replace("=", "").replace("__", "_").replace(",", "_")
+                        out_file = os.path.join(save_path, f"{safe_name}.png")
+                        plt.savefig(out_file, dpi=300)
+                        print(f"📁 Saved: {out_file}")
+
+                    plt.show()
+
+                    print(f"📊 Summary for {sample} - {ref}:")
+                    for bin_label, percent in percentages.items():
+                        print(f"  - {bin_label}: {percent:.1f}%")
+
+                    results.append({
+                        "sample": sample,
+                        "ref": ref,
+                        "hmm_matrix": hmm_matrix,
+                        "gpc_matrix": gpc_matrix,
+                        "cpg_matrix": cpg_matrix,
+                        "row_labels": row_labels,
+                        "bin_labels": bin_labels,
+                        "bin_boundaries": bin_boundaries,
+                        "percentages": percentages
+                    })
+                    
+                    adata.uns['clustermap_results'] = results
+
+            except Exception as e:
+                import traceback
+                traceback.print_exc()
+                continue