pertpy 0.9.4__py3-none-any.whl → 0.10.0__py3-none-any.whl

pertpy/__init__.py

@@ -2,7 +2,7 @@
 
 __author__ = "Lukas Heumos"
 __email__ = "lukas.heumos@posteo.net"
-__version__ = "0.9.4"
+__version__ = "0.10.0"
 
 import warnings
 

pertpy/_doc.py

@@ -0,0 +1,19 @@
+from textwrap import dedent
+
+
+def _doc_params(**kwds):  # pragma: no cover
+    """\
+    Docstrings should start with "\" in the first line for proper formatting.
+    """
+
+    def dec(obj):
+        obj.__orig_doc__ = obj.__doc__
+        obj.__doc__ = dedent(obj.__doc__.format_map(kwds))
+        return obj
+
+    return dec
+
+
+doc_common_plot_args = """\
+return_fig: if `True`, returns figure of the plot, that can be used for saving.\
+"""

pertpy/data/_datasets.py

@@ -66,7 +66,7 @@ def sc_sim_augur() -> AnnData:  # pragma: no cover
     output_file_path = settings.datasetdir / output_file_name
     if not Path(output_file_path).exists():
         _download(
-            url="https://figshare.com/ndownloader/files/31645886",
+            url="https://figshare.com/ndownloader/files/49828902",
             output_file_name=output_file_name,
             output_path=settings.datasetdir,
             is_zip=False,

pertpy/metadata/_cell_line.py

@@ -8,12 +8,15 @@ from lamin_utils import logger
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
+    from matplotlib.pyplot import Figure
+
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 from scanpy import settings
 from scipy import stats
 
+from pertpy._doc import _doc_params, doc_common_plot_args
 from pertpy.data._dataloader import _download
 
 from ._look_up import LookUp
@@ -338,8 +341,8 @@ class CellLine(MetaData):
         # then we can compare these keys and fetch the corresponding metadata.
         if query_id not in adata.obs.columns:
             raise ValueError(
-                f"The specified `query_id` {query_id} can't be found in the `adata.obs`.\n"
-                "Ensure that you are using one of the available query IDs present in the adata.obs for the annotation.\n"
+                f"The specified `query_id` {query_id} can't be found in the `adata.obs`. \n"
+                "Ensure that you are using one of the available query IDs present in the adata.obs for the annotation."
                 "If the desired query ID is not available, you can fetch the cell line metadata "
                 "using the `annotate()` function before calling 'annotate_bulk_rna()'. "
                 "This ensures that the required query ID is included in your data, e.g. stripped_cell_line_name, DepMap ID."
@@ -356,9 +359,8 @@ class CellLine(MetaData):
         else:
             reference_id = "DepMap_ID"
             logger.warning(
-                "To annotate bulk RNA data from Broad Institue, `DepMap_ID` is used as default reference and query identifier if no `reference_id` is given.\n"
-                "Ensure that `DepMap_ID` is available in 'adata.obs'.\n"
-                "Alternatively, use `annotate()` to annotate the cell line first "
+                "To annotate bulk RNA data from Broad Institue, `DepMap_ID` is used as default reference and query identifier if no `reference_id` is given."
+                "If `DepMap_ID` isn't available in 'adata.obs', use `annotate()` to annotate the cell line first."
             )
             if self.bulk_rna_broad is None:
                 self._download_bulk_rna(cell_line_source="broad")
@@ -690,6 +692,7 @@ class CellLine(MetaData):
 
         return corr, pvals, new_corr, new_pvals
 
+    @_doc_params(common_plot_args=doc_common_plot_args)
     def plot_correlation(
         self,
         adata: AnnData,
@@ -700,7 +703,8 @@ class CellLine(MetaData):
         metadata_key: str = "bulk_rna_broad",
         category: str = "cell line",
         subset_identifier: str | int | Iterable[str] | Iterable[int] | None = None,
-    ) -> None:
+        return_fig: bool = False,
+    ) -> Figure | None:
         """Visualise the correlation of cell lines with annotated metadata.
 
         Args:
@@ -713,6 +717,8 @@ class CellLine(MetaData):
             subset_identifier: Selected identifiers for scatter plot visualization between the X matrix and `metadata_key`.
                               If not None, only the chosen cell line will be plotted, either specified as a value in `identifier` (string) or as an index number.
                               If None, all cell lines will be plotted.
+            {common_plot_args}
+
         Returns:
             Pearson correlation coefficients and their corresponding p-values for matched and unmatched cell lines separately.
         """
@@ -740,7 +746,7 @@ class CellLine(MetaData):
                 if all(isinstance(id, str) for id in subset_identifier_list):
                     if set(subset_identifier_list).issubset(adata.obs[identifier].unique()):
                         subset_identifier_list = np.where(
-                            np.in1d(adata.obs[identifier].values, subset_identifier_list)
+                            np.isin(adata.obs[identifier].values, subset_identifier_list)
                         )[0]
                     else:
                         raise ValueError("`Subset_identifier` must be found in adata.obs.`identifier`.")
@@ -790,6 +796,10 @@ class CellLine(MetaData):
                     "edgecolor": "black",
                 },
             )
+
+            if return_fig:
+                return plt.gcf()
             plt.show()
+            return None
         else:
-            raise NotImplementedError
+            raise NotImplementedError("Only 'cell line' category is supported for correlation comparison.")

pertpy/metadata/_compound.py

@@ -42,7 +42,7 @@ class Compound(MetaData):
             adata = adata.copy()
 
         if query_id not in adata.obs.columns:
-            raise ValueError(f"The requested query_id {query_id} is not in `adata.obs`.\n" f"Please check again. ")
+            raise ValueError(f"The requested query_id {query_id} is not in `adata.obs`.\n Please check again.")
 
         query_dict = {}
         not_matched_identifiers = []
@@ -84,7 +84,7 @@ class Compound(MetaData):
         query_df = pd.DataFrame.from_dict(query_dict, orient="index", columns=["pubchem_name", "pubchem_ID", "smiles"])
         # Merge and remove duplicate columns
         # Column is converted to float after merging due to unmatches
-        # Convert back to integers
+        # Convert back to integers afterwards
         if query_id_type == "cid":
             query_df.pubchem_ID = query_df.pubchem_ID.astype("Int64")
             adata.obs = (
@@ -119,8 +119,7 @@ class Compound(MetaData):
 
         The LookUp object provides an overview of the metadata to annotate.
         Each annotate_{metadata} function has a corresponding lookup function in the LookUp object,
-        where users can search the reference_id in the metadata and
-        compare with the query_id in their own data.
+        where users can search the reference_id in the metadata and compare with the query_id in their own data.
 
         Returns:
             Returns a LookUp object specific for compound annotation.

pertpy/metadata/_metadata.py

@@ -62,7 +62,7 @@ class MetaData:
             if verbosity > 0:
                 logger.info(
                     f"There are {total_identifiers} identifiers in `adata.obs`."
-                    f"However, {len(unmatched_identifiers)} identifiers can't be found in the {metadata_type} annotation,"
+                    f"However, {len(unmatched_identifiers)} identifiers can't be found in the {metadata_type} annotation, "
                     "leading to the presence of NA values for their respective metadata.\n"
                     f"Please check again: *unmatched_identifiers[:verbosity]..."
                 )

pertpy/preprocessing/_guide_rna.py

@@ -1,20 +1,27 @@
 from __future__ import annotations
 
 import uuid
-from typing import TYPE_CHECKING
+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 rich.progress import track
+from scipy.sparse import issparse
+
+from pertpy._doc import _doc_params, doc_common_plot_args
+from pertpy.preprocessing._guide_rna_mixture import PoissonGaussMixture
 
 if TYPE_CHECKING:
     from anndata import AnnData
-    from matplotlib.axes import Axes
+    from matplotlib.pyplot import Figure
 
 
 class GuideAssignment:
-    """Offers simple guide assigment based on count thresholds."""
+    """Assign cells to guide RNAs."""
 
     def assign_by_threshold(
         self,
@@ -30,12 +37,12 @@ class GuideAssignment:
         This function expects unnormalized data as input.
 
         Args:
-            adata: Annotated data matrix containing gRNA values
+            adata: AnnData object containing gRNA values.
             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.
+            only_return_results: Whether to input AnnData is not modified and the result is returned as an :class:`np.ndarray`.
 
         Examples:
             Each cell is assigned to gRNA that occurs at least 5 times in the respective cell.
@@ -64,7 +71,7 @@ class GuideAssignment:
         assignment_threshold: float,
         layer: str | None = None,
         output_key: str = "assigned_guide",
-        no_grna_assigned_key: str = "NT",
+        no_grna_assigned_key: str = "Negative",
         only_return_results: bool = False,
     ) -> np.ndarray | None:
         """Simple threshold based max gRNA assignment function.
@@ -73,13 +80,13 @@ class GuideAssignment:
         This function expects unnormalized data as input.
 
         Args:
-            adata: Annotated data matrix containing gRNA values
+            adata: AnnData object containing gRNA values.
             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`.
             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.
+            only_return_results: Whether to 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.
@@ -106,14 +113,103 @@ class GuideAssignment:
 
         return None
 
+    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,
+        uns_key: str = "guide_assignment_params",
+        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.
+            output_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'.")
+
+        if uns_key not in adata.uns:
+            adata.uns[uns_key] = {}
+        elif type(adata.uns[uns_key]) is not dict:
+            raise ValueError(f"adata.uns['{uns_key}'] should be a dictionary. Please remove it or change the key.")
+
+        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
+            adata.uns[uns_key][gene] = mixture_model.params
+
+        # 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 series.values
+
+        adata.obs[assigned_guides_key] = series.values
+
+        return None
+
+    @_doc_params(common_plot_args=doc_common_plot_args)
     def plot_heatmap(
         self,
         adata: AnnData,
+        *,
         layer: str | None = None,
         order_by: np.ndarray | str | None = None,
         key_to_save_order: str = None,
+        return_fig: bool = False,
         **kwargs,
-    ) -> list[Axes]:
+    ) -> Figure | None:
         """Heatmap plotting of guide RNA expression matrix.
 
         Assuming guides have sparse expression, this function reorders cells
@@ -131,11 +227,12 @@ class GuideAssignment:
                       If a string is provided, adata.obs[order_by] will be used as the order.
                       If a numpy array is provided, the array will be used for ordering.
             key_to_save_order: The obs key to save cell orders in the current plot. Only saves if not None.
+            {common_plot_args}
             kwargs: Are passed to sc.pl.heatmap.
 
         Returns:
-            List of Axes. Alternatively you can pass save or show parameters as they will be passed to sc.pl.heatmap.
-            Order of cells in the y-axis will be saved on adata.obs[key_to_save_order] if provided.
+            If `return_fig` is `True`, returns the figure, otherwise `None`.
+            Order of cells in the y-axis will be saved on `adata.obs[key_to_save_order]` if provided.
 
         Examples:
             Each cell is assigned to gRNA that occurs at least 5 times in the respective cell, which is then
@@ -172,7 +269,7 @@ class GuideAssignment:
             adata.obs[key_to_save_order] = pd.Categorical(order)
 
         try:
-            axis_group = sc.pl.heatmap(
+            fig = sc.pl.heatmap(
                 adata[order, :],
                 var_names=adata.var.index.tolist(),
                 groupby=temp_col_name,
@@ -180,9 +277,13 @@ class GuideAssignment:
                 use_raw=False,
                 dendrogram=False,
                 layer=layer,
+                show=False,
                 **kwargs,
             )
         finally:
             del adata.obs[temp_col_name]
 
-        return axis_group
+        if return_fig:
+            return fig
+        plt.show()
+        return None

pertpy/preprocessing/_guide_rna_mixture.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import numpyro
+import numpyro.distributions as dist
+from jax import random
+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.
+        """
+        pass
+
+    @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.
+        """
+        pass
+
+    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(random.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 numpyro.plate("data", data.shape[0]):
+            log_likelihoods = self.log_likelihood(data, params)
+            log_mixture_likelihood = jax.scipy.special.logsumexp(log_likelihoods, axis=-1)
+            numpyro.sample("obs", dist.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.keys()}
+        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
+        numpyro.factor("separation_penalty", +10 * jnp.heaviside(-poisson_rate + gaussian_mean, 0))
+
+        log_likelihoods = jnp.stack(
+            [
+                # Poisson component
+                jnp.log(mix_probs[0]) + dist.Poisson(poisson_rate).log_prob(data),
+                # Gaussian component
+                jnp.log(mix_probs[1]) + dist.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"] = numpyro.sample("poisson_rate", dist.Exponential(self.poisson_rate_prior))
+        params["gaussian_mean"] = numpyro.sample("gaussian_mean", dist.Normal(*self.gaussian_mean_prior))
+        params["gaussian_std"] = numpyro.sample("gaussian_std", dist.HalfNormal(self.gaussian_std_prior))
+        params["mix_probs"] = numpyro.sample(
+            "mix_probs",
+            dist.Dirichlet(jnp.array([1 - self.fraction_positive_expected, self.fraction_positive_expected])),
+        )
+        return params

pertpy/tools/__init__.py

@@ -46,7 +46,7 @@ Sccoda = lazy_import("pertpy.tools._coda._sccoda", "Sccoda", CODA_EXTRAS)
 Tasccoda = lazy_import("pertpy.tools._coda._tasccoda", "Tasccoda", CODA_EXTRAS)
 
 DE_EXTRAS = ["formulaic", "pydeseq2"]
-EdgeR = lazy_import("pertpy.tools._differential_gene_expression", "EdgeR", DE_EXTRAS) # edgeR will be imported via rpy2
+EdgeR = lazy_import("pertpy.tools._differential_gene_expression", "EdgeR", DE_EXTRAS)  # edgeR will be imported via rpy2
 PyDESeq2 = lazy_import("pertpy.tools._differential_gene_expression", "PyDESeq2", DE_EXTRAS)
 Statsmodels = lazy_import("pertpy.tools._differential_gene_expression", "Statsmodels", DE_EXTRAS + ["statsmodels"])
 TTest = lazy_import("pertpy.tools._differential_gene_expression", "TTest", DE_EXTRAS)