sclab 0.2.5__py3-none-any.whl → 0.3.0__py3-none-any.whl

sclab/preprocess/_harmony_integrate.py

@@ -0,0 +1,50 @@
+"""Use harmony to integrate cells from different experiments.
+
+Note: code adapted from scanpy to use a custom version of harmonypy
+
+Harmony:
+Korsunsky, I., Millard, N., Fan, J. et al. Fast, sensitive and accurate integration of single-cell data with Harmony.
+Nat Methods 16, 1289-1296 (2019). https://doi.org/10.1038/s41592-019-0619-0
+
+Scanpy:
+Wolf, F., Angerer, P. & Theis, F. SCANPY: large-scale single-cell gene expression data analysis.
+Genome Biol 19, 15 (2018). https://doi.org/10.1186/s13059-017-1382-0
+
+Scverse:
+Virshup, I., Bredikhin, D., Heumos, L. et al. The scverse project provides a computational ecosystem for single-cell omics data analysis.
+Nat Biotechnol 41, 604-606 (2023). https://doi.org/10.1038/s41587-023-01733-8
+"""
+
+from collections.abc import Sequence
+
+from anndata import AnnData
+import numpy as np
+
+from ._harmony import run_harmony
+
+
+def harmony_integrate(
+    adata: AnnData,
+    key: str | Sequence[str],
+    *,
+    basis: str = "X_pca",
+    adjusted_basis: str = "X_pca_harmony",
+    reference_batch: str | list[str] | None = None,
+    **kwargs,
+):
+    """Use harmonypy :cite:p:`Korsunsky2019` to integrate different experiments."""
+
+    if isinstance(reference_batch, str):
+        reference_batch = [reference_batch]
+
+    if reference_batch is not None:
+        reference_values = np.zeros(adata.n_obs, dtype=bool)
+        for batch in reference_batch:
+            reference_values |= adata.obs[key].values == batch
+        kwargs["reference_values"] = reference_values
+
+    X = adata.obsm[basis].astype(np.float64)
+
+    harmony_out = run_harmony(X, adata.obs, key, **kwargs)
+
+    adata.obsm[adjusted_basis] = harmony_out.Z_corr.T

sclab/preprocess/_normalize_weighted.py

@@ -0,0 +1,61 @@
+import warnings
+
+import numpy as np
+from anndata import AnnData, ImplicitModificationWarning
+from scipy.sparse import csr_matrix, issparse
+
+
+def normalize_weighted(
+    adata: AnnData,
+    target_scale: float | None = None,
+    batch_key: str | None = None,
+) -> None:
+    if batch_key is not None:
+        for _, idx in adata.obs.groupby(batch_key, observed=True).groups.items():
+            with warnings.catch_warnings():
+                warnings.filterwarnings(
+                    "ignore",
+                    category=ImplicitModificationWarning,
+                    message="Modifying `X` on a view results in data being overridden",
+                )
+                normalize_weighted(adata[idx], target_scale, None)
+
+        return
+
+    X: csr_matrix
+    Y: csr_matrix
+    Z: csr_matrix
+
+    X = adata.X
+    assert issparse(X)
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", category=RuntimeWarning, message="divide by zero"
+        )
+        Y = X.multiply(1 / X.sum(axis=0))
+    Y = Y.tocsr()
+    Y.eliminate_zeros()
+    Y.data = -Y.data * np.log(Y.data)
+    entropy = Y.sum(axis=0)
+
+    Z = X.multiply(entropy)
+    Z = Z.tocsr()
+    Z.eliminate_zeros()
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", category=RuntimeWarning, message="divide by zero"
+        )
+        scale = Z.sum(axis=1)
+        Z = Z.multiply(1 / scale)
+    Z = Z.tocsr()
+
+    if target_scale is None:
+        target_scale = np.median(scale.A1[scale.A1 > 0])
+
+    Z = Z * target_scale
+
+    adata.X = Z
+
+    return

sclab/preprocess/_subset.py

@@ -0,0 +1,208 @@
+from typing import Sequence
+
+import numpy as np
+import pandas as pd
+from anndata import AnnData
+
+
+def subset_obs(
+    adata: AnnData,
+    subset: pd.Index | Sequence[str | int | bool] | str,
+) -> None:
+    """Subset observations (rows) in an AnnData object.
+
+    This function modifies the AnnData object in-place by selecting a subset of observations
+    based on the provided subset parameter. The subsetting can be done using observation
+    names, integer indices, a boolean mask, a query string, or a pandas Index.
+
+    Parameters
+    ----------
+    adata : AnnData
+        The annotated data matrix to subset. Will be modified in-place.
+    subset : pd.Index | Sequence[str | int | bool] | str
+        The subset specification. Can be one of:
+        * A pandas Index containing observation names
+        * A sequence of observation names (strings)
+        * A sequence of integer indices
+        * A boolean mask of length `adata.n_obs`
+        * A query string to match observations by their metadata columns
+
+    Examples
+    --------
+    >>> # Create an example AnnData object
+    >>> import anndata
+    >>> import pandas as pd
+    >>> import numpy as np
+    >>>
+    >>> obs = pd.DataFrame(
+    ...     index=['A', 'B', 'C'],
+    ...     data={'cell_type': ['type1', 'type2', 'type2']})
+    >>> adata_ = anndata.AnnData(obs=obs)
+    >>>
+    >>> # Subset using pandas Index
+    >>> adata = adata_.copy()
+    >>> subset_obs(adata, pd.Index(['B', 'C']))
+    >>> adata.obs_names.tolist()
+    ['B', 'C']
+    >>>
+    >>> # Subset using observation names
+    >>> adata = adata_.copy()
+    >>> subset_obs(adata, ['A', 'B'])
+    >>> adata.obs_names.tolist()
+    ['A', 'B']
+    >>>
+    >>> # Subset using integer indices
+    >>> adata = adata_.copy()
+    >>> subset_obs(adata, [0, 1])
+    >>> adata.obs_names.tolist()
+    ['A', 'B']
+    >>>
+    >>> # Subset using boolean mask
+    >>> adata = adata_.copy()
+    >>> subset_obs(adata, [True, False, True])
+    >>> adata.obs_names.tolist()
+    ['A', 'C']
+    >>>
+    >>> # Subset using query string
+    >>> adata = adata_.copy()
+    >>> subset_obs(adata, 'cell_type == "type2"')
+    >>> adata.obs_names.tolist()
+    ['B', 'C']
+
+    Notes
+    -----
+    - The function modifies the AnnData object in-place
+    - When using a boolean mask, its length must match the number of observations
+    - When using integer indices, they must be valid indices for the observations
+    - Invalid observation names or indices will raise KeyError or IndexError respectively
+    - The order of observations in the output will match the order in the subset parameter
+    """
+    if isinstance(subset, str):
+        subset = adata.obs.query(subset).index
+
+    if not isinstance(subset, pd.Index):
+        subset = np.asarray(subset)
+
+    # Handle boolean mask
+    if subset.dtype.kind == "b":
+        if len(subset) != adata.n_obs:
+            raise IndexError(
+                f"Boolean mask length ({len(subset)}) does not match number of "
+                f"observations ({adata.n_obs})"
+            )
+        subset = adata.obs_names[subset]
+
+    # Handle integer indices
+    elif subset.dtype.kind in "iu":
+        if np.any(subset < 0) or np.any(subset >= adata.n_obs):
+            raise IndexError(f"Integer indices must be between 0 and {adata.n_obs - 1}")
+        subset = adata.obs_names[subset]
+
+    if adata.n_obs == subset.size and (subset == adata.obs_names).all():
+        # No need to subset, avoid making a copy. Useful for large AnnData objects
+        return
+
+    adata._inplace_subset_obs(subset)
+
+
+def subset_var(
+    adata: AnnData,
+    subset: pd.Index | Sequence[str | int | bool] | str,
+) -> None:
+    """Subset variables (columns) in an AnnData object.
+
+    This function modifies the AnnData object in-place by selecting a subset of variables
+    based on the provided subset parameter. The subsetting can be done using variable
+    names, integer indices, a boolean mask, a query string, or a pandas Index.
+
+    Parameters
+    ----------
+    adata : AnnData
+        The annotated data matrix to subset. Will be modified in-place.
+    subset : pd.Index | Sequence[str | int | bool] | str
+        The subset specification. Can be one of:
+        * A pandas Index containing variable names
+        * A sequence of variable names (strings)
+        * A sequence of integer indices
+        * A boolean mask of length `adata.n_vars`
+        * A query string to match variables by their metadata columns
+
+    Examples
+    --------
+    >>> # Create an example AnnData object
+    >>> import anndata
+    >>> import pandas as pd
+    >>> import numpy as np
+    >>>
+    >>> var = pd.DataFrame(
+    ...     index=['gene1', 'gene2', 'gene3'],
+    ...     data={'gene_type': ['type1', 'type2', 'type1']})
+    >>> adata_ = anndata.AnnData(var=var)
+    >>>
+    >>> # Subset using pandas Index
+    >>> adata = adata_.copy()
+    >>> subset_var(adata, pd.Index(['gene2', 'gene3']))
+    >>> adata.var_names.tolist()
+    ['gene2', 'gene3']
+    >>>
+    >>> # Subset using variable names
+    >>> adata = adata_.copy()
+    >>> subset_var(adata, ['gene1', 'gene2'])
+    >>> adata.var_names.tolist()
+    ['gene1', 'gene2']
+    >>>
+    >>> # Subset using integer indices
+    >>> adata = adata_.copy()
+    >>> subset_var(adata, [0, 1])
+    >>> adata.var_names.tolist()
+    ['gene1', 'gene2']
+    >>>
+    >>> # Subset using boolean mask
+    >>> adata = adata_.copy()
+    >>> subset_var(adata, [True, False, True])
+    >>> adata.var_names.tolist()
+    ['gene1', 'gene3']
+    >>>
+    >>> # Subset using query string
+    >>> adata = adata_.copy()
+    >>> subset_var(adata, 'gene_type == "type1"')
+    >>> adata.var_names.tolist()
+    ['gene1', 'gene3']
+
+    Notes
+    -----
+    - The function modifies the AnnData object in-place
+    - When using a boolean mask, its length must match the number of variables
+    - When using integer indices, they must be valid indices for the variables
+    - Invalid variable names or indices will raise KeyError or IndexError respectively
+    - The order of variables in the output will match the order in the subset parameter
+    """
+
+    if isinstance(subset, str):
+        subset = adata.var.query(subset).index
+
+    if not isinstance(subset, pd.Index):
+        subset = np.asarray(subset)
+
+    # Handle boolean mask
+    if subset.dtype.kind == "b":
+        if len(subset) != adata.n_vars:
+            raise IndexError(
+                f"Boolean mask length ({len(subset)}) does not match number of "
+                f"variables ({adata.n_vars})"
+            )
+        subset = adata.var_names[subset]
+
+    # Handle integer indices
+    elif subset.dtype.kind in "iu":
+        if np.any(subset < 0) or np.any(subset >= adata.n_vars):
+            raise IndexError(
+                f"Integer indices must be between 0 and {adata.n_vars - 1}"
+            )
+        subset = adata.var_names[subset]
+
+    if adata.n_vars == subset.size and (subset == adata.var_names).all():
+        # No need to subset, avoid making a copy. Useful for large AnnData objects
+        return
+
+    adata._inplace_subset_var(subset)

sclab/preprocess/_transfer_metadata.py

@@ -0,0 +1,137 @@
+from collections import Counter
+from functools import partial
+from typing import 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",
+):
+    D: csr_matrix = adata.obsp["distances"]
+    C: csr_matrix = adata.obsp["connectivities"]
+    D = D.tocsr()
+
+    match weight_by:
+        case "connectivity":
+            W = C.tocsr()
+        case "distance":
+            W = D.tocsr()
+            W.data = 1.0 / W.data
+        case "constant":
+            W = D.tocsr()
+            W.data[:] = 1.0
+        case _:
+            raise ValueError(f"Unsupported weight_by {weight_by}")
+
+    meta_values: pd.Series
+    new_values: pd.Series
+
+    series = adata.obs[column]
+    if isinstance(series.dtype, pd.CategoricalDtype) or is_bool_dtype(series.dtype):
+        assign_value_fn = _assign_categorical
+        new_column = f"transferred_{column}"
+        new_column_err = f"transferred_{column}_proportion"
+    elif is_numeric_dtype(series.dtype) and periodic:
+        assign_value_fn = partial(_assign_numerical_periodic, vmin=vmin, vmax=vmax)
+        new_column = f"transferred_{column}"
+        new_column_err = f"transferred_{column}_error"
+    elif is_numeric_dtype(series.dtype):
+        assign_value_fn = _assign_numerical
+        new_column = f"transferred_{column}"
+        new_column_err = f"transferred_{column}_error"
+    else:
+        raise ValueError(f"Unsupported dtype {series.dtype} for column {column}")
+
+    meta_values = series.copy()
+    meta_values[adata.obs[group_key] != source_group] = np.nan
+    new_values = pd.Series(index=series.index, dtype=series.dtype, name=new_column)
+    new_values_err = pd.Series(index=series.index, dtype=float, name=new_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
+
+    adata.obs[new_column] = new_values.copy()
+    adata.obs[new_column_err] = new_values_err.copy()
+
+
+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()