smftools 0.2.4__py3-none-any.whl → 0.3.0__py3-none-any.whl

smftools/machine_learning/evaluation/evaluators.py

@@ -1,15 +1,26 @@
+from __future__ import annotations
+
 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
-)
+from smftools.optional_imports import require
+
+plt = require("matplotlib.pyplot", extra="plotting", purpose="evaluation plots")
+sklearn_metrics = require("sklearn.metrics", extra="ml-base", purpose="model evaluation")
+
+auc = sklearn_metrics.auc
+confusion_matrix = sklearn_metrics.confusion_matrix
+f1_score = sklearn_metrics.f1_score
+precision_recall_curve = sklearn_metrics.precision_recall_curve
+roc_auc_score = sklearn_metrics.roc_auc_score
+roc_curve = sklearn_metrics.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
@@ -21,41 +32,45 @@ class ModelEvaluator:
         """
         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
+                "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
+                "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']
+            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']]
+        return df[["name", "f1", "auc", "pr_auc", "pr_auc_norm", "num_pos", "pos_freq"]]
 
     def plot_all_curves(self):
         """
@@ -66,30 +81,31 @@ class ModelEvaluator:
         # ROC
         plt.subplot(1, 2, 1)
         for res in self.results:
-            fpr, tpr = res['roc_curve']
+            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.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']
+            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.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):
         """
@@ -179,12 +195,14 @@ class PostInferenceModelEvaluator:
             "pos_freq": pos_freq,
             "confusion_matrix": cm,
             "pr_rc_curve": (pr, rc),
-            "roc_curve": (tpr, fpr)
+            "roc_curve": (tpr, fpr),
         }
 
         return metrics
-    
-    def _subsample_for_fixed_positive_frequency(self, binary_labels, target_freq=0.3, max_positive=None):
+
+    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]
 

smftools/machine_learning/inference/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from .lightning_inference import run_lightning_inference
+from .sklearn_inference import run_sklearn_inference
 from .sliding_window_inference import sliding_window_inference
-from .sklearn_inference import run_sklearn_inference

smftools/machine_learning/inference/inference_utils.py

@@ -1,5 +1,8 @@
+from __future__ import annotations
+
 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.
@@ -8,7 +11,7 @@ def annotate_split_column(adata, model, split_col="split"):
     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:
@@ -20,8 +23,10 @@ def annotate_split_column(adata, model, split_col="split"):
             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"])
-    
+    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

@@ -1,17 +1,16 @@
-import torch
-import pandas as pd
+from __future__ import annotations
+
 import numpy as np
-from pytorch_lightning import Trainer
+import pandas as pd
+
+from smftools.optional_imports import require
+
 from .inference_utils import annotate_split_column
 
-def run_lightning_inference(
-    adata,
-    model,
-    datamodule,
-    trainer,
-    prefix="model",
-    devices=1
-):
+torch = require("torch", extra="ml-base", purpose="Lightning inference")
+
+
+def run_lightning_inference(adata, model, datamodule, trainer, prefix="model", devices=1):
     """
     Run inference on AnnData using TorchClassifierWrapper + AnnDataModule (in inference mode).
     """
@@ -57,7 +56,9 @@ def run_lightning_inference(
     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_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):
@@ -65,4 +66,4 @@ def run_lightning_inference(
 
     adata.obsm[f"{full_prefix}_pred_prob_all"] = probs_all
 
-    print(f"Inference complete: stored under prefix '{full_prefix}'")
+    print(f"Inference complete: stored under prefix '{full_prefix}'")

smftools/machine_learning/inference/sklearn_inference.py

@@ -1,14 +1,12 @@
-import pandas as pd
+from __future__ import annotations
+
 import numpy as np
+import pandas as pd
+
 from .inference_utils import annotate_split_column
 
 
-def run_sklearn_inference(
-    adata,
-    model,
-    datamodule, 
-    prefix="model"
-):
+def run_sklearn_inference(adata, model, datamodule, prefix="model"):
     """
     Run inference on AnnData using SklearnModelWrapper.
     """
@@ -44,7 +42,9 @@ def run_sklearn_inference(
     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_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):

smftools/machine_learning/inference/sliding_window_inference.py

@@ -1,18 +1,21 @@
+from __future__ import annotations
+
 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',
+    adata,
+    trained_results,
+    tensor_source="X",
     tensor_key=None,
-    label_col='activity_status',
+    label_col="activity_status",
     batch_size=64,
     cleanup=False,
-    target_eval_freq=None, 
-    max_eval_positive=None
+    target_eval_freq=None,
+    max_eval_positive=None,
 ):
     """
     Apply trained sliding window models to an AnnData object (Lightning or Sklearn).
@@ -24,11 +27,11 @@ def sliding_window_inference(
         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,
@@ -38,31 +41,31 @@ def sliding_window_inference(
                     batch_size=batch_size,
                     window_start=window_start,
                     window_size=window_size,
-                    inference_mode=True
+                    inference_mode=True,
                 )
                 datamodule.setup()
 
                 # Extract model + detect type
-                model = run['model']
+                model = run["model"]
 
                 # Lightning models
-                if hasattr(run, 'trainer') or 'trainer' in run:
-                    trainer = run['trainer']
+                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}"
+                        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}"
+                        prefix=f"{model_name}_w{window_size}_c{center_varname}",
                     )
 
     print("Inference complete across all models.")
@@ -77,27 +80,36 @@ def sliding_window_inference(
                 prefix = f"{model_name}_w{window_size}_c{center_varname}"
                 # Use full key for uniqueness
                 key = prefix
-                model_wrappers[key] = run['model']
+                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 = 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')
+    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)
+    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()]
+        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:
@@ -111,4 +123,4 @@ def sliding_window_inference(
 
         print(f"Cleaned up {len(prefixes)} model prefixes from AnnData.")
 
-    return df
+    return df

smftools/machine_learning/models/__init__.py

@@ -1,9 +1,16 @@
+from __future__ import annotations
+
 from .base import BaseTorchModel
-from .mlp import MLPClassifier
 from .cnn import CNNClassifier
-from .rnn import RNNClassifier
-from .transformer import BaseTransformer, TransformerClassifier, DANNTransformerClassifier, MaskedTransformerPretrainer
+from .lightning_base import TorchClassifierWrapper
+from .mlp import MLPClassifier
 from .positional import PositionalEncoding
+from .rnn import RNNClassifier
+from .sklearn_models import SklearnModelWrapper
+from .transformer import (
+    BaseTransformer,
+    DANNTransformerClassifier,
+    MaskedTransformerPretrainer,
+    TransformerClassifier,
+)
 from .wrappers import ScaledModel
-from .lightning_base import TorchClassifierWrapper
-from .sklearn_models import SklearnModelWrapper

smftools/machine_learning/models/base.py

@@ -1,17 +1,25 @@
-import torch
-import torch.nn as nn
+from __future__ import annotations
+
 import numpy as np
+
+from smftools.optional_imports import require
+
 from ..utils.device import detect_device
 
+torch = require("torch", extra="ml-base", purpose="ML base models")
+nn = torch.nn
+
+
 class BaseTorchModel(nn.Module):
     """
     Minimal base class for torch models that:
     - Stores device and dropout regularization
     """
+
     def __init__(self, dropout_rate=0.0):
         super().__init__()
-        self.device = detect_device() # detects available devices
-        self.dropout_rate = dropout_rate # default dropout rate to be used in regularization.
+        self.device = detect_device()  # detects available devices
+        self.dropout_rate = dropout_rate  # default dropout rate to be used in regularization.
 
     def compute_saliency(
         self,
@@ -21,11 +29,11 @@ class BaseTorchModel(nn.Module):
         smoothgrad=False,
         smooth_samples=25,
         smooth_noise=0.1,
-        signed=True
+        signed=True,
     ):
         """
         Compute vanilla saliency or SmoothGrad saliency.
-        
+
         Arguments:
         ----------
         x : torch.Tensor
@@ -43,7 +51,7 @@ class BaseTorchModel(nn.Module):
         """
         self.eval()
         x = x.clone().detach().requires_grad_(True)
-        
+
         if smoothgrad:
             saliency_accum = torch.zeros_like(x)
             for i in range(smooth_samples):
@@ -56,7 +64,7 @@ class BaseTorchModel(nn.Module):
                 if logits.shape[1] == 1:
                     scores = logits.squeeze(1)
                 else:
-                    scores = logits[torch.arange(x.shape[0]), target_class]                
+                    scores = logits[torch.arange(x.shape[0]), target_class]
                 scores.sum().backward()
                 saliency_accum += x_noisy.grad.detach()
             saliency = saliency_accum / smooth_samples
@@ -69,17 +77,17 @@ class BaseTorchModel(nn.Module):
                 scores = logits[torch.arange(x.shape[0]), target_class]
             scores.sum().backward()
             saliency = x.grad.detach()
-        
+
         if not signed:
             saliency = saliency.abs()
-    
+
         if reduction == "sum" and x.ndim == 3:
             return saliency.sum(dim=-1)
         elif reduction == "mean" and x.ndim == 3:
             return saliency.mean(dim=-1)
         else:
             return saliency
-        
+
     def compute_gradient_x_input(self, x, target_class=None):
         """
         Computes gradient × input attribution.
@@ -118,22 +126,11 @@ class BaseTorchModel(nn.Module):
             baseline = torch.zeros_like(x)
 
         attributions, delta = ig.attribute(
-            x,
-            baselines=baseline,
-            target=target_class,
-            n_steps=steps,
-            return_convergence_delta=True
+            x, baselines=baseline, target=target_class, n_steps=steps, return_convergence_delta=True
         )
         return attributions, delta
 
-    def compute_deeplift(
-        self,
-        x,
-        baseline=None,
-        target_class=None,
-        reduction="sum",
-        signed=True
-    ):
+    def compute_deeplift(self, x, baseline=None, target_class=None, reduction="sum", signed=True):
         """
         Compute DeepLIFT scores using captum.
 
@@ -158,21 +155,15 @@ class BaseTorchModel(nn.Module):
 
         if not signed:
             attr = attr.abs()
-        
+
         if reduction == "sum" and x.ndim == 3:
             return attr.sum(dim=-1)
         elif reduction == "mean" and x.ndim == 3:
             return attr.mean(dim=-1)
         else:
             return attr
-        
-    def compute_occlusion(
-        self,
-        x,
-        target_class=None,
-        window_size=5,
-        baseline=None
-    ):
+
+    def compute_occlusion(self, x, target_class=None, window_size=5, baseline=None):
         """
         Computes per-sample occlusion attribution.
         Supports 2D [B, S] or 3D [B, S, D] inputs.
@@ -208,9 +199,7 @@ class BaseTorchModel(nn.Module):
                     x_occluded[left:right, :] = baseline[left:right, :]
 
                 x_tensor = torch.tensor(
-                    x_occluded,
-                    device=self.device,
-                    dtype=torch.float32
+                    x_occluded, device=self.device, dtype=torch.float32
                 ).unsqueeze(0)
 
                 logits = self.forward(x_tensor)
@@ -235,7 +224,7 @@ class BaseTorchModel(nn.Module):
         device="cpu",
         target_class=None,
         normalize=True,
-        signed=True
+        signed=True,
     ):
         """
         Apply a chosen attribution method to a dataloader and store results in adata.
@@ -252,7 +241,9 @@ class BaseTorchModel(nn.Module):
                 attr = model.compute_saliency(x, target_class=target_class, signed=signed)
 
             elif method == "smoothgrad":
-                attr = model.compute_saliency(x, smoothgrad=True, target_class=target_class, signed=signed)
+                attr = model.compute_saliency(
+                    x, smoothgrad=True, target_class=target_class, signed=signed
+                )
 
             elif method == "IG":
                 attributions, delta = model.compute_integrated_gradients(
@@ -261,15 +252,15 @@ class BaseTorchModel(nn.Module):
                 attr = attributions
 
             elif method == "deeplift":
-                attr = model.compute_deeplift(x, baseline=baseline, target_class=target_class, signed=signed)
+                attr = model.compute_deeplift(
+                    x, baseline=baseline, target_class=target_class, signed=signed
+                )
 
             elif method == "gradxinput":
                 attr = model.compute_gradient_x_input(x, target_class=target_class)
 
             elif method == "occlusion":
-                attr = model.compute_occlusion(
-                    x, target_class=target_class, baseline=baseline
-                )
+                attr = model.compute_occlusion(x, target_class=target_class, baseline=baseline)
 
             else:
                 raise ValueError(f"Unknown method {method}")
@@ -292,4 +283,4 @@ class BaseTorchModel(nn.Module):
             return target_class
         if logits.shape[1] == 1:
             return (logits > 0).long().squeeze(1)
-        return logits.argmax(dim=1)
+        return logits.argmax(dim=1)