Borzoi profile prediction and variant effect prediction, Update SNP variant effect prediction

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning][].
 
 - `embpy.store.actions.compile_action_table` is now the shared headless action compiler, with explicit RESOLVED / CONTROL / UNRESOLVED status accounting and JSON sidecars.
 - `embpy.resources.gene.control.ControlPolicy` and status-aware embedding APIs add clearer handling for control and unresolved identifiers.
+- `AlphaGenomeWrapper` (Google DeepMind's cloud-API DNA model) and `ScoobyWrapper` (gagneurlab/scooby, single-cell-resolution DNA sequence model) are now registered in `DNA_MODELS`/`MODEL_REGISTRY` under the keys `alphagenome`, `scooby_onek1k`, `scooby_neurips`, and `scooby_epicardioids`, reachable through `BioEmbedder.embed(model=...)` and `list_available_models()` like any other model. Each is gated behind its own optional extra (`embpy[alphagenome]`, `embpy[scooby]`) and standalone pixi environment (`pixi install -e alphagenome` / `-e scooby`) -- neither is pulled in by `embpy[all]` or the default/gpu/dev pixi environments, since AlphaGenome needs an external `ALPHAGENOME_API_KEY` and Scooby installs from a git URL.
 
 ### Changed
 

docs/index.md

@@ -24,4 +24,5 @@ notebooks/genes
 notebooks/proteins
 notebooks/small_molecules
 notebooks/cells
+notebooks/variant_effects
 ```

docs/technical.md

@@ -127,10 +127,12 @@ pixi run -e default verify
 Common environments:
 
 ```bash
-pixi install -e default   # CPU/dev
-pixi install -e mps       # Apple Silicon/PyTorch MPS
-pixi install -e gpu       # Linux CUDA GPU
-pixi install -e boltz     # Boltz-specific examples
+pixi install -e default      # CPU/dev
+pixi install -e mps          # Apple Silicon/PyTorch MPS
+pixi install -e gpu          # Linux CUDA GPU
+pixi install -e boltz        # Boltz-specific examples
+pixi install -e alphagenome  # AlphaGenome cloud API client (opt-in)
+pixi install -e scooby       # Scooby single-cell DNA model (opt-in, git install)
 ```
 
 Pip users can start with:
@@ -142,6 +144,29 @@ pip install embpy
 Optional model families may require extra dependencies. Large GPU models should
 be installed in an environment with compatible PyTorch/CUDA wheels.
 
+Pip extras for individual model families (see `pyproject.toml` for the full,
+up-to-date list and per-extra caveats):
+
+```bash
+pip install "embpy[seqmodels]"     # Borzoi, Enformer
+pip install "embpy[alphagenome]"   # AlphaGenome (cloud API; needs ALPHAGENOME_API_KEY)
+pip install "embpy[scooby]"        # Scooby (installs from git; see pyproject.toml)
+pip install "embpy[esm3]"          # ESM-3 / ESM-C
+pip install "embpy[helical]"       # Helical-backed single-cell models
+pip install "embpy[state]"         # Arc STATE
+pip install "embpy[stack]"         # Arc STACK
+pip install "embpy[all]"           # every extra that installs cleanly via plain pip
+```
+
+`embpy[all]` intentionally excludes extras that need a compiler/CUDA toolchain
+or pre-built C library (`caduceus`, `evo2`, `helical`), need torch
+pre-installed to build (`minimol`), pin incompatible dependency versions
+(`esm3`, `boltz`), install from a git URL (`scooby`), or call an external
+hosted API (`alphagenome`) -- install those explicitly.
+
+`AlphaGenomeWrapper` calls a hosted API rather than loading local weights, so
+set `ALPHAGENOME_API_KEY` in the environment before calling `.load()`.
+
 ## Model Keys
 
 The active model registry depends on the installed optional dependencies.

pixi.toml

@@ -21,6 +21,8 @@
 #   pixi shell -e gpu              # activate the CUDA 12.4 env
 #   pixi run -e dev test           # run the test suite in the dev env
 #   pixi run -e docs build-docs    # build the Sphinx docs
+#   pixi install -e alphagenome    # AlphaGenome cloud API client (opt-in)
+#   pixi install -e scooby         # Scooby single-cell DNA model (opt-in, git install)
 # ---------------------------------------------------------------------------
 
 [workspace]
@@ -158,6 +160,30 @@ transformers = ">=4.45,<4.51"
 borzoi-pytorch = ">=0.4.3"
 enformer-pytorch = ">=0.8.10"
 
+# AlphaGenome (Google DeepMind cloud API client). No local weights are
+# downloaded; inference runs against Google's servers and needs
+# ALPHAGENOME_API_KEY set in the environment. Pure-Python with no CUDA/build
+# requirements -- kept opt-in (rather than folded into default/gpu/dev) so
+# the main envs don't gain a hard dependency on an external API account.
+[feature.alphagenome.pypi-dependencies]
+alphagenome = ">=0.7"
+
+# Scooby (single-cell-resolution DNA sequence model, gagneurlab/scooby).
+# Installed from git because PyPI's `scooby` package name is already taken
+# by an unrelated project (a pyvista/VTK system-report utility) -- the real
+# package has to be pulled from the gagneurlab fork by URL, not by name. It
+# also pulls in a patched `peft` fork via git (scooby's own setup.py) plus
+# `borzoi-pytorch`/`enformer-pytorch`, which pin `transformers<4.51.0` --
+# mirrored on the conda side here like `seqmodels`.
+# Kept as its own standalone environment (see `[environments]` below) so the
+# git-based resolve doesn't slow down or destabilize the main envs.
+[feature.scooby.dependencies]
+transformers = ">=4.45,<4.51"
+
+[feature.scooby.pypi-dependencies]
+snapatac2-scooby = "*"
+scooby = { git = "https://github.com/gagneurlab/scooby.git" }
+
 # ---------------------------------------------------------------------------
 # Evo 2 (Arc Institute long-context DNA LM, StripedHyena 2 architecture via
 # Vortex). Linux-64 + CUDA 12 only.
@@ -584,6 +610,21 @@ caduceus = { features = [
   "jupyter",
 ] }
 
+# Standalone env for AlphaGenome (cloud API client, no local weights). Kept
+# out of default/gpu/dev so those envs don't gain a hard dependency on an
+# external API account; CPU-only since there's no local inference. Usage:
+#   pixi install -e alphagenome && pixi shell -e alphagenome
+#   export ALPHAGENOME_API_KEY=...
+alphagenome = { features = ["cpu", "alphagenome", "morphology", "jupyter"], solve-group = "cpu" }
+
+# Standalone env for Scooby (single-cell-resolution DNA sequence model).
+# Kept separate from `default`/`gpu` because it installs from a git URL (see
+# `[feature.scooby]` above), which would slow down/destabilize the main
+# solve-groups. GPU only in practice -- Scooby shares Borzoi's trunk, which
+# needs a GPU for reasonable runtime. Usage:
+#   pixi install -e scooby && pixi shell -e scooby
+scooby = { features = ["gpu", "scooby", "morphology", "jupyter"] }
+
 # CPU env that also pulls the docs + dev extras. Useful for tutorial work.
 dev = { features = ["cpu", "morphology", "jump", "scanpy", "cell_eval", "esm3", "seqmodels", "jupyter", "dev", "docs"], solve-group = "cpu" }
 

pyproject.toml

@@ -93,6 +93,20 @@ optional-dependencies.seqmodels = [
   "borzoi-pytorch>=0.4.3",
   "enformer-pytorch>=0.8.10",
 ]
+# AlphaGenome (Google DeepMind cloud API client). No local weights are
+# downloaded; inference runs against Google's servers and requires an
+# `ALPHAGENOME_API_KEY` in the environment. Pure-Python, no CUDA/compiler
+# needed -- kept out of "all" anyway (see below) since it implies an
+# external API account rather than a local model.
+optional-dependencies.alphagenome = [
+  "alphagenome>=0.7",
+]
+# Scooby (gagneurlab/scooby): single-cell-resolution DNA sequence model built
+# on the Borzoi trunk, conditioned on a precomputed per-cell embedding. 
+optional-dependencies.scooby = [
+  "snapatac2-scooby",
+  "scooby @ git+https://github.com/gagneurlab/scooby.git",
+]
 # MiniMol molecule GNN. Kept separate because it transitively depends on
 # `torch-sparse`, `torch-scatter` etc. which fail to build via pip unless
 # torch is pre-installed (the packages don't declare torch as a build-time
@@ -163,6 +177,10 @@ optional-dependencies.morphology = [
 #   * helical -- transitively needs the igraph C library (louvain) which
 #     pip cannot build; use conda/pixi for helical instead
 #   * minimol -- transitive torch-sparse build chain; use conda/pixi instead
+#   * scooby -- pulls in `scooby` and `peft` via direct git references (see
+#     the scooby extra above); install it explicitly if you need it
+#   * alphagenome -- calls a hosted API and needs an ALPHAGENOME_API_KEY;
+#     kept opt-in so "all" doesn't imply an external API account
 # Install those explicitly in a separate env if you need them.
 optional-dependencies.all = [
   "embpy[torch,ppi,lamindb,pertpy,scanpy,morphology,seqmodels]",

src/embpy/embedder_registry/dna.py

@@ -8,12 +8,12 @@
 ``BioEmbedder.embed_gene``'s human/mouse guard and should live next to
 the DNA registry; they do, now.
 
-The optional ``EvoWrapper`` / ``Evo2Wrapper`` gating uses the
-``_HAVE_*`` flag pattern so the entire dict can be constructed even
-when the optional dependency is missing -- the wrapper slot for
-unavailable models is ``None`` and the per-call dispatch (see
-``BioEmbedder._discover_models``) drops the entry from the user-facing
-listing.
+The optional ``EvoWrapper`` / ``Evo2Wrapper`` / ``AlphaGenomeWrapper`` /
+``ScoobyWrapper`` gating uses the ``_HAVE_*`` flag pattern so the entire
+dict can be constructed even when the optional dependency is missing --
+the wrapper slot for unavailable models is ``None`` and the per-call
+dispatch (see ``BioEmbedder._discover_models``) drops the entry from the
+user-facing listing.
 """
 
 from __future__ import annotations
@@ -45,6 +45,22 @@
     _HAVE_EVO2 = False
     Evo2Wrapper = None  # type: ignore
 
+try:
+    from ..models.alphagenome_models import AlphaGenomeWrapper
+
+    _HAVE_ALPHAGENOME = True
+except ImportError:
+    _HAVE_ALPHAGENOME = False
+    AlphaGenomeWrapper = None  # type: ignore
+
+try:
+    from ..models.scooby_models import ScoobyWrapper
+
+    _HAVE_SCOOBY = True
+except ImportError:
+    _HAVE_SCOOBY = False
+    ScoobyWrapper = None  # type: ignore
+
 
 DNA_MODELS: dict[str, tuple[type[BaseModelWrapper] | None, str | None]] = {
     "enformer_human_rough": (EnformerWrapper, "EleutherAI/enformer-official-rough"),
@@ -141,6 +157,18 @@
         CaduceusWrapper,
         "kuleshov-group/caduceus-ps_seqlen-131k_d_model-256_n_layer-16",
     ),
+    # AlphaGenome (Google DeepMind cloud API client) -- pip install embpy[alphagenome]
+    "alphagenome": (AlphaGenomeWrapper if _HAVE_ALPHAGENOME else None, "alphagenome"),
+    # Scooby (gagneurlab/scooby) -- pip install embpy[scooby]. Single-cell-resolution
+    # DNA sequence model conditioned on a precomputed per-cell embedding; the model
+    # path selects the checkpoint (and, via ScoobyWrapper.KNOWN_CHECKPOINTS, its
+    # architecture hyperparameters).
+    "scooby_onek1k": (ScoobyWrapper if _HAVE_SCOOBY else None, "lauradmartens/onek1k-scooby"),
+    "scooby_neurips": (ScoobyWrapper if _HAVE_SCOOBY else None, "johahi/neurips-scooby"),
+    "scooby_epicardioids": (
+        ScoobyWrapper if _HAVE_SCOOBY else None,
+        "lauradmartens/epicardioids-scooby",
+    ),
 }
 
 

src/embpy/models/__init__.py

@@ -63,6 +63,18 @@
     singlecell_info,
 )
 
+# AlphaGenome is optional (requires: pip install alphagenome)
+try:
+    from .alphagenome_models import AlphaGenomeWrapper
+except ImportError:
+    AlphaGenomeWrapper = None  # type: ignore
+
+# Scooby is optional (requires: pip install snapatac2-scooby && pip install git+https://github.com/gagneurlab/scooby.git)
+try:
+    from .scooby_models import ScoobyWrapper
+except ImportError:
+    ScoobyWrapper = None  # type: ignore
+
 __all__ = [
     "BorzoiWrapper",
     "EnformerWrapper",
@@ -100,4 +112,7 @@
     "get_singlecell_wrapper",
     "list_singlecell_models",
     "singlecell_info",
+    "AlphaGenomeWrapper",
+    "ScoobyWrapper"
 ]
+