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

smftools/machine_learning/evaluation/evaluators.py

@@ -0,0 +1,223 @@
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+from sklearn.metrics import (
+    roc_auc_score, precision_recall_curve, auc, f1_score, confusion_matrix, roc_curve
+)
+
+class ModelEvaluator:
+    """
+    A model evaluator for consolidating Sklearn and Lightning model evaluation metrics on testing data
+    """
+    def __init__(self):
+        self.results = []
+        self.pos_freq = None
+        self.num_pos = None
+
+    def add_model(self, name, model, is_torch=True):
+        """
+        Add a trained model with its evaluation metrics.
+        """
+        if is_torch:
+            entry = {
+                'name': name,
+                'f1': model.test_f1,
+                'auc': model.test_roc_auc,
+                'pr_auc': model.test_pr_auc,
+                'pr_auc_norm': model.test_pr_auc / model.test_pos_freq if model.test_pos_freq > 0 else np.nan,
+                'pr_curve': model.test_pr_curve,
+                'roc_curve': model.test_roc_curve,
+                'num_pos': model.test_num_pos,
+                'pos_freq': model.test_pos_freq
+            }
+        else:
+            entry = {
+                'name': name,
+                'f1': model.test_f1,
+                'auc': model.test_roc_auc,
+                'pr_auc': model.test_pr_auc,
+                'pr_auc_norm': model.test_pr_auc / model.test_pos_freq if model.test_pos_freq > 0 else np.nan,
+                'pr_curve': model.test_pr_curve,
+                'roc_curve': model.test_roc_curve,
+                'num_pos': model.test_num_pos,
+                'pos_freq': model.test_pos_freq
+            }
+        
+        self.results.append(entry)
+
+        if not self.pos_freq:
+            self.pos_freq = entry['pos_freq']
+            self.num_pos = entry['num_pos']
+
+    def get_metrics_dataframe(self):
+        """
+        Return all metrics as pandas DataFrame.
+        """
+        df = pd.DataFrame(self.results)
+        return df[['name', 'f1', 'auc', 'pr_auc', 'pr_auc_norm', 'num_pos', 'pos_freq']]
+
+    def plot_all_curves(self):
+        """
+        Plot unified ROC and PR curves across all models.
+        """
+        plt.figure(figsize=(12, 5))
+
+        # ROC
+        plt.subplot(1, 2, 1)
+        for res in self.results:
+            fpr, tpr = res['roc_curve']
+            plt.plot(fpr, tpr, label=f"{res['name']} (AUC={res['auc']:.3f})")
+        plt.plot([0, 1], [0, 1], linestyle="--", color="gray")
+        plt.xlabel("False Positive Rate")
+        plt.ylabel("True Positive Rate")
+        plt.ylim(0,1.05)
+        plt.title(f"ROC Curves - {self.num_pos} positive instances")
+        plt.legend()
+
+        # PR
+        plt.subplot(1, 2, 2)
+        for res in self.results:
+            rc, pr = res['pr_curve']
+            plt.plot(rc, pr, label=f"{res['name']} (AUPRC={res['pr_auc']:.3f})")
+        plt.xlabel("Recall")
+        plt.ylabel("Precision")
+        plt.ylim(0,1.05)
+        plt.axhline(self.pos_freq, linestyle='--', color='grey')
+        plt.title(f"Precision-Recall Curves - {self.num_pos} positive instances")
+        plt.legend()
+
+        plt.tight_layout()
+        plt.show()
+
+class PostInferenceModelEvaluator:
+    def __init__(self, adata, models, target_eval_freq=None, max_eval_positive=None):
+        """
+        Initialize evaluator.
+
+        Parameters:
+        -----------
+        adata : AnnData
+            The annotated dataset where predictions are stored in obs/obsm.
+        models : dict
+            Dictionary of models: {model_name: model_instance}.
+            Supports TorchClassifierWrapper and SklearnModelWrapper.
+        """
+        self.adata = adata
+        self.models = models
+        self.target_eval_freq = target_eval_freq
+        self.max_eval_positive = max_eval_positive
+        self.results = {}
+
+    def evaluate_all(self):
+        """
+        Evaluate all models and store results.
+        """
+        for name, model in self.models.items():
+            print(f"Evaluating {name}...")
+            label_col = model.label_col
+            full_prefix = f"{name}_{label_col}"
+            self.results[full_prefix] = self._evaluate_model(name, model)
+
+    def _evaluate_model(self, model_name, model):
+        """
+        Evaluate one model and return metrics.
+        """
+        label_col = model.label_col
+        num_classes = model.num_classes
+        class_names = model.class_names
+        focus_class = model.focus_class
+
+        full_prefix = f"{model_name}_{label_col}"
+
+        # Extract ground truth + predictions
+        y_true = self.adata.obs[label_col].cat.codes.to_numpy()
+        y_pred = self.adata.obs[f"{full_prefix}_pred"].to_numpy()
+        probs_all = self.adata.obsm[f"{full_prefix}_pred_prob_all"]
+
+        binary_focus = (y_true == focus_class).astype(int)
+
+        # OPTIONAL SUBSAMPLING
+        if self.target_eval_freq is not None:
+            indices = self._subsample_for_fixed_positive_frequency(
+                binary_focus, target_freq=self.target_eval_freq, max_positive=self.max_eval_positive
+            )
+            y_true = y_true[indices]
+            y_pred = y_pred[indices]
+            probs_all = probs_all[indices]
+            binary_focus = (y_true == focus_class).astype(int)
+
+        acc = np.mean(y_true == y_pred)
+
+        if num_classes == 2:
+            focus_probs = probs_all[:, focus_class]
+            f1 = f1_score(binary_focus, (y_pred == focus_class).astype(int))
+            roc_auc = roc_auc_score(binary_focus, focus_probs)
+            pr, rc, _ = precision_recall_curve(binary_focus, focus_probs)
+            fpr, tpr, _ = roc_curve(binary_focus, focus_probs)
+            pr_auc = auc(rc, pr)
+            pos_freq = binary_focus.mean()
+            pr_auc_norm = pr_auc / pos_freq if pos_freq > 0 else np.nan
+        else:
+            f1 = f1_score(y_true, y_pred, average="macro")
+            roc_auc = roc_auc_score(y_true, probs_all, multi_class="ovr", average="macro")
+            focus_probs = probs_all[:, focus_class]
+            pr, rc, _ = precision_recall_curve(binary_focus, focus_probs)
+            fpr, tpr, _ = roc_curve(binary_focus, focus_probs)
+            pr_auc = auc(rc, pr)
+            pos_freq = binary_focus.mean()
+            pr_auc_norm = pr_auc / pos_freq if pos_freq > 0 else np.nan
+
+        cm = confusion_matrix(y_true, y_pred)
+
+        metrics = {
+            "accuracy": acc,
+            "f1": f1,
+            "roc_auc": roc_auc,
+            "pr_auc": pr_auc,
+            "pr_auc_norm": pr_auc_norm,
+            "pos_freq": pos_freq,
+            "confusion_matrix": cm,
+            "pr_rc_curve": (pr, rc),
+            "roc_curve": (tpr, fpr)
+        }
+
+        return metrics
+    
+    def _subsample_for_fixed_positive_frequency(self, binary_labels, target_freq=0.3, max_positive=None):
+        pos_idx = np.where(binary_labels == 1)[0]
+        neg_idx = np.where(binary_labels == 0)[0]
+
+        max_pos = len(pos_idx)
+        max_neg = len(neg_idx)
+
+        max_possible_freq = max_pos / (max_pos + max_neg)
+        if target_freq > max_possible_freq:
+            target_freq = max_possible_freq
+
+        num_pos_target = int(target_freq * max_neg / (1 - target_freq))
+        num_pos_target = min(num_pos_target, max_pos)
+        if max_positive is not None:
+            num_pos_target = min(num_pos_target, max_positive)
+
+        num_neg_target = int(num_pos_target * (1 - target_freq) / target_freq)
+        num_neg_target = min(num_neg_target, max_neg)
+
+        pos_sampled = np.random.choice(pos_idx, size=num_pos_target, replace=False)
+        neg_sampled = np.random.choice(neg_idx, size=num_neg_target, replace=False)
+        sampled_idx = np.concatenate([pos_sampled, neg_sampled])
+        np.random.shuffle(sampled_idx)
+        return sampled_idx
+
+    def to_dataframe(self):
+        """
+        Convert results to pandas DataFrame (excluding confusion matrices).
+        """
+        records = []
+        for model_name, metrics in self.results.items():
+            row = {"model": model_name}
+            for k, v in metrics.items():
+                if k not in ["confusion_matrix", "pr_rc_curve", "roc_curve"]:
+                    row[k] = v
+            records.append(row)
+        return pd.DataFrame(records)

smftools/machine_learning/inference/__init__.py

@@ -0,0 +1,3 @@
+from .lightning_inference import run_lightning_inference
+from .sliding_window_inference import sliding_window_inference
+from .sklearn_inference import run_sklearn_inference

smftools/machine_learning/inference/inference_utils.py

@@ -0,0 +1,27 @@
+import pandas as pd
+
+def annotate_split_column(adata, model, split_col="split"):
+    """
+    Annotate adata.obs with train/val/test/new labels based on model's stored obs_names.
+    """
+    # Get sets for fast lookup
+    train_set = set(model.train_obs_names)
+    val_set = set(model.val_obs_names)
+    test_set = set(model.test_obs_names)
+    
+    # Create array for split labels
+    split_labels = []
+    for obs in adata.obs_names:
+        if obs in train_set:
+            split_labels.append("training")
+        elif obs in val_set:
+            split_labels.append("validation")
+        elif obs in test_set:
+            split_labels.append("testing")
+        else:
+            split_labels.append("new")
+    
+    # Store in AnnData.obs
+    adata.obs[split_col] = pd.Categorical(split_labels, categories=["training", "validation", "testing", "new"])
+    
+    print(f"Annotated {split_col} column with training/validation/testing/new status.")

smftools/machine_learning/inference/lightning_inference.py

@@ -0,0 +1,68 @@
+import torch
+import pandas as pd
+import numpy as np
+from pytorch_lightning import Trainer
+from .inference_utils import annotate_split_column
+
+def run_lightning_inference(
+    adata,
+    model,
+    datamodule,
+    trainer,
+    prefix="model",
+    devices=1
+):
+    """
+    Run inference on AnnData using TorchClassifierWrapper + AnnDataModule (in inference mode).
+    """
+
+    # Device logic
+    if torch.cuda.is_available():
+        accelerator = "gpu"
+    elif torch.backends.mps.is_available():
+        accelerator = "mps"
+        devices = 1
+    else:
+        accelerator = "cpu"
+        devices = 1
+
+    label_col = model.label_col
+    num_classes = model.num_classes
+    class_labels = model.class_names
+    focus_class = model.focus_class
+    focus_class_name = model.focus_class_name
+
+    annotate_split_column(adata, model, split_col=f"{prefix}_training_split")
+
+    # Run predictions
+    outputs = trainer.predict(model, datamodule=datamodule)
+
+    preds_list, probs_list = zip(*outputs)
+    preds = torch.cat(preds_list, dim=0).cpu().numpy()
+    probs = torch.cat(probs_list, dim=0).cpu().numpy()
+
+    # Handle binary vs multiclass formats
+    if num_classes == 2:
+        # probs shape: (N,) from sigmoid
+        pred_class_idx = (probs >= 0.5).astype(int)
+        probs_all = np.vstack([1 - probs, probs]).T  # shape (N, 2)
+        pred_class_probs = probs_all[np.arange(len(probs_all)), pred_class_idx]
+    else:
+        pred_class_idx = probs.argmax(axis=1)
+        probs_all = probs
+        pred_class_probs = probs_all[np.arange(len(probs_all)), pred_class_idx]
+
+    pred_class_labels = [class_labels[i] for i in pred_class_idx]
+
+    full_prefix = f"{prefix}_{label_col}"
+
+    adata.obs[f"{full_prefix}_pred"] = pred_class_idx
+    adata.obs[f"{full_prefix}_pred_label"] = pd.Categorical(pred_class_labels, categories=class_labels)
+    adata.obs[f"{full_prefix}_pred_prob"] = pred_class_probs
+
+    for i, class_name in enumerate(class_labels):
+        adata.obs[f"{full_prefix}_prob_{class_name}"] = probs_all[:, i]
+
+    adata.obsm[f"{full_prefix}_pred_prob_all"] = probs_all
+
+    print(f"Inference complete: stored under prefix '{full_prefix}'")

smftools/machine_learning/inference/sklearn_inference.py

@@ -0,0 +1,55 @@
+import pandas as pd
+import numpy as np
+from .inference_utils import annotate_split_column
+
+
+def run_sklearn_inference(
+    adata,
+    model,
+    datamodule, 
+    prefix="model"
+):
+    """
+    Run inference on AnnData using SklearnModelWrapper.
+    """
+
+    label_col = model.label_col
+    num_classes = model.num_classes
+    class_labels = model.class_names
+    focus_class_name = model.focus_class_name
+
+    annotate_split_column(adata, model, split_col=f"{prefix}_training_split")
+
+    datamodule.setup()
+
+    X_infer = datamodule.to_numpy()
+
+    # Run predictions
+    preds = model.predict(X_infer)
+    probs = model.predict_proba(X_infer)
+
+    # Handle binary vs multiclass formats
+    if num_classes == 2:
+        # probs shape: (N, 2) from predict_proba
+        pred_class_idx = preds
+        probs_all = probs
+        pred_class_probs = probs[np.arange(len(probs)), pred_class_idx]
+    else:
+        pred_class_idx = preds
+        probs_all = probs
+        pred_class_probs = probs[np.arange(len(probs)), pred_class_idx]
+
+    pred_class_labels = [class_labels[i] for i in pred_class_idx]
+
+    full_prefix = f"{prefix}_{label_col}"
+
+    adata.obs[f"{full_prefix}_pred"] = pred_class_idx
+    adata.obs[f"{full_prefix}_pred_label"] = pd.Categorical(pred_class_labels, categories=class_labels)
+    adata.obs[f"{full_prefix}_pred_prob"] = pred_class_probs
+
+    for i, class_name in enumerate(class_labels):
+        adata.obs[f"{full_prefix}_prob_{class_name}"] = probs_all[:, i]
+
+    adata.obsm[f"{full_prefix}_pred_prob_all"] = probs_all
+
+    print(f"Inference complete: stored under prefix '{full_prefix}'")

smftools/machine_learning/inference/sliding_window_inference.py

@@ -0,0 +1,114 @@
+from ..data import AnnDataModule
+from ..evaluation import PostInferenceModelEvaluator
+from .lightning_inference import run_lightning_inference
+from .sklearn_inference import run_sklearn_inference
+
+def sliding_window_inference(
+    adata, 
+    trained_results, 
+    tensor_source='X',
+    tensor_key=None,
+    label_col='activity_status',
+    batch_size=64,
+    cleanup=False,
+    target_eval_freq=None, 
+    max_eval_positive=None
+):
+    """
+    Apply trained sliding window models to an AnnData object (Lightning or Sklearn).
+    Evaluate model performance and return a df.
+    Optionally remove the appended inference columns from AnnData to clean up obs namespace.
+    """
+    ## Inference using trained models
+    for model_name, model_dict in trained_results.items():
+        for window_size, window_data in model_dict.items():
+            for center_varname, run in window_data.items():
+                print(f"\nEvaluating {model_name} window {window_size} around {center_varname}")
+                
+                # Extract window start from varname
+                center_idx = adata.var_names.get_loc(center_varname)
+                window_start = center_idx - window_size // 2
+                
+                # Build datamodule for window
+                datamodule = AnnDataModule(
+                    adata,
+                    tensor_source=tensor_source,
+                    tensor_key=tensor_key,
+                    label_col=label_col,
+                    batch_size=batch_size,
+                    window_start=window_start,
+                    window_size=window_size,
+                    inference_mode=True
+                )
+                datamodule.setup()
+
+                # Extract model + detect type
+                model = run['model']
+
+                # Lightning models
+                if hasattr(run, 'trainer') or 'trainer' in run:
+                    trainer = run['trainer']
+                    run_lightning_inference(
+                        adata,
+                        model=model,
+                        datamodule=datamodule,
+                        trainer=trainer,
+                        prefix=f"{model_name}_w{window_size}_c{center_varname}"
+                    )
+                
+                # Sklearn models
+                else:
+                    run_sklearn_inference(
+                        adata,
+                        model=model,
+                        datamodule=datamodule,
+                        prefix=f"{model_name}_w{window_size}_c{center_varname}"
+                    )
+
+    print("Inference complete across all models.")
+
+    ## Post-inference model evaluation
+    model_wrappers = {}
+
+    for model_name, model_dict in trained_results.items():
+        for window_size, window_data in model_dict.items():
+            for center_varname, run in window_data.items():
+                # Reconstruct the prefix string you used in inference
+                prefix = f"{model_name}_w{window_size}_c{center_varname}"
+                # Use full key for uniqueness
+                key = prefix
+                model_wrappers[key] = run['model']
+
+    # Run evaluator
+    evaluator = PostInferenceModelEvaluator(adata, model_wrappers, target_eval_freq=target_eval_freq, max_eval_positive=max_eval_positive)
+    evaluator.evaluate_all()
+
+    # Get results
+    df = evaluator.to_dataframe()
+
+    df[['model_name', 'window_size', 'center']] = df['model'].str.extract(r'(\w+)_w(\d+)_c(\d+)_activity_status')
+
+    # Cast window_size and center to integers for plotting
+    df['window_size'] = df['window_size'].astype(int)
+    df['center'] = df['center'].astype(int)
+
+    ## Optional cleanup:
+    if cleanup:
+        prefixes = [f"{model_name}_w{window_size}_c{center_varname}" 
+                    for model_name, model_dict in trained_results.items() 
+                    for window_size, window_data in model_dict.items() 
+                    for center_varname in window_data.keys()]
+
+        # Remove matching obs columns
+        for prefix in prefixes:
+            to_remove = [col for col in adata.obs.columns if col.startswith(prefix)]
+            adata.obs.drop(columns=to_remove, inplace=True)
+
+            # Remove obsm entries if any
+            obsm_key = f"{prefix}_pred_prob_all"
+            if obsm_key in adata.obsm:
+                del adata.obsm[obsm_key]
+
+        print(f"Cleaned up {len(prefixes)} model prefixes from AnnData.")
+
+    return df