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

smftools/machine_learning/inference/inference_utils.py

@@ -1,5 +1,6 @@
 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 +9,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 +21,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,11 @@
-import torch
-import pandas as pd
 import numpy as np
-from pytorch_lightning import Trainer
+import pandas as pd
+import torch
+
 from .inference_utils import annotate_split_column
 
-def run_lightning_inference(
-    adata,
-    model,
-    datamodule,
-    trainer,
-    prefix="model",
-    devices=1
-):
+
+def run_lightning_inference(adata, model, datamodule, trainer, prefix="model", devices=1):
     """
     Run inference on AnnData using TorchClassifierWrapper + AnnDataModule (in inference mode).
     """
@@ -57,7 +51,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 +61,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,10 @@
-import pandas as pd
 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 +40,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

@@ -3,16 +3,17 @@ 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 +25,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 +39,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 +78,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 +121,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,14 @@
 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,20 @@
+import numpy as np
 import torch
 import torch.nn as nn
-import numpy as np
+
 from ..utils.device import detect_device
 
+
 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 +24,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 +46,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 +59,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 +72,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 +121,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 +150,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 +194,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 +219,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 +236,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 +247,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 +278,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)

smftools/machine_learning/models/cnn.py

@@ -1,7 +1,9 @@
+import numpy as np
 import torch
 import torch.nn as nn
+
 from .base import BaseTorchModel
-import numpy as np
+
 
 class CNNClassifier(BaseTorchModel):
     def __init__(
@@ -15,7 +17,7 @@ class CNNClassifier(BaseTorchModel):
         use_pooling=False,
         dropout=0.2,
         gradcam_layer_idx=-1,
-        **kwargs
+        **kwargs,
     ):
         super().__init__(**kwargs)
         self.name = "CNNClassifier"
@@ -30,7 +32,9 @@ class CNNClassifier(BaseTorchModel):
 
         # Build conv layers
         for out_channels, ksize in zip(conv_channels, kernel_sizes):
-            layers.append(nn.Conv1d(in_channels, out_channels, kernel_size=ksize, padding=ksize // 2))
+            layers.append(
+                nn.Conv1d(in_channels, out_channels, kernel_size=ksize, padding=ksize // 2)
+            )
             if use_batchnorm:
                 layers.append(nn.BatchNorm1d(out_channels))
             layers.append(nn.ReLU())
@@ -76,7 +80,7 @@ class CNNClassifier(BaseTorchModel):
         x = self.conv(x)
         x = x.view(x.size(0), -1)
         return self.fc(x)
-    
+
     def _register_gradcam_hooks(self):
         def forward_hook(module, input, output):
             self.gradcam_activations = output.detach()
@@ -97,15 +101,15 @@ class CNNClassifier(BaseTorchModel):
         self.eval()  # disable dropout etc.
 
         output = self.forward(x)  # shape (B, C) or (B, 1)
-        
+
         if class_idx is None:
             class_idx = output.argmax(dim=1)
-        
+
         if output.shape[1] == 1:
             target = output.view(-1)  # shape (B,)
         else:
             target = output[torch.arange(output.shape[0]), class_idx]
-        
+
         target.sum().backward(retain_graph=True)
 
         # restore training mode
@@ -114,16 +118,16 @@ class CNNClassifier(BaseTorchModel):
 
         # get activations and gradients (set these via forward hook!)
         activations = self.gradcam_activations  # (B, C, L)
-        gradients = self.gradcam_gradients      # (B, C, L)
+        gradients = self.gradcam_gradients  # (B, C, L)
 
         weights = gradients.mean(dim=2, keepdim=True)  # (B, C, 1)
-        cam = (weights * activations).sum(dim=1)       # (B, L)
+        cam = (weights * activations).sum(dim=1)  # (B, L)
 
         cam = torch.relu(cam)
         cam = cam / (cam.max(dim=1, keepdim=True).values + 1e-6)
 
         return cam
-    
+
     def apply_gradcam_to_adata(self, dataloader, adata, obsm_key="gradcam", device="cpu"):
         self.to(device)
         self.eval()
@@ -135,4 +139,4 @@ class CNNClassifier(BaseTorchModel):
             cams.append(cam_batch.cpu().numpy())
 
         cams = np.concatenate(cams, axis=0)  # shape: [n_obs, input_len]
-        adata.obsm[obsm_key] = cams
+        adata.obsm[obsm_key] = cams