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

smftools/machine_learning/models/lightning_base.py

@@ -1,10 +1,16 @@
-import torch
-import pytorch_lightning as pl
 import matplotlib.pyplot as plt
+import numpy as np
+import pytorch_lightning as pl
+import torch
 from sklearn.metrics import (
-    roc_auc_score, precision_recall_curve, auc, f1_score, confusion_matrix, roc_curve
+    auc,
+    confusion_matrix,
+    f1_score,
+    precision_recall_curve,
+    roc_auc_score,
+    roc_curve,
 )
-import numpy as np
+
 
 class TorchClassifierWrapper(pl.LightningModule):
     """
@@ -16,25 +22,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 +64,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 +88,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 +111,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 +160,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 +209,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 +227,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 +238,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 +261,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 +278,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 +318,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 +327,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 +365,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,18 @@
-import torch
 import torch.nn as nn
+
 from .base import BaseTorchModel
-    
+
+
 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 +32,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

@@ -2,6 +2,7 @@ import numpy as np
 import torch
 import torch.nn as nn
 
+
 class PositionalEncoding(nn.Module):
     def __init__(self, d_model, max_len=5000):
         super().__init__()
@@ -14,5 +15,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,8 @@
-import torch
 import torch.nn as nn
+
 from .base import BaseTorchModel
 
+
 class RNNClassifier(BaseTorchModel):
     def __init__(self, input_size, hidden_dim, num_classes, **kwargs):
         super().__init__(**kwargs)
@@ -14,4 +15,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,30 @@
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
 from sklearn.metrics import (
-    roc_auc_score, precision_recall_curve, auc, f1_score, confusion_matrix, roc_curve
+    auc,
+    confusion_matrix,
+    f1_score,
+    precision_recall_curve,
+    roc_auc_score,
+    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 +44,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 +139,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 +175,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 +202,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 +237,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 +246,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 +255,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 +287,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