Custom Oracles

Oracles provide predictions (structure, embeddings, etc.) that energy terms use to compute scores. This guide shows you how to add your own.

The Oracle base class

All oracles inherit from Oracle and must implement a predict() method that takes a list of chains and returns an OracleResult:

from bagel.oracles.base import Oracle, OracleResult
from bagel.chain import Chain


class MyOracle(Oracle):
    result_class = MyResult  # must be set to your result class

    def predict(self, chains: list[Chain]) -> OracleResult:
        # Run your model and return results
        ...

Key points:

result_class must be set to your custom result class (a subclass of OracleResult). This enables type checking when energy terms access results.
Oracles are not copied during system copies — they are shared references. This avoids duplicating heavy model weights in memory.

Defining a result class

Result classes use Pydantic BaseModel and must implement save_attributes():

import pathlib as pl
from bagel.oracles.base import OracleResult
from bagel.chain import Chain


class MyResult(OracleResult):
    input_chains: list[Chain]
    my_score: float
    my_array: list[float]

    class Config:
        arbitrary_types_allowed = True  # needed for numpy arrays or biotite types

    def save_attributes(self, filepath: pl.Path) -> None:
        """Save result data to disk (called by loggers)."""
        with open(f"{filepath}.score.txt", "w") as f:
            f.write(str(self.my_score))

The input_chains field is inherited from OracleResult and is always required — it records which chains were passed to the oracle.

Implementing predict()

The predict() method receives a list of Chain objects and must return an instance of your result class:

class MyOracle(Oracle):
    result_class = MyResult

    def __init__(self, model_path: str):
        self.model = load_model(model_path)

    def predict(self, chains: list[Chain]) -> MyResult:
        # Extract sequences from chains
        sequences = [chain.sequence for chain in chains]

        # Run your model
        score = self.model.score(sequences)

        return MyResult(
            input_chains=chains,
            my_score=score,
            my_array=[0.0] * len(sequences),
        )

Folding vs embedding oracles

For common prediction types, BAGEL provides specialized base classes with more structure:

FoldingOracle

For models that predict 3D structures. Implement fold() instead of predict():

from bagel.oracles.folding import FoldingOracle, FoldingResult
from bagel.chain import Chain
from biotite.structure import AtomArray


class MyFoldingOracle(FoldingOracle):
    def fold(self, chains: list[Chain]) -> FoldingResult:
        # Run your structure prediction model
        structure: AtomArray = run_prediction(chains)

        return FoldingResult(
            input_chains=chains,
            structure=structure,  # biotite AtomArray with 3D coordinates
            # Add any additional metrics your model provides
        )

FoldingOracle automatically routes predict() to fold(). Energy terms access the structure via oracles_result.get_structure(oracle).

EmbeddingOracle

For language models that produce per-residue embeddings. Implement embed(), _pre_process(), and _post_process():

import numpy as np
from bagel.oracles.embedding import EmbeddingOracle, EmbeddingResult
from bagel.chain import Chain


class MyEmbeddingOracle(EmbeddingOracle):
    def embed(self, chains: list[Chain]) -> EmbeddingResult:
        processed = self._pre_process(chains)
        output = self.model(processed)
        return self._post_process(output)

    def _pre_process(self, chains: list[Chain]):
        """Convert chains to model input format."""
        return [chain.sequence for chain in chains]

    def _post_process(self, output) -> EmbeddingResult:
        """Convert model output to EmbeddingResult."""
        return EmbeddingResult(
            input_chains=self._current_chains,
            embeddings=np.array(output),  # shape: (total_residues, embedding_dim)
        )

Energy terms access embeddings via oracles_result.get_embeddings(oracle).

Example: a custom oracle

Here is a simple property oracle that computes a solubility score based on sequence composition:

import pathlib as pl
import numpy as np
from bagel.oracles.base import Oracle, OracleResult
from bagel.chain import Chain


# Amino acid solubility contributions (simplified)
SOLUBILITY_SCORES = {
    "D": 1.0, "E": 1.0, "K": 0.9, "R": 0.8, "N": 0.7, "Q": 0.7,
    "S": 0.5, "T": 0.5, "H": 0.4, "G": 0.3, "A": 0.2, "P": 0.2,
    "Y": -0.1, "W": -0.3, "F": -0.5, "M": -0.4,
    "V": -0.6, "I": -0.7, "L": -0.7, "C": -0.2,
}


class SolubilityResult(OracleResult):
    input_chains: list[Chain]
    per_residue_scores: list[float]
    overall_score: float

    def save_attributes(self, filepath: pl.Path) -> None:
        np.savetxt(f"{filepath}.solubility.txt", self.per_residue_scores)


class SolubilityOracle(Oracle):
    """Predicts sequence solubility from amino acid composition."""

    result_class = SolubilityResult

    def predict(self, chains: list[Chain]) -> SolubilityResult:
        per_residue = []
        for chain in chains:
            for residue in chain.residues:
                per_residue.append(SOLUBILITY_SCORES.get(residue.name, 0.0))

        return SolubilityResult(
            input_chains=chains,
            per_residue_scores=per_residue,
            overall_score=float(np.mean(per_residue)),
        )

You can then write a custom energy term that uses this oracle:

from bagel.energies import EnergyTerm
from bagel.oracles import OraclesResultDict


class SolubilityEnergy(EnergyTerm):
    def __init__(self, oracle: SolubilityOracle, weight: float = 1.0):
        super().__init__(
            name="SolubilityEnergy",
            oracle=oracle,
            inheritable=True,
            weight=weight,
        )

    def compute(self, oracles_result: OraclesResultDict) -> tuple[float, float]:
        result = oracles_result[self.oracle]
        # Invert so that higher solubility = lower energy
        unweighted = -result.overall_score
        return unweighted, unweighted * self.weight

Guides

Case Studies

Customization

Custom Oracles

The Oracle base class

Defining a result class

Implementing predict()

Folding vs embedding oracles

FoldingOracle

EmbeddingOracle

Example: a custom oracle

​The Oracle base class

​Defining a result class

​Implementing predict()

​Folding vs embedding oracles

​FoldingOracle

​EmbeddingOracle

​Example: a custom oracle

The Oracle base class

Defining a result class

Implementing predict()

Folding vs embedding oracles

FoldingOracle

EmbeddingOracle

Example: a custom oracle