Visualise API¶

def plot_cohort_overview(mdata: MuData, show: bool = True) -> None:
    """Plot the subtype distribution and survival-time histogram side by side.

    Reads the clinical column names from `mdata.uns`.

    Args:
        mdata: Multi-omics object with clinical columns registered in `mdata.uns`.
        show: If True, display the figure interactively.
    """
    label = mdata.uns[Keys.LABEL]
    surv_time = mdata.uns[Keys.SURV_TIME]

    counts = mdata.obs[label].value_counts().sort_index()

    _, axes = plt.subplots(1, 2, figsize=(12, 4))
    counts.plot(kind="bar", ax=axes[0], color="steelblue", edgecolor="white")
    axes[0].set_title("Samples per subtype")
    axes[0].set_xlabel("Cluster ID (subtype)")
    axes[0].set_ylabel("Count")
    axes[0].tick_params(axis="x", rotation=0)

    mdata.obs[surv_time].hist(bins=20, ax=axes[1], color="coral", edgecolor="white")
    axes[1].set_title("Overall survival time distribution")
    axes[1].set_xlabel("Days")
    sns.despine(offset=10, trim=True)

    if show:
        plt.show()

def get_latent_representation(
    self,
    mdata: MuData,
) -> np.ndarray:
    """Compute the integrated central latent representation.

    Args:
        mdata: Multi-omics object with all sources (same keys as used in `fit`).

    Returns:
        Latent matrix, shape `(n_samples, central_latent_dim)`. Rows
        correspond to `get_shared_samples(mdata)`, in that order.

    Raises:
        ModelNotFittedError: If called before `fit()`.
    """
    self._require_fitted()
    self._set_eval_mode()
    shared_samples = get_shared_samples(mdata)
    x = [
        torch.tensor(np.asarray(mdata[mod][shared_samples].X), dtype=torch.float32).to(self.device)
        for mod in self.source_names
    ]
    with torch.no_grad():
        z = self._get_central_representation(x)
    return z.cpu().numpy()

def plot_representation(self, mdata: MuData, color: str | None = None, show: bool = True) -> None:
    """Compute latent representations and save a t-SNE scatter plot.

    Args:
        mdata: Multi-omics object.
        color: Column in `mdata.obs` to use for colouring samples. By default, use the label column.
        show: Display the figure interactively.
    """
    from customics.visualization import plot_representation as _plot

    _plot(self, mdata, color or mdata.uns[Keys.LABEL], show=show)

def stratify(self, mdata: MuData, show: bool = True) -> None:
    """Stratify patients by predicted risk and plot Kaplan-Meier curves.

    Args:
        mdata: Multi-omics object.
        show: Display the figure interactively.
    """
    from customics.visualization import plot_survival_stratification as _plot

    _plot(self, mdata, show)

def explain(
    self,
    sample_ids: list[str],
    mdata: MuData,
    source: str,
    subtype: str,
    device: str | None = None,
    show: bool = True,
) -> None:
    """Compute and plot SHAP values for one omics source and one subtype.

    Uses `shap.DeepExplainer` on the single-source forward path.
    SHAP values are computed for `subtype` against all other classes.

    Args:
        sample_ids: Sample IDs to use as the SHAP background and foreground sets.
        mdata: Multi-omics object whose `obs` holds the clinical metadata.
        source: Omics source key to explain.
        subtype: Class label to explain.
        device: Device for SHAP tensors.
        show: Display the SHAP plot interactively.

    Raises:
        ModelNotFittedError: If called before `fit()`.
    """
    import matplotlib.pyplot as plt
    import shap

    from customics.explain.shap import (
        ModelWrapper,
        add_to_tensor,
        process_phenotype_data_for_samples,
        random_training_sample,
        split_expr_and_sample,
    )

    self._require_fitted()

    device = _parse_device(device)

    expr_df = mdata[source].to_df()
    sample_ids = list(set(sample_ids) & set(expr_df.index))
    phenotype = process_phenotype_data_for_samples(mdata.obs, sample_ids)
    condition = phenotype[mdata.uns[Keys.LABEL]] == subtype

    expr_df = expr_df.loc[sample_ids, :]
    background = add_to_tensor(random_training_sample(expr_df, 10), device)
    foreground_df = split_expr_and_sample(condition=condition, sample_size=10, expr=expr_df)
    foreground = add_to_tensor(foreground_df, device)

    class_idx = int(self.label_encoder.transform([subtype])[0])
    explainer = shap.DeepExplainer(ModelWrapper(self, source=source), background)
    shap_values = explainer.shap_values(foreground, ranked_outputs=None)

    # SHAP ≥0.46 stacks class outputs into (n_samples, n_features, n_classes);
    # older versions return a list indexed [class][sample, feature].
    import numpy as np

    if isinstance(shap_values, np.ndarray) and shap_values.ndim == 3:
        sv = shap_values[..., class_idx]
    else:
        sv = shap_values[class_idx]

    shap.summary_plot(
        sv,
        features=foreground_df,
        feature_names=list(expr_df.columns),
        show=False,
        plot_type="violin",
        max_display=10,
        plot_size=[4, 6],
    )

    if show:
        plt.show()

    # SHAP registers forward/backward hook tensors as nn.Parameter on each
    # module. Remove them so state_dict() stays clean for save/load.
    for module in self.modules():
        module._parameters.pop("x", None)
        module._parameters.pop("y", None)