smftools 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

smftools/tools/calculate_nmf.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
+import numpy as np
+
+from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
+
+if TYPE_CHECKING:
+    import anndata as ad
+
+logger = get_logger(__name__)
+
+
+def calculate_nmf(
+    adata: "ad.AnnData",
+    layer: str | None = "nan_half",
+    var_filters: Sequence[str] | None = None,
+    n_components: int = 2,
+    max_iter: int = 200,
+    random_state: int = 0,
+    overwrite: bool = True,
+    embedding_key: str = "X_nmf",
+    components_key: str = "H_nmf",
+    uns_key: str = "nmf",
+) -> "ad.AnnData":
+    """Compute a low-dimensional NMF embedding.
+
+    Args:
+        adata: AnnData object to update.
+        layer: Layer name to use for NMF (``None`` uses ``adata.X``).
+        var_filters: Optional list of var masks to subset features.
+        n_components: Number of NMF components to compute.
+        max_iter: Maximum number of NMF iterations.
+        random_state: Random seed for the NMF initializer.
+        overwrite: Whether to recompute if the embedding already exists.
+        embedding_key: Key for the embedding in ``adata.obsm``.
+        components_key: Key for the components matrix in ``adata.varm``.
+        uns_key: Key for metadata stored in ``adata.uns``.
+
+    Returns:
+        anndata.AnnData: Updated AnnData object.
+    """
+    from scipy.sparse import issparse
+
+    require("sklearn", extra="ml-base", purpose="NMF calculation")
+    from sklearn.decomposition import NMF
+
+    has_embedding = embedding_key in adata.obsm
+    has_components = components_key in adata.varm
+    if has_embedding and has_components and not overwrite:
+        logger.info("NMF embedding and components already present; skipping recomputation.")
+        return adata
+    if has_embedding and not has_components and not overwrite:
+        logger.info("NMF embedding present without components; recomputing to store components.")
+
+    subset_mask = None
+    if var_filters:
+        subset_mask = np.logical_or.reduce([adata.var[f].values for f in var_filters])
+        adata_subset = adata[:, subset_mask].copy()
+        logger.info(
+            "Subsetting adata: retained %s features based on filters %s",
+            adata_subset.shape[1],
+            var_filters,
+        )
+    else:
+        adata_subset = adata.copy()
+        logger.info("No var filters provided. Using all features.")
+
+    data = adata_subset.layers[layer] if layer else adata_subset.X
+    if issparse(data):
+        data = data.copy()
+        if data.data.size and np.isnan(data.data).any():
+            logger.warning("NaNs detected in sparse data, filling with 0.5 before NMF.")
+            data.data = np.nan_to_num(data.data, nan=0.5)
+        if data.data.size and (data.data < 0).any():
+            logger.warning("Negative values detected in sparse data, clipping to 0 for NMF.")
+            data.data[data.data < 0] = 0
+    else:
+        if np.isnan(data).any():
+            logger.warning("NaNs detected, filling with 0.5 before NMF.")
+            data = np.nan_to_num(data, nan=0.5)
+        if (data < 0).any():
+            logger.warning("Negative values detected, clipping to 0 for NMF.")
+            data = np.clip(data, a_min=0, a_max=None)
+
+    model = NMF(
+        n_components=n_components,
+        init="nndsvda",
+        max_iter=max_iter,
+        random_state=random_state,
+    )
+    embedding = model.fit_transform(data)
+    components = model.components_.T
+
+    if subset_mask is not None:
+        components_matrix = np.zeros((adata.shape[1], components.shape[1]))
+        components_matrix[subset_mask, :] = components
+    else:
+        components_matrix = components
+
+    adata.obsm[embedding_key] = embedding
+    adata.varm[components_key] = components_matrix
+    adata.uns[uns_key] = {
+        "n_components": n_components,
+        "max_iter": max_iter,
+        "random_state": random_state,
+        "layer": layer,
+        "var_filters": list(var_filters) if var_filters else None,
+        "components_key": components_key,
+    }
+
+    logger.info(
+        "Stored: adata.obsm['%s'] and adata.varm['%s']",
+        embedding_key,
+        components_key,
+    )
+    return adata

smftools/tools/calculate_umap.py

@@ -19,6 +19,7 @@ def calculate_umap(
     knn_neighbors: int = 100,
     overwrite: bool = True,
     threads: int = 8,
+    random_state: int | None = 0,
 ) -> "ad.AnnData":
     """Compute PCA, neighbors, and UMAP embeddings.
 
@@ -37,9 +38,11 @@ def calculate_umap(
     import os
 
     import numpy as np
+    import scipy.linalg as spla
+    import scipy.sparse as sp
 
-    sc = require("scanpy", extra="scanpy", purpose="UMAP calculation")
-    from scipy.sparse import issparse
+    umap = require("umap", extra="umap", purpose="UMAP calculation")
+    pynndescent = require("pynndescent", extra="umap", purpose="KNN graph computation")
 
     os.environ["OMP_NUM_THREADS"] = str(threads)
 
@@ -59,7 +62,7 @@ def calculate_umap(
     # Step 2: NaN handling inside layer
     if layer:
         data = adata_subset.layers[layer]
-        if not issparse(data):
+        if not sp.issparse(data):
             if np.isnan(data).any():
                 logger.warning("NaNs detected, filling with 0.5 before PCA + neighbors.")
                 data = np.nan_to_num(data, nan=0.5)
@@ -75,18 +78,98 @@ def calculate_umap(
     if "X_umap" not in adata_subset.obsm or overwrite:
         n_pcs = min(adata_subset.shape[1], n_pcs)
         logger.info("Running PCA with n_pcs=%s", n_pcs)
-        sc.pp.pca(adata_subset, layer=layer)
-        logger.info("Running neighborhood graph")
-        sc.pp.neighbors(adata_subset, use_rep="X_pca", n_pcs=n_pcs, n_neighbors=knn_neighbors)
+
+        if layer:
+            matrix = adata_subset.layers[layer]
+        else:
+            matrix = adata_subset.X
+
+        if sp.issparse(matrix):
+            logger.warning("Converting sparse matrix to dense for PCA.")
+            matrix = matrix.toarray()
+
+        matrix = np.asarray(matrix, dtype=float)
+        mean = matrix.mean(axis=0)
+        centered = matrix - mean
+
+        if centered.shape[0] == 0 or centered.shape[1] == 0:
+            raise ValueError("PCA requires a non-empty matrix.")
+
+        if n_pcs <= 0:
+            raise ValueError("n_pcs must be positive.")
+
+        if centered.shape[1] <= n_pcs:
+            n_pcs = centered.shape[1]
+
+        if centered.shape[0] < n_pcs:
+            n_pcs = centered.shape[0]
+
+        u, s, vt = spla.svd(centered, full_matrices=False)
+
+        u = u[:, :n_pcs]
+        s = s[:n_pcs]
+        vt = vt[:n_pcs]
+
+        adata_subset.obsm["X_pca"] = u * s
+        adata_subset.varm["PCs"] = vt.T
+
+        logger.info("Running neighborhood graph with pynndescent (n_neighbors=%s)", knn_neighbors)
+        n_neighbors = min(knn_neighbors, max(1, adata_subset.n_obs - 1))
+        nn_index = pynndescent.NNDescent(
+            adata_subset.obsm["X_pca"],
+            n_neighbors=n_neighbors,
+            metric="euclidean",
+            random_state=random_state,
+            n_jobs=threads,
+        )
+        knn_indices, knn_dists = nn_index.neighbor_graph
+
+        rows = np.repeat(np.arange(adata_subset.n_obs), n_neighbors)
+        cols = knn_indices.reshape(-1)
+        distances = sp.coo_matrix(
+            (knn_dists.reshape(-1), (rows, cols)),
+            shape=(adata_subset.n_obs, adata_subset.n_obs),
+        ).tocsr()
+        adata_subset.obsp["distances"] = distances
+
         logger.info("Running UMAP")
-        sc.tl.umap(adata_subset)
+        umap_model = umap.UMAP(
+            n_neighbors=n_neighbors,
+            n_components=2,
+            metric="euclidean",
+            random_state=random_state,
+        )
+        adata_subset.obsm["X_umap"] = umap_model.fit_transform(adata_subset.obsm["X_pca"])
+
+        try:
+            from umap.umap_ import fuzzy_simplicial_set
+
+            fuzzy_result = fuzzy_simplicial_set(
+                adata_subset.obsm["X_pca"],
+                n_neighbors=n_neighbors,
+                random_state=random_state,
+                metric="euclidean",
+                knn_indices=knn_indices,
+                knn_dists=knn_dists,
+            )
+            connectivities = fuzzy_result[0] if isinstance(fuzzy_result, tuple) else fuzzy_result
+        except TypeError:
+            connectivities = umap_model.graph_
+
+        adata_subset.obsp["connectivities"] = connectivities
 
     # Step 4: Store results in original adata
     adata.obsm["X_pca"] = adata_subset.obsm["X_pca"]
     adata.obsm["X_umap"] = adata_subset.obsm["X_umap"]
     adata.obsp["distances"] = adata_subset.obsp["distances"]
     adata.obsp["connectivities"] = adata_subset.obsp["connectivities"]
-    adata.uns["neighbors"] = adata_subset.uns["neighbors"]
+    adata.uns["neighbors"] = {
+        "params": {
+            "n_neighbors": knn_neighbors,
+            "method": "pynndescent",
+            "metric": "euclidean",
+        }
+    }
 
     # Fix varm["PCs"] shape mismatch
     pc_matrix = np.zeros((adata.shape[1], adata_subset.varm["PCs"].shape[1]))

smftools/tools/rolling_nn_distance.py

@@ -0,0 +1,235 @@
+from __future__ import annotations
+
+import ast
+import json
+from typing import TYPE_CHECKING, Optional, Sequence, Tuple
+
+import numpy as np
+
+from smftools.logging_utils import get_logger
+
+if TYPE_CHECKING:
+    import anndata as ad
+
+logger = get_logger(__name__)
+
+
+def _pack_bool_to_u64(B: np.ndarray) -> np.ndarray:
+    """
+    Pack a boolean (or 0/1) matrix (n, w) into uint64 blocks (n, ceil(w/64)).
+    Safe w.r.t. contiguity/layout.
+    """
+    B = np.asarray(B, dtype=np.uint8)
+    packed_u8 = np.packbits(B, axis=1)  # (n, ceil(w/8)) uint8
+
+    n, nb = packed_u8.shape
+    pad = (-nb) % 8
+    if pad:
+        packed_u8 = np.pad(packed_u8, ((0, 0), (0, pad)), mode="constant", constant_values=0)
+
+    packed_u8 = np.ascontiguousarray(packed_u8)
+
+    # group 8 bytes -> uint64
+    packed_u64 = packed_u8.reshape(n, -1, 8).view(np.uint64).reshape(n, -1)
+    return packed_u64
+
+
+def _popcount_u64_matrix(A_u64: np.ndarray) -> np.ndarray:
+    """
+    Popcount for an array of uint64, vectorized and portable across NumPy versions.
+
+    Returns an integer array with the SAME SHAPE as A_u64.
+    """
+    A_u64 = np.ascontiguousarray(A_u64)
+    # View as bytes; IMPORTANT: reshape to add a trailing byte axis of length 8
+    b = A_u64.view(np.uint8).reshape(A_u64.shape + (8,))
+    # unpack bits within that byte axis -> (..., 64), then sum
+    return np.unpackbits(b, axis=-1).sum(axis=-1)
+
+
+def rolling_window_nn_distance(
+    adata,
+    layer: Optional[str] = None,
+    window: int = 15,
+    step: int = 2,
+    min_overlap: int = 10,
+    return_fraction: bool = True,
+    block_rows: int = 256,
+    block_cols: int = 2048,
+    store_obsm: Optional[str] = "rolling_nn_dist",
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Rolling-window nearest-neighbor distance per read, overlap-aware.
+
+    Distance between reads i,j in a window:
+      - use only positions where BOTH are observed (non-NaN)
+      - require overlap >= min_overlap
+      - mismatch = count(x_i != x_j) over overlapped positions
+      - distance = mismatch/overlap (if return_fraction) else mismatch
+
+    Returns
+    -------
+    out : (n_obs, n_windows) float
+        Nearest-neighbor distance per read per window (NaN if no valid neighbor).
+    starts : (n_windows,) int
+        Window start indices in var-space.
+    """
+    X = adata.layers[layer] if layer is not None else adata.X
+    X = X.toarray() if hasattr(X, "toarray") else np.asarray(X)
+
+    n, p = X.shape
+    if window > p:
+        raise ValueError(f"window={window} is larger than n_vars={p}")
+    if window <= 0:
+        raise ValueError("window must be > 0")
+    if step <= 0:
+        raise ValueError("step must be > 0")
+    if min_overlap <= 0:
+        raise ValueError("min_overlap must be > 0")
+
+    starts = np.arange(0, p - window + 1, step, dtype=int)
+    nW = len(starts)
+    out = np.full((n, nW), np.nan, dtype=float)
+
+    for wi, s in enumerate(starts):
+        wX = X[:, s : s + window]  # (n, window)
+
+        # observed mask; values as 0/1 where observed, 0 elsewhere
+        M = ~np.isnan(wX)
+        V = np.where(M, wX, 0).astype(np.float32)
+
+        # ensure binary 0/1
+        V = (V > 0).astype(np.uint8)
+
+        M64 = _pack_bool_to_u64(M)
+        V64 = _pack_bool_to_u64(V.astype(bool))
+
+        best = np.full(n, np.inf, dtype=float)
+
+        for i0 in range(0, n, block_rows):
+            i1 = min(n, i0 + block_rows)
+            Mi = M64[i0:i1]  # (bi, nb)
+            Vi = V64[i0:i1]
+            bi = i1 - i0
+
+            local_best = np.full(bi, np.inf, dtype=float)
+
+            for j0 in range(0, n, block_cols):
+                j1 = min(n, j0 + block_cols)
+                Mj = M64[j0:j1]  # (bj, nb)
+                Vj = V64[j0:j1]
+                bj = j1 - j0
+
+                overlap_counts = np.zeros((bi, bj), dtype=np.uint16)
+                mismatch_counts = np.zeros((bi, bj), dtype=np.uint16)
+
+                for k in range(Mi.shape[1]):
+                    ob = (Mi[:, k][:, None] & Mj[:, k][None, :]).astype(np.uint64)
+                    overlap_counts += _popcount_u64_matrix(ob).astype(np.uint16)
+
+                    mb = ((Vi[:, k][:, None] ^ Vj[:, k][None, :]) & ob).astype(np.uint64)
+                    mismatch_counts += _popcount_u64_matrix(mb).astype(np.uint16)
+
+                ok = overlap_counts >= min_overlap
+                if not np.any(ok):
+                    continue
+
+                dist = np.full((bi, bj), np.inf, dtype=float)
+                if return_fraction:
+                    dist[ok] = mismatch_counts[ok] / overlap_counts[ok]
+                else:
+                    dist[ok] = mismatch_counts[ok].astype(float)
+
+                # exclude self comparisons (diagonal) when blocks overlap
+                if (i0 <= j1) and (j0 <= i1):
+                    ii = np.arange(i0, i1)
+                    jj = ii[(ii >= j0) & (ii < j1)]
+                    if jj.size:
+                        dist[(jj - i0), (jj - j0)] = np.inf
+
+                local_best = np.minimum(local_best, dist.min(axis=1))
+
+            best[i0:i1] = local_best
+
+        best[~np.isfinite(best)] = np.nan
+        out[:, wi] = best
+
+    if store_obsm is not None:
+        adata.obsm[store_obsm] = out
+        adata.uns[f"{store_obsm}_starts"] = starts
+        adata.uns[f"{store_obsm}_window"] = int(window)
+        adata.uns[f"{store_obsm}_step"] = int(step)
+        adata.uns[f"{store_obsm}_min_overlap"] = int(min_overlap)
+        adata.uns[f"{store_obsm}_return_fraction"] = bool(return_fraction)
+        adata.uns[f"{store_obsm}_layer"] = layer if layer is not None else "X"
+
+    return out, starts
+
+
+def assign_rolling_nn_results(
+    parent_adata: "ad.AnnData",
+    subset_adata: "ad.AnnData",
+    values: np.ndarray,
+    starts: np.ndarray,
+    obsm_key: str,
+    window: int,
+    step: int,
+    min_overlap: int,
+    return_fraction: bool,
+    layer: Optional[str],
+) -> None:
+    """
+    Assign rolling NN results computed on a subset back onto a parent AnnData.
+
+    Parameters
+    ----------
+    parent_adata : AnnData
+        Parent AnnData that should store the combined results.
+    subset_adata : AnnData
+        Subset AnnData used to compute `values`.
+    values : np.ndarray
+        Rolling NN output with shape (n_subset_obs, n_windows).
+    starts : np.ndarray
+        Window start indices corresponding to `values`.
+    obsm_key : str
+        Key to store results under in parent_adata.obsm.
+    window : int
+        Rolling window size (stored in parent_adata.uns).
+    step : int
+        Rolling window step size (stored in parent_adata.uns).
+    min_overlap : int
+        Minimum overlap (stored in parent_adata.uns).
+    return_fraction : bool
+        Whether distances are fractional (stored in parent_adata.uns).
+    layer : str | None
+        Layer used for calculations (stored in parent_adata.uns).
+    """
+    n_obs = parent_adata.n_obs
+    n_windows = values.shape[1]
+
+    if obsm_key not in parent_adata.obsm:
+        parent_adata.obsm[obsm_key] = np.full((n_obs, n_windows), np.nan, dtype=float)
+        parent_adata.uns[f"{obsm_key}_starts"] = starts
+        parent_adata.uns[f"{obsm_key}_window"] = int(window)
+        parent_adata.uns[f"{obsm_key}_step"] = int(step)
+        parent_adata.uns[f"{obsm_key}_min_overlap"] = int(min_overlap)
+        parent_adata.uns[f"{obsm_key}_return_fraction"] = bool(return_fraction)
+        parent_adata.uns[f"{obsm_key}_layer"] = layer if layer is not None else "X"
+    else:
+        existing = parent_adata.obsm[obsm_key]
+        if existing.shape[1] != n_windows:
+            raise ValueError(
+                f"Existing obsm[{obsm_key!r}] has {existing.shape[1]} windows; "
+                f"new values have {n_windows} windows."
+            )
+        existing_starts = parent_adata.uns.get(f"{obsm_key}_starts")
+        if existing_starts is not None and not np.array_equal(existing_starts, starts):
+            raise ValueError(
+                f"Existing obsm[{obsm_key!r}] has different window starts than new values."
+            )
+
+    parent_indexer = parent_adata.obs_names.get_indexer(subset_adata.obs_names)
+    if (parent_indexer < 0).any():
+        raise ValueError("Subset AnnData contains obs not present in parent AnnData.")
+
+    parent_adata.obsm[obsm_key][parent_indexer, :] = values

smftools/tools/tensor_factorization.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Iterable, Sequence
+
+import numpy as np
+
+from smftools.constants import MODKIT_EXTRACT_SEQUENCE_BASE_TO_INT
+from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
+
+if TYPE_CHECKING:
+    import anndata as ad
+
+logger = get_logger(__name__)
+
+
+def build_sequence_one_hot_and_mask(
+    encoded_sequences: np.ndarray,
+    *,
+    bases: Sequence[str] = ("A", "C", "G", "T"),
+    dtype: np.dtype | type[np.floating] = np.float32,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Build one-hot encoded reads and a seen/unseen mask.
+
+    Args:
+        encoded_sequences: Integer-encoded sequences shaped (n_reads, seq_len).
+        bases: Bases to one-hot encode.
+        dtype: Output dtype for the one-hot tensor.
+
+    Returns:
+        Tuple of (one_hot_tensor, mask) where:
+            - one_hot_tensor: (n_reads, seq_len, n_bases)
+            - mask: (n_reads, seq_len) boolean array indicating seen bases.
+    """
+    encoded = np.asarray(encoded_sequences)
+    if encoded.ndim != 2:
+        raise ValueError(
+            f"encoded_sequences must be 2D with shape (n_reads, seq_len); got {encoded.shape}."
+        )
+
+    base_values = np.array(
+        [MODKIT_EXTRACT_SEQUENCE_BASE_TO_INT[base] for base in bases],
+        dtype=encoded.dtype,
+    )
+
+    if np.issubdtype(encoded.dtype, np.floating):
+        encoded = encoded.copy()
+        encoded[np.isnan(encoded)] = -1
+
+    mask = np.isin(encoded, base_values)
+    one_hot = np.zeros((*encoded.shape, len(base_values)), dtype=dtype)
+
+    for idx, base_value in enumerate(base_values):
+        one_hot[..., idx] = encoded == base_value
+
+    return one_hot, mask
+
+
+def calculate_sequence_cp_decomposition(
+    adata: "ad.AnnData",
+    *,
+    layer: str,
+    rank: int = 5,
+    n_iter_max: int = 100,
+    random_state: int = 0,
+    overwrite: bool = True,
+    embedding_key: str = "X_cp_sequence",
+    components_key: str = "H_cp_sequence",
+    uns_key: str = "cp_sequence",
+    bases: Iterable[str] = ("A", "C", "G", "T"),
+    backend: str = "pytorch",
+    show_progress: bool = False,
+    init: str = "random",
+) -> "ad.AnnData":
+    """Compute CP decomposition on one-hot encoded sequence data with masking.
+
+    Args:
+        adata: AnnData object to update.
+        layer: Layer name containing integer-encoded sequences.
+        rank: CP rank.
+        n_iter_max: Maximum number of iterations for the solver.
+        random_state: Random seed for initialization.
+        overwrite: Whether to recompute if the embedding already exists.
+        embedding_key: Key for embedding in ``adata.obsm``.
+        components_key: Key for position factors in ``adata.varm``.
+        uns_key: Key for metadata stored in ``adata.uns``.
+        bases: Bases to one-hot encode (in order).
+        backend: Tensorly backend to use (``numpy`` or ``pytorch``).
+        show_progress: Whether to display progress during factorization if supported.
+
+    Returns:
+        Updated AnnData object containing the CP decomposition outputs.
+    """
+    if embedding_key in adata.obsm and components_key in adata.varm and not overwrite:
+        logger.info("CP embedding and components already present; skipping recomputation.")
+        return adata
+
+    if backend not in {"numpy", "pytorch"}:
+        raise ValueError(f"Unsupported backend '{backend}'. Use 'numpy' or 'pytorch'.")
+
+    tensorly = require("tensorly", extra="ml-base", purpose="CP decomposition")
+    from tensorly.decomposition import parafac
+
+    tensorly.set_backend(backend)
+
+    if layer not in adata.layers:
+        raise KeyError(f"Layer '{layer}' not found in adata.layers.")
+
+    one_hot, mask = build_sequence_one_hot_and_mask(adata.layers[layer], bases=tuple(bases))
+    mask_tensor = np.repeat(mask[:, :, None], one_hot.shape[2], axis=2)
+
+    device = "numpy"
+    if backend == "pytorch":
+        torch = require("torch", extra="ml-base", purpose="CP decomposition backend")
+        if torch.cuda.is_available():
+            device = torch.device("cuda")
+        elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+            device = torch.device("mps")
+        else:
+            device = torch.device("cpu")
+
+        one_hot = torch.tensor(one_hot, dtype=torch.float32, device=device)
+        mask_tensor = torch.tensor(mask_tensor, dtype=torch.float32, device=device)
+
+    parafac_kwargs = {
+        "rank": rank,
+        "n_iter_max": n_iter_max,
+        "init": init,
+        "mask": mask_tensor,
+        "random_state": random_state,
+    }
+    import inspect
+
+    if "verbose" in inspect.signature(parafac).parameters:
+        parafac_kwargs["verbose"] = show_progress
+
+    cp = parafac(one_hot, **parafac_kwargs)
+
+    if backend == "pytorch":
+        weights = cp.weights.detach().cpu().numpy()
+        read_factors, position_factors, base_factors = [
+            factor.detach().cpu().numpy() for factor in cp.factors
+        ]
+    else:
+        weights = np.asarray(cp.weights)
+        read_factors, position_factors, base_factors = [np.asarray(f) for f in cp.factors]
+
+    adata.obsm[embedding_key] = read_factors
+    adata.varm[components_key] = position_factors
+    adata.uns[uns_key] = {
+        "rank": rank,
+        "n_iter_max": n_iter_max,
+        "random_state": random_state,
+        "layer": layer,
+        "components_key": components_key,
+        "weights": weights,
+        "base_factors": base_factors,
+        "base_labels": list(bases),
+        "backend": backend,
+        "device": str(device),
+    }
+
+    logger.info(
+        "Stored: adata.obsm['%s'], adata.varm['%s'], adata.uns['%s']",
+        embedding_key,
+        components_key,
+        uns_key,
+    )
+    return adata