Ottjax>0.5 support, optional dependency handling fixes and dataset download fixes (#288)

* ottjax compat initial fix attempt

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* comment on copied file

* compat errors and dataset error

* scvi is also properly optional

* scvi fixes

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* formatting

* add new dataset tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix draft on scvi

* relax test bounds

* remove unused import

* OTFM change to genot

* formattting

* try out new batching inpredict

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* reformat

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: selmanozleyen

pyproject.toml

@@ -37,12 +37,13 @@ dependencies = [
   "dask",
   "diffrax",
   "flax",
+  "huggingface-hub",
   "orbax",
-  "ott-jax==0.5",
-  "pyarrow",             # required for dask.dataframe
+  "ott-jax[neural]>=0.5",
+  "pyarrow",              # required for dask.dataframe
   "scanpy",
   "scikit-learn==1.5.1",
-  "scipy<1.16",          # see https://github.com/statsmodels/statsmodels/issues/9584
+  "scipy<1.16",           # see https://github.com/statsmodels/statsmodels/issues/9584
   "session-info",
 ]
 
@@ -168,6 +169,7 @@ addopts = [
 ]
 markers = [
   "slow: marks tests as slow (deselect with '-m \"not slow\"')",
+  "internet: marks tests that require internet access (deselect with '-m \"not internet\"')",
 ]
 
 [tool.coverage.run]

src/cellflow/_compat.py

@@ -0,0 +1,118 @@
+"""Compatibility helpers for optional and version-gated dependencies.
+
+``ott-jax>=0.6`` removed ``ott.neural.methods.flows.dynamics`` and the
+``ott.neural.networks.velocity_field.VelocityField`` (flax linen) class.
+This module re-exports the symbols needed by CellFlow so that both
+``ott-jax>=0.5,<0.6`` and ``ott-jax>=0.6`` are supported.
+
+The embedding helpers (``torch``, ``transformers``) are optional and only
+required when using gene-embedding functionality.
+"""
+
+# ---------------------------------------------------------------------------
+# Probability-path dynamics (BaseFlow, ConstantNoiseFlow, BrownianBridge)
+#
+# For ott-jax <0.6 we import directly from ott.  For ott-jax >=0.6 the
+# module was removed, so we provide a vendored copy below.
+#
+# The fallback classes are a verbatim copy of
+#   ott.neural.methods.flows.dynamics
+# from ott-jax 0.5.0 (commit 690b1ae, 2024-12-03).
+# ott-jax is licensed under the Apache License 2.0, which permits
+# reproduction and distribution of derivative works provided the license
+# and copyright notice are retained.  See:
+#   https://github.com/ott-jax/ott/blob/0.5.0/LICENSE
+# ---------------------------------------------------------------------------
+try:
+    from ott.neural.methods.flows.dynamics import (  # ott-jax <0.6
+        BaseFlow,
+        BrownianBridge,
+        ConstantNoiseFlow,
+    )
+except ImportError:
+    # -- Vendored from ott-jax 0.5.0 (Apache-2.0) --------------------------
+    # Source: src/ott/neural/methods/flows/dynamics.py
+    # Copyright OTT-JAX contributors
+    # -----------------------------------------------------------------------
+    import abc
+
+    import jax
+    import jax.numpy as jnp
+
+    class BaseFlow(abc.ABC):
+        """Base class for all flows."""
+
+        def __init__(self, sigma: float):
+            self.sigma = sigma
+
+        @abc.abstractmethod
+        def compute_mu_t(self, t: jnp.ndarray, x0: jnp.ndarray, x1: jnp.ndarray) -> jnp.ndarray: ...
+
+        @abc.abstractmethod
+        def compute_sigma_t(self, t: jnp.ndarray) -> jnp.ndarray: ...
+
+        @abc.abstractmethod
+        def compute_ut(self, t: jnp.ndarray, x: jnp.ndarray, x0: jnp.ndarray, x1: jnp.ndarray) -> jnp.ndarray: ...
+
+        def compute_xt(self, rng: jax.Array, t: jnp.ndarray, x0: jnp.ndarray, x1: jnp.ndarray) -> jnp.ndarray:
+            """Sample from the probability path."""
+            noise = jax.random.normal(rng, shape=x0.shape)
+            mu_t = self.compute_mu_t(t, x0, x1)
+            sigma_t = self.compute_sigma_t(t)
+            return mu_t + sigma_t * noise
+
+    class _StraightFlow(BaseFlow, abc.ABC):
+        def compute_mu_t(self, t, x0, x1):
+            return (1.0 - t) * x0 + t * x1
+
+        def compute_ut(self, t, x, x0, x1):
+            del t, x
+            return x1 - x0
+
+    class ConstantNoiseFlow(_StraightFlow):
+        r"""Flow with straight paths and constant noise :math:`\sigma`."""
+
+        def compute_sigma_t(self, t):
+            return jnp.full_like(t, fill_value=self.sigma)
+
+    class BrownianBridge(_StraightFlow):
+        r"""Brownian Bridge with :math:`\sigma_t = \sigma \sqrt{t(1-t)}`."""
+
+        def compute_sigma_t(self, t):
+            return self.sigma * jnp.sqrt(t * (1.0 - t))
+
+        def compute_ut(self, t, x, x0, x1):
+            drift_term = (1 - 2 * t) / (2 * t * (1 - t)) * (x - (t * x1 + (1 - t) * x0))
+            control_term = x1 - x0
+            return drift_term + control_term
+
+
+# ---------------------------------------------------------------------------
+# Optional embedding dependencies (torch, transformers)
+# ---------------------------------------------------------------------------
+_EMBEDDING_ERR_MSG = (
+    "To use gene embedding, please install `transformers` and `torch` e.g. via `pip install cellflow['embedding']`."
+)
+
+try:
+    import torch  # noqa: F401
+    import transformers  # noqa: F401
+
+    HAS_EMBEDDING_DEPS = True
+except ImportError:
+    HAS_EMBEDDING_DEPS = False
+
+
+def check_embedding_deps() -> None:
+    """Raise a helpful error if torch/transformers are not installed."""
+    if not HAS_EMBEDDING_DEPS:
+        raise ImportError(_EMBEDDING_ERR_MSG)
+
+
+__all__ = [
+    "BaseFlow",
+    "BrownianBridge",
+    "ConstantNoiseFlow",
+    "HAS_EMBEDDING_DEPS",
+    "check_embedding_deps",
+]

src/cellflow/datasets.py

@@ -1,19 +1,17 @@
-import os
 from typing import Any
 
 import anndata as ad
-from scanpy.readwrite import _check_datafile_present_and_download
-
-from cellflow._types import PathLike
+from huggingface_hub import hf_hub_download
 
 __all__ = [
     "ineurons",
     "pbmc_cytokines",
 ]
 
+_HF_REPO = "theislab/cellflow-datasets"
+
 
 def ineurons(
-    path: PathLike = "~/.cache/cellflow/ineurons.h5ad",
     force_download: bool = False,
     **kwargs: Any,
 ) -> ad.AnnData:
@@ -24,28 +22,23 @@ def ineurons(
 
     Parameters
     ----------
-    path
-        Path where to save the file.
     force_download
         Whether to force-download the data.
     kwargs
-        Keyword arguments for :func:`scanpy.read`.
+        Keyword arguments for :func:`anndata.read_h5ad`.
 
     Returns
     -------
     Annotated data object.
     """
-    return _load_dataset_from_url(
-        path,
-        backup_url="https://figshare.com/ndownloader/files/52852961",
-        expected_shape=(54134, 2000),  # TODO: adapt this, and enable check
+    return _load_dataset(
+        filename="ineurons.h5ad",
         force_download=force_download,
         **kwargs,
     )
 
 
 def pbmc_cytokines(
-    path: PathLike = "~/.cache/cellflow/pbmc_parse.h5ad",
     force_download: bool = False,
     **kwargs: Any,
 ) -> ad.AnnData:
@@ -57,28 +50,23 @@ def pbmc_cytokines(
 
     Parameters
     ----------
-    path
-        Path where to save the file.
     force_download
         Whether to force-download the data.
     kwargs
-        Keyword arguments for :func:`scanpy.read`.
+        Keyword arguments for :func:`anndata.read_h5ad`.
 
     Returns
     -------
     Annotated data object.
     """
-    return _load_dataset_from_url(
-        path,
-        backup_url="https://figshare.com/ndownloader/files/53372768",
-        expected_shape=(54134, 2000),  # TODO: adapt this, and enable check
+    return _load_dataset(
+        filename="pbmc_parse.h5ad",
         force_download=force_download,
         **kwargs,
     )
 
 
 def zesta(
-    path: PathLike = "~/.cache/cellflow/zesta.h5ad",
     force_download: bool = False,
     **kwargs: Any,
 ) -> ad.AnnData:
@@ -90,47 +78,33 @@ def zesta(
 
     Parameters
     ----------
-    path
-        Path where to save the file.
     force_download
         Whether to force-download the data.
     kwargs
-        Keyword arguments for :func:`scanpy.read`.
+        Keyword arguments for :func:`anndata.read_h5ad`.
 
     Returns
     -------
     Annotated data object.
     """
-    return _load_dataset_from_url(
-        path,
-        backup_url="https://figshare.com/ndownloader/files/52966469",
-        expected_shape=(54134, 2000),  # TODO: adapt this, and enable check
+    return _load_dataset(
+        filename="zesta.h5ad",
         force_download=force_download,
         **kwargs,
     )
 
 
-def _load_dataset_from_url(
-    fpath: PathLike,
+def _load_dataset(
+    filename: str,
     *,
-    backup_url: str,
-    expected_shape: tuple[int, int],
+    repo_id: str = _HF_REPO,
     force_download: bool = False,
     **kwargs: Any,
 ) -> ad.AnnData:
-    fpath = os.path.expanduser(fpath)
-    if not fpath.endswith(".h5ad"):
-        fpath += ".h5ad"
-    if force_download and os.path.exists(fpath):
-        os.remove(fpath)
-    if not _check_datafile_present_and_download(backup_url=backup_url, path=fpath):
-        raise FileNotFoundError(f"File `{fpath}` not found or download failed.")
-    data = ad.read_h5ad(filename=fpath, **kwargs)
-
-    # TODO: enable the dataset shape check
-    # if data.shape != expected_shape:
-    #    raise ValueError(
-    #        f"Expected AnnData object to have shape `{expected_shape}`, found `{data.shape}`."
-    #    )
-
-    return data
+    fpath = hf_hub_download(
+        repo_id=repo_id,
+        filename=filename,
+        repo_type="dataset",
+        force_download=force_download,
+    )
+    return ad.read_h5ad(fpath, **kwargs)

src/cellflow/external/__init__.py

@@ -1,6 +1 @@
-try:
-    from cellflow.external._scvi import CFJaxSCVI
-except ImportError as e:
-    raise ImportError(
-        "cellflow.external requires more dependencies. Please install via pip install 'cellflow[external]'"
-    ) from e
+from cellflow.external._scvi import CFJaxSCVI