smftools 0.1.6__py3-none-any.whl → 0.2.1__py3-none-any.whl

smftools/machine_learning/__init__.py

@@ -0,0 +1,12 @@
+from . import models
+from . import data
+from . import utils
+from . import evaluation
+from . import inference
+from . import training
+
+__all__ = [
+    "calculate_relative_risk_on_activity",
+    "evaluate_models_by_subgroup",
+    "prepare_melted_model_data",
+]

smftools/machine_learning/data/__init__.py

@@ -0,0 +1,2 @@
+from .anndata_data_module import AnnDataModule, build_anndata_loader
+from .preprocessing import random_fill_nans

smftools/machine_learning/data/anndata_data_module.py

@@ -0,0 +1,234 @@
+import torch
+from torch.utils.data import DataLoader, TensorDataset, random_split, Dataset, Subset
+import pytorch_lightning as pl
+import numpy as np
+import pandas as pd
+from .preprocessing import random_fill_nans
+from sklearn.utils.class_weight import compute_class_weight
+
+    
+class AnnDataDataset(Dataset):
+    """
+    Generic PyTorch Dataset from AnnData.
+    """
+    def __init__(self, adata, tensor_source="X", tensor_key=None, label_col=None, window_start=None, window_size=None):
+        self.adata = adata
+        self.tensor_source = tensor_source
+        self.tensor_key = tensor_key
+        self.label_col = label_col
+        self.window_start = window_start
+        self.window_size = window_size
+        
+        if tensor_source == "X":
+            X = adata.X
+        elif tensor_source == "layers":
+            assert tensor_key in adata.layers
+            X = adata.layers[tensor_key]
+        elif tensor_source == "obsm":
+            assert tensor_key in adata.obsm
+            X = adata.obsm[tensor_key]
+        else:
+            raise ValueError(f"Invalid tensor_source: {tensor_source}")
+        
+        if self.window_start is not None and self.window_size is not None:
+            X = X[:, self.window_start : self.window_start + self.window_size]
+        
+        X = random_fill_nans(X)
+
+        self.X_tensor = torch.tensor(X, dtype=torch.float32)
+
+        if label_col is not None:
+            y = adata.obs[label_col]
+            if y.dtype.name == 'category':
+                y = y.cat.codes
+            self.y_tensor = torch.tensor(y.values, dtype=torch.long)
+        else:
+            self.y_tensor = None
+
+    def numpy(self, indices):
+        return self.X_tensor[indices].numpy(), self.y_tensor[indices].numpy()
+    
+    def __len__(self):
+        return len(self.X_tensor)
+
+    def __getitem__(self, idx):
+        x = self.X_tensor[idx]
+        if self.y_tensor is not None:
+            y = self.y_tensor[idx]
+            return x, y
+        else:
+            return (x,)
+
+
+def split_dataset(adata, dataset, train_frac=0.6, val_frac=0.1, test_frac=0.3, 
+                                 random_seed=42, split_col="train_val_test_split", 
+                                 load_existing_split=False, split_save_path=None):
+    """
+    Perform split and record assignment into adata.obs[split_col].
+    """
+    total_len = len(dataset)
+
+    if load_existing_split:
+        if split_col in adata.obs:
+            pass  # use existing
+        elif split_save_path:
+            split_df = pd.read_csv(split_save_path, index_col=0)
+            adata.obs[split_col] = split_df.loc[adata.obs_names][split_col].values
+        else:
+            raise ValueError("No existing split column found and no file provided.")
+    else:
+        indices = np.arange(total_len)
+        np.random.seed(random_seed)
+        np.random.shuffle(indices)
+
+        n_train = int(train_frac * total_len)
+        n_val = int(val_frac * total_len)
+        n_test = total_len - n_train - n_val
+
+        split_array = np.full(total_len, "test", dtype=object)
+        split_array[indices[:n_train]] = "train"
+        split_array[indices[n_train:n_train + n_val]] = "val"
+        adata.obs[split_col] = split_array
+
+        if split_save_path:
+            adata.obs[[split_col]].to_csv(split_save_path)
+
+    split_labels = adata.obs[split_col].values
+    train_indices = np.where(split_labels == "train")[0]
+    val_indices = np.where(split_labels == "val")[0]
+    test_indices = np.where(split_labels == "test")[0]
+
+    train_set = Subset(dataset, train_indices)
+    val_set = Subset(dataset, val_indices)
+    test_set = Subset(dataset, test_indices)
+
+    return train_set, val_set, test_set
+
+class AnnDataModule(pl.LightningDataModule):
+    """
+    Unified LightningDataModule version of AnnDataDataset + splitting with adata.obs recording.
+    """
+    def __init__(self, adata, tensor_source="X", tensor_key=None, label_col="labels",
+                 batch_size=64, train_frac=0.6, val_frac=0.1, test_frac=0.3, random_seed=42,
+                 inference_mode=False, split_col="train_val_test_split", split_save_path=None,
+                 load_existing_split=False, window_start=None, window_size=None, num_workers=None, persistent_workers=False):
+        super().__init__()
+        self.adata = adata
+        self.tensor_source = tensor_source
+        self.tensor_key = tensor_key
+        self.label_col = label_col
+        self.batch_size = batch_size
+        self.train_frac = train_frac
+        self.val_frac = val_frac
+        self.test_frac = test_frac
+        self.random_seed = random_seed
+        self.inference_mode = inference_mode
+        self.split_col = split_col
+        self.split_save_path = split_save_path
+        self.load_existing_split = load_existing_split
+        self.var_names = adata.var_names.copy()
+        self.window_start = window_start
+        self.window_size = window_size
+        self.num_workers = num_workers
+        self.persistent_workers = persistent_workers
+
+    def setup(self, stage=None):
+        dataset = AnnDataDataset(self.adata, self.tensor_source, self.tensor_key, 
+                                  None if self.inference_mode else self.label_col,
+                                    window_start=self.window_start, window_size=self.window_size)
+
+        if self.inference_mode:
+            self.infer_dataset = dataset
+            return
+
+        self.train_set, self.val_set, self.test_set = split_dataset(
+            self.adata, dataset, train_frac=self.train_frac, val_frac=self.val_frac, 
+            test_frac=self.test_frac, random_seed=self.random_seed, 
+            split_col=self.split_col, split_save_path=self.split_save_path, 
+            load_existing_split=self.load_existing_split
+        )
+
+    def train_dataloader(self):
+        if self.num_workers:
+            return DataLoader(self.train_set, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers, persistent_workers=self.persistent_workers)
+        else:
+            return DataLoader(self.train_set, batch_size=self.batch_size, shuffle=True)
+
+    def val_dataloader(self):
+        if self.num_workers:
+            return DataLoader(self.val_set, batch_size=self.batch_size, num_workers=self.num_workers, persistent_workers=self.persistent_workers)
+        else:
+            return DataLoader(self.train_set, batch_size=self.batch_size, shuffle=False)
+        
+    def test_dataloader(self):
+        if self.num_workers:
+            return DataLoader(self.test_set, batch_size=self.batch_size, num_workers=self.num_workers, persistent_workers=self.persistent_workers)
+        else:
+            return DataLoader(self.train_set, batch_size=self.batch_size, shuffle=False)
+        
+    def predict_dataloader(self):
+        if not self.inference_mode:
+            raise RuntimeError("Only valid in inference mode")
+        return DataLoader(self.infer_dataset, batch_size=self.batch_size)
+    
+    def compute_class_weights(self):
+        train_indices = self.train_set.indices # get the indices of the training set
+        y_all = self.train_set.dataset.y_tensor # get labels for the entire dataset (We are pulling from a Subset object, so this syntax can be confusing)
+        y_train = y_all[train_indices].cpu().numpy() # get the labels for the training set and move to a numpy array
+        
+        class_weights = compute_class_weight('balanced', classes=np.unique(y_train), y=y_train)
+        return torch.tensor(class_weights, dtype=torch.float32)
+    
+    def inference_numpy(self):
+        """
+        Return inference data as numpy for use in sklearn inference.
+        """
+        if not self.inference_mode:
+            raise RuntimeError("Must be in inference_mode=True to use inference_numpy()")
+        X_np = self.infer_dataset.X_tensor.numpy()
+        return X_np
+    
+    def to_numpy(self):
+        """
+        Move the AnnDataModule tensors into numpy arrays
+        """
+        if not self.inference_mode:
+            train_X, train_y = self.train_set.dataset.numpy(self.train_set.indices)
+            val_X, val_y = self.val_set.dataset.numpy(self.val_set.indices)
+            test_X, test_Y = self.test_set.dataset.numpy(self.test_set.indices)
+            return train_X, train_y, val_X, val_y, test_X, test_Y
+        else:
+            return self.inference_numpy()
+
+
+def build_anndata_loader(
+    adata, tensor_source="X", tensor_key=None, label_col=None, train_frac=0.6, val_frac=0.1, 
+    test_frac=0.3, random_seed=42, batch_size=64, lightning=True, inference_mode=False, 
+    split_col="train_val_test_split", split_save_path=None, load_existing_split=False
+):
+    """
+    Unified pipeline for both Lightning and raw PyTorch.
+    The lightning loader works for both Lightning and the Sklearn wrapper.
+    Set lightning to False if you want to make data loaders for base PyTorch or base sklearn models
+    """
+    if lightning:
+        return AnnDataModule(
+            adata, tensor_source=tensor_source, tensor_key=tensor_key, label_col=label_col,
+            batch_size=batch_size, train_frac=train_frac, val_frac=val_frac, test_frac=test_frac,
+            random_seed=random_seed, inference_mode=inference_mode,
+            split_col=split_col, split_save_path=split_save_path, load_existing_split=load_existing_split
+        )
+    else:
+        var_names = adata.var_names.copy()
+        dataset = AnnDataDataset(adata, tensor_source, tensor_key, None if inference_mode else label_col)
+        if inference_mode:
+            return DataLoader(dataset, batch_size=batch_size)
+        else:
+            train_set, val_set, test_set = split_dataset(
+                adata, dataset, train_frac, val_frac, test_frac, random_seed, 
+                split_col, split_save_path, load_existing_split
+            )
+            train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
+            val_loader = DataLoader(val_set, batch_size=batch_size)
+            test_loader = DataLoader(test_set, batch_size=batch_size)
+            return train_loader, val_loader, test_loader

smftools/machine_learning/data/preprocessing.py

@@ -0,0 +1,6 @@
+import numpy as np
+
+def random_fill_nans(X):
+    nan_mask = np.isnan(X)
+    X[nan_mask] = np.random.rand(*X[nan_mask].shape)
+    return X

smftools/machine_learning/evaluation/__init__.py

@@ -0,0 +1,2 @@
+from .evaluators import ModelEvaluator, PostInferenceModelEvaluator
+from .eval_utils import flatten_sliding_window_results

smftools/machine_learning/evaluation/eval_utils.py

@@ -0,0 +1,31 @@
+import pandas as pd
+
+def flatten_sliding_window_results(results_dict):
+    """
+    Flatten nested sliding window results into pandas DataFrame.
+    
+    Expects structure: 
+        results[model_name][window_size][window_center]['metrics'][metric_name]
+    """
+    records = []
+
+    for model_name, model_results in results_dict.items():
+        for window_size, window_results in model_results.items():
+            for center_var, result in window_results.items():
+                metrics = result['metrics']
+                record = {
+                    'model': model_name,
+                    'window_size': window_size,
+                    'center_var': center_var
+                }
+                # Add all metrics
+                record.update(metrics)
+                records.append(record)
+                
+    df = pd.DataFrame.from_records(records)
+    
+    # Convert center_var to numeric if possible (optional but helpful for plotting)
+    df['center_var'] = pd.to_numeric(df['center_var'], errors='coerce')
+    df = df.sort_values(['model', 'window_size', 'center_var'])
+    
+    return df

smftools/machine_learning/evaluation/evaluators.py

@@ -0,0 +1,223 @@
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+from sklearn.metrics import (
+    roc_auc_score, precision_recall_curve, auc, f1_score, confusion_matrix, roc_curve
+)
+
+class ModelEvaluator:
+    """
+    A model evaluator for consolidating Sklearn and Lightning model evaluation metrics on testing data
+    """
+    def __init__(self):
+        self.results = []
+        self.pos_freq = None
+        self.num_pos = None
+
+    def add_model(self, name, model, is_torch=True):
+        """
+        Add a trained model with its evaluation metrics.
+        """
+        if is_torch:
+            entry = {
+                'name': name,
+                'f1': model.test_f1,
+                'auc': model.test_roc_auc,
+                'pr_auc': model.test_pr_auc,
+                'pr_auc_norm': model.test_pr_auc / model.test_pos_freq if model.test_pos_freq > 0 else np.nan,
+                'pr_curve': model.test_pr_curve,
+                'roc_curve': model.test_roc_curve,
+                'num_pos': model.test_num_pos,
+                'pos_freq': model.test_pos_freq
+            }
+        else:
+            entry = {
+                'name': name,
+                'f1': model.test_f1,
+                'auc': model.test_roc_auc,
+                'pr_auc': model.test_pr_auc,
+                'pr_auc_norm': model.test_pr_auc / model.test_pos_freq if model.test_pos_freq > 0 else np.nan,
+                'pr_curve': model.test_pr_curve,
+                'roc_curve': model.test_roc_curve,
+                'num_pos': model.test_num_pos,
+                'pos_freq': model.test_pos_freq
+            }
+        
+        self.results.append(entry)
+
+        if not self.pos_freq:
+            self.pos_freq = entry['pos_freq']
+            self.num_pos = entry['num_pos']
+
+    def get_metrics_dataframe(self):
+        """
+        Return all metrics as pandas DataFrame.
+        """
+        df = pd.DataFrame(self.results)
+        return df[['name', 'f1', 'auc', 'pr_auc', 'pr_auc_norm', 'num_pos', 'pos_freq']]
+
+    def plot_all_curves(self):
+        """
+        Plot unified ROC and PR curves across all models.
+        """
+        plt.figure(figsize=(12, 5))
+
+        # ROC
+        plt.subplot(1, 2, 1)
+        for res in self.results:
+            fpr, tpr = res['roc_curve']
+            plt.plot(fpr, tpr, label=f"{res['name']} (AUC={res['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 Curves - {self.num_pos} positive instances")
+        plt.legend()
+
+        # PR
+        plt.subplot(1, 2, 2)
+        for res in self.results:
+            rc, pr = res['pr_curve']
+            plt.plot(rc, pr, label=f"{res['name']} (AUPRC={res['pr_auc']:.3f})")
+        plt.xlabel("Recall")
+        plt.ylabel("Precision")
+        plt.ylim(0,1.05)
+        plt.axhline(self.pos_freq, linestyle='--', color='grey')
+        plt.title(f"Precision-Recall Curves - {self.num_pos} positive instances")
+        plt.legend()
+
+        plt.tight_layout()
+        plt.show()
+
+class PostInferenceModelEvaluator:
+    def __init__(self, adata, models, target_eval_freq=None, max_eval_positive=None):
+        """
+        Initialize evaluator.
+
+        Parameters:
+        -----------
+        adata : AnnData
+            The annotated dataset where predictions are stored in obs/obsm.
+        models : dict
+            Dictionary of models: {model_name: model_instance}.
+            Supports TorchClassifierWrapper and SklearnModelWrapper.
+        """
+        self.adata = adata
+        self.models = models
+        self.target_eval_freq = target_eval_freq
+        self.max_eval_positive = max_eval_positive
+        self.results = {}
+
+    def evaluate_all(self):
+        """
+        Evaluate all models and store results.
+        """
+        for name, model in self.models.items():
+            print(f"Evaluating {name}...")
+            label_col = model.label_col
+            full_prefix = f"{name}_{label_col}"
+            self.results[full_prefix] = self._evaluate_model(name, model)
+
+    def _evaluate_model(self, model_name, model):
+        """
+        Evaluate one model and return metrics.
+        """
+        label_col = model.label_col
+        num_classes = model.num_classes
+        class_names = model.class_names
+        focus_class = model.focus_class
+
+        full_prefix = f"{model_name}_{label_col}"
+
+        # Extract ground truth + predictions
+        y_true = self.adata.obs[label_col].cat.codes.to_numpy()
+        y_pred = self.adata.obs[f"{full_prefix}_pred"].to_numpy()
+        probs_all = self.adata.obsm[f"{full_prefix}_pred_prob_all"]
+
+        binary_focus = (y_true == focus_class).astype(int)
+
+        # OPTIONAL SUBSAMPLING
+        if self.target_eval_freq is not None:
+            indices = self._subsample_for_fixed_positive_frequency(
+                binary_focus, target_freq=self.target_eval_freq, max_positive=self.max_eval_positive
+            )
+            y_true = y_true[indices]
+            y_pred = y_pred[indices]
+            probs_all = probs_all[indices]
+            binary_focus = (y_true == focus_class).astype(int)
+
+        acc = np.mean(y_true == y_pred)
+
+        if num_classes == 2:
+            focus_probs = probs_all[:, focus_class]
+            f1 = f1_score(binary_focus, (y_pred == focus_class).astype(int))
+            roc_auc = roc_auc_score(binary_focus, focus_probs)
+            pr, rc, _ = precision_recall_curve(binary_focus, focus_probs)
+            fpr, tpr, _ = roc_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
+        else:
+            f1 = f1_score(y_true, y_pred, average="macro")
+            roc_auc = roc_auc_score(y_true, probs_all, multi_class="ovr", average="macro")
+            focus_probs = probs_all[:, focus_class]
+            pr, rc, _ = precision_recall_curve(binary_focus, focus_probs)
+            fpr, tpr, _ = roc_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, y_pred)
+
+        metrics = {
+            "accuracy": acc,
+            "f1": f1,
+            "roc_auc": roc_auc,
+            "pr_auc": pr_auc,
+            "pr_auc_norm": pr_auc_norm,
+            "pos_freq": pos_freq,
+            "confusion_matrix": cm,
+            "pr_rc_curve": (pr, rc),
+            "roc_curve": (tpr, fpr)
+        }
+
+        return metrics
+    
+    def _subsample_for_fixed_positive_frequency(self, binary_labels, target_freq=0.3, max_positive=None):
+        pos_idx = np.where(binary_labels == 1)[0]
+        neg_idx = np.where(binary_labels == 0)[0]
+
+        max_pos = len(pos_idx)
+        max_neg = len(neg_idx)
+
+        max_possible_freq = max_pos / (max_pos + max_neg)
+        if target_freq > max_possible_freq:
+            target_freq = max_possible_freq
+
+        num_pos_target = int(target_freq * max_neg / (1 - target_freq))
+        num_pos_target = min(num_pos_target, max_pos)
+        if max_positive is not None:
+            num_pos_target = min(num_pos_target, max_positive)
+
+        num_neg_target = int(num_pos_target * (1 - target_freq) / target_freq)
+        num_neg_target = min(num_neg_target, max_neg)
+
+        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 to_dataframe(self):
+        """
+        Convert results to pandas DataFrame (excluding confusion matrices).
+        """
+        records = []
+        for model_name, metrics in self.results.items():
+            row = {"model": model_name}
+            for k, v in metrics.items():
+                if k not in ["confusion_matrix", "pr_rc_curve", "roc_curve"]:
+                    row[k] = v
+            records.append(row)
+        return pd.DataFrame(records)

smftools/machine_learning/inference/__init__.py

@@ -0,0 +1,3 @@
+from .lightning_inference import run_lightning_inference
+from .sliding_window_inference import sliding_window_inference
+from .sklearn_inference import run_sklearn_inference

smftools/machine_learning/inference/inference_utils.py

@@ -0,0 +1,27 @@
+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.
+    """
+    # Get sets for fast lookup
+    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:
+        if obs in train_set:
+            split_labels.append("training")
+        elif obs in val_set:
+            split_labels.append("validation")
+        elif obs in test_set:
+            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"])
+    
+    print(f"Annotated {split_col} column with training/validation/testing/new status.")

smftools/machine_learning/inference/lightning_inference.py

@@ -0,0 +1,68 @@
+import torch
+import pandas as pd
+import numpy as np
+from pytorch_lightning import Trainer
+from .inference_utils import annotate_split_column
+
+def run_lightning_inference(
+    adata,
+    model,
+    datamodule,
+    trainer,
+    prefix="model",
+    devices=1
+):
+    """
+    Run inference on AnnData using TorchClassifierWrapper + AnnDataModule (in inference mode).
+    """
+
+    # Device logic
+    if torch.cuda.is_available():
+        accelerator = "gpu"
+    elif torch.backends.mps.is_available():
+        accelerator = "mps"
+        devices = 1
+    else:
+        accelerator = "cpu"
+        devices = 1
+
+    label_col = model.label_col
+    num_classes = model.num_classes
+    class_labels = model.class_names
+    focus_class = model.focus_class
+    focus_class_name = model.focus_class_name
+
+    annotate_split_column(adata, model, split_col=f"{prefix}_training_split")
+
+    # Run predictions
+    outputs = trainer.predict(model, datamodule=datamodule)
+
+    preds_list, probs_list = zip(*outputs)
+    preds = torch.cat(preds_list, dim=0).cpu().numpy()
+    probs = torch.cat(probs_list, dim=0).cpu().numpy()
+
+    # Handle binary vs multiclass formats
+    if num_classes == 2:
+        # probs shape: (N,) from sigmoid
+        pred_class_idx = (probs >= 0.5).astype(int)
+        probs_all = np.vstack([1 - probs, probs]).T  # shape (N, 2)
+        pred_class_probs = probs_all[np.arange(len(probs_all)), pred_class_idx]
+    else:
+        pred_class_idx = probs.argmax(axis=1)
+        probs_all = probs
+        pred_class_probs = probs_all[np.arange(len(probs_all)), pred_class_idx]
+
+    pred_class_labels = [class_labels[i] for i in pred_class_idx]
+
+    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_prob"] = pred_class_probs
+
+    for i, class_name in enumerate(class_labels):
+        adata.obs[f"{full_prefix}_prob_{class_name}"] = probs_all[:, i]
+
+    adata.obsm[f"{full_prefix}_pred_prob_all"] = probs_all
+
+    print(f"Inference complete: stored under prefix '{full_prefix}'")

smftools/machine_learning/inference/sklearn_inference.py

@@ -0,0 +1,55 @@
+import pandas as pd
+import numpy as np
+from .inference_utils import annotate_split_column
+
+
+def run_sklearn_inference(
+    adata,
+    model,
+    datamodule, 
+    prefix="model"
+):
+    """
+    Run inference on AnnData using SklearnModelWrapper.
+    """
+
+    label_col = model.label_col
+    num_classes = model.num_classes
+    class_labels = model.class_names
+    focus_class_name = model.focus_class_name
+
+    annotate_split_column(adata, model, split_col=f"{prefix}_training_split")
+
+    datamodule.setup()
+
+    X_infer = datamodule.to_numpy()
+
+    # Run predictions
+    preds = model.predict(X_infer)
+    probs = model.predict_proba(X_infer)
+
+    # Handle binary vs multiclass formats
+    if num_classes == 2:
+        # probs shape: (N, 2) from predict_proba
+        pred_class_idx = preds
+        probs_all = probs
+        pred_class_probs = probs[np.arange(len(probs)), pred_class_idx]
+    else:
+        pred_class_idx = preds
+        probs_all = probs
+        pred_class_probs = probs[np.arange(len(probs)), pred_class_idx]
+
+    pred_class_labels = [class_labels[i] for i in pred_class_idx]
+
+    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_prob"] = pred_class_probs
+
+    for i, class_name in enumerate(class_labels):
+        adata.obs[f"{full_prefix}_prob_{class_name}"] = probs_all[:, i]
+
+    adata.obsm[f"{full_prefix}_pred_prob_all"] = probs_all
+
+    print(f"Inference complete: stored under prefix '{full_prefix}'")