pertpy 1.0.0__py3-none-any.whl → 1.0.2__py3-none-any.whl

pertpy/tools/_coda/_tasccoda.py

@@ -120,8 +120,10 @@ class Tasccoda(CompositionalModel2):
                 covariate_df=covariate_df,
             )
             mdata = MuData({modality_key_1: adata, modality_key_2: adata_coda})
-        else:
+        elif type == "sample_level":
             mdata = MuData({modality_key_1: AnnData(), modality_key_2: adata})
+        else:
+            raise ValueError(f'{type} is not a supported type, expected "cell_level" or "sample_level".')
         import_tree(
             data=mdata,
             modality_1=modality_key_1,
@@ -464,7 +466,7 @@ class Tasccoda(CompositionalModel2):
         self,
         data: AnnData | MuData,
         modality_key: str = "coda",
-        rng_key=None,
+        rng_key: int | None = None,
         num_prior_samples: int = 500,
         use_posterior_predictive: bool = True,
     ) -> az.InferenceData:
@@ -547,6 +549,8 @@ class Tasccoda(CompositionalModel2):
         if rng_key is None:
             rng = np.random.default_rng()
             rng_key = random.key(rng.integers(0, 10000))
+        else:
+            rng_key = random.key(rng_key)
 
         if use_posterior_predictive:
             posterior_predictive = Predictive(self.model, self.mcmc.get_samples())(
@@ -557,6 +561,15 @@ class Tasccoda(CompositionalModel2):
                 ref_index=ref_index,
                 sample_adata=sample_adata,
             )
+            # Remove problematic posterior predictive arrays with wrong dimensions
+            if posterior_predictive and "counts" in posterior_predictive:
+                counts_shape = posterior_predictive["counts"].shape
+                expected_dims = 2  # ['sample', 'cell_type']
+                if len(counts_shape) != expected_dims:
+                    posterior_predictive = {k: v for k, v in posterior_predictive.items() if k != "counts"}
+                    logger.warning(
+                        f"Removed 'counts' from posterior_predictive due to dimension mismatch: got {len(counts_shape)}D, expected {expected_dims}D"
+                    )
         else:
             posterior_predictive = None
 
@@ -569,6 +582,15 @@ class Tasccoda(CompositionalModel2):
                 ref_index=ref_index,
                 sample_adata=sample_adata,
             )
+            # Remove problematic prior arrays with wrong dimensions
+            if prior and "counts" in prior:
+                counts_shape = prior["counts"].shape
+                expected_dims = 2  # ['sample', 'cell_type']
+                if len(counts_shape) != expected_dims:
+                    prior = {k: v for k, v in prior.items() if k != "counts"}
+                    logger.warning(
+                        f"Removed 'counts' from prior due to dimension mismatch: got {len(counts_shape)}D, expected {expected_dims}D"
+                    )
         else:
             prior = None
 
@@ -592,80 +614,88 @@ 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).
-        """  # noqa: D205
+        """
+
+        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, D212
         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).
-        """  # noqa: D205
+        """
+
+        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, D212
         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).
-        """  # noqa: D205
+        """
+
+        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, D212
         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).
-        """  # noqa: D205
+        """
+
+        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, D212
         return super().set_fdr(data, est_fdr, modality_key, *args, **kwargs)
 
     set_fdr.__doc__ = CompositionalModel2.set_fdr.__doc__ + set_fdr.__doc__

pertpy/tools/_dialogue.py

@@ -882,9 +882,9 @@ class Dialogue:
             if len(conditions_compare) != 2:
                 raise ValueError("Please specify conditions to compare or supply an object with only 2 conditions")
 
-        pvals = pd.DataFrame(1, adata.obs[celltype_label].unique(), ["mcp_" + str(n) for n in range(n_mcps)])
-        tstats = pd.DataFrame(1, adata.obs[celltype_label].unique(), ["mcp_" + str(n) for n in range(n_mcps)])
-        pvals_adj = pd.DataFrame(1, adata.obs[celltype_label].unique(), ["mcp_" + str(n) for n in range(n_mcps)])
+        pvals = pd.DataFrame(1.0, adata.obs[celltype_label].unique(), ["mcp_" + str(n) for n in range(n_mcps)])
+        tstats = pd.DataFrame(1.0, adata.obs[celltype_label].unique(), ["mcp_" + str(n) for n in range(n_mcps)])
+        pvals_adj = pd.DataFrame(1.0, adata.obs[celltype_label].unique(), ["mcp_" + str(n) for n in range(n_mcps)])
 
         response = adata.obs.groupby(sample_label)[condition_label].agg(pd.Series.mode)
         for celltype in adata.obs[celltype_label].unique():

pertpy/tools/_differential_gene_expression/__init__.py

@@ -1,9 +1,52 @@
+import contextlib
+from importlib import import_module
+from importlib.util import find_spec
+
 from ._base import LinearModelBase, MethodBase
 from ._dge_comparison import DGEEVAL
 from ._edger import EdgeR
-from ._pydeseq2 import PyDESeq2
 from ._simple_tests import SimpleComparisonBase, TTest, WilcoxonTest
-from ._statsmodels import Statsmodels
+
+
+def __getattr__(name: str):
+    deps = {
+        "PyDESeq2": ["pydeseq2", "formulaic_contrasts", "formulaic"],
+        "EdgeR": ["rpy2", "formulaic_contrasts", "formulaic"],
+        "Statsmodels": ["formulaic_contrasts", "formulaic"],
+    }
+
+    if name in deps:
+        for dep in deps[name]:
+            if find_spec(dep) is None:
+                raise ImportError(f"{dep} is required but not installed")
+
+        module_map = {
+            "PyDESeq2": "pertpy.tools._differential_gene_expression._pydeseq2",
+            "EdgeR": "pertpy.tools._differential_gene_expression._edger",
+            "Statsmodels": "pertpy.tools._differential_gene_expression._statsmodels",
+        }
+
+        module = import_module(module_map[name])
+        return getattr(module, name)
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def _get_available_methods():
+    methods = [WilcoxonTest, TTest]
+    from importlib.util import find_spec
+
+    for name in ["Statsmodels", "PyDESeq2", "EdgeR"]:
+        with contextlib.suppress(ImportError):
+            methods.append(__getattr__(name))
+
+    return methods
+
+
+AVAILABLE_METHODS = _get_available_methods()
+
+
+AVAILABLE_METHODS = _get_available_methods()
 
 __all__ = [
     "MethodBase",
@@ -15,5 +58,3 @@ __all__ = [
     "WilcoxonTest",
     "TTest",
 ]
-
-AVAILABLE_METHODS = [Statsmodels, EdgeR, PyDESeq2, WilcoxonTest, TTest]

pertpy/tools/_differential_gene_expression/_base.py

@@ -12,7 +12,6 @@ import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import seaborn as sns
-from formulaic_contrasts import FormulaicContrasts
 from lamin_utils import logger
 from matplotlib.pyplot import Figure
 from matplotlib.ticker import MaxNLocator
@@ -881,6 +880,8 @@ class LinearModelBase(MethodBase):
         super().__init__(adata, mask=mask, layer=layer)
         self._check_counts()
 
+        from formulaic_contrasts import FormulaicContrasts
+
         self.formulaic_contrasts = None
         if isinstance(design, str):
             self.formulaic_contrasts = FormulaicContrasts(adata.obs, design)

pertpy/tools/_differential_gene_expression/_edger.py

@@ -23,9 +23,6 @@ class EdgeR(LinearModelBase):
         Args:
             **kwargs: Keyword arguments specific to glmQLFit()
         """
-        # For running in notebook
-        # pandas2ri.activate()
-        # rpy2.robjects.numpy2ri.activate()
         try:
             from rpy2 import robjects as ro
             from rpy2.robjects import numpy2ri, pandas2ri
@@ -47,17 +44,17 @@ class EdgeR(LinearModelBase):
             expr = self.adata.X if self.layer is None else self.adata.layers[self.layer]
             expr = expr.T.toarray() if issparse(expr) else expr.T
 
-        with localconverter(get_conversion() + pandas2ri.converter):
-            expr_r = ro.conversion.py2rpy(pd.DataFrame(expr, index=self.adata.var_names, columns=self.adata.obs_names))
-            samples_r = ro.conversion.py2rpy(self.adata.obs)
+        with localconverter(get_conversion() + pandas2ri.converter) as cv:
+            expr_r = cv.py2rpy(pd.DataFrame(expr, index=self.adata.var_names, columns=self.adata.obs_names))
+            samples_r = cv.py2rpy(self.adata.obs)
 
         dge = edger.DGEList(counts=expr_r, samples=samples_r)
 
         logger.info("Calculating NormFactors")
         dge = edger.calcNormFactors(dge)
 
-        with localconverter(get_conversion() + numpy2ri.converter):
-            design_r = ro.conversion.py2rpy(self.design.values)
+        with localconverter(get_conversion() + numpy2ri.converter) as cv:
+            design_r = cv.py2rpy(self.design.values)
 
         logger.info("Estimating Dispersions")
         dge = edger.estimateDisp(dge, design=design_r)
@@ -100,8 +97,8 @@ class EdgeR(LinearModelBase):
             ) from None
 
         # Convert vector to R, which drops a category like `self.design_matrix` to use the intercept for the left out.
-        with localconverter(get_conversion() + numpy2ri.converter):
-            contrast_vec_r = ro.conversion.py2rpy(np.asarray(contrast))
+        with localconverter(get_conversion() + numpy2ri.converter) as cv:
+            contrast_vec_r = cv.py2rpy(np.asarray(contrast))
         ro.globalenv["contrast_vec"] = contrast_vec_r
 
         # Test contrast with R
@@ -121,8 +118,8 @@ class EdgeR(LinearModelBase):
             return de_res.reset_index().rename(columns={"PValue": "p_value", "logFC": "log_fc", "FDR": "adj_p_value"})
 
         # Convert to Pandas DataFrame if still an R object
-        with localconverter(get_conversion() + pandas2ri.converter):
-            de_res = ro.conversion.rpy2py(de_res)
+        with localconverter(get_conversion() + pandas2ri.converter) as cv:
+            de_res = cv.rpy2py(de_res)
 
         de_res.index.name = "variable"
         de_res = de_res.reset_index()

pertpy/tools/_differential_gene_expression/_pydeseq2.py

@@ -1,6 +1,4 @@
 import os
-import re
-import warnings
 
 import numpy as np
 import pandas as pd

pertpy/tools/_distances/_distance_tests.py

@@ -8,7 +8,7 @@ from rich.progress import track
 from sklearn.metrics import pairwise_distances
 from statsmodels.stats.multitest import multipletests
 
-from ._distances import Distance
+from ._distances import Distance, Metric
 
 if TYPE_CHECKING:
     from anndata import AnnData
@@ -43,7 +43,7 @@ class DistanceTest:
 
     def __init__(
         self,
-        metric: str,
+        metric: Metric,
         n_perms: int = 1000,
         layer_key: str = None,
         obsm_key: str = None,

pertpy/tools/_distances/_distances.py

@@ -34,6 +34,31 @@ class MeanVar(NamedTuple):
     variance: float
 
 
+Metric = Literal[
+    "edistance",
+    "euclidean",
+    "root_mean_squared_error",
+    "mse",
+    "mean_absolute_error",
+    "pearson_distance",
+    "spearman_distance",
+    "kendalltau_distance",
+    "cosine_distance",
+    "r2_distance",
+    "mean_pairwise",
+    "mmd",
+    "wasserstein",
+    "sym_kldiv",
+    "t_test",
+    "ks_test",
+    "nb_ll",
+    "classifier_proba",
+    "classifier_cp",
+    "mean_var_distribution",
+    "mahalanobis",
+]
+
+
 class Distance:
     """Distance class, used to compute distances between groups of cells.
 
@@ -112,7 +137,7 @@ class Distance:
 
     def __init__(
         self,
-        metric: str = "edistance",
+        metric: Metric = "edistance",
         agg_fct: Callable = np.mean,
         layer_key: str = None,
         obsm_key: str = None,
@@ -660,19 +685,19 @@ class MMD(AbstractDistance):
         super().__init__()
         self.accepts_precomputed = False
 
-    def __call__(self, X: np.ndarray, Y: np.ndarray, kernel="linear", **kwargs) -> float:
+    def __call__(self, X: np.ndarray, Y: np.ndarray, *, kernel="linear", gamma=1.0, degree=2, **kwargs) -> float:
         if kernel == "linear":
             XX = np.dot(X, X.T)
             YY = np.dot(Y, Y.T)
             XY = np.dot(X, Y.T)
         elif kernel == "rbf":
-            XX = rbf_kernel(X, X, gamma=1.0)
-            YY = rbf_kernel(Y, Y, gamma=1.0)
-            XY = rbf_kernel(X, Y, gamma=1.0)
+            XX = rbf_kernel(X, X, gamma=gamma)
+            YY = rbf_kernel(Y, Y, gamma=gamma)
+            XY = rbf_kernel(X, Y, gamma=gamma)
         elif kernel == "poly":
-            XX = polynomial_kernel(X, X, degree=2, gamma=1.0, coef0=0)
-            YY = polynomial_kernel(Y, Y, degree=2, gamma=1.0, coef0=0)
-            XY = polynomial_kernel(X, Y, degree=2, gamma=1.0, coef0=0)
+            XX = polynomial_kernel(X, X, degree=degree, gamma=gamma, coef0=0)
+            YY = polynomial_kernel(Y, Y, degree=degree, gamma=gamma, coef0=0)
+            XY = polynomial_kernel(X, Y, degree=degree, gamma=gamma, coef0=0)
         else:
             raise ValueError(f"Kernel {kernel} not recognized.")
 

pertpy/tools/_milo.py

@@ -411,6 +411,8 @@ class Milo:
                     res = base.as_data_frame(
                         edgeR.topTags(edgeR.glmQLFTest(fit, coef=n_coef), sort_by="none", n=np.inf)
                     )
+            if res is None:
+                raise ValueError("Unable to generate results with edgeR. Is your installation correct?")
             if not isinstance(res, pd.DataFrame):
                 res = pd.DataFrame(res)
             # The columns of res looks like e.g. table.A, table.B, so remove the prefix
@@ -530,7 +532,7 @@ class Milo:
 
         anno_frac_dataframe = pd.DataFrame(anno_frac, columns=anno_dummies.columns, index=sample_adata.var_names)
         sample_adata.varm["frac_annotation"] = anno_frac_dataframe.values
-        sample_adata.uns["annotation_labels"] = anno_frac_dataframe.columns
+        sample_adata.uns["annotation_labels"] = anno_frac_dataframe.columns.to_list()
         sample_adata.uns["annotation_obs"] = anno_col
         sample_adata.var["nhood_annotation"] = anno_frac_dataframe.idxmax(1)
         sample_adata.var["nhood_annotation_frac"] = anno_frac_dataframe.max(1)

pertpy/tools/_perturbation_space/_discriminator_classifiers.py

@@ -1,13 +1,11 @@
 from __future__ import annotations
 
-import warnings
-
 import anndata
 import numpy as np
-import pandas as pd
 import scipy
 import torch
 from anndata import AnnData
+from fast_array_utils.conv import to_dense
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import EarlyStopping
 from sklearn.linear_model import LogisticRegression
@@ -112,18 +110,6 @@ class LRClassifierSpace(PerturbationSpace):
         return pert_adata
 
 
-# Ensure backward compatibility with DiscriminatorClassifierSpace
-def DiscriminatorClassifierSpace():
-    warnings.warn(
-        "The DiscriminatorClassifierSpace class is deprecated and will be removed in the future."
-        "Please use the MLPClassifierSpace or the LRClassifierSpace class instead.",
-        DeprecationWarning,
-        stacklevel=2,
-    )
-
-    return MLPClassifierSpace()
-
-
 class MLPClassifierSpace(PerturbationSpace):
     """Fits an ANN classifier to the data and takes the feature space (weights in the last layer) as embedding.
 
@@ -202,7 +188,7 @@ class MLPClassifierSpace(PerturbationSpace):
         labels = adata.obs[target_col].values.reshape(-1, 1)
         encoder = OneHotEncoder()
         encoded_labels = encoder.fit_transform(labels).toarray()
-        adata.obs["encoded_perturbations"] = [np.float32(label) for label in encoded_labels]
+        adata.obsm["encoded_perturbations"] = encoded_labels.astype(np.float32)
 
         # Split the data in train, test and validation
         X = list(range(adata.n_obs))
@@ -226,7 +212,7 @@ class MLPClassifierSpace(PerturbationSpace):
         # Fix class unbalance (likely to happen in perturbation datasets)
         # Usually control cells are overrepresented such that predicting control all time would give good results
         # Cells with rare perturbations are sampled more
-        train_weights = 1 / (1 + torch.sum(torch.tensor(train_dataset.labels.to_list()), dim=1))
+        train_weights = 1 / (1 + torch.sum(torch.tensor(train_dataset.labels), dim=1))
         train_sampler = WeightedRandomSampler(train_weights, len(train_weights))
 
         self.train_dataloader = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler, num_workers=4)
@@ -278,11 +264,10 @@ class MLPClassifierSpace(PerturbationSpace):
         pert_adata.obs = pert_adata.obs.reset_index(drop=True)
         if "perturbations" in self.adata_obs.columns:
             self.adata_obs = self.adata_obs.drop("perturbations", axis=1)
-        pert_adata.obs = pd.concat([pert_adata.obs, self.adata_obs], axis=1)
-
-        # Drop the 'encoded_perturbations' colums, since this stores the one-hot encoded labels as numpy arrays,
-        # which would cause errors in the downstream processing of the AnnData object (e.g. when plotting)
-        pert_adata.obs = pert_adata.obs.drop("encoded_perturbations", axis=1)
+        obs_subset = self.adata_obs.iloc[: len(pert_adata.obs)].copy()
+        for col in obs_subset.columns:
+            if col not in ["perturbations", "encoded_perturbations"]:
+                pert_adata.obs[col] = obs_subset[col].values
 
         return pert_adata
 
@@ -397,7 +382,13 @@ class PLDataset(Dataset):
         else:
             self.data = adata.X
 
-        self.labels = adata.obs[target_col]
+        if target_col in adata.obs.columns:
+            self.labels = adata.obs[target_col]
+        elif target_col in adata.obsm:
+            self.labels = adata.obsm[target_col]
+        else:
+            raise ValueError(f"Target column {target_col} not found in obs or obsm")
+
         self.pert_labels = adata.obs[label_col]
 
     def __len__(self):
@@ -405,8 +396,8 @@ class PLDataset(Dataset):
 
     def __getitem__(self, idx):
         """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]
+        sample = to_dense(self.data[idx]).squeeze() if scipy.sparse.issparse(self.data) else self.data[idx]
+        num_label = self.labels.iloc[idx] if hasattr(self.labels, "iloc") else self.labels[idx]
         str_label = self.pert_labels.iloc[idx]
 
         return sample, num_label, str_label

pertpy/tools/_perturbation_space/_simple.py

@@ -161,12 +161,20 @@ class PseudobulkSpace(PerturbationSpace):
                 adata = adata_emb
 
         adata.obs[target_col] = adata.obs[target_col].astype("category")
+        grouping_cols = [target_col] if groups_col is None else [target_col, groups_col]
+        original_obs = adata.obs.copy()
         ps_adata = sc.get.aggregate(
             adata, by=[target_col] if groups_col is None else [target_col, groups_col], func=mode, layer=layer_key
         )
+
         if mode in ps_adata.layers:
             ps_adata.X = ps_adata.layers[mode]
 
+        for col in original_obs.columns:
+            if col not in ps_adata.obs.columns:
+                grouped_values = original_obs.groupby(grouping_cols)[col].first()
+                ps_adata.obs[col] = grouped_values.reindex(ps_adata.obs.index).values
+
         ps_adata.obs[target_col] = ps_adata.obs[target_col].astype("category")
 
         return ps_adata