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

smftools/machine_learning/models/cnn.py

@@ -1,8 +1,15 @@
-import torch
-import torch.nn as nn
-from .base import BaseTorchModel
+from __future__ import annotations
+
 import numpy as np
 
+from smftools.optional_imports import require
+
+from .base import BaseTorchModel
+
+torch = require("torch", extra="ml-base", purpose="CNN models")
+nn = torch.nn
+
+
 class CNNClassifier(BaseTorchModel):
     def __init__(
         self,
@@ -15,7 +22,7 @@ class CNNClassifier(BaseTorchModel):
         use_pooling=False,
         dropout=0.2,
         gradcam_layer_idx=-1,
-        **kwargs
+        **kwargs,
     ):
         super().__init__(**kwargs)
         self.name = "CNNClassifier"
@@ -30,7 +37,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 +85,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 +106,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 +123,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 +144,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

smftools/machine_learning/models/lightning_base.py

@@ -1,11 +1,22 @@
-import torch
-import pytorch_lightning as pl
-import matplotlib.pyplot as plt
-from sklearn.metrics import (
-    roc_auc_score, precision_recall_curve, auc, f1_score, confusion_matrix, roc_curve
-)
+from __future__ import annotations
+
 import numpy as np
 
+from smftools.optional_imports import require
+
+plt = require("matplotlib.pyplot", extra="plotting", purpose="model evaluation plots")
+pl = require("pytorch_lightning", extra="ml-extended", purpose="Lightning models")
+torch = require("torch", extra="ml-base", purpose="Lightning models")
+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 TorchClassifierWrapper(pl.LightningModule):
     """
     A Pytorch Lightning wrapper for PyTorch classifiers.
@@ -16,25 +27,26 @@ class TorchClassifierWrapper(pl.LightningModule):
     - Can pass the index of the class label to use as the focus class when calculating precision/recall.
     - Contains a prediction step to run inference with.
     """
+
     def __init__(
         self,
         model: torch.nn.Module,
         label_col: str,
         num_classes: int,
-        class_names: list=None,
+        class_names: list = None,
         optimizer_cls=torch.optim.AdamW,
         optimizer_kwargs=None,
         criterion_kwargs=None,
         lr: float = 1e-3,
         focus_class: int = 1,  # used for binary or multiclass precision-recall
         class_weights=None,
-        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,
     ):
         super().__init__()
         self.model = model
-        self.save_hyperparameters(ignore=['model'])  # logs all except actual model instance
+        self.save_hyperparameters(ignore=["model"])  # logs all except actual model instance
         self.optimizer_cls = optimizer_cls
         self.optimizer_kwargs = optimizer_kwargs or {"weight_decay": 1e-4}
         self.criterion = None
@@ -57,14 +69,17 @@ class TorchClassifierWrapper(pl.LightningModule):
                 if torch.is_tensor(class_weights[self.focus_class]):
                     self.criterion_kwargs["pos_weight"] = class_weights[self.focus_class]
                 else:
-                    self.criterion_kwargs["pos_weight"] = torch.tensor(class_weights[self.focus_class], dtype=torch.float32, device=self.device)
+                    self.criterion_kwargs["pos_weight"] = torch.tensor(
+                        class_weights[self.focus_class], dtype=torch.float32, device=self.device
+                    )
             else:
                 # CrossEntropyLoss expects weight tensor of size C
                 if torch.is_tensor(class_weights):
                     self.criterion_kwargs["weight"] = class_weights
                 else:
-                    self.criterion_kwargs["weight"] = torch.tensor(class_weights, dtype=torch.float32)
-
+                    self.criterion_kwargs["weight"] = torch.tensor(
+                        class_weights, dtype=torch.float32
+                    )
 
         self._val_outputs = []
         self._test_outputs = []
@@ -78,12 +93,20 @@ class TorchClassifierWrapper(pl.LightningModule):
 
     def _init_criterion(self):
         if self.num_classes == 2:
-            if "pos_weight" in self.criterion_kwargs and not torch.is_tensor(self.criterion_kwargs["pos_weight"]):
-                self.criterion_kwargs["pos_weight"] = torch.tensor(self.criterion_kwargs["pos_weight"], dtype=torch.float32, device=self.device)
+            if "pos_weight" in self.criterion_kwargs and not torch.is_tensor(
+                self.criterion_kwargs["pos_weight"]
+            ):
+                self.criterion_kwargs["pos_weight"] = torch.tensor(
+                    self.criterion_kwargs["pos_weight"], dtype=torch.float32, device=self.device
+                )
             self.criterion = torch.nn.BCEWithLogitsLoss(**self.criterion_kwargs)
         else:
-            if "weight" in self.criterion_kwargs and not torch.is_tensor(self.criterion_kwargs["weight"]):
-                self.criterion_kwargs["weight"] = torch.tensor(self.criterion_kwargs["weight"], dtype=torch.float32, device=self.device)
+            if "weight" in self.criterion_kwargs and not torch.is_tensor(
+                self.criterion_kwargs["weight"]
+            ):
+                self.criterion_kwargs["weight"] = torch.tensor(
+                    self.criterion_kwargs["weight"], dtype=torch.float32, device=self.device
+                )
             self.criterion = torch.nn.CrossEntropyLoss(**self.criterion_kwargs)
 
     def _resolve_focus_class(self, focus_class):
@@ -93,11 +116,13 @@ class TorchClassifierWrapper(pl.LightningModule):
             if self.class_names is None:
                 raise ValueError("class_names must be provided if focus_class is a string.")
             if focus_class not in self.class_names:
-                raise ValueError(f"focus_class '{focus_class}' not found in class_names {self.class_names}.")
+                raise ValueError(
+                    f"focus_class '{focus_class}' not found in class_names {self.class_names}."
+                )
             return self.class_names.index(focus_class)
         else:
             raise ValueError(f"focus_class must be int or str, got {type(focus_class)}")
-        
+
     def set_training_indices(self, datamodule):
         """
         Store obs_names for train/val/test subsets used during training.
@@ -140,7 +165,7 @@ class TorchClassifierWrapper(pl.LightningModule):
         self.log_dict({"val_loss": loss, "val_acc": acc}, prog_bar=False)
         self._val_outputs.append((logits.detach(), y.detach()))
         return loss
-    
+
     def test_step(self, batch, batch_idx):
         """
         Test step for a batch through the Lightning Trainer.
@@ -189,7 +214,7 @@ class TorchClassifierWrapper(pl.LightningModule):
             return self.criterion(logits.view(-1, 1), y)
         else:
             return self.criterion(logits, y)
-        
+
     def _get_probs(self, logits):
         """
         A helper function for getting class probabilities for binary vs multiclass classifications.
@@ -207,8 +232,10 @@ class TorchClassifierWrapper(pl.LightningModule):
             return (torch.sigmoid(logits.view(-1)) >= 0.5).long()
         else:
             return logits.argmax(dim=1)
-        
-    def _subsample_for_fixed_positive_frequency(self, y_true, probs, target_freq=0.3, max_positive=None):
+
+    def _subsample_for_fixed_positive_frequency(
+        self, y_true, probs, target_freq=0.3, max_positive=None
+    ):
         pos_idx = np.where(y_true == self.focus_class)[0]
         neg_idx = np.where(y_true != self.focus_class)[0]
 
@@ -216,16 +243,20 @@ class TorchClassifierWrapper(pl.LightningModule):
         max_positives_possible = len(pos_idx)
 
         # maximum achievable positive class frequency
-        max_possible_freq = max_positives_possible / (max_positives_possible + max_negatives_possible)
+        max_possible_freq = max_positives_possible / (
+            max_positives_possible + max_negatives_possible
+        )
 
         if target_freq > max_possible_freq:
             target_freq = max_possible_freq  # clip if you ask for impossible freq
 
         # now calculate positive count
-        num_pos_target = min(int(target_freq * max_negatives_possible / (1 - target_freq)), max_positives_possible)
+        num_pos_target = min(
+            int(target_freq * max_negatives_possible / (1 - target_freq)), max_positives_possible
+        )
         num_neg_target = int(num_pos_target * (1 - target_freq) / target_freq)
         num_neg_target = min(num_neg_target, max_negatives_possible)
-        
+
         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)
 
@@ -235,7 +266,7 @@ class TorchClassifierWrapper(pl.LightningModule):
         actual_freq = len(pos_sampled) / len(sampled_idx)
 
         return sampled_idx
-        
+
     def _log_classification_metrics(self, logits, targets, prefix="val"):
         """
         A helper function for logging validation and testing split model evaluations.
@@ -252,9 +283,12 @@ class TorchClassifierWrapper(pl.LightningModule):
         num_pos = binary_focus.sum()
 
         # Subsample if you want to enforce a fixed proportion of the positive class
-        if prefix == 'test' and self.enforce_eval_balance:
+        if prefix == "test" and self.enforce_eval_balance:
             sampled_idx = self._subsample_for_fixed_positive_frequency(
-                y_true, probs, target_freq=self.target_eval_freq, max_positive=self.max_eval_positive
+                y_true,
+                probs,
+                target_freq=self.target_eval_freq,
+                max_positive=self.max_eval_positive,
             )
             y_true = y_true[sampled_idx]
             probs = probs[sampled_idx]
@@ -289,7 +323,7 @@ class TorchClassifierWrapper(pl.LightningModule):
         cm = confusion_matrix(y_true, preds)
 
         # Save attributes for later plotting
-        if prefix == 'test':
+        if prefix == "test":
             self.test_roc_curve = (fpr, tpr)
             self.test_pr_curve = (rc, pr)
             self.test_roc_auc = roc_auc
@@ -298,19 +332,21 @@ class TorchClassifierWrapper(pl.LightningModule):
             self.test_num_pos = num_pos
             self.test_acc = acc
             self.test_f1 = f1
-        elif prefix == 'val':
+        elif prefix == "val":
             pass
 
         # Logging
-        self.log_dict({
-            f"{prefix}_acc": acc,
-            f"{prefix}_f1": f1,
-            f"{prefix}_auc": roc_auc,
-            f"{prefix}_pr_auc": pr_auc,
-            f"{prefix}_pr_auc_norm": pr_auc_norm,
-            f"{prefix}_pos_freq": pos_freq,
-            f"{prefix}_num_pos": num_pos
-        })
+        self.log_dict(
+            {
+                f"{prefix}_acc": acc,
+                f"{prefix}_f1": f1,
+                f"{prefix}_auc": roc_auc,
+                f"{prefix}_pr_auc": pr_auc,
+                f"{prefix}_pr_auc_norm": pr_auc_norm,
+                f"{prefix}_pos_freq": pos_freq,
+                f"{prefix}_num_pos": num_pos,
+            }
+        )
         setattr(self, f"{prefix}_confusion_matrix", cm)
 
     def _plot_roc_pr_curves(self, logits, targets):
@@ -334,7 +370,7 @@ class TorchClassifierWrapper(pl.LightningModule):
         pos_freq = self.test_pos_freq
         plt.subplot(1, 2, 2)
         plt.plot(rc, pr, label=f"PR AUC={pr_auc:.3f}")
-        plt.axhline(pos_freq, linestyle='--', color="gray")
+        plt.axhline(pos_freq, linestyle="--", color="gray")
         plt.xlabel("Recall")
         plt.ylabel("Precision")
         plt.ylim(0, 1.05)

smftools/machine_learning/models/mlp.py

@@ -1,9 +1,22 @@
-import torch
-import torch.nn as nn
+from __future__ import annotations
+
+from smftools.optional_imports import require
+
 from .base import BaseTorchModel
-    
+
+nn = require("torch.nn", extra="ml-base", purpose="MLP models")
+
+
 class MLPClassifier(BaseTorchModel):
-    def __init__(self, input_dim, num_classes=2, hidden_dims=[64, 64], dropout=0.2, use_batchnorm=True, **kwargs):
+    def __init__(
+        self,
+        input_dim,
+        num_classes=2,
+        hidden_dims=[64, 64],
+        dropout=0.2,
+        use_batchnorm=True,
+        **kwargs,
+    ):
         super().__init__(**kwargs)
         layers = []
         in_dim = input_dim
@@ -23,4 +36,4 @@ class MLPClassifier(BaseTorchModel):
         self.model = nn.Sequential(*layers)
 
     def forward(self, x):
-        return self.model(x)
+        return self.model(x)

smftools/machine_learning/models/positional.py

@@ -1,6 +1,12 @@
+from __future__ import annotations
+
 import numpy as np
-import torch
-import torch.nn as nn
+
+from smftools.optional_imports import require
+
+torch = require("torch", extra="ml-base", purpose="positional encoding")
+nn = torch.nn
+
 
 class PositionalEncoding(nn.Module):
     def __init__(self, d_model, max_len=5000):
@@ -14,5 +20,5 @@ class PositionalEncoding(nn.Module):
         self.register_buffer("pe", pe)
 
     def forward(self, x):
-        x = x + self.pe[:, :x.size(1)]
-        return x
+        x = x + self.pe[:, : x.size(1)]
+        return x

smftools/machine_learning/models/rnn.py

@@ -1,7 +1,12 @@
-import torch
-import torch.nn as nn
+from __future__ import annotations
+
+from smftools.optional_imports import require
+
 from .base import BaseTorchModel
 
+nn = require("torch.nn", extra="ml-base", purpose="RNN models")
+
+
 class RNNClassifier(BaseTorchModel):
     def __init__(self, input_size, hidden_dim, num_classes, **kwargs):
         super().__init__(**kwargs)
@@ -14,4 +19,4 @@ class RNNClassifier(BaseTorchModel):
     def forward(self, x):
         x = x.unsqueeze(1)  # [B, 1, L] → for LSTM expecting batch_first
         _, (h_n, _) = self.lstm(x)  # h_n: [1, B, H]
-        return self.fc(h_n.squeeze(0))  # [B, H] → [B, num_classes]
+        return self.fc(h_n.squeeze(0))  # [B, H] → [B, num_classes]

smftools/machine_learning/models/sklearn_models.py

@@ -1,23 +1,35 @@
+from __future__ import annotations
+
 import numpy as np
-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="model 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 SklearnModelWrapper:
     """
     Unified sklearn wrapper matching TorchClassifierWrapper interface.
     """
+
     def __init__(
-        self, 
-        model, 
+        self,
+        model,
         label_col: str,
-        num_classes: int, 
-        class_names=None, 
-        focus_class: int=1,
-        enforce_eval_balance: bool=False,
-        target_eval_freq: float=0.3,
-        max_eval_positive=None
+        num_classes: int,
+        class_names=None,
+        focus_class: int = 1,
+        enforce_eval_balance: bool = False,
+        target_eval_freq: float = 0.3,
+        max_eval_positive=None,
     ):
         self.model = model
         self.label_col = label_col
@@ -37,7 +49,9 @@ class SklearnModelWrapper:
             if self.class_names is None:
                 raise ValueError("class_names must be provided if focus_class is a string.")
             if focus_class not in self.class_names:
-                raise ValueError(f"focus_class '{focus_class}' not found in class_names {self.class_names}.")
+                raise ValueError(
+                    f"focus_class '{focus_class}' not found in class_names {self.class_names}."
+                )
             return self.class_names.index(focus_class)
         else:
             raise ValueError(f"focus_class must be int or str, got {type(focus_class)}")
@@ -130,7 +144,7 @@ class SklearnModelWrapper:
             f"{prefix}_pr_auc": pr_auc,
             f"{prefix}_pr_auc_norm": pr_auc_norm,
             f"{prefix}_pos_freq": pos_freq,
-            f"{prefix}_num_pos": num_pos
+            f"{prefix}_num_pos": num_pos,
         }
 
         return self.metrics
@@ -166,7 +180,10 @@ class SklearnModelWrapper:
 
     def fit_from_datamodule(self, datamodule):
         datamodule.setup()
-        X_tensor, y_tensor = datamodule.train_set.dataset.X_tensor, datamodule.train_set.dataset.y_tensor
+        X_tensor, y_tensor = (
+            datamodule.train_set.dataset.X_tensor,
+            datamodule.train_set.dataset.y_tensor,
+        )
         indices = datamodule.train_set.indices
         X_train = X_tensor[indices].numpy()
         y_train = y_tensor[indices].numpy()
@@ -190,11 +207,11 @@ class SklearnModelWrapper:
         y_eval = y_tensor[indices].numpy()
 
         return self.evaluate(X_eval, y_eval, prefix=split)
-    
+
     def compute_shap(self, X, background=None, nsamples=100, target_class=None):
         """
         Compute SHAP values on input X, optionally for a specified target class.
-        
+
         Parameters
         ----------
         X : array-like
@@ -225,7 +242,7 @@ class SklearnModelWrapper:
             shap_values = explainer.shap_values(X)
         else:
             shap_values = explainer.shap_values(X, nsamples=nsamples)
-        
+
         if isinstance(shap_values, np.ndarray):
             if shap_values.ndim == 3:
                 if isinstance(target_class, int):
@@ -234,10 +251,7 @@ class SklearnModelWrapper:
                     # target_class is per-sample
                     if np.any(target_class >= shap_values.shape[2]):
                         raise ValueError(f"target_class values exceed {shap_values.shape[2]}")
-                    selected = np.array([
-                        shap_values[i, :, c]
-                        for i, c in enumerate(target_class)
-                    ])
+                    selected = np.array([shap_values[i, :, c] for i, c in enumerate(target_class)])
                     return selected
                 else:
                     # fallback to class 0
@@ -246,7 +260,15 @@ class SklearnModelWrapper:
                 # 2D shape (samples, features), no class dimension
                 return shap_values
 
-    def apply_shap_to_adata(self, dataloader, adata, background=None, adata_key="shap_values", target_class=None, normalize=True):
+    def apply_shap_to_adata(
+        self,
+        dataloader,
+        adata,
+        background=None,
+        adata_key="shap_values",
+        target_class=None,
+        normalize=True,
+    ):
         """
         Compute SHAP from a DataLoader and store in AnnData if provided.
         """
@@ -270,4 +292,4 @@ class SklearnModelWrapper:
             row_max[row_max == 0] = 1  # avoid divide by zero
             normalized = arr / row_max
 
-            adata.obsm[f"{adata_key}_normalized"] = normalized
+            adata.obsm[f"{adata_key}_normalized"] = normalized