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

pertpy/tools/_scgen/_jax_scgen.py

@@ -1,370 +0,0 @@
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, Any
-
-import jax.numpy as jnp
-import numpy as np
-from anndata import AnnData
-from jax import Array
-from scvi import REGISTRY_KEYS
-from scvi.data import AnnDataManager
-from scvi.data.fields import CategoricalObsField, LayerField
-from scvi.model.base import BaseModelClass, JaxTrainingMixin
-from scvi.utils import setup_anndata_dsp
-
-from ._jax_scgenvae import JaxSCGENVAE
-from ._utils import balancer, extractor
-
-if TYPE_CHECKING:
-    from collections.abc import Sequence
-
-font = {"family": "Arial", "size": 14}
-
-
-class SCGEN(JaxTrainingMixin, BaseModelClass):
-    """Jax Implementation of scGen model for batch removal and perturbation prediction."""
-
-    def __init__(
-        self,
-        adata: AnnData,
-        n_hidden: int = 800,
-        n_latent: int = 100,
-        n_layers: int = 2,
-        dropout_rate: float = 0.2,
-        **model_kwargs,
-    ):
-        super().__init__(adata)
-
-        self.module = JaxSCGENVAE(
-            n_input=self.summary_stats.n_vars,
-            n_hidden=n_hidden,
-            n_latent=n_latent,
-            n_layers=n_layers,
-            dropout_rate=dropout_rate,
-            **model_kwargs,
-        )
-        self._model_summary_string = (
-            "SCGEN Model with the following params: \nn_hidden: {}, n_latent: {}, n_layers: {}, dropout_rate: " "{}"
-        ).format(
-            n_hidden,
-            n_latent,
-            n_layers,
-            dropout_rate,
-        )
-        self.init_params_ = self._get_init_params(locals())
-
-    def predict(
-        self,
-        ctrl_key=None,
-        stim_key=None,
-        adata_to_predict=None,
-        celltype_to_predict=None,
-        restrict_arithmetic_to="all",
-    ) -> tuple[AnnData, Any]:
-        """Predicts the cell type provided by the user in stimulated condition.
-
-        Args:
-            ctrl_key: Key for `control` part of the `data` found in `condition_key`.
-            stim_key: Key for `stimulated` part of the `data` found in `condition_key`.
-            adata_to_predict: Adata for unperturbed cells you want to be predicted.
-            celltype_to_predict: The cell type you want to be predicted.
-            restrict_arithmetic_to: Dictionary of celltypes you want to be observed for prediction.
-
-        Returns:
-            `np nd-array` of predicted cells in primary space.
-        delta: float
-            Difference between stimulated and control cells in latent space
-
-        Examples:
-            >>> import pertpy as pt
-            >>> data = pt.dt.kang_2018()
-            >>> pt.tl.SCGEN.setup_anndata(data, batch_key="label", labels_key="cell_type")
-            >>> model = pt.tl.SCGEN(data)
-            >>> model.train(max_epochs=10, batch_size=64, early_stopping=True, early_stopping_patience=5)
-            >>> pred, delta = model.predict(ctrl_key='ctrl', stim_key='stim', celltype_to_predict='CD4 T cells')
-        """
-        # use keys registered from `setup_anndata()`
-        cell_type_key = self.adata_manager.get_state_registry(REGISTRY_KEYS.LABELS_KEY).original_key
-        condition_key = self.adata_manager.get_state_registry(REGISTRY_KEYS.BATCH_KEY).original_key
-
-        if restrict_arithmetic_to == "all":
-            ctrl_x = self.adata[self.adata.obs[condition_key] == ctrl_key, :]
-            stim_x = self.adata[self.adata.obs[condition_key] == stim_key, :]
-            ctrl_x = balancer(ctrl_x, cell_type_key)
-            stim_x = balancer(stim_x, cell_type_key)
-        else:
-            key = list(restrict_arithmetic_to.keys())[0]
-            values = restrict_arithmetic_to[key]
-            subset = self.adata[self.adata.obs[key].isin(values)]
-            ctrl_x = subset[subset.obs[condition_key] == ctrl_key, :]
-            stim_x = subset[subset.obs[condition_key] == stim_key, :]
-            if len(values) > 1:
-                ctrl_x = balancer(ctrl_x, cell_type_key)
-                stim_x = balancer(stim_x, cell_type_key)
-        if celltype_to_predict is not None and adata_to_predict is not None:
-            raise Exception("Please provide either a cell type or adata not both!")
-        if celltype_to_predict is None and adata_to_predict is None:
-            raise Exception("Please provide a cell type name or adata for your unperturbed cells")
-        if celltype_to_predict is not None:
-            ctrl_pred = extractor(
-                self.adata,
-                celltype_to_predict,
-                condition_key,
-                cell_type_key,
-                ctrl_key,
-                stim_key,
-            )[1]
-        else:
-            ctrl_pred = adata_to_predict
-
-        eq = min(ctrl_x.X.shape[0], stim_x.X.shape[0])
-        rng = np.random.default_rng()
-        cd_ind = rng.choice(range(ctrl_x.shape[0]), size=eq, replace=False)
-        stim_ind = rng.choice(range(stim_x.shape[0]), size=eq, replace=False)
-        ctrl_adata = ctrl_x[cd_ind, :]
-        stim_adata = stim_x[stim_ind, :]
-
-        latent_ctrl = self._avg_vector(ctrl_adata)
-        latent_stim = self._avg_vector(stim_adata)
-
-        delta = latent_stim - latent_ctrl
-
-        latent_cd = self.get_latent_representation(ctrl_pred)
-
-        stim_pred = delta + latent_cd
-        predicted_cells = self.module.as_bound().generative(stim_pred)["px"]
-
-        predicted_adata = AnnData(
-            X=np.array(predicted_cells),
-            obs=ctrl_pred.obs.copy(),
-            var=ctrl_pred.var.copy(),
-            obsm=ctrl_pred.obsm.copy(),
-        )
-        return predicted_adata, delta
-
-    def _avg_vector(self, adata):
-        return np.mean(self.get_latent_representation(adata), axis=0)
-
-    def get_decoded_expression(
-        self,
-        adata: AnnData | None = None,
-        indices: Sequence[int] | None = None,
-        batch_size: int | None = None,
-    ) -> Array:
-        """Get decoded expression.
-
-        Args:
-            adata: AnnData object with equivalent structure to initial AnnData. If `None`, defaults to the
-                   AnnData object used to initialize the model.
-            indices: Indices of cells in adata to use. If `None`, all cells are used.
-            batch_size: Minibatch size for data loading into model. Defaults to `scvi.settings.batch_size`.
-
-        Returns:
-            Decoded expression for each cell
-
-        Examples:
-            >>> import pertpy as pt
-            >>> data = pt.dt.kang_2018()
-            >>> pt.tl.SCGEN.setup_anndata(data, batch_key="label", labels_key="cell_type")
-            >>> model = pt.tl.SCGEN(data)
-            >>> model.train(max_epochs=10, batch_size=64, early_stopping=True, early_stopping_patience=5)
-            >>> decoded_X = model.get_decoded_expression()
-        """
-        if self.is_trained_ is False:
-            raise RuntimeError("Please train the model first.")
-
-        adata = self._validate_anndata(adata)
-        scdl = self._make_data_loader(adata=adata, indices=indices, batch_size=batch_size)
-        decoded = []
-        for tensors in scdl:
-            _, generative_outputs = self.module.as_bound()(tensors, compute_loss=False)
-            px = generative_outputs["px"]
-            decoded.append(px)
-
-        return jnp.concatenate(decoded)
-
-    def batch_removal(self, adata: AnnData | None = None) -> AnnData:
-        """Removes batch effects.
-
-        Args:
-            adata: AnnData object with equivalent structure to initial AnnData. If `None`, defaults to the
-                   AnnData object used to initialize the model. Must have been setup with `batch_key` and `labels_key`,
-                   corresponding to batch and cell type metadata, respectively.
-
-        Returns:
-            corrected: `~anndata.AnnData`
-            AnnData of corrected gene expression in adata.X and corrected latent space in adata.obsm["latent"].
-            A reference to the original AnnData is in `corrected.raw` if the input adata had no `raw` attribute.
-
-        Examples:
-            >>> import pertpy as pt
-            >>> data = pt.dt.kang_2018()
-            >>> pt.tl.SCGEN.setup_anndata(data, batch_key="label", labels_key="cell_type")
-            >>> model = pt.tl.SCGEN(data)
-            >>> model.train(max_epochs=10, batch_size=64, early_stopping=True, early_stopping_patience=5)
-            >>> corrected_adata = model.batch_removal()
-        """
-        adata = self._validate_anndata(adata)
-        latent_all = self.get_latent_representation(adata)
-        # use keys registered from `setup_anndata()`
-        cell_label_key = self.adata_manager.get_state_registry(REGISTRY_KEYS.LABELS_KEY).original_key
-        batch_key = self.adata_manager.get_state_registry(REGISTRY_KEYS.BATCH_KEY).original_key
-
-        adata_latent = AnnData(latent_all)
-        adata_latent.obs = adata.obs.copy(deep=True)
-        unique_cell_types = np.unique(adata_latent.obs[cell_label_key])
-        shared_ct = []
-        not_shared_ct = []
-        for cell_type in unique_cell_types:
-            temp_cell = adata_latent[adata_latent.obs[cell_label_key] == cell_type].copy()
-            if len(np.unique(temp_cell.obs[batch_key])) < 2:
-                cell_type_ann = adata_latent[adata_latent.obs[cell_label_key] == cell_type]
-                not_shared_ct.append(cell_type_ann)
-                continue
-            temp_cell = adata_latent[adata_latent.obs[cell_label_key] == cell_type].copy()
-            batch_list = {}
-            batch_ind = {}
-            max_batch = 0
-            max_batch_ind = ""
-            batches = np.unique(temp_cell.obs[batch_key])
-            for i in batches:
-                temp = temp_cell[temp_cell.obs[batch_key] == i]
-                temp_ind = temp_cell.obs[batch_key] == i
-                if max_batch < len(temp):
-                    max_batch = len(temp)
-                    max_batch_ind = i
-                batch_list[i] = temp
-                batch_ind[i] = temp_ind
-            max_batch_ann = batch_list[max_batch_ind]
-            for study in batch_list:
-                delta = np.average(max_batch_ann.X, axis=0) - np.average(batch_list[study].X, axis=0)
-                batch_list[study].X = delta + batch_list[study].X
-                temp_cell[batch_ind[study]].X = batch_list[study].X
-            shared_ct.append(temp_cell)
-
-        all_shared_ann = AnnData.concatenate(*shared_ct, batch_key="concat_batch", index_unique=None)
-        if "concat_batch" in all_shared_ann.obs.columns:
-            del all_shared_ann.obs["concat_batch"]
-        if len(not_shared_ct) < 1:
-            corrected = AnnData(
-                np.array(self.module.as_bound().generative(all_shared_ann.X)["px"]),
-                obs=all_shared_ann.obs,
-            )
-            corrected.var_names = adata.var_names.tolist()
-            corrected = corrected[adata.obs_names]
-            if adata.raw is not None:
-                adata_raw = AnnData(X=adata.raw.X, var=adata.raw.var)
-                adata_raw.obs_names = adata.obs_names
-                corrected.raw = adata_raw
-            corrected.obsm["latent"] = all_shared_ann.X
-            corrected.obsm["corrected_latent"] = self.get_latent_representation(corrected)
-            return corrected
-        else:
-            all_not_shared_ann = AnnData.concatenate(*not_shared_ct, batch_key="concat_batch", index_unique=None)
-            all_corrected_data = AnnData.concatenate(
-                all_shared_ann,
-                all_not_shared_ann,
-                batch_key="concat_batch",
-                index_unique=None,
-            )
-            if "concat_batch" in all_shared_ann.obs.columns:
-                del all_corrected_data.obs["concat_batch"]
-            corrected = AnnData(
-                np.array(self.module.as_bound().generative(all_corrected_data.X)["px"]),
-                obs=all_corrected_data.obs,
-            )
-            corrected.var_names = adata.var_names.tolist()
-            corrected = corrected[adata.obs_names]
-            if adata.raw is not None:
-                adata_raw = AnnData(X=adata.raw.X, var=adata.raw.var)
-                adata_raw.obs_names = adata.obs_names
-                corrected.raw = adata_raw
-            corrected.obsm["latent"] = all_corrected_data.X
-            corrected.obsm["corrected_latent"] = self.get_latent_representation(corrected)
-
-            return corrected
-
-    @classmethod
-    @setup_anndata_dsp.dedent
-    def setup_anndata(
-        cls,
-        adata: AnnData,
-        batch_key: str | None = None,
-        labels_key: str | None = None,
-        **kwargs,
-    ):
-        """%(summary)s.
-
-        scGen expects the expression data to come from `adata.X`
-
-        %(param_batch_key)s
-        %(param_labels_key)s
-
-        Examples:
-            >>> import pertpy as pt
-            >>> data = pt.dt.kang_2018()
-            >>> pt.tl.SCGEN.setup_anndata(data, batch_key="label", labels_key="cell_type")
-        """
-        setup_method_args = cls._get_setup_method_args(**locals())
-        anndata_fields = [
-            LayerField(REGISTRY_KEYS.X_KEY, None, is_count_data=False),
-            CategoricalObsField(REGISTRY_KEYS.BATCH_KEY, batch_key),
-            CategoricalObsField(REGISTRY_KEYS.LABELS_KEY, labels_key),
-        ]
-        adata_manager = AnnDataManager(fields=anndata_fields, setup_method_args=setup_method_args)
-        adata_manager.register_fields(adata, **kwargs)
-        cls.register_manager(adata_manager)
-
-    def to_device(self, device):
-        pass
-
-    @property
-    def device(self):
-        return self.module.device
-
-    def get_latent_representation(
-        self,
-        adata: AnnData | None = None,
-        indices: Sequence[int] | None = None,
-        give_mean: bool = True,
-        n_samples: int = 1,
-        batch_size: int | None = None,
-    ) -> np.ndarray:
-        """Return the latent representation for each cell.
-
-        Args:
-            adata: AnnData object with equivalent structure to initial AnnData. If `None`, defaults to the
-                   AnnData object used to initialize the model.
-            indices: Indices of cells in adata to use. If `None`, all cells are used.
-            batch_size: Minibatch size for data loading into model. Defaults to `scvi.settings.batch_size`.
-
-        Returns:
-            Low-dimensional representation for each cell
-
-        Examples:
-            >>> import pertpy as pt
-            >>> data = pt.dt.kang_2018()
-            >>> pt.tl.SCGEN.setup_anndata(data, batch_key="label", labels_key="cell_type")
-            >>> model = pt.tl.SCGEN(data)
-            >>> model.train(max_epochs=10, batch_size=64, early_stopping=True, early_stopping_patience=5)
-            >>> latent_X = model.get_latent_representation()
-        """
-        self._check_if_trained(warn=False)
-
-        adata = self._validate_anndata(adata)
-        scdl = self._make_data_loader(adata=adata, indices=indices, batch_size=batch_size, iter_ndarray=True)
-
-        jit_inference_fn = self.module.get_jit_inference_fn(inference_kwargs={"n_samples": n_samples})
-
-        latent = []
-        for array_dict in scdl:
-            out = jit_inference_fn(self.module.rngs, array_dict)
-            if give_mean:
-                z = out["qz"].mean
-            else:
-                z = out["z"]
-            latent.append(z)
-        concat_axis = 0 if ((n_samples == 1) or give_mean) else 1
-        latent = jnp.concatenate(latent, axis=concat_axis)  # type: ignore
-
-        return self.module.as_numpy_array(latent)

pertpy-0.6.0.dist-info/RECORD

@@ -1,50 +0,0 @@
-pertpy/__init__.py,sha256=3__crpMVG7ky5lmD91Pq9qIGWgUuZQTH8xpiM5qcUJA,546
-pertpy/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pertpy/data/__init__.py,sha256=dvFUk-vAVelA65esA4EbIAVEQoE3s9K6LmE31-j2fC0,1197
-pertpy/data/_dataloader.py,sha256=pNDXLSNzOLeFM_mqf9nNvN_6Y4uA4gJfG3Y7VS-03ko,2397
-pertpy/data/_datasets.py,sha256=q20-f7MT2neWTplN300QOBlu-ihWa8IRTKxUgLgemIw,59496
-pertpy/plot/__init__.py,sha256=HB6nEBfOPmOVRHOJsJ7IJcxx2j6-6oQ__sJRaszBKuk,455
-pertpy/plot/_augur.py,sha256=pRhgc1RdRhXp6xl7-y8Z4o8beUBfltJY3XUeN9GJKbs,9064
-pertpy/plot/_cinemaot.py,sha256=tPTab-5jqalGLfa1NNeevG3_ExbKRfnIE8RRnt8Eecc,3199
-pertpy/plot/_coda.py,sha256=Ma24jc5KhuY3dtIJ6xO-pp0JpW7vWc-TPhSKJMXBEmQ,43650
-pertpy/plot/_dialogue.py,sha256=TGv_fb5f1zPEaJA8SgCue77IJkHKsQLR8f8oIz9SEcE,3881
-pertpy/plot/_guide_rna.py,sha256=Z-_vjHcOIK-DXLDTZGl5HmG6A2TnJBHv9L8VK7L3_fA,3286
-pertpy/plot/_milopy.py,sha256=6K9DtmHiCh6FUb5xScUZTxXUZoRCwD0oyfAMu0SmRGA,10994
-pertpy/plot/_mixscape.py,sha256=KeLCqWRcn2092VqB94PqBtP_wxD_OY4uS8GcZ2RXc7Y,27903
-pertpy/plot/_scgen.py,sha256=KnPe8iOqDDZw0MpSxOU7Xr-2t1UtHKehYgBQ7_4O8d4,15125
-pertpy/preprocessing/__init__.py,sha256=uja9T469LLYQAGgrTyFa4MudXci6NXnAgOn97FHXcxA,40
-pertpy/preprocessing/_guide_rna.py,sha256=EYSrsMP7FpztS0NQhn1xg0oBZZ5RT5fz6YBFvmOab58,4247
-pertpy/tools/__init__.py,sha256=QiFFM1IL7K47vuTbQqjgB8rVzauWmn6JVVpQG9AikvA,1108
-pertpy/tools/_augur.py,sha256=EUe-aRGO-PzszTS8vMfUJtzpfC3CmUSorSJTkEEU60w,45193
-pertpy/tools/_cinemaot.py,sha256=bqbxc88AH4vo2--Y5yLH3anuu1prWDAxoRZaiNvOgtQ,33374
-pertpy/tools/_dialogue.py,sha256=OUSjPzTRi46WG5QARoj2_fpmr7IQ2ftTlXT3-OiiWJc,48116
-pertpy/tools/_differential_gene_expression.py,sha256=mR06huO71KRLcU32ktCWzL-XxA9IGz8OYiRZA26eH0E,3681
-pertpy/tools/_kernel_pca.py,sha256=3S1D_wrp4vlHUPiRbCAoRbUyY-rVs112Qh-BZHSmTxE,1578
-pertpy/tools/_milo.py,sha256=OyLztlNO4Jt1c2aN3WsBbcA0UKVXVvWAnTaKwjPwJ2I,30737
-pertpy/tools/_mixscape.py,sha256=l3YHeyaUUrtuP9P8L5Z7gH47lJpzb0glszMX84DyJBI,23559
-pertpy/tools/transferlearning_MMD_LICENSE,sha256=MUvDA-o_j9htRpI8fStVdCRuyLdPkQUuIH0a_EIc57w,1069
-pertpy/tools/_coda/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pertpy/tools/_coda/_base_coda.py,sha256=mxNe5PT1XvIlZmvjQg50kh_bSmeTGVzOC63XLw2TdiI,66859
-pertpy/tools/_coda/_sccoda.py,sha256=cxaqGsXxeLf4guTU1HApAzXN2maQPexsGXIJOlW8UTM,21616
-pertpy/tools/_coda/_tasccoda.py,sha256=q0I7zM_hGjPrpy5dF2Z9trw6u8OqdkrypGgeuAhi26k,30721
-pertpy/tools/_distances/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pertpy/tools/_distances/_distance_tests.py,sha256=zRcOeLc18mRnUJ-_usUdVxWn3cZqZ8gLhglt77SaF9k,13604
-pertpy/tools/_distances/_distances.py,sha256=RMNtCD1zkORDE35XWcrh_6mw1c03hOQflmXNfoNtSRA,29780
-pertpy/tools/_metadata/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pertpy/tools/_metadata/_cell_line.py,sha256=4sUULdmxQ3TFUZDCwikN9TcHG5hf2hzlEO6gOglGl-A,33830
-pertpy/tools/_metadata/_look_up.py,sha256=H7kp9MgfgYMVdxyg3Qpf3_QmqNUkKFNMsswWeA_e1rQ,18200
-pertpy/tools/_perturbation_space/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pertpy/tools/_perturbation_space/_clustering.py,sha256=ha0TfRKUIFJmL6LE-xIfENAlYyQf4nfTpgg47X_2pHA,3237
-pertpy/tools/_perturbation_space/_discriminator_classifier.py,sha256=hTEAKTnLH4ToSdEHYuJnwui3B8L-zlSR667oG3yb49M,13861
-pertpy/tools/_perturbation_space/_metrics.py,sha256=y8-baP8WRdB1iDgvP3uuQxSCDxA2lcxvEHHM2C_vWHY,3248
-pertpy/tools/_perturbation_space/_perturbation_space.py,sha256=_A96OFbpjZULcQGfbsDhXiBjvD0chBl6c-4FoQNoV3w,14169
-pertpy/tools/_perturbation_space/_simple.py,sha256=AZx8GaNJV67evSi5oUkY11QcUkq3EcL0mtkCipjcx6c,10367
-pertpy/tools/_scgen/__init__.py,sha256=bMQ_2QbB4nnzQ7TzhI4DEFfuCDUNbZkL5xDClhQjhcA,49
-pertpy/tools/_scgen/_base_components.py,sha256=dIw-_7Z8iCietPF4tnpM7bFHtDksjnaHXwUjp9GoCIQ,2936
-pertpy/tools/_scgen/_jax_scgen.py,sha256=6fmen3zQm54Yprmd3r7zJK3GIWqpMd034DLGmi-krrs,15368
-pertpy/tools/_scgen/_jax_scgenvae.py,sha256=v_6tZ4wY-JjdMH1QVd_wG4_N0PoaqB-FM8zC2JsDu1o,3935
-pertpy/tools/_scgen/_utils.py,sha256=_G9cxBVcTIOs4wN0pgtOSkCsPJoohkeRDIb_anUqSfY,2871
-pertpy-0.6.0.dist-info/METADATA,sha256=bmYUVV99CMPm870ehtSiTbB6lPsYg0kSrmK1aoCvuu8,5046
-pertpy-0.6.0.dist-info/WHEEL,sha256=9QBuHhg6FNW7lppboF2vKVbCGTVzsFykgRQjjlajrhA,87
-pertpy-0.6.0.dist-info/licenses/LICENSE,sha256=OZ-ZkXM5CmExJiEMM90b_7dGNNvRpj7kdE-49AnrLuI,1070
-pertpy-0.6.0.dist-info/RECORD,,