Data API¶

def toy_dataset() -> MuData:
    """Load toy multi-omics dataset from GitHub.

    Returns:
        MuData object:
            - mdata.mod: Multi-omics dictionary (modality name → AnnData).
            - mdata.obs: Clinical metadata with sample IDs as index.
    """
    PREFIX = "https://raw.githubusercontent.com/prism-oncology/customics/refs/heads/main/data"

    protein_df = pd.read_csv(f"{PREFIX}/toy_data/protein.txt", sep="\t", index_col=0).T
    gene_exp_df = pd.read_csv(f"{PREFIX}/toy_data/gene_exp.txt", sep="\t", index_col=0).T
    methyl_df = pd.read_csv(f"{PREFIX}/toy_data/methyl.txt", sep="\t", index_col=0).T

    clinical_df = pd.read_csv(f"{PREFIX}/toy_data/labels.txt", sep="\t", index_col=1, header=0)

    rng = np.random.default_rng(42)
    clinical_df["OS"] = rng.integers(0, 2, size=len(clinical_df))  # 0 = censored, 1 = event
    clinical_df["OS.time"] = rng.integers(200, 3000, size=len(clinical_df))  # days

    mdata = MuData({
        "rna": AnnData(gene_exp_df),
        "protein": AnnData(protein_df),
        "methyl": AnnData(methyl_df),
    })

    mdata.obs = clinical_df
    mdata.obs["cluster.id"] = mdata.obs["cluster.id"].astype(str)

    return mdata

def prepare_input(mdata: MuData, label: str, event: str, surv_time: str) -> None:
    """Validate clinical columns and register them in `mdata.uns`.

    Args:
        mdata: Multi-omics object whose `obs` holds the clinical annotations.
        label: Name of the `mdata.obs` column used as the classification target.
        event: Name of the `mdata.obs` column holding the survival event indicator
            (1 = event, 0 = censored).
        surv_time: Name of the `mdata.obs` column holding the survival time.

    Raises:
        DataValidationError: If any of the given columns is missing from `mdata.obs`.
    """
    for key, column in {Keys.LABEL: label, Keys.EVENT: event, Keys.SURV_TIME: surv_time}.items():
        if column not in mdata.obs:
            raise DataValidationError(f"Column '{column}' not found in mdata.obs")
        mdata.uns[key] = column

def get_shared_samples(mdata: MuData) -> list[str]:
    """Return sample IDs present in every modality.

    Args:
        mdata: Multi-omics object whose modalities' `obs_names` are sample IDs.

    Returns:
        Sorted list of common sample IDs.
    """
    adatas = list(mdata.mod.values())
    common = set(adatas[0].obs_names)
    for adata in adatas[1:]:
        common &= set(adata.obs_names)
    return sorted(common)

def get_sub_mudata(mdata: MuData, shared_samples: list[str]) -> MuData:
    """Subset a MuData to the given samples across every modality.

    Each modality is restricted to the requested samples that it actually
    contains (their intersection), and the clinical `obs`/`uns` are carried
    over so the result is ready to pass to `CustOMICS.fit`.

    Args:
        mdata: Multi-omics object to subset.
        shared_samples: Sample IDs to keep.

    Returns:
        A new MuData containing only `shared_samples`.
    """
    sub = MuData({
        name: adata[[s for s in shared_samples if s in adata.obs_names]].copy() for name, adata in mdata.mod.items()
    })
    sub.obs = mdata.obs.loc[[s for s in shared_samples if s in mdata.obs_names]]
    sub.uns = dict(mdata.uns)
    return sub

def split_mudata(
    mdata: MuData,
    test_size: float = 0.20,
    val_size: float = 0.15,
    random_state: int | None = 42,
) -> tuple[MuData, MuData, MuData]:
    """Split a MuData into train/validation/test subsets by shared samples.

    Samples present in every modality are split into a test set, then the
    remaining samples are split into train and validation sets.

    Args:
        mdata: Multi-omics object to split.
        test_size: Fraction of shared samples held out for the test set.
        val_size: Fraction of the remaining (non-test) samples used for validation.
        random_state: Seed for reproducible splits.

    Returns:
        `(mdata_train, mdata_val, mdata_test)`.
    """
    shared_samples = get_shared_samples(mdata)

    samples_train, samples_test = train_test_split(shared_samples, test_size=test_size, random_state=random_state)
    samples_train, samples_val = train_test_split(samples_train, test_size=val_size, random_state=random_state)

    return (
        get_sub_mudata(mdata, samples_train),
        get_sub_mudata(mdata, samples_val),
        get_sub_mudata(mdata, samples_test),
    )