"""Terminal-state records, validation summaries and constrained beam decode. ``terminal_record`` scores a finished trajectory and flattens the result into one auditable row; ``summarize`` aggregates those rows into the validation metrics, including the selection metric; ``beam_decode`` runs plan-conditioned constrained decoding. """ from __future__ import annotations import json import math from collections import Counter from pathlib import Path from typing import Any, Iterable import numpy as np import torch from staplebridge.chemistry.edit_distance import weighted_edit_distance from staplebridge.chemistry.state import StapleState from staplebridge.data.schemas import LeadExample from staplebridge.hydrocarbon.actions import ( FailureReason, validate_hydrocarbon_staple, ) from staplebridge.hydrocarbon.terminal_energy import HydrocarbonTerminalEnergy from staplebridge.hydrocarbon.tokenizer import tokenize_sequence from staplebridge.models.control_kernel import ControlledKernel def state_key(state: StapleState) -> tuple[Any, ...]: return (tuple(state.sequence_tokens), state.anchor_pair, state.block_id, state.topology) def terminal_record( z0: StapleState, terminal: StapleState, lead: LeadExample, energy_fn: HydrocarbonTerminalEnergy, stack: dict[str, Any], **extra: Any, ) -> dict[str, Any]: block = stack["catalog_index"].get(terminal.block_id) if terminal.block_id else None verdict = validate_hydrocarbon_staple( terminal.sequence_tokens, None if terminal.anchor_pair is None else tuple(terminal.anchor_pair), block, stack["catalog"], ) energy, info = energy_fn(z0, terminal, lead) topology_label = None if verdict is FailureReason.OK and terminal.anchor_pair is not None: i, j = terminal.anchor_pair topology_label = f"{terminal.sequence_tokens[i]}-{terminal.sequence_tokens[j]}/i,i+{j-i}" record: dict[str, Any] = { "example_id": lead.example_id, "terminal_sequence": terminal.to_sequence(), "anchor_pair": terminal.anchor_pair, "block_id": terminal.block_id, "topology_label": topology_label, "stapled": terminal.topology == "stapled", "chemistry_valid": verdict is FailureReason.OK, "topology_status": verdict.value, "weighted_edit_distance": float(weighted_edit_distance(terminal, z0)), "terminal_energy": float(energy), "base_terminal_energy": float(info.get("base_terminal_energy", 0.0)), "endpoint_prior_energy": float(info.get("E_endpoint_total", 0.0)), "property_status": info.get("hydrocarbon_property_status"), } for key, value in info.items(): if key.startswith("hydrocarbon_"): record[key] = value scorer = getattr(energy_fn, "property_scorer", None) if scorer is not None: original = scorer.score_original_linear( z0.sequence_tokens, lead_key=str(lead.example_id) ) record.update(original) product = record.get( "hydrocarbon_permeability_penetrance_product_mean" ) baseline = record.get( "hydrocarbon_permeability_penetrance_original_linear" ) if product is not None and baseline is not None: record[ "hydrocarbon_permeability_penetrance_delta_vs_original_lead" ] = float(product) - float(baseline) for prop in ("solubility", "toxicity", "halflife"): product = record.get(f"hydrocarbon_{prop}_product_mean") baseline = record.get(f"hydrocarbon_{prop}_original_linear") if product is not None and baseline is not None: record[f"hydrocarbon_{prop}_delta_vs_original_lead"] = ( float(product) - float(baseline) ) record.update(extra) return record def finite_values(rows: Iterable[dict[str, Any]], key: str) -> list[float]: values: list[float] = [] for row in rows: value = row.get(key) if isinstance(value, (int, float)) and math.isfinite(float(value)): values.append(float(value)) return values def average(rows: Iterable[dict[str, Any]], key: str) -> float | None: values = finite_values(rows, key) return float(np.mean(values)) if values else None def summarize(rows: list[dict[str, Any]], selected: list[dict[str, Any]], n_leads: int) -> dict[str, Any]: selected_scored = [row for row in selected if row.get("property_status") == "scored"] failures = Counter() for row in selected: if not row["stapled"]: failures["not_stapled"] += 1 elif not row["chemistry_valid"]: failures[f"invalid:{row['topology_status']}"] += 1 elif row.get("property_status") != "scored": failures[str(row.get("property_status", "property_not_scored"))] += 1 mode_counts = Counter(row.get("topology_label") for row in selected if row.get("topology_label")) tox_threshold = 0.49 tox_values = finite_values(selected_scored, "hydrocarbon_toxicity_product_mean") return { "n_candidates": len(rows), "n_selected": len(selected), "chemistry_valid_rate": sum(bool(row["chemistry_valid"]) for row in rows) / len(rows) if rows else 0.0, "stapled_rate": sum(bool(row["stapled"]) for row in rows) / len(rows) if rows else 0.0, "accepted_at_lead_count": sum(bool(row["chemistry_valid"] and row["stapled"] and row.get("property_status") == "scored") for row in selected), "accepted_at_lead_rate": sum(bool(row["chemistry_valid"] and row["stapled"] and row.get("property_status") == "scored") for row in selected) / n_leads, "mean_original_linear_penetrance": average( selected_scored, "hydrocarbon_permeability_penetrance_original_linear", ), "mean_product_penetrance": average( selected_scored, "hydrocarbon_permeability_penetrance_product_mean", ), "mean_delta_penetrance_vs_original_lead": average( selected_scored, "hydrocarbon_permeability_penetrance_delta_vs_original_lead", ), "mean_delta_penetrance_vs_plan_precursor": average( selected_scored, "hydrocarbon_permeability_penetrance_delta_vs_plan_precursor", ), "mean_delta_penetrance": average( selected_scored, "hydrocarbon_permeability_penetrance_delta_vs_original_lead", ), "mean_delta_penetrance_definition": "vs_original_lead", "mean_delta_toxicity_log_only": average(selected_scored, "hydrocarbon_toxicity_delta_product_linear"), "mean_product_toxicity_log_only": average(selected_scored, "hydrocarbon_toxicity_product_mean"), "toxicity_violation_rate_log_only": sum(value > tox_threshold for value in tox_values) / len(tox_values) if tox_values else None, "mean_hemolysis_monitor": average(selected_scored, "hydrocarbon_hemolysis_product_mean"), "mean_halflife_monitor": average(selected_scored, "hydrocarbon_halflife_product_mean"), "mean_weighted_edit_distance": average(selected, "weighted_edit_distance"), "mean_penetrance_ez_sensitivity": average(selected_scored, "hydrocarbon_permeability_penetrance_EZ_abs_diff"), "max_penetrance_ez_sensitivity": max(finite_values(selected_scored, "hydrocarbon_permeability_penetrance_EZ_abs_diff"), default=None), "s5_s5_i4_count": mode_counts.get("S5-S5/i,i+4", 0), "s5_s5_i4_rate": mode_counts.get("S5-S5/i,i+4", 0) / n_leads, "r8_s5_i7_count": mode_counts.get("R8-S5/i,i+7", 0), "r8_s5_i7_rate": mode_counts.get("R8-S5/i,i+7", 0) / n_leads, "topology_distribution": dict(mode_counts), "failure_reasons": dict(failures), "backend_distribution": dict(Counter(row.get("hydrocarbon_property_backend", "not_scored") for row in selected)), "mode_distribution": dict(Counter(row.get("hydrocarbon_property_mode", "not_scored") for row in selected)), "endpoint_prior_energy_max_abs": max((abs(float(row["endpoint_prior_energy"])) for row in rows), default=0.0), } def select_lowest_energy(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: grouped: dict[str, list[dict[str, Any]]] = {} for row in rows: grouped.setdefault(str(row["example_id"]), []).append(row) return [min(group, key=lambda row: (row["terminal_energy"], row.get("trajectory_index", 0))) for _, group in sorted(grouped.items())] def beam_decode( lead: LeadExample, stack: dict[str, Any], kernel: ControlledKernel, horizon: int, beam_size: int, committed_plan: Any | None = None, ) -> list[tuple[StapleState, float, int]]: z0 = StapleState(sequence_tokens=tokenize_sequence(lead.linear_sequence)) peptide_ca = (lead.target_context or {}).get("peptide_ca") if committed_plan is None: plan_probe = stack["sampler"].sample_trajectory( z0, protected_positions=lead.protected_positions, context={ "protected_positions": lead.protected_positions, "peptide_ca": peptide_ca, }, horizon=0, ) committed_plan = plan_probe.plan context = { "protected_positions": lead.protected_positions, "peptide_ca": peptide_ca, "hydrocarbon_plan": committed_plan, } beam: list[tuple[StapleState, float, int]] = [(z0, 0.0, 0)] for t in range(horizon): expanded: list[tuple[StapleState, float, int]] = [] for state, log_probability, path_length in beam: candidates = stack["graph"].neighbors(state, protected_positions=lead.protected_positions) if not candidates: expanded.append((state, log_probability, path_length)) continue log_probs = torch.log_softmax(kernel.controlled_logits(state, candidates, t=t, context=context), dim=0) keep = min(beam_size, len(candidates)) values, indices = torch.topk(log_probs, k=keep) expanded.extend((candidates[int(index)], log_probability + float(value), path_length + 1) for value, index in zip(values.cpu(), indices.cpu())) best: dict[tuple[Any, ...], tuple[StapleState, float, int]] = {} for item in expanded: key = state_key(item[0]) if key not in best or item[1] > best[key][1]: best[key] = item beam = sorted(best.values(), key=lambda item: item[1], reverse=True)[:beam_size] return beam def write_json(path: Path, value: Any) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") def write_jsonl(path: Path, rows: Iterable[dict[str, Any]], mode: str = "w") -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open(mode, encoding="utf-8") as handle: for row in rows: handle.write(json.dumps(row, ensure_ascii=False) + "\n")