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

smftools/machine_learning/models/transformer.py

@@ -1,22 +1,31 @@
-import torch
-import torch.nn as nn
-from .base import BaseTorchModel
-from .positional import PositionalEncoding 
-from ..utils.grl import grad_reverse
+from __future__ import annotations
+
 import numpy as np
 
+from smftools.optional_imports import require
+
+from ..utils.grl import grad_reverse
+from .base import BaseTorchModel
+from .positional import PositionalEncoding
+
+torch = require("torch", extra="ml-base", purpose="Transformer models")
+nn = torch.nn
+
+
 class TransformerEncoderLayerWithAttn(nn.TransformerEncoderLayer):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     def forward(self, src, src_mask=None, is_causal=False, src_key_padding_mask=None):
         self_attn_output, attn_weights = self.self_attn(
-            src, src, src,
+            src,
+            src,
+            src,
             attn_mask=src_mask,
             key_padding_mask=src_key_padding_mask,
             need_weights=True,
             average_attn_weights=False,  # preserve [B, num_heads, S, S]
-            is_causal=is_causal
+            is_causal=is_causal,
         )
         src = src + self.dropout1(self_attn_output)
         src = self.norm1(src)
@@ -27,18 +36,21 @@ class TransformerEncoderLayerWithAttn(nn.TransformerEncoderLayer):
         # Save attention weights to module
         self.attn_weights = attn_weights  # Save to layer
         return src
-    
+
+
 class BaseTransformer(BaseTorchModel):
-    def __init__(self, 
-                 input_dim=1, 
-                 model_dim=64, 
-                 num_heads=4, 
-                 num_layers=2, 
-                 dropout=0.2,
-                 seq_len=None, 
-                 use_learnable_pos=False, 
-                 use_cls_token=True,
-                 **kwargs):
+    def __init__(
+        self,
+        input_dim=1,
+        model_dim=64,
+        num_heads=4,
+        num_layers=2,
+        dropout=0.2,
+        seq_len=None,
+        use_learnable_pos=False,
+        use_cls_token=True,
+        **kwargs,
+    ):
         super().__init__(**kwargs)
         # Input FC layer to map D_input to D_model
         self.model_dim = model_dim
@@ -52,7 +64,9 @@ class BaseTransformer(BaseTorchModel):
 
         if use_learnable_pos:
             assert seq_len is not None, "Must provide seq_len if use_learnable_pos=True"
-            self.pos_embed = nn.Parameter(torch.randn(seq_len + (1 if use_cls_token else 0), model_dim))
+            self.pos_embed = nn.Parameter(
+                torch.randn(seq_len + (1 if use_cls_token else 0), model_dim)
+            )
             self.pos_encoder = None
         else:
             self.pos_encoder = PositionalEncoding(model_dim)
@@ -62,7 +76,13 @@ class BaseTransformer(BaseTorchModel):
             self.cls_token = nn.Parameter(torch.zeros(1, 1, model_dim))  # (1, 1, D)
 
         # Specify the transformer encoder structure
-        encoder_layer = TransformerEncoderLayerWithAttn(d_model=model_dim, nhead=num_heads, batch_first=True, dim_feedforward=self.ff_dim, dropout=self.dropout)
+        encoder_layer = TransformerEncoderLayerWithAttn(
+            d_model=model_dim,
+            nhead=num_heads,
+            batch_first=True,
+            dim_feedforward=self.ff_dim,
+            dropout=self.dropout,
+        )
         # Stack the transformer encoder layers
         self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
 
@@ -95,7 +115,7 @@ class BaseTransformer(BaseTorchModel):
             x = torch.cat([cls, x], dim=1)  # (B, S+1, D)
 
         if self.pos_embed is not None:
-            x = x + self.pos_embed.unsqueeze(0)[:, :x.shape[1], :]
+            x = x + self.pos_embed.unsqueeze(0)[:, : x.shape[1], :]
         elif self.pos_encoder is not None:
             x = self.pos_encoder(x)
 
@@ -106,8 +126,8 @@ class BaseTransformer(BaseTorchModel):
 
         encoded = self.transformer(x)
         return encoded
-    
-    def compute_attn_grad(self, reduction='mean'):
+
+    def compute_attn_grad(self, reduction="mean"):
         """
         Computes attention × gradient scores across layers.
         Returns: [B, S] tensor of importance scores
@@ -116,19 +136,19 @@ class BaseTransformer(BaseTorchModel):
         for attn, grad in zip(self.attn_weights, self.attn_grads):
             # attn: [B, H, S, S]
             # grad: [B, S, D]
-            attn = attn.mean(dim=1)            # [B, S, S]
-            grad_norm = grad.norm(dim=-1)      # [B, S]
+            attn = attn.mean(dim=1)  # [B, S, S]
+            grad_norm = grad.norm(dim=-1)  # [B, S]
             attn_grad_score = (attn * grad_norm.unsqueeze(1)).sum(dim=-1)  # [B, S]
             scores.append(attn_grad_score)
 
         # Combine across layers
         stacked = torch.stack(scores, dim=0)  # [L, B, S]
         if reduction == "mean":
-            return stacked.mean(dim=0)        # [B, S]
+            return stacked.mean(dim=0)  # [B, S]
         elif reduction == "sum":
-            return stacked.sum(dim=0)         # [B, S]
+            return stacked.sum(dim=0)  # [B, S]
         else:
-            return stacked                    # [L, B, S]
+            return stacked  # [L, B, S]
 
     def compute_rollout(self):
         """
@@ -143,9 +163,9 @@ class BaseTransformer(BaseTorchModel):
             attn_heads = attn_heads + torch.eye(S, device=device).unsqueeze(0)  # add residual
             attn_heads = attn_heads / attn_heads.sum(dim=-1, keepdim=True).clamp(min=1e-6)
             rollout = torch.bmm(attn_heads, rollout)  # [B, S, S]
-        
+
         return rollout  # [B, S, S]
-    
+
     def reset_attn_buffers(self):
         self.attn_weights = []
         self.attn_grads = []
@@ -158,11 +178,15 @@ class BaseTransformer(BaseTorchModel):
         if head_idx is not None:
             attn = attn[:, head_idx]  # [B, S, S]
         return attn
-    
-    def apply_attn_interpretations_to_adata(self, dataloader, adata, 
-                                            obsm_key_grad="attn_grad", 
-                                            obsm_key_rollout="attn_rollout",
-                                            device="cpu"):
+
+    def apply_attn_interpretations_to_adata(
+        self,
+        dataloader,
+        adata,
+        obsm_key_grad="attn_grad",
+        obsm_key_rollout="attn_rollout",
+        device="cpu",
+    ):
         self.to(device)
         self.eval()
         grad_maps = []
@@ -193,12 +217,10 @@ class BaseTransformer(BaseTorchModel):
         # add per-row normalized version
         grad_normed = grad_concat / (np.max(grad_concat, axis=1, keepdims=True) + 1e-8)
         adata.obsm[f"{obsm_key_grad}_normalized"] = grad_normed
-    
+
+
 class TransformerClassifier(BaseTransformer):
-    def __init__(self, 
-                 input_dim, 
-                 num_classes, 
-                 **kwargs):
+    def __init__(self, input_dim, num_classes, **kwargs):
         super().__init__(input_dim, **kwargs)
         # Classification head
         output_size = 1 if num_classes == 2 else num_classes
@@ -215,7 +237,7 @@ class TransformerClassifier(BaseTransformer):
             x = x.unsqueeze(0).unsqueeze(-1)  # just in case (S,) → (1, S, 1)
         else:
             pass
-        encoded = self.encode(x) # -> (B, S, D_model)
+        encoded = self.encode(x)  # -> (B, S, D_model)
         if self.use_cls_token:
             pooled = encoded[:, 0]  # (B, D)
         else:
@@ -223,14 +245,13 @@ class TransformerClassifier(BaseTransformer):
 
         out = self.cls_head(pooled)  # (B, C)
         return out
-    
+
+
 class DANNTransformerClassifier(TransformerClassifier):
     def __init__(self, input_dim, model_dim, num_classes, n_domains, **kwargs):
         super().__init__(input_dim, model_dim, num_classes, **kwargs)
         self.domain_classifier = nn.Sequential(
-            nn.Linear(model_dim, 128),
-            nn.ReLU(),
-            nn.Linear(128, n_domains)
+            nn.Linear(model_dim, 128), nn.ReLU(), nn.Linear(128, n_domains)
         )
 
     def forward(self, x, alpha=1.0):
@@ -242,6 +263,7 @@ class DANNTransformerClassifier(TransformerClassifier):
 
         return class_logits, domain_logits
 
+
 class MaskedTransformerPretrainer(BaseTransformer):
     def __init__(self, input_dim, model_dim, num_heads=4, num_layers=2, **kwargs):
         super().__init__(input_dim, model_dim, num_heads, num_layers, **kwargs)
@@ -254,12 +276,13 @@ class MaskedTransformerPretrainer(BaseTransformer):
         """
         if x.dim() == 2:
             x = x.unsqueeze(-1)
-        encoded = self.encode(x, mask=mask) # -> (B, S, D_model)
-        return self.decoder(encoded) # -> (B, D_input)
-    
+        encoded = self.encode(x, mask=mask)  # -> (B, S, D_model)
+        return self.decoder(encoded)  # -> (B, D_input)
+
+
 class DANNTransformer(BaseTransformer):
-    """
-    """
+    """ """
+
     def __init__(self, seq_len, model_dim, n_heads, n_layers, n_domains):
         super().__init__(
             input_dim=1,  # 1D scalar input per token
@@ -267,7 +290,7 @@ class DANNTransformer(BaseTransformer):
             num_heads=n_heads,
             num_layers=n_layers,
             seq_len=seq_len,
-            use_learnable_pos=True  # enables learnable pos_embed in base
+            use_learnable_pos=True,  # enables learnable pos_embed in base
         )
 
         # Reconstruction head
@@ -275,9 +298,7 @@ class DANNTransformer(BaseTransformer):
 
         # Domain classification head
         self.domain_classifier = nn.Sequential(
-            nn.Linear(model_dim, 128),
-            nn.ReLU(),
-            nn.Linear(128, n_domains)
+            nn.Linear(model_dim, 128), nn.ReLU(), nn.Linear(128, n_domains)
         )
 
     def forward(self, x, alpha=1.0):
@@ -300,4 +321,3 @@ class DANNTransformer(BaseTransformer):
         domain_logits = self.domain_classifier(rev)  # (B, n_batches)
 
         return recon, domain_logits
-    

smftools/machine_learning/models/wrappers.py

@@ -1,5 +1,10 @@
-import torch
-import torch.nn as nn
+from __future__ import annotations
+
+from smftools.optional_imports import require
+
+torch = require("torch", extra="ml-base", purpose="model wrappers")
+nn = torch.nn
+
 
 class ScaledModel(nn.Module):
     def __init__(self, model, mean, std):
@@ -17,4 +22,4 @@ class ScaledModel(nn.Module):
             x = (x - mean[None, None, :]) / std[None, None, :]
         else:
             raise ValueError(f"Unsupported input shape {x.shape}")
-        return self.model(x)
+        return self.model(x)

smftools/machine_learning/training/__init__.py

@@ -1,2 +1,4 @@
-from .train_lightning_model import train_lightning_model, run_sliding_window_lightning_training
-from .train_sklearn_model import train_sklearn_model, run_sliding_window_sklearn_training
+from __future__ import annotations
+
+from .train_lightning_model import run_sliding_window_lightning_training, train_lightning_model
+from .train_sklearn_model import run_sliding_window_sklearn_training, train_sklearn_model

smftools/machine_learning/training/train_lightning_model.py

@@ -1,9 +1,21 @@
-import torch
-from pytorch_lightning import Trainer
-from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+from __future__ import annotations
+
+from smftools.optional_imports import require
+
 from ..data import AnnDataModule
 from ..models import TorchClassifierWrapper
 
+torch = require("torch", extra="ml-base", purpose="Lightning training")
+pytorch_lightning = require("pytorch_lightning", extra="ml-extended", purpose="Lightning training")
+pl_callbacks = require(
+    "pytorch_lightning.callbacks", extra="ml-extended", purpose="Lightning training"
+)
+
+Trainer = pytorch_lightning.Trainer
+EarlyStopping = pl_callbacks.EarlyStopping
+ModelCheckpoint = pl_callbacks.ModelCheckpoint
+
+
 def train_lightning_model(
     model,
     datamodule,
@@ -12,7 +24,7 @@ def train_lightning_model(
     monitor_metric="val_loss",
     checkpoint_path=None,
     evaluate_test=True,
-    devices=1
+    devices=1,
 ):
     """
     Takes a PyTorch Lightning Model and a Lightning DataLoader module to define a Lightning Trainer.
@@ -39,13 +51,15 @@ def train_lightning_model(
         EarlyStopping(monitor=monitor_metric, patience=patience, mode="min"),
     ]
     if checkpoint_path:
-        callbacks.append(ModelCheckpoint(
-            dirpath=checkpoint_path,
-            filename="{epoch}-{val_loss:.4f}",
-            monitor=monitor_metric,
-            save_top_k=1,
-            mode="min",
-        ))
+        callbacks.append(
+            ModelCheckpoint(
+                dirpath=checkpoint_path,
+                filename="{epoch}-{val_loss:.4f}",
+                monitor=monitor_metric,
+                save_top_k=1,
+                mode="min",
+            )
+        )
 
     # Trainer setup
     trainer = Trainer(
@@ -54,7 +68,7 @@ def train_lightning_model(
         accelerator=accelerator,
         devices=devices,
         log_every_n_steps=10,
-        enable_progress_bar=False
+        enable_progress_bar=False,
     )
 
     # Fit model with trainer
@@ -63,7 +77,7 @@ def train_lightning_model(
     # Test model (if applicable)
     if evaluate_test and hasattr(datamodule, "test_dataloader"):
         trainer.test(model, datamodule=datamodule)
-    
+
     # Return best checkpoint path
     best_ckpt = None
     for cb in callbacks:
@@ -72,6 +86,7 @@ def train_lightning_model(
 
     return trainer, best_ckpt
 
+
 def run_sliding_window_lightning_training(
     adata,
     tensor_source,
@@ -86,13 +101,13 @@ def run_sliding_window_lightning_training(
     stride,
     max_epochs=30,
     patience=5,
-    enforce_eval_balance: bool=False,
-    target_eval_freq: float=0.3,
-    max_eval_positive: int=None
+    enforce_eval_balance: bool = False,
+    target_eval_freq: float = 0.3,
+    max_eval_positive: int = None,
 ):
     input_len = adata.shape[1]
     results = {}
-    
+
     for start in range(0, input_len - window_size + 1, stride):
         center_idx = start + window_size // 2
         center_varname = adata.var_names[center_idx]
@@ -106,18 +121,22 @@ def run_sliding_window_lightning_training(
             label_col=label_col,
             batch_size=64,
             window_start=start,
-            window_size=window_size
+            window_size=window_size,
         )
         datamodule.setup()
 
         # Build model for this window
         model = model_class(window_size, num_classes)
         wrapper = TorchClassifierWrapper(
-            model, label_col=label_col, num_classes=num_classes,
+            model,
+            label_col=label_col,
+            num_classes=num_classes,
             class_names=class_names,
             class_weights=class_weights,
-            focus_class=focus_class, enforce_eval_balance=enforce_eval_balance,
-            target_eval_freq=target_eval_freq, max_eval_positive=max_eval_positive
+            focus_class=focus_class,
+            enforce_eval_balance=enforce_eval_balance,
+            target_eval_freq=target_eval_freq,
+            max_eval_positive=max_eval_positive,
         )
 
         # Train model
@@ -129,7 +148,7 @@ def run_sliding_window_lightning_training(
             "model": wrapper,
             "trainer": trainer,
             "checkpoint": ckpt,
-            "metrics": trainer.callback_metrics
+            "metrics": trainer.callback_metrics,
         }
-    
+
     return results

smftools/machine_learning/training/train_sklearn_model.py

@@ -1,16 +1,14 @@
+from __future__ import annotations
+
 from ..data import AnnDataModule
 from ..models import SklearnModelWrapper
 
-def train_sklearn_model(
-    model_wrapper, 
-    datamodule, 
-    evaluate_test=True, 
-    evaluate_val=False
-):
+
+def train_sklearn_model(model_wrapper, datamodule, evaluate_test=True, evaluate_val=False):
     """
     Fits a SklearnModelWrapper on the train split from datamodule.
     Evaluates on test and/or val set.
-    
+
     Parameters:
         model_wrapper: SklearnModelWrapper instance
         datamodule: AnnDataModule instance (with setup() method)
@@ -39,6 +37,7 @@ def train_sklearn_model(
 
     return metrics
 
+
 def run_sliding_window_sklearn_training(
     adata,
     tensor_source,
@@ -58,7 +57,7 @@ def run_sliding_window_sklearn_training(
     enforce_eval_balance=False,
     target_eval_freq=0.3,
     max_eval_positive=None,
-    **model_kwargs
+    **model_kwargs,
 ):
     """
     Sliding window training for sklearn models using AnnData.
@@ -86,29 +85,26 @@ def run_sliding_window_sklearn_training(
             train_frac=train_frac,
             val_frac=val_frac,
             test_frac=test_frac,
-            random_seed=random_seed
+            random_seed=random_seed,
         )
         datamodule.setup()
 
         # Build model wrapper
         sklearn_model = model_class(**model_kwargs)
         wrapper = SklearnModelWrapper(
-            sklearn_model, 
+            sklearn_model,
             num_classes=num_classes,
             label_col=label_col,
             class_names=class_names,
             focus_class=focus_class,
             enforce_eval_balance=enforce_eval_balance,
             target_eval_freq=target_eval_freq,
-            max_eval_positive=max_eval_positive
+            max_eval_positive=max_eval_positive,
         )
 
         # Fit and evaluate
         metrics = train_sklearn_model(wrapper, datamodule, evaluate_test=True, evaluate_val=False)
 
-        results[center_varname] = {
-            "model": wrapper,
-            "metrics": metrics
-        }
+        results[center_varname] = {"model": wrapper, "metrics": metrics}
 
     return results

smftools/machine_learning/utils/__init__.py

@@ -1,2 +1,4 @@
+from __future__ import annotations
+
 from .device import detect_device
-from .grl import GradReverse
+from .grl import GradReverse

smftools/machine_learning/utils/device.py

@@ -1,10 +1,17 @@
-import torch
+from __future__ import annotations
+
+from smftools.optional_imports import require
+
+torch = require("torch", extra="ml-base", purpose="device selection")
+
 
 def detect_device():
     device = (
-        torch.device('cuda') if torch.cuda.is_available() else
-        torch.device('mps') if torch.backends.mps.is_available() else
-        torch.device('cpu')
+        torch.device("cuda")
+        if torch.cuda.is_available()
+        else torch.device("mps")
+        if torch.backends.mps.is_available()
+        else torch.device("cpu")
     )
     print(f"Detected device: {device}")
-    return device
+    return device

smftools/machine_learning/utils/grl.py

@@ -1,4 +1,9 @@
-import torch
+from __future__ import annotations
+
+from smftools.optional_imports import require
+
+torch = require("torch", extra="ml-base", purpose="gradient reversal layers")
+
 
 class GradReverse(torch.autograd.Function):
     @staticmethod
@@ -10,5 +15,6 @@ class GradReverse(torch.autograd.Function):
     def backward(ctx, grad_output):
         return -ctx.alpha * grad_output, None
 
+
 def grad_reverse(x, alpha=1.0):
-    return GradReverse.apply(x, alpha)
+    return GradReverse.apply(x, alpha)