diff --git a/src/cellink/tl/external/__init__.py b/src/cellink/tl/external/__init__.py index afc8f5c..3637e47 100644 --- a/src/cellink/tl/external/__init__.py +++ b/src/cellink/tl/external/__init__.py @@ -186,6 +186,8 @@ "convert_scooby_lora_checkpoint", "predict_scooby_profile", "score_variant_effects_scooby", + "score_variant_effects_scooby_batched", + "load_scooby_wrapper_and_embedder", "resolve_snp_and_exon_bins", ] @@ -220,6 +222,8 @@ "convert_scooby_lora_checkpoint", "predict_scooby_profile", "score_variant_effects_scooby", + "score_variant_effects_scooby_batched", + "load_scooby_wrapper_and_embedder", "resolve_snp_and_exon_bins", } diff --git a/src/cellink/tl/external/_scooby.py b/src/cellink/tl/external/_scooby.py index f8757c0..f8b5fb5 100644 --- a/src/cellink/tl/external/_scooby.py +++ b/src/cellink/tl/external/_scooby.py @@ -998,6 +998,136 @@ def score_variant_effects_scooby( ) +def load_scooby_wrapper_and_embedder( + model_path_or_name: str, + *, + pooling_strategy: str | None = None, + runner: ScoobyRunner | None = None, + **checkpoint_kwargs, +): + """Load a Scooby checkpoint once into a ``(wrapper, embedder)`` pair for + repeated batched scoring via ``score_variant_effects_scooby_batched``. + + Call this once per checkpoint, then pass the same ``wrapper``/``embedder`` + into multiple ``score_variant_effects_scooby_batched`` calls (e.g. one + per credible set, with each call's own ``cell_embeddings``) instead of + reloading the checkpoint each time. + + Parameters + ---------- + model_path_or_name : str + Passed to ``load_scooby_checkpoint``. + pooling_strategy : str, optional + Forwarded to ``SNPEmbedder`` if given; ``None`` uses embpy's own + default. + + Returns + ------- + (embpy.models.scooby_models.ScoobyWrapper, embpy.tl.genomics.SNPEmbedder) + """ + try: + from embpy.models.scooby_models import ScoobyWrapper + from embpy.tl.genomics import SNPEmbedder + except ImportError as e: + raise ImportError( + "embpy is required for Scooby variant-effect scoring. Install with:\n\n" + " pip install cellink[embpy]" + ) from e + + runner = runner or get_scooby_runner() + wrapper = ScoobyWrapper(model_path_or_name, device=runner.resolve_device(), **checkpoint_kwargs) + wrapper.load(runner.resolve_device()) + embedder = SNPEmbedder(wrapper) if pooling_strategy is None else SNPEmbedder(wrapper, pooling_strategy=pooling_strategy) + return wrapper, embedder + + +def score_variant_effects_scooby_batched( + variants: list[tuple[Any, str]], + model_path_or_name: str, + cell_embeddings: np.ndarray, + *, + bin_indices: list | None = None, + pseudocount: float = 1.0, + aggregate: Literal["pseudobulk", "none"] = "pseudobulk", + pooling_strategy: str | None = None, + runner: ScoobyRunner | None = None, + wrapper: Any | None = None, + embedder: Any | None = None, + **checkpoint_kwargs, +): + """Score many variants against ONE loaded Scooby checkpoint. + + ``score_variant_effects_scooby`` reconstructs and reloads the checkpoint + from disk on every call, which is the right shape for a single + exploratory lookup but does not scale to scoring a batch of candidates + (e.g. every member of a fine-mapped credible set, or a control-SNP set) + in one job. This function scores a whole batch against one load. + + By default it loads the checkpoint itself (like + ``score_variant_effects_scooby``, just once for the whole ``variants`` + list instead of once per variant). For a run that needs to score many + separate batches with the same checkpoint but different + ``cell_embeddings`` per batch (e.g. one call per credible set), load + once via ``load_scooby_wrapper_and_embedder`` and pass the result in as + ``wrapper``/``embedder`` to skip the reload on every call. + + Parameters + ---------- + variants : list of (snp, chromosome_sequence) + Each element is one variant to score: an + ``embpy.tl.genomics.SNPContext`` and the reference sequence window + it should be scored against (same meaning as + ``score_variant_effects_scooby``'s ``snp``/``chromosome_sequence``). + model_path_or_name : str + Passed to ``load_scooby_checkpoint`` if ``wrapper``/``embedder`` + aren't given; ignored otherwise. + cell_embeddings : np.ndarray + Per-cell embeddings to pseudobulk over (or per-cell if + ``aggregate="none"``), shared across every variant in this batch. + bin_indices : list, optional + Per-variant bin restriction, same length and order as ``variants`` + (e.g. each variant's own gene's exon bins). ``None`` scores every + variant over all bins; an individual entry may also be ``None``. + pseudocount : float + Added before the log2-fold-change computation. + pooling_strategy : str, optional + Forwarded to ``SNPEmbedder`` if given and ``wrapper``/``embedder`` + aren't; ``None`` uses embpy's own default. + wrapper, embedder : optional + A ``(ScoobyWrapper, SNPEmbedder)`` pair already loaded via + ``load_scooby_wrapper_and_embedder``. If given, both must be given, + and ``model_path_or_name``/``pooling_strategy``/``runner``/ + ``checkpoint_kwargs`` are ignored (the checkpoint is not reloaded). + + Returns + ------- + list of embpy.tl.genomics.snp_utils.VariantEffectResult + One result per element of ``variants``, in the same order. + """ + if bin_indices is not None and len(bin_indices) != len(variants): + raise ValueError(f"bin_indices has {len(bin_indices)} entries, expected one per variant ({len(variants)}).") + bin_indices_per_variant = bin_indices if bin_indices is not None else [None] * len(variants) + + if (wrapper is None) != (embedder is None): + raise ValueError("wrapper and embedder must be given together, or not at all.") + if wrapper is None: + _wrapper, embedder = load_scooby_wrapper_and_embedder( + model_path_or_name, pooling_strategy=pooling_strategy, runner=runner, **checkpoint_kwargs + ) + + return [ + embedder.predict_variant_effect( + snp, + chromosome_sequence, + bin_indices=variant_bin_indices, + pseudocount=pseudocount, + aggregate=aggregate, + cell_embeddings=cell_embeddings, + ) + for (snp, chromosome_sequence), variant_bin_indices in zip(variants, bin_indices_per_variant) + ] + + def resolve_snp_and_exon_bins( *, chrom: str,