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

pertpy/tools/_perturbation_space/_discriminator_classifiers.py

@@ -5,10 +5,10 @@ import warnings
 import anndata
 import numpy as np
 import pandas as pd
-import pytorch_lightning as pl
 import scipy
 import torch
 from anndata import AnnData
+from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import EarlyStopping
 from sklearn.linear_model import LogisticRegression
 from sklearn.model_selection import train_test_split
@@ -35,9 +35,7 @@ class LRClassifierSpace(PerturbationSpace):
         test_split_size: float = 0.2,
         max_iter: int = 1000,
     ):
-        """
-        Fits a logistic regression model to the data and takes the coefficients of the logistic regression
-        model as perturbation embedding.
+        """Fits a logistic regression model to the data and takes the coefficients of the logistic regression model as perturbation embedding.
 
         Args:
             adata: AnnData object of size cells x genes
@@ -60,7 +58,7 @@ class LRClassifierSpace(PerturbationSpace):
         if layer_key is not None and layer_key not in adata.obs.columns:
             raise ValueError(f"Layer key {layer_key} not found in adata.")
 
-        if embedding_key is not None and embedding_key not in adata.obsm.keys():
+        if embedding_key is not None and embedding_key not in adata.obsm:
             raise ValueError(f"Embedding key {embedding_key} not found in adata.obsm.")
 
         if layer_key is not None and embedding_key is not None:
@@ -207,7 +205,7 @@ class MLPClassifierSpace(PerturbationSpace):
         adata.obs["encoded_perturbations"] = [np.float32(label) for label in encoded_labels]
 
         # Split the data in train, test and validation
-        X = list(range(0, adata.n_obs))
+        X = list(range(adata.n_obs))
         y = adata.obs[target_col]
 
         X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_split_size, stratify=y)
@@ -248,7 +246,7 @@ class MLPClassifierSpace(PerturbationSpace):
         # Save adata observations for embedding annotations in get_embeddings
         self.adata_obs = adata.obs.reset_index(drop=True)
 
-        self.trainer = pl.Trainer(
+        self.trainer = Trainer(
             min_epochs=1,
             max_epochs=max_epochs,
             check_val_every_n_epoch=val_epochs_check,
@@ -273,7 +271,7 @@ class MLPClassifierSpace(PerturbationSpace):
                 if dataset_count == 0:
                     pert_adata = batch_adata
                 else:
-                    pert_adata = anndata.concat([pert_adata, batch_adata])
+                    pert_adata = batch_adata if dataset_count == 0 else anndata.concat([pert_adata, batch_adata])
 
         # Add .obs annotations to the pert_adata. Because shuffle=False and num_workers=0, the order of the data is stable
         # and we can just add the annotations from the original AnnData object
@@ -308,9 +306,7 @@ class MLPClassifierSpace(PerturbationSpace):
 
 
 class MLP(torch.nn.Module):
-    """
-    A multilayer perceptron with ReLU activations, optional Dropout and optional BatchNorm.
-    """
+    """A multilayer perceptron with ReLU activations, optional Dropout and optional BatchNorm."""
 
     def __init__(
         self,
@@ -320,7 +316,8 @@ class MLP(torch.nn.Module):
         layer_norm: bool = False,
         last_layer_act: str = "linear",
     ) -> None:
-        """
+        """Multilayer perceptron with ReLU activations, optional Dropout and optional BatchNorm.
+
         Args:
             sizes: size of layers.
             dropout: Dropout probability.
@@ -375,8 +372,8 @@ def init_weights(m):
 
 
 class PLDataset(Dataset):
-    """
-    Dataset for perturbation classification.
+    """Dataset for perturbation classification.
+
     Needed for training a model that classifies the perturbed cells and takes as perturbation embedding the second to last layer.
     """
 
@@ -387,14 +384,14 @@ class PLDataset(Dataset):
         label_col: str = "perturbations",
         layer_key: str = None,
     ):
-        """
+        """PyTorch lightning Dataset for perturbation classification.
+
         Args:
             adata: AnnData object with observations and labels.
             target_col: key with the perturbation labels numerically encoded.
             label_col: key with the perturbation labels.
-            layer_key: key of the layer to be used as data, otherwise .X
+            layer_key: key of the layer to be used as data, otherwise .X.
         """
-
         if layer_key:
             self.data = adata.layers[layer_key]
         else:
@@ -407,7 +404,7 @@ class PLDataset(Dataset):
         return self.data.shape[0]
 
     def __getitem__(self, idx):
-        """Returns a sample and corresponding perturbations applied (labels)"""
+        """Returns a sample and corresponding perturbations applied (labels)."""
         sample = self.data[idx].toarray().squeeze() if scipy.sparse.issparse(self.data) else self.data[idx]
         num_label = self.labels.iloc[idx]
         str_label = self.pert_labels.iloc[idx]
@@ -415,7 +412,7 @@ class PLDataset(Dataset):
         return sample, num_label, str_label
 
 
-class PerturbationClassifier(pl.LightningModule):
+class PerturbationClassifier(LightningModule):
     def __init__(
         self,
         model: torch.nn.Module,
@@ -428,7 +425,8 @@ class PerturbationClassifier(pl.LightningModule):
         lr=1e-4,
         seed=42,
     ):
-        """
+        """Perturbation Classifier.
+
         Args:
             model: model to be trained
             batch_size: batch size
@@ -438,7 +436,7 @@ class PerturbationClassifier(pl.LightningModule):
             layer_norm: whether to apply layer norm
             last_layer_act: activation function of last layer
             lr: learning rate
-            seed: random seed
+            seed: random seed.
         """
         super().__init__()
         self.batch_size = batch_size
@@ -457,16 +455,37 @@ class PerturbationClassifier(pl.LightningModule):
             last_layer_act=self.hparams.last_layer_act,
         )
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the network.
+
+        Args:
+            x: Input tensor
+
+        Returns:
+            Network output tensor
+        """
         x = self.net(x)
         return x
 
-    def configure_optimizers(self):
-        optimizer = optim.Adam(self.parameters(), lr=self.hparams.lr, weight_decay=0.1)
+    def configure_optimizers(self) -> optim.Adam:
+        """Configure optimizer for the model.
 
+        Returns:
+            Adam optimizer with weight decay
+        """
+        optimizer = optim.Adam(self.parameters(), lr=self.hparams.lr, weight_decay=0.1)
         return optimizer
 
-    def training_step(self, batch, batch_idx):
+    def training_step(self, batch: tuple[torch.Tensor, torch.Tensor, torch.Tensor], batch_idx: int) -> torch.Tensor:
+        """Perform a training step.
+
+        Args:
+            batch: Tuple of (input, target, metadata)
+            batch_idx: Index of the current batch
+
+        Returns:
+            Loss value
+        """
         x, y, _ = batch
         x = x.to(torch.float32)
 
@@ -480,7 +499,16 @@ class PerturbationClassifier(pl.LightningModule):
 
         return loss
 
-    def validation_step(self, batch, batch_idx):
+    def validation_step(self, batch: tuple[torch.Tensor, torch.Tensor, torch.Tensor], batch_idx: int) -> torch.Tensor:
+        """Perform a validation step.
+
+        Args:
+            batch: Tuple of (input, target, metadata)
+            batch_idx: Index of the current batch
+
+        Returns:
+            Loss value
+        """
         x, y, _ = batch
         x = x.to(torch.float32)
 
@@ -494,7 +522,16 @@ class PerturbationClassifier(pl.LightningModule):
 
         return loss
 
-    def test_step(self, batch, batch_idx):
+    def test_step(self, batch: tuple[torch.Tensor, torch.Tensor, torch.Tensor], batch_idx: int) -> torch.Tensor:
+        """Perform a test step.
+
+        Args:
+            batch: Tuple of (input, target, metadata)
+            batch_idx: Index of the current batch
+
+        Returns:
+            Loss value
+        """
         x, y, _ = batch
         x = x.to(torch.float32)
 
@@ -508,15 +545,29 @@ class PerturbationClassifier(pl.LightningModule):
 
         return loss
 
-    def embedding(self, x):
-        """
-        Inputs:
-            x: Input features of shape [Batch, SeqLen, 1]
+    def embedding(self, x: torch.Tensor) -> torch.Tensor:
+        """Extract embeddings from input features.
+
+        Args:
+            x: Input tensor of shape [Batch, SeqLen, 1]
+
+        Returns:
+            Embedded representation of the input
         """
         x = self.net.embedding(x)
         return x
 
-    def get_embeddings(self, batch):
+    def get_embeddings(
+        self, batch: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Extract embeddings from a batch.
+
+        Args:
+            batch: Tuple of (input, target, metadata)
+
+        Returns:
+            Tuple of (embeddings, metadata)
+        """
         x, _, y = batch
         x = x.to(torch.float32)
 

pertpy/tools/_perturbation_space/_perturbation_space.py

@@ -70,7 +70,7 @@ class PerturbationSpace:
         if embedding_key is not None and embedding_key not in adata.obsm_keys():
             raise ValueError(f"Embedding key {embedding_key} not found in obsm keys of the anndata.")
 
-        if layer_key is not None and layer_key not in adata.layers.keys():
+        if layer_key is not None and layer_key not in adata.layers:
             raise ValueError(f"Layer {layer_key!r} does not exist in the anndata.")
 
         if copy:
@@ -123,7 +123,7 @@ class PerturbationSpace:
         if all_data:
             layers_keys = list(adata.layers.keys())
             for local_layer_key in layers_keys:
-                if local_layer_key != layer_key and local_layer_key != new_layer_key:
+                if local_layer_key not in (layer_key, new_layer_key):
                     adata.layers[local_layer_key + "_control_diff"] = np.zeros((adata.n_obs, adata.n_vars))
                     for mask in group_masks:
                         adata.layers[local_layer_key + "_control_diff"][mask, :] = adata.layers[local_layer_key][
@@ -132,7 +132,7 @@ class PerturbationSpace:
 
             embedding_keys = list(adata.obsm_keys())
             for local_embedding_key in embedding_keys:
-                if local_embedding_key != embedding_key and local_embedding_key != new_embedding_key:
+                if local_embedding_key not in (embedding_key, new_embedding_key):
                     adata.obsm[local_embedding_key + "_control_diff"] = np.zeros(adata.obsm[local_embedding_key].shape)
                     for mask in group_masks:
                         adata.obsm[local_embedding_key + "_control_diff"][mask, :] = adata.obsm[local_embedding_key][
@@ -193,7 +193,7 @@ class PerturbationSpace:
 
         data: dict[str, np.array] = {}
 
-        for local_layer_key in adata.layers.keys():
+        for local_layer_key in adata.layers:
             data["layers"] = {}
             control_local = adata[reference_key].layers[local_layer_key].copy()
             for perturbation in perturbations:
@@ -231,14 +231,14 @@ class PerturbationSpace:
         new_obs.loc[new_pert_name[:-1]] = new_pert_obs
         new_perturbation.obs = new_obs
 
-        if "layers" in data.keys():
+        if "layers" in data:
             for key in data["layers"]:
                 key_name = key
                 if key.endswith("_control_diff"):
                     key_name = key.removesuffix("_control_diff")
                 new_perturbation.layers[key_name] = data["layers"][key]
 
-        if "embeddings" in data.keys():
+        if "embeddings" in data:
             key_name = key
             for key in data["embeddings"]:
                 if key.endswith("_control_diff"):
@@ -260,7 +260,7 @@ class PerturbationSpace:
         ensure_consistency: bool = False,
         target_col: str = "perturbation",
     ) -> tuple[AnnData, AnnData] | AnnData:
-        """Subtract perturbations linearly. Assumes input of size n_perts x dimensionality
+        """Subtract perturbations linearly. Assumes input of size n_perts x dimensionality.
 
         Args:
             adata: Anndata object of size n_perts x dim.
@@ -302,7 +302,7 @@ class PerturbationSpace:
 
         data: dict[str, np.array] = {}
 
-        for local_layer_key in adata.layers.keys():
+        for local_layer_key in adata.layers:
             data["layers"] = {}
             control_local = adata[reference_key].layers[local_layer_key].copy()
             for perturbation in perturbations:
@@ -340,14 +340,14 @@ class PerturbationSpace:
         new_obs.loc[new_pert_name[:-1]] = new_pert_obs
         new_perturbation.obs = new_obs
 
-        if "layers" in data.keys():
+        if "layers" in data:
             for key in data["layers"]:
                 key_name = key
                 if key.endswith("_control_diff"):
                     key_name = key.removesuffix("_control_diff")
                 new_perturbation.layers[key_name] = data["layers"][key]
 
-        if "embeddings" in data.keys():
+        if "embeddings" in data:
             key_name = key
             for key in data["embeddings"]:
                 if key.endswith("_control_diff"):

pertpy/tools/_perturbation_space/_simple.py

@@ -2,10 +2,11 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-import decoupler as dc
 import matplotlib.pyplot as plt
 import numpy as np
 from anndata import AnnData
+from decoupler import get_pseudobulk as dc_get_pseudobulk
+from decoupler import plot_psbulk_samples as dc_plot_psbulk_samples
 from sklearn.cluster import DBSCAN, KMeans
 
 from pertpy._doc import _doc_params, doc_common_plot_args
@@ -53,7 +54,6 @@ class CentroidSpace(PerturbationSpace):
             >>> cs = pt.tl.CentroidSpace()
             >>> cs_adata = cs.compute(mdata["rna"], target_col="gene_target")
         """
-
         X = None
         if layer_key is not None and embedding_key is not None:
             raise ValueError("Please, select just either layer or embedding for computation.")
@@ -65,7 +65,7 @@ class CentroidSpace(PerturbationSpace):
                 X = np.empty((len(adata.obs[target_col].unique()), adata.obsm[embedding_key].shape[1]))
 
         if layer_key is not None:
-            if layer_key not in adata.layers.keys():
+            if layer_key not in adata.layers:
                 raise ValueError(f"Layer {layer_key!r} does not exist in the .layers attribute.")
             else:
                 X = np.empty((len(adata.obs[target_col].unique()), adata.layers[layer_key].shape[1]))
@@ -79,8 +79,7 @@ class CentroidSpace(PerturbationSpace):
             X = np.empty((len(adata.obs[target_col].unique()), adata.obsm[embedding_key].shape[1]))
 
         index = []
-        pert_index = 0
-        for group_name, group_data in grouped:
+        for pert_index, (group_name, group_data) in enumerate(grouped):
             indices = group_data.index
             if layer_key is not None:
                 points = adata[indices].layers[layer_key]
@@ -94,7 +93,6 @@ class CentroidSpace(PerturbationSpace):
                 points, key=lambda point: np.linalg.norm(point - centroid)
             )  # Find the point in the array closest to the centroid
             X[pert_index, :] = closest_point
-            pert_index += 1
 
         ps_adata = AnnData(X=X)
         ps_adata.obs_names = index
@@ -153,7 +151,7 @@ class PseudobulkSpace(PerturbationSpace):
         if layer_key is not None and embedding_key is not None:
             raise ValueError("Please, select just either layer or embedding for computation.")
 
-        if layer_key is not None and layer_key not in adata.layers.keys():
+        if layer_key is not None and layer_key not in adata.layers:
             raise ValueError(f"Layer {layer_key!r} does not exist in the .layers attribute.")
 
         if target_col not in adata.obs:
@@ -169,14 +167,14 @@ class PseudobulkSpace(PerturbationSpace):
                 adata = adata_emb
 
         adata.obs[target_col] = adata.obs[target_col].astype("category")
-        ps_adata = dc.get_pseudobulk(adata, sample_col=target_col, layer=layer_key, groups_col=groups_col, **kwargs)  # type: ignore
+        ps_adata = dc_get_pseudobulk(adata, sample_col=target_col, layer=layer_key, groups_col=groups_col, **kwargs)  # type: ignore
 
         ps_adata.obs[target_col] = ps_adata.obs[target_col].astype("category")
 
         return ps_adata
 
     @_doc_params(common_plot_args=doc_common_plot_args)
-    def plot_psbulk_samples(
+    def plot_psbulk_samples(  # pragma: no cover # noqa: D417
         self,
         adata: AnnData,
         groupby: str,
@@ -209,7 +207,7 @@ class PseudobulkSpace(PerturbationSpace):
         Preview:
             .. image:: /_static/docstring_previews/pseudobulk_samples.png
         """
-        fig = dc.plot_psbulk_samples(adata, groupby, return_fig=True, **kwargs)
+        fig = dc_plot_psbulk_samples(adata, groupby, return_fig=True, **kwargs)
 
         if return_fig:
             return fig
@@ -244,7 +242,7 @@ class KMeansSpace(ClusteringSpace):
         Returns:
             If return_object is True, the adata and the clustering object is returned.
             Otherwise, only the adata is returned. The adata is updated with a new .obs column as specified in cluster_key,
-             that stores the cluster labels.
+            that stores the cluster labels.
 
         Examples:
             >>> import pertpy as pt
@@ -265,7 +263,7 @@ class KMeansSpace(ClusteringSpace):
                 self.X = adata.obsm[embedding_key]
 
         elif layer_key is not None:
-            if layer_key not in adata.layers.keys():
+            if layer_key not in adata.layers:
                 raise ValueError(f"Layer {layer_key!r} does not exist in the anndata.")
             else:
                 self.X = adata.layers[layer_key]
@@ -284,7 +282,7 @@ class KMeansSpace(ClusteringSpace):
 
 
 class DBSCANSpace(ClusteringSpace):
-    """Cluster the given data using DBSCAN"""
+    """Cluster the given data using DBSCAN."""
 
     def compute(  # type: ignore
         self,
@@ -328,7 +326,7 @@ class DBSCANSpace(ClusteringSpace):
                 self.X = adata.obsm[embedding_key]
 
         elif layer_key is not None:
-            if layer_key not in adata.layers.keys():
+            if layer_key not in adata.layers:
                 raise ValueError(f"Layer {layer_key!r} does not exist in the anndata.")
             else:
                 self.X = adata.layers[layer_key]

pertpy/tools/_scgen/_scgen.py

@@ -77,7 +77,7 @@ class Scgen(JaxTrainingMixin, BaseModelClass):
             restrict_arithmetic_to: Dictionary of celltypes you want to be observed for prediction.
 
         Returns:
-            `np nd-array` of predicted cells in primary space.
+            :class:`numpy.ndarray` of predicted cells in primary space.
         delta: float
             Difference between stimulated and control cells in latent space
 
@@ -198,7 +198,7 @@ class Scgen(JaxTrainingMixin, BaseModelClass):
                    corresponding to batch and cell type metadata, respectively.
 
         Returns:
-            corrected: `~anndata.AnnData`
+            A corrected `~anndata.AnnData` object.
             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.
 
@@ -343,6 +343,8 @@ class Scgen(JaxTrainingMixin, BaseModelClass):
                    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`.
+            give_mean: Whether to return the mean
+            n_samples: The number of samples to use.
 
         Returns:
             Low-dimensional representation for each cell
@@ -365,17 +367,14 @@ class Scgen(JaxTrainingMixin, BaseModelClass):
         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"]
+            z = out["qz"].mean if give_mean else 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)
 
-    def plot_reg_mean_plot(
+    def plot_reg_mean_plot(  # pragma: no cover # noqa: D417
         self,
         adata,
         condition_key: str,
@@ -495,14 +494,14 @@ class Scgen(JaxTrainingMixin, BaseModelClass):
         ax.text(
             max(x) - max(x) * x_coeff,
             max(y) - y_coeff * max(y),
-            r"$\mathrm{R^2_{\mathrm{\mathsf{all\ genes}}}}$= " + f"{r_value ** 2:.2f}",
+            r"$\mathrm{R^2_{\mathrm{\mathsf{all\ genes}}}}$= " + f"{r_value**2:.2f}",
             fontsize=kwargs.get("textsize", fontsize),
         )
         if diff_genes is not None:
             ax.text(
                 max(x) - max(x) * x_coeff,
                 max(y) - (y_coeff + 0.15) * max(y),
-                r"$\mathrm{R^2_{\mathrm{\mathsf{top\ 100\ DEGs}}}}$= " + f"{r_value_diff ** 2:.2f}",
+                r"$\mathrm{R^2_{\mathrm{\mathsf{top\ 100\ DEGs}}}}$= " + f"{r_value_diff**2:.2f}",
                 fontsize=kwargs.get("textsize", fontsize),
             )
 
@@ -516,7 +515,7 @@ class Scgen(JaxTrainingMixin, BaseModelClass):
         else:
             return r_value**2
 
-    def plot_reg_var_plot(
+    def plot_reg_var_plot(  # pragma: no cover # noqa: D417
         self,
         adata,
         condition_key: str,
@@ -576,7 +575,7 @@ class Scgen(JaxTrainingMixin, BaseModelClass):
             m, b, r_value_diff, p_value_diff, std_err_diff = stats.linregress(x_diff, y_diff)
             if verbose:
                 logger.info("Top 100 DEGs var: ", r_value_diff**2)
-        if "y1" in axis_keys.keys():
+        if "y1" in axis_keys:
             real_stim = adata[adata.obs[condition_key] == axis_keys["y1"]]
         x = np.asarray(np.var(ctrl.X, axis=0)).ravel()
         y = np.asarray(np.var(stim.X, axis=0)).ravel()
@@ -594,7 +593,7 @@ class Scgen(JaxTrainingMixin, BaseModelClass):
         # plt.plot(x, m * x + b, "-", color="green")
         ax.set_xlabel(labels["x"], fontsize=fontsize)
         ax.set_ylabel(labels["y"], fontsize=fontsize)
-        if "y1" in axis_keys.keys():
+        if "y1" in axis_keys:
             y1 = np.asarray(np.var(real_stim.X, axis=0)).ravel()
             _ = plt.scatter(
                 x,
@@ -611,7 +610,7 @@ class Scgen(JaxTrainingMixin, BaseModelClass):
                 y_bar = y[j]
                 plt.text(x_bar, y_bar, i, fontsize=11, color="black")
                 plt.plot(x_bar, y_bar, "o", color="red", markersize=5)
-                if "y1" in axis_keys.keys():
+                if "y1" in axis_keys:
                     y1_bar = y1[j]
                     plt.text(x_bar, y1_bar, "*", color="black", alpha=0.5)
         if legend:
@@ -623,14 +622,14 @@ class Scgen(JaxTrainingMixin, BaseModelClass):
         ax.text(
             max(x) - max(x) * x_coeff,
             max(y) - y_coeff * max(y),
-            r"$\mathrm{R^2_{\mathrm{\mathsf{all\ genes}}}}$= " + f"{r_value ** 2:.2f}",
+            r"$\mathrm{R^2_{\mathrm{\mathsf{all\ genes}}}}$= " + f"{r_value**2:.2f}",
             fontsize=kwargs.get("textsize", fontsize),
         )
         if diff_genes is not None:
             ax.text(
                 max(x) - max(x) * x_coeff,
                 max(y) - (y_coeff + 0.15) * max(y),
-                r"$\mathrm{R^2_{\mathrm{\mathsf{top\ 100\ DEGs}}}}$= " + f"{r_value_diff ** 2:.2f}",
+                r"$\mathrm{R^2_{\mathrm{\mathsf{top\ 100\ DEGs}}}}$= " + f"{r_value_diff**2:.2f}",
                 fontsize=kwargs.get("textsize", fontsize),
             )
 
@@ -645,7 +644,7 @@ class Scgen(JaxTrainingMixin, BaseModelClass):
             return r_value**2
 
     @_doc_params(common_plot_args=doc_common_plot_args)
-    def plot_binary_classifier(
+    def plot_binary_classifier(  # pragma: no cover # noqa: D417
         self,
         scgen: Scgen,
         adata: AnnData | None,
@@ -665,7 +664,7 @@ class Scgen(JaxTrainingMixin, BaseModelClass):
         Args:
             scgen: ScGen object that was trained.
             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`,
+                   AnnData object used to initialize the model. Must have been set up with `batch_key` and `labels_key`,
                    corresponding to batch and cell type metadata, respectively.
             delta: Difference between stimulated and control cells in latent space
             ctrl_key: Key for `control` part of the `data` found in `condition_key`.

pertpy/tools/_scgen/_scgenvae.py

@@ -24,8 +24,8 @@ class JaxSCGENVAE(JaxBaseModuleClass):
     training: bool = True
 
     def setup(self):
-        use_batch_norm_encoder = self.use_batch_norm == "encoder" or self.use_batch_norm == "both"
-        use_layer_norm_encoder = self.use_layer_norm == "encoder" or self.use_layer_norm == "both"
+        use_batch_norm_encoder = self.use_batch_norm in ("encoder", "both")
+        use_layer_norm_encoder = self.use_layer_norm in ("encoder", "both")
 
         self.encoder = FlaxEncoder(
             n_latent=self.n_latent,

pertpy/tools/_scgen/_utils.py

@@ -32,7 +32,9 @@ def extractor(
 
             train_data = anndata.read("./data/train.h5ad")
             test_data = anndata.read("./data/test.h5ad")
-            train_data_extracted_list = extractor(train_data, "CD4T", "conditions", "cell_type", "control", "stimulated")
+            train_data_extracted_list = extractor(
+                train_data, "CD4T", "conditions", "cell_type", "control", "stimulated"
+            )
     """
     cell_with_both_condition = data[data.obs[cell_type_key] == cell_type]
     condition_1 = data[(data.obs[cell_type_key] == cell_type) & (data.obs[condition_key] == ctrl_key)]