from __future__ import annotations from dataclasses import dataclass, field from typing import Any from staplebridge.chemistry.state import StapleState from staplebridge.data.schemas import BuildingBlock from staplebridge.oracles.base import AnchorPriorBase, BlockPriorBase, GeometryOracleBase, PeptidePriorBase from staplebridge.utils.profiling import STAGE_TIMER ACTION_NOOP = "noop" ACTION_RESIDUE_SUBSTITUTION_MOTIF = "residue_substitution_motif" ACTION_RESIDUE_SUBSTITUTION_OTHER = "residue_substitution_other" ACTION_ASSIGN_ANCHOR = "assign_anchor" ACTION_REASSIGN_ANCHOR = "reassign_anchor" ACTION_ASSIGN_BLOCK = "assign_block" ACTION_ACTIVATE_TOPOLOGY = "activate_topology" ACTION_UNKNOWN = "unknown" @dataclass class ComponentWeights: """Per-component multipliers applied on top of the eta_ scalars. Defaults are all 1.0 → identical behavior to pre-weight code. Turning any weight away from 1.0 lets the config section ``reference.weights.{seq,anchor,block,geom}`` re-scale the reference energy without changing the base eta_cost / eta_spps / eta_type / eta_geom knobs. """ seq: float = 1.0 anchor: float = 1.0 block: float = 1.0 geom: float = 1.0 progress: float = 1.0 @dataclass class ActionProgressWeights: """Reward/penalty on the action-type level. All values are *positive contributions to the reference energy*, so negative numbers reward the action and positive numbers penalize it. Defaults were picked so that motif-creating actions clearly beat unrelated substitutions but do not overpower the peptide/anchor/block priors. """ noop: float = 0.5 residue_substitution_motif: float = -1.5 residue_substitution_other: float = 1.0 assign_anchor: float = -3.0 reassign_anchor: float = -0.5 assign_block: float = -2.0 activate_topology: float = -6.0 unknown: float = 0.0 def as_dict(self) -> dict[str, float]: return { ACTION_NOOP: self.noop, ACTION_RESIDUE_SUBSTITUTION_MOTIF: self.residue_substitution_motif, ACTION_RESIDUE_SUBSTITUTION_OTHER: self.residue_substitution_other, ACTION_ASSIGN_ANCHOR: self.assign_anchor, ACTION_REASSIGN_ANCHOR: self.reassign_anchor, ACTION_ASSIGN_BLOCK: self.assign_block, ACTION_ACTIVATE_TOPOLOGY: self.activate_topology, ACTION_UNKNOWN: self.unknown, } @dataclass class ReferenceEnergyConfig: eta_cost: float = 0.5 eta_spps: float = 0.3 eta_type: float = 5.0 eta_geom: float = 1.0 eta_progress: float = 1.0 weights: dict[str, float] | None = None action_weights: dict[str, float] | None = None normalize: dict[str, Any] | None = None # reserved; disabled by default # Stage-aware geometry: # True (default): if candidate lacks anchor OR block, drop full E_geom / # E_type from the reference energy and record geom_status="not_applicable". # False: legacy behavior (any missing anchor/block is treated as full # geometry failure via the sentinel cgeom=10 in the geometry oracle). stage_aware_geometry: bool = True def resolved_weights(self) -> ComponentWeights: w = ComponentWeights() if self.weights: for k in ("seq", "anchor", "block", "geom", "progress"): if k in self.weights: setattr(w, k, float(self.weights[k])) return w def resolved_action_weights(self) -> ActionProgressWeights: aw = ActionProgressWeights() if self.action_weights: for k, v in self.action_weights.items(): if hasattr(aw, k): setattr(aw, k, float(v)) return aw # --------------------------------------------------------------------------- # Action-type classification # --------------------------------------------------------------------------- def _spans_lactam_motif(seq: list[str], i: int, j: int) -> bool: """True if (i, j) is a valid K-(D|E) motif at spacing 3 or 4.""" if i > j: i, j = j, i sp = j - i if sp not in (3, 4): return False if i < 0 or j >= len(seq): return False if sp == 3: return seq[i] == "K" and seq[j] == "D" return seq[i] == "K" and seq[j] == "E" def _has_lactam_partner(seq: list[str], pos: int, tok: str) -> bool: """True if placing `tok` at `pos` completes a K-(D|E) i,i+3 / i,i+4 motif given the *current* residues at the partner positions.""" n = len(seq) if tok == "K": for sp, partner in ((3, "D"), (4, "E")): j = pos + sp if j < n and seq[j] == partner: return True k = pos - sp if k >= 0 and seq[k] == partner: return True return False if tok in ("D", "E"): sp = 3 if tok == "D" else 4 i = pos - sp if i >= 0 and seq[i] == "K": return True # K on the other side (K would be j-side, unusual, but be permissive) j = pos + sp if j < n and seq[j] == "K": return True return False return False def classify_action(z: StapleState, z_next: StapleState) -> str: """Assign an action-type label to the transition z -> z_next. Robust to composite actions in :mod:`staplebridge.graph.neighbors` (anchor-assign additionally sets block_id in the same transition). """ if z.topology != z_next.topology: if z.topology == "linear" and z_next.topology == "stapled": return ACTION_ACTIVATE_TOPOLOGY return ACTION_UNKNOWN if z.sequence_tokens != z_next.sequence_tokens: # Which position(s) changed? n = min(len(z.sequence_tokens), len(z_next.sequence_tokens)) diffs = [i for i in range(n) if z.sequence_tokens[i] != z_next.sequence_tokens[i]] # Residue substitution action always changes exactly one position. for pos in diffs: tok = z_next.sequence_tokens[pos] if _has_lactam_partner(z.sequence_tokens, pos, tok): return ACTION_RESIDUE_SUBSTITUTION_MOTIF return ACTION_RESIDUE_SUBSTITUTION_OTHER if z.anchor_pair != z_next.anchor_pair: if z.anchor_pair is None and z_next.anchor_pair is not None: return ACTION_ASSIGN_ANCHOR return ACTION_REASSIGN_ANCHOR if z.block_id != z_next.block_id: return ACTION_ASSIGN_BLOCK return ACTION_NOOP class ReferenceEnergy: def __init__( self, peptide_prior: PeptidePriorBase, anchor_prior: AnchorPriorBase, block_prior: BlockPriorBase, geometry_oracle: GeometryOracleBase, catalog_index: dict[str, BuildingBlock], config: ReferenceEnergyConfig, ) -> None: self.peptide_prior = peptide_prior self.anchor_prior = anchor_prior self.block_prior = block_prior self.geometry_oracle = geometry_oracle self.catalog_index = catalog_index self.cfg = config self._weights = config.resolved_weights() self._action_weights = config.resolved_action_weights() # Optional debug ring buffer, off by default. self._debug_last: dict[str, float] | None = None # ------------------------------------------------------------------ # Decomposition (used by diagnostics; structure mirrors the paper's # reference process: peptide prior + anchor prior + block prior + # synthesis cost/SPPS + geometry term + action-type progress prior). # ------------------------------------------------------------------ def decompose( self, z: StapleState, z_next: StapleState, context: dict[str, Any] | None = None, precomputed_seq_score: float | None = None, ) -> dict[str, float]: context = context or {} block = self.catalog_index.get(z_next.block_id) if z_next.block_id else None peptide_ca = context.get("peptide_ca") w = self._weights if precomputed_seq_score is not None: e_seq_score = float(precomputed_seq_score) else: with STAGE_TIMER.section("ESM2_prior_time"): e_seq_score = self.peptide_prior.score_transition(z, z_next, context) e_seq = -w.seq * e_seq_score with STAGE_TIMER.section("anchor_prior_time"): anchor_score = self.anchor_prior.score_anchor( z_next.sequence_tokens, z_next.anchor_pair, context ) e_anchor = -w.anchor * anchor_score with STAGE_TIMER.section("block_prior_time"): block_score = self.block_prior.score_block( z_next.sequence_tokens, z_next.anchor_pair, block, context ) cost = block.cost_score if block else 1.0 spps = block.spps_score if block else 1.0 e_block = -w.block * block_score e_cost = self.cfg.eta_cost * cost e_spps = self.cfg.eta_spps * spps # Stage-aware geometry: only apply full ctype / cgeom when the # candidate is a real "stapling geometry" state — anchor + block set. # Otherwise the sentinel cgeom=10 for missing anchor/block would # dominate the reference energy and trap intermediate transitions. stage_aware = bool(getattr(self.cfg, "stage_aware_geometry", True)) geom_applicable = ( z_next.anchor_pair is not None and z_next.block_id is not None ) if stage_aware and not geom_applicable: ctype_bool = False cgeom_val = 0.0 e_type = 0.0 e_geom = 0.0 geom_status = "not_applicable" else: with STAGE_TIMER.section("geometry_time"): ctype_bool = bool( self.geometry_oracle.ctype( z_next.sequence_tokens, z_next.anchor_pair, block, peptide_ca=peptide_ca, ) ) cgeom_val = float( self.geometry_oracle.cgeom( z_next.sequence_tokens, z_next.anchor_pair, block, peptide_ca=peptide_ca, ) ) e_type = self.cfg.eta_type * (1.0 - float(ctype_bool)) e_geom = w.geom * self.cfg.eta_geom * cgeom_val geom_status = "applied" # Action-type progress prior — encourages the reference process to # actually build a stapled Khard state instead of drifting through # unrelated residue substitutions. action_type = classify_action(z, z_next) action_bias = self._action_weights.as_dict().get(action_type, 0.0) # For activate_topology, only reward if geometry is truly applicable # AND ctype passes; otherwise fall back to a small positive penalty. if action_type == ACTION_ACTIVATE_TOPOLOGY and not (geom_applicable and ctype_bool): action_bias = abs(self._action_weights.as_dict().get(action_type, 0.0)) * 0.5 e_progress = w.progress * self.cfg.eta_progress * action_bias # Decomposition convention: every term is a *positive contribution* # to the reference energy (so that "dominant" is just argmax). decomp = { "E_seq": float(e_seq), "E_anchor": float(e_anchor), "E_block": float(e_block), "E_cost": float(e_cost), "E_spps": float(e_spps), "E_type": float(e_type), "E_geom": float(e_geom), "E_progress": float(e_progress), "raw_cost": float(cost), "raw_spps": float(spps), "raw_cgeom": float(cgeom_val), "ctype_ok": bool(ctype_bool), "geom_status": geom_status, "action_type": action_type, } decomp["E_total"] = float( decomp["E_seq"] + decomp["E_anchor"] + decomp["E_block"] + decomp["E_cost"] + decomp["E_spps"] + decomp["E_type"] + decomp["E_geom"] + decomp["E_progress"] ) self._debug_last = decomp return decomp def compute_reference_energy( self, z: StapleState, z_next: StapleState, context: dict[str, Any] | None = None, ) -> float: return float(self.decompose(z, z_next, context)["E_total"]) def decompose_batch( self, z: StapleState, candidates: list[StapleState], context: dict[str, Any] | None = None, ) -> list[dict[str, float]]: """Decompose all candidates, sharing one batched peptide-prior call. Motivation: with the ESM2 delta prior each candidate would otherwise trigger its own model forward, even though most candidates in a neighborhood share the same source sequence. ``batch_score_transitions`` lets the prior amortize over the whole neighborhood. """ if not candidates: return [] with STAGE_TIMER.section("ESM2_prior_time"): seq_scores = self.peptide_prior.batch_score_transitions( z, candidates, context ) STAGE_TIMER.bump("candidates_scored", len(candidates)) return [ self.decompose(z, c, context, precomputed_seq_score=seq_scores[i]) for i, c in enumerate(candidates) ] @property def component_weights(self) -> ComponentWeights: return self._weights @property def action_weights(self) -> ActionProgressWeights: return self._action_weights @property def last_decomposition(self) -> dict[str, float] | None: return dict(self._debug_last) if self._debug_last else None