File size: 5,541 Bytes
bb6d2aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | from __future__ import annotations
import math
from typing import Any
import torch
from staplebridge.chemistry.state import StapleState
from staplebridge.reference.energy import (
ACTION_ACTIVATE_TOPOLOGY,
ACTION_ASSIGN_ANCHOR,
ACTION_ASSIGN_BLOCK,
ACTION_NOOP,
ACTION_REASSIGN_ANCHOR,
ACTION_RESIDUE_SUBSTITUTION_MOTIF,
ACTION_RESIDUE_SUBSTITUTION_OTHER,
ACTION_UNKNOWN,
ReferenceEnergy,
classify_action,
)
from staplebridge.utils.profiling import STAGE_TIMER
# Action groups for optional group-normalized sampling. Substitutions form one
# group; anchor/block/topology actions each form their own so that a single
# high-value transition is not swamped by a large fan-out of substitutions.
ACTION_GROUP_MAP = {
ACTION_NOOP: "noop",
ACTION_RESIDUE_SUBSTITUTION_MOTIF: "substitution",
ACTION_RESIDUE_SUBSTITUTION_OTHER: "substitution",
ACTION_ASSIGN_ANCHOR: "anchor",
ACTION_REASSIGN_ANCHOR: "anchor",
ACTION_ASSIGN_BLOCK: "block",
ACTION_ACTIVATE_TOPOLOGY: "topology",
ACTION_UNKNOWN: "substitution",
}
class ReferenceKernel:
def __init__(
self,
energy_model: ReferenceEnergy,
group_normalize: bool = False,
substitution_downweight: float = 1.0,
) -> None:
self.energy_model = energy_model
self.group_normalize = group_normalize
# When both structural (anchor/block/topology) and substitution
# candidates exist, multiply substitution probabilities by this factor
# (0..1) before renormalizing. 1.0 = no downweight (default).
self.substitution_downweight = float(substitution_downweight)
def _decompose_all(
self,
z: StapleState,
candidates: list[StapleState],
context: dict[str, Any] | None,
) -> list[dict[str, Any]]:
return self.energy_model.decompose_batch(z, candidates, context=context)
def reference_logits(
self,
z: StapleState,
candidates: list[StapleState],
context: dict[str, Any] | None = None,
) -> torch.Tensor:
with STAGE_TIMER.section("reference_logits_time"):
decomps = self._decompose_all(z, candidates, context)
out = torch.tensor([-d["E_total"] for d in decomps], dtype=torch.float32)
STAGE_TIMER.bump("reference_logits_candidates", len(candidates))
return out
def reference_probs(
self,
z: StapleState,
candidates: list[StapleState],
context: dict[str, Any] | None = None,
) -> torch.Tensor:
with STAGE_TIMER.section("reference_logits_time"):
decomps = self._decompose_all(z, candidates, context)
logits = torch.tensor([-d["E_total"] for d in decomps], dtype=torch.float32)
STAGE_TIMER.bump("reference_logits_candidates", len(candidates))
probs = torch.softmax(logits, dim=0)
if self.group_normalize:
# First sample action *group* uniformly over the groups actually
# present in the neighborhood (weighted by aggregated group logit),
# then softmax within group. Preserves the reference-energy
# ordering while preventing large substitution fan-outs from
# dominating the pmf.
groups: dict[str, list[int]] = {}
for idx, d in enumerate(decomps):
g = ACTION_GROUP_MAP.get(d.get("action_type", ACTION_UNKNOWN), "substitution")
groups.setdefault(g, []).append(idx)
group_probs = probs.clone()
group_probs.zero_()
# Aggregate group score = logsumexp of member logits
group_scores = {}
for g, idxs in groups.items():
gl = logits[idxs]
group_scores[g] = float(torch.logsumexp(gl, dim=0).item())
# Softmax over groups
g_keys = list(group_scores.keys())
g_logits = torch.tensor([group_scores[k] for k in g_keys], dtype=torch.float32)
g_pmf = torch.softmax(g_logits, dim=0)
for gi, gk in enumerate(g_keys):
idxs = groups[gk]
sub_logits = logits[idxs]
sub_pmf = torch.softmax(sub_logits, dim=0)
group_probs[idxs] = sub_pmf * g_pmf[gi]
probs = group_probs
if self.substitution_downweight != 1.0:
structural_present = any(
ACTION_GROUP_MAP.get(d.get("action_type", ACTION_UNKNOWN), "substitution")
in ("anchor", "block", "topology")
for d in decomps
)
if structural_present:
factors = torch.tensor(
[
self.substitution_downweight
if ACTION_GROUP_MAP.get(d.get("action_type", ACTION_UNKNOWN), "substitution")
== "substitution"
else 1.0
for d in decomps
],
dtype=torch.float32,
)
probs = probs * factors
s = probs.sum()
if float(s.item()) > 0.0:
probs = probs / s
return probs
def sample_next(
self,
z: StapleState,
candidates: list[StapleState],
context: dict[str, Any] | None = None,
) -> StapleState:
probs = self.reference_probs(z, candidates, context=context)
idx = torch.multinomial(probs, num_samples=1).item()
return candidates[idx]
|