Search

Search strategies

Search algorithms for protein design.

Search strategies optimize protein designs by exploring the generation space with reward-guided algorithms like beam search and vanilla generation.

class rfproteina.search.BeamSearchStrategy(checkpoints: dict[int, dict[str, Any]], n_branch: int = 4, beam_width: int = 4, save_rollouts: bool = False, save_criteria: dict[str, dict[str, Any]] | None = None, save_states: bool = False, resume_from: str | None = None, resume_seed: int | None = None)

Bases: SearchStrategy

Branch and prune at self-contained scoring checkpoints.

Parameters:

• checkpoints – Positive step keys with metrics, rewards, use_clean_prediction, and continuation_criteria. Zero is implicit. Without a checkpoint at the schedule end, survivors finish without selection.

• n_branch – Default children per state; checkpoints may override it for their incoming interval.

• beam_width – Candidates retained per input after pruning.

• save_rollouts – Keep all completed scoring rollouts as additional outputs.

• save_criteria – Optional raw-metric thresholds for saving completed rollouts.

• save_states – Save selected candidates at interior checkpoints for later branching.

• resume_from – Path to a saved candidate batch; start a new beam at its saved step.

• resume_seed – Seed for the new branches, independent of input preprocessing.

configure_metrics(model: Proteina) → None

Validate explicitly selected metrics without inferring dependencies.

generate(model: Proteina, batch: dict) → list[dict[str, Any]]

Generate structures using beam search.

Implements beam search with reward-guided pruning.

Parameters:

• model – Proteina model instance

• batch – Input batch from dataloader

Returns:

• generated_atom_array: Final protein structure

• reward: Reward score, if scored at a checkpoint

• (optionally) intermediate_samples

Return type:

List of dictionaries, each containing

Raises:

ValueError – If a checkpoint’s reward or save criterion is unevaluable.

property metric_names: set[str]

Explicit search metrics to initialize for the validation epoch.

class rfproteina.search.SearchStrategy

Bases: ABC

Abstract base class for search strategies.

Search strategies guide protein generation by exploring the generation space with reward models. Different strategies implement different exploration/exploitation trade-offs.

Subclasses must implement:

• generate(): Main generation method that uses the model and rewards

configure_metrics(model: Any) → None

Validate metric selection before generation, when needed.

abstractmethod generate(model: Any, batch: dict) → list[dict[str, Any]]

Generate structures using the search strategy.

Parameters:

• model – Proteina model instance

• batch – Input batch from dataloader

Returns:

List of generated-structure dicts (formatted).

property metric_names: set[str]

Explicit search metrics to initialize for the validation epoch.

class rfproteina.search.VanillaStrategy

Bases: SearchStrategy

Vanilla generation without search.

This strategy simply generates structures without any search or optimization. It performs standard flow matching sampling and returns the generated structures.

Example

>>> strategy = VanillaStrategy()
>>> results = strategy.generate(model, batch)

generate(model: Any, batch: dict) → list[dict[str, Any]]

Generate structures without search.

Parameters:

• model – Proteina model instance

• batch – Input batch from dataloader

Returns:

List of generated-structure dicts.

Beam search

Beam search strategy for protein design.

Beam search maintains a fixed-size set of top-N partial trajectories (the beam), advancing them iteratively while pruning based on reward scores. This balances exploration (branching) with exploitation (keeping only top candidates).

Algorithm:

1. Initialize beam with N noise samples at t=0

2. While t < 1:

   • For each state in beam, generate L branches by running K denoising steps

   • Roll out each branch to completion to compute rewards R(x, s)

   • Select top-N branches by reward to form new beam

   • Update timestep t += K

3. Return all samples that meet success criteria

Completed rollouts can be retained as designs. Selected nonfinal checkpoints can instead score predicted clean structures, without retaining predictions or calling refolding models.

class rfproteina.search.beam_search.BeamSearchStrategy(checkpoints: dict[int, dict[str, Any]], n_branch: int = 4, beam_width: int = 4, save_rollouts: bool = False, save_criteria: dict[str, dict[str, Any]] | None = None, save_states: bool = False, resume_from: str | None = None, resume_seed: int | None = None)

Bases: SearchStrategy

Branch and prune at self-contained scoring checkpoints.

Parameters:

• checkpoints – Positive step keys with metrics, rewards, use_clean_prediction, and continuation_criteria. Zero is implicit. Without a checkpoint at the schedule end, survivors finish without selection.

• n_branch – Default children per state; checkpoints may override it for their incoming interval.

• beam_width – Candidates retained per input after pruning.

• save_rollouts – Keep all completed scoring rollouts as additional outputs.

• save_criteria – Optional raw-metric thresholds for saving completed rollouts.

• save_states – Save selected candidates at interior checkpoints for later branching.

• resume_from – Path to a saved candidate batch; start a new beam at its saved step.

• resume_seed – Seed for the new branches, independent of input preprocessing.

property metric_names: set[str]

Explicit search metrics to initialize for the validation epoch.

configure_metrics(model: Proteina) → None

Validate explicitly selected metrics without inferring dependencies.

generate(model: Proteina, batch: dict) → list[dict[str, Any]]

Generate structures using beam search.

Implements beam search with reward-guided pruning.

Parameters:

• model – Proteina model instance

• batch – Input batch from dataloader

Returns:

• generated_atom_array: Final protein structure

• reward: Reward score, if scored at a checkpoint

• (optionally) intermediate_samples

Return type:

List of dictionaries, each containing

Raises:

ValueError – If a checkpoint’s reward or save criterion is unevaluable.