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

smftools/schema/anndata_schema_v1.yaml

@@ -60,6 +60,20 @@ stages:
         notes: "Mapping quality score."
         requires: []
         optional_inputs: []
+      reference_start:
+        dtype: "float"
+        created_by: "smftools.informatics.h5ad_functions.add_read_length_and_mapping_qc"
+        modified_by: []
+        notes: "0-based reference start position for the alignment."
+        requires: []
+        optional_inputs: []
+      reference_end:
+        dtype: "float"
+        created_by: "smftools.informatics.h5ad_functions.add_read_length_and_mapping_qc"
+        modified_by: []
+        notes: "0-based reference end position (exclusive) for the alignment."
+        requires: []
+        optional_inputs: []
       read_length_to_reference_length_ratio:
         dtype: "float"
         created_by: "smftools.informatics.h5ad_functions.add_read_length_and_mapping_qc"
@@ -179,7 +193,7 @@ stages:
     obs:
       leiden:
         dtype: "category"
-        created_by: "smftools.tools.calculate_umap"
+        created_by: "smftools.tools.calculate_leiden"
         modified_by: []
         notes: "Leiden cluster assignments."
         requires: [["obsm.X_umap"]]

smftools/tools/__init__.py

@@ -3,6 +3,11 @@ from __future__ import annotations
 from importlib import import_module
 
 _LAZY_ATTRS = {
+    "calculate_leiden": "smftools.tools.calculate_leiden",
+    "calculate_nmf": "smftools.tools.calculate_nmf",
+    "calculate_sequence_cp_decomposition": "smftools.tools.tensor_factorization",
+    "calculate_pca": "smftools.tools.calculate_pca",
+    "calculate_knn": "smftools.tools.calculate_knn",
     "calculate_umap": "smftools.tools.calculate_umap",
     "cluster_adata_on_methylation": "smftools.tools.cluster_adata_on_methylation",
     "combine_layers": "smftools.tools.general_tools",
@@ -11,6 +16,11 @@ _LAZY_ATTRS = {
     "calculate_relative_risk_on_activity": "smftools.tools.position_stats",
     "compute_positionwise_statistics": "smftools.tools.position_stats",
     "calculate_row_entropy": "smftools.tools.read_stats",
+    "align_sequences_with_mismatches": "smftools.tools.sequence_alignment",
+    "rolling_window_nn_distance": "smftools.tools.rolling_nn_distance",
+    "annotate_zero_hamming_segments": "smftools.tools.rolling_nn_distance",
+    "assign_per_read_segments_layer": "smftools.tools.rolling_nn_distance",
+    "select_top_segments_per_read": "smftools.tools.rolling_nn_distance",
     "subset_adata": "smftools.tools.subset_adata",
 }
 

smftools/tools/calculate_knn.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+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_knn(
+    adata: "ad.AnnData",
+    obsm: str = "X_pca",
+    knn_neighbors: int = 100,
+    overwrite: bool = True,
+    threads: int = 8,
+    random_state: int | None = 0,
+    symmetrize: bool = True,
+) -> "ad.AnnData":
+    """Compute a KNN distance graph on an embedding in `adata.obsm[obsm]`.
+
+    Stores:
+      - adata.obsp[f"knn_distances_{obsm}"] : CSR sparse matrix of distances
+      - adata.uns[f"knn_distances_{obsm}"]["params"] : metadata
+
+    Args:
+        adata: AnnData object to update.
+        obsm: Key in `adata.obsm` to use as the embedding.
+        knn_neighbors: Target number of neighbors (will be clipped to n_obs-1).
+        overwrite: If False and graph exists, do nothing.
+        threads: Parallel jobs for pynndescent.
+        random_state: Seed for pynndescent.
+        symmetrize: If True, make distance graph symmetric via min(A, A.T).
+
+    Returns:
+        Updated AnnData.
+    """
+    import numpy as np
+    import scipy.sparse as sp
+
+    if obsm not in adata.obsm:
+        raise KeyError(f"`{obsm}` not found in adata.obsm. Available: {list(adata.obsm.keys())}")
+
+    out_key = f"knn_distances_{obsm}"
+    if not overwrite and out_key in adata.obsp:
+        logger.info("KNN graph %r already exists and overwrite=False; skipping.", out_key)
+        return adata
+
+    data = adata.obsm[obsm]
+
+    if sp.issparse(data):
+        # Convert to float32 for pynndescent/numba friendliness if needed
+        data = data.astype(np.float32)
+        logger.info(
+            "Sparse embedding detected (%s). Proceeding without NaN check.", type(data).__name__
+        )
+    else:
+        data = np.asarray(data)
+        if np.isnan(data).any():
+            logger.warning("NaNs detected in %s; filling NaNs with 0.5 before KNN.", obsm)
+            data = np.nan_to_num(data, nan=0.5)
+        data = data.astype(np.float32, copy=False)
+
+    pynndescent = require("pynndescent", extra="umap", purpose="KNN graph computation")
+
+    n_obs = data.shape[0]
+    if n_obs < 2:
+        raise ValueError(f"Need at least 2 observations for KNN; got n_obs={n_obs}")
+
+    n_neighbors = min(int(knn_neighbors), n_obs - 1)
+    if n_neighbors < 1:
+        raise ValueError(f"Computed n_neighbors={n_neighbors}; check knn_neighbors and n_obs.")
+
+    logger.info(
+        "Running pynndescent KNN (obsm=%s, n_neighbors=%d, metric=euclidean, n_jobs=%d)",
+        obsm,
+        n_neighbors,
+        threads,
+    )
+
+    nn_index = pynndescent.NNDescent(
+        data,
+        n_neighbors=n_neighbors,
+        metric="euclidean",
+        random_state=random_state,
+        n_jobs=threads,
+    )
+    knn_indices, knn_dists = nn_index.neighbor_graph  # shapes: (n_obs, n_neighbors)
+
+    rows = np.repeat(np.arange(n_obs, dtype=np.int64), n_neighbors)
+    cols = knn_indices.reshape(-1).astype(np.int64, copy=False)
+    vals = knn_dists.reshape(-1).astype(np.float32, copy=False)
+
+    distances = sp.coo_matrix((vals, (rows, cols)), shape=(n_obs, n_obs)).tocsr()
+
+    # Optional: ensure diagonal is 0 and (optionally) symmetrize
+    distances.setdiag(0.0)
+    distances.eliminate_zeros()
+
+    if symmetrize:
+        # Keep the smaller directed distance for each undirected edge
+        distances = distances.minimum(distances.T)
+
+    adata.obsp[out_key] = distances
+    adata.uns[out_key] = {
+        "params": {
+            "obsm": obsm,
+            "n_neighbors_requested": int(knn_neighbors),
+            "n_neighbors_used": int(n_neighbors),
+            "method": "pynndescent",
+            "metric": "euclidean",
+            "random_state": random_state,
+            "n_jobs": int(threads),
+            "symmetrize": bool(symmetrize),
+        }
+    }
+
+    return adata

smftools/tools/calculate_leiden.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pandas as pd
+
+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_leiden(
+    adata: "ad.AnnData",
+    *,
+    resolution: float = 0.1,
+    key_added: str = "leiden",
+    connectivities_key: str = "connectivities",
+) -> "ad.AnnData":
+    """Compute Leiden clusters from a connectivity graph.
+
+    Args:
+        adata: AnnData object with ``obsp[connectivities_key]`` set.
+        resolution: Resolution parameter for Leiden clustering.
+        key_added: Column name to store cluster assignments in ``adata.obs``.
+        connectivities_key: Key in ``adata.obsp`` containing a sparse adjacency matrix.
+
+    Returns:
+        Updated AnnData object with Leiden labels in ``adata.obs``.
+    """
+    if connectivities_key not in adata.obsp:
+        raise KeyError(f"Missing connectivities '{connectivities_key}' in adata.obsp.")
+
+    igraph = require("igraph", extra="cluster", purpose="Leiden clustering")
+    leidenalg = require("leidenalg", extra="cluster", purpose="Leiden clustering")
+
+    connectivities = adata.obsp[connectivities_key]
+    coo = connectivities.tocoo()
+    edges = list(zip(coo.row.tolist(), coo.col.tolist()))
+    graph = igraph.Graph(n=connectivities.shape[0], edges=edges, directed=False)
+    graph.es["weight"] = coo.data.tolist()
+
+    partition = leidenalg.find_partition(
+        graph,
+        leidenalg.RBConfigurationVertexPartition,
+        weights=graph.es["weight"],
+        resolution_parameter=resolution,
+    )
+
+    labels = np.array(partition.membership, dtype=str)
+    adata.obs[key_added] = pd.Categorical(labels)
+    logger.info("Stored Leiden clusters in adata.obs['%s'].", key_added)
+    return adata

smftools/tools/calculate_nmf.py

@@ -0,0 +1,130 @@
+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
+    import numpy as np
+
+logger = get_logger(__name__)
+
+
+def calculate_nmf(
+    adata: "ad.AnnData",
+    layer: str | None = "nan_half",
+    var_mask: "np.ndarray | Sequence[bool] | 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",
+    suffix: str | None = None,
+) -> "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_mask: Optional boolean mask 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
+
+    if suffix:
+        embedding_key = f"{embedding_key}_{suffix}"
+        components_key = f"{components_key}_{suffix}"
+        uns_key = f"{uns_key}_{suffix}"
+
+    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_mask is not None:
+        subset_mask = np.asarray(var_mask, dtype=bool)
+        if subset_mask.ndim != 1 or subset_mask.shape[0] != adata.n_vars:
+            raise ValueError(
+                "var_mask must be a 1D boolean array with length matching adata.n_vars."
+            )
+        adata_subset = adata[:, subset_mask].copy()
+        logger.info(
+            "Subsetting adata: retained %s features based on filters %s",
+            adata_subset.shape[1],
+            "var_mask",
+        )
+    else:
+        adata_subset = adata.copy()
+        logger.info("No var_mask 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_mask_provided": var_mask is not None,
+        "components_key": components_key,
+    }
+
+    logger.info(
+        "Stored: adata.obsm['%s'] and adata.varm['%s']",
+        embedding_key,
+        components_key,
+    )
+    return adata

smftools/tools/calculate_pca.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
+from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
+
+if TYPE_CHECKING:
+    import anndata as ad
+    import numpy as np
+
+logger = get_logger(__name__)
+
+
+def calculate_pca(
+    adata: "ad.AnnData",
+    layer: str | None = "nan_half",
+    var_mask: "np.ndarray | Sequence[bool] | None" = None,
+    n_pcs: int = 15,
+    overwrite: bool = True,
+    output_suffix: str | None = None,
+    fill_nan: float | None = 0.5,
+) -> "ad.AnnData":
+    """Compute PCA and store scores in `.obsm` and loadings in `.varm`."""
+
+    import numpy as np
+    import scipy.sparse as sp
+
+    obsm_output = f"X_pca_{output_suffix}" if output_suffix else "X_pca"
+    varm_output = f"PCs_{output_suffix}" if output_suffix else "PCs"
+
+    if not overwrite and obsm_output in adata.obsm and varm_output in adata.varm:
+        logger.info(
+            "PCA outputs already exist and overwrite=False; skipping (%s, %s).",
+            obsm_output,
+            varm_output,
+        )
+        return adata
+
+    # --- Build feature subset mask (over vars) ---
+    if var_mask is not None:
+        subset_mask = np.asarray(var_mask, dtype=bool)
+        if subset_mask.ndim != 1 or subset_mask.shape[0] != adata.n_vars:
+            raise ValueError(
+                "var_mask must be a 1D boolean array with length matching adata.n_vars."
+            )
+        n_vars_used = int(subset_mask.sum())
+        if n_vars_used == 0:
+            raise ValueError("var_mask retained 0 features.")
+        logger.info(
+            "Subsetting vars: retained %d / %d features from var_mask",
+            n_vars_used,
+            adata.n_vars,
+        )
+    else:
+        subset_mask = slice(None)
+        n_vars_used = adata.n_vars
+        logger.info("No var_mask provided; using all %d features.", adata.n_vars)
+
+    # --- Pull matrix view ---
+    if layer is None:
+        matrix = adata.X
+        layer_used = None
+    else:
+        if layer not in adata.layers:
+            raise KeyError(
+                f"Layer {layer!r} not found in adata.layers. Available: {list(adata.layers.keys())}"
+            )
+        matrix = adata.layers[layer]
+        layer_used = layer
+
+    matrix = matrix[:, subset_mask]  # slice view (sparse OK)
+
+    n_obs = matrix.shape[0]
+    if n_obs < 2:
+        raise ValueError(f"PCA requires at least 2 observations; got n_obs={n_obs}")
+    if n_vars_used < 1:
+        raise ValueError("PCA requires at least 1 feature.")
+
+    n_pcs_requested = int(n_pcs)
+    n_pcs_used = min(n_pcs_requested, n_obs, n_vars_used)
+    if n_pcs_used < 1:
+        raise ValueError(f"n_pcs_used became {n_pcs_used}; check inputs.")
+
+    # --- NaN handling (dense only; sparse usually won’t store NaNs) ---
+    if not sp.issparse(matrix):
+        X = np.asarray(matrix, dtype=np.float32)
+        if fill_nan is not None and np.isnan(X).any():
+            logger.warning("NaNs detected; filling NaNs with %s before PCA.", fill_nan)
+            X = np.nan_to_num(X, nan=float(fill_nan))
+    else:
+        X = matrix  # keep sparse
+
+    # --- PCA ---
+    # Prefer sklearn's randomized PCA for speed on big matrices.
+    used_sklearn = False
+    try:
+        sklearn = require("sklearn", extra="ml", purpose="PCA computation")
+        from sklearn.decomposition import PCA, TruncatedSVD
+
+        if sp.issparse(X):
+            # TruncatedSVD works on sparse without centering; good approximation.
+            # If you *need* centered PCA on sparse, you'd need different machinery.
+            logger.info("Running TruncatedSVD (sparse) with n_components=%d", n_pcs_used)
+            model = TruncatedSVD(n_components=n_pcs_used, random_state=0)
+            scores = model.fit_transform(X)  # (n_obs, n_pcs)
+            loadings = model.components_.T  # (n_vars_used, n_pcs)
+            mean = None
+            explained_variance_ratio = getattr(model, "explained_variance_ratio_", None)
+        else:
+            logger.info(
+                "Running sklearn PCA with n_components=%d (svd_solver=randomized)", n_pcs_used
+            )
+            model = PCA(n_components=n_pcs_used, svd_solver="randomized", random_state=0)
+            scores = model.fit_transform(X)  # (n_obs, n_pcs)
+            loadings = model.components_.T  # (n_vars_used, n_pcs)
+            mean = model.mean_
+            explained_variance_ratio = model.explained_variance_ratio_
+
+        used_sklearn = True
+
+    except Exception as e:
+        # Fallback to your manual SVD (dense only)
+        if sp.issparse(X):
+            raise RuntimeError(
+                "Sparse input PCA fallback is not implemented without sklearn. "
+                "Install scikit-learn (extra 'ml') or densify upstream."
+            ) from e
+
+        import scipy.linalg as spla
+
+        logger.warning(
+            "sklearn PCA unavailable; falling back to full SVD (can be slow). Reason: %s", e
+        )
+        Xd = np.asarray(X, dtype=np.float64)
+        mean = Xd.mean(axis=0)
+        centered = Xd - mean
+        u, s, vt = spla.svd(centered, full_matrices=False)
+        u = u[:, :n_pcs_used]
+        s = s[:n_pcs_used]
+        vt = vt[:n_pcs_used]
+        scores = u * s
+        loadings = vt.T
+        explained_variance_ratio = None
+
+    # --- Store scores (obsm) ---
+    adata.obsm[obsm_output] = scores
+
+    # --- Store loadings (varm) with original var dimension ---
+    pc_matrix = np.zeros((adata.n_vars, n_pcs_used), dtype=np.float32)
+    if isinstance(subset_mask, slice):
+        pc_matrix[:, :] = loadings
+    else:
+        pc_matrix[subset_mask, :] = loadings.astype(np.float32, copy=False)
+
+    adata.varm[varm_output] = pc_matrix
+
+    # --- Metadata ---
+    adata.uns[obsm_output] = {
+        "params": {
+            "layer": layer_used,
+            "var_mask_provided": var_mask is not None,
+            "n_pcs_requested": n_pcs_requested,
+            "n_pcs_used": int(n_pcs_used),
+            "used_sklearn": used_sklearn,
+            "fill_nan": fill_nan,
+            "note_sparse": bool(sp.issparse(matrix)),
+        },
+        "explained_variance_ratio": explained_variance_ratio,
+        "mean": mean.tolist() if (mean is not None and isinstance(mean, np.ndarray)) else None,
+    }
+
+    logger.info(
+        "Stored PCA: adata.obsm[%s] (%s) and adata.varm[%s] (%s)",
+        obsm_output,
+        scores.shape,
+        varm_output,
+        pc_matrix.shape,
+    )
+    return adata