Evaluate API¶

def evaluate(
    self,
    mdata: MuData,
    task: str,
    batch_size: int = 32,
    plot_roc: bool = False,
    figsize: tuple[float, float] = (4, 3),
) -> float | dict[str, float]:
    """Evaluate the model on held-out data.

    Args:
        mdata: Multi-omics object whose `obs` holds the clinical annotations.
        task: `'classification'` or `'survival'`.
        batch_size: Evaluation batch size.
        plot_roc: Save a ROC curve image (classification only).
        figsize: Figure size for the ROC curve.

    Returns:
        Concordance index (float) for `task='survival'`, or a metrics dict
        (Accuracy, F1-score, Precision, Recall, AUC) for
        `task='classification'`.

    Raises:
        ModelNotFittedError: If called before `fit()`.
        ValueError: If `task` is not `'classification'` or `'survival'`.
    """
    self._require_fitted()
    if task not in ("classification", "survival"):
        raise ValueError(f"task must be 'classification' or 'survival', got '{task}'.")

    label = mdata.uns[Keys.LABEL]

    encoded_labels = pd.Series(self.label_encoder.transform(mdata.obs[label].values), index=mdata.obs_names)

    loader_kw: dict = {"num_workers": 2, "pin_memory": True} if self.device.type == "cuda" else {}
    shared_samples = get_shared_samples(mdata)
    test_loader = DataLoader(
        MultiOmicsDataset(mdata, shared_samples, encoded_labels),
        batch_size=batch_size,
        shuffle=False,
        **loader_kw,
    )

    self._set_eval_mode()
    all_y_true, all_y_pred, all_y_proba = [], [], []
    all_hazard, all_os_time, all_os_event = [], [], []

    with torch.no_grad():
        for x, labels, os_time, os_event in test_loader:
            x = [xi.to(self.device) for xi in x]
            z = self._get_central_representation(x)

            if task == "survival":
                hazard = self.survival_predictor(z).cpu().numpy().reshape(-1, 1)
                all_hazard.append(hazard)
                all_os_time.append(os_time.numpy())
                all_os_event.append(os_event.numpy())
            else:
                logits = self.classifier(z)
                all_y_pred.append(torch.argmax(logits, dim=1).cpu().numpy())
                all_y_proba.append(torch.softmax(logits, dim=1).cpu().numpy())
                all_y_true.append(labels.cpu().numpy())

    if task == "survival":
        hazard_cat = np.vstack(all_hazard)
        return float(
            CIndex_lifeline(
                hazard_cat,
                np.concatenate(all_os_event),
                np.concatenate(all_os_time),
            )
        )

    y_true = np.concatenate(all_y_true)
    y_pred = np.concatenate(all_y_pred)
    y_proba = np.vstack(all_y_proba)

    if plot_roc:
        plot_roc_multiclass(
            y_test=y_true,
            y_pred_proba=y_proba,
            filename="test",
            n_classes=self.num_classes,
            var_names=np.unique(mdata.obs[label].values.tolist()).tolist(),
            figsize=figsize,
        )

    return multi_classification_evaluation(y_true, y_pred, y_proba, ohe=self.one_hot_encoder)

def predict(self, mdata: MuData) -> np.ndarray:
    """Predict class labels.

    Args:
        mdata: Multi-omics object matching the sources used in `fit`.

    Returns:
        Integer class predictions, shape `(n_samples,)`.

    Raises:
        ModelNotFittedError: If called before `fit()`.
    """
    self._require_fitted()
    self._set_eval_mode()
    z = torch.tensor(self.get_latent_representation(mdata), dtype=torch.float32).to(self.device)
    with torch.no_grad():
        logits = self.classifier(z)
    return torch.argmax(logits, dim=1).cpu().numpy()

def predict_survival(self, mdata: MuData) -> dict[str, pd.DataFrame]:
    """Compute patient-level estimated survival functions.

    Args:
        mdata: Multi-omics object matching the sources used in `fit`.

    Returns:
        Maps sample ID → estimated survival function DataFrame.

    Raises:
        ModelNotFittedError: If called before `fit()`.
    """
    self._require_fitted()
    shared_samples = get_shared_samples(mdata)
    z = torch.tensor(self.get_latent_representation(mdata), dtype=torch.float32).to(self.device)
    self._set_eval_mode()
    with torch.no_grad():
        risk_scores = self.survival_predictor(z).cpu().numpy()
    return {s: self.baseline * np.exp(r[0]) for s, r in zip(shared_samples, risk_scores)}