sonata-learn 0.1.0__py3-none-any.whl

sonata/__init__.py

@@ -0,0 +1,19 @@
+"""
+Sonata: a non-negative matrix factorization toolkit for signature analysis
+==========================================================================
+"""
+
+from . import models
+from . import plot as pl
+from . import tools as tl
+
+__version__ = "0.1.0"
+
+pl.set_sonata_style()
+
+__all__ = [
+    "__version__",
+    "models",
+    "pl",
+    "tl",
+]

sonata/consts.py

@@ -0,0 +1,95 @@
+from matplotlib.colors import LinearSegmentedColormap
+
+NUCLEOTIDES = ["A", "C", "G", "T"]
+
+SBS_TYPES_6 = ["C>A", "C>G", "C>T", "T>A", "T>C", "T>G"]
+SBS_TYPES_96 = [
+    f"{n1}[{sbs_6}]{n2}"
+    for sbs_6 in SBS_TYPES_6
+    for n1 in NUCLEOTIDES
+    for n2 in NUCLEOTIDES
+]
+
+# fmt: off
+INDEL_TYPES_83 = [
+    "DEL.C.1.1", "DEL.C.1.2", 'DEL.C.1.3', "DEL.C.1.4", "DEL.C.1.5", "DEL.C.1.6+",
+    "DEL.T.1.1", "DEL.T.1.2", 'DEL.T.1.3', "DEL.T.1.4", "DEL.T.1.5", "DEL.T.1.6+",
+    "INS.C.1.0", "INS.C.1.1", 'INS.C.1.2', "INS.C.1.3", "INS.C.1.4", "INS.C.1.5+",
+    "INS.T.1.0", "INS.T.1.1", 'INS.T.1.2', "INS.T.1.3", "INS.T.1.4", "INS.T.1.5+",
+    "DEL.repeats.2.1", "DEL.repeats.2.2", "DEL.repeats.2.3",
+    "DEL.repeats.2.4", "DEL.repeats.2.5", "DEL.repeats.2.6+",
+    "DEL.repeats.3.1", "DEL.repeats.3.2", "DEL.repeats.3.3",
+    "DEL.repeats.3.4", "DEL.repeats.3.5", "DEL.repeats.3.6+",
+    "DEL.repeats.4.1", "DEL.repeats.4.2", "DEL.repeats.4.3",
+    "DEL.repeats.4.4", "DEL.repeats.4.5", "DEL.repeats.4.6+",
+    "DEL.repeats.5+.1", "DEL.repeats.5+.2", "DEL.repeats.5+.3",
+    "DEL.repeats.5+.4", "DEL.repeats.5+.5", "DEL.repeats.5+.6+",
+    "INS.repeats.2.0", "INS.repeats.2.1", "INS.repeats.2.2",
+    "INS.repeats.2.3", "INS.repeats.2.4", "INS.repeats.2.5+",
+    "INS.repeats.3.0", "INS.repeats.3.1", "INS.repeats.3.2",
+    "INS.repeats.3.3", "INS.repeats.3.4", "INS.repeats.3.5+",
+    "INS.repeats.4.0", "INS.repeats.4.1", "INS.repeats.4.2",
+    "INS.repeats.4.3", "INS.repeats.4.4", "INS.repeats.4.5+",
+    "INS.repeats.5+.0", "INS.repeats.5+.1", "INS.repeats.5+.2",
+    "INS.repeats.5+.3", "INS.repeats.5+.4", "INS.repeats.5+.5+",
+    "DEL.MH.2.1",
+    "DEL.MH.3.1", "DEL.MH.3.2",
+    "DEL.MH.4.1", "DEL.MH.4.2", "DEL.MH.4.3",
+    "DEL.MH.5+.1", "DEL.MH.5+.2", "DEL.MH.5+.3", "DEL.MH.5+.4", "DEL.MH.5+.5+"
+]
+# fmt: on
+
+# 10 colors
+COLORS_MATHEMATICA = [
+    (0.368417, 0.506779, 0.709798),
+    (0.880722, 0.611041, 0.142051),
+    (0.560181, 0.691569, 0.194885),
+    (0.922526, 0.385626, 0.209179),
+    (0.528288, 0.470624, 0.701351),
+    (0.772079, 0.431554, 0.102387),
+    (0.363898, 0.618501, 0.782349),
+    (1.0, 0.75, 0.0),
+    (0.280264, 0.715, 0.429209),
+    (0.0, 0.0, 0.0),
+]
+
+# Trinucleotide colors for the 96 dimensional mutation spectrum
+COLORS_TRINUCLEOTIDES = [
+    "#427aa1ff",
+    (0.0, 0.0, 0.0),
+    "#d1664aff",
+    (0.78, 0.78, 0.78),
+    "#64b3aaff",
+    (0.89, 0.67, 0.72),
+]
+
+COLORS_SBS96 = [COLORS_TRINUCLEOTIDES[i // 16] for i in range(96)]
+
+DIVERGING_PALETTE = LinearSegmentedColormap.from_list(
+    "sonata_diverging",
+    ["#427aa1ff", "#FAFAFA", "#e07a5fff"],
+)
+
+COLORS_INDEL = [
+    "#FCBD6F",  # 1bp Del C
+    "#FD8001",  # 1bp Del T
+    "#B0DC8B",  # 1bp Ins C
+    "#35A02E",  # 1bp Ins T
+    "#FCC9B4",  # 2bp Del Repeats
+    "#FC896B",  # 3bp Del Repeats
+    "#F04432",  # 4bp Del Repeats
+    "#BC1A1A",  # 5+ bp Del Repeats
+    "#CFE0F0",  # 2bp Ins Repeats
+    "#94C3DF",  # 3bp Ins Repeats
+    "#4A98C8",  # 4bp Ins Repeats
+    "#1665AA",  # 5+ bp Ins Repeats
+    "#E1E0ED",  # 2bp Del MH
+    "#B5B5D8",  # 3bp Del MH
+    "#8683BC",  # 4bp Del MH
+    "#624099",  # 5+bp Del MH
+]
+
+# 12 * 6 + 11 = 83 colors
+n_times = 12 * [6] + [1, 2, 3, 5]
+COLORS_INDEL83 = [n * [col] for n, col in zip(n_times, COLORS_INDEL)]
+COLORS_INDEL83 = [col for color_list in COLORS_INDEL83 for col in color_list]

sonata/initialization/__init__.py

@@ -0,0 +1 @@
+""""""

sonata/initialization/initialize.py

@@ -0,0 +1,382 @@
+"""
+Initialization methods for non-negative matrix factorization (NMF) models.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import anndata as ad
+import numpy as np
+
+from ..utils import (
+    dict_checker,
+    normalize_WH,
+    shape_checker,
+    type_checker,
+    value_checker,
+)
+from .methods import (
+    _INIT_METHODS,
+    _Init_methods,
+    init_custom,
+    init_flat,
+    init_nndsvd,
+    init_random,
+    init_separableNMF,
+)
+
+EPSILON = np.finfo(np.float32).eps
+
+# allowed given parameters
+GIVEN_PARAMETERS_STANDARD_NMF = ["asignatures"]
+GIVEN_PARAMETERS_CORNET = [
+    "asignatures",
+    "signature_offsets",
+    "sample_offsets",
+    "signature_embeddings",
+    "sample_embeddings",
+    "variance",
+]
+
+
+def initialize_mat(
+    data_mat: np.ndarray,
+    n_signatures: int,
+    method: _Init_methods = "nndsvd",
+    given_signatures_mat: np.ndarray | None = None,
+    **kwargs,
+) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Initialize the signature and exposure matrices.
+
+    Inputs
+    ------
+    data_mat : np.ndarray
+        shape (n_samples, n_features)
+
+    n_signatures : int
+
+    method : str
+        initialization method. One of 'custom', 'flat',
+        'nndsvd', 'nndsvda', 'nndsvdar', 'random', 'separableNMF'.
+
+    given_signatures_mat : np.ndarray, optional
+        At most 'n_signatures' many signatures can be provided to
+        overwrite some of the initialized signatures. This does not
+        change the initialized exposures.
+
+    kwargs : dict
+        Any keyword arguments to be passed to the initialization method.
+        This includes, for example, a possible 'seed' keyword argument
+        for all stochastic methods.
+
+    Returns
+    -------
+    signatures_mat : np.ndarray
+        shape (n_signatures, n_features)
+
+    exposures_mat : np.ndarray
+        shape (n_samples, n_signatures)
+    """
+    value_checker("method", method, _INIT_METHODS)
+
+    if method == "custom":
+        matrices = init_custom(data_mat, n_signatures, **kwargs)
+    elif method == "flat":
+        matrices = init_flat(data_mat, n_signatures)
+    elif method in ["nndsvd", "nndsvda", "nndsvdar"]:
+        # mypy does not recognize that 'method' is compatible
+        # with Literal["nndsvd", "nndsvda", "nndsvdar"]
+        matrices = init_nndsvd(
+            data_mat, n_signatures, method=method, **kwargs  # type: ignore[arg-type] # noqa: E501
+        )
+    elif method == "random":
+        matrices = init_random(data_mat, n_signatures, **kwargs)
+    else:
+        matrices = init_separableNMF(data_mat, n_signatures, **kwargs)
+
+    signatures_mat, exposures_mat = matrices
+
+    if given_signatures_mat is not None:
+        type_checker("given_signatures_mat", given_signatures_mat, np.ndarray)
+        given_n_signatures, given_n_features = given_signatures_mat.shape
+
+        if given_n_features != data_mat.shape[1]:
+            raise ValueError(
+                "The given signature matrix has a different number of features "
+                "than the data."
+            )
+        if given_n_signatures > n_signatures:
+            raise ValueError("The given signature matrix contains too many signatures.")
+
+        signatures_mat[:given_n_signatures, :] = given_signatures_mat.copy()
+
+    W, H = normalize_WH(signatures_mat.T, exposures_mat.T)
+    W, H = W.clip(EPSILON), H.clip(EPSILON)
+    signatures_mat, exposures_mat = W.T, H.T
+    return signatures_mat, exposures_mat
+
+
+def check_given_asignatures(
+    given_asignatures: ad.AnnData, adata: ad.AnnData, n_signatures: int
+) -> None:
+    """
+    Check if the given signatures are compatible with
+    the input data and the number of signatures to be initialized.
+    The number of given signatures can be less or equal to the number of
+    signatures specified.
+
+    Inputs
+    ------
+    given_asignatures: AnnData
+        Known signatures that should be fixed by the algorithm.
+
+    adata: ad.AnnData
+        Input data.
+
+    n_signatures: int
+        The number of signatures to initialize.
+    """
+    type_checker("given_asignatures", given_asignatures, ad.AnnData)
+    if given_asignatures.n_vars != adata.n_vars:
+        raise ValueError(
+            "The given signatures have a different number of features than the data."
+        )
+    if not all(given_asignatures.var_names == adata.var_names):
+        raise ValueError(
+            "The features of the given signatures and the data are not identical."
+        )
+    if given_asignatures.n_obs > n_signatures:
+        raise ValueError(
+            "The number of given signatures exceeds "
+            "the number of signatures to initialize."
+        )
+
+
+def initialize_base(
+    adata: ad.AnnData,
+    n_signatures: int,
+    method: _Init_methods = "nndsvd",
+    given_asignatures: ad.AnnData | None = None,
+    **kwargs,
+) -> tuple[ad.AnnData, np.ndarray]:
+    """
+    Initialize the signature anndata object and the exposure matrix.
+    The anndata object is unchanged and the exposure matrix is returned.
+
+    Inputs
+    ------
+    adata : ad.AnnData
+
+    n_signatures : int
+
+    method : str
+        initialization method. One of 'custom', 'flat',
+        'nndsvd', 'nndsvda', 'nndsvdar', 'random', 'separableNMF'.
+
+    given_asignatures : ad.AnnData, optional
+        At most 'n_signatures' many signatures can be provided to
+        overwrite some of the initialized signatures. This does not
+        change the initialized exposures.
+
+    kwargs : dict
+        Any keyword arguments to be passed to the initialization method.
+        This includes, for example, a possible 'seed' keyword argument
+        for all stochastic methods.
+
+    Returns
+    -------
+    asignatures : ad.AnnData
+        Annotated signature matrix of shape (n_signatures, n_features)
+
+    exposures_mat : np.ndarray
+        shape (n_samples, n_signatures)
+    """
+    if given_asignatures is not None:
+        check_given_asignatures(given_asignatures, adata, n_signatures)
+        given_signatures_mat = given_asignatures.X
+    else:
+        given_signatures_mat = None
+
+    signatures_mat, exposures_mat = initialize_mat(
+        adata.X, n_signatures, method, given_signatures_mat, **kwargs
+    )
+    asignatures = ad.AnnData(signatures_mat)
+    asignatures.var_names = adata.var_names
+    asignatures.obs_names = [f"Sig{k+1}" for k in range(n_signatures)]
+
+    # keep signature annotations
+    if given_asignatures is not None:
+        n_given_signatures = given_asignatures.n_obs
+        asignatures.obs_names = np.roll(asignatures.obs_names, n_given_signatures)
+        asignatures = ad.concat(
+            [given_asignatures, asignatures[n_given_signatures:, :]], join="outer"
+        )
+
+    return asignatures, exposures_mat
+
+
+def check_given_parameters_standard_nmf(
+    adata: ad.AnnData,
+    n_signatures: int,
+    given_parameters: dict[str, Any],
+) -> None:
+    dict_checker("given_parameters", given_parameters, GIVEN_PARAMETERS_STANDARD_NMF)
+
+    if "asignatures" in given_parameters:
+        check_given_asignatures(given_parameters["asignatures"], adata, n_signatures)
+
+
+def initialize_standard_nmf(
+    adata: ad.AnnData,
+    n_signatures: int,
+    method: _Init_methods = "nndsvd",
+    given_parameters: dict[str, Any] | None = None,
+    **kwargs,
+) -> ad.AnnData:
+    given_parameters = {} if given_parameters is None else given_parameters.copy()
+    check_given_parameters_standard_nmf(adata, n_signatures, given_parameters)
+
+    if "asignatures" in given_parameters:
+        given_asignatures = given_parameters["asignatures"]
+    else:
+        given_asignatures = None
+
+    asignatures, exposures_mat = initialize_base(
+        adata,
+        n_signatures,
+        method,
+        given_asignatures,
+        **kwargs,
+    )
+    adata.obsm["exposures"] = exposures_mat
+    return asignatures
+
+
+def check_given_offsets_cornet(
+    given_offsets: np.ndarray, n_offsets_expected: int, name: str
+) -> None:
+    """
+    Check if the given sample or signature offsets match the expected shape.
+    """
+    type_checker(name, given_offsets, np.ndarray)
+    shape_checker(name, given_offsets, (n_offsets_expected,))
+
+
+def check_given_embeddings_cornet(
+    given_embeddings: np.ndarray,
+    n_embeddings_expected: int,
+    dim_embeddings_expected: int,
+    name: str,
+) -> None:
+    type_checker(name, given_embeddings, np.ndarray)
+    shape_checker(
+        name, given_embeddings, (n_embeddings_expected, dim_embeddings_expected)
+    )
+
+
+def check_given_parameters_cornet(
+    adata: ad.AnnData,
+    n_signatures: int,
+    dim_embeddings: int,
+    given_parameters: dict[str, Any],
+) -> None:
+    dict_checker("given_parameters", given_parameters, GIVEN_PARAMETERS_CORNET)
+
+    if "asignatures" in given_parameters:
+        check_given_asignatures(given_parameters["asignatures"], adata, n_signatures)
+
+    if "signature_offsets" in given_parameters:
+        check_given_offsets_cornet(
+            given_parameters["signature_offsets"],
+            n_signatures,
+            "given_signature_offsets",
+        )
+    if "sample_offsets" in given_parameters:
+        check_given_offsets_cornet(
+            given_parameters["sample_offsets"], adata.n_obs, "given_sample_offsets"
+        )
+    if "signature_embeddings" in given_parameters:
+        check_given_embeddings_cornet(
+            given_parameters["signature_embeddings"],
+            n_signatures,
+            dim_embeddings,
+            "given_signature_embeddings",
+        )
+    if "sample_embeddings" in given_parameters:
+        check_given_embeddings_cornet(
+            given_parameters["sample_embeddings"],
+            adata.n_obs,
+            dim_embeddings,
+            "given_sample_embeddings",
+        )
+    if "variance" in given_parameters:
+        given_variance = given_parameters["variance"]
+        type_checker("given_variance", given_variance, [float, int])
+        if given_variance <= 0.0:
+            raise ValueError("The variance has to be a positive real number.")
+
+
+def initialize_cornet(
+    adata: ad.AnnData,
+    n_signatures: int,
+    dim_embeddings: int,
+    method: _Init_methods = "nndsvd",
+    given_parameters: dict[str, Any] | None = None,
+    **kwargs,
+) -> tuple[ad.AnnData, float]:
+    if method == "custom":
+        raise ValueError(
+            "Custom parameter initializations are currently not supported "
+            "for correlated NMF."
+        )
+
+    given_parameters = {} if given_parameters is None else given_parameters.copy()
+    check_given_parameters_cornet(adata, n_signatures, dim_embeddings, given_parameters)
+
+    if "asignatures" in given_parameters:
+        given_asignatures = given_parameters["asignatures"]
+    else:
+        given_asignatures = None
+
+    asignatures, _ = initialize_base(
+        adata,
+        n_signatures,
+        method,
+        given_asignatures,
+        **kwargs,
+    )
+
+    if "signature_offsets" in given_parameters:
+        asignatures.obs["offsets"] = given_parameters["signature_offsets"]
+    else:
+        asignatures.obs["offsets"] = np.zeros(n_signatures)
+
+    if "sample_offsets" in given_parameters:
+        adata.obs["offsets"] = given_parameters["sample_offsets"]
+    else:
+        adata.obs["offsets"] = np.zeros(adata.n_obs)
+
+    if "signature_embeddings" in given_parameters:
+        asignatures.obsm["embeddings"] = given_parameters["signature_embeddings"]
+    else:
+        asignatures.obsm["embeddings"] = np.random.multivariate_normal(
+            np.zeros(dim_embeddings), np.identity(dim_embeddings), size=n_signatures
+        )
+
+    if "sample_embeddings" in given_parameters:
+        adata.obsm["embeddings"] = given_parameters["sample_embeddings"]
+    else:
+        adata.obsm["embeddings"] = np.random.multivariate_normal(
+            np.zeros(dim_embeddings),
+            np.identity(dim_embeddings),
+            size=adata.n_obs,
+        )
+
+    if "variance" in given_parameters:
+        variance = float(given_parameters["variance"])
+    else:
+        variance = 1.0
+
+    return asignatures, variance

sonata/initialization/methods.py

@@ -0,0 +1,135 @@
+"""
+Initialization methods for non-negative matrix factorization (NMF)
+"""
+
+from __future__ import annotations
+
+from typing import Literal, get_args
+
+import numpy as np
+from sklearn.decomposition import _nmf as sknmf
+
+from ..utils import shape_checker, type_checker
+
+EPSILON = np.finfo(np.float32).eps
+_Init_methods = Literal[
+    "custom",
+    "flat",
+    "nndsvd",
+    "nndsvda",
+    "nndsvdar",
+    "random",
+    "separableNMF",
+]
+_INIT_METHODS = get_args(_Init_methods)
+
+
+def init_custom(
+    data_mat: np.ndarray,
+    n_signatures: int,
+    signatures_mat: np.ndarray,
+    exposures_mat: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Perform type and shape checks on custom signature and
+    exposure matrix initializations.
+
+    Inputs:
+    -------
+    data_mat: np.ndarray
+        shape (n_samples, n_features)
+
+    n_signatures: int
+
+    signatures_mat: np.ndarray
+        shape (n_signatures, n_features)
+
+    exposures_mat: np.ndarray
+        shape (n_samples, n_signatures)
+    """
+    type_checker("signatures_mat", signatures_mat, np.ndarray)
+    type_checker("exposures_mat", exposures_mat, np.ndarray)
+    n_samples, n_features = data_mat.shape
+    shape_checker("signatures_mat", signatures_mat, (n_signatures, n_features))
+    shape_checker("exposures_mat", exposures_mat, (n_samples, n_signatures))
+    return signatures_mat, exposures_mat
+
+
+def init_flat(data_mat: np.ndarray, n_signatures: int) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Initialize the signature and exposure matrices with one float, respectively.
+    """
+    n_features = data_mat.shape[1]
+    signatures_mat = np.full((n_signatures, n_features), 1 / n_features)
+    exposures = np.sum(data_mat, axis=1) / n_signatures
+    exposures_mat = np.tile(exposures, (n_signatures, 1)).T
+    return signatures_mat, exposures_mat
+
+
+def init_nndsvd(
+    data_mat: np.ndarray,
+    n_signatures: int,
+    method: Literal["nndsvd", "nndsvda", "nndsvdar"] = "nndsvd",
+    seed: int | None = None,
+) -> tuple[np.ndarray, np.ndarray]:
+    """
+    A wrapper around the non-negative double singular value decomposition (NNDSVD)
+    initialization methods "nndsvd", "nndsvda" and "nndsvdar" from scikit-learn.
+    """
+    if seed is not None:
+        np.random.seed(seed)
+
+    # pylint: disable-next=W0212
+    exposures_mat, signatures_mat = sknmf._initialize_nmf(
+        data_mat, n_signatures, init=method
+    )
+    return signatures_mat, exposures_mat
+
+
+def init_random(
+    data_mat: np.ndarray, n_signatures: int, seed: int | None = None
+) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Initialize each signature by drawing from the uniform
+    distribution on the simplex.
+    Initialize the exposures of each sample as a scaled sample
+    from the uniform distribution on a simplex.
+    The scaling is chosen such that the expected total exposure is equal to
+    the column sum of that sample in the count matrix X.
+    """
+    if seed is not None:
+        np.random.seed(seed)
+
+    n_samples, n_features = data_mat.shape
+    signatures_mat = np.random.dirichlet(np.ones(n_features), size=n_signatures)
+    scaling = np.sum(data_mat, axis=1)
+    exposures_mat = scaling[:, np.newaxis] * np.random.dirichlet(
+        np.ones(n_signatures), size=n_samples
+    )
+    return signatures_mat, exposures_mat
+
+
+def init_separableNMF(
+    data_mat: np.ndarray, n_signatures: int, seed: int | None = None
+) -> tuple[np.ndarray, np.ndarray]:
+    r"""
+    This code is following Algorithm 1 from "Fast and Robust Recursive
+    Algorithms for Separable Nonnegative Matrix Factorization"
+    (Gillis and Vavasis, 2013), with the canonical choice of
+    f(x) = \| x \|_2^2 as the strongly convex function f satisfying
+    Assumption 2 from the paper.
+    """
+    signature_indices = np.empty(n_signatures, dtype=int)
+    R = data_mat.T / np.sum(data_mat.T, axis=0)
+
+    for k in range(n_signatures):
+        column_norms = np.sum(R**2, axis=0)
+        kstar = np.argmax(column_norms)
+        u = R[:, kstar]
+        R = (np.identity(R.shape[0]) - np.outer(u, u) / column_norms[kstar]) @ R
+        signature_indices[k] = kstar
+
+    signatures_mat = data_mat[signature_indices, :].astype(float)
+    signatures_mat /= signatures_mat.sum(axis=1)[:, np.newaxis]
+    _, exposures_mat = init_random(data_mat, n_signatures, seed=seed)
+    return signatures_mat, exposures_mat

sonata/models/__init__.py

@@ -0,0 +1,13 @@
+"""
+A collection of NMF algorithms
+"""
+
+from .cornet import Cornet
+from .mvnmf import MvNMF
+from .nmf import NMF
+
+__all__ = [
+    "Cornet",
+    "NMF",
+    "MvNMF",
+]