Merge pull request #262 from theislab/fix/control_split

Fix control splitting

Author: MUCDK

src/cellflow/data/_datamanager.py

@@ -250,12 +250,18 @@ def get_prediction_data(
             adata=adata, split_cov_combs=split_cov_combs
         )
 
+        control_to_perturbation = self._get_control_to_perturbation(
+            covariate_data=covariate_data,
+            perturbation_idx_to_covariates=cond_data.perturbation_idx_to_covariates,
+            split_cov_combs=split_cov_combs,
+        )
+
         return PredictionData(
             cell_data=cell_data,
             split_covariates_mask=split_covariates_mask,
             split_idx_to_covariates=split_idx_to_covariates,
             condition_data=cond_data.condition_data,
-            control_to_perturbation=cond_data.control_to_perturbation,
+            control_to_perturbation=control_to_perturbation,
             perturbation_idx_to_covariates=cond_data.perturbation_idx_to_covariates,
             perturbation_idx_to_id=cond_data.perturbation_idx_to_id,
             max_combination_length=cond_data.max_combination_length,
@@ -830,6 +836,35 @@ def _get_split_covariates_mask(
             src_counter += 1
         return np.asarray(split_covariates_mask), split_idx_to_covariates
 
+    def _get_control_to_perturbation(
+        self,
+        covariate_data: pd.DataFrame,
+        perturbation_idx_to_covariates: dict[int, tuple[Any]],
+        split_cov_combs: np.ndarray | list[list[Any]],
+    ) -> dict[int, np.ndarray]:
+        control_to_perturbation = {}
+
+        if len(self._split_covariates) == 0:
+            control_to_perturbation[0] = sorted(perturbation_idx_to_covariates.keys())
+        else:
+            for control_idx, split_combination in enumerate(split_cov_combs):
+                filter_dict = dict(zip(self.split_covariates, split_combination, strict=False))
+                split_cov_mask = (covariate_data[list(filter_dict.keys())] == list(filter_dict.values())).all(axis=1)
+                # Get subset of covariate_data that matches this split combination
+                matching_data = covariate_data[split_cov_mask]
+                # Find perturbation indices that correspond to this split combination
+                perturbation_indices = []
+                for pert_idx, pert_covariates in perturbation_idx_to_covariates.items():
+                    for _, row in matching_data.iterrows():
+                        pert_values = tuple(row[self.perturb_covar_keys])
+                        if pert_values == pert_covariates:
+                            perturbation_indices.append(pert_idx)
+                            break
+
+                control_to_perturbation[control_idx] = sorted(perturbation_indices)
+
+        return control_to_perturbation
+
     @staticmethod
     def _verify_perturbation_covariates(data: dict[str, Sequence[str]] | None) -> dict[str, list[str]]:
         if data is None:

src/cellflow/model/_cellflow.py

@@ -232,6 +232,16 @@ def prepare_validation_data(
             n_conditions_on_train_end=n_conditions_on_train_end,
         )
         self._validation_data[name] = val_data
+        # Batched prediction is not compatible with split covariates
+        # as all conditions need to be the same size
+        split_val = len(val_data.control_to_perturbation) > 1
+        predict_kwargs = predict_kwargs or {}
+        # Check if predict_kwargs is alreday provided from an earlier call
+        if "predict_kwargs" in self._validation_data:
+            predict_kwargs = self._validation_data["predict_kwargs"].update(predict_kwargs)
+        # Set batched prediction to False if split_val is True
+        if split_val:
+            predict_kwargs["batched"] = False
         self._validation_data["predict_kwargs"] = predict_kwargs
 
     def prepare_model(
@@ -494,10 +504,8 @@ def prepare_model(
             )
         else:
             raise NotImplementedError(f"Solver must be an instance of OTFlowMatching or GENOT, got {type(self.solver)}")
-        if "predict_kwargs" in self.validation_data:
-            self._trainer = CellFlowTrainer(solver=self.solver, predict_kwargs=self.validation_data["predict_kwargs"])  # type: ignore[arg-type]
-        else:
-            self._trainer = CellFlowTrainer(solver=self.solver)  # type: ignore[arg-type]
+
+        self._trainer = CellFlowTrainer(solver=self.solver, predict_kwargs=self.validation_data["predict_kwargs"])  # type: ignore[arg-type]
 
     def train(
         self,

src/cellflow/solvers/_otfm.py

@@ -1,4 +1,5 @@
 from collections.abc import Callable
+from functools import partial
 from typing import Any
 
 import diffrax
@@ -263,6 +264,12 @@ def predict(
 
             pred_targets = batched_predict(src_inputs, batched_conditions)
             return {k: pred_targets[i] for i, k in enumerate(keys)}
+        elif isinstance(x, dict):
+            return jax.tree.map(
+                partial(self._predict_jit, rng=rng, **kwargs),
+                x,
+                condition,  # type: ignore[attr-defined]
+            )
         else:
             x_pred = self._predict_jit(x, condition, rng, **kwargs)
             return np.array(x_pred)

src/cellflow/training/_trainer.py

@@ -67,9 +67,7 @@ def _validation_step(
             condition = batch.get("condition", None)
             true_tgt = batch["target"]
             valid_source_data[val_key] = src
-            valid_pred_data[val_key] = self.solver.predict(
-                src, condition=condition, batched=True, **self.predict_kwargs
-            )
+            valid_pred_data[val_key] = self.solver.predict(src, condition=condition, **self.predict_kwargs)
             valid_true_data[val_key] = true_tgt
 
         return valid_source_data, valid_true_data, valid_pred_data