simplify

Author: flying-sheep

src/scanpy/preprocessing/_combat.py

@@ -202,18 +202,14 @@ def combat(  # noqa: PLR0915
     sanitize_anndata(adata)
 
     # construct a pandas series of the batch annotation
-    model = adata.obs[[key, *(covariates if covariates else [])]]
-    batch_info = model.groupby(key, observed=True).indices.values()
+    model: pd.DataFrame = adata.obs[[key, *(covariates if covariates else [])]]
+    batch_info = model.groupby(key, observed=True).indices
     n_batch = len(batch_info)
-    n_batches = np.array([len(v) for v in batch_info])
+    n_batches = np.array([len(v) for v in batch_info.values()])
 
     # check for batches with fewer than 2 cells
     small_batches = [
-        batch
-        for batch, size in zip(
-            model.groupby(key, observed=True).indices, n_batches, strict=True
-        )
-        if size < 2
+        batch for batch, size in zip(batch_info, n_batches, strict=True) if size < 2
     ]
     if small_batches:
         msg = (
@@ -236,7 +232,9 @@ def combat(  # noqa: PLR0915
         la.inv(batch_design.T @ batch_design) @ batch_design.T @ s_data.T
     ).values
     # first estimate for the multiplicative batch effect
-    delta_hat = [s_data.iloc[:, batch_idxs].var(axis=1) for batch_idxs in batch_info]
+    delta_hat = [
+        s_data.iloc[:, batch_idxs].var(axis=1) for batch_idxs in batch_info.values()
+    ]
 
     # empirically fix the prior hyperparameters
     gamma_bar = gamma_hat.mean(axis=1)
@@ -249,7 +247,7 @@ def combat(  # noqa: PLR0915
     # gamma star and delta star will be our empirical bayes (EB) estimators
     # for the additive and multiplicative batch effect per batch and cell
     gamma_star, delta_star = [], []
-    for i, batch_idxs in enumerate(batch_info):
+    for i, batch_idxs in enumerate(batch_info.values()):
         # temp stores our estimates for the batch effect parameters.
         # temp[0] is the additive batch effect
         # temp[1] is the multiplicative batch effect
@@ -273,7 +271,7 @@ def combat(  # noqa: PLR0915
 
     # we now apply the parametric adjustment to the standardized data from above
     # loop over all batches in the data
-    for j, batch_idxs in enumerate(batch_info):
+    for j, batch_idxs in enumerate(batch_info.values()):
         # we basically subtract the additive batch effect, rescale by the ratio
         # of multiplicative batch effect to pooled variance and add the overall gene
         # wise mean