pertpy 0.6.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

pertpy/tools/_coda/_sccoda.py

@@ -8,8 +8,8 @@ import numpy as np
 import numpyro as npy
 import numpyro.distributions as npd
 from anndata import AnnData
-from jax import random
-from jax.config import config
+from jax import config, random
+from lamin_utils import logger
 from mudata import MuData
 from numpyro.infer import Predictive
 from rich import print
@@ -23,7 +23,6 @@ config.update("jax_enable_x64", True)
 
 
 class Sccoda(CompositionalModel2):
-
     """
     Statistical model for single-cell differential composition analysis with specification of a reference cell type.
     This is the standard scCODA model and recommended for all uses.
@@ -75,13 +74,13 @@ class Sccoda(CompositionalModel2):
             adata: AnnData object.
             type : Specify the input adata type, which could be either a cell-level AnnData or an aggregated sample-level AnnData.
             generate_sample_level: Whether to generate an AnnData object on the sample level or create an empty AnnData object.
-            cell_type_identifier: If type is "cell_level", specify column name in adata.obs that specifies the cell types. Defaults to None.
-            sample_identifier: If type is "cell_level", specify column name in adata.obs that specifies the sample. Defaults to None.
-            covariate_uns: If type is "cell_level", specify key for adata.uns, where covariate values are stored. Defaults to None.
-            covariate_obs: If type is "cell_level", specify list of keys for adata.obs, where covariate values are stored. Defaults to None.
-            covariate_df: If type is "cell_level", specify dataFrame with covariates. Defaults to None.
-            modality_key_1: Key to the cell-level AnnData in the MuData object. Defaults to "rna".
-            modality_key_2: Key to the aggregated sample-level AnnData object in the MuData object. Defaults to "coda".
+            cell_type_identifier: If type is "cell_level", specify column name in adata.obs that specifies the cell types.
+            sample_identifier: If type is "cell_level", specify column name in adata.obs that specifies the sample.
+            covariate_uns: If type is "cell_level", specify key for adata.uns, where covariate values are stored.
+            covariate_obs: If type is "cell_level", specify list of keys for adata.obs, where covariate values are stored.
+            covariate_df: If type is "cell_level", specify dataFrame with covariates.
+            modality_key_1: Key to the cell-level AnnData in the MuData object.
+            modality_key_2: Key to the aggregated sample-level AnnData object in the MuData object.
 
         Returns:
             MuData: MuData object with cell-level AnnData (`mudata[modality_key_1]`) and aggregated sample-level AnnData (`mudata[modality_key_2]`).
@@ -90,8 +89,11 @@ class Sccoda(CompositionalModel2):
             >>> import pertpy as pt
             >>> haber_cells = pt.dt.haber_2017_regions()
             >>> sccoda = pt.tl.Sccoda()
-            >>> mdata = sccoda.load(haber_cells, type="cell_level", generate_sample_level=True, cell_type_identifier="cell_label", \
-                sample_identifier="batch", covariate_obs=["condition"])
+            >>> mdata = sccoda.load(haber_cells,
+            >>>                     type="cell_level",
+            >>>                     generate_sample_level=True,
+            >>>                     cell_type_identifier="cell_label",
+            >>>                     sample_identifier="batch", covariate_obs=["condition"])
         """
         if type == "cell_level":
             if generate_sample_level:
@@ -126,10 +128,10 @@ class Sccoda(CompositionalModel2):
                 Categorical covariates are handled automatically, with the covariate value of the first sample being used as the reference category.
                 To set a different level as the base category for a categorical covariate, use "C(<CovariateName>, Treatment('<ReferenceLevelName>'))"
             reference_cell_type: Column name that sets the reference cell type.
-                Reference the name of a column. If "automatic", the cell type with the lowest dispersion in relative abundance that is present in at least 90% of samlpes will be chosen. Defaults to "automatic".
+                Reference the name of a column. If "automatic", the cell type with the lowest dispersion in relative abundance that is present in at least 90% of samlpes will be chosen.
             automatic_reference_absence_threshold: If using reference_cell_type = "automatic", determine the maximum fraction of zero entries for a cell type
-                to be considered as a possible reference cell type. Defaults to 0.05.
-            modality_key: If data is a MuData object, specify key to the aggregated sample-level AnnData object in the MuData object. Defaults to "coda".
+                to be considered as a possible reference cell type.
+            modality_key: If data is a MuData object, specify key to the aggregated sample-level AnnData object in the MuData object.
 
         Returns:
             Return an AnnData (if input data is an AnnData object) or return a MuData (if input data is a MuData object)
@@ -144,8 +146,12 @@ class Sccoda(CompositionalModel2):
             >>> import pertpy as pt
             >>> haber_cells = pt.dt.haber_2017_regions()
             >>> sccoda = pt.tl.Sccoda()
-            >>> mdata = sccoda.load(haber_cells, type="cell_level", generate_sample_level=True, cell_type_identifier="cell_label", \
-                sample_identifier="batch", covariate_obs=["condition"])
+            >>> mdata = sccoda.load(haber_cells,
+            >>>                     type="cell_level",
+            >>>                     generate_sample_level=True,
+            >>>                     cell_type_identifier="cell_label",
+            >>>                     sample_identifier="batch",
+            >>>                     covariate_obs=["condition"])
             >>> mdata = sccoda.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
         """
         if isinstance(data, MuData):
@@ -193,10 +199,14 @@ class Sccoda(CompositionalModel2):
             >>> import pertpy as pt
             >>> haber_cells = pt.dt.haber_2017_regions()
             >>> sccoda = pt.tl.Sccoda()
-            >>> mdata = sccoda.load(haber_cells, type="cell_level", generate_sample_level=True, cell_type_identifier="cell_label", \
-                sample_identifier="batch", covariate_obs=["condition"])
+            >>> mdata = sccoda.load(haber_cells,
+            >>>                     type="cell_level",
+            >>>                     generate_sample_level=True,
+            >>>                     cell_type_identifier="cell_label",
+            >>>                     sample_identifier="batch",
+            >>>                     covariate_obs=["condition"])
             >>> mdata = sccoda.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
-            >>> adata = sccoda.set_init_mcmc_states(rng_key=42, ref_index=0, sample_adata=mdata['coda'])
+            >>> adata = sccoda.set_init_mcmc_states(rng_key=42, ref_index=0, sample_adata=mdata["coda"])
         """
         # data dimensions
         N, D = sample_adata.obsm["covariate_matrix"].shape
@@ -300,10 +310,10 @@ class Sccoda(CompositionalModel2):
 
         Args:
             data: AnnData object or MuData object.
-            modality_key: If data is a MuData object, specify which modality to use. Defaults to "coda".
-            rng_key: The rng state used for the prior simulation. If None, a random state will be selected. Defaults to None.
-            num_prior_samples: Number of prior samples calculated. Defaults to 500.
-            use_posterior_predictive: If True, the posterior predictive will be calculated. Defaults to True.
+            modality_key: If data is a MuData object, specify which modality to use.
+            rng_key: The rng state used for the prior simulation. If None, a random state will be selected.
+            num_prior_samples: Number of prior samples calculated.
+            use_posterior_predictive: If True, the posterior predictive will be calculated.
 
         Returns:
             az.InferenceData: arviz_data with all MCMC information
@@ -312,8 +322,12 @@ class Sccoda(CompositionalModel2):
             >>> import pertpy as pt
             >>> haber_cells = pt.dt.haber_2017_regions()
             >>> sccoda = pt.tl.Sccoda()
-            >>> mdata = sccoda.load(haber_cells, type="cell_level", generate_sample_level=True, cell_type_identifier="cell_label", \
-                sample_identifier="batch", covariate_obs=["condition"])
+            >>> mdata = sccoda.load(haber_cells,
+            >>>                     type="cell_level",
+            >>>                     generate_sample_level=True,
+            >>>                     cell_type_identifier="cell_label",
+            >>>                     sample_identifier="batch",
+            >>>                     covariate_obs=["condition"])
             >>> mdata = sccoda.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
             >>> sccoda.run_nuts(mdata, num_warmup=100, num_samples=1000, rng_key=42)
             >>> arviz_data = sccoda.make_arviz(mdata, num_prior_samples=100)
@@ -322,7 +336,7 @@ class Sccoda(CompositionalModel2):
             try:
                 sample_adata = data[modality_key]
             except IndexError:
-                print("When data is a MuData object, modality_key must be specified!")
+                logger.error("When data is a MuData object, modality_key must be specified!")
                 raise
         if isinstance(data, AnnData):
             sample_adata = data
@@ -365,7 +379,7 @@ class Sccoda(CompositionalModel2):
 
         if rng_key is None:
             rng = np.random.default_rng()
-            rng_key = random.PRNGKey(rng.integers(0, 10000))
+            rng_key = random.key(rng.integers(0, 10000))
 
         if use_posterior_predictive:
             posterior_predictive = Predictive(self.model, self.mcmc.get_samples())(
@@ -414,8 +428,12 @@ class Sccoda(CompositionalModel2):
             >>> import pertpy as pt
             >>> haber_cells = pt.dt.haber_2017_regions()
             >>> sccoda = pt.tl.Sccoda()
-            >>> mdata = sccoda.load(haber_cells, type="cell_level", generate_sample_level=True, cell_type_identifier="cell_label", \
-                sample_identifier="batch", covariate_obs=["condition"])
+            >>> mdata = sccoda.load(haber_cells,
+            >>>                     type="cell_level",
+            >>>                     generate_sample_level=True,
+            >>>                     cell_type_identifier="cell_label",
+            >>>                     sample_identifier="batch",
+            >>>                     covariate_obs=["condition"])
             >>> mdata = sccoda.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
             >>> sccoda.run_nuts(mdata, num_warmup=100, num_samples=1000, rng_key=42)
         """
@@ -429,8 +447,12 @@ class Sccoda(CompositionalModel2):
             >>> import pertpy as pt
             >>> haber_cells = pt.dt.haber_2017_regions()
             >>> sccoda = pt.tl.Sccoda()
-            >>> mdata = sccoda.load(haber_cells, type="cell_level", generate_sample_level=True, cell_type_identifier="cell_label", \
-                sample_identifier="batch", covariate_obs=["condition"])
+            >>> mdata = sccoda.load(haber_cells,
+            >>>                     type="cell_level",
+            >>>                     generate_sample_level=True,
+            >>>                     cell_type_identifier="cell_label",
+            >>>                     sample_identifier="batch",
+            >>>                     covariate_obs=["condition"])
             >>> mdata = sccoda.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
             >>> sccoda.run_nuts(mdata, num_warmup=100, num_samples=1000, rng_key=42)
             >>> credible_effects = sccoda.credible_effects(mdata)
@@ -445,8 +467,12 @@ class Sccoda(CompositionalModel2):
             >>> import pertpy as pt
             >>> haber_cells = pt.dt.haber_2017_regions()
             >>> sccoda = pt.tl.Sccoda()
-            >>> mdata = sccoda.load(haber_cells, type="cell_level", generate_sample_level=True, cell_type_identifier="cell_label", \
-                sample_identifier="batch", covariate_obs=["condition"])
+            >>> mdata = sccoda.load(haber_cells,
+            >>>                     type="cell_level",
+            >>>                     generate_sample_level=True,
+            >>>                     cell_type_identifier="cell_label",
+            >>>                     sample_identifier="batch",
+            >>>                     covariate_obs=["condition"])
             >>> mdata = sccoda.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
             >>> sccoda.run_nuts(mdata, num_warmup=100, num_samples=1000, rng_key=42)
             >>> sccoda.summary(mdata)
@@ -461,8 +487,12 @@ class Sccoda(CompositionalModel2):
             >>> import pertpy as pt
             >>> haber_cells = pt.dt.haber_2017_regions()
             >>> sccoda = pt.tl.Sccoda()
-            >>> mdata = sccoda.load(haber_cells, type="cell_level", generate_sample_level=True, cell_type_identifier="cell_label", \
-                sample_identifier="batch", covariate_obs=["condition"])
+            >>> mdata = sccoda.load(haber_cells,
+            >>>                     type="cell_level",
+            >>>                     generate_sample_level=True,
+            >>>                     cell_type_identifier="cell_label",
+            >>>                     sample_identifier="batch",
+            >>>                     covariate_obs=["condition"])
             >>> mdata = sccoda.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
             >>> sccoda.run_nuts(mdata, num_warmup=100, num_samples=1000, rng_key=42)
             >>> sccoda.set_fdr(mdata, est_fdr=0.4)

pertpy/tools/_coda/_tasccoda.py

@@ -3,18 +3,16 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Literal
 
 import arviz as az
-import ete3 as ete
 import jax.numpy as jnp
 import numpy as np
 import numpyro as npy
 import numpyro.distributions as npd
 import toytree as tt
 from anndata import AnnData
-from jax import random
-from jax.config import config
+from jax import config, random
+from lamin_utils import logger
 from mudata import MuData
 from numpyro.infer import Predictive
-from rich import print
 
 from pertpy.tools._coda._base_coda import (
     CompositionalModel2,
@@ -87,25 +85,25 @@ class Tasccoda(CompositionalModel2):
         Args:
             adata: AnnData object.
             type: Specify the input adata type, which could be either a cell-level AnnData or an aggregated sample-level AnnData.
-            cell_type_identifier: If type is "cell_level", specify column name in adata.obs that specifies the cell types. Defaults to None.
-            sample_identifier: If type is "cell_level", specify column name in adata.obs that specifies the sample. Defaults to None.
-            covariate_uns: If type is "cell_level", specify key for adata.uns, where covariate values are stored. Defaults to None.
-            covariate_obs: If type is "cell_level", specify list of keys for adata.obs, where covariate values are stored. Defaults to None.
-            covariate_df: If type is "cell_level", specify dataFrame with covariates. Defaults to None.
-            dendrogram_key: Key to the scanpy.tl.dendrogram result in `.uns` of original cell level anndata object. Defaults to None.
-            levels_orig: List that indicates which columns in `.obs` of the original data correspond to tree levels. The list must begin with the root level, and end with the leaf level. Defaults to None.
-            levels_agg: List that indicates which columns in `.var` of the aggregated data correspond to tree levels. The list must begin with the root level, and end with the leaf level. Defaults to None.
-            add_level_name: If True, internal nodes in the tree will be named as "{level_name}_{node_name}" instead of just {level_name}. Defaults to False.
-            key_added: If not specified, the tree is stored in .uns[‘tree’]. If `data` is AnnData, save tree in `data`. If `data` is MuData, save tree in data[modality_2]. Defaults to "tree".
-            modality_key_1: Key to the cell-level AnnData in the MuData object. Defaults to "rna".
-            modality_key_2: Key to the aggregated sample-level AnnData object in the MuData object. Defaults to "coda".
+            cell_type_identifier: If type is "cell_level", specify column name in adata.obs that specifies the cell types.
+            sample_identifier: If type is "cell_level", specify column name in adata.obs that specifies the sample.
+            covariate_uns: If type is "cell_level", specify key for adata.uns, where covariate values are stored.
+            covariate_obs: If type is "cell_level", specify list of keys for adata.obs, where covariate values are stored.
+            covariate_df: If type is "cell_level", specify dataFrame with covariates.
+            dendrogram_key: Key to the scanpy.tl.dendrogram result in `.uns` of original cell level anndata object.
+            levels_orig: List that indicates which columns in `.obs` of the original data correspond to tree levels. The list must begin with the root level, and end with the leaf level.
+            levels_agg: List that indicates which columns in `.var` of the aggregated data correspond to tree levels. The list must begin with the root level, and end with the leaf level.
+            add_level_name: If True, internal nodes in the tree will be named as "{level_name}_{node_name}" instead of just {level_name}.
+            key_added: If not specified, the tree is stored in .uns[‘tree’]. If `data` is AnnData, save tree in `data`. If `data` is MuData, save tree in data[modality_2].
+            modality_key_1: Key to the cell-level AnnData in the MuData object.
+            modality_key_2: Key to the aggregated sample-level AnnData object in the MuData object.
 
         Returns:
             MuData: MuData object with cell-level AnnData (`mudata[modality_key_1]`) and aggregated sample-level AnnData (`mudata[modality_key_2]`).
 
         Examples:
             >>> import pertpy as pt
-            >>> adata = pt.dt.smillie()
+            >>> adata = pt.dt.tasccoda_example()
             >>> tasccoda = pt.tl.Tasccoda()
             >>> mdata = tasccoda.load(
             >>>     adata, type="sample_level",
@@ -148,21 +146,22 @@ class Tasccoda(CompositionalModel2):
         pen_args: dict = None,
         modality_key: str = "coda",
     ) -> AnnData | MuData:
-        """Handles data preprocessing, covariate matrix creation, reference selection, and zero count replacement for tascCODA. Also sets model parameters, model type (tree_agg), effect selection type (sslaso) and performs tree processing.
+        """Handles data preprocessing, covariate matrix creation, reference selection, and zero count replacement for tascCODA.
 
         Args:
             data: Anndata object with cell counts as .X and covariates saved in .obs or a MuData object.
             formula: R-style formula for building the covariate matrix.
-                Categorical covariates are handled automatically, with the covariate value of the first sample being used as the reference category.
-                To set a different level as the base category for a categorical covariate, use "C(<CovariateName>, Treatment('<ReferenceLevelName>'))"
+                     Categorical covariates are handled automatically, with the covariate value of the first sample being used as the reference category.
+                     To set a different level as the base category for a categorical covariate, use "C(<CovariateName>, Treatment('<ReferenceLevelName>'))"
             reference_cell_type: Column name that sets the reference cell type.
-                Reference the name of a column. If "automatic", the cell type with the lowest dispersion in relative abundance that is present in at least 90% of samlpes will be chosen. Defaults to "automatic".
-            automatic_reference_absence_threshold: If using reference_cell_type = "automatic", determine the maximum fraction of zero entries for a cell type
-                to be considered as a possible reference cell type. Defaults to 0.05.
+                                 If "automatic", the cell type with the lowest dispersion in relative abundance that is present in at least 90% of samlpes will be chosen.
+            automatic_reference_absence_threshold: If using reference_cell_type = "automatic",
+                                                   determine the maximum fraction of zero entries for a cell type
+                                                   to be considered as a possible reference cell type.
             tree_key: Key in `adata.uns` that contains the tree structure
             pen_args: Dictionary with penalty arguments. With `reg="scaled_3"`, the parameters phi (aggregation bias), lambda_1, lambda_0 can be set here.
                 See the tascCODA paper for an explanation of these parameters. Default: lambda_0 = 50, lambda_1 = 5, phi = 0.
-            modality_key: If data is a MuData object, specify key to the aggregated sample-level AnnData object in the MuData object. Defaults to "coda".
+            modality_key: If data is a MuData object, specify key to the aggregated sample-level AnnData object in the MuData object.
 
         Returns:
             Return an AnnData (if input data is an AnnData object) or return a MuData (if input data is a MuData object)
@@ -175,7 +174,7 @@ class Tasccoda(CompositionalModel2):
 
         Examples:
             >>> import pertpy as pt
-            >>> adata = pt.dt.smillie()
+            >>> adata = pt.dt.tasccoda_example()
             >>> tasccoda = pt.tl.Tasccoda()
             >>> mdata = tasccoda.load(
             >>>     adata, type="sample_level",
@@ -199,8 +198,16 @@ class Tasccoda(CompositionalModel2):
         if tree_key is None:
             raise ValueError("Please specify the key in .uns that contains the tree structure!")
 
+        # Scoped import due to installation issues
+        try:
+            import ete3 as ete
+        except ImportError:
+            raise ImportError(
+                "To use tasccoda please install additional dependencies as `pip install pertpy[coda]`"
+            ) from None
+
         # toytree tree - only for legacy reasons, can be removed in the final version
-        if isinstance(adata.uns[tree_key], tt.tree):
+        if isinstance(adata.uns[tree_key], tt.core.ToyTree):
             # Collapse singularities in the tree
             phy_tree = collapse_singularities(adata.uns[tree_key])
 
@@ -315,7 +322,7 @@ class Tasccoda(CompositionalModel2):
 
         Examples:
             >>> import pertpy as pt
-            >>> adata = pt.dt.smillie()
+            >>> adata = pt.dt.tasccoda_example()
             >>> tasccoda = pt.tl.Tasccoda()
             >>> mdata = tasccoda.load(
             >>>     adata, type="sample_level",
@@ -325,7 +332,7 @@ class Tasccoda(CompositionalModel2):
             >>> mdata = tasccoda.prepare(
             >>>     mdata, formula="Health", reference_cell_type="automatic", tree_key="lineage", pen_args={"phi": 0}
             >>> )
-            >>> adata = tasccoda.set_init_mcmc_states(rng_key=42, ref_index=[0,1], sample_adata=mdata['coda'])
+            >>> adata = tasccoda.set_init_mcmc_states(rng_key=42, ref_index=[0, 1], sample_adata=mdata["coda"])
         """
         N, D = sample_adata.obsm["covariate_matrix"].shape
         P = sample_adata.X.shape[1]
@@ -469,17 +476,17 @@ class Tasccoda(CompositionalModel2):
 
         Args:
             data: AnnData object or MuData object.
-            modality_key: If data is a MuData object, specify which modality to use. Defaults to "coda".
-            rng_key: The rng state used for the prior simulation. If None, a random state will be selected. Defaults to None.
-            num_prior_samples: Number of prior samples calculated. Defaults to 500.
-            use_posterior_predictive: If True, the posterior predictive will be calculated. Defaults to True.
+            modality_key: If data is a MuData object, specify which modality to use.
+            rng_key: The rng state used for the prior simulation. If None, a random state will be selected.
+            num_prior_samples: Number of prior samples calculated.
+            use_posterior_predictive: If True, the posterior predictive will be calculated.
 
         Returns:
             arviz.InferenceData: arviz_data
 
         Examples:
             >>> import pertpy as pt
-            >>> adata = pt.dt.smillie()
+            >>> adata = pt.dt.tasccoda_example()
             >>> tasccoda = pt.tl.Tasccoda()
             >>> mdata = tasccoda.load(
             >>>     adata, type="sample_level",
@@ -496,7 +503,7 @@ class Tasccoda(CompositionalModel2):
             try:
                 sample_adata = data[modality_key]
             except IndexError:
-                print("When data is a MuData object, modality_key must be specified!")
+                logger.error("When data is a MuData object, modality_key must be specified!")
                 raise
         if isinstance(data, AnnData):
             sample_adata = data
@@ -543,7 +550,7 @@ class Tasccoda(CompositionalModel2):
 
         if rng_key is None:
             rng = np.random.default_rng()
-            rng_key = random.PRNGKey(rng.integers(0, 10000))
+            rng_key = random.key(rng.integers(0, 10000))
 
         if use_posterior_predictive:
             posterior_predictive = Predictive(self.model, self.mcmc.get_samples())(
@@ -590,7 +597,7 @@ class Tasccoda(CompositionalModel2):
         """
         Examples:
             >>> import pertpy as pt
-            >>> adata = pt.dt.smillie()
+            >>> adata = pt.dt.tasccoda_example()
             >>> tasccoda = pt.tl.Tasccoda()
             >>> mdata = tasccoda.load(
             >>>     adata, type="sample_level",
@@ -610,7 +617,7 @@ class Tasccoda(CompositionalModel2):
         """
         Examples:
             >>> import pertpy as pt
-            >>> adata = pt.dt.smillie()
+            >>> adata = pt.dt.tasccoda_example()
             >>> tasccoda = pt.tl.Tasccoda()
             >>> mdata = tasccoda.load(
             >>>     adata, type="sample_level",
@@ -631,7 +638,7 @@ class Tasccoda(CompositionalModel2):
         """
         Examples:
             >>> import pertpy as pt
-            >>> adata = pt.dt.smillie()
+            >>> adata = pt.dt.tasccoda_example()
             >>> tasccoda = pt.tl.Tasccoda()
             >>> mdata = tasccoda.load(
             >>>     adata, type="sample_level",
@@ -652,7 +659,7 @@ class Tasccoda(CompositionalModel2):
         """
         Examples:
             >>> import pertpy as pt
-            >>> adata = pt.dt.smillie()
+            >>> adata = pt.dt.tasccoda_example()
             >>> tasccoda = pt.tl.Tasccoda()
             >>> mdata = tasccoda.load(
             >>>     adata, type="sample_level",