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

smftools/machine_learning/models/lightning_base.py

@@ -0,0 +1,345 @@
+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
+)
+import numpy as np
+
+class TorchClassifierWrapper(pl.LightningModule):
+    """
+    A Pytorch Lightning wrapper for PyTorch classifiers.
+    - Takes a PyTorch model as input.
+    - Number of classes should be passed.
+    - Optimizer is set as default to AdamW without any keyword arguments.
+    - Loss criterion is automatically detected based on if it's a binary of multi-class classifier.
+    - 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,
+        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
+    ):
+        super().__init__()
+        self.model = model
+        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
+        self.lr = lr
+        self.label_col = label_col
+        self.num_classes = num_classes
+        self.class_names = class_names
+        self.focus_class = self._resolve_focus_class(focus_class)
+        self.focus_class_name = focus_class
+        self.enforce_eval_balance = enforce_eval_balance
+        self.target_eval_freq = target_eval_freq
+        self.max_eval_positive = max_eval_positive
+
+        # Handle class weights
+        self.criterion_kwargs = criterion_kwargs or {}
+
+        if class_weights is not None:
+            if num_classes == 2:
+                # BCEWithLogits uses pos_weight, expects a scalar or tensor
+                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)
+            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._val_outputs = []
+        self._test_outputs = []
+
+    def setup(self, stage=None):
+        """
+        Sets the loss criterion.
+        """
+        if self.criterion is None and self.num_classes is not None:
+            self._init_criterion()
+
+    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)
+            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)
+            self.criterion = torch.nn.CrossEntropyLoss(**self.criterion_kwargs)
+
+    def _resolve_focus_class(self, focus_class):
+        if isinstance(focus_class, int):
+            return focus_class
+        elif isinstance(focus_class, str):
+            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}.")
+            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.
+        """
+        self.train_obs_names = datamodule.adata.obs_names[datamodule.train_set.indices].tolist()
+        self.val_obs_names = datamodule.adata.obs_names[datamodule.val_set.indices].tolist()
+        self.test_obs_names = datamodule.adata.obs_names[datamodule.test_set.indices].tolist()
+
+    def configure_optimizers(self):
+        return self.optimizer_cls(self.parameters(), lr=self.lr, **self.optimizer_kwargs)
+
+    def forward(self, x):
+        """
+        Forward pass through the model.
+        """
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        """
+        Training step for a batch through the Lightning Trainer.
+        """
+        x, y = batch
+        if self.num_classes is None:
+            self.num_classes = int(torch.max(y).item()) + 1
+            self._init_criterion()
+        logits = self(x)
+        loss = self._compute_loss(logits, y)
+        self.log("train_loss", loss, prog_bar=False)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        """
+        Validation step for a batch through the Lightning Trainer.
+        """
+        x, y = batch
+        logits = self(x)
+        loss = self._compute_loss(logits, y)
+        preds = self._get_preds(logits)
+        acc = (preds == y).float().mean()
+        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.
+        """
+        x, y = batch
+        logits = self(x)
+        self._test_outputs.append((logits.detach(), y.detach()))
+
+    def predict_step(self, batch, batch_idx):
+        """
+        Gets predictions and prediction probabilities for the batch using the trained Lightning model.
+        """
+        x = batch[0]
+        logits = self(x)
+        probs = self._get_probs(logits)
+        preds = self._get_preds(logits)
+        return preds, probs
+
+    def on_validation_epoch_end(self):
+        """
+        Final logging of all validation steps
+        """
+        if not self._val_outputs:
+            return
+        logits, targets = zip(*self._val_outputs)
+        self._val_outputs.clear()
+        self._log_classification_metrics(logits, targets, prefix="val")
+
+    def on_test_epoch_end(self):
+        """
+        Final logging of all testing steps
+        """
+        if not self._test_outputs:
+            return
+        logits, targets = zip(*self._test_outputs)
+        self._test_outputs.clear()
+        self._log_classification_metrics(logits, targets, prefix="test")
+        self._plot_roc_pr_curves(logits, targets)
+
+    def _compute_loss(self, logits, y):
+        """
+        A helper function for computing loss for binary vs multiclass classifications.
+        """
+        if self.num_classes == 2:
+            y = y.float().view(-1, 1)  # shape [B, 1]
+            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.
+        """
+        if self.num_classes == 2:
+            return torch.sigmoid(logits.view(-1))
+        else:
+            return torch.softmax(logits, dim=1)
+
+    def _get_preds(self, logits):
+        """
+        A helper function for getting class predictions for binary vs multiclass classifications.
+        """
+        if self.num_classes == 2:
+            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):
+        pos_idx = np.where(y_true == self.focus_class)[0]
+        neg_idx = np.where(y_true != self.focus_class)[0]
+
+        max_negatives_possible = len(neg_idx)
+        max_positives_possible = len(pos_idx)
+
+        # maximum achievable positive class frequency
+        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_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)
+
+        sampled_idx = np.concatenate([pos_sampled, neg_sampled])
+        np.random.shuffle(sampled_idx)
+
+        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.
+        """
+        logits = torch.cat(logits).cpu()
+        y_true = torch.cat(targets).cpu().numpy()
+
+        probs = self._get_probs(logits).numpy()
+        preds = self._get_preds(logits).cpu().numpy()
+
+        # remap binary focus class correctly:
+        binary_focus = (y_true == self.focus_class).astype(int)
+
+        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:
+            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 = y_true[sampled_idx]
+            probs = probs[sampled_idx]
+            preds = preds[sampled_idx]
+            binary_focus = (y_true == self.focus_class).astype(int)
+            num_pos = binary_focus.sum()
+
+        # Accuracy
+        acc = np.mean(preds == y_true)
+
+        # F1 & ROC-AUC
+        if self.num_classes == 2:
+            if self.focus_class == 1:
+                focus_probs = probs
+            else:
+                focus_probs = 1 - probs
+            f1 = f1_score(y_true, preds)
+            fpr, tpr, _ = roc_curve((y_true == self.focus_class).astype(int), focus_probs)
+            roc_auc = roc_auc_score((y_true == self.focus_class).astype(int), focus_probs)
+        else:
+            f1 = f1_score(y_true, preds, average="macro")
+            roc_auc = roc_auc_score(y_true, probs, multi_class="ovr", average="macro")
+            focus_probs = probs[:, self.focus_class]
+            fpr, tpr, _ = roc_curve((y_true == self.focus_class).astype(int), focus_probs)
+
+        # PR AUC for focus class
+        pr, rc, _ = precision_recall_curve(binary_focus, focus_probs)
+        pr_auc = auc(rc, pr)
+        pos_freq = binary_focus.mean()
+        pr_auc_norm = pr_auc / pos_freq if pos_freq > 0 else np.nan
+
+        cm = confusion_matrix(y_true, preds)
+
+        # Save attributes for later plotting
+        if prefix == 'test':
+            self.test_roc_curve = (fpr, tpr)
+            self.test_pr_curve = (rc, pr)
+            self.test_roc_auc = roc_auc
+            self.test_pr_auc = pr_auc
+            self.test_pos_freq = pos_freq
+            self.test_num_pos = num_pos
+            self.test_acc = acc
+            self.test_f1 = f1
+        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
+        })
+        setattr(self, f"{prefix}_confusion_matrix", cm)
+
+    def _plot_roc_pr_curves(self, logits, targets):
+        plt.figure(figsize=(12, 5))
+
+        # ROC Curve
+        fpr, tpr = self.test_roc_curve
+        roc_auc = self.test_roc_auc
+        plt.subplot(1, 2, 1)
+        plt.plot(fpr, tpr, label=f"ROC AUC={roc_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.title(f"Test ROC Curve - {self.test_num_pos} positive class instances")
+        plt.legend()
+
+        # PR Curve
+        rc, pr = self.test_pr_curve
+        pr_auc = self.test_pr_auc
+        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.xlabel("Recall")
+        plt.ylabel("Precision")
+        plt.ylim(0, 1.05)
+        plt.title(f"Test Precision-Recall Curve - {self.test_num_pos} positive class instances")
+        plt.legend()
+
+        plt.tight_layout()
+        plt.show()

smftools/machine_learning/models/mlp.py

@@ -0,0 +1,26 @@
+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):
+        super().__init__(**kwargs)
+        layers = []
+        in_dim = input_dim
+
+        for h in hidden_dims:
+            layers.append(nn.Linear(in_dim, h))
+            if use_batchnorm:
+                layers.append(nn.BatchNorm1d(h))
+            layers.append(nn.ReLU())
+            if dropout > 0:
+                layers.append(nn.Dropout(dropout))
+            in_dim = h
+
+        output_size = 1 if num_classes == 2 else num_classes
+
+        layers.append(nn.Linear(in_dim, output_size))
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)

smftools/{tools → machine_learning}/models/positional.py

@@ -10,8 +10,9 @@ class PositionalEncoding(nn.Module):
         div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
         pe[:, 0::2] = torch.sin(position * div_term)
         pe[:, 1::2] = torch.cos(position * div_term)
-        self.pe = pe.unsqueeze(0)  # (1, max_len, d_model)
+        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
+        self.register_buffer("pe", pe)
 
     def forward(self, x):
-        x = x + self.pe[:, :x.size(1)].to(x.device)
+        x = x + self.pe[:, :x.size(1)]
         return x

smftools/{tools → machine_learning}/models/rnn.py

@@ -8,7 +8,8 @@ class RNNClassifier(BaseTorchModel):
         # Define LSTM layer
         self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_dim, batch_first=True)
         # Define fully connected output layer
-        self.fc = nn.Linear(hidden_dim, num_classes)
+        output_size = 1 if num_classes == 2 else num_classes
+        self.fc = nn.Linear(hidden_dim, output_size)
 
     def forward(self, x):
         x = x.unsqueeze(1)  # [B, 1, L] → for LSTM expecting batch_first

smftools/machine_learning/models/sklearn_models.py

@@ -0,0 +1,273 @@
+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
+)
+
+class SklearnModelWrapper:
+    """
+    Unified sklearn wrapper matching TorchClassifierWrapper interface.
+    """
+    def __init__(
+        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
+    ):
+        self.model = model
+        self.label_col = label_col
+        self.num_classes = num_classes
+        self.class_names = class_names
+        self.focus_class = self._resolve_focus_class(focus_class)
+        self.focus_class_name = focus_class
+        self.enforce_eval_balance = enforce_eval_balance
+        self.target_eval_freq = target_eval_freq
+        self.max_eval_positive = max_eval_positive
+        self.metrics = {}
+
+    def _resolve_focus_class(self, focus_class):
+        if isinstance(focus_class, int):
+            return focus_class
+        elif isinstance(focus_class, str):
+            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}.")
+            return self.class_names.index(focus_class)
+        else:
+            raise ValueError(f"focus_class must be int or str, got {type(focus_class)}")
+
+    def fit(self, X, y):
+        self.model.fit(X, y)
+
+    def predict(self, X):
+        return self.model.predict(X)
+
+    def predict_proba(self, X):
+        return self.model.predict_proba(X)
+
+    def _subsample_for_fixed_positive_frequency(self, y_true):
+        pos_idx = np.where(y_true == self.focus_class)[0]
+        neg_idx = np.where(y_true != self.focus_class)[0]
+
+        max_neg = len(neg_idx)
+        max_pos = len(pos_idx)
+        max_possible_freq = max_pos / (max_pos + max_neg)
+
+        target_freq = min(self.target_eval_freq, max_possible_freq)
+        num_pos_target = min(int(target_freq * max_neg / (1 - target_freq)), max_pos)
+        num_neg_target = int(num_pos_target * (1 - target_freq) / target_freq)
+        num_neg_target = min(num_neg_target, max_neg)
+
+        if self.max_eval_positive is not None:
+            num_pos_target = min(num_pos_target, self.max_eval_positive)
+
+        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)
+
+        sampled_idx = np.concatenate([pos_sampled, neg_sampled])
+        np.random.shuffle(sampled_idx)
+
+        return sampled_idx
+
+    def evaluate(self, X, y, prefix="test"):
+        y_true = y
+        y_prob = self.predict_proba(X)
+        y_pred = self.predict(X)
+
+        if self.enforce_eval_balance:
+            sampled_idx = self._subsample_for_fixed_positive_frequency(y_true)
+            y_true = y_true[sampled_idx]
+            y_prob = y_prob[sampled_idx]
+            y_pred = y_pred[sampled_idx]
+
+        binary_focus = (y_true == self.focus_class).astype(int)
+        num_pos = binary_focus.sum()
+
+        is_binary = self.num_classes == 2
+
+        if is_binary:
+            if self.focus_class == 1:
+                focus_probs = y_prob[:, 1]
+            else:
+                focus_probs = y_prob[:, 0]
+            preds_focus = (y_pred == self.focus_class).astype(int)
+        else:
+            focus_probs = y_prob[:, self.focus_class]
+            preds_focus = (y_pred == self.focus_class).astype(int)
+
+        f1 = f1_score(binary_focus, preds_focus)
+        roc_auc = roc_auc_score(binary_focus, focus_probs)
+        pr, rc, _ = precision_recall_curve(binary_focus, focus_probs)
+        pr_auc = auc(rc, pr)
+        pos_freq = binary_focus.mean()
+        pr_auc_norm = pr_auc / pos_freq if pos_freq > 0 else np.nan
+        fpr, tpr, _ = roc_curve(binary_focus, focus_probs)
+        cm = confusion_matrix(y_true, y_pred)
+        acc = np.mean(y_pred == y_true)
+
+        # store metrics as attributes for plotting later
+        setattr(self, f"{prefix}_f1", f1)
+        setattr(self, f"{prefix}_roc_curve", (fpr, tpr))
+        setattr(self, f"{prefix}_pr_curve", (rc, pr))
+        setattr(self, f"{prefix}_roc_auc", roc_auc)
+        setattr(self, f"{prefix}_pr_auc", pr_auc)
+        setattr(self, f"{prefix}_pos_freq", pos_freq)
+        setattr(self, f"{prefix}_num_pos", num_pos)
+        setattr(self, f"{prefix}_confusion_matrix", cm)
+        setattr(self, f"{prefix}_acc", acc)
+
+        # also store a metrics dict
+        self.metrics = {
+            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
+        }
+
+        return self.metrics
+
+    def plot_roc_pr_curves(self, prefix="test"):
+        plt.figure(figsize=(12, 5))
+
+        fpr, tpr = getattr(self, f"{prefix}_roc_curve")
+        roc_auc = getattr(self, f"{prefix}_roc_auc")
+        plt.subplot(1, 2, 1)
+        plt.plot(fpr, tpr, label=f"ROC AUC={roc_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.title(f"ROC Curve - {getattr(self, f'{prefix}_num_pos')} positives")
+        plt.legend()
+
+        rc, pr = getattr(self, f"{prefix}_pr_curve")
+        pr_auc = getattr(self, f"{prefix}_pr_auc")
+        pos_freq = getattr(self, f"{prefix}_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.xlabel("Recall")
+        plt.ylabel("Precision")
+        plt.ylim(0, 1.05)
+        plt.title(f"PR Curve - {getattr(self, f'{prefix}_num_pos')} positives")
+        plt.legend()
+
+        plt.tight_layout()
+        plt.show()
+
+    def fit_from_datamodule(self, datamodule):
+        datamodule.setup()
+        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()
+        self.fit(X_train, y_train)
+        self.train_obs_names = datamodule.adata.obs_names[datamodule.train_set.indices].tolist()
+        self.val_obs_names = datamodule.adata.obs_names[datamodule.val_set.indices].tolist()
+        self.test_obs_names = datamodule.adata.obs_names[datamodule.test_set.indices].tolist()
+
+    def evaluate_from_datamodule(self, datamodule, split="test"):
+        datamodule.setup()
+        if split == "val":
+            subset = datamodule.val_set
+        elif split == "test":
+            subset = datamodule.test_set
+        else:
+            raise ValueError(f"Invalid split '{split}'")
+
+        X_tensor, y_tensor = subset.dataset.X_tensor, subset.dataset.y_tensor
+        indices = subset.indices
+        X_eval = X_tensor[indices].numpy()
+        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
+            Input features
+        background : array-like
+            SHAP background
+        nsamples : int
+            Number of samples for kernel approximation
+        target_class : int, optional
+            If None, uses model predicted class
+        """
+        import shap
+
+        # choose explainer
+        if hasattr(self.model, "tree_") or hasattr(self.model, "estimators_"):
+            explainer = shap.TreeExplainer(self.model, data=background)
+        else:
+            if background is None:
+                background = shap.kmeans(X, 10)
+            explainer = shap.KernelExplainer(self.model.predict_proba, background)
+
+        # determine class
+        if target_class is None:
+            preds = self.model.predict(X)
+            target_class = preds
+
+        if isinstance(explainer, shap.TreeExplainer):
+            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):
+                    return shap_values[:, :, target_class]
+                elif isinstance(target_class, np.ndarray):
+                    # 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)
+                    ])
+                    return selected
+                else:
+                    # fallback to class 0
+                    return shap_values[:, :, 0]
+            else:
+                # 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):
+        """
+        Compute SHAP from a DataLoader and store in AnnData if provided.
+        """
+        X_batches = []
+
+        for batch in dataloader:
+            X = batch[0].detach().cpu().numpy()
+            X_batches.append(X)
+
+        X_full = np.concatenate(X_batches, axis=0)
+
+        shap_values = self.compute_shap(X_full, background=background, target_class=target_class)
+
+        if adata is not None:
+            adata.obsm[adata_key] = shap_values
+
+        if normalize:
+            arr = shap_values
+            # row-wise normalization
+            row_max = np.max(np.abs(arr), axis=1, keepdims=True)
+            row_max[row_max == 0] = 1  # avoid divide by zero
+            normalized = arr / row_max
+
+            adata.obsm[f"{adata_key}_normalized"] = normalized