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

pertpy/tools/_coda/_sccoda.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Literal
 
-import arviz as az
 import jax.numpy as jnp
 import numpy as np
 import numpyro as npy
@@ -17,14 +16,15 @@ from rich import print
 from pertpy.tools._coda._base_coda import CompositionalModel2, from_scanpy
 
 if TYPE_CHECKING:
+    import arviz as az
     import pandas as pd
 
 config.update("jax_enable_x64", True)
 
 
 class Sccoda(CompositionalModel2):
-    """
-    Statistical model for single-cell differential composition analysis with specification of a reference cell type.
+    r"""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.
 
     The hierarchical formulation of the model for one sample is:
@@ -45,7 +45,6 @@ class Sccoda(CompositionalModel2):
 
     For further information, see `scCODA is a Bayesian model for compositional single-cell data analysis`
     (Büttner, Ostner et al., NatComms, 2021)
-
     """
 
     def __init__(self, *args, **kwargs):
@@ -83,7 +82,7 @@ class Sccoda(CompositionalModel2):
             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]`).
+            :class:`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
@@ -108,8 +107,12 @@ class Sccoda(CompositionalModel2):
             else:
                 adata_coda = AnnData()
             mdata = MuData({modality_key_1: adata, modality_key_2: adata_coda})
+            mdata.pull_obs()
+            mdata.pull_var()
         else:
             mdata = MuData({modality_key_1: AnnData(), modality_key_2: adata})
+            mdata.pull_obs()
+            mdata.pull_var()
         return mdata
 
     def prepare(
@@ -184,8 +187,7 @@ class Sccoda(CompositionalModel2):
             return adata
 
     def set_init_mcmc_states(self, rng_key: None, ref_index: np.ndarray, sample_adata: AnnData) -> AnnData:  # type: ignore
-        """
-        Sets initial MCMC state values for scCODA model
+        """Sets initial MCMC state values for scCODA model.
 
         Args:
             rng_key: RNG value to be set
@@ -237,8 +239,7 @@ class Sccoda(CompositionalModel2):
         ref_index,
         sample_adata: AnnData,
     ):
-        """
-        Implements scCODA model in numpyro
+        """Implements scCODA model in numpyro.
 
         Args:
             counts: Count data array
@@ -306,7 +307,7 @@ class Sccoda(CompositionalModel2):
         num_prior_samples: int = 500,
         use_posterior_predictive: bool = True,
     ) -> az.InferenceData:
-        """Creates arviz object from model results for MCMC diagnosis
+        """Creates arviz object from model results for MCMC diagnosis.
 
         Args:
             data: AnnData object or MuData object.
@@ -316,7 +317,7 @@ class Sccoda(CompositionalModel2):
             use_posterior_predictive: If True, the posterior predictive will be calculated.
 
         Returns:
-            az.InferenceData: arviz_data with all MCMC information
+            :class:`arviz.InferenceData`: arviz_data with all MCMC information
 
         Examples:
             >>> import pertpy as pt
@@ -405,6 +406,8 @@ class Sccoda(CompositionalModel2):
         else:
             prior = None
 
+        import arviz as az
+
         # Create arviz object
         arviz_data = az.from_numpyro(
             self.mcmc, prior=prior, posterior_predictive=posterior_predictive, dims=dims, coords=coords
@@ -423,80 +426,76 @@ class Sccoda(CompositionalModel2):
         *args,
         **kwargs,
     ):
-        """
-        Examples:
-            >>> 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.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
-            >>> sccoda.run_nuts(mdata, num_warmup=100, num_samples=1000, rng_key=42)
-        """
+        """Examples:
+        >>> 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.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
+        >>> sccoda.run_nuts(mdata, num_warmup=100, num_samples=1000, rng_key=42).
+        """  # noqa: D205
         return super().run_nuts(data, modality_key, num_samples, num_warmup, rng_key, copy, *args, **kwargs)
 
     run_nuts.__doc__ = CompositionalModel2.run_nuts.__doc__ + run_nuts.__doc__
 
     def credible_effects(self, data: AnnData | MuData, modality_key: str = "coda", est_fdr: float = None) -> pd.Series:
-        """
-        Examples:
-            >>> 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.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)
-        """
+        """Examples:
+        >>> 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.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).
+        """  # noqa: D205
         return super().credible_effects(data, modality_key, est_fdr)
 
     credible_effects.__doc__ = CompositionalModel2.credible_effects.__doc__ + credible_effects.__doc__
 
     def summary(self, data: AnnData | MuData, extended: bool = False, modality_key: str = "coda", *args, **kwargs):
-        """
-        Examples:
-            >>> 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.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
-            >>> sccoda.run_nuts(mdata, num_warmup=100, num_samples=1000, rng_key=42)
-            >>> sccoda.summary(mdata)
-        """
+        """Examples:
+        >>> 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.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
+        >>> sccoda.run_nuts(mdata, num_warmup=100, num_samples=1000, rng_key=42)
+        >>> sccoda.summary(mdata).
+        """  # noqa: D205
         return super().summary(data, extended, modality_key, *args, **kwargs)
 
     summary.__doc__ = CompositionalModel2.summary.__doc__ + summary.__doc__
 
     def set_fdr(self, data: AnnData | MuData, est_fdr: float, modality_key: str = "coda", *args, **kwargs):
-        """
-        Examples:
-            >>> 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.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)
-        """
+        """Examples:
+        >>> 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.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).
+        """  # noqa: D205
         return super().set_fdr(data, est_fdr, modality_key, *args, **kwargs)
 
     set_fdr.__doc__ = CompositionalModel2.set_fdr.__doc__ + set_fdr.__doc__

pertpy/tools/_coda/_tasccoda.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Literal
 
-import arviz as az
 import jax.numpy as jnp
 import numpy as np
 import numpyro as npy
@@ -25,14 +24,14 @@ from pertpy.tools._coda._base_coda import (
 )
 
 if TYPE_CHECKING:
+    import arviz as az
     import pandas as pd
 
 config.update("jax_enable_x64", True)
 
 
 class Tasccoda(CompositionalModel2):
-    """
-    Statistical model for tree-aggregated differential composition analysis (tascCODA, Ostner et al., 2021).
+    r"""Statistical model for tree-aggregated differential composition analysis (tascCODA, Ostner et al., 2021).
 
     The hierarchical formulation of the model for one sample is:
 
@@ -99,7 +98,7 @@ class Tasccoda(CompositionalModel2):
             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]`).
+            :class:`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
@@ -111,7 +110,6 @@ class Tasccoda(CompositionalModel2):
             >>>     key_added="lineage", add_level_name=True
             >>> )
         """
-
         if type == "cell_level":
             adata_coda = from_scanpy(
                 adata=adata,
@@ -198,9 +196,8 @@ 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
+            import ete4 as ete
         except ImportError:
             raise ImportError(
                 "To use tasccoda please install additional dependencies as `pip install pertpy[coda]`"
@@ -236,7 +233,7 @@ class Tasccoda(CompositionalModel2):
 
             # number of leaves for each internal node (important for aggregation penalty lambda_1)
             if "node_leaves" not in pen_args:
-                node_leaves = [len(n.get_leaves()) for n in phy_tree.idx_dict.values()]
+                node_leaves = [len(n.leaves()) for n in phy_tree.idx_dict.values()]
                 node_leaves.reverse()
                 pen_args["node_leaves"] = np.delete(np.array(node_leaves[:-1]), refs)
 
@@ -245,7 +242,7 @@ class Tasccoda(CompositionalModel2):
             # Collapse singularities in the tree
             phy_tree = collapse_singularities_2(adata.uns[tree_key])
 
-            node_names = [n.name for n in phy_tree.iter_descendants()]
+            node_names = [n.name for n in phy_tree.descendants()]
 
             # Get ancestor matrix
             A, T = get_a_2(phy_tree, leaf_order=adata.var.index.tolist(), node_order=node_names)
@@ -256,7 +253,7 @@ class Tasccoda(CompositionalModel2):
             # Ancestors of reference are a reference, too!
             # Get names of reference nodes
             reference_cell_type = adata.uns["scCODA_params"]["reference_cell_type"]
-            ref_nodes = [n.name for n in phy_tree.search_nodes(name=reference_cell_type)[0].get_ancestors()[:-1]]
+            ref_nodes = [n.name for n in list(next(phy_tree.search_nodes(name=reference_cell_type)).ancestors())[:-1]]
             ref_nodes = [reference_cell_type] + ref_nodes
             adata.uns["scCODA_params"]["reference_nodes"] = ref_nodes
 
@@ -267,12 +264,12 @@ class Tasccoda(CompositionalModel2):
 
             # number of leaves for each internal node (important for aggregation penalty lambda_1)
             if "node_leaves" not in pen_args:
-                node_leaves = [len(n.get_leaves()) for n in phy_tree.iter_descendants()]
+                node_leaves = [len(list(n.leaves())) for n in phy_tree.descendants()]
                 pen_args["node_leaves"] = np.delete(np.array(node_leaves), ref_idxs)
 
         # No valid tree structure
         else:
-            raise ValueError("Tree structure is not a toytree or ete3 tree object")
+            raise ValueError("Tree structure is not a toytree or ete4 tree object")
 
         # Default spike-and-slab LASSO parameters
         if "lambda_0" not in pen_args:
@@ -309,8 +306,7 @@ class Tasccoda(CompositionalModel2):
             return adata
 
     def set_init_mcmc_states(self, rng_key: None, ref_index: np.ndarray, sample_adata: AnnData) -> AnnData:  # type: ignore
-        """
-        Sets initial MCMC state values for tascCODA model
+        """Sets initial MCMC state values for tascCODA model.
 
         Args:
             rng_key: RNG value to be set
@@ -387,7 +383,7 @@ class Tasccoda(CompositionalModel2):
         ref_index: np.ndarray,
         sample_adata: AnnData,
     ):
-        """Implements tascCODA model in numpyro
+        """Implements tascCODA model in numpyro.
 
         Args:
             counts: Count data array
@@ -472,7 +468,7 @@ class Tasccoda(CompositionalModel2):
         num_prior_samples: int = 500,
         use_posterior_predictive: bool = True,
     ) -> az.InferenceData:
-        """Creates arviz object from model results for MCMC diagnosis
+        """Creates arviz object from model results for MCMC diagnosis.
 
         Args:
             data: AnnData object or MuData object.
@@ -482,7 +478,7 @@ class Tasccoda(CompositionalModel2):
             use_posterior_predictive: If True, the posterior predictive will be calculated.
 
         Returns:
-            arviz.InferenceData: arviz_data
+            :class:`arviz.InferenceData`: arviz_data
 
         Examples:
             >>> import pertpy as pt
@@ -576,6 +572,8 @@ class Tasccoda(CompositionalModel2):
         else:
             prior = None
 
+        import arviz as az
+
         # Create arviz object
         arviz_data = az.from_numpyro(
             self.mcmc, prior=prior, posterior_predictive=posterior_predictive, dims=dims, coords=coords
@@ -594,84 +592,80 @@ class Tasccoda(CompositionalModel2):
         *args,
         **kwargs,
     ):
-        """
-        Examples:
-            >>> import pertpy as pt
-            >>> adata = pt.dt.tasccoda_example()
-            >>> tasccoda = pt.tl.Tasccoda()
-            >>> mdata = tasccoda.load(
-            >>>     adata, type="sample_level",
-            >>>     levels_agg=["Major_l1", "Major_l2", "Major_l3", "Major_l4", "Cluster"],
-            >>>     key_added="lineage", add_level_name=True
-            >>> )
-            >>> mdata = tasccoda.prepare(
-            >>>     mdata, formula="Health", reference_cell_type="automatic", tree_key="lineage", pen_args={"phi": 0}
-            >>> )
-            >>> tasccoda.run_nuts(mdata, num_samples=1000, num_warmup=100, rng_key=42)
-        """
+        """Examples:
+        >>> import pertpy as pt
+        >>> adata = pt.dt.tasccoda_example()
+        >>> tasccoda = pt.tl.Tasccoda()
+        >>> mdata = tasccoda.load(
+        >>>     adata, type="sample_level",
+        >>>     levels_agg=["Major_l1", "Major_l2", "Major_l3", "Major_l4", "Cluster"],
+        >>>     key_added="lineage", add_level_name=True
+        >>> )
+        >>> mdata = tasccoda.prepare(
+        >>>     mdata, formula="Health", reference_cell_type="automatic", tree_key="lineage", pen_args={"phi": 0}
+        >>> )
+        >>> tasccoda.run_nuts(mdata, num_samples=1000, num_warmup=100, rng_key=42).
+        """  # noqa: D205
         return super().run_nuts(data, modality_key, num_samples, num_warmup, rng_key, copy, *args, **kwargs)
 
     run_nuts.__doc__ = CompositionalModel2.run_nuts.__doc__ + run_nuts.__doc__
 
     def summary(self, data: AnnData | MuData, extended: bool = False, modality_key: str = "coda", *args, **kwargs):
-        """
-        Examples:
-            >>> import pertpy as pt
-            >>> adata = pt.dt.tasccoda_example()
-            >>> tasccoda = pt.tl.Tasccoda()
-            >>> mdata = tasccoda.load(
-            >>>     adata, type="sample_level",
-            >>>     levels_agg=["Major_l1", "Major_l2", "Major_l3", "Major_l4", "Cluster"],
-            >>>     key_added="lineage", add_level_name=True
-            >>> )
-            >>> mdata = tasccoda.prepare(
-            >>>     mdata, formula="Health", reference_cell_type="automatic", tree_key="lineage", pen_args={"phi": 0}
-            >>> )
-            >>> tasccoda.run_nuts(mdata, num_samples=1000, num_warmup=100, rng_key=42)
-            >>> tasccoda.summary(mdata)
-        """
+        """Examples:
+        >>> import pertpy as pt
+        >>> adata = pt.dt.tasccoda_example()
+        >>> tasccoda = pt.tl.Tasccoda()
+        >>> mdata = tasccoda.load(
+        >>>     adata, type="sample_level",
+        >>>     levels_agg=["Major_l1", "Major_l2", "Major_l3", "Major_l4", "Cluster"],
+        >>>     key_added="lineage", add_level_name=True
+        >>> )
+        >>> mdata = tasccoda.prepare(
+        >>>     mdata, formula="Health", reference_cell_type="automatic", tree_key="lineage", pen_args={"phi": 0}
+        >>> )
+        >>> tasccoda.run_nuts(mdata, num_samples=1000, num_warmup=100, rng_key=42)
+        >>> tasccoda.summary(mdata).
+        """  # noqa: D205
         return super().summary(data, extended, modality_key, *args, **kwargs)
 
     summary.__doc__ = CompositionalModel2.summary.__doc__ + summary.__doc__
 
     def credible_effects(self, data: AnnData | MuData, modality_key: str = "coda", est_fdr: float = None) -> pd.Series:
-        """
-        Examples:
-            >>> import pertpy as pt
-            >>> adata = pt.dt.tasccoda_example()
-            >>> tasccoda = pt.tl.Tasccoda()
-            >>> mdata = tasccoda.load(
-            >>>     adata, type="sample_level",
-            >>>     levels_agg=["Major_l1", "Major_l2", "Major_l3", "Major_l4", "Cluster"],
-            >>>     key_added="lineage", add_level_name=True
-            >>> )
-            >>> mdata = tasccoda.prepare(
-            >>>     mdata, formula="Health", reference_cell_type="automatic", tree_key="lineage", pen_args={"phi": 0}
-            >>> )
-            >>> tasccoda.run_nuts(mdata, num_samples=1000, num_warmup=100, rng_key=42)
-            >>> tasccoda.credible_effects(mdata)
-        """
+        """Examples:
+        >>> import pertpy as pt
+        >>> adata = pt.dt.tasccoda_example()
+        >>> tasccoda = pt.tl.Tasccoda()
+        >>> mdata = tasccoda.load(
+        >>>     adata, type="sample_level",
+        >>>     levels_agg=["Major_l1", "Major_l2", "Major_l3", "Major_l4", "Cluster"],
+        >>>     key_added="lineage", add_level_name=True
+        >>> )
+        >>> mdata = tasccoda.prepare(
+        >>>     mdata, formula="Health", reference_cell_type="automatic", tree_key="lineage", pen_args={"phi": 0}
+        >>> )
+        >>> tasccoda.run_nuts(mdata, num_samples=1000, num_warmup=100, rng_key=42)
+        >>> tasccoda.credible_effects(mdata).
+        """  # noqa: D205
         return super().credible_effects(data, modality_key, est_fdr)
 
     credible_effects.__doc__ = CompositionalModel2.credible_effects.__doc__ + credible_effects.__doc__
 
     def set_fdr(self, data: AnnData | MuData, est_fdr: float, modality_key: str = "coda", *args, **kwargs):
-        """
-        Examples:
-            >>> import pertpy as pt
-            >>> adata = pt.dt.tasccoda_example()
-            >>> tasccoda = pt.tl.Tasccoda()
-            >>> mdata = tasccoda.load(
-            >>>     adata, type="sample_level",
-            >>>     levels_agg=["Major_l1", "Major_l2", "Major_l3", "Major_l4", "Cluster"],
-            >>>     key_added="lineage", add_level_name=True
-            >>> )
-            >>> mdata = tasccoda.prepare(
-            >>>     mdata, formula="Health", reference_cell_type="automatic", tree_key="lineage", pen_args={"phi": 0}
-            >>> )
-            >>> tasccoda.run_nuts(mdata, num_samples=1000, num_warmup=100, rng_key=42)
-            >>> tasccoda.set_fdr(mdata, est_fdr=0.4)
-        """
+        """Examples:
+        >>> import pertpy as pt
+        >>> adata = pt.dt.tasccoda_example()
+        >>> tasccoda = pt.tl.Tasccoda()
+        >>> mdata = tasccoda.load(
+        >>>     adata, type="sample_level",
+        >>>     levels_agg=["Major_l1", "Major_l2", "Major_l3", "Major_l4", "Cluster"],
+        >>>     key_added="lineage", add_level_name=True
+        >>> )
+        >>> mdata = tasccoda.prepare(
+        >>>     mdata, formula="Health", reference_cell_type="automatic", tree_key="lineage", pen_args={"phi": 0}
+        >>> )
+        >>> tasccoda.run_nuts(mdata, num_samples=1000, num_warmup=100, rng_key=42)
+        >>> tasccoda.set_fdr(mdata, est_fdr=0.4).
+        """  # noqa: D205
         return super().set_fdr(data, est_fdr, modality_key, *args, **kwargs)
 
     set_fdr.__doc__ = CompositionalModel2.set_fdr.__doc__ + set_fdr.__doc__