sclab 0.1.7__py3-none-any.whl → 0.3.4__py3-none-any.whl

sclab/preprocess/_transfer_metadata.py

@@ -0,0 +1,196 @@
+from collections import Counter
+from functools import partial
+from typing import Callable, Literal
+
+import numpy as np
+import pandas as pd
+from anndata import AnnData
+from numpy.typing import NDArray
+from pandas.api.types import is_bool_dtype, is_numeric_dtype
+from scipy.sparse import csr_matrix
+from scipy.special import gamma
+from tqdm.auto import tqdm
+
+
+def transfer_metadata(
+    adata: AnnData,
+    group_key: str,
+    source_group: str,
+    column: str,
+    periodic: bool = False,
+    vmin: float = 0,
+    vmax: float = 1,
+    min_neighs: int = 5,
+    weight_by: Literal["connectivity", "distance", "constant"] = "connectivity",
+):
+    new_values, new_values_err = _propagate_metadata(
+        adata,
+        column=column,
+        periodic=periodic,
+        vmin=vmin,
+        vmax=vmax,
+        min_neighs=min_neighs,
+        weight_by=weight_by,
+        mask=adata.obs[group_key] != source_group,
+    )
+
+    adata.obs[f"transferred_{new_values.name}"] = new_values
+    adata.obs[f"transferred_{new_values_err.name}"] = new_values_err
+
+
+def propagate_metadata(
+    adata: AnnData,
+    column: str,
+    periodic: bool = False,
+    vmin: float = 0,
+    vmax: float = 1,
+    min_neighs: int = 5,
+    weight_by: Literal["connectivity", "distance", "constant"] = "connectivity",
+):
+    new_values, new_values_err = _propagate_metadata(
+        adata,
+        column=column,
+        periodic=periodic,
+        vmin=vmin,
+        vmax=vmax,
+        min_neighs=min_neighs,
+        weight_by=weight_by,
+    )
+
+    mask = adata.obs[column].isna()
+    adata.obs.loc[mask, column] = new_values.loc[mask]
+    adata.obs.loc[mask, new_values_err.name] = new_values_err.loc[mask]
+
+
+def _propagate_metadata(
+    adata: AnnData,
+    column: str,
+    periodic: bool = False,
+    vmin: float = 0,
+    vmax: float = 1,
+    min_neighs: int = 5,
+    weight_by: Literal["connectivity", "distance", "constant"] = "connectivity",
+    mask: np.ndarray | pd.Series | None = None,
+) -> tuple[pd.Series, pd.Series]:
+    D, W = _get_neighbors_and_weights(adata, weight_by=weight_by)
+
+    assign_value_fn: Callable
+    series = adata.obs[column]
+    if isinstance(series.dtype, pd.CategoricalDtype) or is_bool_dtype(series.dtype):
+        assign_value_fn = _assign_categorical
+    elif is_numeric_dtype(series.dtype) and periodic:
+        assign_value_fn = partial(_assign_numerical_periodic, vmin=vmin, vmax=vmax)
+    elif is_numeric_dtype(series.dtype):
+        assign_value_fn = _assign_numerical
+    else:
+        raise ValueError(f"Unsupported dtype {series.dtype} for column {column}")
+
+    if isinstance(series.dtype, pd.CategoricalDtype) or is_bool_dtype(series.dtype):
+        column_err = f"{column}_proportion"
+    else:
+        column_err = f"{column}_error"
+
+    meta_values: pd.Series = series.copy()
+    if mask is not None:
+        meta_values[mask] = pd.NA
+
+    new_values = pd.Series(index=series.index, dtype=series.dtype, name=column)
+    new_values_err = pd.Series(index=series.index, dtype=float, name=column_err)
+
+    for i, (d, w) in tqdm(enumerate(zip(D, W)), total=D.shape[0]):
+        if not pd.isna(meta_values.iloc[i]):
+            continue
+
+        d = d.tocoo()
+        w = w.toarray().ravel()
+        neighs = d.coords[1]
+
+        values: pd.Series = meta_values.iloc[neighs]
+        msk = pd.notna(values)
+        if msk.sum() < min_neighs:
+            continue
+
+        values = values.loc[msk]
+        weights = w[neighs][msk]
+
+        if np.allclose(weights, 0):
+            continue
+
+        assigned_value, assigned_value_err = assign_value_fn(values, weights)
+        new_values.iloc[i] = assigned_value
+        new_values_err.iloc[i] = assigned_value_err
+
+    new_values = pd.concat([new_values, meta_values], axis=1).bfill(axis=1).iloc[:, 0]
+
+    return new_values, new_values_err
+
+
+def _get_neighbors_and_weights(
+    adata: AnnData,
+    weight_by: Literal["connectivity", "distance", "constant"] = "connectivity",
+):
+    D: csr_matrix = adata.obsp["distances"].copy()
+    C: csr_matrix = adata.obsp["connectivities"].copy()
+    D = D.tocsr()
+    W: csr_matrix
+
+    match weight_by:
+        case "connectivity":
+            W = C.tocsr().copy()
+        case "distance":
+            W = D.tocsr().copy()
+            W.data = 1.0 / W.data
+        case "constant":
+            W = D.tocsr().copy()
+            W.data[:] = 1.0
+        case _:
+            raise ValueError(f"Unsupported weight_by {weight_by}")
+
+    return D, W
+
+
+def _assign_categorical(values: pd.Series, weights: NDArray):
+    # weighted majority and proportion of votes
+    tally = Counter()
+    for v, w in zip(values, weights):
+        tally[v] += w
+
+    winner, shares = tally.most_common()[0]
+    return winner, shares / weights.sum()
+
+
+def _assign_numerical(values: pd.Series, weights: NDArray):
+    # weighted mean and standard error
+    sum_w: float = weights.sum()
+    sum2_w: float = weights.sum() ** 2
+    sum_w2: float = (weights**2).sum()
+    n_eff: float = sum2_w / sum_w2
+
+    mean_x: float = (values * weights).sum() / sum_w
+    var_x: float = ((values - mean_x) ** 2 * weights).sum() * sum_w / (sum2_w - sum_w2)
+    err_x: float = np.sqrt(var_x / n_eff)
+
+    return mean_x, err_x
+
+
+def _assign_numerical_periodic(
+    values: pd.Series, weights: NDArray, vmin: float, vmax: float
+):
+    vspan = vmax - vmin
+
+    values = values - vmin
+    offset = np.median(values)
+    values = values - offset + vspan / 2
+    values = values % vspan
+    assigned_value, assigned_value_err = _assign_numerical(values, weights)
+    assigned_value = assigned_value + offset - vspan / 2
+    assigned_value = assigned_value % vspan
+    assigned_value = assigned_value + vmin
+
+    return assigned_value, assigned_value_err
+
+
+def _c4(n: float):
+    # correct for bias
+    nm1 = n - 1
+    return np.sqrt(2 / nm1) * gamma(n / 2) / gamma(nm1 / 2)

sclab/preprocess/_transform.py

@@ -0,0 +1,82 @@
+from typing import Optional
+
+from anndata import AnnData
+from numpy import ndarray
+from scipy.sparse import csr_matrix
+
+from ._utils import get_neighbors_adjacency_matrix
+
+
+def pool_neighbors(
+    adata: AnnData,
+    *,
+    layer: Optional[str] = None,
+    n_neighbors: Optional[int] = None,
+    neighbors_key: Optional[str] = None,
+    weighted: bool = False,
+    directed: bool = True,
+    key_added: Optional[str] = None,
+    copy: bool = False,
+) -> csr_matrix | ndarray | None:
+    """
+    Given an adjacency matrix, pool cell features using a weighted sum of feature counts
+    from neighboring cells. The weights are the normalized connectivities from the
+    adjacency matrix.
+
+    Parameters
+    ----------
+    adata : AnnData
+        Annotated data matrix.
+    layer : str, optional
+        Layer in AnnData object to use for pooling. Defaults to None.
+    n_neighbors : int, optional
+        Number of neighbors to consider. Defaults to None.
+    neighbors_key : str, optional
+        Key in AnnData object to use for neighbors. Defaults to None.
+    weighted : bool, optional
+        Whether to weight neighbors by their connectivities in the adjacency matrix.
+        Defaults to False.
+    directed : bool, optional
+        Whether to use directed or undirected neighbors. Defaults to True.
+    key_added : str, optional
+        Key to use in AnnData object for the pooled features. Defaults to None.
+    copy : bool, optional
+        Whether to return a copy of the pooled features instead of modifying the
+        original AnnData object. Defaults to False.
+
+    Returns
+    -------
+    csr_matrix | ndarray | None
+        The pooled features if copy is True, otherwise None.
+    """
+    if layer is None or layer == "X":
+        X = adata.X
+    else:
+        X = adata.layers[layer]
+
+    adjacency = get_neighbors_adjacency_matrix(
+        adata,
+        key=neighbors_key,
+        n_neighbors=n_neighbors,
+        weighted=weighted,
+        directed=directed,
+    )
+
+    W = adjacency.tolil()
+    W.setdiag(1)
+
+    W = W / W.sum(axis=1)
+
+    pooled = W.dot(X)
+
+    if copy:
+        return pooled
+
+    if key_added is not None:
+        adata.layers[key_added] = pooled
+        return
+
+    if layer is None or layer == "X":
+        adata.X = pooled
+    else:
+        adata.layers[layer] = pooled

sclab/preprocess/_utils.py

@@ -0,0 +1,96 @@
+from typing import Literal, Optional
+
+import numpy as np
+from anndata import AnnData
+from scanpy import Neighbors
+from scipy.sparse import coo_matrix, csr_matrix
+
+
+def get_neighbors_adjacency_matrix(
+    adata: AnnData,
+    *,
+    key: Optional[str] = "neighbors",
+    n_neighbors: Optional[int] = None,
+    weighted: bool = False,
+    directed: bool = True,
+) -> csr_matrix:
+    # get the current neighbors
+    neigh = Neighbors(adata, neighbors_key=key)
+    params = adata.uns[key]["params"]
+
+    if n_neighbors is None:
+        n_neighbors = neigh.n_neighbors
+
+    if n_neighbors < neigh.n_neighbors and not weighted:
+        distances = _filter_knn_matrix(
+            neigh.distances, n_neighbors=n_neighbors, mode="distances"
+        )
+
+    elif n_neighbors != neigh.n_neighbors:
+        neigh.compute_neighbors(**{**params, "n_neighbors": n_neighbors})
+        distances = neigh.distances
+
+    else:
+        distances = neigh.distances
+
+    adjacency = distances.copy()
+    adjacency.data = np.ones_like(adjacency.data)
+
+    if not directed:
+        # make the adjacency matrix symmetric
+        adjacency = _symmetrize_sparse_matrix(adjacency)
+
+    if weighted:
+        # use the connectivities to assign weights
+        adjacency = adjacency.multiply(neigh.connectivities)
+
+    return adjacency
+
+
+def _filter_knn_matrix(
+    matrix: csr_matrix, *, n_neighbors: int, mode: Literal["distances", "weights"]
+) -> csr_matrix:
+    assert mode in ["distances", "weights"]
+    nrows, _ = matrix.shape
+
+    # Initialize arrays for new sparse matrix with pre-allocated size
+    indptr = np.arange(0, (n_neighbors - 1) * (nrows + 1), n_neighbors - 1)
+    data = np.zeros(nrows * (n_neighbors - 1), dtype=float)
+    indices = np.zeros(nrows * (n_neighbors - 1), dtype=int)
+
+    # Process each row to keep top n_neighbors-1 connections
+    for i in range(nrows):
+        start, end = matrix.indptr[i : i + 2]
+        idxs = matrix.indices[start:end]
+        vals = matrix.data[start:end]
+
+        # Sort by values and keep top n_neighbors-1
+        if mode == "weights":
+            # Sort in descending order (keep largest weights)
+            o = np.argsort(-vals)[: n_neighbors - 1]
+        else:
+            # Sort in ascending order (keep smallest distances)
+            o = np.argsort(vals)[: n_neighbors - 1]
+
+        # Maintain original order within top neighbors
+        oo = np.argsort(idxs[o])
+        start, end = indptr[i : i + 2]
+        indices[start:end] = idxs[o][oo]
+        data[start:end] = vals[o][oo]
+
+    return csr_matrix((data, indices, indptr))
+
+
+def _symmetrize_sparse_matrix(matrix: csr_matrix) -> csr_matrix:
+    A = matrix.tocoo()
+
+    # Make matrix symmetric by duplicating entries in both directions
+    coords = np.array([[*A.row, *A.col], [*A.col, *A.row]])
+    data = np.array([*A.data, *A.data])
+
+    # Remove duplicate entries that might occur in symmetrization
+    idxs = np.unique(coords, axis=1, return_index=True)[1]
+    coords, data = coords[:, idxs], data[idxs]
+    A = coo_matrix((data, coords), shape=matrix.shape)
+
+    return A.tocsr()

sclab/scanpy/_compat.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass, field
+from functools import cache, partial
+from importlib.util import find_spec
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from packaging.version import Version
+
+if TYPE_CHECKING:
+    from importlib.metadata import PackageMetadata
+
+
+if TYPE_CHECKING:
+    # type checkers are confused and can only see …core.Array
+    from dask.array.core import Array as DaskArray
+elif find_spec("dask"):
+    from dask.array import Array as DaskArray
+else:
+
+    class DaskArray:
+        pass
+
+
+if find_spec("zappy") or TYPE_CHECKING:
+    from zappy.base import ZappyArray
+else:
+
+    class ZappyArray:
+        pass
+
+
+__all__ = [
+    "DaskArray",
+    "ZappyArray",
+    "fullname",
+    "pkg_metadata",
+    "pkg_version",
+]
+
+
+def fullname(typ: type) -> str:
+    module = typ.__module__
+    name = typ.__qualname__
+    if module == "builtins" or module is None:
+        return name
+    return f"{module}.{name}"
+
+
+if sys.version_info >= (3, 11):
+    from contextlib import chdir
+else:
+    import os
+    from contextlib import AbstractContextManager
+
+    @dataclass
+    class chdir(AbstractContextManager):
+        path: Path
+        _old_cwd: list[Path] = field(default_factory=list)
+
+        def __enter__(self) -> None:
+            self._old_cwd.append(Path.cwd())
+            os.chdir(self.path)
+
+        def __exit__(self, *_excinfo) -> None:
+            os.chdir(self._old_cwd.pop())
+
+
+def pkg_metadata(package: str) -> PackageMetadata:
+    from importlib.metadata import metadata
+
+    return metadata(package)
+
+
+@cache
+def pkg_version(package: str) -> Version:
+    from importlib.metadata import version
+
+    return Version(version(package))
+
+
+if find_spec("legacy_api_wrap") or TYPE_CHECKING:
+    from legacy_api_wrap import legacy_api  # noqa: TID251
+
+    old_positionals = partial(legacy_api, category=FutureWarning)
+else:
+    # legacy_api_wrap is currently a hard dependency,
+    # but this code makes it possible to run scanpy without it.
+    def old_positionals(*old_positionals: str):
+        return lambda func: func