pertpy 0.9.5__py3-none-any.whl → 0.11.0__py3-none-any.whl

pertpy/preprocessing/_guide_rna.py

@@ -1,44 +1,53 @@
 from __future__ import annotations
 
 import uuid
-from typing import TYPE_CHECKING
+from functools import singledispatchmethod
+from typing import TYPE_CHECKING, Literal
+from warnings import warn
 
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import scanpy as sc
-import scipy
+from anndata import AnnData
+from numba import njit, prange
+from rich.progress import track
+from scanpy.get import _get_obs_rep, _set_obs_rep
+from scipy.sparse import csr_matrix, issparse
 
 from pertpy._doc import _doc_params, doc_common_plot_args
+from pertpy._types import CSRBase
+from pertpy.preprocessing._guide_rna_mixture import PoissonGaussMixture
 
 if TYPE_CHECKING:
-    from anndata import AnnData
     from matplotlib.pyplot import Figure
 
 
 class GuideAssignment:
-    """Offers simple guide assigment based on count thresholds."""
+    """Assign cells to guide RNAs."""
 
+    @singledispatchmethod
     def assign_by_threshold(
         self,
-        adata: AnnData,
+        data: AnnData | np.ndarray | CSRBase,
+        /,
+        *,
         assignment_threshold: float,
         layer: str | None = None,
         output_layer: str = "assigned_guides",
-        only_return_results: bool = False,
-    ) -> np.ndarray | None:
+    ):
         """Simple threshold based gRNA assignment function.
 
         Each cell is assigned to gRNA with at least `assignment_threshold` counts.
         This function expects unnormalized data as input.
 
         Args:
-            adata: Annotated data matrix containing gRNA values
+            data: The (annotated) data matrix of shape `n_obs` × `n_vars`.
+                  Rows correspond to cells and columns to genes.
             assignment_threshold: The count threshold that is required for an assignment to be viable.
             layer: Key to the layer containing raw count values of the gRNAs.
                    adata.X is used if layer is None. Expects count data.
             output_layer: Assigned guide will be saved on adata.layers[output_key].
-            only_return_results: If True, input AnnData is not modified and the result is returned as an np.ndarray.
 
         Examples:
             Each cell is assigned to gRNA that occurs at least 5 times in the respective cell.
@@ -49,26 +58,52 @@ class GuideAssignment:
             >>> ga = pt.pp.GuideAssignment()
             >>> ga.assign_by_threshold(gdo, assignment_threshold=5)
         """
-        counts = adata.X if layer is None else adata.layers[layer]
-        if scipy.sparse.issparse(counts):
-            counts = counts.toarray()
-
-        assigned_grnas = np.where(counts >= assignment_threshold, 1, 0)
-        assigned_grnas = scipy.sparse.csr_matrix(assigned_grnas)
-        if only_return_results:
-            return assigned_grnas
-        adata.layers[output_layer] = assigned_grnas
-
-        return None
+        raise NotImplementedError(
+            f"No implementation found for {type(data)}. Must be numpy array, sparse matrix, or AnnData object."
+        )
 
-    def assign_to_max_guide(
+    @assign_by_threshold.register(AnnData)
+    def _assign_by_threshold_anndata(
         self,
         adata: AnnData,
+        /,
+        *,
         assignment_threshold: float,
         layer: str | None = None,
-        output_key: str = "assigned_guide",
-        no_grna_assigned_key: str = "NT",
-        only_return_results: bool = False,
+        output_layer: str = "assigned_guides",
+    ) -> None:
+        X = _get_obs_rep(adata, layer=layer)
+        guide_assignments = self.assign_by_threshold(X, assignment_threshold=assignment_threshold)
+        _set_obs_rep(adata, guide_assignments, layer=output_layer)
+
+    @assign_by_threshold.register(np.ndarray)
+    def _assign_by_threshold_numpy(self, X: np.ndarray, /, *, assignment_threshold: float) -> np.ndarray:
+        return np.where(assignment_threshold <= X, 1, 0)
+
+    @staticmethod
+    @njit(parallel=True)
+    def _threshold_sparse_numba(data: np.ndarray, threshold: float) -> np.ndarray:
+        out = np.zeros_like(data, dtype=np.int8)
+        for i in prange(data.shape[0]):
+            if data[i] >= threshold:
+                out[i] = 1
+        return out
+
+    @assign_by_threshold.register(CSRBase)
+    def _assign_by_threshold_sparse(self, X: CSRBase, /, *, assignment_threshold: float) -> CSRBase:
+        new_data = self._threshold_sparse_numba(X.data, assignment_threshold)
+        return csr_matrix((new_data, X.indices, X.indptr), shape=X.shape)
+
+    @singledispatchmethod
+    def assign_to_max_guide(
+        self,
+        data: AnnData | np.ndarray | CSRBase,
+        /,
+        *,
+        assignment_threshold: float,
+        layer: str | None = None,
+        obs_key: str = "assigned_guide",
+        no_grna_assigned_key: str = "Negative",
     ) -> np.ndarray | None:
         """Simple threshold based max gRNA assignment function.
 
@@ -76,13 +111,13 @@ class GuideAssignment:
         This function expects unnormalized data as input.
 
         Args:
-            adata: Annotated data matrix containing gRNA values
+            data: The (annotated) data matrix of shape `n_obs` × `n_vars`.
+                  Rows correspond to cells and columns to genes.
             assignment_threshold: The count threshold that is required for an assignment to be viable.
             layer: Key to the layer containing raw count values of the gRNAs.
                    adata.X is used if layer is None. Expects count data.
-            output_key: Assigned guide will be saved on adata.obs[output_key]. default value is `assigned_guide`.
+            obs_key: Assigned guide will be saved on adata.obs[output_key].
             no_grna_assigned_key: The key to return if no gRNA is expressed enough.
-            only_return_results: If True, input AnnData is not modified and the result is returned as an np.ndarray.
 
         Examples:
             Each cell is assigned to the most expressed gRNA if it has at least 5 counts.
@@ -93,31 +128,179 @@ class GuideAssignment:
             >>> ga = pt.pp.GuideAssignment()
             >>> ga.assign_to_max_guide(gdo, assignment_threshold=5)
         """
-        counts = adata.X if layer is None else adata.layers[layer]
-        if scipy.sparse.issparse(counts):
-            counts = counts.toarray()
+        raise NotImplementedError(
+            f"No implementation found for {type(data)}. Must be numpy array, sparse matrix, or AnnData object."
+        )
+
+    @assign_to_max_guide.register(AnnData)
+    def assign_to_max_guide_anndata(
+        self,
+        adata: AnnData,
+        /,
+        *,
+        assignment_threshold: float,
+        layer: str | None = None,
+        obs_key: str = "assigned_guide",
+        no_grna_assigned_key: str = "Negative",
+    ) -> None:
+        X = _get_obs_rep(adata, layer=layer)
+        guide_assignments = self.assign_to_max_guide(
+            X, var=adata.var, assignment_threshold=assignment_threshold, no_grna_assigned_key=no_grna_assigned_key
+        )
+        adata.obs[obs_key] = guide_assignments
 
+    @assign_to_max_guide.register(np.ndarray)
+    def assign_to_max_guide_numpy(
+        self,
+        X: np.ndarray,
+        /,
+        *,
+        var: pd.DataFrame,
+        assignment_threshold: float,
+        no_grna_assigned_key: str = "Negative",
+    ) -> np.ndarray:
         assigned_grna = np.where(
-            counts.max(axis=1).squeeze() >= assignment_threshold,
-            adata.var.index[counts.argmax(axis=1).squeeze()],
+            X.max(axis=1).squeeze() >= assignment_threshold,
+            var.index[X.argmax(axis=1).squeeze()],
             no_grna_assigned_key,
         )
 
+        return assigned_grna
+
+    @staticmethod
+    @njit(parallel=True)
+    def _assign_max_guide_sparse(indptr, data, indices, assignment_threshold, assigned_grna):
+        n_rows = len(indptr) - 1
+        for i in range(n_rows):
+            row_start = indptr[i]
+            row_end = indptr[i + 1]
+
+            if row_end > row_start:
+                data_row = data[row_start:row_end]
+                indices_row = indices[row_start:row_end]
+                max_pos = np.argmax(data_row)
+                if data_row[max_pos] >= assignment_threshold:
+                    assigned_grna[i] = indices_row[max_pos]
+        return assigned_grna
+
+    @assign_to_max_guide.register(CSRBase)
+    def assign_to_max_guide_sparse(
+        self, X: CSRBase, /, *, var: pd.DataFrame, assignment_threshold: float, no_grna_assigned_key: str = "Negative"
+    ) -> np.ndarray:
+        n_rows = X.shape[0]
+
+        assigned_positions = np.zeros(n_rows, dtype=np.int32) - 1  # -1 means not assigned
+        assigned_positions = self._assign_max_guide_sparse(
+            X.indptr, X.data, X.indices, assignment_threshold, assigned_positions
+        )
+
+        assigned_grna = np.full(n_rows, no_grna_assigned_key, dtype=object)
+        mask = assigned_positions >= 0
+        var_index_array = np.array(var.index)
+        if np.any(mask):
+            assigned_grna[mask] = var_index_array[assigned_positions[mask]]
+
+        return assigned_grna
+
+    def assign_mixture_model(
+        self,
+        adata: AnnData,
+        model: Literal["poisson_gauss_mixture"] = "poisson_gauss_mixture",
+        assigned_guides_key: str = "assigned_guide",
+        no_grna_assigned_key: str = "negative",
+        max_assignments_per_cell: int = 5,
+        multiple_grna_assigned_key: str = "multiple",
+        multiple_grna_assignment_string: str = "+",
+        only_return_results: bool = False,
+        show_progress: bool = False,
+        **mixture_model_kwargs,
+    ) -> np.ndarray | None:
+        """Assigns gRNAs to cells using a mixture model.
+
+        Args:
+            adata: AnnData object containing gRNA values.
+            model: The model to use for the mixture model. Currently only `Poisson_Gauss_Mixture` is supported.
+            assigned_guides_key: Assigned guide will be saved on adata.obs[output_key].
+            no_grna_assigned_key: The key to return if a cell is negative for all gRNAs.
+            max_assignments_per_cell: The maximum number of gRNAs that can be assigned to a cell.
+            multiple_grna_assigned_key: The key to return if multiple gRNAs are assigned to a cell.
+            multiple_grna_assignment_string: The string to use to join multiple gRNAs assigned to a cell.
+            only_return_results: Whether input AnnData is not modified and the result is returned as an np.ndarray.
+            show_progress: Whether to shows progress bar.
+            mixture_model_kwargs: Are passed to the mixture model.
+
+        Examples:
+            >>> import pertpy as pt
+            >>> mdata = pt.dt.papalexi_2021()
+            >>> gdo = mdata.mod["gdo"]
+            >>> ga = pt.pp.GuideAssignment()
+            >>> ga.assign_mixture_model(gdo)
+        """
+        if model == "poisson_gauss_mixture":
+            mixture_model = PoissonGaussMixture(**mixture_model_kwargs)
+        else:
+            raise ValueError("Model not implemented. Please use 'poisson_gauss_mixture'.")
+
+        res = pd.DataFrame(0, index=adata.obs_names, columns=adata.var_names)
+        fct = track if show_progress else lambda iterable: iterable
+        for gene in fct(adata.var_names):
+            is_nonzero = (
+                np.ravel((adata[:, gene].X != 0).todense()) if issparse(adata.X) else np.ravel(adata[:, gene].X != 0)
+            )
+            if sum(is_nonzero) < 2:
+                warn(f"Skipping {gene} as there are less than 2 cells expressing the guide at all.", stacklevel=2)
+                continue
+            # We are only fitting the model to the non-zero values, the rest is
+            # automatically assigned to the negative class
+            data = adata[is_nonzero, gene].X.todense().A1 if issparse(adata.X) else adata[is_nonzero, gene].X
+            data = np.ravel(data)
+
+            if np.any(data < 0):
+                raise ValueError(
+                    "Data contains negative values. Please use non-negative data for guide assignment with the Mixture Model."
+                )
+
+            # Log2 transform the data so positive population is approximately normal
+            data = np.log2(data)
+            assignments = mixture_model.run_model(data)
+            res.loc[adata.obs_names[is_nonzero][assignments == "Positive"], gene] = 1
+
+            # Add the parameters to the adata.var DataFrame
+            for params_name, param in mixture_model.params.items():
+                if param.ndim == 0:
+                    if params_name not in adata.var.columns:
+                        adata.var[params_name] = np.nan
+                    adata.var.loc[gene, params_name] = param.item()
+                else:
+                    for i, p in enumerate(param):
+                        if f"{params_name}_{i}" not in adata.var.columns:
+                            adata.var[f"{params_name}_{i}"] = np.nan
+                        adata.var.loc[gene, f"{params_name}_{i}"] = p
+
+        # Assign guides to cells
+        # Some cells might have multiple guides assigned
+        series = pd.Series(no_grna_assigned_key, index=adata.obs_names)
+        num_guides_assigned = res.sum(1)
+        series.loc[(num_guides_assigned <= max_assignments_per_cell) & (num_guides_assigned != 0)] = res.apply(
+            lambda row: row.index[row == 1].tolist(), axis=1
+        ).str.join(multiple_grna_assignment_string)
+        series.loc[num_guides_assigned > max_assignments_per_cell] = multiple_grna_assigned_key
+
         if only_return_results:
-            return assigned_grna
-        adata.obs[output_key] = assigned_grna
+            return series.values
+
+        adata.obs[assigned_guides_key] = series.values
 
         return None
 
     @_doc_params(common_plot_args=doc_common_plot_args)
-    def plot_heatmap(
+    def plot_heatmap(  # pragma: no cover # noqa: D417
         self,
         adata: AnnData,
         *,
         layer: str | None = None,
         order_by: np.ndarray | str | None = None,
         key_to_save_order: str = None,
-        show: bool = True,
         return_fig: bool = False,
         **kwargs,
     ) -> Figure | None:
@@ -159,7 +342,7 @@ class GuideAssignment:
         data = adata.X if layer is None else adata.layers[layer]
 
         if order_by is None:
-            if scipy.sparse.issparse(data):
+            if issparse(data):
                 max_values = data.max(axis=1).toarray().squeeze()
                 data_argmax = data.argmax(axis=1).A.squeeze()
                 max_guide_index = np.where(max_values != data.min(axis=1).toarray().squeeze(), data_argmax, -1)
@@ -194,8 +377,7 @@ class GuideAssignment:
         finally:
             del adata.obs[temp_col_name]
 
-        if show:
-            plt.show()
         if return_fig:
             return fig
+        plt.show()
         return None

pertpy/preprocessing/_guide_rna_mixture.py

@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+
+import jax.numpy as jnp
+import numpy as np
+from jax.random import PRNGKey
+from jax.scipy.special import logsumexp
+from numpyro import factor, plate, sample
+from numpyro.distributions import Dirichlet, Exponential, HalfNormal, Normal, Poisson
+from numpyro.infer import MCMC, NUTS
+
+ParamsDict = Mapping[str, jnp.ndarray]
+
+
+class MixtureModel(ABC):
+    """Abstract base class for 2-component mixture models.
+
+    Args:
+        num_warmup: Number of warmup steps for MCMC sampling.
+        num_samples: Number of samples to draw after warmup.
+        fraction_positive_expected: Prior belief about fraction of positive components.
+        poisson_rate_prior: Rate parameter for exponential prior on Poisson component.
+        gaussian_mean_prior: Mean and standard deviation for Gaussian prior on positive component mean.
+        gaussian_std_prior: Scale parameter for half-normal prior on positive component std.
+    """
+
+    def __init__(
+        self,
+        num_warmup: int = 50,
+        num_samples: int = 100,
+        fraction_positive_expected: float = 0.15,
+        poisson_rate_prior: float = 0.2,
+        gaussian_mean_prior: tuple[float, float] = (3, 2),
+        gaussian_std_prior: float = 1,
+    ) -> None:
+        self.num_warmup = num_warmup
+        self.num_samples = num_samples
+        self.fraction_positive_expected = fraction_positive_expected
+        self.poisson_rate_prior = poisson_rate_prior
+        self.gaussian_mean_prior = gaussian_mean_prior
+        self.gaussian_std_prior = gaussian_std_prior
+
+    @abstractmethod
+    def initialize_params(self) -> ParamsDict:
+        """Initialize model parameters via sampling from priors.
+
+        Returns:
+            Dictionary of sampled parameter values.
+        """
+
+    @abstractmethod
+    def log_likelihood(self, data: jnp.ndarray, params: ParamsDict) -> jnp.ndarray:
+        """Calculate log likelihood of data under current parameters.
+
+        Args:
+            data: Input data array.
+            params: Current parameter values.
+
+        Returns:
+            Log likelihood values for each datapoint.
+        """
+
+    def fit_model(self, data: jnp.ndarray, seed: int = 0) -> MCMC:
+        """Fit the mixture model using MCMC.
+
+        Args:
+            data: Input data to fit.
+            seed: Random seed for reproducibility.
+
+        Returns:
+            Fitted MCMC object containing samples.
+        """
+        nuts_kernel = NUTS(self.mixture_model)
+        mcmc = MCMC(nuts_kernel, num_warmup=self.num_warmup, num_samples=self.num_samples, progress_bar=False)
+        mcmc.run(PRNGKey(seed), data=data)
+        return mcmc
+
+    def run_model(self, data: jnp.ndarray, seed: int = 0) -> np.ndarray:
+        """Run model fitting and assign components.
+
+        Args:
+            data: Input data array.
+            seed: Random seed.
+
+        Returns:
+            Array of "Positive"/"Negative" assignments for each datapoint.
+        """
+        self.mcmc = self.fit_model(data, seed)
+        self.samples = self.mcmc.get_samples()
+        self.assignments = self.assignment(self.samples, data)
+        return self.assignments
+
+    def mixture_model(self, data: jnp.ndarray) -> None:
+        """Define mixture model structure for NumPyro.
+
+        Args:
+            data: Input data array.
+        """
+        params = self.initialize_params()
+
+        with plate("data", data.shape[0]):
+            log_likelihoods = self.log_likelihood(data, params)
+            log_mixture_likelihood = logsumexp(log_likelihoods, axis=-1)
+            sample("obs", Normal(log_mixture_likelihood, 1.0), obs=data)
+
+    def assignment(self, samples: ParamsDict, data: jnp.ndarray) -> np.ndarray:
+        """Assign data points to mixture components.
+
+        Args:
+            samples: MCMC samples of parameters.
+            data: Input data array.
+
+        Returns:
+            Array of component assignments.
+        """
+        params = {key: samples[key].mean(axis=0) for key in samples}
+        self.params = params
+
+        log_likelihoods = self.log_likelihood(data, params)
+        guide_assignments = jnp.argmax(log_likelihoods, axis=-1)
+
+        assignments = ["Negative" if assign == 0 else "Positive" for assign in guide_assignments]
+        return np.array(assignments)
+
+
+class PoissonGaussMixture(MixtureModel):
+    """Mixture model combining Poisson and Gaussian distributions."""
+
+    def log_likelihood(self, data: np.ndarray, params: ParamsDict) -> jnp.ndarray:
+        """Calculate component-wise log likelihoods.
+
+        Args:
+            data: Input data array.
+            params: Current parameter values.
+
+        Returns:
+            Log likelihood values for each component.
+        """
+        poisson_rate = params["poisson_rate"]
+        gaussian_mean = params["gaussian_mean"]
+        gaussian_std = params["gaussian_std"]
+        mix_probs = params["mix_probs"]
+
+        # We penalize the model for positioning the Poisson component to the right of the Gaussian component
+        # by imposing a soft constraint to penalize the Poisson rate being larger than the Gaussian mean
+        # Heuristic regularization term to prevent flipping of the components
+        factor("separation_penalty", +10 * jnp.heaviside(-poisson_rate + gaussian_mean, 0))
+
+        log_likelihoods = jnp.stack(
+            [
+                # Poisson component
+                jnp.log(mix_probs[0]) + Poisson(poisson_rate).log_prob(data),
+                # Gaussian component
+                jnp.log(mix_probs[1]) + Normal(gaussian_mean, gaussian_std).log_prob(data),
+            ],
+            axis=-1,
+        )
+
+        return log_likelihoods
+
+    def initialize_params(self) -> ParamsDict:
+        """Initialize model parameters via prior sampling.
+
+        Returns:
+            Dictionary of sampled parameter values.
+        """
+        params = {}
+        params["poisson_rate"] = sample("poisson_rate", Exponential(self.poisson_rate_prior))
+        params["gaussian_mean"] = sample("gaussian_mean", Normal(*self.gaussian_mean_prior))
+        params["gaussian_std"] = sample("gaussian_std", HalfNormal(self.gaussian_std_prior))
+        params["mix_probs"] = sample(
+            "mix_probs",
+            Dirichlet(jnp.array([1 - self.fraction_positive_expected, self.fraction_positive_expected])),
+        )
+        return params

pertpy/tools/__init__.py

@@ -41,7 +41,7 @@ from pertpy.tools._perturbation_space._simple import (
 )
 from pertpy.tools._scgen import Scgen
 
-CODA_EXTRAS = ["toytree", "arviz", "ete3"]  # also pyqt5 technically
+CODA_EXTRAS = ["toytree", "arviz", "ete4"]  # also pyqt6 technically
 Sccoda = lazy_import("pertpy.tools._coda._sccoda", "Sccoda", CODA_EXTRAS)
 Tasccoda = lazy_import("pertpy.tools._coda._tasccoda", "Tasccoda", CODA_EXTRAS)