theislab
diff --git a/‎.github/workflows/test-slow.yaml‎
Lines changed: 64 additions & 0 deletions b/‎.github/workflows/test-slow.yaml‎
Lines changed: 64 additions & 0 deletions
diff --git a/‎.github/workflows/test.yaml‎
Lines changed: 1 addition & 1 deletion b/‎.github/workflows/test.yaml‎
Lines changed: 1 addition & 1 deletion
diff --git a/‎pyproject.toml‎
Lines changed: 7 additions & 1 deletion b/‎pyproject.toml‎
Lines changed: 7 additions & 1 deletion
diff --git a/‎src/cellflow/data/_data.py‎
Lines changed: 2 additions & 2 deletions b/‎src/cellflow/data/_data.py‎
Lines changed: 2 additions & 2 deletions
diff --git a/‎src/cellflow/data/_dataloader.py‎
Lines changed: 2 additions & 3 deletions b/‎src/cellflow/data/_dataloader.py‎
Lines changed: 2 additions & 3 deletions
@@ -0,0 +1,64 @@
+name: Test Slow Functions
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+  schedule:
+    - cron: "0 5 1,15 * *"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test-slow:
+    runs-on: ${{ matrix.os }}
+    defaults:
+      run:
+        shell: bash -e {0} # -e to fail on error
+
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - os: ubuntu-latest
+            python: "3.12"
+          - os: macos-latest
+            python: "3.12"
+
+    name: ${{ matrix.name }} Python ${{ matrix.python }}
+
+    env:
+      OS: ${{ matrix.os }}
+      PYTHON: ${{ matrix.python }}
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python }}
+          cache: "pip"
+          cache-dependency-path: "**/pyproject.toml"
+
+      - name: Install test dependencies
+        run: |
+          python -m pip install --upgrade pip wheel
+      - name: Install dependencies
+        run: |
+          pip install tox
+          pip install ${{ matrix.pip-flags }} ".[optional-dependencies.dev,optional-dependencies.test,optional-dependencies.external,optional-dependencies.docs,optional-dependencies.embedding]"
+          pip install wandb omegaconf
+      - name: Test slow functions
+        run: |
+          tox -e py-${{ matrix.python }} -- -m "slow"
+        env:
+          PYTEST_ADDOPTS: -vv -n 2
+
+      - name: Upload coverage reports to Codecov
+        uses: codecov/codecov-action@v5
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          slug: theislab/cellflow
@@ -57,7 +57,7 @@ jobs:
           pip install wandb omegaconf
       - name: Test
         run: |
-          tox -e py-${{ matrix.python }}
+          tox -e py-${{ matrix.python }} -- -m "not slow"
         env:
           PYTEST_ADDOPTS: -vv -n 2
 
 
@@ -34,12 +34,15 @@ dependencies = [
   "anndata",
   "cloudpickle",
   "coverage",
+  "dask",
   "diffrax",
   "flax",
   "orbax",
   "ott-jax==0.5",
+  "pyarrow",             # required for dask.dataframe
   "scanpy",
   "scikit-learn==1.5.1",
+  "scipy<1.16",          # see https://github.com/statsmodels/statsmodels/issues/9584
   "session-info",
 ]
 
@@ -137,7 +140,7 @@ lint.ignore = [
   # Missing docstring in __init__
   "D107",
   ## Disable one in each pair of mutually incompatible rules
-  # We don’t want a blank line before a class docstring
+  # We don't want a blank line before a class docstring
   "D203",
   # We want docstrings to start immediately after the opening triple quote
   "D213",
@@ -163,6 +166,9 @@ xfail_strict = true
 addopts = [
   "--import-mode=importlib", # allow using test files with same name
 ]
+markers = [
+  "slow: marks tests as slow (deselect with '-m \"not slow\"')",
+]
 
 [tool.coverage.run]
 branch = true
 
@@ -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
 
@@ -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]