Optimization with Huge Speedup: Dask Dataframes Instead of Manual Dataframe Iterations

pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
   "anndata",
   "cloudpickle",
   "coverage",
+  "dask",
   "diffrax",
   "flax",
   "orbax",

src/cellflow/data/_data.py

@@ -162,8 +162,8 @@ class ValidationData(BaseDataMixin):
         int, tuple[str, ...]
     ]  # (n_targets,), dictionary explaining perturbation_covariates_mask
     perturbation_idx_to_id: dict[int, Any]
-    condition_data: dict[str, ArrayLike]  # (n_targets,) all embeddings for conditions
-    control_to_perturbation: dict[int, jax.Array]  # mapping from control idx to target distribution idcs
+    condition_data: dict[str, np.ndarray]  # (n_targets,) all embeddings for conditions
+    control_to_perturbation: dict[int, np.ndarray]  # mapping from control idx to target distribution idcs
     max_combination_length: int
     null_value: Any
     data_manager: Any

src/cellflow/data/_dataloader.py

@@ -5,7 +5,6 @@
 from typing import Any, Literal
 
 import jax
-import jax.numpy as jnp
 import numpy as np
 
 from cellflow.data._data import PredictionData, TrainingData, ValidationData
@@ -111,14 +110,14 @@ def _get_key(self, cond_idx: int) -> tuple[str, ...]:
         cov_combination = self._data.perturbation_idx_to_covariates[cond_idx]  # type: ignore[attr-defined]
         return tuple(cov_combination[i] for i in range(len(cov_combination)))
 
-    def _get_perturbation_to_control(self, data: ValidationData | PredictionData) -> dict[int, int]:
+    def _get_perturbation_to_control(self, data: ValidationData | PredictionData) -> dict[int, np.ndarray]:
         d = {}
         for k, v in data.control_to_perturbation.items():
             for el in v:
                 d[el] = k
         return d
 
-    def _get_condition_data(self, cond_idx: int) -> jnp.ndarray:
+    def _get_condition_data(self, cond_idx: int) -> dict[str, np.ndarray]:
         return {k: v[[cond_idx], ...] for k, v in self._data.condition_data.items()}  # type: ignore[attr-defined]
 
 

tests/conftest.py

@@ -137,7 +137,7 @@ def adata_perturbation() -> ad.AnnData:
     for drug in adata.obs["drug1"].cat.categories:
         drug_emb[drug] = np.random.randn(5, 1)
     adata.uns["drug"] = drug_emb
-
+    print(adata.uns["drug"])
     cell_type_emb = {}
     for cell_type in adata.obs["cell_type"].cat.categories:
         cell_type_emb[cell_type] = np.random.randn(3, 1)

tests/data/test_datamanager.py

@@ -1,5 +1,5 @@
 import anndata as ad
-import jax
+import numpy as np
 import pytest
 
 from cellflow.data._datamanager import DataManager
@@ -168,6 +168,7 @@ def test_get_train_data(
             perturbation_covariate_reps=perturbation_covariate_reps,
             sample_covariates=sample_covariates,
         )
+
         assert isinstance(dm, DataManager)
         assert dm._sample_rep == sample_rep
         assert dm._control_key == "control"
@@ -185,18 +186,18 @@ def test_get_train_data(
             assert train_data.n_controls == len(adata_perturbation.obs["cell_type"].cat.categories)
 
         assert isinstance(train_data.condition_data, dict)
-        assert isinstance(list(train_data.condition_data.values())[0], jax.Array)
+        assert isinstance(list(train_data.condition_data.values())[0], np.ndarray)
         assert train_data.max_combination_length == 1
 
         if sample_covariates == [] and perturbation_covariates == {"drug": ("drug1",)}:
             assert (
                 train_data.n_perturbations
                 == (len(adata_perturbation.obs["drug1"].cat.categories) - 1) * train_data.n_controls
             )
-        assert isinstance(train_data.cell_data, jax.Array)
-        assert isinstance(train_data.split_covariates_mask, jax.Array)
+        assert isinstance(train_data.cell_data, np.ndarray)
+        assert isinstance(train_data.split_covariates_mask, np.ndarray)
         assert isinstance(train_data.split_idx_to_covariates, dict)
-        assert isinstance(train_data.perturbation_covariates_mask, jax.Array)
+        assert isinstance(train_data.perturbation_covariates_mask, np.ndarray)
         assert isinstance(train_data.perturbation_idx_to_covariates, dict)
         assert isinstance(train_data.control_to_perturbation, dict)
 
@@ -222,7 +223,6 @@ def test_get_train_data_with_combinations(
             sample_covariates=["cell_type"],
             sample_covariate_reps={"cell_type": "cell_type"},
         )
-
         train_data = dm.get_train_data(adata_perturbation)
 
         assert ((train_data.perturbation_covariates_mask == -1) + (train_data.split_covariates_mask == -1)).all()
@@ -233,7 +233,7 @@ def test_get_train_data_with_combinations(
             assert train_data.n_controls == len(adata_perturbation.obs["cell_type"].cat.categories)
 
         assert isinstance(train_data.condition_data, dict)
-        assert isinstance(list(train_data.condition_data.values())[0], jax.Array)
+        assert isinstance(list(train_data.condition_data.values())[0], np.ndarray)
         assert train_data.max_combination_length == len(perturbation_covariates["drug"])
 
         for k in perturbation_covariates.keys():
@@ -253,10 +253,10 @@ def test_get_train_data_with_combinations(
                 cov_name = adata_perturbation.obs[cov_key].values[0]
             assert train_data.condition_data[v].shape[2] == adata_perturbation.uns[k][cov_name].shape[0]
 
-        assert isinstance(train_data.cell_data, jax.Array)
-        assert isinstance(train_data.split_covariates_mask, jax.Array)
+        assert isinstance(train_data.cell_data, np.ndarray)
+        assert isinstance(train_data.split_covariates_mask, np.ndarray)
         assert isinstance(train_data.split_idx_to_covariates, dict)
-        assert isinstance(train_data.perturbation_covariates_mask, jax.Array)
+        assert isinstance(train_data.perturbation_covariates_mask, np.ndarray)
         assert isinstance(train_data.perturbation_idx_to_covariates, dict)
         assert isinstance(train_data.control_to_perturbation, dict)
 
@@ -319,16 +319,16 @@ def test_get_validation_data(
 
         val_data = dm.get_validation_data(adata_perturbation)
 
-        assert isinstance(val_data.cell_data, jax.Array)
-        assert isinstance(val_data.split_covariates_mask, jax.Array)
+        assert isinstance(val_data.cell_data, np.ndarray)
+        assert isinstance(val_data.split_covariates_mask, np.ndarray)
         assert isinstance(val_data.split_idx_to_covariates, dict)
-        assert isinstance(val_data.perturbation_covariates_mask, jax.Array)
+        assert isinstance(val_data.perturbation_covariates_mask, np.ndarray)
         assert isinstance(val_data.perturbation_idx_to_covariates, dict)
         assert isinstance(val_data.control_to_perturbation, dict)
         assert val_data.max_combination_length == len(perturbation_covariates["drug"])
 
         assert isinstance(val_data.condition_data, dict)
-        assert isinstance(list(val_data.condition_data.values())[0], jax.Array)
+        assert isinstance(list(val_data.condition_data.values())[0], np.ndarray)
 
         if sample_covariates == [] and perturbation_covariates == {"drug": ("drug1",)}:
             assert (
@@ -399,15 +399,15 @@ def test_get_prediction_data(
         adata_pred.obs["control"] = True
         pred_data = dm.get_prediction_data(adata_pred, covariate_data=adata_pred.obs, sample_rep=sample_rep)
 
-        assert isinstance(pred_data.cell_data, jax.Array)
-        assert isinstance(pred_data.split_covariates_mask, jax.Array)
+        assert isinstance(pred_data.cell_data, np.ndarray)
+        assert isinstance(pred_data.split_covariates_mask, np.ndarray)
         assert isinstance(pred_data.split_idx_to_covariates, dict)
         assert isinstance(pred_data.perturbation_idx_to_covariates, dict)
         assert isinstance(pred_data.control_to_perturbation, dict)
         assert pred_data.max_combination_length == len(perturbation_covariates["drug"])
 
         assert isinstance(pred_data.condition_data, dict)
-        assert isinstance(list(pred_data.condition_data.values())[0], jax.Array)
+        assert isinstance(list(pred_data.condition_data.values())[0], np.ndarray)
 
         if sample_covariates == [] and perturbation_covariates == {"drug": ("drug1",)}:
             assert (