diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..4901f9881b97eb5b90a9a7b11078b918721e03cb 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +figure/framework.pdf filter=lfs diff=lfs merge=lfs -text +figure/framework.png filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md index 154df8298fab5ecf322016157858e08cd1bccbe1..e943ffaa1ef8b63b34e2af42f19004506c629d65 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,335 @@ +# StapleBridge + +Official training code for **StapleBridge**, a chemistry-aware framework for optimizing existing peptide leads through hydrocarbon stapling. + +StapleBridge constructs a finite set of chemically and geometrically feasible stapling interventions for each peptide, learns to rank these interventions, and executes the selected plan with minimal sequence edits. + +This repository contains the **main StapleBridge training pipeline** and **one representative pretrained checkpoint**. + +## Framework + +[![StapleBridge framework: feasible stapling plans, exact finite-support control target, and plan execution](figure/framework.png)](figure/framework.pdf) + --- -license: apache-2.0 ---- + +## 1. What is included + +The code required to **train the main StapleBridge model**, and one +representative pretrained checkpoint. + +``` +release/staplebridge_training/ +├── README.md +├── THIRD_PARTY_NOTICES.md +├── requirements.txt +├── .gitignore +├── configs/ +│ └── staplebridge_main.yaml # configuration used for the checkpoint +├── scripts/ +│ └──train.py # canonical training entry point +├── staplebridge/ +│ ├── chemistry/ # design state, actions, edit distance +│ ├── data/ # split loading, schemas, catalog, vocab +│ ├── hydrocarbon/ # plans, catalog, q_ref, q*, q_theta, geometry +│ ├── models/ # policy / value nets, controlled kernel +│ ├── oracles/ # ESM2 + anchor/block reference priors +│ ├── reference/ # reference energy and kernel +│ ├── training/ # stack construction, main loop, losses +│ ├── integrations/ # PeptiVerse wrapper +│ └── utils/ # paths, profiling +├── data/ +│ └── README.md # expected input schema and file layout +└── checkpoints/ + └── staplebridge_seed42_best.pt # representative seed=42 model +``` + +The training code in `staplebridge/training/` is numerically identical to the +run that produced the shipped checkpoint. + + + +## 2. Environment setup + +The released checkpoint was trained under: + +| | Version | +| --- | --- | +| Python | 3.10.20 | +| PyTorch | 2.12.1+cu130 (CUDA 13.0) | +| NumPy | 2.0.2 | +| PyYAML | 6.0.3 | +| RDKit | 2026.03.5 | +| transformers | 4.46.0 | + +```bash +python -m venv .venv && source .venv/bin/activate +# Install torch first, matched to your CUDA build: https://pytorch.org +pip install -r requirements.txt +``` + + + +## 3. External model dependencies + +**No third-party model weights are bundled.** Four external resources must be +provided and referenced from the config: the ESM2-650M snapshot, the PeptiVerse +distribution, and the two SMILES encoders PeptiVerse depends on. Paths may be +absolute, or relative to the package root. + +### ESM2-650M (frozen sequence context) + +`facebook/esm2_t33_650M_UR50D`, used frozen — never fine-tuned. It supplies the +reference-process peptide prior and the V2 plan head's anchor/local-context +features. It is also the feature source for the plan head, which refuses to +build without it. + +```bash +huggingface-cli download facebook/esm2_t33_650M_UR50D \ + --local-dir models/esm2_t33_650M_UR50D +``` + +Then set, in `configs/staplebridge_main.yaml`: + +```yaml +reference_priors: + peptide: + model_name_or_path: models/esm2_t33_650M_UR50D +property_predictor: + esm_model_name_or_path: models/esm2_t33_650M_UR50D +``` + +The config runs the prior with `offline: true` and `strict_runtime: true`, so the +snapshot must already be on disk; training fails fast rather than downloading or +silently substituting a fallback. Licensed by Meta under the ESM2 terms. + +### PeptiVerse (property oracles) + +The main training objective optimises the PeptiVerse +permeability-penetrance E/Z product mean. Obtain the PeptiVerse checkout and its +classifier weights separately, then set: + +```yaml +property_predictor: + peptiverse_root: external/PeptiVerse + classifier_weight_root: external/PeptiVerse + manifest_path: external/PeptiVerse/basic_models.txt +``` + +Scoring is **strict**: `strict: true`, `enable_fallback: false`, +`allow_wt_token_fallback: false`. If the oracle stack cannot load, training +aborts — it never degrades to a heuristic. + +Toxicity, hemolysis and half-life are monitored only. Solubility and binding +affinity are excluded from the objective. + +### PeptideCLM-23M and ChemBERTa-77M (required) + +The `basic_models.txt` manifest selects predictors embedded with PeptideCLM and +ChemBERTa, so **both are required** — not optional. The permeability-penetrance +predictor that defines the objective is itself ChemBERTa-embedded. Loading fails +fast without them. + +```bash +huggingface-cli download aaronfeller/PeptideCLM-23M-all \ + --local-dir models/PeptideCLM-23M-all +huggingface-cli download DeepChem/ChemBERTa-77M-MLM \ + --local-dir models/ChemBERTa-77M-MLM +``` + +```yaml +property_predictor: + peptideclm_model_name_or_path: models/PeptideCLM-23M-all + chemberta_model_name_or_path: models/ChemBERTa-77M-MLM +``` + +`scripts/check_config.py` verifies all four before training starts. + +## 4. Data + +**The processed training and validation data are not included in this release** +and will be handled separately. No preprocessing, download or reconstruction +utilities are provided. + +Training reads two JSON Lines files, resolved from the config: + +``` +data/real/ # data.root +├── train.jsonl # data.train_file +└── valid.jsonl # data.valid_file +``` + +See [`data/README.md`](data/README.md) for the expected input schema — in +particular the required per-residue Cα coordinates, which staple-geometry +feasibility depends on. + +The reference protocol uses 4020 training and 111 validation leads; +`scripts/check_config.py` asserts those counts, so substituting a differently +sized dataset requires relaxing the check. + +## 5. Training command + +```bash + +python scripts/train.py \ + --config configs/staplebridge_main.yaml \ + --out-dir outputs/main_seed42 +``` + + + + + + diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000000000000000000000000000000000000..c9d425a20b2ef41e23efb7b2c66b5255cf684c18 --- /dev/null +++ b/THIRD_PARTY_NOTICES.md @@ -0,0 +1,70 @@ +# Third-party notices + +This release bundles no third-party model weights and no third-party source +code. The components below are **required at runtime** and must be obtained by +the user under their own licenses. See README.md §4. + +## ESM2-650M — `facebook/esm2_t33_650M_UR50D` + +Used frozen (never fine-tuned) as the sequence-context model: it provides the +reference-process peptide prior and the plan head's anchor/local-context +features. + +- Publisher: Meta AI (Fundamental AI Research Protein Team) +- Weights: not redistributed here; download from Hugging Face. +- License: the ESM2 model license from Meta. Review it before redistributing + weights or derivatives. +- Reference: Lin et al., "Evolutionary-scale prediction of atomic-level protein + structure with a language model", *Science* 379 (2023). + +## PeptiVerse — peptide property oracles + +Supplies the permeability-penetrance predictor that defines the main training +objective, plus the monitored toxicity / hemolysis / half-life predictors. + +- Weights and source: not redistributed here; obtain the PeptiVerse + distribution separately. +- License: as specified by the PeptiVerse authors. + +## PeptideCLM-23M — `aaronfeller/PeptideCLM-23M-all` + +Required. Supplies SMILES embeddings for several PeptiVerse predictors selected +by the official `basic_models.txt` manifest (including the half-life and +nonfouling models). + +- Weights: not redistributed here; download from Hugging Face. +- License: as published with the model. + +## ChemBERTa-77M — `DeepChem/ChemBERTa-77M-MLM` + +Required. Supplies SMILES embeddings for the permeability-penetrance predictor +that defines the main training objective, plus the toxicity, PAMPA and Caco-2 +models. + +- Weights: not redistributed here; download from Hugging Face. +- License: as published by DeepChem. +- Reference: Chithrananda et al., "ChemBERTa: Large-Scale Self-Supervised + Pretraining for Molecular Property Prediction" (2020). + +## Python dependencies + +Declared in `requirements.txt` and installed from PyPI, each under its own +license: + +| Package | License | +| --- | --- | +| PyTorch | BSD-3-Clause | +| NumPy | BSD-3-Clause | +| PyYAML | MIT | +| RDKit | BSD-3-Clause | +| transformers (Hugging Face) | Apache-2.0 | + +The PeptiVerse distribution brings its own further dependencies (scikit-learn, +XGBoost, MAPIE, pandas, joblib and others); those are governed by their +respective licenses and are not declared by this package. + +## Data + +No dataset is included in this release. The processed training and validation +splits are handled separately; nothing here downloads, reconstructs or +redistributes data. See `data/README.md`. diff --git a/checkpoints/staplebridge_seed42_best.pt b/checkpoints/staplebridge_seed42_best.pt new file mode 100644 index 0000000000000000000000000000000000000000..ea774510018f02a1f955956424ce91545f96e278 --- /dev/null +++ b/checkpoints/staplebridge_seed42_best.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:22d273164a3132d43617c51649947028f11859f90ca0bee28540edc9ad62a298 +size 1394383 diff --git a/configs/staplebridge_main.yaml b/configs/staplebridge_main.yaml new file mode 100644 index 0000000000000000000000000000000000000000..686598c69dee57999412b3e82dbfb7a662bc5309 --- /dev/null +++ b/configs/staplebridge_main.yaml @@ -0,0 +1,244 @@ +# Main StapleBridge training configuration (Full Exact-SB). +# +# This is the configuration that produced checkpoints/staplebridge_seed42_best.pt. +# Scale, objective, chemistry and checkpoint-selection settings are those of the +# reported run and should not be changed to reproduce it. +# +# Paths marked EXTERNAL ASSET must point at resources you provide (see README.md +# section 4 and data/README.md). Relative paths resolve against the package root, +# absolute paths are used as given. +chemistry: hydrocarbon +seed: 42 +exp_name: staplebridge_main_seed42 +train_n: 4020 +max_neighbors: 128 +data: + # EXTERNAL ASSET: official train/valid splits. See data/README.md. + root: data/real + train_file: train.jsonl + valid_file: valid.jsonl + max_lead_length: 25 +model: + mode: policy_tilt + emb_dim: 32 +hydrocarbon: + reference: + factorized_plan_reference: false + decode: + hierarchical_plan_ranking: true + catalog: + include_optional: false + actions: + max_neighbors: 128 + allow_noop: true + plan_control: + enabled: true + plans_per_lead: 4 + hidden_dim: 128 + loss_weight: 1.0 + target_temperature: 0.1 + exact_sb_objective: true + exact_sb_beta: 1.0 + plan_encoder_v2: true + validation_all_plans: true + exact_sb_cache: + enabled: true + path: outputs/cache/exact_sb_targets_full.sqlite + curriculum: + enabled: false + max_anchor_edits: 2 + num_trajectories_per_lead: 4 + prefer_existing_anchors: true + protect_positions: true + require_valid_terminal: true + plan_reference: + enabled: true + mode_prior: + prior_dir: staplebridge/hydrocarbon/data + dedup_version: sequence_deduplicated + use_smoothed: true + beta: 0.75 + unobserved_probability: 0.001 + bias: + first_anchor: 3.0 + second_anchor: 3.5 + anchor_assign: 4.0 + block_assign: 4.0 + topology_activation: 4.5 + off_plan_topology_activation: -4.5 + off_plan_substitution: -3.0 + off_plan_anchor_selection: -3.0 + off_plan_block_assign: -1.0 + noop: -1.0 + endpoint_prior: + enabled: false + # Unused: this prior is disabled (and force-disabled in build_stack); the + # empirical endpoint evidence enters only via plan_reference.mode_prior above. + prior_dir: null + dedup_version: sequence_deduplicated + use_smoothed: true + weight_pair: 0.0 + use_length: false + weight_length: 0.0 + use_relative_position: false + weight_relative_position: 0.0 + use_local_context: false + weight_local_context: 0.0 + local_context_window: 2 + max_component_energy: 25.0 + terminal_energy: + invalid_topology_penalty: 10.0 + penalize_unstapled: true + property: + enabled: true + enable_developability_constraints: false + enable_halflife_preservation: false + penetrance_weight: 5.0 + toxicity_guard_enabled: false + toxicity_threshold: 0.49 + toxicity_guard_weight: 0.0 + input_convention: neutral_canonical + product_geometries: + - E + - Z + product_aggregation: mean + ez_uncertainty: absolute_difference + objective: permeability_penetrance + log_only: + - toxicity + - hemolysis + - halflife + excluded: + - solubility + - binding_affinity + geometry: + sentinel_cgeom: 10.0 +reference_priors: + strict_no_mock: true + peptide: + backend: esm2_delta + # EXTERNAL ASSET: local frozen ESM2-650M snapshot directory (not redistributed). + model_name_or_path: models/esm2_t33_650M_UR50D + # Persistent ESM2 prior cache (SQLite, WAL). Created on first use; reusing a + # warm cache across runs is a pure speedup and does not change any value. + cache_path: outputs/cache/esm2_peptide_prior.sqlite + device: cuda:0 + offline: true + strict_runtime: true + temperature: 1.0 + max_batch_size: 128 + ncaa_policy: + mode: canonical_surrogate + surrogates: + S5: A + R8: A + R5: A + S8: A + unknown_penalty: 0.2 + anchor: + backend: motif_support_geometry + allowed_spacings: + - 4 + - 7 + max_motif_edits: 2 + avoid_protected: true + project_motif_for_geometry: true + weights: + valid_spacing: 1.0 + existing_motif: 2.0 + one_edit_motif: 1.2 + two_edit_motif: 0.4 + protected_penalty: 4.0 + geometry_surrogate: 1.0 + block: + backend: catalog_scored + max_motif_edits: 2 + weights: + spacing: 2.0 + residue_compatibility: 2.0 + motif_edit_distance: 1.0 + synthetic_accessibility: 0.5 + spps: 0.5 + cost: 0.2 +reference: + eta_cost: 0.5 + eta_spps: 0.3 + eta_type: 5.0 + eta_geom: 1.0 +training: + epochs: 10 + lr: 0.001 + device: cuda:0 + horizon: 8 + trajectories_per_lead: 4 + chunk_size: 32 + eps_geom: 2.5 + lambda_close: 1.0 + lambda_edit: 0.2 + lambda_cost: 0.2 + infeasible_penalty: 20.0 + grad_clip_norm: 1.0 +property_predictor: + backend: peptiverse + # EXTERNAL ASSET: PeptiVerse checkout + its classifier weights (not redistributed). + peptiverse_root: external/PeptiVerse + classifier_weight_root: external/PeptiVerse + manifest_path: external/PeptiVerse/basic_models.txt + hf_cache_dir: models/hf + esm_model_name_or_path: models/esm2_t33_650M_UR50D + peptideclm_model_name_or_path: models/PeptideCLM-23M-all + chemberta_model_name_or_path: models/ChemBERTa-77M-MLM + device: cuda:0 + strict: true + enable_fallback: false + allow_wt_token_fallback: false + cache_enabled: true + offline: true + uncertainty: false + mode: smiles +validation: + enabled: true + every_epoch: true + n_leads: 111 + max_lead_length: 25 + beam_size: 8 + horizon: 8 + seed: 42 +checkpointing: + save_last: true + save_best: true + save_every_epoch: true + save_optimizer: true + save_rng_state: true + best_metric: q_star_vs_q_theta_kl + mode: min +early_stopping: + # All 10 epochs must complete. scripts/train.py ignores this block entirely + # and never early-stops. + enabled: false + metric: q_star_vs_q_theta_kl + mode: min + patience: 3 + min_delta: 0.001 + restore_best: false +guardrails: + chemistry_valid_min: 0.9 + stapled_min: 0.9 + edit_distance_target: 4.0 + edit_distance_tolerance: 0.5 + require_both_topologies: true + no_block_set_worsening_tolerance: 0.01 +edit_constraints: + max_edit_budget: 6.0 + min_sequence_identity: 0.6 + allow_protected_edits: false +wandb: + # Disabled in the official run, and no released code reads this block. + # Retained only so the config matches the one that produced the checkpoint. + enabled: false + mode: disabled +profiling: + # Per-32-lead chunk timing breakdown printed as `[chunk timing] {...}`. + # Wall-clock instrumentation only; does not affect loss, RNG, sampling, + # batch/chunk size, or model. See train_enabled_epoch for the CUDA-sync note. + chunk_timing: true diff --git a/data/README.md b/data/README.md new file mode 100644 index 0000000000000000000000000000000000000000..985fffb7e0197a2bde37238f140c94b16f103be8 --- /dev/null +++ b/data/README.md @@ -0,0 +1,53 @@ +# Data + +**The processed training and validation data are not included in this release.** +They will be handled separately. This directory is a placeholder; nothing here +downloads, reconstructs, or redistributes a dataset. + +## Where to place files + +Training reads exactly two files, resolved from the config +(`configs/staplebridge_main.yaml`): + +```yaml +data: + root: data/real # relative paths resolve against the release root + train_file: train.jsonl + valid_file: valid.jsonl + max_lead_length: 25 +``` + +giving the default layout: + +``` +data/real/ +├── train.jsonl +└── valid.jsonl +``` + +Any location works — set `data.root` to an absolute path if you prefer. + +## Expected input schema + +JSON Lines: one object per line, loaded into +`staplebridge.data.schemas.LeadExample` by +`staplebridge.data.dataset.load_leads`. + +| Field | Type | Required | Meaning | +| --- | --- | --- | --- | +| `example_id` | str | yes | Unique lead identifier; also the property-cache key. | +| `linear_sequence` | str | yes | Linear lead peptide, one letter per residue. | +| `protected_positions` | list[int] | yes (may be `[]`) | 0-based positions edits must not touch; enforced as a hard constraint. | +| `target_context` | object | yes in practice | Must contain `peptide_ca` (below). | +| `target_id` | str \| null | no | Binding-partner identifier. | +| `preferred_property_direction` | object | no | Not used by the main objective. | +| `thresholds` | object | no | Not used by the main objective. | +| `known_active_motif_positions` | list[int] \| null | no | Optional motif annotation. | + +`target_context.peptide_ca` must be a list of `[x, y, z]` Cα coordinates, one +per residue of `linear_sequence`, in order. Staple-geometry feasibility (the +`ca_window` span check gating every candidate plan) is computed from these; a +lead without them cannot yield a feasible plan support. Other keys in +`target_context` are ignored. + +Leads longer than `max_lead_length` are filtered out before use. diff --git a/figure/framework.pdf b/figure/framework.pdf new file mode 100644 index 0000000000000000000000000000000000000000..2b7bdd605f29c6e0229cd11a0f86c383afdcccc1 --- /dev/null +++ b/figure/framework.pdf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d7fabc5efb91b9cef419d11bb646493004ef796d043c4863c19b678dcce470a2 +size 884444 diff --git a/figure/framework.png b/figure/framework.png new file mode 100644 index 0000000000000000000000000000000000000000..2488cb8f97ddfa2bd3493bcc1b975b8df714112d --- /dev/null +++ b/figure/framework.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6d6032d2c83de1513c5aa4f29dfdf5f106fe1e0df019319e65e69a7c79ce8eb7 +size 444427 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..f23788967058942fb4e7726c7cc3a79d226335bb --- /dev/null +++ b/requirements.txt @@ -0,0 +1,21 @@ +# Direct dependencies of this package, pinned to the versions the shipped +# checkpoint was trained under. Install torch first, matched to your CUDA build: +# https://pytorch.org +# +# Reference environment: Python 3.10.20, torch 2.12.1+cu130 (CUDA 13.0). + +torch==2.12.1 +numpy==2.0.2 +PyYAML==6.0.3 + +# Strict SMILES construction and validation of stapled products. +rdkit==2026.3.5 + +# Loads the frozen ESM2-650M sequence-context model. Weights are not bundled; +# see README.md section 4. +transformers==4.46.0 + +# Not listed here on purpose: the PeptiVerse property oracles are an external +# dependency that brings its own requirements (scikit-learn, XGBoost, MAPIE, +# pandas, joblib). Nothing in this package imports them directly, so install +# them from the PeptiVerse distribution instead. See README.md section 4. diff --git a/scripts/train.py b/scripts/train.py new file mode 100644 index 0000000000000000000000000000000000000000..e15bc6f050bab2b9b160911ba6a6824b9a952f2e --- /dev/null +++ b/scripts/train.py @@ -0,0 +1,413 @@ +#!/usr/bin/env python +"""Training entry point for the main StapleBridge model (Full Exact-SB). + +This is the orchestrator that produced the released checkpoint: + +* full 4020-lead training each epoch; +* full validation on every validation lead (111), every epoch; +* a fixed 10 epochs, no early stopping (all 10 always run); +* two checkpoints maintained independently, each updated after that epoch's + validation pass: + - ``best_kl.pt``: minimum ``q_star_vs_q_theta_kl`` -- the selection rule; + - ``best_pv.pt``: maximum ``mean_delta_penetrance_vs_original_lead``, + recorded for monitoring only and not used to select the released model; +* per-epoch logging of both metrics with their running bests. + +``checkpoints/staplebridge_seed42_best.pt`` is the ``best_kl.pt`` of this run: +the epoch minimising ``q_star_vs_q_theta_kl`` on the validation split. Model, +loss, Exact-SB, property scoring, decoding and every other training setting are +read from the config. +""" +from __future__ import annotations + +import argparse +import gc +import json +import random +import sys +import time +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import yaml + +PACKAGE_ROOT = Path(__file__).resolve().parents[1] +if str(PACKAGE_ROOT) not in sys.path: + sys.path.insert(0, str(PACKAGE_ROOT)) + +from staplebridge.data.dataset import load_leads # noqa: E402 +from staplebridge.hydrocarbon.exact_sb_cache import build_cache_from_config # noqa: E402 +from staplebridge.hydrocarbon.plan_control import ( # noqa: E402 + HydrocarbonPlanControlConfig, build_hydrocarbon_plan_head, +) +from staplebridge.hydrocarbon.property_energy import ( # noqa: E402 + HydrocarbonPropertyEnergyConfig, + HydrocarbonPropertyScorer, + required_original_lead_properties, +) +from staplebridge.hydrocarbon.tokenizer import tokenize_sequence # noqa: E402 +from staplebridge.training.main_loop import ( # noqa: E402 + train_enabled_epoch, validate_enabled, +) +from staplebridge.training.records import write_json, write_jsonl # noqa: E402 +from staplebridge.training.stack import ( # noqa: E402 + build_energy, build_models, build_predictor, build_stack, load_config, + seed_everything, select_leads, +) + +FULL_TRAIN_N = 4020 +SEED = 42 +EPOCHS = 10 +KL_KEY = "q_star_vs_q_theta_kl" +PV_KEY = "mean_delta_penetrance_vs_original_lead" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, + default=PACKAGE_ROOT / "configs/staplebridge_main.yaml") + parser.add_argument("--out-dir", type=Path, + default=PACKAGE_ROOT / "outputs/main_seed42") + parser.add_argument("--resume", type=Path, default=None, + help="Checkpoint to resume from (e.g. checkpoints/epoch_008.pt). " + "Restores model/optimizer/RNG so remaining epochs are identical " + "to an uninterrupted run; appends to the existing metrics/log.") + return parser.parse_args() + + +def require(condition: bool, message: str) -> None: + if not condition: + raise SystemExit(message) + + +def rng_payload() -> dict[str, Any]: + return { + "python_random_state": random.getstate(), + "numpy_random_state": np.random.get_state(), + "torch_rng_state": torch.get_rng_state(), + "cuda_rng_state_all": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, + } + + +def build_original_cache(config: dict[str, Any], leads: list[Any], out_dir: Path) -> dict: + """Compute the configured unedited-lead property cache once.""" + path = out_dir / "original_linear_cache.jsonl" + wrapper, _ = build_predictor(config) + cache: dict[tuple[Any, ...], dict[str, Any]] = {} + scorer = HydrocarbonPropertyScorer(wrapper, original_linear_cache=cache) + property_cfg = HydrocarbonPropertyEnergyConfig.from_dict( + (((config.get("hydrocarbon") or {}).get("terminal_energy") or {}).get("property")) + ) + properties = required_original_lead_properties(property_cfg) + started = time.perf_counter() + if ( + property_cfg.enable_developability_constraints + or property_cfg.enable_halflife_preservation + or property_cfg.enable_joint_perm_halflife_support + ): + scorer.prefetch( + properties, + [ + scorer.original_linear_smiles(tokenize_sequence(lead.linear_sequence)) + for lead in leads + ], + ) + for index, lead in enumerate(leads): + scorer.score_original_linear( + tokenize_sequence(lead.linear_sequence), + lead_key=str(lead.example_id), + properties=properties, + ) + if (index + 1) % 512 == 0 or index + 1 == len(leads): + print(f"[original baseline] {index + 1}/{len(leads)}", flush=True) + rows = [{"lead_key": k[0], "tokens": list(k[1]), "scores": v} for k, v in cache.items()] + write_jsonl(path, rows) + del scorer, wrapper + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + print(f"[original baseline] built {len(cache)} entries in {time.perf_counter() - started:.1f}s", flush=True) + return cache + + +def load_original_cache(path: Path) -> dict[tuple[Any, ...], dict[str, Any]]: + cache: dict[tuple[Any, ...], dict[str, Any]] = {} + with path.open(encoding="utf-8") as handle: + for line in handle: + row = json.loads(line) + cache[(row["lead_key"], tuple(row["tokens"]))] = dict(row["scores"]) + return cache + + +def save_checkpoint(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(path.suffix + ".tmp") + torch.save(payload, tmp) + tmp.replace(path) + + +def main() -> None: + args = parse_args() + config = load_config(args.config.resolve()) + out_dir = args.out_dir.resolve() + resuming = args.resume is not None + + if resuming: + require(out_dir.is_dir(), f"resume requires existing output dir: {out_dir}") + require(args.resume.is_file(), f"resume checkpoint not found: {args.resume}") + else: + # Never overwrite prior outputs: require a fresh/empty directory. + if out_dir.exists(): + leftovers = [p for p in out_dir.iterdir()] + require(not leftovers, f"output directory not empty (refusing to overwrite): {out_dir}") + out_dir.mkdir(parents=True, exist_ok=True) + + require(int(config.get("train_n", -1)) == FULL_TRAIN_N, "train_n must be 4020") + require(int(config["training"]["epochs"]) == EPOCHS, f"epochs must be {EPOCHS}") + device = torch.device(str(config["training"]["device"])) + if device.type == "cuda": + require(torch.cuda.is_available(), f"CUDA unavailable for {device}") + torch.cuda.set_device(device) + + (out_dir / "resolved_config.yaml").write_text(yaml.safe_dump(config, sort_keys=False)) + seed_everything(SEED) + + # ---- data: full train (4020) + FULL valid (all 111, not the 32 subset) -- + train_leads = select_leads( + Path(config["data"]["root"]) / config["data"]["train_file"], + FULL_TRAIN_N, int(config["data"]["max_lead_length"]), + ) + valid_cap = int(config["validation"]["max_lead_length"]) + valid_leads = [ + lead for lead in load_leads(Path(config["data"]["root"]) / config["data"]["valid_file"]) + if len(lead.linear_sequence) <= valid_cap + ] + print(f"[data] train={len(train_leads)} full_valid={len(valid_leads)} (cap len<={valid_cap})", flush=True) + + # ---- original-lead penetrance cache (train + valid) for delta-PV -------- + original_path = out_dir / "original_linear_cache.jsonl" + if resuming and original_path.is_file(): + original_cache = load_original_cache(original_path) + print(f"[original baseline] reused {len(original_cache)} entries from {original_path}", flush=True) + else: + original_cache = build_original_cache(config, train_leads + valid_leads, out_dir) + + # ---- model + energy stack (identical construction to the standard run) -- + stack = build_stack(config, SEED) + wrapper, _ = build_predictor(config) + scorer = HydrocarbonPropertyScorer(wrapper, original_linear_cache=original_cache) + energy_fn = build_energy(config, stack, scorer) + policy, value, kernel, optimizer = build_models(config, stack, device) + head = build_hydrocarbon_plan_head( + config, int(config["model"]["emb_dim"]), device, + esm2_prior=stack["reference_priors"].peptide, + ) + optimizer.add_param_group({"params": list(head.parameters())}) + parameters = list(policy.parameters()) + list(value.parameters()) + list(head.parameters()) + plan_cfg = HydrocarbonPlanControlConfig.from_config(config) + plan_rng = random.Random(SEED) + + # ---- reuse the existing persistent q* cache (fingerprint unchanged) ----- + exact_sb_cache = build_cache_from_config(config, catalog=stack["catalog"], repo_root=PACKAGE_ROOT) + if exact_sb_cache.enabled: + print(f"[exact-sb cache] {json.dumps(exact_sb_cache.describe(), ensure_ascii=False)}", flush=True) + + def make_payload(epoch: int) -> dict[str, Any]: + return { + "epoch": epoch, "config": config, "plan_control_enabled": True, + "policy_state_dict": policy.state_dict(), "value_state_dict": value.state_dict(), + "plan_head_state_dict": head.state_dict(), "optimizer_state_dict": optimizer.state_dict(), + "plan_rng_state": plan_rng.getstate(), **rng_payload(), + } + + metrics_path = out_dir / "metrics.jsonl" + history: list[dict[str, Any]] = [] + best_kl = {"value": float("inf"), "epoch": None} + best_pv = {"value": float("-inf"), "epoch": None} + start_epoch = 0 + + if resuming: + ckpt = torch.load(args.resume, map_location=device, weights_only=False) + policy.load_state_dict(ckpt["policy_state_dict"]) + value.load_state_dict(ckpt["value_state_dict"]) + head.load_state_dict(ckpt["plan_head_state_dict"]) + optimizer.load_state_dict(ckpt["optimizer_state_dict"]) + plan_rng.setstate(ckpt["plan_rng_state"]) + random.setstate(ckpt["python_random_state"]) + np.random.set_state(ckpt["numpy_random_state"]) + torch.set_rng_state(ckpt["torch_rng_state"].cpu()) + if torch.cuda.is_available() and ckpt.get("cuda_rng_state_all") is not None: + torch.cuda.set_rng_state_all([s.cpu() for s in ckpt["cuda_rng_state_all"]]) + start_epoch = int(ckpt["epoch"]) + # Rebuild running bests + history from the persisted per-epoch metrics so + # best_kl/best_pv provenance carries across the resume boundary. + for line in metrics_path.read_text().splitlines(): + rec = json.loads(line) + if int(rec["epoch"]) <= start_epoch: + history.append(rec) + if history: + last = history[-1] + best_kl = {"value": float(last["best_kl_value"]), "epoch": int(last["best_kl_epoch"])} + best_pv = {"value": float(last["best_pv_value"]), "epoch": int(last["best_pv_epoch"])} + print( + f"[resume] from {args.resume} completed_epoch={start_epoch}; " + f"best_kl={best_kl['value']:.6f}@ep{best_kl['epoch']} " + f"best_pv={best_pv['value']:.6f}@ep{best_pv['epoch']}", + flush=True, + ) + run_started = time.perf_counter() + + for epoch in range(start_epoch, EPOCHS): + policy.train(); value.train(); head.train() + train_started = time.perf_counter() + rows, train_metrics = train_enabled_epoch( + train_leads, config, stack, energy_fn, policy, kernel, optimizer, + parameters, head, plan_cfg, plan_rng, epoch, exact_sb_cache=exact_sb_cache, + ) + train_seconds = time.perf_counter() - train_started + write_jsonl(out_dir / "training" / f"epoch_{epoch + 1:03d}_candidates.jsonl", rows) + joint_train_audit = train_metrics.get("joint_perm_halflife_support") + if joint_train_audit and epoch == 0: + write_jsonl( + out_dir / "joint_support_audit" / "train4020_per_lead.jsonl", + joint_train_audit["per_lead"], + ) + write_json( + out_dir / "joint_support_audit" / "train4020_summary.json", + {k: v for k, v in joint_train_audit.items() if k != "per_lead"}, + ) + + policy.eval(); value.eval(); head.eval() + valid_started = time.perf_counter() + selected, validation = validate_enabled( + valid_leads, config, stack, energy_fn, policy, kernel, head, + exact_sb_cache=exact_sb_cache, + ) + valid_seconds = time.perf_counter() - valid_started + valid_dir = out_dir / "validation" / f"epoch_{epoch + 1:03d}" + write_jsonl(valid_dir / "selected.jsonl", selected) + write_json(valid_dir / "summary.json", validation) + joint_valid_audit = validation.get("joint_perm_halflife_support") + if joint_valid_audit: + write_jsonl( + valid_dir / "joint_support_per_lead.jsonl", + joint_valid_audit["per_lead"], + ) + write_json( + valid_dir / "joint_support_summary.json", + {k: v for k, v in joint_valid_audit.items() if k != "per_lead"}, + ) + + kl = float(validation[KL_KEY]) + pv = float(validation[PV_KEY]) + guard_cfg = dict(config.get("guardrails") or {}) + both_topologies_present = bool( + validation.get("s5_s5_i4_count", 0) > 0 + and validation.get("r8_s5_i7_count", 0) > 0 + ) + checkpoint_eligible = bool( + not guard_cfg.get("require_both_topologies", False) + or both_topologies_present + ) + improved_kl = checkpoint_eligible and kl < best_kl["value"] + improved_pv = checkpoint_eligible and pv > best_pv["value"] + if improved_kl: + best_kl = {"value": kl, "epoch": epoch + 1} + save_checkpoint(out_dir / "checkpoints" / "best_kl.pt", make_payload(epoch + 1)) + if improved_pv: + best_pv = {"value": pv, "epoch": epoch + 1} + save_checkpoint(out_dir / "checkpoints" / "best_pv.pt", make_payload(epoch + 1)) + # Per-epoch + latest snapshots so no epoch is lost (fresh dir; nothing + # is overwritten across runs). + payload = make_payload(epoch + 1) + save_checkpoint(out_dir / "checkpoints" / f"epoch_{epoch + 1:03d}.pt", payload) + save_checkpoint(out_dir / "checkpoints" / "latest.pt", payload) + + record = { + "epoch": epoch + 1, + "valid_kl": kl, + "valid_delta_penetrance": pv, + "best_kl_epoch": best_kl["epoch"], "best_kl_value": best_kl["value"], "kl_improved": improved_kl, + "best_pv_epoch": best_pv["epoch"], "best_pv_value": best_pv["value"], "pv_improved": improved_pv, + "valid_product_penetrance": validation.get("mean_product_penetrance"), + "valid_top1_chemistry_valid_rate": validation.get("top1_chemistry_valid_rate"), + "valid_top1_stapled_rate": validation.get("top1_stapled_rate"), + "valid_q_star_top1_agreement": validation.get("q_star_top1_agreement"), + "valid_q_star_spearman": validation.get("q_star_spearman"), + "valid_edit_distance": validation.get("mean_weighted_edit_distance"), + "both_topologies_present": both_topologies_present, + "checkpoint_eligible": checkpoint_eligible, + "train_loss": train_metrics.get("loss"), + "train_plan_loss": train_metrics.get("plan_loss"), + "train_q_star_vs_q_theta_kl": train_metrics.get("q_star_vs_q_theta_kl"), + "train_seconds": train_seconds, + "valid_seconds": valid_seconds, + "epoch_seconds": train_seconds + valid_seconds, + "train_stage_seconds": train_metrics.get("stage_seconds"), + "exact_sb_cache": exact_sb_cache.describe() if exact_sb_cache.enabled else {"enabled": False}, + "train_joint_perm_halflife_support": ( + {k: v for k, v in (joint_train_audit or {}).items() if k != "per_lead"} + if joint_train_audit + else None + ), + "valid_joint_perm_halflife_support": ( + { + k: v + for k, v in (joint_valid_audit or {}).items() + if k != "per_lead" + } + if joint_valid_audit + else None + ), + } + history.append(record) + write_jsonl(metrics_path, [record], mode="a") + print( + f"[epoch {epoch + 1}/{EPOCHS}] " + f"valid_KL={kl:.6f} (best {best_kl['value']:.6f} @ep{best_kl['epoch']}" + f"{' NEW' if improved_kl else ''}) " + f"valid_deltaPV={pv:.6f} (best {best_pv['value']:.6f} @ep{best_pv['epoch']}" + f"{' NEW' if improved_pv else ''}) " + f"epoch_seconds={record['epoch_seconds']:.1f}", + flush=True, + ) + # Semantics-neutral memory hygiene: reclaim inter-epoch CUDA cache / + # Python garbage so allocator fragmentation does not accumulate across + # the 10 epochs. Does not touch weights, RNG, or any cached value. + gc.collect() + if device.type == "cuda": + torch.cuda.empty_cache() + print( + f"[gpu] epoch {epoch + 1} allocated={torch.cuda.memory_allocated(device) / 2**20:.0f}MiB " + f"reserved={torch.cuda.memory_reserved(device) / 2**20:.0f}MiB", + flush=True, + ) + + summary = { + "exp_name": config.get("exp_name"), + "epochs": EPOCHS, + "early_stopping": False, + "train_n": FULL_TRAIN_N, + "n_valid_leads": len(valid_leads), + "seed": SEED, + "best_kl": best_kl, + "best_pv": best_pv, + "kl_curve": [(r["epoch"], r["valid_kl"]) for r in history], + "pv_curve": [(r["epoch"], r["valid_delta_penetrance"]) for r in history], + "runtime_seconds": time.perf_counter() - run_started, + "gpu_peak_mib": (torch.cuda.max_memory_allocated(device) / 2 ** 20) if device.type == "cuda" else 0, + "checkpoints": { + "best_kl": str(out_dir / "checkpoints" / "best_kl.pt"), + "best_pv": str(out_dir / "checkpoints" / "best_pv.pt"), + }, + "history": history, + } + write_json(out_dir / "run_summary.json", summary) + print("\n" + json.dumps({k: v for k, v in summary.items() if k != "history"}, indent=2, default=str), flush=True) + + +if __name__ == "__main__": + main() diff --git a/staplebridge/.DS_Store b/staplebridge/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..31edbff8f50ebad4c5cdeb862f97ffa6c10ee070 Binary files /dev/null and b/staplebridge/.DS_Store differ diff --git a/staplebridge/__init__.py b/staplebridge/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..289ba634c18da10775bbee05d1214c2de5e8cee1 --- /dev/null +++ b/staplebridge/__init__.py @@ -0,0 +1,4 @@ +"""StapleBridge: reference-guided discrete Schrodinger bridge for stapled peptides.""" + +__all__ = ["__version__"] +__version__ = "0.1.0" diff --git a/staplebridge/chemistry/__init__.py b/staplebridge/chemistry/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f3b0d32f802e05f5e1927433a19f527fd5a43b92 --- /dev/null +++ b/staplebridge/chemistry/__init__.py @@ -0,0 +1 @@ +"""""" diff --git a/staplebridge/chemistry/actions.py b/staplebridge/chemistry/actions.py new file mode 100644 index 0000000000000000000000000000000000000000..bf408b7feabffa0b4bd68e41ccdb35b2ee58f1e3 --- /dev/null +++ b/staplebridge/chemistry/actions.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from staplebridge.chemistry.state import StapleState + + +class Action: + """Base action for editing StapleState.""" + + def apply(self, state: StapleState) -> StapleState: + raise NotImplementedError + + +@dataclass +class ResidueSubstitutionAction(Action): + position: int + new_token: str + + def apply(self, state: StapleState) -> StapleState: + ns = state.copy() + ns.sequence_tokens[self.position] = self.new_token + return ns + + +@dataclass +class NcAASubstitutionAction(Action): + position: int + block_or_token: str + + def apply(self, state: StapleState) -> StapleState: + ns = state.copy() + ns.sequence_tokens[self.position] = self.block_or_token + return ns + + +@dataclass +class AnchorAssignAction(Action): + i: int + j: int + + def apply(self, state: StapleState) -> StapleState: + ns = state.copy() + ns.anchor_pair = (self.i, self.j) + return ns + + +@dataclass +class AnchorReassignAction(Action): + i: int + j: int + + def apply(self, state: StapleState) -> StapleState: + ns = state.copy() + ns.anchor_pair = (self.i, self.j) + return ns + + +@dataclass +class BlockAssignAction(Action): + block_id: str + + def apply(self, state: StapleState) -> StapleState: + ns = state.copy() + ns.block_id = self.block_id + return ns + + +@dataclass +class TopologyActivationAction(Action): + def apply(self, state: StapleState) -> StapleState: + ns = state.copy() + ns.topology = "stapled" + return ns + + +@dataclass +class NoOpAction(Action): + def apply(self, state: StapleState) -> StapleState: + return state.copy() diff --git a/staplebridge/chemistry/edit_distance.py b/staplebridge/chemistry/edit_distance.py new file mode 100644 index 0000000000000000000000000000000000000000..e37f674c7b3c4f1d951551d7a350cc170041b84d --- /dev/null +++ b/staplebridge/chemistry/edit_distance.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from staplebridge.chemistry.state import StapleState + + +def weighted_edit_distance( + state: StapleState, + lead: StapleState, + token_weight: float = 1.0, + anchor_weight: float = 1.0, + topology_weight: float = 0.5, + block_weight: float = 0.5, +) -> float: + token_diff = sum(a != b for a, b in zip(state.sequence_tokens, lead.sequence_tokens)) + anchor_pen = 0.0 if state.anchor_pair == lead.anchor_pair else 1.0 + top_pen = 0.0 if state.topology == lead.topology else 1.0 + block_pen = 0.0 if state.block_id == lead.block_id else 1.0 + return token_weight * token_diff + anchor_weight * anchor_pen + topology_weight * top_pen + block_weight * block_pen diff --git a/staplebridge/chemistry/edit_metrics.py b/staplebridge/chemistry/edit_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..8a331f1fd9b412d51c12146cc5f683f292c27812 --- /dev/null +++ b/staplebridge/chemistry/edit_metrics.py @@ -0,0 +1,97 @@ +"""Edit-distance + minimal-action diagnostics. + +These helpers extend the simple ``weighted_edit_distance`` with structured +information needed by the diagnostics protocol: + + edit_distance : raw token-level edits (positions where seq differs from lead) + edited_positions : list[int] + sequence_identity : 1 - (edit_distance / max(len(lead), len(seq))) + protected_edit_violation: any edited position falls in the protected set + anchor_created_by_edit : the anchor's i/j residues were edited to create the motif + anchor_used_existing_motif: the anchor's i/j residues were already correct in the lead + edit_script : list of {pos, lead_aa, designed_aa} + +This module is read-only with respect to StapleState; it does not mutate +anything. +""" + +from __future__ import annotations + +from typing import Any + +from staplebridge.chemistry.state import StapleState +from staplebridge.data.catalog import default_catalog +from staplebridge.data.schemas import BuildingBlock + + +def _motif_for_block(block_id: str | None, catalog: dict[str, BuildingBlock] | None) -> dict[str, Any] | None: + if block_id is None: + return None + if catalog is None: + for b in default_catalog(): + if b.block_id == block_id: + return b.motif + return None + block = catalog.get(block_id) + return block.motif if block else None + + +def edit_diagnostics( + state: StapleState, + lead: StapleState, + *, + protected_positions: list[int] | None = None, + catalog: dict[str, BuildingBlock] | None = None, +) -> dict[str, Any]: + """Compute structured edit metrics for ``state`` relative to ``lead``.""" + seq = state.sequence_tokens + lead_seq = lead.sequence_tokens + n = max(len(seq), len(lead_seq)) + edited: list[int] = [] + edit_script: list[dict[str, Any]] = [] + for i in range(n): + a = seq[i] if i < len(seq) else None + b = lead_seq[i] if i < len(lead_seq) else None + if a != b: + edited.append(i) + edit_script.append({"pos": i, "lead_aa": b, "designed_aa": a}) + edit_distance = float(len(edited)) + seq_identity = 1.0 - (edit_distance / max(n, 1)) + + protected = set(protected_positions or []) + protected_edits = [p for p in edited if p in protected] + protected_violation = len(protected_edits) > 0 + + anchor_created = False + anchor_existing = False + if state.anchor_pair is not None: + i, j = state.anchor_pair + i_edited = i in edited + j_edited = j in edited + motif = _motif_for_block(state.block_id, catalog) + i_aa_ok = True + j_aa_ok = True + if motif: + if motif.get("i_aa") and 0 <= i < len(seq): + i_aa_ok = seq[i] in motif["i_aa"] + if motif.get("j_aa") and 0 <= j < len(seq): + j_aa_ok = seq[j] in motif["j_aa"] + # "Created by edit": at least one of the anchor sites was changed AND + # the resulting site satisfies the motif. + if (i_edited or j_edited) and i_aa_ok and j_aa_ok: + anchor_created = True + # "Used existing motif": both anchor sites were already correct in + # the lead (no edits at i/j) and the motif matches. + if (not i_edited) and (not j_edited) and i_aa_ok and j_aa_ok: + anchor_existing = True + + return { + "edit_distance": edit_distance, + "edited_positions": edited, + "sequence_identity": float(seq_identity), + "protected_edit_violation": bool(protected_violation), + "protected_edited_positions": protected_edits, + "anchor_created_by_edit": bool(anchor_created), + "anchor_used_existing_motif": bool(anchor_existing), + "edit_script": edit_script, + } diff --git a/staplebridge/chemistry/protected.py b/staplebridge/chemistry/protected.py new file mode 100644 index 0000000000000000000000000000000000000000..ee5f36559fcbb78bff9743373802de3db475ceca --- /dev/null +++ b/staplebridge/chemistry/protected.py @@ -0,0 +1,5 @@ +from __future__ import annotations + + +def can_edit_position(position: int, protected_positions: list[int]) -> bool: + return position not in set(protected_positions) diff --git a/staplebridge/chemistry/state.py b/staplebridge/chemistry/state.py new file mode 100644 index 0000000000000000000000000000000000000000..703f5430658d545a3ce33d16fa6a4292082049ee --- /dev/null +++ b/staplebridge/chemistry/state.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass(eq=False) +class StapleState: + sequence_tokens: list[str] + anchor_pair: tuple[int, int] | None = None + block_id: str | None = None + topology: str = "linear" + metadata: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_sequence(cls, seq: str, **kwargs: Any) -> "StapleState": + return cls(sequence_tokens=list(seq), **kwargs) + + def to_sequence(self) -> str: + return "".join(self.sequence_tokens) + + def copy(self) -> "StapleState": + return StapleState( + sequence_tokens=self.sequence_tokens.copy(), + anchor_pair=None if self.anchor_pair is None else tuple(self.anchor_pair), + block_id=self.block_id, + topology=self.topology, + metadata=dict(self.metadata), + ) + + def __hash__(self) -> int: + return hash((tuple(self.sequence_tokens), self.anchor_pair, self.block_id, self.topology)) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, StapleState): + return False + return ( + self.sequence_tokens == other.sequence_tokens + and self.anchor_pair == other.anchor_pair + and self.block_id == other.block_id + and self.topology == other.topology + ) + + def __str__(self) -> str: + return f"StapleState(seq={self.to_sequence()}, anchor={self.anchor_pair}, block={self.block_id}, topo={self.topology})" diff --git a/staplebridge/data/__init__.py b/staplebridge/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f3b0d32f802e05f5e1927433a19f527fd5a43b92 --- /dev/null +++ b/staplebridge/data/__init__.py @@ -0,0 +1 @@ +"""""" diff --git a/staplebridge/data/catalog.py b/staplebridge/data/catalog.py new file mode 100644 index 0000000000000000000000000000000000000000..51af46861121a722204d4baa779873af943145f9 --- /dev/null +++ b/staplebridge/data/catalog.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Iterable + +from staplebridge.data.schemas import BuildingBlock + + +def default_catalog() -> list[BuildingBlock]: + """Single CP-Composer-style stapled block: K ↔ D/E lactam at i,i+3 or i,i+4. + + Motif and Cα window come straight from CP-Composer's success criterion + (`evaluate_utils/success_utils.ipynb`): + - peptide[i] == 'K' and peptide[i+3] in {'D','E'} or peptide[i+4] in {'D','E'} + - 4.0 ≤ Cα(i)-Cα(j) ≤ 6.5 Å + """ + return [ + BuildingBlock( + block_id="STAPLE_LACTAM", + name="K-(D/E) lactam staple", + chemistry_class="stapled", + synthetic_accessibility_score=0.85, + cost_score=0.5, + spps_score=0.7, + motif={"i_aa": ["K"], "j_aa": ["D", "E"], "spacings": [3, 4]}, + ca_window=(4.0, 6.5), + ), + ] + + +def save_catalog(blocks: Iterable[BuildingBlock], out_path: str | Path) -> None: + path = Path(out_path) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + json.dump([b.to_dict() for b in blocks], f, indent=2) + + +def load_catalog(path: str | Path) -> list[BuildingBlock]: + with Path(path).open("r", encoding="utf-8") as f: + data = json.load(f) + return [BuildingBlock.from_dict(x) for x in data] + + +def catalog_index(blocks: Iterable[BuildingBlock]) -> dict[str, BuildingBlock]: + return {b.block_id: b for b in blocks} diff --git a/staplebridge/data/dataset.py b/staplebridge/data/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..4f1e683e5af13016a81f8e1d5235982550722c13 --- /dev/null +++ b/staplebridge/data/dataset.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +import json +from pathlib import Path + +from staplebridge.data.schemas import LeadExample + + +def save_leads(leads: list[LeadExample], out_path: str | Path) -> None: + path = Path(out_path) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + for lead in leads: + f.write(json.dumps(lead.to_dict()) + "\n") + + +def load_leads(path: str | Path) -> list[LeadExample]: + leads: list[LeadExample] = [] + with Path(path).open("r", encoding="utf-8") as f: + for line in f: + if line.strip(): + leads.append(LeadExample(**json.loads(line))) + return leads diff --git a/staplebridge/data/schemas.py b/staplebridge/data/schemas.py new file mode 100644 index 0000000000000000000000000000000000000000..1522313e75149b2dff1a762fa2bce45810c823b1 --- /dev/null +++ b/staplebridge/data/schemas.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from dataclasses import dataclass, field, asdict +from typing import Any, Optional + + +@dataclass +class LeadExample: + example_id: str + linear_sequence: str + target_id: Optional[str] = None + target_context: Optional[dict[str, Any]] = None + protected_positions: list[int] = field(default_factory=list) + preferred_property_direction: dict[str, str] = field(default_factory=dict) + thresholds: dict[str, float] = field(default_factory=dict) + known_active_motif_positions: Optional[list[int]] = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class BuildingBlock: + """Stapling building block. + + The geometry / motif fields drive CP-Composer-style feasibility: + - `motif`: sequence-level pattern; for K↔D/E lactam stapling this is + {"i_aa": ["K"], "j_aa": ["D","E"], "spacings": [3, 4]}. + - `ca_window`: allowed Cα(i)-Cα(j) distance window in Å. + + `chemistry_class` is now restricted to {"stapled"} (head_to_tail / + disulfide / bicycle were removed; the previous hydrocarbon i,i+4 / i,i+7 + blocks were also removed because they fall outside CP-Composer's scope). + """ + + block_id: str + name: str + chemistry_class: str + synthetic_accessibility_score: float + cost_score: float + spps_score: float + motif: Optional[dict[str, Any]] = None + ca_window: tuple[float, float] = (4.0, 6.5) + + def to_dict(self) -> dict[str, Any]: + d = asdict(self) + d["ca_window"] = list(self.ca_window) + return d + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> "BuildingBlock": + d = dict(d) + if "ca_window" in d and isinstance(d["ca_window"], list): + d["ca_window"] = tuple(d["ca_window"]) + # tolerate legacy catalogs by dropping retired fields + for legacy in ("allowed_anchor_spacings", "compatible_residue_types", "token_substitution"): + d.pop(legacy, None) + return cls(**d) diff --git a/staplebridge/data/vocab.py b/staplebridge/data/vocab.py new file mode 100644 index 0000000000000000000000000000000000000000..d7b84f7045923c9e6c0a8e34782f4c149bc377de --- /dev/null +++ b/staplebridge/data/vocab.py @@ -0,0 +1,11 @@ +AMINO_ACIDS = list("ACDEFGHIKLMNPQRSTVWY") +NCAA_TOKENS = ["X", "B"] +PAD_TOKEN = "" +UNK_TOKEN = "" + +ALL_TOKENS = [PAD_TOKEN, UNK_TOKEN] + AMINO_ACIDS + NCAA_TOKENS +TOKEN_TO_ID = {t: i for i, t in enumerate(ALL_TOKENS)} +ID_TO_TOKEN = {i: t for t, i in TOKEN_TO_ID.items()} + +HYDROPHOBIC = set("AILMFWVY") +RISKY_TOKENS = {"W", "F", "B"} diff --git a/staplebridge/hydrocarbon/__init__.py b/staplebridge/hydrocarbon/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ede60c8e53a5e48080b7526c0b5e701b111e1807 --- /dev/null +++ b/staplebridge/hydrocarbon/__init__.py @@ -0,0 +1,39 @@ +"""Hydrocarbon (RCM) stapling branch. + +Entry points used by the main StapleBridge training pipeline: + :mod:`staplebridge.hydrocarbon.catalog` - hydrocarbon building blocks. + :mod:`staplebridge.hydrocarbon.tokenizer` - multi-character monomer + tokenizer (S5, R8, ...) plus the model-vocab projection. + :mod:`staplebridge.hydrocarbon.actions` - hydrocarbon action generator. + :mod:`staplebridge.hydrocarbon.curriculum` - hydrocarbon demonstration paths. + :mod:`staplebridge.hydrocarbon.endpoint_prior` - empirical endpoint prior. + :mod:`staplebridge.hydrocarbon.monomers` - monomer SMILES fragment library. + :mod:`staplebridge.hydrocarbon.smiles_builder` - linear precursor and + RCM-stapled product SMILES construction with strict RDKit validation. + :mod:`staplebridge.hydrocarbon.plan_reference` - plan-aware empirical + reference process (select a whole staple plan, then complete it). + :mod:`staplebridge.hydrocarbon.plan_control` - finite feasible plan support, + q_ref, plan head q_theta and committed-plan trajectory sampling. + :mod:`staplebridge.hydrocarbon.exact_sb_cache` - exact finite-support + teacher q*(p|x) construction and its persistent cache. + :mod:`staplebridge.hydrocarbon.terminal_energy` - hydrocarbon terminal + energy augmentation. +""" + +from __future__ import annotations + +__all__ = [ + "actions", + "catalog", + "curriculum", + "endpoint_prior", + "exact_sb_cache", + "monomers", + "plan_control", + "plan_reference", + "plan_validation", + "property_energy", + "smiles_builder", + "terminal_energy", + "tokenizer", +] diff --git a/staplebridge/hydrocarbon/actions.py b/staplebridge/hydrocarbon/actions.py new file mode 100644 index 0000000000000000000000000000000000000000..9e12eeca6b2bcd5999006c660c9f223774b93ba8 --- /dev/null +++ b/staplebridge/hydrocarbon/actions.py @@ -0,0 +1,384 @@ +"""Hydrocarbon action generator. + +A separate generator from :func:`staplebridge.graph.neighbors.enumerate_neighbors`, +which keeps its exact lactam behaviour (K/D/E residue substitution, lactam motif +anchors, ``STAPLE_LACTAM``). Neither generator can produce the other's actions: + +* this generator substitutes ncAA anchor monomers (``S5``, ``R8``, ...) and + never proposes a K-D/E lactam block; +* the lactam generator substitutes only ``K``/``D``/``E``/``A`` and never emits + a multi-character monomer. + +Action families (mirroring the lactam generator's shape so the trainer, graph +and beam search need no special-casing): + +1. ncAA substitution - install an anchor monomer at an editable position +2. anchor assignment - select the (i, j) pair for a catalog topology +3. block assignment - attach the hydrocarbon block +4. topology activation - close the staple + +Every rejection carries an explicit reason (:class:`FailureReason`) rather than +being silently dropped, so an empty neighbour set can always be explained. +""" + +from __future__ import annotations + +from collections import Counter +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Final + +from staplebridge.chemistry.actions import ( + AnchorAssignAction, + AnchorReassignAction, + BlockAssignAction, + NcAASubstitutionAction, + NoOpAction, + TopologyActivationAction, +) +from staplebridge.chemistry.protected import can_edit_position +from staplebridge.chemistry.state import StapleState +from staplebridge.data.schemas import BuildingBlock +from staplebridge.hydrocarbon.catalog import ( + UNSUPPORTED_ANCHOR_TOKENS, + block_topology, + is_hydrocarbon_block, +) +from staplebridge.hydrocarbon.tokenizer import is_anchor_token + + +class FailureReason(str, Enum): + """Explicit reasons a hydrocarbon action or state was rejected.""" + + OK = "ok" + WRONG_CHEMISTRY = "wrong_chemistry" + NON_HYDROCARBON_BLOCK = "non_hydrocarbon_block" + UNSUPPORTED_ANCHOR_TOKEN = "unsupported_anchor_token" + UNSUPPORTED_PAIR = "unsupported_anchor_pair" + UNSUPPORTED_SPACING = "unsupported_spacing" + UNSUPPORTED_PAIR_SPACING = "unsupported_pair_spacing_combination" + WRONG_ANCHOR_COUNT = "wrong_anchor_count" + ANCHOR_OUT_OF_RANGE = "anchor_out_of_range" + ANCHOR_POSITION_PROTECTED = "anchor_position_protected" + ANCHOR_TOKEN_MISMATCH = "anchor_token_mismatch" + NO_ANCHOR_PAIR = "no_anchor_pair_set" + NO_BLOCK = "no_block_set" + ALREADY_STAPLED = "already_stapled" + DOUBLE_STAPLE_UNSUPPORTED = "double_staple_unsupported" + SEQUENCE_TOO_SHORT = "sequence_too_short" + + +#: Exactly two anchors define a single hydrocarbon staple. Double staples are +#: out of scope, so any other count is a failure rather than a truncation. +REQUIRED_ANCHOR_COUNT: Final[int] = 2 + + +@dataclass +class ActionGenerationReport: + """Diagnostics for one call to :func:`enumerate_hydrocarbon_neighbors`.""" + + n_candidates: int = 0 + by_family: Counter = field(default_factory=Counter) + failures: Counter = field(default_factory=Counter) + + def record_failure(self, reason: FailureReason, count: int = 1) -> None: + """Tally one rejection.""" + self.failures[reason.value] += count + + def as_dict(self) -> dict[str, Any]: + """JSON-serialisable view.""" + return { + "n_candidates": int(self.n_candidates), + "by_family": {k: int(v) for k, v in sorted(self.by_family.items())}, + "failures": {k: int(v) for k, v in sorted(self.failures.items())}, + } + + +def validate_hydrocarbon_staple( + tokens: list[str], + anchor_pair: tuple[int, int] | None, + block: BuildingBlock | None, + catalog: list[BuildingBlock], +) -> FailureReason: + """Validate a candidate hydrocarbon staple. + + Checks, in order: block provenance, anchor presence, index range, anchor + monomer identity, and whether the (pair, spacing) combination is one the + catalog actually describes. + + Returns: + :attr:`FailureReason.OK` when the staple is legal, else the specific + reason it is not. + """ + if block is None: + return FailureReason.NO_BLOCK + if not is_hydrocarbon_block(block): + return FailureReason.NON_HYDROCARBON_BLOCK + if anchor_pair is None: + return FailureReason.NO_ANCHOR_PAIR + + i, j = anchor_pair + if not (0 <= i < len(tokens)) or not (0 <= j < len(tokens)): + return FailureReason.ANCHOR_OUT_OF_RANGE + if i == j: + return FailureReason.WRONG_ANCHOR_COUNT + + token_i = tokens[i].upper() + token_j = tokens[j].upper() + for token in (token_i, token_j): + if token in {t.upper() for t in UNSUPPORTED_ANCHOR_TOKENS}: + return FailureReason.UNSUPPORTED_ANCHOR_TOKEN + + # More than two anchor monomers in the chain would be a double staple. + if sum(1 for t in tokens if is_anchor_token(t)) != REQUIRED_ANCHOR_COUNT: + return FailureReason.DOUBLE_STAPLE_UNSUPPORTED + + if not is_anchor_token(token_i) or not is_anchor_token(token_j): + return FailureReason.ANCHOR_TOKEN_MISMATCH + + pair = f"{token_i}-{token_j}" + spacing = j - i + catalog_topologies = {block_topology(b) for b in catalog if is_hydrocarbon_block(b)} + block_pair, block_spacing = block_topology(block) + + if pair != block_pair: + # Distinguish "this pair exists elsewhere in the catalog" from "this pair + # is unknown", which are different problems for the caller. + if any(p == pair for p, _ in catalog_topologies): + return FailureReason.UNSUPPORTED_PAIR_SPACING + return FailureReason.UNSUPPORTED_PAIR + if spacing != block_spacing: + if (pair, spacing) in catalog_topologies: + return FailureReason.UNSUPPORTED_PAIR_SPACING + return FailureReason.UNSUPPORTED_SPACING + if (pair, spacing) not in catalog_topologies: + return FailureReason.UNSUPPORTED_PAIR_SPACING + + return FailureReason.OK + + +def _existing_anchor_pairs( + tokens: list[str], block: BuildingBlock +) -> list[tuple[int, int]]: + """(i, j) pairs already realised in ``tokens`` for this block's topology.""" + pair, spacing = block_topology(block) + i_token, j_token = pair.split("-") + pairs: list[tuple[int, int]] = [] + for i in range(len(tokens)): + j = i + spacing + if j >= len(tokens): + continue + if tokens[i].upper() == i_token and tokens[j].upper() == j_token: + pairs.append((i, j)) + return pairs + + +def _substitution_sites( + tokens: list[str], + block: BuildingBlock, + protected_positions: list[int], + report: ActionGenerationReport, +) -> list[tuple[int, str]]: + """Positions where installing an anchor monomer opens a legal topology. + + A substitution is only proposed when it could actually contribute to a + staple this block supports: either it completes a pair against an existing + partner anchor, or it starts one at a position whose partner slot is + editable. + """ + pair, spacing = block_topology(block) + i_token, j_token = pair.split("-") + sites: list[tuple[int, str]] = [] + + for position in range(len(tokens)): + if not can_edit_position(position, protected_positions): + report.record_failure(FailureReason.ANCHOR_POSITION_PROTECTED) + continue + + # Install the i-side monomer when the j-side partner is reachable. + partner_j = position + spacing + if partner_j < len(tokens) and tokens[position].upper() != i_token: + partner_ok = tokens[partner_j].upper() == j_token or can_edit_position( + partner_j, protected_positions + ) + if partner_ok: + sites.append((position, i_token)) + + # Install the j-side monomer when the i-side partner is reachable. + partner_i = position - spacing + if partner_i >= 0 and tokens[position].upper() != j_token: + partner_ok = tokens[partner_i].upper() == i_token or can_edit_position( + partner_i, protected_positions + ) + if partner_ok: + sites.append((position, j_token)) + + return sites + + +def enumerate_hydrocarbon_neighbors( + state: StapleState, + catalog: list[BuildingBlock], + protected_positions: list[int], + allow_noop: bool = True, + max_neighbors: int = 128, + report: ActionGenerationReport | None = None, +) -> list[StapleState]: + """Enumerate hydrocarbon edit neighbours of ``state``. + + Only topologies present in ``catalog`` are generated: no unsupported pair, + no unsupported spacing, and never a lactam block. + + Args: + state: current state, whose tokens are hydrocarbon monomers. + catalog: hydrocarbon blocks (non-hydrocarbon blocks are rejected). + protected_positions: positions that must not be edited. + allow_noop: include the no-op self transition. + max_neighbors: cap on returned candidates. + report: optional diagnostics sink recording failure reasons. + + Returns: + Candidate states, ordered so that structurally meaningful actions are + never crowded out by substitutions (same priority discipline as the + lactam generator). + """ + report = report if report is not None else ActionGenerationReport() + + hydrocarbon_blocks: list[BuildingBlock] = [] + for block in catalog: + if is_hydrocarbon_block(block): + hydrocarbon_blocks.append(block) + else: + report.record_failure(FailureReason.NON_HYDROCARBON_BLOCK) + + # Absorbing terminal, matching the lactam generator's contract: an empty + # neighbour set is how every consumer recognises a terminal state. + if state.topology == "stapled": + report.record_failure(FailureReason.ALREADY_STAPLED) + return [] + + tokens = state.sequence_tokens + if len(tokens) < 2: + report.record_failure(FailureReason.SEQUENCE_TOO_SHORT) + return [] + + noop_candidates: list[StapleState] = [] + if allow_noop: + noop_candidates.append(NoOpAction().apply(state)) + + anchor_candidates: list[StapleState] = [] + topology_candidates: list[StapleState] = [] + block_candidates: list[StapleState] = [] + substitution_candidates: list[StapleState] = [] + + for block in hydrocarbon_blocks: + # -- anchor assignment over already-installed monomer pairs --------- + for (i, j) in _existing_anchor_pairs(tokens, block): + verdict = validate_hydrocarbon_staple(tokens, (i, j), block, hydrocarbon_blocks) + if verdict is not FailureReason.OK: + report.record_failure(verdict) + continue + if state.anchor_pair is None: + candidate = AnchorAssignAction(i, j).apply(state) + elif (i, j) != tuple(state.anchor_pair): + candidate = AnchorReassignAction(i, j).apply(state) + else: + candidate = None + if candidate is not None: + candidate.block_id = block.block_id + anchor_candidates.append(candidate) + + # -- block assignment: anchors already right, block not yet set -- + if ( + state.anchor_pair is not None + and tuple(state.anchor_pair) == (i, j) + and state.block_id != block.block_id + ): + block_candidates.append(BlockAssignAction(block.block_id).apply(state)) + + # -- ncAA substitution --------------------------------------------- + for position, monomer in _substitution_sites( + tokens, block, protected_positions, report + ): + substitution_candidates.append( + NcAASubstitutionAction(position, monomer).apply(state) + ) + + # -- topology activation: needs a valid anchor pair and a valid block ---- + if state.topology == "linear" and state.anchor_pair is not None and state.block_id: + current_block = next( + (b for b in hydrocarbon_blocks if b.block_id == state.block_id), None + ) + verdict = validate_hydrocarbon_staple( + tokens, tuple(state.anchor_pair), current_block, hydrocarbon_blocks + ) + if verdict is FailureReason.OK: + topology_candidates.append(TopologyActivationAction().apply(state)) + else: + report.record_failure(verdict) + + kept = _dedup( + noop_candidates + topology_candidates + anchor_candidates + block_candidates + )[:max_neighbors] + + seen = set(kept) + for candidate in substitution_candidates: + if len(kept) >= max_neighbors: + break + if candidate not in seen: + seen.add(candidate) + kept.append(candidate) + + report.n_candidates = len(kept) + report.by_family.update( + { + "noop": len(noop_candidates), + "topology_activation": len(topology_candidates), + "anchor_assignment": len(anchor_candidates), + "block_assignment": len(block_candidates), + "ncaa_substitution": max(0, len(kept) - len(_dedup( + noop_candidates + topology_candidates + anchor_candidates + block_candidates + )[:max_neighbors])), + } + ) + return kept + + +def _dedup(states: list[StapleState]) -> list[StapleState]: + """Order-preserving dedup, matching the lactam generator's helper.""" + seen: set[StapleState] = set() + out: list[StapleState] = [] + for state in states: + if state not in seen: + seen.add(state) + out.append(state) + return out + + +class HydrocarbonTransitionGraph: + """Drop-in graph for the hydrocarbon branch. + + Mirrors :class:`staplebridge.graph.transition_graph.TransitionGraph`'s + interface so the trainer and decoders can use it unchanged, while routing to + the hydrocarbon generator. The lactam graph class is untouched. + """ + + def __init__(self, catalog: list[BuildingBlock], max_neighbors: int = 128) -> None: + self.catalog = catalog + self.max_neighbors = max_neighbors + self.last_report: ActionGenerationReport | None = None + + def neighbors( + self, state: StapleState, protected_positions: list[int] + ) -> list[StapleState]: + """Hydrocarbon neighbours of ``state``; diagnostics land in ``last_report``.""" + report = ActionGenerationReport() + out = enumerate_hydrocarbon_neighbors( + state, + catalog=self.catalog, + protected_positions=protected_positions, + max_neighbors=self.max_neighbors, + report=report, + ) + self.last_report = report + return out diff --git a/staplebridge/hydrocarbon/catalog.py b/staplebridge/hydrocarbon/catalog.py new file mode 100644 index 0000000000000000000000000000000000000000..38874089117d3036c7c754dab621b78bc5351f22 --- /dev/null +++ b/staplebridge/hydrocarbon/catalog.py @@ -0,0 +1,198 @@ +"""Hydrocarbon (RCM) staple catalog. + +Separate from :func:`staplebridge.data.catalog.default_catalog`, which still +returns exactly one block (``STAPLE_LACTAM``) and is not touched here. No +hydrocarbon block is ever registered into the lactam catalog, and no block here +carries the ``STAPLE_LACTAM`` id. + +Supported topologies +-------------------- +========================= ========= ======= =================================== +block pair spacing status +========================= ========= ======= =================================== +``STAPLE_HC_S5S5_I4`` S5-S5 i,i+4 enabled +``STAPLE_HC_R8S5_I7`` R8-S5 i,i+7 enabled +``STAPLE_HC_R5S8_I7`` R5-S8 i,i+7 disabled by default (low frequency) +========================= ========= ======= =================================== + +The two enabled topologies are the dominant modes in the observed data +(``analysis/hydrocarbon_endpoint_distribution``: S5-S5 at spacing 4 is 221/224 +of its pair, R8-S5 at spacing 7 is 100/103). R5-S8/i,i+7 has only 7 observations +and is therefore an opt-in extension rather than a default. + +Deliberately unsupported for now: S3/R3 anchors (zero observations), any other +spacing, and double staples (more than one anchor pair). +""" + +from __future__ import annotations + +from typing import Final + +from staplebridge.data.schemas import BuildingBlock + +#: ``chemistry_class`` marker for every hydrocarbon block. Distinct from the +#: lactam blocks' ``"stapled"`` so the two can never be confused by class alone. +HYDROCARBON_CHEMISTRY_CLASS: Final[str] = "hydrocarbon_stapled" + +BLOCK_S5S5_I4: Final[str] = "STAPLE_HC_S5S5_I4" +BLOCK_R8S5_I7: Final[str] = "STAPLE_HC_R8S5_I7" +BLOCK_R5S8_I7: Final[str] = "STAPLE_HC_R5S8_I7" + +#: Enabled unless a config opts in to the extensions. +DEFAULT_ENABLED_BLOCKS: Final[tuple[str, ...]] = (BLOCK_S5S5_I4, BLOCK_R8S5_I7) + +#: Low-frequency topologies, off unless explicitly requested. +OPTIONAL_BLOCKS: Final[tuple[str, ...]] = (BLOCK_R5S8_I7,) + +#: Anchor monomers this catalog can place. +SUPPORTED_ANCHOR_TOKENS: Final[tuple[str, ...]] = ("S5", "R8", "R5", "S8") + +#: Anchors that exist in StaPep syntax but are not supported here. +UNSUPPORTED_ANCHOR_TOKENS: Final[tuple[str, ...]] = ("S3", "R3") + +#: Cα(i)-Cα(j) windows. i,i+4 spans one helical turn and i,i+7 spans two, so the +#: longer staple gets the wider, longer-distance window. +_CA_WINDOW_I4: Final[tuple[float, float]] = (4.5, 7.5) +_CA_WINDOW_I7: Final[tuple[float, float]] = (8.5, 13.0) + + +def _block( + block_id: str, + name: str, + i_aa: str, + j_aa: str, + spacing: int, + ca_window: tuple[float, float], + cost_score: float, + spps_score: float, + sa_score: float, +) -> BuildingBlock: + """Build one hydrocarbon block with a single (pair, spacing) topology. + + Each block pins exactly one spacing, so ``motif['spacings']`` has length 1. + That is what keeps the action generator from proposing an i,i+4 R8-S5 staple + or any other combination the catalog does not describe. + """ + return BuildingBlock( + block_id=block_id, + name=name, + chemistry_class=HYDROCARBON_CHEMISTRY_CLASS, + synthetic_accessibility_score=sa_score, + cost_score=cost_score, + spps_score=spps_score, + motif={ + "i_aa": [i_aa], + "j_aa": [j_aa], + "spacings": [spacing], + # Extra descriptive keys, ignored by the generic motif matcher but + # read by the hydrocarbon action generator and the endpoint prior. + "anchor_pair_ordered": f"{i_aa}-{j_aa}", + "anchor_pair_unordered": "-".join(sorted((i_aa, j_aa))), + "staple_chemistry": "hydrocarbon_rcm", + }, + ca_window=ca_window, + ) + + +def hydrocarbon_catalog( + include_optional: bool = False, + enabled_blocks: list[str] | tuple[str, ...] | None = None, +) -> list[BuildingBlock]: + """Return the hydrocarbon building blocks. + + Args: + include_optional: also return the low-frequency R5-S8/i,i+7 extension. + enabled_blocks: explicit allow-list of block ids. When given it wins + over ``include_optional``. + + Returns: + Blocks in a stable order. Never includes any lactam block. + + Raises: + ValueError: if ``enabled_blocks`` names an unknown block id, so a typo + cannot silently yield an empty or partial catalog. + """ + all_blocks = { + BLOCK_S5S5_I4: _block( + BLOCK_S5S5_I4, + "S5-S5 hydrocarbon staple (i, i+4)", + i_aa="S5", + j_aa="S5", + spacing=4, + ca_window=_CA_WINDOW_I4, + cost_score=0.55, + spps_score=0.65, + sa_score=0.80, + ), + BLOCK_R8S5_I7: _block( + BLOCK_R8S5_I7, + "R8-S5 hydrocarbon staple (i, i+7)", + i_aa="R8", + j_aa="S5", + spacing=7, + ca_window=_CA_WINDOW_I7, + cost_score=0.70, + spps_score=0.55, + sa_score=0.70, + ), + BLOCK_R5S8_I7: _block( + BLOCK_R5S8_I7, + "R5-S8 hydrocarbon staple (i, i+7), low-frequency extension", + i_aa="R5", + j_aa="S8", + spacing=7, + ca_window=_CA_WINDOW_I7, + cost_score=0.80, + spps_score=0.50, + sa_score=0.60, + ), + } + + if enabled_blocks is not None: + requested = list(enabled_blocks) + unknown = [b for b in requested if b not in all_blocks] + if unknown: + raise ValueError( + f"unknown hydrocarbon block id(s): {unknown}; " + f"known: {sorted(all_blocks)}" + ) + order = list(DEFAULT_ENABLED_BLOCKS) + list(OPTIONAL_BLOCKS) + return [all_blocks[b] for b in order if b in set(requested)] + + selected = list(DEFAULT_ENABLED_BLOCKS) + if include_optional: + selected += list(OPTIONAL_BLOCKS) + return [all_blocks[b] for b in selected] + + +def hydrocarbon_catalog_from_config(config: dict | None) -> list[BuildingBlock]: + """Build the catalog from a ``hydrocarbon.catalog`` config section.""" + section = dict((config or {}).get("catalog") or {}) + return hydrocarbon_catalog( + include_optional=bool(section.get("include_optional", False)), + enabled_blocks=section.get("enabled_blocks"), + ) + + +def is_hydrocarbon_block(block: BuildingBlock | None) -> bool: + """True when ``block`` came from this catalog.""" + return bool(block is not None and block.chemistry_class == HYDROCARBON_CHEMISTRY_CLASS) + + +def block_topology(block: BuildingBlock) -> tuple[str, int]: + """Return ``(ordered_pair, spacing)`` for a hydrocarbon block.""" + motif = block.motif or {} + spacings = motif.get("spacings") or [] + if len(spacings) != 1: + raise ValueError( + f"hydrocarbon block {block.block_id} must pin exactly one spacing, " + f"got {spacings!r}" + ) + i_aa = (motif.get("i_aa") or [""])[0] + j_aa = (motif.get("j_aa") or [""])[0] + return f"{i_aa}-{j_aa}", int(spacings[0]) + + +def supported_topologies(blocks: list[BuildingBlock]) -> list[tuple[str, int]]: + """All ``(ordered_pair, spacing)`` topologies the given blocks allow.""" + return [block_topology(b) for b in blocks] diff --git a/staplebridge/hydrocarbon/curriculum.py b/staplebridge/hydrocarbon/curriculum.py new file mode 100644 index 0000000000000000000000000000000000000000..da5edaeddc37b263f9a5f292564d1b2b429f5d41 --- /dev/null +++ b/staplebridge/hydrocarbon/curriculum.py @@ -0,0 +1,249 @@ +"""Hydrocarbon demonstration paths (curriculum). + +The lactam curriculum +(:class:`staplebridge.training.curriculum_sampler.CurriculumTrajectorySampler`, +which wraps ``propose_minimal_lactam_motif_edits``) is untouched and still +defaults to ``block_id="STAPLE_LACTAM"``. This module is the hydrocarbon +counterpart and never emits a lactam block. + +A demonstration path installs the two anchor monomers with the fewest possible +substitutions, then assigns the anchor pair, then activates the topology: + + linear -> [ncAA substitution]* -> anchor assign -> topology activation + +Plans are ranked so that sites already carrying the right monomer are preferred +(zero edits beats one edit), and every rejected plan carries a +:class:`~staplebridge.hydrocarbon.actions.FailureReason`. +""" + +from __future__ import annotations + +from collections import Counter +from dataclasses import dataclass, field +from typing import Any + +from staplebridge.chemistry.actions import ( + AnchorAssignAction, + NcAASubstitutionAction, + TopologyActivationAction, +) +from staplebridge.chemistry.protected import can_edit_position +from staplebridge.chemistry.state import StapleState +from staplebridge.data.schemas import BuildingBlock +from staplebridge.hydrocarbon.actions import ( + FailureReason, + validate_hydrocarbon_staple, +) +from staplebridge.hydrocarbon.catalog import block_topology, is_hydrocarbon_block + + +@dataclass +class HydrocarbonStaplePlan: + """One concrete plan for building a hydrocarbon staple.""" + + block_id: str + anchor_pair: tuple[int, int] + ordered_pair: str + spacing: int + substitutions: list[tuple[int, str]] = field(default_factory=list) + + @property + def n_edits(self) -> int: + """Number of ncAA substitutions this plan requires.""" + return len(self.substitutions) + + def as_dict(self) -> dict[str, Any]: + """JSON-serialisable view.""" + return { + "block_id": self.block_id, + "anchor_pair": list(self.anchor_pair), + "ordered_pair": self.ordered_pair, + "spacing": self.spacing, + "substitutions": [[p, t] for p, t in self.substitutions], + "n_edits": self.n_edits, + } + + +@dataclass +class HydrocarbonCurriculumConfig: + """Config for the hydrocarbon curriculum. + + ``block_id`` is deliberately absent: the block is chosen per plan from the + hydrocarbon catalog, so a lactam id cannot leak in through config. + """ + + enabled: bool = False + max_anchor_edits: int = 2 + num_trajectories_per_lead: int = 4 + prefer_existing_anchors: bool = True + protect_positions: bool = True + require_valid_terminal: bool = True + + +@dataclass +class HydrocarbonCurriculumDiagnostics: + """Bookkeeping mirroring the lactam curriculum's diagnostics shape.""" + + trajectory_count: int = 0 + valid_terminal_count: int = 0 + num_edits_acc: list[int] = field(default_factory=list) + failure_reasons: Counter = field(default_factory=Counter) + leads_seen: int = 0 + leads_with_plan: int = 0 + + def as_dict(self) -> dict[str, Any]: + """JSON-serialisable view.""" + mean_edits = ( + sum(self.num_edits_acc) / len(self.num_edits_acc) + if self.num_edits_acc + else 0.0 + ) + return { + "trajectory_count": int(self.trajectory_count), + "valid_terminal_count": int(self.valid_terminal_count), + "mean_num_edits": float(mean_edits), + "leads_seen": int(self.leads_seen), + "leads_with_plan": int(self.leads_with_plan), + "failure_reasons": {k: int(v) for k, v in sorted(self.failure_reasons.items())}, + } + + +def propose_hydrocarbon_staple_plans( + tokens: list[str], + catalog: list[BuildingBlock], + protected_positions: list[int] | None = None, + max_anchor_edits: int = 2, + diagnostics: HydrocarbonCurriculumDiagnostics | None = None, +) -> list[HydrocarbonStaplePlan]: + """Enumerate minimal-edit plans for every catalog topology. + + Args: + tokens: current monomer tokens. + catalog: hydrocarbon blocks; non-hydrocarbon blocks are skipped. + protected_positions: positions that must not be edited. + max_anchor_edits: reject plans needing more substitutions than this. + diagnostics: optional sink for failure reasons. + + Returns: + Plans sorted by edit count ascending, then by position, so the caller can + take the cheapest demonstrations first. + """ + protected = list(protected_positions or []) + plans: list[HydrocarbonStaplePlan] = [] + + for block in catalog: + if not is_hydrocarbon_block(block): + if diagnostics is not None: + diagnostics.failure_reasons[FailureReason.NON_HYDROCARBON_BLOCK.value] += 1 + continue + + ordered_pair, spacing = block_topology(block) + i_token, j_token = ordered_pair.split("-") + + for i in range(len(tokens)): + j = i + spacing + if j >= len(tokens): + continue + + substitutions: list[tuple[int, str]] = [] + blocked = False + for position, wanted in ((i, i_token), (j, j_token)): + if tokens[position].upper() == wanted: + continue + if not can_edit_position(position, protected): + if diagnostics is not None: + diagnostics.failure_reasons[ + FailureReason.ANCHOR_POSITION_PROTECTED.value + ] += 1 + blocked = True + break + substitutions.append((position, wanted)) + if blocked: + continue + + if len(substitutions) > max_anchor_edits: + if diagnostics is not None: + diagnostics.failure_reasons["exceeds_max_anchor_edits"] += 1 + continue + + # Reject up front any plan whose terminal state would not be a legal + # single staple. The common case: the lead already carries anchor + # monomers elsewhere, so installing this pair would leave more than + # two anchors in the chain (a double staple, which is out of scope). + # Validating here rather than at build time keeps the invalid plan + # out of the ranked list instead of raising deep inside path + # construction. + projected = list(tokens) + for position, wanted in substitutions: + projected[position] = wanted + verdict = validate_hydrocarbon_staple(projected, (i, j), block, catalog) + if verdict is not FailureReason.OK: + if diagnostics is not None: + diagnostics.failure_reasons[verdict.value] += 1 + continue + + plans.append( + HydrocarbonStaplePlan( + block_id=block.block_id, + anchor_pair=(i, j), + ordered_pair=ordered_pair, + spacing=spacing, + substitutions=substitutions, + ) + ) + + return rank_hydrocarbon_plans(plans) + + +def rank_hydrocarbon_plans( + plans: list[HydrocarbonStaplePlan], +) -> list[HydrocarbonStaplePlan]: + """Cheapest-first ordering: fewest edits, then earliest anchor, then block id.""" + return sorted( + plans, + key=lambda p: (p.n_edits, p.anchor_pair[0], p.anchor_pair[1], p.block_id), + ) + + +def build_hydrocarbon_demonstration_path( + initial: StapleState, + plan: HydrocarbonStaplePlan, + catalog: list[BuildingBlock], +) -> list[StapleState]: + """Materialise ``plan`` as a state path from ``initial`` to a stapled terminal. + + Returns: + The full path including ``initial``. The last state has + ``topology == "stapled"``. + + Raises: + ValueError: if the resulting terminal state would not be a legal + hydrocarbon staple, so an invalid demonstration is never emitted. + """ + path = [initial.copy()] + current = initial + + for position, monomer in plan.substitutions: + current = NcAASubstitutionAction(position, monomer).apply(current) + path.append(current) + + block = next((b for b in catalog if b.block_id == plan.block_id), None) + if block is None: + raise ValueError(f"plan references unknown block {plan.block_id!r}") + + current = AnchorAssignAction(*plan.anchor_pair).apply(current) + current.block_id = plan.block_id + path.append(current) + + verdict = validate_hydrocarbon_staple( + current.sequence_tokens, plan.anchor_pair, block, catalog + ) + if verdict is not FailureReason.OK: + raise ValueError( + f"plan for block {plan.block_id} at {plan.anchor_pair} is invalid: " + f"{verdict.value}" + ) + + current = TopologyActivationAction().apply(current) + path.append(current) + return path diff --git a/staplebridge/hydrocarbon/data/pair_spacing_probabilities.json b/staplebridge/hydrocarbon/data/pair_spacing_probabilities.json new file mode 100644 index 0000000000000000000000000000000000000000..8f4db2b0f575808dbffb16d7757ea8673571a1a9 --- /dev/null +++ b/staplebridge/hydrocarbon/data/pair_spacing_probabilities.json @@ -0,0 +1,28 @@ +{ + "description": "Versioned pair/spacing subset used by the default plan-aware reference.", + "source": "analysis/hydrocarbon_endpoint_distribution", + "probabilities_by_version": { + "sequence_deduplicated": { + "alpha": 1.0, + "total": 317, + "full_support_size": 10, + "support": [ + "S5-S5|4", + "R8-S5|7" + ], + "categories": { + "S5-S5|4": { + "count": 210, + "raw_probability": 0.6624605678233438, + "laplace_smoothed_probability": 0.6452599388379205 + }, + "R8-S5|7": { + "count": 93, + "raw_probability": 0.29337539432176657, + "laplace_smoothed_probability": 0.2874617737003058 + } + } + } + }, + "usage_note": "Empirical structural prior; not a trained model and does not use permeability labels." +} diff --git a/staplebridge/hydrocarbon/endpoint_prior.py b/staplebridge/hydrocarbon/endpoint_prior.py new file mode 100644 index 0000000000000000000000000000000000000000..1d6cbe23a8a5764a97ed4474b3db7d953c4f8916 --- /dev/null +++ b/staplebridge/hydrocarbon/endpoint_prior.py @@ -0,0 +1,442 @@ +"""Empirical hydrocarbon endpoint prior. + +Reads the frequency tables produced by +``analysis/hydrocarbon_endpoint_distribution`` and turns them into an energy +term. Base form: + + E_pair = -log p(pair, spacing) + +Optional, config-gated additions: peptide length, relative anchor position and +local sequence context around the anchors. + +Two properties this prior deliberately has: + +* **It never sees the permeability label.** The upstream endpoint definition is + purely structural (tokenizes cleanly, exactly two hydrocarbon anchors), so + this is a prior over *what real stapled designs look like*, not over what is + reported permeable. No classifier is trained or consulted. +* **It is hydrocarbon-only.** It is consumed exclusively by + :mod:`staplebridge.hydrocarbon.terminal_energy`; the lactam terminal energy in + ``BridgeTrainer.terminal_energy`` is not reachable from here, whether the prior + is on or off. + +Smoothed probabilities are preferred over raw ones so that a legal-but-unobserved +topology gets finite energy instead of ``inf``. +""" + +from __future__ import annotations + +import json +import math +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Final + +from staplebridge.chemistry.state import StapleState +from staplebridge.data.schemas import BuildingBlock +from staplebridge.hydrocarbon.catalog import block_topology, is_hydrocarbon_block +from staplebridge.hydrocarbon.tokenizer import is_anchor_token + +#: Default location of the empirical tables, relative to the repository root. +DEFAULT_PRIOR_DIR: Final[str] = ( + "analysis/hydrocarbon_endpoint_distribution/outputs/priors" +) + +#: Which deduplication version to read. ``raw_records`` keeps replicate rows; +#: ``sequence_deduplicated`` / ``molecule_deduplicated`` collapse them. +DEFAULT_DEDUP_VERSION: Final[str] = "sequence_deduplicated" + +#: Energy assigned when a topology is absent even from the smoothed table. Large +#: but finite, so an unsupported design is strongly penalised without producing +#: ``inf`` and poisoning downstream arithmetic. +FALLBACK_ENERGY: Final[float] = 25.0 + + +class PriorFilesMissingError(FileNotFoundError): + """Raised when the empirical prior directory or a required file is absent.""" + + +@dataclass +class EndpointPriorConfig: + """Config for :class:`EmpiricalHydrocarbonEndpointPrior`. + + Only ``E_pair`` is on by default. The optional terms are additive and each + has its own weight, so they can be enabled one at a time. + """ + + enabled: bool = False + prior_dir: str = DEFAULT_PRIOR_DIR + dedup_version: str = DEFAULT_DEDUP_VERSION + use_smoothed: bool = True + + weight_pair: float = 1.0 + use_length: bool = False + weight_length: float = 0.5 + use_relative_position: bool = False + weight_relative_position: float = 0.5 + use_local_context: bool = False + weight_local_context: float = 0.25 + local_context_window: int = 2 + + #: Cap on any single component, keeping one missing bin from dominating. + max_component_energy: float = FALLBACK_ENERGY + + @classmethod + def from_dict(cls, data: dict[str, Any] | None) -> "EndpointPriorConfig": + """Build from a ``hydrocarbon.endpoint_prior`` config section.""" + section = dict(data or {}) + cfg = cls() + for key, value in section.items(): + if hasattr(cfg, key): + current = getattr(cfg, key) + if isinstance(current, bool): + setattr(cfg, key, bool(value)) + elif isinstance(current, float): + setattr(cfg, key, float(value)) + elif isinstance(current, int) and not isinstance(current, bool): + setattr(cfg, key, int(value)) + else: + setattr(cfg, key, value) + return cfg + + +@dataclass +class EndpointEnergyBreakdown: + """Per-component decomposition of one endpoint prior evaluation.""" + + e_pair: float = 0.0 + e_length: float = 0.0 + e_relative_position: float = 0.0 + e_local_context: float = 0.0 + total: float = 0.0 + pair_key: str | None = None + status: str = "ok" + details: dict[str, Any] = field(default_factory=dict) + + def as_dict(self) -> dict[str, Any]: + """JSON-serialisable view.""" + return { + "E_endpoint_pair": float(self.e_pair), + "E_endpoint_length": float(self.e_length), + "E_endpoint_relative_position": float(self.e_relative_position), + "E_endpoint_local_context": float(self.e_local_context), + "E_endpoint_total": float(self.total), + "endpoint_pair_key": self.pair_key, + "endpoint_prior_status": self.status, + **({"endpoint_prior_details": self.details} if self.details else {}), + } + + +class EmpiricalHydrocarbonEndpointPrior: + """Energy over hydrocarbon endpoint realism, from observed frequencies. + + Args: + config: prior configuration. + root: repository root used to resolve a relative ``prior_dir``. + + Raises: + PriorFilesMissingError: if the prior is enabled but its files are absent. + An enabled-but-broken prior fails loudly rather than silently + scoring zero. + """ + + def __init__( + self, config: EndpointPriorConfig | None = None, root: Path | None = None + ) -> None: + self.cfg = config or EndpointPriorConfig() + self._root = Path(root) if root is not None else _repository_root() + self._pair_spacing: dict[str, dict[str, Any]] = {} + self._length_probabilities: dict[int, float] = {} + self._position_histogram: dict[str, Any] = {} + self._local_context: dict[str, Any] = {} + self._loaded = False + self._load_failure: str | None = None + + if self.cfg.enabled: + self.load() + + # -- loading --------------------------------------------------------- + + @property + def prior_dir(self) -> Path: + """Resolved directory holding the empirical JSON tables.""" + candidate = Path(self.cfg.prior_dir) + return candidate if candidate.is_absolute() else self._root / candidate + + @property + def is_active(self) -> bool: + """True when the prior is enabled and its tables loaded.""" + return bool(self.cfg.enabled and self._loaded) + + def load(self) -> None: + """Read the empirical tables from :attr:`prior_dir`.""" + directory = self.prior_dir + if not directory.is_dir(): + raise PriorFilesMissingError( + f"hydrocarbon endpoint prior is enabled but {directory} does not " + "exist. The main StapleBridge configuration keeps this prior " + "disabled (hydrocarbon.endpoint_prior.enabled=false); the " + "empirical endpoint evidence enters via " + "hydrocarbon.plan_reference.mode_prior instead." + ) + + pair_file = directory / "pair_spacing_probabilities.json" + if not pair_file.is_file(): + raise PriorFilesMissingError(f"missing required prior file: {pair_file}") + self._pair_spacing = self._read_probability_file(pair_file) + + if self.cfg.use_length: + self._length_probabilities = self._read_length( + directory / "length_distribution.json" + ) + if self.cfg.use_relative_position: + self._position_histogram = _read_json( + directory / "relative_anchor_position_histogram.json" + ) + if self.cfg.use_local_context: + self._local_context = _read_json(directory / "local_context_counts.json") + + self._loaded = True + + def _read_probability_file(self, path: Path) -> dict[str, dict[str, Any]]: + """Extract the selected dedup version's category table.""" + payload = _read_json(path) + versions = payload.get("probabilities_by_version") or {} + version = self.cfg.dedup_version + if version not in versions: + available = sorted(versions) + raise PriorFilesMissingError( + f"dedup version {version!r} not in {path.name}; available: {available}" + ) + return dict(versions[version].get("categories") or {}) + + def _read_length(self, path: Path) -> dict[int, float]: + """Normalise the unbinned length counts into a probability table.""" + payload = _read_json(path) + counts = (payload.get("unbinned") or {}).get(self.cfg.dedup_version) or {} + total = sum(int(v) for v in counts.values()) + if total <= 0: + return {} + # Add-one smoothing over the observed support so a nearby unobserved + # length is penalised rather than treated as impossible. + support = len(counts) + return { + int(k): (int(v) + 1.0) / (total + support) + for k, v in counts.items() + } + + # -- scoring --------------------------------------------------------- + + def _probability(self, key: str) -> float | None: + """Look up ``key`` in the pair x spacing table.""" + entry = self._pair_spacing.get(key) + if entry is None: + return None + field_name = ( + "laplace_smoothed_probability" if self.cfg.use_smoothed else "raw_probability" + ) + value = entry.get(field_name) + return None if value is None else float(value) + + def score_endpoint( + self, + state: StapleState, + block: BuildingBlock | None, + context: dict[str, Any] | None = None, + ) -> EndpointEnergyBreakdown: + """Compute the endpoint prior energy for a terminal hydrocarbon state. + + Returns a zero breakdown with an explanatory ``status`` when the prior is + disabled or the state is not a scorable hydrocarbon endpoint, so callers + can always add ``total`` unconditionally. + """ + del context + breakdown = EndpointEnergyBreakdown() + + if not self.cfg.enabled: + breakdown.status = "disabled" + return breakdown + if not self._loaded: + breakdown.status = "not_loaded" + return breakdown + if block is None or not is_hydrocarbon_block(block): + breakdown.status = "not_hydrocarbon_block" + return breakdown + if state.anchor_pair is None: + breakdown.status = "no_anchor_pair" + return breakdown + + tokens = state.sequence_tokens + i, j = int(state.anchor_pair[0]), int(state.anchor_pair[1]) + if not (0 <= i < len(tokens)) or not (0 <= j < len(tokens)): + breakdown.status = "anchor_out_of_range" + return breakdown + + # Read the pair off the *state*, not the block, so a mislabelled block + # cannot silently score as its intended topology. + ordered_pair = f"{tokens[i].upper()}-{tokens[j].upper()}" + spacing = j - i + key = f"{ordered_pair}|{spacing}" + breakdown.pair_key = key + + probability = self._probability(key) + if probability is None or probability <= 0.0: + breakdown.e_pair = float(self.cfg.max_component_energy) + breakdown.status = "pair_spacing_unobserved" + breakdown.details["catalog_topology"] = "-".join( + str(x) for x in block_topology(block) + ) + else: + breakdown.e_pair = min( + -math.log(probability), float(self.cfg.max_component_energy) + ) + + breakdown.e_pair *= float(self.cfg.weight_pair) + + if self.cfg.use_length: + breakdown.e_length = self.cfg.weight_length * self._length_energy(len(tokens)) + if self.cfg.use_relative_position: + breakdown.e_relative_position = ( + self.cfg.weight_relative_position + * self._relative_position_energy(i, j, len(tokens)) + ) + if self.cfg.use_local_context: + breakdown.e_local_context = ( + self.cfg.weight_local_context + * self._local_context_energy(tokens, (i, j), ordered_pair) + ) + + breakdown.total = float( + breakdown.e_pair + + breakdown.e_length + + breakdown.e_relative_position + + breakdown.e_local_context + ) + return breakdown + + def _length_energy(self, length: int) -> float: + """-log p(peptide length).""" + if not self._length_probabilities: + return 0.0 + probability = self._length_probabilities.get(int(length)) + if probability is None or probability <= 0.0: + return float(self.cfg.max_component_energy) + return min(-math.log(probability), float(self.cfg.max_component_energy)) + + def _relative_position_energy(self, i: int, j: int, length: int) -> float: + """Mean -log p over the two anchors' relative positions.""" + histogram = self._position_histogram.get("all_anchors") or {} + probabilities = histogram.get("probabilities") or [] + edges = self._position_histogram.get("bin_edges") or [] + if not probabilities or len(edges) < 2 or length < 2: + return 0.0 + + n_bins = len(probabilities) + total = 0.0 + for index in (i, j): + relative = index / (length - 1) + bin_index = min(int(relative * n_bins), n_bins - 1) + probability = float(probabilities[bin_index]) + if probability <= 0.0: + total += float(self.cfg.max_component_energy) + else: + total += min( + -math.log(probability), float(self.cfg.max_component_energy) + ) + return total / 2.0 + + def _local_context_energy( + self, tokens: list[str], anchor_pair: tuple[int, int], ordered_pair: str + ) -> float: + """Mean -log p of the residues flanking each anchor. + + Uses the pair-specific table when the observed pair has one, otherwise + the pooled ``all_anchors`` table. + """ + groups = self._local_context.get("groups") or {} + group = groups.get(f"pair_{ordered_pair}") or groups.get("all_anchors") + if not group: + return 0.0 + offsets = group.get("offsets") or {} + padding = self._local_context.get("padding_tokens") or {} + n_pad = padding.get("n_terminal", "") + c_pad = padding.get("c_terminal", "") + + window = int(self.cfg.local_context_window) + energies: list[float] = [] + for anchor in anchor_pair: + for offset in range(-window, window + 1): + if offset == 0: + continue + entry = offsets.get(str(offset)) + if not entry: + continue + frequencies = entry.get("frequencies") or {} + target = anchor + offset + if target < 0: + token = n_pad + elif target >= len(tokens): + token = c_pad + else: + token = tokens[target].upper() + probability = float(frequencies.get(token, 0.0)) + if probability <= 0.0: + # Back off to a floor rather than +inf: an unseen flank is + # unusual, not impossible. + total_observed = max(int(entry.get("total", 0)), 1) + probability = 1.0 / (total_observed + 1.0) + energies.append( + min(-math.log(probability), float(self.cfg.max_component_energy)) + ) + return sum(energies) / len(energies) if energies else 0.0 + + # -- introspection --------------------------------------------------- + + def describe(self) -> dict[str, Any]: + """Summary of what was loaded, for logging and audits.""" + return { + "enabled": bool(self.cfg.enabled), + "loaded": bool(self._loaded), + "prior_dir": str(self.prior_dir), + "dedup_version": self.cfg.dedup_version, + "use_smoothed": bool(self.cfg.use_smoothed), + "n_pair_spacing_categories": len(self._pair_spacing), + "n_length_bins": len(self._length_probabilities), + "components": { + "pair": True, + "length": bool(self.cfg.use_length), + "relative_position": bool(self.cfg.use_relative_position), + "local_context": bool(self.cfg.use_local_context), + }, + "uses_permeability_label": False, + "is_trained_classifier": False, + } + + def observed_topologies(self) -> dict[str, float]: + """Observed ``PAIR|SPACING`` keys mapped to their energies, cheapest first.""" + out: dict[str, float] = {} + for key in self._pair_spacing: + probability = self._probability(key) + if probability and probability > 0.0: + out[key] = min( + -math.log(probability), float(self.cfg.max_component_energy) + ) + return dict(sorted(out.items(), key=lambda kv: kv[1])) + + +def count_anchor_monomers(tokens: list[str]) -> int: + """Number of hydrocarbon anchor monomers present in ``tokens``.""" + return sum(1 for t in tokens if is_anchor_token(t)) + + +def _read_json(path: Path) -> dict[str, Any]: + """Read one JSON file, with a clear error when it is missing.""" + if not path.is_file(): + raise PriorFilesMissingError(f"missing required prior file: {path}") + with path.open("r", encoding="utf-8") as handle: + return json.load(handle) + + +def _repository_root() -> Path: + """Repository root, derived from this file's location.""" + return Path(__file__).resolve().parents[2] diff --git a/staplebridge/hydrocarbon/exact_sb_cache.py b/staplebridge/hydrocarbon/exact_sb_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..7aa1b6a8f089dec5e8348bb6b9ef79a9a962115d --- /dev/null +++ b/staplebridge/hydrocarbon/exact_sb_cache.py @@ -0,0 +1,692 @@ +"""Persistent cache for the deterministic part of the Exact-SB plan target. + +For a fixed lead and a fixed configuration, everything on the right-hand side of + + q*(p|x) ∝ q_ref(p|x) · exp[-beta · E_T(x, p)] + +is deterministic: the legal-plan enumeration, ``q_ref``, the PeptiVerse-backed +terminal energies, and therefore ``log q*`` itself. Only ``q_theta`` changes as +the plan head trains. This module persists the deterministic half so each epoch +recomputes just ``q_theta`` and then ``KL(q*||q_theta)``. + +Deliberately **not** cached: + +* ``q_theta`` — it is a function of the live model weights. Caching or freezing + it would silently stop training the plan head. Nothing in this module reads, + writes, or accepts ``q_theta``. +* Anything used for candidate ranking, decoding, or the legacy validation + metrics. The cache only replays plan order, ``reference_logp``, terminal + energies and ``log q*``; every consumer recomputes ``q_theta`` itself. + +Correctness rests on the fingerprint: any change to the lead, catalog, +empirical prior, ``exact_sb_beta``, terminal-energy/property configuration, +PeptiVerse model set, or geometry/edit settings produces a different +fingerprint, so a stale entry is never served. Entries also store the plan +signatures they were built from and are rejected if the live enumeration +disagrees, which catches drift the fingerprint alone would miss. +""" + +from __future__ import annotations + +import hashlib +import inspect +import json +import os +import sqlite3 +import threading +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Iterable, Sequence + +import torch + +from staplebridge.hydrocarbon.curriculum import HydrocarbonStaplePlan + +#: Bumped whenever the stored payload's meaning changes. Old rows then miss. +CACHE_SCHEMA_VERSION = 1 + + +def _stable_json(payload: Any) -> str: + return json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + + +def _sha256(text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def plan_signature(plan: HydrocarbonStaplePlan) -> str: + """Identity of a plan: everything that changes its terminal state.""" + return _stable_json( + [ + str(plan.block_id), + [int(plan.anchor_pair[0]), int(plan.anchor_pair[1])], + str(plan.ordered_pair), + int(plan.spacing), + [[int(position), str(monomer)] for position, monomer in plan.substitutions], + ] + ) + + +def plan_signatures(plans: Sequence[HydrocarbonStaplePlan]) -> list[str]: + return [plan_signature(plan) for plan in plans] + + +def _directory_digest(root: Path, suffixes: tuple[str, ...]) -> str: + """Digest of (relative path, size, mtime) for weight files under ``root``. + + Content hashing 400 MB of PeptiVerse weights on every run would cost more + than the cache saves, so identity is (path, size, mtime-ns). Touching or + swapping a weight file therefore invalidates the cache. + """ + if not root.is_dir(): + return f"missing:{root}" + entries: list[tuple[str, int, int]] = [] + for path in sorted(root.rglob("*")): + if not path.is_file() or path.suffix.lower() not in suffixes: + continue + stat = path.stat() + entries.append((str(path.relative_to(root)), int(stat.st_size), int(stat.st_mtime_ns))) + return _sha256(_stable_json(entries)) + + +def _prior_digest(prior_dir: Path) -> str: + if not prior_dir.is_dir(): + return f"missing:{prior_dir}" + entries = [] + for path in sorted(prior_dir.rglob("*")): + if path.is_file() and path.suffix.lower() in (".json", ".csv", ".tsv", ".yaml", ".yml"): + stat = path.stat() + entries.append((str(path.relative_to(prior_dir)), int(stat.st_size), int(stat.st_mtime_ns))) + return _sha256(_stable_json(entries)) + + +def build_fingerprint( + config: dict[str, Any], + *, + catalog: Iterable[Any], + repo_root: Path | None = None, +) -> dict[str, Any]: + """Everything ``log q*`` depends on, other than the lead itself. + + Any mismatch in any component invalidates cached entries automatically, + because the fingerprint hash is part of every row's key. + """ + repo_root = repo_root or Path.cwd() + hydro = dict(config.get("hydrocarbon") or {}) + plan_control = dict(hydro.get("plan_control") or {}) + terminal = dict(hydro.get("terminal_energy") or {}) + property_cfg = dict(terminal.get("property") or {}) + predictor = dict(config.get("property_predictor") or {}) + priors = dict(config.get("reference_priors") or {}) + peptide_prior = dict(priors.get("peptide") or {}) + plan_reference = dict(hydro.get("plan_reference") or {}) + mode_prior = dict(plan_reference.get("mode_prior") or {}) + training = dict(config.get("training") or {}) + + # --- catalog identity ------------------------------------------------- + catalog_entries = [] + for block in catalog: + catalog_entries.append( + { + "block_id": getattr(block, "block_id", None), + "name": getattr(block, "name", None), + "chemistry_class": getattr(block, "chemistry_class", None), + "motif": getattr(block, "motif", None), + "ca_window": getattr(block, "ca_window", None), + "cost_score": getattr(block, "cost_score", None), + "spps_score": getattr(block, "spps_score", None), + } + ) + + prior_dir = mode_prior.get("prior_dir") + components = { + "schema_version": CACHE_SCHEMA_VERSION, + # --- exact-SB target definition --- + "exact_sb_beta": float(plan_control.get("exact_sb_beta", 1.0)), + "exact_sb_objective": bool(plan_control.get("exact_sb_objective", False)), + # Hard constraints define the support on which q_ref, q*, and q_theta + # are conditioned. Version this explicitly so pre-mask cache rows can + # never collide with post-mask targets. + "property_free_hard_plan_predicate": "v1", + "hard_plan_geometry_eps": float( + config.get("eps_geom", training.get("eps_geom", 2.5)) + ), + # --- catalog / empirical prior version --- + "catalog": catalog_entries, + "catalog_config": dict(hydro.get("catalog") or {}), + "mode_prior": mode_prior, + "mode_prior_digest": _prior_digest( + (repo_root / str(prior_dir)) if prior_dir else repo_root / "__missing__" + ), + "plan_reference_enabled": bool(plan_reference.get("enabled", False)), + "plan_reference_bias": dict(plan_reference.get("bias") or {}), + "factorized_plan_reference": bool( + (hydro.get("reference") or {}).get("factorized_plan_reference", False) + ), + # --- terminal energy / property config --- + "terminal_energy": { + key: terminal[key] for key in sorted(terminal) if key != "property" + }, + "property": property_cfg, + "endpoint_prior": dict(hydro.get("endpoint_prior") or {}), + # These base-terminal coefficients feed E_T through base_terminal_factory. + "base_terminal_coefficients": { + key: training.get(key) + for key in ("lambda_close", "lambda_edit", "lambda_cost", "infeasible_penalty") + }, + # --- PeptiVerse model / config --- + "peptiverse": { + key: predictor.get(key) + for key in ( + "backend", "mode", "strict", "enable_fallback", + "allow_wt_token_fallback", "uncertainty", "offline", + "peptiverse_root", "classifier_weight_root", "manifest_path", + ) + }, + "peptiverse_manifest_digest": ( + _sha256(Path(str(predictor["manifest_path"])).read_text()) + if predictor.get("manifest_path") and Path(str(predictor["manifest_path"])).is_file() + else "missing_manifest" + ), + "peptiverse_weights_digest": _directory_digest( + Path(str(predictor.get("classifier_weight_root", ""))) / "training_classifiers", + (".pt", ".json", ".joblib", ".bin", ".safetensors", ".txt"), + ), + # --- reference priors (ESM-2 identity affects q_ref via the sampler) --- + "peptide_prior": { + key: peptide_prior.get(key) + for key in ("backend", "model_name_or_path", "temperature", "offline", + "strict_runtime", "ncaa_policy") + }, + "anchor_prior": dict(priors.get("anchor") or {}), + "block_prior": dict(priors.get("block") or {}), + "reference_energy": dict(config.get("reference") or {}), + # --- geometry / edit settings --- + "geometry": dict(hydro.get("geometry") or {}), + "edit_constraints": dict(config.get("edit_constraints") or {}), + "curriculum_limits": { + key: (hydro.get("curriculum") or {}).get(key) + for key in ("max_anchor_edits", "prefer_existing_anchors", "protect_positions", + "require_valid_terminal") + }, + "actions": dict(hydro.get("actions") or {}), + "max_neighbors": config.get("max_neighbors"), + "chemistry": config.get("chemistry"), + } + return {"hash": _sha256(_stable_json(components)), "components": components} + + +def lead_key(lead: Any) -> str: + """Lead identity *and* content, so an edited sequence cannot reuse a row.""" + return _stable_json( + { + "example_id": str(getattr(lead, "example_id", "")), + "linear_sequence": str(getattr(lead, "linear_sequence", "")), + "protected_positions": sorted(int(p) for p in (getattr(lead, "protected_positions", None) or [])), + # peptide_ca participates in the geometry term of E_T. + "peptide_ca": (getattr(lead, "target_context", None) or {}).get("peptide_ca"), + } + ) + + +@dataclass +class ExactSBTargetEntry: + """The deterministic half of the Exact-SB target for one lead.""" + + plan_signatures: list[str] + reference_logp: list[float] + terminal_energies: list[float] + log_q_star: list[float] + target_support_mask: list[bool] | None = None + joint_plan_count: int | None = None + + def as_json(self) -> str: + payload: dict[str, Any] = { + "plan_signatures": self.plan_signatures, + "reference_logp": self.reference_logp, + "terminal_energies": self.terminal_energies, + "log_q_star": self.log_q_star, + } + # Preserve the byte-level shape of historical false-flag cache rows. + if self.target_support_mask is not None or self.joint_plan_count is not None: + payload["target_support_mask"] = self.target_support_mask + payload["joint_plan_count"] = self.joint_plan_count + return _stable_json(payload) + + @classmethod + def from_json(cls, text: str) -> "ExactSBTargetEntry": + payload = json.loads(text) + return cls( + plan_signatures=list(payload["plan_signatures"]), + reference_logp=[float(v) for v in payload["reference_logp"]], + terminal_energies=[float(v) for v in payload["terminal_energies"]], + log_q_star=[float(v) for v in payload["log_q_star"]], + target_support_mask=( + [bool(v) for v in payload["target_support_mask"]] + if payload.get("target_support_mask") is not None + else None + ), + joint_plan_count=( + int(payload["joint_plan_count"]) + if payload.get("joint_plan_count") is not None + else None + ), + ) + + +@dataclass +class ExactSBCacheStats: + hits: int = 0 + misses: int = 0 + signature_mismatches: int = 0 + writes: int = 0 + plans_recomputed: int = 0 + plans_served_from_cache: int = 0 + energy_calls_saved: int = 0 + + def as_dict(self) -> dict[str, Any]: + total = self.hits + self.misses + return { + "hits": self.hits, + "misses": self.misses, + "lookups": total, + "hit_rate": (self.hits / total) if total else None, + "signature_mismatches": self.signature_mismatches, + "writes": self.writes, + "plans_recomputed": self.plans_recomputed, + "plans_served_from_cache": self.plans_served_from_cache, + "terminal_energy_calls_saved": self.energy_calls_saved, + } + + +class ExactSBTargetCache: + """SQLite-backed store for per-lead Exact-SB targets. + + SQLite (WAL, one row per lead) is used rather than one file per lead so a + 4020-lead run does not create 4020 files, and so concurrent readers during + resume are safe. ``read_only=True`` gives a disabled/A-B arm that never + writes. + """ + + def __init__( + self, + path: Path | str | None, + fingerprint: dict[str, Any], + *, + enabled: bool = True, + read_only: bool = False, + ) -> None: + self.enabled = bool(enabled and path is not None) + self.read_only = bool(read_only) + self.fingerprint_hash = str(fingerprint["hash"]) + self.fingerprint_components = fingerprint.get("components", {}) + self.path = Path(path) if path is not None else None + self.stats = ExactSBCacheStats() + self._lock = threading.Lock() + self._connection: sqlite3.Connection | None = None + if self.enabled: + self._open() + + # -- storage --------------------------------------------------------- + def _open(self) -> None: + assert self.path is not None + self.path.parent.mkdir(parents=True, exist_ok=True) + self._connection = sqlite3.connect(str(self.path), check_same_thread=False) + self._connection.execute("PRAGMA journal_mode=WAL") + self._connection.execute("PRAGMA synchronous=NORMAL") + self._connection.executescript( + """ + CREATE TABLE IF NOT EXISTS exact_sb_targets ( + fingerprint TEXT NOT NULL, + lead_key TEXT NOT NULL, + payload TEXT NOT NULL, + n_plans INTEGER NOT NULL, + PRIMARY KEY (fingerprint, lead_key) + ); + CREATE TABLE IF NOT EXISTS fingerprints ( + fingerprint TEXT PRIMARY KEY, + components TEXT NOT NULL + ); + """ + ) + if not self.read_only: + self._connection.execute( + "INSERT OR REPLACE INTO fingerprints (fingerprint, components) VALUES (?, ?)", + (self.fingerprint_hash, _stable_json(self.fingerprint_components)), + ) + self._connection.commit() + + def close(self) -> None: + with self._lock: + if self._connection is not None: + self._connection.commit() + self._connection.close() + self._connection = None + + # -- lookup / store -------------------------------------------------- + def get(self, lead: Any, plans: Sequence[HydrocarbonStaplePlan]) -> ExactSBTargetEntry | None: + """Return the cached target, or ``None`` on miss or plan drift.""" + if not self.enabled or self._connection is None: + self.stats.misses += 1 + return None + key = lead_key(lead) + with self._lock: + row = self._connection.execute( + "SELECT payload FROM exact_sb_targets WHERE fingerprint = ? AND lead_key = ?", + (self.fingerprint_hash, key), + ).fetchone() + if row is None: + self.stats.misses += 1 + return None + entry = ExactSBTargetEntry.from_json(row[0]) + # Defence in depth: even on a fingerprint match, the live enumeration + # must produce exactly the same plans in the same order. + if entry.plan_signatures != plan_signatures(plans): + self.stats.signature_mismatches += 1 + self.stats.misses += 1 + return None + self.stats.hits += 1 + self.stats.plans_served_from_cache += len(entry.plan_signatures) + self.stats.energy_calls_saved += len(entry.plan_signatures) + return entry + + def put(self, lead: Any, entry: ExactSBTargetEntry) -> None: + if not self.enabled or self.read_only or self._connection is None: + return + with self._lock: + self._connection.execute( + "INSERT OR REPLACE INTO exact_sb_targets " + "(fingerprint, lead_key, payload, n_plans) VALUES (?, ?, ?, ?)", + ( + self.fingerprint_hash, + lead_key(lead), + entry.as_json(), + len(entry.plan_signatures), + ), + ) + self._connection.commit() + self.stats.writes += 1 + + # -- diagnostics ----------------------------------------------------- + def disk_bytes(self) -> int: + if self.path is None or not self.path.is_file(): + return 0 + total = self.path.stat().st_size + for suffix in ("-wal", "-shm"): + side = self.path.with_name(self.path.name + suffix) + if side.is_file(): + total += side.stat().st_size + return int(total) + + def row_count(self) -> int: + if not self.enabled or self._connection is None: + return 0 + with self._lock: + return int( + self._connection.execute( + "SELECT COUNT(*) FROM exact_sb_targets WHERE fingerprint = ?", + (self.fingerprint_hash,), + ).fetchone()[0] + ) + + def describe(self) -> dict[str, Any]: + return { + "enabled": self.enabled, + "read_only": self.read_only, + "path": None if self.path is None else str(self.path), + "fingerprint": self.fingerprint_hash, + "rows_for_fingerprint": self.row_count(), + "disk_bytes": self.disk_bytes(), + **self.stats.as_dict(), + } + + +def cache_path_from_config(config: dict[str, Any], repo_root: Path | None = None) -> Path | None: + """Resolve ``hydrocarbon.plan_control.exact_sb_cache.path``; ``None`` = off.""" + plan_control = dict(((config.get("hydrocarbon") or {}).get("plan_control") or {})) + section = dict(plan_control.get("exact_sb_cache") or {}) + if not section or not bool(section.get("enabled", False)): + return None + raw = str(section.get("path") or "outputs/cache/exact_sb_targets.sqlite") + path = Path(raw) + if not path.is_absolute(): + path = (repo_root or Path.cwd()) / path + return path + + +def energy_only_from_config(config: dict[str, Any]) -> bool: + """Read ``hydrocarbon.plan_control.exact_sb_cache.energy_only`` (default on). + + Off gives the original all-property scalar construction, which the A/B + benchmark uses as the reference arm. + """ + plan_control = dict(((config.get("hydrocarbon") or {}).get("plan_control") or {})) + section = dict(plan_control.get("exact_sb_cache") or {}) + return bool(section.get("energy_only", True)) + + +def build_cache_from_config( + config: dict[str, Any], + *, + catalog: Iterable[Any], + repo_root: Path | None = None, + read_only: bool = False, + override_path: Path | None = None, +) -> ExactSBTargetCache: + """Construct the cache declared by ``config`` (disabled when absent).""" + catalog = list(catalog) + fingerprint = build_fingerprint(config, catalog=catalog, repo_root=repo_root) + path = override_path if override_path is not None else cache_path_from_config(config, repo_root) + return ExactSBTargetCache(path, fingerprint, enabled=path is not None, read_only=read_only) + + +def _accepts_energy_only(energy_fn: Any) -> bool: + """Whether ``energy_fn`` takes the ``energy_only`` keyword.""" + target = energy_fn.__call__ if not inspect.isfunction(energy_fn) else energy_fn + try: + signature = inspect.signature(target) + except (TypeError, ValueError): + return False + parameters = signature.parameters + if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values()): + return True + return "energy_only" in parameters + + +def resolve_exact_sb_target( + *, + lead: Any, + plans: Sequence[HydrocarbonStaplePlan], + reference_log_probabilities: torch.Tensor, + beta: float, + energy_fn: Any, + initial_state: Any, + build_terminal: Any, + cache: ExactSBTargetCache | None, + energy_only: bool = False, + scorer: Any = None, + scorer_config: Any = None, +) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Return ``(terminal_energies, log_q_star, info)`` for one lead. + + On a cache hit the PeptiVerse-backed terminal energies are replayed from + disk. On a miss they are computed exactly as before and then stored. + ``log q*`` is always produced by the shared + :func:`exact_sb_target_log_probabilities`, so the cached and recomputed + paths cannot diverge in definition. + + ``energy_only`` restricts property prediction to the properties that reach + the terminal energy and batch-prefetches their SMILES across all of this + lead's plans. It is a pure speedup: the energies, and therefore ``q*``, are + unchanged, so a cache row written by either path is interchangeable. It + requires ``scorer``/``scorer_config`` and an ``energy_fn`` accepting the + ``energy_only`` keyword; without them the original scalar path runs. + """ + # Imported here to keep plan_control free of a dependency on this module. + from staplebridge.hydrocarbon.plan_control import ( + exact_sb_target_log_probabilities, + ) + from staplebridge.hydrocarbon.property_energy import required_energy_properties + + device = reference_log_probabilities.device + dtype = reference_log_probabilities.dtype + if scorer_config is None: + scorer_config = getattr(energy_fn, "property_cfg", None) + joint_enabled = bool( + getattr(scorer_config, "enable_joint_perm_halflife_support", False) + ) + + def support_info( + mask_values: list[bool] | None, joint_count: int | None + ) -> dict[str, Any]: + if not joint_enabled: + return { + "joint_perm_halflife_support_enabled": False, + "target_support_mask": None, + "joint_plan_count": None, + "joint_nonempty": None, + "joint_fallback": None, + "q_star_support_size": len(plans), + } + count = int(joint_count or 0) + fallback = count == 0 + effective = None if fallback else list(mask_values or []) + return { + "joint_perm_halflife_support_enabled": True, + "target_support_mask": effective, + "joint_plan_count": count, + "joint_nonempty": not fallback, + "joint_fallback": fallback, + "q_star_support_size": count if count else len(plans), + } + + entry = None if cache is None else cache.get(lead, plans) + if entry is not None: + if joint_enabled and entry.joint_plan_count is None: + # Backward-compatible payloads do not carry enough information to + # reconstruct a strict target support. A matching new-arm + # fingerprint should make this impossible, but fail closed. + entry = None + if entry is not None: + energies = torch.tensor(entry.terminal_energies, dtype=dtype, device=device) + mask = ( + torch.tensor(entry.target_support_mask, dtype=torch.bool, device=device) + if joint_enabled + and int(entry.joint_plan_count or 0) > 0 + and entry.target_support_mask is not None + else None + ) + # Recomputed from the cached energies rather than trusting the stored + # log_q_star blindly; the stored copy is then verified against it. + log_q_star = exact_sb_target_log_probabilities( + reference_log_probabilities, + energies, + beta, + target_support_mask=mask, + ) + stored = torch.tensor(entry.log_q_star, dtype=dtype, device=device) + finite = torch.isfinite(log_q_star) & torch.isfinite(stored) + max_drift = ( + float((log_q_star[finite] - stored[finite]).abs().max().item()) + if bool(finite.any().item()) + else 0.0 + ) + if not torch.equal(torch.isneginf(log_q_star), torch.isneginf(stored)): + max_drift = float("inf") + return energies, log_q_star, { + "source": "cache", + "log_q_star_drift": max_drift, + **support_info(entry.target_support_mask, entry.joint_plan_count), + } + + energies_list: list[float] = [] + terminals = [build_terminal(initial_state, plan) for plan in plans] + + # Energy-only + batched prefetch. Both are pure accelerations: the prefetch + # only warms caches, and energy_only skips predictions that are provably not + # summed into the energy. When either is unavailable the loop below is the + # original scalar path, so the stored energies are the same either way. + prefetch_info: dict[str, Any] = {} + if energy_only: + # Only the hydrocarbon terminal energy accepts the keyword and carries a + # scorer; any other callable (tests, lactam-style stubs) keeps the + # original scalar path rather than being handed an argument it rejects. + if not _accepts_energy_only(energy_fn): + energy_only = False + if energy_only: + # HydrocarbonTerminalEnergy carries both; taking them from energy_fn + # keeps the two knobs in sync with the energy that will actually run. + if scorer is None: + scorer = getattr(energy_fn, "property_scorer", None) + if scorer_config is None: + scorer_config = getattr(energy_fn, "property_cfg", None) + if scorer is None: + energy_only = False + if energy_only and scorer is not None: + properties = required_energy_properties(scorer_config) if scorer_config else () + if properties: + smiles: list[str] = [] + for terminal in terminals: + try: + smiles.extend(scorer.energy_only_smiles(terminal)) + except Exception: # noqa: BLE001 + # An unscorable terminal is handled by the energy itself + # (topology gate); nothing to prefetch for it. + continue + if smiles: + prefetch_info = scorer.prefetch(properties, smiles) + + joint_mask_values: list[bool] = [] + for terminal in terminals: + if energy_only: + energy, terms = energy_fn(initial_state, terminal, lead, energy_only=True) + else: + energy, terms = energy_fn(initial_state, terminal, lead) + energies_list.append(float(energy)) + if joint_enabled: + if "hydrocarbon_joint_perm_halflife_condition" not in terms: + raise RuntimeError( + "joint Exact-SB support enabled but terminal energy did not " + "report the joint condition" + ) + joint_mask_values.append( + bool(terms["hydrocarbon_joint_perm_halflife_condition"]) + ) + energies = torch.tensor(energies_list, dtype=dtype, device=device) + joint_count = sum(joint_mask_values) if joint_enabled else None + target_mask = ( + torch.tensor(joint_mask_values, dtype=torch.bool, device=device) + if joint_enabled and int(joint_count or 0) > 0 + else None + ) + log_q_star = exact_sb_target_log_probabilities( + reference_log_probabilities, + energies, + beta, + target_support_mask=target_mask, + ) + if cache is not None: + cache.stats.plans_recomputed += len(plans) + cache.put( + lead, + ExactSBTargetEntry( + plan_signatures=plan_signatures(plans), + reference_logp=[float(v) for v in reference_log_probabilities.detach().cpu().tolist()], + terminal_energies=energies_list, + log_q_star=[float(v) for v in log_q_star.detach().cpu().tolist()], + target_support_mask=( + list(joint_mask_values) if joint_enabled else None + ), + joint_plan_count=joint_count, + ), + ) + return energies, log_q_star, { + "source": "computed", + "log_q_star_drift": 0.0, + "energy_only": bool(energy_only), + "prefetch": prefetch_info, + **support_info( + list(joint_mask_values) if joint_enabled else None, + joint_count, + ), + } diff --git a/staplebridge/hydrocarbon/factorized_plan_reference.py b/staplebridge/hydrocarbon/factorized_plan_reference.py new file mode 100644 index 0000000000000000000000000000000000000000..296ae501966ec812ddf3f85bfff2f82263c2ff72 --- /dev/null +++ b/staplebridge/hydrocarbon/factorized_plan_reference.py @@ -0,0 +1,220 @@ +"""Topology-mass-preserving hydrocarbon plan reference. + +This module is hydrocarbon-only. It does not import or modify the lactam +catalog, decoder, property model, SMILES builder, plan space, or loss. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +from staplebridge.chemistry.state import StapleState +from staplebridge.data.schemas import BuildingBlock +from staplebridge.hydrocarbon.curriculum import ( + HydrocarbonStaplePlan, + build_hydrocarbon_demonstration_path, +) + + +@dataclass +class FactorizedPlanReferenceConfig: + """Configuration read from ``hydrocarbon.reference``. + + Defaults were preregistered without property labels from the component + scale audit: MotifSupportAnchorPrior/geometry is primary and frozen ESM2 + delta is only a weak regularizer. CatalogBlockPrior is a legality check and + diagnostic, not a soft ranking term. + """ + + enabled: bool = False + geometry_coefficient: float = 1.0 + esm2_coefficient: float = 0.1 + temperature: float = 1.0 + block_legality_only: bool = True + + @classmethod + def from_config( + cls, root_cfg: dict[str, Any] | None + ) -> "FactorizedPlanReferenceConfig": + root_cfg = dict(root_cfg or {}) + hydro = dict(root_cfg.get("hydrocarbon") or {}) + section = dict(hydro.get("reference") or {}) + within = dict(section.get("within_mode") or {}) + return cls( + enabled=bool(section.get("factorized_plan_reference", False)), + geometry_coefficient=float(within.get("geometry_coefficient", 1.0)), + esm2_coefficient=float(within.get("esm2_coefficient", 0.1)), + temperature=max(float(within.get("temperature", 1.0)), 1e-8), + block_legality_only=bool(within.get("block_legality_only", True)), + ) + + def describe(self) -> dict[str, Any]: + return { + "factorized_plan_reference": bool(self.enabled), + "geometry_coefficient": float(self.geometry_coefficient), + "esm2_coefficient": float(self.esm2_coefficient), + "temperature": float(self.temperature), + "block_legality_only": bool(self.block_legality_only), + "uses_property_labels": False, + "normalization": "softmax separately within each feasible mode", + } + + +class FactorizedPlanReference: + """Compute ``q_mode(mode|lead) * q_within(plan|lead,mode)``. + + The supplied ``mode_prior`` owns the StaPep probability and empirical + beta. Within-mode components are normalized separately, so they cannot + change the total mass assigned to a topology. + """ + + def __init__( + self, + mode_prior: Any, + catalog: list[BuildingBlock], + peptide_prior: Any, + anchor_prior: Any, + block_prior: Any, + config: FactorizedPlanReferenceConfig, + ) -> None: + self.mode_prior = mode_prior + self.catalog = list(catalog) + self.catalog_index = {block.block_id: block for block in catalog} + self.peptide_prior = peptide_prior + self.anchor_prior = anchor_prior + self.block_prior = block_prior + self.cfg = config + self.last_diagnostics: list[dict[str, Any]] = [] + + @staticmethod + def _softmax(values: list[float]) -> list[float]: + if not values: + return [] + peak = max(values) + exponentials = [math.exp(value - peak) for value in values] + total = sum(exponentials) + if total <= 0.0 or not math.isfinite(total): + return [1.0 / len(values)] * len(values) + probabilities = [value / total for value in exponentials] + if len(probabilities) > 1: + probabilities[-1] = 1.0 - sum(probabilities[:-1]) + return probabilities + + def weights( + self, + initial: StapleState, + plans: list[HydrocarbonStaplePlan], + context: dict[str, Any] | None = None, + ) -> list[float]: + """Return normalized factorized probabilities in ``plans`` order.""" + if not plans: + self.last_diagnostics = [] + return [] + + context = dict(context or {}) + terminals = [ + build_hydrocarbon_demonstration_path(initial, plan, self.catalog)[-1] + for plan in plans + ] + esm2_scores = self.peptide_prior.batch_score_transitions( + initial, terminals, context + ) + + feasible_modes: list[tuple[str, int]] = [] + for plan in plans: + mode = (plan.ordered_pair, plan.spacing) + if mode not in feasible_modes: + feasible_modes.append(mode) + tilted = [self.mode_prior.tilted_weight(mode) for mode in feasible_modes] + tilted_total = sum(tilted) + if tilted_total <= 0.0: + self.last_diagnostics = [] + return [0.0] * len(plans) + q_mode = { + mode: weight / tilted_total for mode, weight in zip(feasible_modes, tilted) + } + if len(feasible_modes) > 1: + q_mode[feasible_modes[-1]] = 1.0 - sum( + q_mode[mode] for mode in feasible_modes[:-1] + ) + + logits: list[float] = [] + diagnostics: list[dict[str, Any]] = [] + for plan, terminal, esm2_score in zip(plans, terminals, esm2_scores): + block = self.catalog_index.get(plan.block_id) + if block is None: + anchor_score = float("-inf") + block_score = float("-inf") + anchor_components: dict[str, float] = {} + legal = False + else: + anchor_context = dict(context) + anchor_context["return_components"] = True + anchor_score = float( + self.anchor_prior.score_anchor( + terminal.sequence_tokens, plan.anchor_pair, anchor_context + ) + ) + anchor_components = dict(anchor_context.get("_components") or {}) + block_score = float( + self.block_prior.score_block( + terminal.sequence_tokens, plan.anchor_pair, block, context + ) + ) + legal = math.isfinite(anchor_score) and math.isfinite(block_score) + + logit = ( + self.cfg.geometry_coefficient * anchor_score + + self.cfg.esm2_coefficient * float(esm2_score) + ) / self.cfg.temperature + if not legal: + logit = float("-inf") + logits.append(float(logit)) + diagnostics.append( + { + "mode": f"{plan.ordered_pair}/i,i+{plan.spacing}", + "anchor_pair": list(plan.anchor_pair), + "block_id": plan.block_id, + "stapep_probability": self.mode_prior.probability( + (plan.ordered_pair, plan.spacing) + ), + "stapep_tilted_weight": self.mode_prior.tilted_weight( + (plan.ordered_pair, plan.spacing) + ), + "q_mode_target": q_mode[(plan.ordered_pair, plan.spacing)], + "anchor_score_raw": anchor_score, + "anchor_components": anchor_components, + "esm2_delta_raw": float(esm2_score), + "block_score_diagnostic_only": block_score, + "block_legal": bool(legal), + "within_mode_logit": float(logit), + } + ) + + weights = [0.0] * len(plans) + for mode in feasible_modes: + indices = [ + index + for index, plan in enumerate(plans) + if (plan.ordered_pair, plan.spacing) == mode + ] + mode_logits = [logits[index] for index in indices] + finite = [math.isfinite(value) for value in mode_logits] + if not any(finite): + continue + masked = [value if ok else -1e30 for value, ok in zip(mode_logits, finite)] + within = self._softmax(masked) + for local_index, plan_index in enumerate(indices): + diagnostics[plan_index]["q_within_mode"] = float(within[local_index]) + weights[plan_index] = q_mode[mode] * within[local_index] + if len(indices) > 1: + weights[indices[-1]] = q_mode[mode] - sum( + weights[index] for index in indices[:-1] + ) + + for index, weight in enumerate(weights): + diagnostics[index]["q_ref"] = float(weight) + self.last_diagnostics = diagnostics + return weights diff --git a/staplebridge/hydrocarbon/geometry.py b/staplebridge/hydrocarbon/geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..0ee3863d542e6936db6562c27f1a94d2f28124fb --- /dev/null +++ b/staplebridge/hydrocarbon/geometry.py @@ -0,0 +1,102 @@ +"""Hydrocarbon geometry oracle. + +The lactam oracles in :mod:`staplebridge.oracles.geometry` (``MockGeometryOracle``, +``CPComposerStapledOracle``) are **not modified**. This is a separate oracle for +the hydrocarbon branch, because the two chemistries have genuinely different +feasibility criteria: + +* lactam K-D/E closes at i,i+3 / i,i+4 with a 4.0-6.5 A Cα window; +* hydrocarbon RCM staples span one helical turn (i,i+4) or two (i,i+7), so the + i,i+7 window sits much further out. + +``ctype`` here additionally requires that the anchor residues really are the +anchor monomers the block names, which the lactam oracle has no notion of. +""" + +from __future__ import annotations + +import math +from typing import Any + +from staplebridge.data.schemas import BuildingBlock +from staplebridge.hydrocarbon.actions import FailureReason, validate_hydrocarbon_staple +from staplebridge.hydrocarbon.catalog import is_hydrocarbon_block +from staplebridge.oracles.base import GeometryOracleBase + + +def _ca_distance( + coords: list[tuple[float, float, float]] | None, i: int, j: int +) -> float | None: + """Cα(i)-Cα(j) distance, or ``None`` when coordinates are unavailable.""" + if coords is None or i < 0 or j < 0 or i >= len(coords) or j >= len(coords): + return None + xi, yi, zi = coords[i] + xj, yj, zj = coords[j] + return math.sqrt((xi - xj) ** 2 + (yi - yj) ** 2 + (zi - zj) ** 2) + + +class HydrocarbonGeometryOracle(GeometryOracleBase): + """Feasibility for hydrocarbon staples. + + Args: + catalog: the hydrocarbon blocks in play, used for topology validation. + sentinel_cgeom: penalty returned when the staple is not even + type-feasible. Matches the lactam oracles' convention of a large + sentinel so that stage-aware geometry can detect it. + """ + + def __init__( + self, + catalog: list[BuildingBlock] | None = None, + sentinel_cgeom: float = 10.0, + ) -> None: + self.catalog = list(catalog or []) + self.sentinel_cgeom = float(sentinel_cgeom) + + def ctype( + self, + sequence: list[str], + anchor_pair: tuple[int, int] | None, + block: BuildingBlock | None, + *, + peptide_ca: Any = None, + ) -> bool: + """True when ``(sequence, anchor_pair, block)`` is a legal hydrocarbon staple. + + Sequence-level only; the Cα window is scored by :meth:`cgeom`. + """ + del peptide_ca + if block is None or not is_hydrocarbon_block(block): + return False + catalog = self.catalog or [block] + verdict = validate_hydrocarbon_staple(sequence, anchor_pair, block, catalog) + return verdict is FailureReason.OK + + def cgeom( + self, + sequence: list[str], + anchor_pair: tuple[int, int] | None, + block: BuildingBlock | None, + *, + peptide_ca: Any = None, + ) -> float: + """Distance of Cα(i)-Cα(j) from the block's allowed window, in Å. + + Returns ``0.0`` inside the window. Without coordinates, falls back to a + sequence-only verdict: ``0.0`` when type-feasible, else the sentinel. + """ + if not self.ctype(sequence, anchor_pair, block): + return self.sentinel_cgeom + assert anchor_pair is not None and block is not None # ctype guarantees + + i, j = anchor_pair + distance = _ca_distance(peptide_ca, i, j) + if distance is None: + return 0.0 + + low, high = block.ca_window + if distance < low: + return float(low - distance) + if distance > high: + return float(distance - high) + return 0.0 diff --git a/staplebridge/hydrocarbon/monomers.py b/staplebridge/hydrocarbon/monomers.py new file mode 100644 index 0000000000000000000000000000000000000000..44583c35ab9fd425f91a1d76f217e2a91ceee090 --- /dev/null +++ b/staplebridge/hydrocarbon/monomers.py @@ -0,0 +1,419 @@ +"""Hydrocarbon monomer library. + +Self-contained SMILES fragments for building hydrocarbon-stapled peptides. This +module is **additive and hydrocarbon-only**: the lactam path has no SMILES +builder at all, so nothing here replaces or shadows existing behaviour. + +Fragment convention +------------------- +Every residue is stored as a *backbone-open* fragment with two named attachment +points, written so that concatenating them forms a normal peptide chain: + + N-terminus of residue k <-- amide bond --> C-terminus of residue k-1 + +Each entry carries the fragment as an explicit SMILES with a dummy atom at each +end (``[*:1]`` for the incoming amide nitrogen side, ``[*:2]`` for the outgoing +carbonyl side). Bonds are then formed programmatically in +:mod:`staplebridge.hydrocarbon.smiles_builder` rather than by string +concatenation, because string splicing cannot keep ring-closure digits or +stereo-bond parities consistent across a macrocyclisation. + +Stereochemistry +--------------- +All natural residues are L (CIP *S* at Cα, except cysteine which is *R* by the +standard CIP quirk, and glycine which is achiral). In this module's fragment +ordering — ``[*:1]`` first, then N, Cα, then the carbonyl — the verified tags are: + + in-chain L residue ``[*:1]N[C@@H](R)C(=O)[*:2]`` -> CIP S + S5 anchor (α-Me, (S)) ``[*:1]N[C@@](C)(CCCC=C)C(=O)[*:2]`` -> CIP S + R8 anchor (α-Me, (R)) ``[*:1]N[C@](C)(CCCCCCC=C)C(=O)[*:2]`` -> CIP R + +A ``@``/``@@`` symbol denotes a parity over the neighbour list *in written +order*, so the same symbol means different absolute configurations depending on +where the attachment dummies sit. Reasoning about these tags by analogy is +therefore unreliable, and every tag in this module is instead machine-checked by +:func:`verify_monomer_stereochemistry` against +``rdkit.Chem.rdCIPLabeler.AssignCIPLabels``. That check is exercised by the +tests; it caught all four anchors being inverted during development, and it +would catch a future edit that reintroduces the same slip. + +Anchor naming follows the StaPep convention used throughout this branch: + + ========== ============================ ========= ======================= + token monomer Cα config olefin tether + ========== ============================ ========= ======================= + ``S5`` (S)-α-Me-α-(4-pentenyl)Gly S ``CCCC=C`` (4 C + ene) + ``R8`` (R)-α-Me-α-(7-octenyl)Gly R ``CCCCCCC=C`` (7 C + ene) + ``R5`` (R)-α-Me-α-(4-pentenyl)Gly R ``CCCC=C`` + ``S8`` (S)-α-Me-α-(7-octenyl)Gly S ``CCCCCCC=C`` + ========== ============================ ========= ======================= + +``S5``/``R8`` are the two anchors the enabled catalog uses; ``R5``/``S8`` exist +here only so the default-off ``R5-S8/i,i+7`` extension can be built when +explicitly requested. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + +#: Dummy-atom map number marking the amide-nitrogen (N-terminal) attachment. +ATTACH_N: Final[int] = 1 +#: Dummy-atom map number marking the carbonyl-carbon (C-terminal) attachment. +ATTACH_C: Final[int] = 2 + + +class UnknownMonomerError(KeyError): + """Raised when a token has no fragment in this library.""" + + +@dataclass(frozen=True) +class Monomer: + """One residue fragment. + + Attributes: + token: canonical monomer token (``"A"``, ``"S5"``, ``"AIB"``, ...). + name: human-readable name. + smiles: backbone-open fragment carrying ``[*:1]`` and ``[*:2]``. + is_anchor: True for the olefin-bearing staple anchors. + olefin_carbons: number of carbons in the tether *including* both alkene + carbons. ``S5`` has 5, ``R8`` has 8. Zero for non-anchors. + cip_code: expected CIP label at Cα, or ``None`` when achiral. + three_letter: conventional three-letter code, for reporting. + """ + + token: str + name: str + smiles: str + is_anchor: bool = False + olefin_carbons: int = 0 + cip_code: str | None = None + three_letter: str = "" + + +def _residue( + token: str, + name: str, + side_chain: str, + three_letter: str, + cip_code: str | None = "S", +) -> Monomer: + """Build a standard L-α-amino-acid fragment from its side chain. + + ``side_chain`` is spliced into ``[*:1]N[C@@H]()C(=O)[*:2]``. + Glycine passes ``side_chain=""`` and uses the achiral ``CH2`` form. + """ + if side_chain: + smiles = f"[*:{ATTACH_N}]N[C@@H]({side_chain})C(=O)[*:{ATTACH_C}]" + else: + smiles = f"[*:{ATTACH_N}]NCC(=O)[*:{ATTACH_C}]" + return Monomer( + token=token, + name=name, + smiles=smiles, + cip_code=cip_code, + three_letter=three_letter, + ) + + +def _anchor( + token: str, + name: str, + tether: str, + olefin_carbons: int, + chirality: str, + cip_code: str, +) -> Monomer: + """Build an α-methyl α-alkenyl glycine anchor fragment. + + ``chirality`` is the raw ``@``/``@@`` tag; the resulting CIP label is + asserted against ``cip_code`` by :func:`verify_monomer_stereochemistry` + rather than trusted. + """ + return Monomer( + token=token, + name=name, + smiles=( + f"[*:{ATTACH_N}]N[C{chirality}](C)({tether})C(=O)[*:{ATTACH_C}]" + ), + is_anchor=True, + olefin_carbons=olefin_carbons, + cip_code=cip_code, + three_letter=token, + ) + + +#: Olefin tether for the "5" series: 4 saturated carbons then a terminal alkene. +TETHER_PENTENYL: Final[str] = "CCCC=C" +#: Olefin tether for the "8" series: 7 saturated carbons then a terminal alkene. +TETHER_OCTENYL: Final[str] = "CCCCCCC=C" + +#: The 20 proteinogenic residues, L configuration. +#: +#: Cysteine is CIP *R* despite being L — the sulfur outranks the carboxyl, so +#: the same spatial arrangement earns the opposite descriptor. Encoding it as +#: ``"S"`` would make :func:`verify_monomer_stereochemistry` fail correctly, so +#: the exception is recorded here rather than special-cased downstream. +NATURAL_MONOMERS: Final[dict[str, Monomer]] = { + "A": _residue("A", "L-alanine", "C", "Ala"), + "R": _residue("R", "L-arginine", "CCCNC(N)=N", "Arg"), + "N": _residue("N", "L-asparagine", "CC(N)=O", "Asn"), + "D": _residue("D", "L-aspartic acid", "CC(=O)O", "Asp"), + "C": _residue("C", "L-cysteine", "CS", "Cys", cip_code="R"), + "Q": _residue("Q", "L-glutamine", "CCC(N)=O", "Gln"), + "E": _residue("E", "L-glutamic acid", "CCC(=O)O", "Glu"), + "G": _residue("G", "glycine", "", "Gly", cip_code=None), + "H": _residue("H", "L-histidine", "Cc1c[nH]cn1", "His"), + "I": _residue("I", "L-isoleucine", "[C@@H](C)CC", "Ile"), + "L": _residue("L", "L-leucine", "CC(C)C", "Leu"), + "K": _residue("K", "L-lysine", "CCCCN", "Lys"), + "M": _residue("M", "L-methionine", "CCSC", "Met"), + "F": _residue("F", "L-phenylalanine", "Cc1ccccc1", "Phe"), + "P": _residue("P", "L-proline", "", "Pro", cip_code="S"), + "S": _residue("S", "L-serine", "CO", "Ser"), + "T": _residue("T", "L-threonine", "[C@H](C)O", "Thr"), + "W": _residue("W", "L-tryptophan", "Cc1c[nH]c2ccccc12", "Trp"), + "Y": _residue("Y", "L-tyrosine", "Cc1ccc(O)cc1", "Tyr"), + "V": _residue("V", "L-valine", "C(C)C", "Val"), +} + +# Proline is a secondary amine whose side chain closes back onto the backbone +# nitrogen, so the generic template cannot express it. Written out explicitly. +NATURAL_MONOMERS["P"] = Monomer( + token="P", + name="L-proline", + smiles=f"[*:{ATTACH_N}]N1[C@@H](CCC1)C(=O)[*:{ATTACH_C}]", + cip_code="S", + three_letter="Pro", +) + +#: The hydrocarbon staple anchors. +#: +#: The ``@``/``@@`` tags below are in the *fragment's* atom ordering +#: (``[*:1]N[C?](C)(tether)C(=O)[*:2]``), where ``@@`` is (S) and ``@`` is (R). +#: That is the opposite of the bare ``N[C?](C)(tether)C(=O)NC`` ordering, which +#: is exactly the kind of slip :func:`verify_monomer_stereochemistry` exists to +#: catch — it did catch it during development, on all four anchors at once. +ANCHOR_MONOMERS: Final[dict[str, Monomer]] = { + "S5": _anchor( + "S5", + "(S)-2-(4-pentenyl)alanine", + TETHER_PENTENYL, + olefin_carbons=5, + chirality="@@", + cip_code="S", + ), + "R8": _anchor( + "R8", + "(R)-2-(7-octenyl)alanine", + TETHER_OCTENYL, + olefin_carbons=8, + chirality="@", + cip_code="R", + ), + "R5": _anchor( + "R5", + "(R)-2-(4-pentenyl)alanine", + TETHER_PENTENYL, + olefin_carbons=5, + chirality="@", + cip_code="R", + ), + "S8": _anchor( + "S8", + "(S)-2-(7-octenyl)alanine", + TETHER_OCTENYL, + olefin_carbons=8, + chirality="@@", + cip_code="S", + ), +} + +#: Non-anchor non-natural monomers the branch tokenizer already accepts. +OTHER_MONOMERS: Final[dict[str, Monomer]] = { + "AIB": Monomer( + token="AIB", + name="2-aminoisobutyric acid", + smiles=f"[*:{ATTACH_N}]NC(C)(C)C(=O)[*:{ATTACH_C}]", + cip_code=None, + three_letter="Aib", + ), + "NLE": Monomer( + token="NLE", + name="L-norleucine", + smiles=f"[*:{ATTACH_N}]N[C@@H](CCCC)C(=O)[*:{ATTACH_C}]", + cip_code="S", + three_letter="Nle", + ), +} + +#: Every monomer this library can place. +MONOMER_LIBRARY: Final[dict[str, Monomer]] = { + **NATURAL_MONOMERS, + **ANCHOR_MONOMERS, + **OTHER_MONOMERS, +} + +# -- terminal capping -------------------------------------------------------- + +#: Free N-terminus: the attachment dummy simply becomes a hydrogen. +N_TERM_FREE: Final[str] = "free_amine" +#: Acetylated N-terminus (``Ac-``). +N_TERM_ACETYL: Final[str] = "acetyl" +#: Free C-terminus carboxylic acid. +C_TERM_ACID: Final[str] = "free_acid" +#: C-terminal primary amide (``-NH2``). +C_TERM_AMIDE: Final[str] = "amide" + +#: Cap fragments, each with a single attachment dummy. +N_TERMINAL_CAPS: Final[dict[str, str | None]] = { + N_TERM_FREE: None, + N_TERM_ACETYL: f"CC(=O)[*:{ATTACH_N}]", +} +C_TERMINAL_CAPS: Final[dict[str, str | None]] = { + C_TERM_ACID: f"O[*:{ATTACH_C}]", + C_TERM_AMIDE: f"N[*:{ATTACH_C}]", +} + + +def get_monomer(token: str) -> Monomer: + """Look up one monomer by token. + + Raises: + UnknownMonomerError: if the token is not in the library. Guessing a + fragment would silently emit the wrong molecule. + """ + key = token.upper() + monomer = MONOMER_LIBRARY.get(key) + if monomer is None: + raise UnknownMonomerError( + f"no SMILES fragment for monomer {token!r}; known tokens: " + f"{sorted(MONOMER_LIBRARY)}" + ) + return monomer + + +def is_anchor_monomer(token: str) -> bool: + """True when ``token`` is one of the olefin-bearing staple anchors.""" + return token.upper() in ANCHOR_MONOMERS + + +def anchor_olefin_carbons(token: str) -> int: + """Tether carbon count for an anchor token, including both alkene carbons.""" + return get_monomer(token).olefin_carbons + + +def staple_carbon_count(i_token: str, j_token: str) -> int: + """Carbons in the closed staple bridge after RCM. + + Ring-closing metathesis joins the two terminal alkenes and expels ethene, so + two carbons — one from each partner's terminal ``=CH2`` — leave the molecule: + + S5 + S5 -> 5 + 5 - 2 = 8 bridge carbons + R8 + S5 -> 8 + 5 - 2 = 11 bridge carbons + + Those counts are what make the observed macrocycle sizes add up (13 backbone + atoms + 8 = 21 for i,i+4; 22 + 11 = 33 for i,i+7), which is the arithmetic + used to identify the reference topologies in the StaPep table. + """ + return anchor_olefin_carbons(i_token) + anchor_olefin_carbons(j_token) - 2 + + +def expected_macrocycle_size(spacing: int, i_token: str, j_token: str) -> int: + """Ring size of the stapled macrocycle. + + The ring runs Cα(i) -> backbone -> Cα(j) -> staple bridge -> back to Cα(i). + Between Cα(i) and Cα(j) there are ``3 * spacing - 1`` intervening backbone + atoms (each residue contributes N, Cα, C'), so including both Cα atoms the + backbone arc is ``3 * spacing + 1`` atoms. Adding the bridge gives the ring. + """ + return 3 * spacing + 1 + staple_carbon_count(i_token, j_token) + + +def verify_monomer_stereochemistry() -> dict[str, dict[str, str | None]]: + """Round-trip every chiral monomer through RDKit's CIP labeller. + + Each fragment is capped into a small Ac-Xaa-NHMe model peptide (the dummy + atoms cannot themselves be sanitised into a stereo-perceivable molecule) and + the observed CIP label at Cα is compared with the declared + :attr:`Monomer.cip_code`. + + Returns: + ``{token: {"expected": ..., "observed": ..., "status": ...}}``. Status is + ``"ok"``, ``"mismatch"``, ``"achiral"`` or ``"unparseable"``. + + This exists so the hand-written ``@``/``@@`` tags are machine-checked. A + silently inverted anchor would otherwise produce a molecule that parses, + sanitises and looks entirely reasonable while being the wrong enantiomer. + """ + from rdkit import Chem + from rdkit.Chem import rdCIPLabeler + + results: dict[str, dict[str, str | None]] = {} + for token, monomer in MONOMER_LIBRARY.items(): + model = monomer.smiles.replace(f"[*:{ATTACH_N}]", "CC(=O)").replace( + f"[*:{ATTACH_C}]", "NC" + ) + mol = Chem.MolFromSmiles(model) + if mol is None: + results[token] = { + "expected": monomer.cip_code, + "observed": None, + "status": "unparseable", + "model_smiles": model, + } + continue + + Chem.AssignStereochemistry(mol, cleanIt=True, force=True) + rdCIPLabeler.AssignCIPLabels(mol) + labels = [ + atom.GetProp("_CIPCode") + for atom in mol.GetAtoms() + if atom.HasProp("_CIPCode") + ] + + if monomer.cip_code is None: + status = "achiral" if not labels else "mismatch" + observed = labels[0] if labels else None + else: + # Ile and Thr carry a second stereocentre in the side chain; the Cα + # label is the one being checked, so require it to be present rather + # than requiring exactly one label overall. + observed = labels[0] if labels else None + status = "ok" if monomer.cip_code in labels else "mismatch" + + results[token] = { + "expected": monomer.cip_code, + "observed": observed, + "all_labels": ",".join(labels) if labels else None, + "status": status, + "model_smiles": model, + } + return results + + +def describe_library() -> dict[str, object]: + """Summary of the library, for reports and audits.""" + return { + "n_monomers": len(MONOMER_LIBRARY), + "n_natural": len(NATURAL_MONOMERS), + "n_anchors": len(ANCHOR_MONOMERS), + "n_other": len(OTHER_MONOMERS), + "anchor_tokens": sorted(ANCHOR_MONOMERS), + "anchor_olefin_carbons": { + token: monomer.olefin_carbons + for token, monomer in sorted(ANCHOR_MONOMERS.items()) + }, + "n_terminal_caps": sorted(N_TERMINAL_CAPS), + "c_terminal_caps": sorted(C_TERMINAL_CAPS), + "staple_bridge_carbons": { + "S5-S5": staple_carbon_count("S5", "S5"), + "R8-S5": staple_carbon_count("R8", "S5"), + "R5-S8": staple_carbon_count("R5", "S8"), + }, + "expected_macrocycle_size": { + "S5-S5/i,i+4": expected_macrocycle_size(4, "S5", "S5"), + "R8-S5/i,i+7": expected_macrocycle_size(7, "R8", "S5"), + }, + } diff --git a/staplebridge/hydrocarbon/plan_control.py b/staplebridge/hydrocarbon/plan_control.py new file mode 100644 index 0000000000000000000000000000000000000000..b9e494aaca37d52087449e7a85198c0f52c97a15 --- /dev/null +++ b/staplebridge/hydrocarbon/plan_control.py @@ -0,0 +1,962 @@ +"""Hydrocarbon-only learnable control over empirical staple plans. + +The empirical distribution remains the reference measure. When enabled, a +small MLP tilts that measure before the existing plan-conditioned action policy +runs. Nothing in this module is imported by the lactam training path. +""" + +from __future__ import annotations + +import math +import random +from collections import OrderedDict +from dataclasses import dataclass +from typing import Any + +import torch +from torch import nn + +from staplebridge.chemistry.edit_distance import weighted_edit_distance +from staplebridge.chemistry.edit_metrics import edit_diagnostics +from staplebridge.chemistry.state import StapleState +from staplebridge.hydrocarbon.actions import FailureReason, validate_hydrocarbon_staple +from staplebridge.hydrocarbon.catalog import hydrocarbon_catalog_from_config +from staplebridge.hydrocarbon.curriculum import ( + HydrocarbonStaplePlan, + build_hydrocarbon_demonstration_path, +) +from staplebridge.hydrocarbon.plan_reference import ( + OFF_PLAN_ANCHOR, + OFF_PLAN_SUBSTITUTION, + OFF_PLAN_TOPOLOGY, + ON_PLAN_ANCHOR_ASSIGN, + ON_PLAN_BLOCK_ASSIGN, + ON_PLAN_FIRST_ANCHOR, + ON_PLAN_LABELS, + ON_PLAN_SECOND_ANCHOR, + ON_PLAN_TOPOLOGY, + PlanAwareReferenceSampler, + PlanAwareTrajectory, + PlanProgress, +) + + +@dataclass +class HydrocarbonPlanControlConfig: + """Configuration for the optional plan-level control head.""" + + enabled: bool = True + plans_per_lead: int = 4 + hidden_dim: int = 32 + loss_weight: float = 1.0 + target_temperature: float = 0.10 + # False preserves the historical sampled-plan Penetrance CE exactly. + # True opts into the lead-local coarse-grained SB objective. + exact_sb_objective: bool = False + exact_sb_beta: float = 1.0 + # Opt-in stable representation repair. The representation is independent + # of which plan-level objective (Exact-SB or Legacy) supervises the head. + # False preserves old V1 checkpoints bit-for-bit. + plan_encoder_v2: bool = False + validation_all_plans: bool = True + + @classmethod + def from_config(cls, root_config: dict[str, Any] | None) -> "HydrocarbonPlanControlConfig": + section = dict(((root_config or {}).get("hydrocarbon") or {}).get("plan_control") or {}) + cfg = cls() + for key, value in section.items(): + if not hasattr(cfg, key): + continue + current = getattr(cfg, key) + if isinstance(current, bool): + setattr(cfg, key, bool(value)) + elif isinstance(current, int): + setattr(cfg, key, int(value)) + else: + setattr(cfg, key, float(value)) + if cfg.plans_per_lead < 2: + raise ValueError("hydrocarbon.plan_control.plans_per_lead must be >= 2") + if ( + cfg.hidden_dim < 1 + or cfg.loss_weight < 0.0 + or cfg.target_temperature <= 0.0 + or cfg.exact_sb_beta <= 0.0 + ): + raise ValueError("invalid hydrocarbon.plan_control numeric setting") + return cfg + + +def property_free_candidate_constraints( + *, + state: StapleState, + initial_state: StapleState, + lead: Any, + plan: HydrocarbonStaplePlan, + catalog: list[Any], + catalog_index: dict[str, Any], + geometry: Any, + config: dict[str, Any], +) -> dict[str, Any]: + """One property-free hard predicate shared by train/valid/test. + + No model or PeptiVerse value participates. The predicate covers exact + plan completion plus chemistry, topology, geometry, edit budget, sequence + identity, and protected-position constraints. + """ + block = catalog_index.get(state.block_id) if state.block_id else None + anchor_pair = None if state.anchor_pair is None else tuple(state.anchor_pair) + verdict = validate_hydrocarbon_staple( + state.sequence_tokens, anchor_pair, block, catalog + ) + chemistry_valid = verdict is FailureReason.OK + stapled = state.topology == "stapled" + peptide_ca = (lead.target_context or {}).get("peptide_ca") + cgeom = float( + geometry.cgeom( + state.sequence_tokens, anchor_pair, block, peptide_ca=peptide_ca + ) + ) + eps_geom = float( + config.get( + "eps_geom", (config.get("training") or {}).get("eps_geom", 2.5) + ) + ) + geometry_feasible = bool( + geometry.ctype( + state.sequence_tokens, anchor_pair, block, peptide_ca=peptide_ca + ) + and cgeom <= eps_geom + ) + edits = edit_diagnostics( + state, + initial_state, + protected_positions=list(lead.protected_positions or []), + catalog=catalog_index, + ) + edit_cfg = dict(config.get("edit_constraints") or {}) + within_edit_budget = bool( + float(edits["edit_distance"]) + <= float(edit_cfg.get("max_edit_budget", float("inf"))) + ) + sequence_identity_ok = bool( + float(edits["sequence_identity"]) + >= float(edit_cfg.get("min_sequence_identity", 0.0)) + ) + protected_positions_ok = bool( + edit_cfg.get("allow_protected_edits", False) + or not edits["protected_edit_violation"] + ) + exact_completion = bool(completes_committed_plan(state, plan)) + strict_feasible = bool( + chemistry_valid + and stapled + and geometry_feasible + and exact_completion + and within_edit_budget + and sequence_identity_ok + and protected_positions_ok + ) + ordered_failures = [ + ("chemistry_invalid", chemistry_valid), + ("not_stapled", stapled), + ("geometry_infeasible", geometry_feasible), + ("inexact_committed_plan_completion", exact_completion), + ("edit_budget_exceeded", within_edit_budget), + ("sequence_identity_below_minimum", sequence_identity_ok), + ("protected_position_violation", protected_positions_ok), + ] + failure_reasons = [name for name, passed in ordered_failures if not passed] + return { + "chemistry_valid": chemistry_valid, + "stapled": stapled, + "topology_status": verdict.value, + "geometry_cgeom": cgeom, + "geometry_eps": eps_geom, + "geometry_feasible": geometry_feasible, + "exact_committed_plan_completion": exact_completion, + "edit_distance": float(edits["edit_distance"]), + "weighted_edit_distance": float( + weighted_edit_distance(state, initial_state) + ), + "sequence_identity": float(edits["sequence_identity"]), + "protected_edit_violation": bool(edits["protected_edit_violation"]), + "within_edit_budget": within_edit_budget, + "sequence_identity_constraint_satisfied": sequence_identity_ok, + "protected_position_constraint_satisfied": protected_positions_ok, + "edit_constraints_satisfied": bool( + within_edit_budget and sequence_identity_ok and protected_positions_ok + ), + "strict_feasible": strict_feasible, + "hard_plan_failure_reasons": failure_reasons, + "hard_plan_primary_failure_reason": ( + failure_reasons[0] if failure_reasons else None + ), + } + + +def property_free_hard_plan_support( + *, + initial_state: StapleState, + lead: Any, + plans: list[HydrocarbonStaplePlan], + catalog: list[Any], + catalog_index: dict[str, Any], + geometry: Any, + config: dict[str, Any], +) -> tuple[list[int], list[dict[str, Any]]]: + """Return strict-support indices and auditable per-plan verdicts.""" + support: list[int] = [] + verdicts: list[dict[str, Any]] = [] + for index, plan in enumerate(plans): + try: + terminal = build_hydrocarbon_demonstration_path( + initial_state, plan, catalog + )[-1] + verdict = property_free_candidate_constraints( + state=terminal, + initial_state=initial_state, + lead=lead, + plan=plan, + catalog=catalog, + catalog_index=catalog_index, + geometry=geometry, + config=config, + ) + except (KeyError, ValueError) as exc: + verdict = { + "chemistry_valid": False, + "stapled": False, + "geometry_feasible": False, + "exact_committed_plan_completion": False, + "edit_constraints_satisfied": False, + "strict_feasible": False, + "hard_plan_failure_reasons": ["plan_materialization_failed"], + "hard_plan_primary_failure_reason": "plan_materialization_failed", + "hard_plan_error": f"{type(exc).__name__}: {exc}", + } + verdicts.append(verdict) + if verdict["strict_feasible"]: + support.append(index) + return support, verdicts + + +def mask_and_renormalize_plan_log_probabilities( + log_probabilities: torch.Tensor, + support_indices: list[int], +) -> torch.Tensor: + """Hard-mask q_theta and renormalize on property-free strict support.""" + if log_probabilities.ndim != 1: + raise ValueError("plan log probabilities must be one-dimensional") + if not support_indices: + raise ValueError("cannot renormalize q_theta on empty strict support") + indices = torch.tensor( + support_indices, dtype=torch.long, device=log_probabilities.device + ) + if int(indices.min().item()) < 0 or int(indices.max().item()) >= int( + log_probabilities.numel() + ): + raise IndexError("strict-support index outside plan distribution") + if int(torch.unique(indices).numel()) != len(support_indices): + raise ValueError("strict-support indices must be unique") + masked = torch.full_like(log_probabilities, -torch.inf) + masked[indices] = log_probabilities[indices] + return masked - torch.logsumexp(masked, dim=0) + + +class HydrocarbonPlanHead(nn.Module): + """Two-layer MLP that tilts empirical plan probabilities.""" + + PLAN_FEATURE_DIM = 8 + + def __init__(self, lead_embedding_dim: int, hidden_dim: int = 32) -> None: + super().__init__() + self.lead_embedding_dim = int(lead_embedding_dim) + self.net = nn.Sequential( + nn.Linear(self.lead_embedding_dim + self.PLAN_FEATURE_DIM, hidden_dim), + nn.SiLU(), + nn.Linear(hidden_dim, 1), + ) + # Exact empirical q_ref at initialization. Any later disagreement is + # learned rather than an arbitrary random tilt. + nn.init.zeros_(self.net[-1].weight) + nn.init.zeros_(self.net[-1].bias) + + @staticmethod + def plan_features( + plans: list[HydrocarbonStaplePlan], + lead_length: int, + device: torch.device, + ) -> torch.Tensor: + denom = float(max(lead_length - 1, 1)) + rows: list[list[float]] = [] + for plan in plans: + i, j = plan.anchor_pair + is_s5 = float(plan.ordered_pair == "S5-S5") + is_r8 = float(plan.ordered_pair == "R8-S5") + rows.append( + [ + is_s5, + is_r8, + i / denom, + j / denom, + (0.5 * (i + j)) / denom, + plan.spacing / denom, + plan.n_edits / 2.0, + lead_length / 32.0, + ] + ) + return torch.tensor(rows, dtype=torch.float32, device=device) + + def forward( + self, + lead_embedding: torch.Tensor, + plans: list[HydrocarbonStaplePlan], + lead_length: int, + ) -> torch.Tensor: + if not plans: + return torch.empty(0, device=lead_embedding.device) + if lead_embedding.ndim == 1: + lead_embedding = lead_embedding.unsqueeze(0) + features = self.plan_features(plans, lead_length, lead_embedding.device) + expanded = lead_embedding.expand(len(plans), -1) + return self.net(torch.cat([expanded, features], dim=-1)).squeeze(-1) + + + +class FrozenESM2PlanFeatureEncoder: + """Stable local features from the frozen ESM2 prior's vector cache. + + The existing peptide prior stores the complete 20-AA masked-token + log-probability vector for ``(sequence, position)``. Reusing those vectors + gives a deterministic contextual feature at every anchor/local position, + avoids registering the 650M model in the head, and makes cache-hit probes + require no ESM2 forward at all. Missing vectors still go through the same + frozen strict ESM2 provider and are cached there. + """ + + VECTOR_DIM = 20 + RAW_FEATURE_COUNT = 11 + _AA_ORDER = "ACDEFGHIKLMNPQRSTVWY" + + def __init__(self, esm2_prior: Any, cache_size: int = 65536) -> None: + if esm2_prior is None: + raise ValueError("plan_encoder_v2 requires the real frozen ESM2 prior") + get_vectors = getattr(esm2_prior, "_get_vectors", None) + if not callable(get_vectors): + raise TypeError("plan_encoder_v2 requires ESM2DeltaPeptidePrior") + self.esm2_prior = esm2_prior + self.hidden_dim = self.VECTOR_DIM + self.cache_size = max(int(cache_size), 1) + self._cache: "OrderedDict[str, torch.Tensor]" = OrderedDict() + self.cache_hits = 0 + self.cache_misses = 0 + self.forward_batches = 0 + + def _canonicalize(self, sequence_tokens: list[str]) -> str: + canonicalize = getattr(self.esm2_prior, "_canonicalize", None) + if not callable(canonicalize): + raise RuntimeError("ESM2 prior lacks canonical-surrogate conversion") + sequence, _ = canonicalize(list(sequence_tokens)) + if len(sequence) != len(sequence_tokens): + raise RuntimeError("ESM2 surrogate changed peptide length") + return sequence + + def _put_cache(self, sequence: str, tokens: torch.Tensor) -> None: + self._cache[sequence] = tokens.detach().to(device="cpu", dtype=torch.float16) + self._cache.move_to_end(sequence) + while len(self._cache) > self.cache_size: + self._cache.popitem(last=False) + + def _encode_missing(self, sequences: list[str]) -> None: + if not sequences: + return + requests = [ + (sequence, position) + for sequence in sequences + for position in range(len(sequence)) + ] + vectors = self.esm2_prior._get_vectors(requests) + for sequence in sequences: + rows = [vectors.get((sequence, position)) for position in range(len(sequence))] + if any(row is None or len(row) != self.VECTOR_DIM for row in rows): + raise RuntimeError("ESM2 local-vector cache returned incomplete features") + self._put_cache(sequence, torch.tensor(rows, dtype=torch.float32)) + + def _fetch_many( + self, sequences: list[str] + ) -> dict[str, torch.Tensor]: + unique = list(dict.fromkeys(sequences)) + missing: list[str] = [] + for sequence in unique: + if sequence in self._cache: + self.cache_hits += 1 + self._cache.move_to_end(sequence) + else: + self.cache_misses += 1 + missing.append(sequence) + self._encode_missing(missing) + return {sequence: self._cache[sequence] for sequence in unique} + + @classmethod + def _one_hot(cls, token: str) -> torch.Tensor: + row = torch.zeros(cls.VECTOR_DIM, dtype=torch.float16) + token = str(token).upper() + if token in cls._AA_ORDER: + row[cls._AA_ORDER.index(token)] = 1.0 + return row + + @staticmethod + def _local_context(tokens: torch.Tensor, position: int, radius: int) -> torch.Tensor: + indices = [ + index + for index in range( + max(0, position - radius), min(tokens.shape[0], position + radius + 1) + ) + if index != position + ] + if not indices: + return tokens[position] + return tokens[indices].mean(dim=0) + + def context_features( + self, + sequence_tokens: list[str], + plans: list[HydrocarbonStaplePlan], + device: torch.device, + ) -> torch.Tensor: + if not plans: + return torch.empty( + (0, self.RAW_FEATURE_COUNT, self.hidden_dim), + dtype=torch.float32, + device=device, + ) + base_sequence = self._canonicalize(sequence_tokens) + cached = self._fetch_many([base_sequence]) + base_tokens = cached[base_sequence] + base_pooled = base_tokens.mean(dim=0) + rows: list[torch.Tensor] = [] + for plan in plans: + i, j = (int(plan.anchor_pair[0]), int(plan.anchor_pair[1])) + if i < 0 or j >= base_tokens.shape[0]: + raise ValueError("plan anchor outside ESM2 token representation") + h_i = base_tokens[i] + h_j = base_tokens[j] + rows.append( + torch.stack( + [ + base_pooled, + h_i, + h_j, + h_i - h_j, + h_i * h_j, + self._local_context(base_tokens, i, 2), + self._local_context(base_tokens, j, 2), + self._local_context(base_tokens, i, 3), + self._local_context(base_tokens, j, 3), + self._one_hot(sequence_tokens[i]), + self._one_hot(sequence_tokens[j]), + ], + dim=0, + ) + ) + return torch.stack(rows, dim=0).to(device=device, dtype=torch.float32) + + def prefetch(self, sequences: list[list[str]]) -> None: + """Warm stable base-sequence representations in efficient batches.""" + canonical = [self._canonicalize(tokens) for tokens in sequences] + self._fetch_many(canonical) + + def diagnostics(self) -> dict[str, int]: + return { + "cache_entries": len(self._cache), + "cache_hits": self.cache_hits, + "cache_misses": self.cache_misses, + "forward_batches": self.forward_batches, + } + + +class PlanGeometryFeatureEncoder: + """Deterministic plan geometry available identically at train/inference.""" + + FEATURE_DIM = 8 + + def __init__(self, ca_windows: dict[str, tuple[float, float]]) -> None: + self.ca_windows = { + str(block_id): (float(window[0]), float(window[1])) + for block_id, window in ca_windows.items() + } + + def features( + self, + plans: list[HydrocarbonStaplePlan], + peptide_ca: list[tuple[float, float, float]] | list[list[float]], + device: torch.device, + ) -> torch.Tensor: + if peptide_ca is None: + raise ValueError("plan_encoder_v2 requires peptide_ca") + rows: list[list[float]] = [] + for plan in plans: + i, j = (int(plan.anchor_pair[0]), int(plan.anchor_pair[1])) + if i < 0 or j < 0 or i >= len(peptide_ca) or j >= len(peptide_ca): + raise ValueError("plan anchor outside peptide_ca") + if plan.block_id not in self.ca_windows: + raise ValueError(f"missing geometry window for {plan.block_id}") + xi, yi, zi = (float(value) for value in peptide_ca[i]) + xj, yj, zj = (float(value) for value in peptide_ca[j]) + distance = math.sqrt( + (xi - xj) ** 2 + (yi - yj) ** 2 + (zi - zj) ** 2 + ) + low, high = self.ca_windows[plan.block_id] + below = max(low - distance, 0.0) + above = max(distance - high, 0.0) + center = 0.5 * (low + high) + rows.append( + [ + 1.0, + distance / 20.0, + low / 20.0, + high / 20.0, + below / 10.0, + above / 10.0, + float(below == 0.0 and above == 0.0), + (distance - center) / 10.0, + ] + ) + return torch.tensor(rows, dtype=torch.float32, device=device) + + +class HydrocarbonPlanHeadV2(nn.Module): + """Exact-SB plan head with stable frozen sequence + geometry features. + + ``lead_embedding`` remains in the call signature solely for V1/V2 API and + checkpoint orchestration compatibility. V2 deliberately ignores it: that + embedding belongs to the fine policy encoder and moves under the path loss. + """ + + PROJECTION_DIM = 32 + + def __init__( + self, + lead_embedding_dim: int, + hidden_dim: int, + feature_encoder: FrozenESM2PlanFeatureEncoder, + geometry_encoder: PlanGeometryFeatureEncoder, + ) -> None: + super().__init__() + self.lead_embedding_dim = int(lead_embedding_dim) + self.feature_encoder = feature_encoder + self.geometry_encoder = geometry_encoder + self.feature_projection = nn.Linear( + feature_encoder.hidden_dim, self.PROJECTION_DIM + ) + input_dim = ( + HydrocarbonPlanHead.PLAN_FEATURE_DIM + + feature_encoder.RAW_FEATURE_COUNT * self.PROJECTION_DIM + + geometry_encoder.FEATURE_DIM + ) + self.net = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.SiLU(), + nn.Linear(hidden_dim, 1), + ) + nn.init.zeros_(self.net[-1].weight) + nn.init.zeros_(self.net[-1].bias) + + def forward( + self, + lead_embedding: torch.Tensor, + plans: list[HydrocarbonStaplePlan], + lead_length: int, + sequence_tokens: list[str], + peptide_ca: list[tuple[float, float, float]] | list[list[float]], + ) -> torch.Tensor: + if not plans: + return torch.empty(0, device=lead_embedding.device) + # The fine-policy embedding is intentionally excluded from V2. + device = lead_embedding.device + old_features = HydrocarbonPlanHead.plan_features( + plans, lead_length, device + ) + context = self.feature_encoder.context_features( + sequence_tokens, plans, device + ) + projected = self.feature_projection(context).flatten(start_dim=1) + geometry = self.geometry_encoder.features(plans, peptide_ca, device) + return self.net( + torch.cat([old_features, projected, geometry], dim=-1) + ).squeeze(-1) + +def build_hydrocarbon_plan_head( + root_config: dict[str, Any] | None, + lead_embedding_dim: int, + device: torch.device | str, + *, + esm2_prior: Any = None, +) -> HydrocarbonPlanHead | HydrocarbonPlanHeadV2 | None: + """Construct the selected head; disabled and V1 paths retain zero V2 cost.""" + cfg = HydrocarbonPlanControlConfig.from_config(root_config) + if not cfg.enabled: + return None + if not cfg.plan_encoder_v2: + return HydrocarbonPlanHead(lead_embedding_dim, cfg.hidden_dim).to(device) + feature_encoder = FrozenESM2PlanFeatureEncoder(esm2_prior) + catalog = hydrocarbon_catalog_from_config( + dict(((root_config or {}).get("hydrocarbon") or {})) + ) + geometry_encoder = PlanGeometryFeatureEncoder( + {block.block_id: tuple(block.ca_window) for block in catalog} + ) + return HydrocarbonPlanHeadV2( + lead_embedding_dim, cfg.hidden_dim, feature_encoder, geometry_encoder + ).to(device) + + +def empirical_log_probabilities(weights: list[float], device: torch.device) -> torch.Tensor: + values = torch.tensor(weights, dtype=torch.float32, device=device) + if values.numel() == 0 or float(values.sum().item()) <= 0.0: + raise ValueError("empirical plan weights must have positive mass") + values = values / values.sum() + return torch.log(values.clamp_min(torch.finfo(values.dtype).tiny)) + + +def controlled_plan_log_probabilities( + head: HydrocarbonPlanHead | HydrocarbonPlanHeadV2, + lead_embedding: torch.Tensor, + plans: list[HydrocarbonStaplePlan], + empirical_weights: list[float], + lead_length: int, + *, + sequence_tokens: list[str] | None = None, + peptide_ca: list[tuple[float, float, float]] | list[list[float]] | None = None, +) -> torch.Tensor: + reference = empirical_log_probabilities(empirical_weights, lead_embedding.device) + if isinstance(head, HydrocarbonPlanHeadV2): + if sequence_tokens is None: + raise ValueError("plan_encoder_v2 requires sequence_tokens") + if peptide_ca is None: + raise ValueError("plan_encoder_v2 requires peptide_ca") + scores = head( + lead_embedding.detach(), plans, lead_length, list(sequence_tokens), peptide_ca + ) + else: + # This is intentionally the historical expression, unchanged. + scores = head(lead_embedding.detach(), plans, lead_length) + return torch.log_softmax(reference + scores, dim=0) + + +def best_committed_plan_trajectory( + decoded: list[tuple[StapleState, float, int]], + plan: HydrocarbonStaplePlan, +) -> tuple[StapleState, float, int] | None: + """Return the highest-probability trajectory completing the selected plan.""" + for item in decoded: + state = item[0] + if completes_committed_plan(state, plan): + return item + return None + + +def completes_committed_plan( + state: StapleState, plan: HydrocarbonStaplePlan +) -> bool: + """The single exact-completion predicate shared by train and inference.""" + return bool( + state.topology == "stapled" + and state.anchor_pair is not None + and tuple(state.anchor_pair) == tuple(plan.anchor_pair) + and state.block_id == plan.block_id + ) + + +def hierarchical_plan_ranking_enabled(root_config: dict[str, Any] | None) -> bool: + """Return the hydrocarbon-only plan-first decoding switch (default: on).""" + decode = dict(((root_config or {}).get("hydrocarbon") or {}).get("decode") or {}) + return bool(decode.get("hierarchical_plan_ranking", True)) + + +def select_plan_and_trajectory( + plan_log_probabilities: torch.Tensor, + decoded_by_plan: list[tuple[Any, float, int] | None], + *, + hierarchical: bool, +) -> tuple[float, int, Any, float, int]: + """Select a plan first, then its best plan-conditioned trajectory. + + In hierarchical mode raw path probability never participates in a + comparison between plans. ``hierarchical=False`` preserves the historical + ``log q(plan) + log P(path | plan)`` ranking exactly for ablations. + """ + if len(decoded_by_plan) != int(plan_log_probabilities.numel()): + raise ValueError("decoded plan count must match plan probabilities") + if not decoded_by_plan: + raise ValueError("no plans to decode") + + if hierarchical: + plan_index = int(torch.argmax(plan_log_probabilities).item()) + decoded = decoded_by_plan[plan_index] + if decoded is None: + raise RuntimeError("top-ranked hydrocarbon plan has no decoded trajectory") + terminal, path_logp, path_length = decoded + # Retain the joint value for logging only. It is not a selection score. + joint = float(plan_log_probabilities[plan_index].item()) + float(path_logp) + return joint, plan_index, terminal, float(path_logp), int(path_length) + + candidates = [] + for plan_index, decoded in enumerate(decoded_by_plan): + if decoded is None: + continue + terminal, path_logp, path_length = decoded + joint = float(plan_log_probabilities[plan_index].item()) + float(path_logp) + candidates.append((joint, plan_index, terminal, float(path_logp), int(path_length))) + if not candidates: + raise RuntimeError("no decoded hydrocarbon plans") + return max(candidates, key=lambda item: item[0]) + + +def sample_distinct_plans( + plans: list[HydrocarbonStaplePlan], + empirical_weights: list[float], + count: int, + rng: random.Random, +) -> list[HydrocarbonStaplePlan]: + """Weighted sampling without replacement from empirical plan support.""" + remaining = list(zip(plans, empirical_weights)) + selected: list[HydrocarbonStaplePlan] = [] + for _ in range(min(int(count), len(remaining))): + total = sum(max(float(weight), 0.0) for _, weight in remaining) + if total <= 0.0: + break + threshold = rng.random() * total + cumulative = 0.0 + chosen = len(remaining) - 1 + for index, (_, weight) in enumerate(remaining): + cumulative += max(float(weight), 0.0) + if cumulative >= threshold: + chosen = index + break + plan, _ = remaining.pop(chosen) + selected.append(plan) + return selected + + +def sample_committed_plan_trajectory( + sampler: PlanAwareReferenceSampler, + init_state: StapleState, + plan: HydrocarbonStaplePlan, + protected_positions: list[int], + context: dict[str, Any], + horizon: int, + early_stop: bool = True, +) -> PlanAwareTrajectory: + """Use the existing plan-aware action kernel with a caller-supplied plan.""" + states = [init_state] + labels: list[str] = [] + current = init_state + progress = PlanProgress(plan_selected=True) + no_neighbor = False + for _ in range(horizon): + candidates = sampler.graph.neighbors( + current, protected_positions=protected_positions + ) + if not candidates: + no_neighbor = True + break + nxt, label = sampler.kernel.sample_next( + current, candidates, plan, context=context + ) + states.append(nxt) + labels.append(label) + progress.n_actions += 1 + if label in ON_PLAN_LABELS: + progress.n_on_plan_actions += 1 + if label == ON_PLAN_FIRST_ANCHOR: + progress.first_anchor_installed = True + elif label == ON_PLAN_SECOND_ANCHOR: + progress.second_anchor_installed = True + elif label == ON_PLAN_ANCHOR_ASSIGN: + progress.anchor_assigned = True + elif label == ON_PLAN_BLOCK_ASSIGN: + progress.block_assigned = True + elif label in (ON_PLAN_TOPOLOGY, OFF_PLAN_TOPOLOGY): + progress.topology_activated = True + elif label == OFF_PLAN_SUBSTITUTION: + progress.n_off_plan_substitutions += 1 + elif label == OFF_PLAN_ANCHOR: + progress.n_off_plan_anchor_selections += 1 + current = nxt + if early_stop and current.topology == "stapled": + break + progress.plan_completed = completes_committed_plan(current, plan) + return PlanAwareTrajectory( + states=states, + plan=plan, + progress=progress, + action_labels=labels, + no_neighbor=no_neighbor, + ) + + +def plan_entropy(log_probabilities: torch.Tensor) -> torch.Tensor: + probabilities = log_probabilities.exp() + return -(probabilities * log_probabilities).sum() + + +def penetrance_target_weights(values: list[float], temperature: float, device: torch.device) -> torch.Tensor: + scores = torch.tensor(values, dtype=torch.float32, device=device) + return torch.softmax(scores / float(temperature), dim=0) + + +def plan_level_loss( + controlled_log_probabilities: torch.Tensor, + sampled_indices: list[int], + penetrance_values: list[float], + temperature: float, +) -> torch.Tensor: + if not sampled_indices: + return controlled_log_probabilities.sum() * 0.0 + target = penetrance_target_weights( + penetrance_values, temperature, controlled_log_probabilities.device + ) + index = torch.tensor(sampled_indices, dtype=torch.long, device=controlled_log_probabilities.device) + return -(target * controlled_log_probabilities[index]).sum() + + +def configured_plan_level_objective( + controlled_log_probabilities: torch.Tensor, + sampled_indices: list[int], + penetrance_values: list[float], + legacy_temperature: float, + *, + exact_sb_objective: bool = False, + reference_log_probabilities: torch.Tensor | None = None, + terminal_energies: torch.Tensor | None = None, + exact_sb_beta: float = 1.0, + target_support_mask: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Dispatch without perturbing the historical false branch. + + Returns ``(optimized_loss, direct_forward_objective, log_q_star)``. + """ + if not exact_sb_objective: + return ( + plan_level_loss( + controlled_log_probabilities, + sampled_indices, + penetrance_values, + legacy_temperature, + ), + None, + None, + ) + if reference_log_probabilities is None or terminal_energies is None: + raise ValueError("exact-SB requires q_ref and all-plan terminal energies") + return exact_sb_plan_objectives( + controlled_log_probabilities, + reference_log_probabilities, + terminal_energies, + exact_sb_beta, + target_support_mask=target_support_mask, + ) + + +def exact_sb_target_log_probabilities( + reference_log_probabilities: torch.Tensor, + terminal_energies: torch.Tensor, + beta: float, + *, + target_support_mask: torch.Tensor | None = None, +) -> torch.Tensor: + """Lead-local Gibbs target, optionally conditioned on a strict support. + + A false mask entry receives ``log q* = -inf`` and therefore probability + exactly zero. Callers must pass ``None`` for the empty-set fallback; this + preserves the historical permeability-only computation bit-for-bit. + """ + if reference_log_probabilities.ndim != 1 or terminal_energies.ndim != 1: + raise ValueError("exact-SB plan tensors must be one-dimensional for one lead") + if reference_log_probabilities.shape != terminal_energies.shape: + raise ValueError("q_ref and terminal energies must cover the same plans") + if reference_log_probabilities.numel() == 0: + raise ValueError("exact-SB requires at least one legal plan") + if float(beta) <= 0.0: + raise ValueError("exact-SB beta must be positive") + energies = terminal_energies.to( + device=reference_log_probabilities.device, + dtype=reference_log_probabilities.dtype, + ) + logits = reference_log_probabilities - float(beta) * energies + if target_support_mask is not None: + if target_support_mask.shape != logits.shape: + raise ValueError("target support mask must cover the same plans as q*") + mask = target_support_mask.to(device=logits.device, dtype=torch.bool) + if not bool(mask.any().item()): + raise ValueError("target support mask must be non-empty; use None for fallback") + logits = logits.masked_fill(~mask, float("-inf")) + return torch.log_softmax(logits, dim=0) + + +def exact_sb_plan_objectives( + controlled_log_probabilities: torch.Tensor, + reference_log_probabilities: torch.Tensor, + terminal_energies: torch.Tensor, + beta: float, + *, + target_support_mask: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return ``KL(q*||q_theta)``, direct SB objective, and ``log q*``. + + The reverse target KL is optimized. The direct forward variational + objective is returned only as a consistency diagnostic. + """ + if controlled_log_probabilities.ndim != 1: + raise ValueError("q_theta must be one-dimensional for one lead") + if controlled_log_probabilities.shape != reference_log_probabilities.shape: + raise ValueError("q_theta and q_ref must cover the same legal plans") + log_q_star = exact_sb_target_log_probabilities( + reference_log_probabilities, + terminal_energies, + beta, + target_support_mask=target_support_mask, + ) + q_star = log_q_star.exp().detach() + support = ( + torch.ones_like(q_star, dtype=torch.bool) + if target_support_mask is None + else target_support_mask.to(device=q_star.device, dtype=torch.bool) + ) + reverse_kl = torch.sum( + q_star[support] + * (log_q_star[support].detach() - controlled_log_probabilities[support]) + ) + energies = terminal_energies.to( + device=controlled_log_probabilities.device, + dtype=controlled_log_probabilities.dtype, + ) + if target_support_mask is None: + # Historical false/fallback branch: keep its arithmetic untouched. + q_theta = controlled_log_probabilities.exp() + forward_objective = torch.sum( + q_theta + * ( + controlled_log_probabilities + - reference_log_probabilities + + float(beta) * energies + ) + ) + else: + # Diagnostic-only conditional forward objective. q_theta is + # renormalized on the same strict support because KL(q_theta||q*) over + # the full plan set would be infinite whenever q_theta has outside + # mass. The trained objective remains reverse_kl above. + theta_log = torch.log_softmax(controlled_log_probabilities[support], dim=0) + ref_log = torch.log_softmax(reference_log_probabilities[support], dim=0) + theta = theta_log.exp() + forward_objective = torch.sum( + theta * (theta_log - ref_log + float(beta) * energies[support]) + ) + return reverse_kl, forward_objective, log_q_star + + +def describe_plan(plan: HydrocarbonStaplePlan) -> str: + i, j = plan.anchor_pair + return f"{plan.ordered_pair}/i,i+{plan.spacing}@{i}:{j}" diff --git a/staplebridge/hydrocarbon/plan_reference.py b/staplebridge/hydrocarbon/plan_reference.py new file mode 100644 index 0000000000000000000000000000000000000000..e68ba1973c502b239a19338a478a6895b357a3c8 --- /dev/null +++ b/staplebridge/hydrocarbon/plan_reference.py @@ -0,0 +1,932 @@ +"""Plan-aware empirical hydrocarbon reference process. + +Why this module exists +---------------------- +The hard-only reference in :mod:`staplebridge.hydrocarbon.actions` + +:class:`staplebridge.reference.kernel.ReferenceKernel` reaches a stapled terminal +on only ~21% of rollouts. The measured cause is **anchor overshoot**, not a +scoring problem: the action generator offers an anchor-monomer substitution at +almost every editable position, each individually legal, so an unguided walk +installs 5-7 anchor monomers. ``validate_hydrocarbon_staple`` then returns +``DOUBLE_STAPLE_UNSUPPORTED``, anchor assignment is never offered, and the +trajectory dead-ends with ``no_anchor_pair``. On a 32-lead probe, 98 of 101 +failures had >2 anchors installed and no anchor pair. + +The fix is to commit to a *whole staple plan* before walking, then bias the walk +toward finishing that plan: + + 1. enumerate every legal plan on the lead (S5-S5/i,i+4 and R8-S5/i,i+7); + 2. filter on protected positions, anchor conflicts, edit budget and catalog; + 3. draw one plan from q(plan | x) ∝ p_empirical(mode)^beta / n_mode(x); + 4. bias the per-step kernel toward first anchor -> second anchor -> + anchor/block assign -> topology activation for *that* plan; + 5. downweight substitutions and anchor re-selection unrelated to the plan. + +The ``1 / n_mode(x)`` factor is the point of step 3: i,i+4 admits more anchor +positions than i,i+7 on the same lead (8 vs 5 on a 12-mer), so weighting plans +by the raw mode probability would amplify i,i+4 purely by opportunity count. +Dividing by the per-lead legal-plan count of that mode makes the *mode* mass +exactly ``p^beta`` and the choice *within* a mode uniform. + +The empirical mode prior is consumed **once, here, at plan selection**. It is +deliberately not multiplied into every action and not re-counted in the terminal +energy; ``configs/hydrocarbon_empirical_reference.yaml`` therefore sets +``endpoint_prior.weight_pair: 0.0`` so the same table cannot be charged twice. + +Isolation +--------- +Additive and hydrocarbon-only. Nothing here is imported by the lactam path: +:class:`staplebridge.reference.kernel.ReferenceKernel`, +:class:`staplebridge.reference.sampler.ReferenceTrajectorySampler`, +``staplebridge.graph.neighbors`` and ``BridgeTrainer`` are wrapped, never +modified. The original hydrocarbon hard-only reference stays reachable exactly +as before, so it remains available as the ablation baseline. +""" + +from __future__ import annotations + +import json +import math +import random +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Final + +import torch + +from staplebridge.chemistry.state import StapleState +from staplebridge.data.schemas import BuildingBlock +from staplebridge.hydrocarbon.catalog import block_topology, is_hydrocarbon_block +from staplebridge.hydrocarbon.curriculum import ( + HydrocarbonStaplePlan, + propose_hydrocarbon_staple_plans, +) +from staplebridge.hydrocarbon.factorized_plan_reference import ( + FactorizedPlanReference, + FactorizedPlanReferenceConfig, +) +from staplebridge.hydrocarbon.tokenizer import is_anchor_token +from staplebridge.reference.kernel import ReferenceKernel + +#: Versioned subset of the generated empirical priors needed for plan +#: selection. Keeping it in the package makes defaults work in a clean clone; +#: the complete analysis output remains optional and generated. +DEFAULT_MODE_PRIOR_DIR: Final[str] = "staplebridge/hydrocarbon/data" + +#: Structural cost ``weighted_edit_distance`` charges for any completed staple: +#: anchor 1.0 + topology 0.5 + block 0.5. A plan's terminal weighted edit +#: distance is therefore ``n_edits + 2.0``, which is what the edit budget filter +#: has to compare against. +STAPLE_STRUCTURAL_EDIT_COST: Final[float] = 2.0 + +# -- action labels, relative to the committed plan --------------------------- +ON_PLAN_FIRST_ANCHOR: Final[str] = "on_plan_first_anchor" +ON_PLAN_SECOND_ANCHOR: Final[str] = "on_plan_second_anchor" +ON_PLAN_ANCHOR_ASSIGN: Final[str] = "on_plan_anchor_assign" +ON_PLAN_BLOCK_ASSIGN: Final[str] = "on_plan_block_assign" +ON_PLAN_TOPOLOGY: Final[str] = "on_plan_topology_activation" +OFF_PLAN_TOPOLOGY: Final[str] = "off_plan_topology_activation" +OFF_PLAN_SUBSTITUTION: Final[str] = "off_plan_substitution" +OFF_PLAN_ANCHOR: Final[str] = "off_plan_anchor_selection" +OFF_PLAN_BLOCK: Final[str] = "off_plan_block_assign" +PLAN_NOOP: Final[str] = "noop" + +#: Labels that count as progress on the committed plan. +ON_PLAN_LABELS: Final[frozenset[str]] = frozenset( + { + ON_PLAN_FIRST_ANCHOR, + ON_PLAN_SECOND_ANCHOR, + ON_PLAN_ANCHOR_ASSIGN, + ON_PLAN_BLOCK_ASSIGN, + ON_PLAN_TOPOLOGY, + } +) + + +class PlanSelectionError(RuntimeError): + """Raised when the empirical mode prior cannot be loaded.""" + + +# --------------------------------------------------------------------------- +# Empirical mode prior +# --------------------------------------------------------------------------- + + +@dataclass +class ModePriorConfig: + """Config for :class:`EmpiricalModePrior`. + + Only the modes the catalog actually supports are kept, and their + probabilities are renormalised over that restricted support. Without the + renormalisation the ``beta`` exponent would act on a distribution whose mass + partly sits on topologies the hard catalog forbids. + """ + + prior_dir: str = DEFAULT_MODE_PRIOR_DIR + dedup_version: str = "sequence_deduplicated" + use_smoothed: bool = True + #: Temperature on the empirical mode probabilities: ``p^beta``. 1.0 follows + #: the data exactly, 0.0 is uniform over modes. + beta: float = 0.75 + #: Floor for a catalog mode absent from the table, so an enabled topology is + #: never assigned probability zero. + unobserved_probability: float = 1e-3 + + @classmethod + def from_dict(cls, data: dict[str, Any] | None) -> "ModePriorConfig": + """Build from a ``hydrocarbon.plan_reference.mode_prior`` section.""" + cfg = cls() + for key, value in dict(data or {}).items(): + if not hasattr(cfg, key): + continue + current = getattr(cfg, key) + if isinstance(current, bool): + setattr(cfg, key, bool(value)) + elif isinstance(current, float): + setattr(cfg, key, float(value)) + else: + setattr(cfg, key, value) + return cfg + + +class EmpiricalModePrior: + """``p_empirical(mode)`` over the catalog's ``(pair, spacing)`` topologies. + + Args: + catalog: the hydrocarbon blocks in play. Defines the support. + config: prior configuration. + root: repository root used to resolve a relative ``prior_dir``. + + Raises: + PlanSelectionError: if the empirical table is missing or names no + catalog mode. Failing loudly beats silently falling back to uniform, + because "plan-aware *empirical* reference" would then be a misnomer. + """ + + def __init__( + self, + catalog: list[BuildingBlock], + config: ModePriorConfig | None = None, + root: Path | None = None, + ) -> None: + self.cfg = config or ModePriorConfig() + self._root = Path(root) if root is not None else Path(__file__).resolve().parents[2] + self.modes: list[tuple[str, int]] = [ + block_topology(b) for b in catalog if is_hydrocarbon_block(b) + ] + self._raw: dict[tuple[str, int], float] = {} + self._probabilities: dict[tuple[str, int], float] = {} + self._load() + + @property + def prior_dir(self) -> Path: + """Resolved directory holding the empirical JSON tables.""" + candidate = Path(self.cfg.prior_dir) + return candidate if candidate.is_absolute() else self._root / candidate + + def _load(self) -> None: + """Read ``pair_spacing_probabilities.json`` and restrict to the catalog.""" + path = self.prior_dir / "pair_spacing_probabilities.json" + if not path.is_file(): + raise PlanSelectionError( + f"plan-aware reference needs the empirical mode table at {path}. " + "It ships with this release at " + "staplebridge/hydrocarbon/data/pair_spacing_probabilities.json; " + "check hydrocarbon.plan_reference.mode_prior.prior_dir." + ) + with path.open("r", encoding="utf-8") as handle: + payload = json.load(handle) + versions = payload.get("probabilities_by_version") or {} + if self.cfg.dedup_version not in versions: + raise PlanSelectionError( + f"dedup version {self.cfg.dedup_version!r} not in {path.name}; " + f"available: {sorted(versions)}" + ) + categories = dict(versions[self.cfg.dedup_version].get("categories") or {}) + field_name = ( + "laplace_smoothed_probability" if self.cfg.use_smoothed else "raw_probability" + ) + + for pair, spacing in self.modes: + entry = categories.get(f"{pair}|{spacing}") or {} + value = entry.get(field_name) + self._raw[(pair, spacing)] = ( + float(self.cfg.unobserved_probability) + if value is None or float(value) <= 0.0 + else float(value) + ) + + total = sum(self._raw.values()) + if total <= 0.0: + raise PlanSelectionError( + f"no catalog mode has positive empirical probability in {path.name}; " + f"catalog modes: {self.modes}" + ) + self._probabilities = {k: v / total for k, v in self._raw.items()} + + def probability(self, mode: tuple[str, int]) -> float: + """Renormalised ``p_empirical(mode)``; 0.0 for a non-catalog mode.""" + return float(self._probabilities.get(mode, 0.0)) + + def tilted_weight(self, mode: tuple[str, int]) -> float: + """``p_empirical(mode) ** beta``, the weight used at plan selection.""" + probability = self.probability(mode) + return 0.0 if probability <= 0.0 else probability ** float(self.cfg.beta) + + def describe(self) -> dict[str, Any]: + """Summary for logging and audits.""" + return { + "prior_dir": str(self.prior_dir), + "dedup_version": self.cfg.dedup_version, + "use_smoothed": bool(self.cfg.use_smoothed), + "beta": float(self.cfg.beta), + "modes": [f"{p}/i,i+{s}" for p, s in self.modes], + "p_empirical": { + f"{p}/i,i+{s}": self.probability((p, s)) for p, s in self.modes + }, + "p_tilted": { + f"{p}/i,i+{s}": self.tilted_weight((p, s)) for p, s in self.modes + }, + "uses_permeability_label": False, + "is_trained_classifier": False, + "consumed": "once, at plan selection", + } + + +# --------------------------------------------------------------------------- +# Plan enumeration, filtering and selection +# --------------------------------------------------------------------------- + + +@dataclass +class PlanFilterConfig: + """Feasibility filters applied to enumerated plans.""" + + #: Reject plans needing more anchor substitutions than this. + max_anchor_edits: int = 2 + #: Terminal weighted-edit-distance ceiling (``edit_constraints.max_edit_budget``). + max_edit_budget: float = 6.0 + #: Minimum surviving sequence identity (``edit_constraints.min_sequence_identity``). + min_sequence_identity: float = 0.60 + + @classmethod + def from_config( + cls, hydro_cfg: dict[str, Any] | None, root_cfg: dict[str, Any] | None + ) -> "PlanFilterConfig": + """Read the curriculum and edit-constraint sections of a full config.""" + curriculum = dict((hydro_cfg or {}).get("curriculum") or {}) + edits = dict((root_cfg or {}).get("edit_constraints") or {}) + return cls( + max_anchor_edits=int(curriculum.get("max_anchor_edits", 2)), + max_edit_budget=float(edits.get("max_edit_budget", 6.0)), + min_sequence_identity=float(edits.get("min_sequence_identity", 0.60)), + ) + + +@dataclass +class PlanEnumerationReport: + """Why plans were rejected, and what the surviving mode mix looks like. + + Every counter accumulates, so one report can be threaded through a whole + batch of leads. ``n_enumerated`` and ``n_kept`` are therefore totals over all + enumeration calls, not per-lead values — mixing the two conventions in one + object would make the per-mode counts unreadable against them. + """ + + n_calls: int = 0 + n_enumerated: int = 0 + n_kept: int = 0 + rejected: dict[str, int] = field(default_factory=dict) + per_mode_counts: dict[str, int] = field(default_factory=dict) + + def reject(self, reason: str) -> None: + """Tally one rejection.""" + self.rejected[reason] = self.rejected.get(reason, 0) + 1 + + def as_dict(self) -> dict[str, Any]: + """JSON-serialisable view, with per-call means alongside the totals.""" + calls = max(self.n_calls, 1) + return { + "n_calls": int(self.n_calls), + "n_enumerated_total": int(self.n_enumerated), + "n_kept_total": int(self.n_kept), + "mean_enumerated_per_lead": float(self.n_enumerated / calls), + "mean_kept_per_lead": float(self.n_kept / calls), + "rejected": dict(sorted(self.rejected.items())), + "per_mode_counts": dict(sorted(self.per_mode_counts.items())), + } + + +def enumerate_legal_plans( + tokens: list[str], + catalog: list[BuildingBlock], + protected_positions: list[int] | None = None, + filters: PlanFilterConfig | None = None, + report: PlanEnumerationReport | None = None, +) -> list[HydrocarbonStaplePlan]: + """Every legal staple plan on ``tokens``, after feasibility filtering. + + Delegates catalog/protected/anchor-conflict/double-staple filtering to + :func:`~staplebridge.hydrocarbon.curriculum.propose_hydrocarbon_staple_plans` + (so the plan-aware reference and the curriculum oracle agree on what is + legal by construction), then applies the edit-budget and sequence-identity + constraints the curriculum does not check. + + Returns: + Plans in the curriculum's cheapest-first order. + """ + filters = filters or PlanFilterConfig() + report = report if report is not None else PlanEnumerationReport() + + plans = propose_hydrocarbon_staple_plans( + tokens, + catalog, + protected_positions=protected_positions, + max_anchor_edits=filters.max_anchor_edits, + ) + report.n_calls += 1 + report.n_enumerated += len(plans) + + kept: list[HydrocarbonStaplePlan] = [] + for plan in plans: + # Terminal weighted edit distance the plan would incur, including the + # fixed structural cost of closing a staple. + projected_edit = float(plan.n_edits) + STAPLE_STRUCTURAL_EDIT_COST + if projected_edit > filters.max_edit_budget: + report.reject("edit_budget_exhausted") + continue + identity = 1.0 - (plan.n_edits / len(tokens)) if tokens else 0.0 + if identity < filters.min_sequence_identity: + report.reject("below_min_sequence_identity") + continue + kept.append(plan) + mode = f"{plan.ordered_pair}/i,i+{plan.spacing}" + report.per_mode_counts[mode] = report.per_mode_counts.get(mode, 0) + 1 + + report.n_kept += len(kept) + return kept + + +def plan_selection_weights( + plans: list[HydrocarbonStaplePlan], mode_prior: EmpiricalModePrior +) -> list[float]: + """``q(plan | x) ∝ p_empirical(mode)^beta / n_mode(x)``, unnormalised. + + Dividing by ``n_mode(x)`` — the number of legal plans of that mode *on this + lead* — is what keeps i,i+4 from being amplified simply because it has more + admissible anchor positions than i,i+7. The resulting mode marginal is + exactly ``p^beta`` and the within-mode choice is uniform. + """ + counts: dict[tuple[str, int], int] = {} + for plan in plans: + key = (plan.ordered_pair, plan.spacing) + counts[key] = counts.get(key, 0) + 1 + + weights: list[float] = [] + for plan in plans: + key = (plan.ordered_pair, plan.spacing) + n_mode = counts[key] + weights.append(mode_prior.tilted_weight(key) / float(n_mode) if n_mode else 0.0) + return weights + + +def select_plan( + plans: list[HydrocarbonStaplePlan], + mode_prior: EmpiricalModePrior, + rng: random.Random, +) -> HydrocarbonStaplePlan | None: + """Draw one plan from ``q(plan | x)``. + + Returns ``None`` when there is no legal plan, or when every legal plan's mode + has zero empirical weight. + """ + if not plans: + return None + weights = plan_selection_weights(plans, mode_prior) + total = sum(weights) + if total <= 0.0: + return None + threshold = rng.random() * total + cumulative = 0.0 + for plan, weight in zip(plans, weights): + cumulative += weight + if cumulative >= threshold: + return plan + return plans[-1] + + +# --------------------------------------------------------------------------- +# Plan-conditional action labelling and biasing +# --------------------------------------------------------------------------- + + +@dataclass +class PlanBiasConfig: + """Log-space bonuses applied to the reference pmf, per plan-relative label. + + Positive values favour an action, negative values suppress it. The four + on-plan structural bonuses increase along the build order (first anchor -> + second anchor -> assign -> activate) so that a partially built plan is + always pulled forward rather than left to compete with a fresh restart. + + The off-plan substitution penalty is the load-bearing one: the action + generator offers an anchor substitution at nearly every editable position, + and unguided that is what installs a third anchor and kills the trajectory. + """ + + first_anchor: float = 3.0 + second_anchor: float = 3.5 + anchor_assign: float = 4.0 + block_assign: float = 4.0 + topology_activation: float = 4.5 + #: Closing a different pair contradicts the committed plan and strict + #: hierarchical inference. Keep it a failure, not an alternative positive. + off_plan_topology_activation: float = -4.5 + off_plan_substitution: float = -3.0 + off_plan_anchor_selection: float = -3.0 + off_plan_block_assign: float = -1.0 + noop: float = -1.0 + + @classmethod + def from_dict(cls, data: dict[str, Any] | None) -> "PlanBiasConfig": + """Build from a ``hydrocarbon.plan_reference.bias`` section.""" + cfg = cls() + for key, value in dict(data or {}).items(): + if hasattr(cfg, key): + setattr(cfg, key, float(value)) + return cfg + + def as_dict(self) -> dict[str, float]: + """Label -> bonus mapping used by the kernel.""" + return { + ON_PLAN_FIRST_ANCHOR: self.first_anchor, + ON_PLAN_SECOND_ANCHOR: self.second_anchor, + ON_PLAN_ANCHOR_ASSIGN: self.anchor_assign, + ON_PLAN_BLOCK_ASSIGN: self.block_assign, + ON_PLAN_TOPOLOGY: self.topology_activation, + OFF_PLAN_TOPOLOGY: self.off_plan_topology_activation, + OFF_PLAN_SUBSTITUTION: self.off_plan_substitution, + OFF_PLAN_ANCHOR: self.off_plan_anchor_selection, + OFF_PLAN_BLOCK: self.off_plan_block_assign, + PLAN_NOOP: self.noop, + } + + +def plan_positions_satisfied( + tokens: list[str], plan: HydrocarbonStaplePlan +) -> tuple[bool, bool]: + """Whether the plan's ``i`` and ``j`` anchor monomers are already installed.""" + i, j = plan.anchor_pair + i_token, j_token = plan.ordered_pair.split("-") + have_i = 0 <= i < len(tokens) and tokens[i].upper() == i_token + have_j = 0 <= j < len(tokens) and tokens[j].upper() == j_token + return have_i, have_j + + +def classify_against_plan( + state: StapleState, candidate: StapleState, plan: HydrocarbonStaplePlan +) -> str: + """Label the transition ``state -> candidate`` relative to ``plan``. + + Checked in the same order the build proceeds, so a composite transition + (the action generator sets ``block_id`` in the same step as the anchor + assignment) is attributed to its most advanced effect. + """ + plan_i, plan_j = plan.anchor_pair + i_token, j_token = plan.ordered_pair.split("-") + + # -- topology activation -------------------------------------------- + if state.topology != candidate.topology: + if candidate.topology != "stapled": + return PLAN_NOOP + on_plan = ( + candidate.anchor_pair is not None + and tuple(candidate.anchor_pair) == (plan_i, plan_j) + and candidate.block_id == plan.block_id + ) + return ON_PLAN_TOPOLOGY if on_plan else OFF_PLAN_TOPOLOGY + + # -- sequence edit --------------------------------------------------- + if state.sequence_tokens != candidate.sequence_tokens: + changed = [ + position + for position in range(min(len(state.sequence_tokens), len(candidate.sequence_tokens))) + if state.sequence_tokens[position] != candidate.sequence_tokens[position] + ] + if len(changed) != 1: + return OFF_PLAN_SUBSTITUTION + position = changed[0] + installed = candidate.sequence_tokens[position].upper() + wanted = ( + i_token if position == plan_i else j_token if position == plan_j else None + ) + if wanted is None or installed != wanted: + return OFF_PLAN_SUBSTITUTION + # Ordering is by *progress*, not by index: whichever of the two plan + # anchors lands first is the "first anchor" install. + have_i, have_j = plan_positions_satisfied(state.sequence_tokens, plan) + return ( + ON_PLAN_SECOND_ANCHOR if (have_i or have_j) else ON_PLAN_FIRST_ANCHOR + ) + + # -- anchor selection ------------------------------------------------ + if state.anchor_pair != candidate.anchor_pair: + if ( + candidate.anchor_pair is not None + and tuple(candidate.anchor_pair) == (plan_i, plan_j) + and candidate.block_id in (None, plan.block_id) + ): + return ON_PLAN_ANCHOR_ASSIGN + return OFF_PLAN_ANCHOR + + # -- block assignment ------------------------------------------------ + if state.block_id != candidate.block_id: + if ( + candidate.block_id == plan.block_id + and candidate.anchor_pair is not None + and tuple(candidate.anchor_pair) == (plan_i, plan_j) + ): + return ON_PLAN_BLOCK_ASSIGN + return OFF_PLAN_BLOCK + + return PLAN_NOOP + + +class PlanAwareReferenceKernel: + """Reference kernel that conditions on a committed staple plan. + + Wraps an unmodified :class:`~staplebridge.reference.kernel.ReferenceKernel`: + the base pmf (peptide/anchor/block priors, cost, geometry, action-progress, + group normalisation, substitution downweight) is computed exactly as today, + then reweighted by ``exp(bonus(label))`` and renormalised. With no plan + committed, or with all bonuses at zero, this is the base kernel. + + Reweighting in probability space rather than editing the base logits keeps + the two kernels directly comparable for the ablation: the only difference is + a plan-conditional multiplicative factor. + """ + + def __init__( + self, base_kernel: ReferenceKernel, bias: PlanBiasConfig | None = None + ) -> None: + self.base_kernel = base_kernel + self.bias = bias or PlanBiasConfig() + self._bonuses = self.bias.as_dict() + + def labels( + self, + state: StapleState, + candidates: list[StapleState], + plan: HydrocarbonStaplePlan | None, + ) -> list[str]: + """Plan-relative label for each candidate.""" + if plan is None: + return [PLAN_NOOP] * len(candidates) + return [classify_against_plan(state, c, plan) for c in candidates] + + def plan_probs( + self, + state: StapleState, + candidates: list[StapleState], + plan: HydrocarbonStaplePlan | None, + context: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, list[str]]: + """Plan-conditional pmf over ``candidates``, plus their labels.""" + probs = self.base_kernel.reference_probs(state, candidates, context=context) + if plan is None: + return probs, [PLAN_NOOP] * len(candidates) + + labels = self.labels(state, candidates, plan) + factors = torch.tensor( + [math.exp(self._bonuses.get(label, 0.0)) for label in labels], + dtype=torch.float32, + ) + tilted = probs * factors + total = float(tilted.sum().item()) + if total <= 0.0: + # Every candidate had zero base mass; fall back rather than emit a + # degenerate pmf that ``torch.multinomial`` would reject. + return probs, labels + return tilted / total, labels + + def sample_next( + self, + state: StapleState, + candidates: list[StapleState], + plan: HydrocarbonStaplePlan | None, + context: dict[str, Any] | None = None, + ) -> tuple[StapleState, str]: + """Draw one candidate from the plan-conditional pmf.""" + probs, labels = self.plan_probs(state, candidates, plan, context=context) + index = int(torch.multinomial(probs, num_samples=1).item()) + return candidates[index], labels[index] + + +# --------------------------------------------------------------------------- +# Plan-aware trajectory sampler +# --------------------------------------------------------------------------- + + +@dataclass +class PlanProgress: + """Which stages of the committed plan a trajectory actually reached. + + Recorded per stage rather than as a single success flag, because when the + stapled rate disappoints the question is always *which* stage lost the + trajectory. + """ + + plan_selected: bool = False + first_anchor_installed: bool = False + second_anchor_installed: bool = False + anchor_assigned: bool = False + block_assigned: bool = False + topology_activated: bool = False + plan_completed: bool = False + n_on_plan_actions: int = 0 + n_off_plan_substitutions: int = 0 + n_off_plan_anchor_selections: int = 0 + n_actions: int = 0 + + @property + def unrelated_substitution_rate(self) -> float: + """Share of this trajectory's actions that were off-plan substitutions.""" + return ( + self.n_off_plan_substitutions / self.n_actions if self.n_actions else 0.0 + ) + + def as_dict(self) -> dict[str, Any]: + """JSON-serialisable view.""" + return { + "plan_selected": bool(self.plan_selected), + "first_anchor_installed": bool(self.first_anchor_installed), + "second_anchor_installed": bool(self.second_anchor_installed), + "anchor_assigned": bool(self.anchor_assigned), + "block_assigned": bool(self.block_assigned), + "topology_activated": bool(self.topology_activated), + "plan_completed": bool(self.plan_completed), + "n_on_plan_actions": int(self.n_on_plan_actions), + "n_off_plan_substitutions": int(self.n_off_plan_substitutions), + "n_off_plan_anchor_selections": int(self.n_off_plan_anchor_selections), + "n_actions": int(self.n_actions), + "unrelated_substitution_rate": float(self.unrelated_substitution_rate), + } + + +@dataclass +class PlanAwareTrajectory: + """One plan-aware rollout.""" + + states: list[StapleState] + plan: HydrocarbonStaplePlan | None + progress: PlanProgress + action_labels: list[str] = field(default_factory=list) + no_plan_reason: str | None = None + #: True when the rollout stopped because the graph offered no neighbour. + no_neighbor: bool = False + + +class PlanAwareReferenceSampler: + """Reference sampler that commits to a plan, then completes it. + + Args: + graph: the hydrocarbon transition graph (unmodified). + kernel: the plan-conditional kernel. + mode_prior: empirical mode prior, consumed once per trajectory. + filters: plan feasibility filters. + seed: base seed for plan selection, kept separate from the global torch + RNG so plan draws are reproducible independently of the pmf draws. + """ + + def __init__( + self, + graph: Any, + kernel: PlanAwareReferenceKernel, + mode_prior: EmpiricalModePrior, + filters: PlanFilterConfig | None = None, + seed: int = 42, + factorized_reference: FactorizedPlanReference | None = None, + ) -> None: + self.graph = graph + self.kernel = kernel + self.mode_prior = mode_prior + self.filters = filters or PlanFilterConfig() + self.factorized_reference = factorized_reference + self._rng = random.Random(seed) + + @property + def factorized_plan_reference_enabled(self) -> bool: + return self.factorized_reference is not None + + def plan_selection_weights( + self, + initial: StapleState, + plans: list[HydrocarbonStaplePlan], + context: dict[str, Any] | None = None, + ) -> list[float]: + """Active plan-reference weights, with a bit-exact legacy branch.""" + if self.factorized_reference is None: + return plan_selection_weights(plans, self.mode_prior) + return self.factorized_reference.weights(initial, plans, context) + + def select_plan( + self, + initial: StapleState, + plans: list[HydrocarbonStaplePlan], + context: dict[str, Any] | None = None, + ) -> HydrocarbonStaplePlan | None: + if not plans: + return None + weights = self.plan_selection_weights(initial, plans, context) + total = sum(weights) + if total <= 0.0: + return None + threshold = self._rng.random() * total + cumulative = 0.0 + for plan, weight in zip(plans, weights): + cumulative += weight + if cumulative >= threshold: + return plan + return plans[-1] + + def sample_trajectory( + self, + init_state: StapleState, + protected_positions: list[int], + context: dict[str, Any], + horizon: int, + early_stop: bool = True, + report: PlanEnumerationReport | None = None, + ) -> PlanAwareTrajectory: + """Select a plan for ``init_state``, then walk toward completing it.""" + plans = enumerate_legal_plans( + init_state.sequence_tokens, + self.graph.catalog, + protected_positions=protected_positions, + filters=self.filters, + report=report, + ) + plan = self.select_plan(init_state, plans, context) + progress = PlanProgress(plan_selected=plan is not None) + if plan is None: + reason = "no_legal_plan" if not plans else "no_mode_weight" + return PlanAwareTrajectory( + states=[init_state], plan=None, progress=progress, no_plan_reason=reason + ) + + states = [init_state] + labels: list[str] = [] + current = init_state + no_neighbor = False + + for _ in range(horizon): + candidates = self.graph.neighbors( + current, protected_positions=protected_positions + ) + if not candidates: + no_neighbor = True + break + nxt, label = self.kernel.sample_next( + current, candidates, plan, context=context + ) + states.append(nxt) + labels.append(label) + + progress.n_actions += 1 + if label in ON_PLAN_LABELS: + progress.n_on_plan_actions += 1 + if label == ON_PLAN_FIRST_ANCHOR: + progress.first_anchor_installed = True + elif label == ON_PLAN_SECOND_ANCHOR: + progress.second_anchor_installed = True + elif label == ON_PLAN_ANCHOR_ASSIGN: + progress.anchor_assigned = True + elif label == ON_PLAN_BLOCK_ASSIGN: + progress.block_assigned = True + elif label == ON_PLAN_TOPOLOGY: + progress.topology_activated = True + elif label == OFF_PLAN_TOPOLOGY: + progress.topology_activated = True + elif label == OFF_PLAN_SUBSTITUTION: + progress.n_off_plan_substitutions += 1 + elif label == OFF_PLAN_ANCHOR: + progress.n_off_plan_anchor_selections += 1 + + current = nxt + if early_stop and current.topology == "stapled": + break + + # The anchor assignment is composite (it sets block_id in the same + # transition), so credit block assignment from the terminal state rather + # than requiring a separate labelled step. + if current.block_id == plan.block_id and tuple( + current.anchor_pair or (-1, -1) + ) == plan.anchor_pair: + progress.block_assigned = True + have_i, have_j = plan_positions_satisfied(current.sequence_tokens, plan) + if have_i and have_j: + progress.first_anchor_installed = True + progress.second_anchor_installed = True + elif have_i or have_j: + progress.first_anchor_installed = True + + progress.plan_completed = bool( + current.topology == "stapled" + and current.anchor_pair is not None + and tuple(current.anchor_pair) == plan.anchor_pair + and current.block_id == plan.block_id + ) + + return PlanAwareTrajectory( + states=states, + plan=plan, + progress=progress, + action_labels=labels, + no_neighbor=no_neighbor, + ) + + def sample_batch( + self, + init_state: StapleState, + protected_positions: list[int], + context: dict[str, Any], + horizon: int, + n: int, + report: PlanEnumerationReport | None = None, + ) -> list[PlanAwareTrajectory]: + """``n`` independent plan-aware rollouts from ``init_state``.""" + return [ + self.sample_trajectory( + init_state, + protected_positions=protected_positions, + context=context, + horizon=horizon, + report=report, + ) + for _ in range(n) + ] + + +@dataclass +class PlanReferenceConfig: + """Full config for the plan-aware reference, from a ``hydrocarbon`` section.""" + + enabled: bool = False + mode_prior: ModePriorConfig = field(default_factory=ModePriorConfig) + bias: PlanBiasConfig = field(default_factory=PlanBiasConfig) + filters: PlanFilterConfig = field(default_factory=PlanFilterConfig) + factorized: FactorizedPlanReferenceConfig = field( + default_factory=FactorizedPlanReferenceConfig + ) + + @classmethod + def from_config(cls, root_cfg: dict[str, Any] | None) -> "PlanReferenceConfig": + """Read ``hydrocarbon.plan_reference`` plus the shared edit constraints.""" + root_cfg = dict(root_cfg or {}) + hydro_cfg = dict(root_cfg.get("hydrocarbon") or {}) + section = dict(hydro_cfg.get("plan_reference") or {}) + return cls( + enabled=bool(section.get("enabled", False)), + mode_prior=ModePriorConfig.from_dict(section.get("mode_prior")), + bias=PlanBiasConfig.from_dict(section.get("bias")), + filters=PlanFilterConfig.from_config(hydro_cfg, root_cfg), + factorized=FactorizedPlanReferenceConfig.from_config(root_cfg), + ) + + +def build_plan_aware_sampler( + graph: Any, + base_kernel: ReferenceKernel, + root_cfg: dict[str, Any] | None, + seed: int = 42, + root: Path | None = None, +) -> tuple[PlanAwareReferenceSampler, PlanReferenceConfig]: + """Assemble the plan-aware sampler from a full config mapping.""" + cfg = PlanReferenceConfig.from_config(root_cfg) + mode_prior = EmpiricalModePrior(graph.catalog, cfg.mode_prior, root=root) + kernel = PlanAwareReferenceKernel(base_kernel, cfg.bias) + factorized_reference = None + if cfg.factorized.enabled: + energy = base_kernel.energy_model + factorized_reference = FactorizedPlanReference( + mode_prior=mode_prior, + catalog=graph.catalog, + peptide_prior=energy.peptide_prior, + anchor_prior=energy.anchor_prior, + block_prior=energy.block_prior, + config=cfg.factorized, + ) + sampler = PlanAwareReferenceSampler( + graph, + kernel, + mode_prior, + filters=cfg.filters, + seed=seed, + factorized_reference=factorized_reference, + ) + return sampler, cfg + + +def count_anchor_monomers(tokens: list[str]) -> int: + """Number of hydrocarbon anchor monomers in ``tokens``.""" + return sum(1 for t in tokens if is_anchor_token(t)) diff --git a/staplebridge/hydrocarbon/plan_validation.py b/staplebridge/hydrocarbon/plan_validation.py new file mode 100644 index 0000000000000000000000000000000000000000..a1f779304f42b2476410df0597986f276b6baea2 --- /dev/null +++ b/staplebridge/hydrocarbon/plan_validation.py @@ -0,0 +1,214 @@ +"""Exact-SB validation diagnostics for the hydrocarbon plan-control head. + +The lead-local Gibbs target is computed with exactly the training definition, +by delegating to :func:`configured_plan_level_objective` with +``exact_sb_objective=True``. Nothing here participates in candidate +generation, plan selection, or reranking: ``q*`` is a diagnostic and a +model-selection signal only. +""" + +from __future__ import annotations + +from collections import Counter +from typing import Any, Sequence + +import torch + +from staplebridge.hydrocarbon.plan_control import ( + HydrocarbonPlanControlConfig, + configured_plan_level_objective, +) + +# The validation metrics this module contributes. They are additive: no +# existing validation metric is replaced or removed. +Q_STAR_VALIDATION_METRICS = ( + "q_star_vs_q_theta_kl", + "q_star_top1_agreement", + "q_star_spearman", + "q_star_entropy", + "q_theta_entropy", + "exact_sb_forward_objective", +) + +S5_MODE = "S5-S5/i,i+4" +R8_MODE = "R8-S5/i,i+7" + + +def exact_sb_validation_enabled(root_config: dict[str, Any] | None) -> bool: + """Only ``exact_sb_objective`` **and** ``validation_all_plans`` opt in.""" + cfg = HydrocarbonPlanControlConfig.from_config(root_config) + return bool(cfg.enabled and cfg.exact_sb_objective and cfg.validation_all_plans) + + +def _average_ranks(values: torch.Tensor) -> torch.Tensor: + """Ascending ranks with ties averaged, as Spearman's rho requires.""" + order = torch.argsort(values) + sorted_values = values[order] + ranks = torch.empty_like(sorted_values, dtype=torch.float64) + start = 0 + for index in range(1, sorted_values.numel() + 1): + if index == sorted_values.numel() or sorted_values[index] != sorted_values[start]: + ranks[start:index] = 0.5 * float(start + index - 1) + 1.0 + start = index + out = torch.empty_like(ranks) + out[order] = ranks + return out + + +def spearman_correlation(left: torch.Tensor, right: torch.Tensor) -> float | None: + """Rank correlation, or ``None`` when it is undefined for this lead. + + Undefined means fewer than two plans, or one side constant (every plan + tied), in which case there is no ordering to agree or disagree with. + """ + if left.shape != right.shape: + raise ValueError("spearman inputs must cover the same plans") + if left.numel() < 2: + return None + left_ranks = _average_ranks(left.detach().to(dtype=torch.float64).flatten()) + right_ranks = _average_ranks(right.detach().to(dtype=torch.float64).flatten()) + left_centered = left_ranks - left_ranks.mean() + right_centered = right_ranks - right_ranks.mean() + denominator = float( + torch.sqrt((left_centered**2).sum() * (right_centered**2).sum()).item() + ) + if denominator <= 0.0: + return None + return float((left_centered * right_centered).sum().item() / denominator) + + +def lead_q_star_diagnostics( + controlled_log_probabilities: torch.Tensor, + reference_log_probabilities: torch.Tensor, + terminal_energies: torch.Tensor, + beta: float, + plan_modes: Sequence[str] | None = None, + target_support_mask: torch.Tensor | None = None, +) -> dict[str, Any]: + """Per-lead ``q*`` diagnostics over the hard-supported legal plans. + + ``q*`` and the forward objective come from the same + :func:`configured_plan_level_objective` call the training loop uses, so a + validation KL can never drift from the trained objective's definition. + """ + reverse_kl, forward_objective, log_q_star = configured_plan_level_objective( + controlled_log_probabilities, + [], + [], + # Unused by the exact-SB branch; passed only to satisfy the shared + # signature. Any positive value leaves the exact-SB result untouched. + 1.0, + exact_sb_objective=True, + reference_log_probabilities=reference_log_probabilities, + terminal_energies=terminal_energies, + exact_sb_beta=beta, + target_support_mask=target_support_mask, + ) + if forward_objective is None or log_q_star is None: + raise RuntimeError("exact-SB validation diagnostics require the exact-SB branch") + q_star = log_q_star.exp() + q_theta = controlled_log_probabilities.exp() + star_top = int(torch.argmax(log_q_star).item()) + theta_top = int(torch.argmax(controlled_log_probabilities).item()) + if target_support_mask is None: + # Historical false/fallback branch is deliberately bit-for-bit. + q_star_entropy = -(q_star * log_q_star).sum() + else: + q_star_entropy = -torch.where( + q_star > 0.0, q_star * log_q_star, torch.zeros_like(q_star) + ).sum() + result: dict[str, Any] = { + "n_plans": int(controlled_log_probabilities.numel()), + "q_star_vs_q_theta_kl": float(reverse_kl.detach().cpu().item()), + "exact_sb_forward_objective": float(forward_objective.detach().cpu().item()), + "q_star_entropy": float(q_star_entropy.detach().cpu().item()), + "q_theta_entropy": float( + (-(q_theta * controlled_log_probabilities).sum()).detach().cpu().item() + ), + "q_star_top1_index": star_top, + "q_theta_top1_index": theta_top, + "q_star_top1_agreement": bool(star_top == theta_top), + "q_star_spearman": spearman_correlation( + log_q_star.detach(), controlled_log_probabilities.detach() + ), + "exact_sb_beta": float(beta), + "q_star_support_size": int( + target_support_mask.to(dtype=torch.bool).sum().item() + if target_support_mask is not None + else controlled_log_probabilities.numel() + ), + } + if plan_modes is not None: + if len(plan_modes) != controlled_log_probabilities.numel(): + raise ValueError("plan_modes must cover the same plans as q* and q_theta") + for distribution_name, probabilities in ( + ("q_star", q_star), + ("q_theta", q_theta), + ): + for metric_name, mode in ( + ("s5_s5_i4", S5_MODE), + ("r8_s5_i7", R8_MODE), + ): + indices = [index for index, value in enumerate(plan_modes) if value == mode] + mass = ( + float(probabilities[indices].sum().detach().cpu().item()) + if indices + else 0.0 + ) + result[f"{distribution_name}_{metric_name}_probability_mass"] = mass + result["q_star_top1_mode"] = str(plan_modes[star_top]) + result["q_theta_top1_mode"] = str(plan_modes[theta_top]) + return result + + +def _mean(values: Sequence[float]) -> float | None: + return float(sum(values) / len(values)) if values else None + + +def aggregate_q_star_diagnostics(rows: Sequence[dict[str, Any]]) -> dict[str, Any]: + """Average per-lead diagnostics into the validation summary metrics.""" + spearmans = [ + float(row["q_star_spearman"]) + for row in rows + if row.get("q_star_spearman") is not None + ] + result: dict[str, Any] = { + "q_star_vs_q_theta_kl": _mean([float(row["q_star_vs_q_theta_kl"]) for row in rows]), + "q_star_top1_agreement": _mean( + [float(bool(row["q_star_top1_agreement"])) for row in rows] + ), + "q_star_spearman": _mean(spearmans), + "q_star_entropy": _mean([float(row["q_star_entropy"]) for row in rows]), + "q_theta_entropy": _mean( + [float(row["q_theta_entropy"]) for row in rows if row.get("q_theta_entropy") is not None] + ), + "exact_sb_forward_objective": _mean( + [float(row["exact_sb_forward_objective"]) for row in rows] + ), + "q_star_leads_scored": len(rows), + "q_star_plans_scored": int(sum(int(row["n_plans"]) for row in rows)), + "q_star_spearman_defined_leads": len(spearmans), + # Stated explicitly because it is the invariant this module must hold: + # diagnostics never touch the decoded candidate set or its ranking. + "q_star_used_for_candidate_reranking": False, + } + for distribution_name in ("q_star", "q_theta"): + for metric_name in ("s5_s5_i4", "r8_s5_i7"): + key = f"{distribution_name}_{metric_name}_probability_mass" + values = [float(row[key]) for row in rows if row.get(key) is not None] + if values: + result[key] = _mean(values) + top1_key = f"{distribution_name}_top1_mode" + modes = Counter(str(row[top1_key]) for row in rows if row.get(top1_key)) + if modes: + result[f"{distribution_name}_top1_counts"] = dict(modes) + result[f"{distribution_name}_top1_mix"] = { + mode: count / len(rows) for mode, count in modes.items() + } + result[f"{distribution_name}_s5_s5_i4_top1_rate"] = ( + modes.get(S5_MODE, 0) / len(rows) + ) + result[f"{distribution_name}_r8_s5_i7_top1_rate"] = ( + modes.get(R8_MODE, 0) / len(rows) + ) + return result diff --git a/staplebridge/hydrocarbon/property_energy.py b/staplebridge/hydrocarbon/property_energy.py new file mode 100644 index 0000000000000000000000000000000000000000..1d38b8c7abe707ef3588fec553da779e01185156 --- /dev/null +++ b/staplebridge/hydrocarbon/property_energy.py @@ -0,0 +1,703 @@ +"""Strict PeptiVerse SMILES scoring for hydrocarbon terminal states. + +Hydrocarbon-only: no lactam module imports this file. A legal terminal is +converted to the neutral canonical linear precursor and both E/Z products. +Penetrance remains the primary objective. Optional lead-relative preservation +adds either the legacy Toxicity/Solubility/Half-life hinges or the independent +Half-life-only hinge without changing the property-free feasible support. A +third independent switch exposes a permeability/half-life condition used only +to mask the Exact-SB teacher support; it never changes terminal energy itself. + +Two scoring depths exist: + +* :meth:`HydrocarbonPropertyScorer.score` — every property on every SMILES. + This is what validation, candidate ranking and final evaluation consume, and + it is unchanged. +* :meth:`HydrocarbonPropertyScorer.score_energy_only` — only the properties + :func:`required_energy_properties` says actually reach the terminal energy, + and only on the E/Z products whose mean the energy reads. Used when the + caller discards the info dict, as the Exact-SB ``q*`` construction does. + The resulting energy is identical because the skipped predictions are + provably not summed into it. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Protocol, Sequence + +from rdkit import Chem + +from staplebridge.chemistry.state import StapleState +from staplebridge.hydrocarbon.smiles_builder import ( + OlefinGeometry, + StapleSmilesPlan, + _assemble_linear, + build_stapled_smiles, +) +from staplebridge.hydrocarbon.monomers import C_TERM_ACID, N_TERM_FREE + + +class StrictSmilesPredictor(Protocol): + """Minimal PeptiVerse wrapper interface used by the scorer.""" + + available: bool + + def predict_property( + self, + property_key: str, + input_str: str, + mode: str = "wt", + uncertainty: bool | None = None, + ) -> dict[str, Any]: + ... + + def predict_property_batch( + self, + property_key: str, + input_strs: Sequence[str], + mode: str = "wt", + uncertainty: bool | None = None, + ) -> list[dict[str, Any]]: + """Same contract as :meth:`predict_property`, in input order.""" + ... + + +@dataclass +class HydrocarbonPropertyEnergyConfig: + """Property objective and optional lead-relative preservation penalties. + + Every developability field defaults to an inert value behind + ``enable_developability_constraints=False``. This preserves the historical + permeability-only energy and its exact PeptiVerse call set. + """ + + enabled: bool = False + penetrance_weight: float = 5.0 + toxicity_guard_enabled: bool = False + toxicity_threshold: float = 0.49 + toxicity_guard_weight: float = 10.0 + enable_developability_constraints: bool = False + developability_lambda_toxicity: float = 1.0 + developability_lambda_solubility: float = 1.0 + developability_lambda_halflife: float = 1.0 + developability_normalization_toxicity_scale: float = 1.0 + developability_normalization_solubility_scale: float = 1.0 + developability_normalization_halflife_scale: float = 1.0 + developability_normalization_artifact: str = "" + enable_halflife_preservation: bool = False + enable_joint_perm_halflife_support: bool = False + halflife_preservation_lambda: float = 1.0 + halflife_preservation_scale: float = 1.0 + halflife_preservation_scale_method: str = "" + halflife_preservation_train_sample_count: float = 0.0 + halflife_preservation_normalization_artifact: str = "" + + @classmethod + def from_dict( + cls, data: dict[str, Any] | None + ) -> "HydrocarbonPropertyEnergyConfig": + cfg = cls() + for key, value in dict(data or {}).items(): + if not hasattr(cfg, key): + continue + current = getattr(cfg, key) + if isinstance(current, bool): + converted = bool(value) + elif isinstance(current, str): + converted = str(value) + else: + converted = float(value) + setattr(cfg, key, converted) + if cfg.enable_developability_constraints: + scales = { + "toxicity": cfg.developability_normalization_toxicity_scale, + "solubility": cfg.developability_normalization_solubility_scale, + "halflife": cfg.developability_normalization_halflife_scale, + } + invalid = { + key: value + for key, value in scales.items() + if not math.isfinite(float(value)) or float(value) <= 0.0 + } + if invalid: + raise ValueError( + "developability normalization scales must be finite and > 0: " + f"{invalid}" + ) + weights = { + "toxicity": cfg.developability_lambda_toxicity, + "solubility": cfg.developability_lambda_solubility, + "halflife": cfg.developability_lambda_halflife, + } + invalid_weights = { + key: value + for key, value in weights.items() + if not math.isfinite(float(value)) or float(value) < 0.0 + } + if invalid_weights: + raise ValueError( + "developability lambdas must be finite and >= 0: " + f"{invalid_weights}" + ) + if cfg.enable_halflife_preservation: + if ( + not math.isfinite(float(cfg.halflife_preservation_scale)) + or float(cfg.halflife_preservation_scale) <= 0.0 + ): + raise ValueError( + "halflife preservation scale must be finite and > 0: " + f"{cfg.halflife_preservation_scale}" + ) + if ( + not math.isfinite(float(cfg.halflife_preservation_lambda)) + or float(cfg.halflife_preservation_lambda) < 0.0 + ): + raise ValueError( + "halflife preservation lambda must be finite and >= 0: " + f"{cfg.halflife_preservation_lambda}" + ) + if ( + not math.isfinite( + float(cfg.halflife_preservation_train_sample_count) + ) + or float(cfg.halflife_preservation_train_sample_count) < 2.0 + ): + raise ValueError( + "halflife preservation train sample count must be finite and >= 2" + ) + if cfg.enable_joint_perm_halflife_support and ( + cfg.enable_developability_constraints or cfg.enable_halflife_preservation + ): + raise ValueError( + "joint permeability/half-life support requires both legacy " + "developability and half-life soft-preservation flags to be false" + ) + return cfg + + def describe(self) -> dict[str, Any]: + return { + "enabled": bool(self.enabled), + "penetrance_weight": float(self.penetrance_weight), + "toxicity_guard_enabled": bool(self.toxicity_guard_enabled), + "toxicity_threshold": float(self.toxicity_threshold), + "toxicity_guard_weight": float(self.toxicity_guard_weight), + "enable_developability_constraints": bool( + self.enable_developability_constraints + ), + "developability_lambda_toxicity": float( + self.developability_lambda_toxicity + ), + "developability_lambda_solubility": float( + self.developability_lambda_solubility + ), + "developability_lambda_halflife": float( + self.developability_lambda_halflife + ), + "developability_normalization": { + "toxicity_scale": float( + self.developability_normalization_toxicity_scale + ), + "solubility_scale": float( + self.developability_normalization_solubility_scale + ), + "halflife_scale": float( + self.developability_normalization_halflife_scale + ), + "artifact": self.developability_normalization_artifact, + }, + "enable_halflife_preservation": bool( + self.enable_halflife_preservation + ), + "enable_joint_perm_halflife_support": bool( + self.enable_joint_perm_halflife_support + ), + "joint_perm_halflife_support": { + "condition": "delta_permeability > 0 and delta_halflife >= 0", + "empty_set_fallback": "permeability-only Full q_star", + "changes_terminal_energy": False, + }, + "halflife_preservation": { + "lambda": float(self.halflife_preservation_lambda), + "scale": float(self.halflife_preservation_scale), + "scale_method": self.halflife_preservation_scale_method, + "train_sample_count": int( + self.halflife_preservation_train_sample_count + ), + "normalization_artifact": ( + self.halflife_preservation_normalization_artifact + ), + "penalty": "lambda * max(0, -delta_halflife / scale)", + }, + "input_convention": "neutral_canonical", + "product_geometries": ["E", "Z"], + "product_aggregation": "mean", + "ez_uncertainty": "absolute_difference", + "energy_properties": list(required_energy_properties(self)), + "monitor_only": ( + ["hemolysis"] + if self.enable_developability_constraints + else ( + ["toxicity", "hemolysis"] + if ( + self.enable_halflife_preservation + or self.enable_joint_perm_halflife_support + ) + else ["toxicity", "hemolysis", "halflife"] + ) + ), + "excluded": ( + ["binding_affinity"] + if self.enable_developability_constraints + else ["solubility", "binding_affinity"] + ), + } + + +class HydrocarbonPropertyScoringError(RuntimeError): + """Construction/backend/mode/non-finite failure in strict scoring.""" + + +def required_energy_properties( + cfg: HydrocarbonPropertyEnergyConfig, +) -> tuple[str, ...]: + """Properties needed by terminal energy or the Exact-SB target support. + + Mirrors :func:`property_energy_terms` exactly: it reads + ``permeability_penetrance`` unconditionally and ``toxicity`` only when the + guard is on. Anything else the scorer reports is a monitor, so omitting it + cannot move the energy. Keep this in lockstep with + :func:`property_energy_terms` — the equivalence test asserts they agree. + """ + if not cfg.enabled: + return () + required = ["permeability_penetrance"] + if cfg.toxicity_guard_enabled: + required.append("toxicity") + if cfg.enable_developability_constraints: + required.extend(("toxicity", "solubility", "halflife")) + if cfg.enable_halflife_preservation: + required.append("halflife") + if cfg.enable_joint_perm_halflife_support: + required.append("halflife") + return tuple(dict.fromkeys(required)) + + +def required_original_lead_properties( + cfg: HydrocarbonPropertyEnergyConfig, +) -> tuple[str, ...]: + """Original-lead scores needed by the configured terminal objective. + + The historical path requests penetrance only. The developability path also + needs all three preservation baselines and retains penetrance for the + existing delta-PV evaluator. + """ + required = ["permeability_penetrance"] + if cfg.enable_developability_constraints: + required.extend(("toxicity", "solubility", "halflife")) + if cfg.enable_halflife_preservation: + required.append("halflife") + if cfg.enable_joint_perm_halflife_support: + required.append("halflife") + return tuple(dict.fromkeys(required)) + + +def full_scoring_properties( + cfg: HydrocarbonPropertyEnergyConfig, +) -> tuple[str, ...]: + """Properties recorded by full validation/ranking for this arm.""" + if cfg.enable_developability_constraints: + return HydrocarbonPropertyScorer.ALL_PROPERTIES + ("solubility",) + return HydrocarbonPropertyScorer.ALL_PROPERTIES + + +class HydrocarbonPropertyScorer: + """Score a legal hydrocarbon terminal with strict PeptiVerse.""" + + ENERGY_PROPERTIES = ("permeability_penetrance", "toxicity") + MONITOR_PROPERTIES = ("hemolysis", "halflife") + ALL_PROPERTIES = ENERGY_PROPERTIES + MONITOR_PROPERTIES + + def __init__( + self, + predictor: StrictSmilesPredictor, + original_linear_cache: dict[tuple[Any, ...], dict[str, Any]] | None = None, + ) -> None: + if not getattr(predictor, "available", False): + raise HydrocarbonPropertyScoringError( + "strict scorer requires available PeptiVerse" + ) + self.predictor = predictor + self._cache: dict[tuple[Any, ...], dict[str, Any]] = {} + self._original_linear_cache = ( + {} if original_linear_cache is None else original_linear_cache + ) + self.cache_hits = 0 + self.cache_misses = 0 + self.original_linear_cache_hits = 0 + self.original_linear_cache_misses = 0 + + @staticmethod + def _canonical_neutral(smiles: str) -> str: + mol = Chem.MolFromSmiles(smiles) + if mol is None: + raise HydrocarbonPropertyScoringError("RDKit could not parse built SMILES") + charge = sum(atom.GetFormalCharge() for atom in mol.GetAtoms()) + if charge != 0: + raise HydrocarbonPropertyScoringError( + f"neutral input convention violated: formal charge={charge}" + ) + return Chem.MolToSmiles(mol, isomericSmiles=True, canonical=True) + + def _predict(self, property_key: str, smiles: str) -> float: + result = self.predictor.predict_property( + property_key, input_str=smiles, mode="smiles", uncertainty=False + ) + backend = result.get("backend_used") + mode = result.get("mode") + try: + score = float(result["score"]) + except (KeyError, TypeError, ValueError) as exc: + raise HydrocarbonPropertyScoringError( + f"{property_key}: missing numeric score in {result!r}" + ) from exc + if backend != "peptiverse": + raise HydrocarbonPropertyScoringError( + f"{property_key}: forbidden backend/fallback {backend!r}" + ) + if mode != "smiles": + raise HydrocarbonPropertyScoringError( + f"{property_key}: forbidden mode {mode!r}" + ) + if not math.isfinite(score): + raise HydrocarbonPropertyScoringError( + f"{property_key}: non-finite score {score!r}" + ) + return score + + def prefetch( + self, properties: Sequence[str], smiles: Sequence[str] + ) -> dict[str, int]: + """Batch-predict ``properties`` over ``smiles`` to warm the cache. + + Purely a warming step: the results land in the predictor's own + memo, so the subsequent per-terminal ``_predict`` calls hit it and + return the same floats they would have computed one at a time. Callers + that lack a batch-capable predictor are served by the scalar loop, so + this is always safe to call. + """ + batch = getattr(self.predictor, "predict_property_batch", None) + unique = list(dict.fromkeys(s for s in smiles if s)) + if not unique or not callable(batch): + return {"properties": len(properties), "smiles": len(unique), "batched": 0} + for prop in properties: + batch(prop, unique, mode="smiles", uncertainty=False) + return { + "properties": len(properties), + "smiles": len(unique), + "batched": len(properties) * len(unique), + } + + def original_linear_smiles(self, sequence_tokens: Sequence[str]) -> str: + """Build the same neutral canonical unedited-lead SMILES as production.""" + tokens = tuple(str(token).upper() for token in sequence_tokens) + linear_rw, _ = _assemble_linear(list(tokens), N_TERM_FREE, C_TERM_ACID) + linear_mol = linear_rw.GetMol() + Chem.SanitizeMol(linear_mol) + return self._canonical_neutral( + Chem.MolToSmiles(linear_mol, isomericSmiles=True) + ) + + def score_original_linear( + self, + sequence_tokens: list[str], + *, + lead_key: str | None = None, + properties: Sequence[str] | None = None, + ) -> dict[str, Any]: + """Score the unedited original lead once, independently of endpoint plan.""" + tokens = tuple(str(token).upper() for token in sequence_tokens) + key = (lead_key, tokens) + requested = tuple(properties or ("permeability_penetrance",)) + info = dict(self._original_linear_cache.get(key, {})) + missing = [ + prop + for prop in requested + if f"hydrocarbon_{prop}_original_linear" not in info + ] + if not missing: + self.original_linear_cache_hits += 1 + return info + + smiles = str(info.get("hydrocarbon_original_linear_smiles") or "") + if not smiles: + smiles = self.original_linear_smiles(tokens) + info["hydrocarbon_original_linear_smiles"] = smiles + for prop in missing: + info[f"hydrocarbon_{prop}_original_linear"] = self._predict(prop, smiles) + self.original_linear_cache_misses += 1 + self._original_linear_cache[key] = dict(info) + return dict(info) + + def _build_smiles(self, terminal: StapleState) -> tuple[str, str, str]: + """Return ``(linear, product_E, product_Z)`` neutral canonical SMILES.""" + pair = None if terminal.anchor_pair is None else tuple(terminal.anchor_pair) + if pair is None: + raise HydrocarbonPropertyScoringError("terminal has no anchor pair") + plan = StapleSmilesPlan.from_tokens(terminal.sequence_tokens, pair) + built_e = build_stapled_smiles( + terminal.sequence_tokens, plan, olefin_geometry=OlefinGeometry.E, strict=True + ) + built_z = build_stapled_smiles( + terminal.sequence_tokens, plan, olefin_geometry=OlefinGeometry.Z, strict=True + ) + return ( + self._canonical_neutral(built_e.linear_smiles), + self._canonical_neutral(built_e.stapled_smiles), + self._canonical_neutral(built_z.stapled_smiles), + ) + + @staticmethod + def _state_key(terminal: StapleState) -> tuple[Any, ...]: + pair = None if terminal.anchor_pair is None else tuple(terminal.anchor_pair) + return (tuple(terminal.sequence_tokens), pair, terminal.block_id, terminal.topology) + + def energy_only_smiles(self, terminal: StapleState) -> tuple[str, str]: + """The E/Z product SMILES the terminal energy needs, nothing else. + + Exposed so a caller can collect SMILES across many plans and prefetch + them through :meth:`StrictSmilesPredictor.predict_property_batch` + before scoring. + """ + _, product_e, product_z = self._build_smiles(terminal) + return product_e, product_z + + def score_energy_only( + self, terminal: StapleState, properties: Sequence[str] + ) -> dict[str, Any]: + """Score only ``properties``, and only on the E/Z products. + + Returns just the ``*_product_mean`` keys that + :func:`property_energy_terms` reads. The linear precursor and the + monitor properties are deliberately not predicted: neither is summed + into the energy, so skipping them leaves it bit-for-bit identical while + removing most of the PeptiVerse work. + """ + key = ("energy_only", tuple(properties), self._state_key(terminal)) + if key in self._cache: + self.cache_hits += 1 + return dict(self._cache[key]) + + product_e, product_z = self.energy_only_smiles(terminal) + info: dict[str, Any] = { + "hydrocarbon_property_status": "scored_energy_only", + "hydrocarbon_property_backend": "peptiverse", + "hydrocarbon_property_mode": "smiles", + "hydrocarbon_property_input_convention": "neutral_canonical", + } + for prop in properties: + e_score = self._predict(prop, product_e) + z_score = self._predict(prop, product_z) + info[f"hydrocarbon_{prop}_product_mean"] = 0.5 * (e_score + z_score) + + self.cache_misses += 1 + self._cache[key] = dict(info) + return info + + def score( + self, terminal: StapleState, properties: Sequence[str] | None = None + ) -> dict[str, Any]: + """Return namespaced linear/E/Z metrics for a terminal state.""" + selected_properties = tuple(properties or self.ALL_PROPERTIES) + key = ( + self._state_key(terminal) + if properties is None + else ("full", selected_properties, self._state_key(terminal)) + ) + if key in self._cache: + self.cache_hits += 1 + return dict(self._cache[key]) + + linear, product_e, product_z = self._build_smiles(terminal) + + info: dict[str, Any] = { + "hydrocarbon_property_status": "scored", + "hydrocarbon_property_backend": "peptiverse", + "hydrocarbon_property_mode": "smiles", + "hydrocarbon_property_input_convention": "neutral_canonical", + "hydrocarbon_linear_smiles": linear, + "hydrocarbon_plan_precursor_smiles": linear, + "hydrocarbon_product_E_smiles": product_e, + "hydrocarbon_product_Z_smiles": product_z, + } + for prop in selected_properties: + linear_score = self._predict(prop, linear) + e_score = self._predict(prop, product_e) + z_score = self._predict(prop, product_z) + mean_score = 0.5 * (e_score + z_score) + prefix = f"hydrocarbon_{prop}" + info.update( + { + f"{prefix}_linear": linear_score, + f"{prefix}_plan_precursor": linear_score, + f"{prefix}_E": e_score, + f"{prefix}_Z": z_score, + f"{prefix}_product_mean": mean_score, + f"{prefix}_E_minus_Z": e_score - z_score, + f"{prefix}_EZ_abs_diff": abs(e_score - z_score), + f"{prefix}_delta_product_linear": mean_score - linear_score, + f"{prefix}_delta_vs_plan_precursor": mean_score - linear_score, + } + ) + + self.cache_misses += 1 + self._cache[key] = dict(info) + return info + + +def property_energy_terms( + scores: dict[str, Any], cfg: HydrocarbonPropertyEnergyConfig +) -> tuple[float, dict[str, Any]]: + """Convert scores to Penetrance objective plus optional toxicity guard.""" + if not cfg.enabled: + return 0.0, { + "hydrocarbon_property_energy": 0.0, + "hydrocarbon_penetrance_energy": 0.0, + "hydrocarbon_toxicity_guard_energy": 0.0, + "hydrocarbon_toxicity_violation": False, + } + + def score(key: str) -> float: + value = float(scores[key]) + if not math.isfinite(value): + raise HydrocarbonPropertyScoringError(f"non-finite property energy input {key}={value}") + return value + + penetrance = score("hydrocarbon_permeability_penetrance_product_mean") + penetrance_energy = float(cfg.penetrance_weight) * (1.0 - penetrance) + if cfg.toxicity_guard_enabled: + # Energy-bearing, so the score must be present regardless of caller. + toxicity = score("hydrocarbon_toxicity_product_mean") + over_threshold = toxicity > float(cfg.toxicity_threshold) + toxicity_energy = float(cfg.toxicity_guard_weight) * max( + 0.0, toxicity - float(cfg.toxicity_threshold) + ) + else: + # Guard off: toxicity contributes 0.0 to the energy, so an energy-only + # caller may legitimately not have predicted it. Report it when it is + # there (unchanged for every full-scoring caller) and stay silent when + # it is not, rather than forcing a prediction the energy never reads. + raw_toxicity = scores.get("hydrocarbon_toxicity_product_mean") + over_threshold = ( + float(raw_toxicity) > float(cfg.toxicity_threshold) + if raw_toxicity is not None + else False + ) + toxicity_energy = 0.0 + violation = bool(cfg.toxicity_guard_enabled and over_threshold) + + developability_terms: dict[str, Any] = {} + developability_energy = 0.0 + if cfg.enable_developability_constraints: + developability_terms = { + "hydrocarbon_developability_constraints_enabled": True, + "hydrocarbon_developability_toxicity_energy": 0.0, + "hydrocarbon_developability_solubility_energy": 0.0, + "hydrocarbon_developability_halflife_energy": 0.0, + "hydrocarbon_developability_energy": 0.0, + } + deltas: dict[str, float] = {} + penalties: dict[str, float] = {} + definitions = { + "toxicity": ( + cfg.developability_lambda_toxicity, + cfg.developability_normalization_toxicity_scale, + 1.0, + ), + "solubility": ( + cfg.developability_lambda_solubility, + cfg.developability_normalization_solubility_scale, + -1.0, + ), + "halflife": ( + cfg.developability_lambda_halflife, + cfg.developability_normalization_halflife_scale, + -1.0, + ), + } + for prop, (weight, scale, adverse_sign) in definitions.items(): + product = score(f"hydrocarbon_{prop}_product_mean") + original = score(f"hydrocarbon_{prop}_original_linear") + delta = product - original + normalized_hinge = max(0.0, adverse_sign * delta) / float(scale) + penalty = float(weight) * normalized_hinge + deltas[prop] = delta + penalties[prop] = penalty + developability_terms[f"hydrocarbon_{prop}_delta_vs_original_lead"] = delta + developability_terms[ + f"hydrocarbon_developability_{prop}_normalized_hinge" + ] = normalized_hinge + developability_terms[ + f"hydrocarbon_developability_{prop}_energy" + ] = penalty + developability_energy = float(sum(penalties.values())) + developability_terms["hydrocarbon_developability_energy"] = ( + developability_energy + ) + + halflife_preservation_terms: dict[str, Any] = {} + halflife_preservation_energy = 0.0 + if cfg.enable_halflife_preservation: + half_product = score("hydrocarbon_halflife_product_mean") + half_original = score("hydrocarbon_halflife_original_linear") + delta_half = half_product - half_original + normalized_hinge = max(0.0, -delta_half) / float( + cfg.halflife_preservation_scale + ) + halflife_preservation_energy = ( + float(cfg.halflife_preservation_lambda) * normalized_hinge + ) + halflife_preservation_terms = { + "hydrocarbon_halflife_preservation_enabled": True, + "hydrocarbon_halflife_delta_vs_original_lead": delta_half, + "hydrocarbon_halflife_preservation_normalized_hinge": normalized_hinge, + "hydrocarbon_halflife_preservation_energy": halflife_preservation_energy, + } + + # This condition is consumed by the plan-level Exact-SB target builder. + # It is deliberately absent from ``total``: inside the allowed set q* + # remains permeability-only, and an empty set falls back exactly. + joint_support_terms: dict[str, Any] = {} + if cfg.enable_joint_perm_halflife_support: + half_product = score("hydrocarbon_halflife_product_mean") + half_original = score("hydrocarbon_halflife_original_linear") + perm_original = score("hydrocarbon_permeability_penetrance_original_linear") + delta_half = half_product - half_original + delta_perm = penetrance - perm_original + joint_support_terms = { + "hydrocarbon_joint_perm_halflife_support_enabled": True, + "hydrocarbon_joint_delta_permeability_vs_original_lead": delta_perm, + "hydrocarbon_joint_delta_halflife_vs_original_lead": delta_half, + "hydrocarbon_joint_perm_halflife_condition": bool( + delta_perm > 0.0 and delta_half >= 0.0 + ), + } + + total = ( + penetrance_energy + + toxicity_energy + + developability_energy + + halflife_preservation_energy + ) + return total, { + "hydrocarbon_property_energy": float(total), + "hydrocarbon_penetrance_energy": float(penetrance_energy), + "hydrocarbon_toxicity_guard_energy": float(toxicity_energy), + "hydrocarbon_toxicity_threshold": float(cfg.toxicity_threshold), + "hydrocarbon_toxicity_over_threshold": bool(over_threshold), + "hydrocarbon_toxicity_violation": violation, + **developability_terms, + **halflife_preservation_terms, + **joint_support_terms, + } diff --git a/staplebridge/hydrocarbon/smiles_builder.py b/staplebridge/hydrocarbon/smiles_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..1b4d8fac11bb3e0e7a8c0eee2e8eba376617cffa --- /dev/null +++ b/staplebridge/hydrocarbon/smiles_builder.py @@ -0,0 +1,1071 @@ +"""Hydrocarbon-stapled product SMILES builder. + +Builds, from a linear sequence plus a staple plan: + + linear precursor the uncyclised peptide with both olefin tethers intact + stapled product the RCM macrocycle, ethene expelled, C=C bond formed + +Both are produced as real molecular graphs. The stapled product is **not** the +linear precursor with a label attached: a carbon-carbon bond is created between +the two tether alkene carbons, the two terminal ``=CH2`` groups are deleted, and +the remaining bond is set to double. :func:`validate_stapled_product` then +re-derives the crosslink from the graph and refuses to certify a molecule whose +macrocycle is not actually closed, so a topology flag alone cannot pass. + +Chemistry +--------- +Ring-closing metathesis joins two terminal alkenes and releases ethene: + + R-CH=CH2 + R'-CH=CH2 -> R-CH=CH-R' + CH2=CH2 + +So each partner contributes its tether minus one carbon, and the resulting +bridge holds ``n_i + n_j - 2`` carbons. The default product is the *E* (trans) +alkene, which is the major RCM product for these staples; ``olefin_geometry`` +can request *Z* or an unspecified double bond instead. + +Isolation +--------- +Additive and hydrocarbon-only. The lactam path has no SMILES builder, so nothing +here overrides existing behaviour, and no lactam module is imported. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Final + +from rdkit import Chem, RDLogger +from rdkit.Chem import Descriptors, rdCIPLabeler, rdMolDescriptors + +from staplebridge.hydrocarbon.catalog import block_topology, is_hydrocarbon_block +from staplebridge.hydrocarbon.monomers import ( + ATTACH_C, + ATTACH_N, + C_TERM_ACID, + C_TERM_AMIDE, + C_TERMINAL_CAPS, + N_TERM_ACETYL, + N_TERM_FREE, + N_TERMINAL_CAPS, + UnknownMonomerError, + expected_macrocycle_size, + get_monomer, + is_anchor_monomer, + staple_carbon_count, +) + +RDLogger.DisableLog("rdApp.*") + +#: Exactly two anchors define a single staple; see +#: :data:`staplebridge.hydrocarbon.actions.REQUIRED_ANCHOR_COUNT`. +REQUIRED_ANCHOR_COUNT: Final[int] = 2 + +#: (pair, spacing) topologies this builder will close, mirroring the enabled +#: catalog. ``R5-S8/i,i+7`` is present but must be opted into explicitly. +SUPPORTED_TOPOLOGIES: Final[frozenset[tuple[str, int]]] = frozenset( + {("S5-S5", 4), ("R8-S5", 7)} +) +OPTIONAL_TOPOLOGIES: Final[frozenset[tuple[str, int]]] = frozenset({("R5-S8", 7)}) + + +class OlefinGeometry(str, Enum): + """Requested geometry of the product double bond.""" + + E = "E" + Z = "Z" + UNSPECIFIED = "unspecified" + + +class BuildFailure(str, Enum): + """Explicit reasons a build or validation was rejected.""" + + OK = "ok" + UNKNOWN_MONOMER = "unknown_monomer" + SEQUENCE_TOO_SHORT = "sequence_too_short" + ANCHOR_OUT_OF_RANGE = "anchor_out_of_range" + ANCHOR_NOT_OLEFINIC = "anchor_not_olefinic" + WRONG_ANCHOR_COUNT = "wrong_anchor_count" + THIRD_ANCHOR_PRESENT = "third_anchor_present" + DOUBLE_STAPLE_UNSUPPORTED = "double_staple_unsupported" + UNSUPPORTED_PAIR_SPACING = "unsupported_pair_spacing_combination" + ANCHOR_TOKEN_MISMATCH = "anchor_token_mismatch" + SANITIZATION_FAILED = "sanitization_failed" + VALENCE_ERROR = "valence_error" + MULTIPLE_FRAGMENTS = "multiple_fragments" + NO_CROSSLINK_BOND = "no_crosslink_bond" + MACROCYCLE_NOT_CLOSED = "macrocycle_not_closed" + UNEXPECTED_RING_SIZE = "unexpected_ring_size" + LINEAR_EQUALS_PRODUCT = "linear_equals_product" + RESIDUAL_TERMINAL_ALKENE = "residual_terminal_alkene" + UNKNOWN_CAP = "unknown_terminal_cap" + + +class SmilesBuildError(ValueError): + """Raised when a hydrocarbon SMILES build cannot be completed. + + Carries the specific :class:`BuildFailure` so callers can distinguish an + illegal request (third anchor, unsupported spacing) from a chemistry bug. + """ + + def __init__(self, reason: BuildFailure, message: str) -> None: + super().__init__(f"{reason.value}: {message}") + self.reason = reason + self.message = message + + +@dataclass +class StapleSmilesPlan: + """The minimum a builder needs: which two positions to join, and with what. + + Deliberately decoupled from + :class:`staplebridge.hydrocarbon.curriculum.HydrocarbonStaplePlan` so the + builder can be driven directly from a sequence in tests and audits, without + constructing a reference-process plan. :meth:`from_hydrocarbon_plan` converts. + """ + + anchor_pair: tuple[int, int] + ordered_pair: str + spacing: int + block_id: str | None = None + + @classmethod + def from_hydrocarbon_plan(cls, plan: Any) -> "StapleSmilesPlan": + """Convert a curriculum/reference ``HydrocarbonStaplePlan``.""" + return cls( + anchor_pair=tuple(plan.anchor_pair), + ordered_pair=plan.ordered_pair, + spacing=int(plan.spacing), + block_id=getattr(plan, "block_id", None), + ) + + @classmethod + def from_block(cls, block: Any, anchor_pair: tuple[int, int]) -> "StapleSmilesPlan": + """Derive the plan from a hydrocarbon catalog block.""" + if not is_hydrocarbon_block(block): + raise SmilesBuildError( + BuildFailure.UNSUPPORTED_PAIR_SPACING, + f"block {getattr(block, 'block_id', block)!r} is not a hydrocarbon block", + ) + pair, spacing = block_topology(block) + return cls( + anchor_pair=tuple(anchor_pair), + ordered_pair=pair, + spacing=spacing, + block_id=block.block_id, + ) + + @classmethod + def from_tokens(cls, tokens: list[str], anchor_pair: tuple[int, int]) -> "StapleSmilesPlan": + """Read the plan off the sequence itself.""" + i, j = int(anchor_pair[0]), int(anchor_pair[1]) + if not (0 <= i < len(tokens)) or not (0 <= j < len(tokens)): + raise SmilesBuildError( + BuildFailure.ANCHOR_OUT_OF_RANGE, + f"anchor pair {(i, j)} outside sequence of length {len(tokens)}", + ) + return cls( + anchor_pair=(i, j), + ordered_pair=f"{tokens[i].upper()}-{tokens[j].upper()}", + spacing=j - i, + ) + + @property + def mode(self) -> str: + """Display label, e.g. ``"S5-S5/i,i+4"``.""" + return f"{self.ordered_pair}/i,i+{self.spacing}" + + +@dataclass +class AtomMapping: + """Where each residue's atoms ended up in the built molecule. + + Indices refer to the *stapled product* unless stated otherwise. This is what + makes the crosslink auditable: the validator re-finds the C=C bond from + :attr:`crosslink_carbons` rather than trusting that the build succeeded. + """ + + #: residue index -> its atom indices in the product + residue_atoms: dict[int, list[int]] = field(default_factory=dict) + #: residue index -> its Cα atom index + alpha_carbons: dict[int, int] = field(default_factory=dict) + #: residue index -> backbone amide N atom index + backbone_nitrogens: dict[int, int] = field(default_factory=dict) + #: residue index -> backbone carbonyl C atom index + carbonyl_carbons: dict[int, int] = field(default_factory=dict) + #: the two carbons joined by the new C-C bond, (i-side, j-side) + crosslink_carbons: tuple[int, int] | None = None + #: every atom of the staple bridge, N-to-C order + staple_atoms: list[int] = field(default_factory=list) + #: atom indices deleted as expelled ethene, in the *linear* numbering + expelled_atoms: list[int] = field(default_factory=list) + + def as_dict(self) -> dict[str, Any]: + """JSON-serialisable view.""" + return { + "residue_atoms": {str(k): v for k, v in sorted(self.residue_atoms.items())}, + "alpha_carbons": {str(k): v for k, v in sorted(self.alpha_carbons.items())}, + "backbone_nitrogens": { + str(k): v for k, v in sorted(self.backbone_nitrogens.items()) + }, + "carbonyl_carbons": { + str(k): v for k, v in sorted(self.carbonyl_carbons.items()) + }, + "crosslink_carbons": ( + list(self.crosslink_carbons) if self.crosslink_carbons else None + ), + "staple_atoms": list(self.staple_atoms), + "expelled_atoms_linear_numbering": list(self.expelled_atoms), + } + + +@dataclass +class StapledSmilesResult: + """Everything one build produced.""" + + sequence_tokens: list[str] + plan: StapleSmilesPlan + linear_smiles: str + stapled_smiles: str + linear_canonical: str + stapled_canonical: str + atom_mapping: AtomMapping + n_terminus: str + c_terminus: str + olefin_geometry: str + #: Per-anchor CIP labels observed in the product. + anchor_cip: dict[int, str | None] = field(default_factory=dict) + linear_formula: str = "" + stapled_formula: str = "" + linear_mw: float = 0.0 + stapled_mw: float = 0.0 + macrocycle_size: int | None = None + expected_macrocycle_size: int | None = None + staple_bridge_carbons: int = 0 + validation: dict[str, Any] = field(default_factory=dict) + + @property + def is_valid(self) -> bool: + """True when every strict validation check passed.""" + return bool(self.validation.get("valid")) + + def as_dict(self) -> dict[str, Any]: + """JSON-serialisable view.""" + return { + "sequence": "-".join(self.sequence_tokens), + "n_residues": len(self.sequence_tokens), + "mode": self.plan.mode, + "anchor_pair": list(self.plan.anchor_pair), + "spacing": self.plan.spacing, + "block_id": self.plan.block_id, + "n_terminus": self.n_terminus, + "c_terminus": self.c_terminus, + "olefin_geometry": self.olefin_geometry, + "linear_smiles": self.linear_smiles, + "stapled_smiles": self.stapled_smiles, + "linear_canonical": self.linear_canonical, + "stapled_canonical": self.stapled_canonical, + "linear_formula": self.linear_formula, + "stapled_formula": self.stapled_formula, + "linear_mw": round(float(self.linear_mw), 4), + "stapled_mw": round(float(self.stapled_mw), 4), + "mw_loss_vs_linear": round(float(self.linear_mw - self.stapled_mw), 4), + "macrocycle_size": self.macrocycle_size, + "expected_macrocycle_size": self.expected_macrocycle_size, + "staple_bridge_carbons": self.staple_bridge_carbons, + "anchor_cip": {str(k): v for k, v in sorted(self.anchor_cip.items())}, + "atom_mapping": self.atom_mapping.as_dict(), + "validation": self.validation, + "linear_differs_from_product": self.linear_canonical != self.stapled_canonical, + } + + +# --------------------------------------------------------------------------- +# Assembly +# --------------------------------------------------------------------------- + + +def _fragment_mol(smiles: str) -> Chem.Mol: + """Parse one fragment without sanitising (dummy atoms defeat valence checks).""" + mol = Chem.MolFromSmiles(smiles, sanitize=False) + if mol is None: + raise SmilesBuildError( + BuildFailure.UNKNOWN_MONOMER, f"fragment SMILES did not parse: {smiles!r}" + ) + return mol + + +def _find_dummy(mol: Chem.Mol, map_number: int) -> int: + """Index of the dummy atom carrying ``map_number``.""" + for atom in mol.GetAtoms(): + if atom.GetAtomicNum() == 0 and atom.GetAtomMapNum() == map_number: + return atom.GetIdx() + raise SmilesBuildError( + BuildFailure.UNKNOWN_MONOMER, + f"fragment has no [*:{map_number}] attachment point", + ) + + +def _validate_request( + tokens: list[str], plan: StapleSmilesPlan, include_optional: bool +) -> None: + """Reject an illegal build request before any atom is placed. + + Every check here is a *hard* failure with a named reason. In particular a + third anchor monomer anywhere in the chain is rejected as a double staple, + which is what stops a silently-wrong molecule from being emitted for a + sequence the catalog cannot express. + """ + if len(tokens) < 2: + raise SmilesBuildError( + BuildFailure.SEQUENCE_TOO_SHORT, + f"need at least 2 residues, got {len(tokens)}", + ) + + for token in tokens: + try: + get_monomer(token) + except UnknownMonomerError as exc: + raise SmilesBuildError(BuildFailure.UNKNOWN_MONOMER, str(exc)) from exc + + i, j = plan.anchor_pair + if not (0 <= i < len(tokens)) or not (0 <= j < len(tokens)): + raise SmilesBuildError( + BuildFailure.ANCHOR_OUT_OF_RANGE, + f"anchor pair {(i, j)} outside sequence of length {len(tokens)}", + ) + if i == j: + raise SmilesBuildError( + BuildFailure.WRONG_ANCHOR_COUNT, f"anchor pair must be distinct, got {(i, j)}" + ) + + # Both named positions must actually carry olefinic anchors. + for position in (i, j): + if not is_anchor_monomer(tokens[position]): + raise SmilesBuildError( + BuildFailure.ANCHOR_NOT_OLEFINIC, + f"position {position} carries {tokens[position]!r}, which has no " + "olefin tether to metathesise", + ) + + # A third anchor would give a double staple, which is out of scope. + anchor_positions = [k for k, t in enumerate(tokens) if is_anchor_monomer(t)] + if len(anchor_positions) != REQUIRED_ANCHOR_COUNT: + reason = ( + BuildFailure.THIRD_ANCHOR_PRESENT + if len(anchor_positions) > REQUIRED_ANCHOR_COUNT + else BuildFailure.WRONG_ANCHOR_COUNT + ) + raise SmilesBuildError( + reason, + f"a single staple needs exactly {REQUIRED_ANCHOR_COUNT} anchor " + f"monomers, found {len(anchor_positions)} at {anchor_positions}", + ) + if set(anchor_positions) != {i, j}: + raise SmilesBuildError( + BuildFailure.DOUBLE_STAPLE_UNSUPPORTED, + f"anchors at {anchor_positions} do not match the requested pair {(i, j)}", + ) + + # The declared pair must match what the sequence says, so a mislabelled plan + # cannot be built as though it were its intended topology. + observed_pair = f"{tokens[i].upper()}-{tokens[j].upper()}" + if observed_pair != plan.ordered_pair: + raise SmilesBuildError( + BuildFailure.ANCHOR_TOKEN_MISMATCH, + f"plan declares {plan.ordered_pair!r} but sequence carries " + f"{observed_pair!r} at {(i, j)}", + ) + if plan.spacing != j - i: + raise SmilesBuildError( + BuildFailure.ANCHOR_TOKEN_MISMATCH, + f"plan declares spacing {plan.spacing} but anchors are {j - i} apart", + ) + + allowed = set(SUPPORTED_TOPOLOGIES) + if include_optional: + allowed |= set(OPTIONAL_TOPOLOGIES) + if (plan.ordered_pair, plan.spacing) not in allowed: + raise SmilesBuildError( + BuildFailure.UNSUPPORTED_PAIR_SPACING, + f"topology {plan.mode} is not supported; enabled: " + f"{sorted(f'{p}/i,i+{s}' for p, s in allowed)}", + ) + + +def _assemble_linear( + tokens: list[str], n_terminus: str, c_terminus: str +) -> tuple[Chem.RWMol, AtomMapping]: + """Build the linear peptide graph, recording per-residue atom indices.""" + if n_terminus not in N_TERMINAL_CAPS: + raise SmilesBuildError( + BuildFailure.UNKNOWN_CAP, + f"unknown N-terminal cap {n_terminus!r}; known: {sorted(N_TERMINAL_CAPS)}", + ) + if c_terminus not in C_TERMINAL_CAPS: + raise SmilesBuildError( + BuildFailure.UNKNOWN_CAP, + f"unknown C-terminal cap {c_terminus!r}; known: {sorted(C_TERMINAL_CAPS)}", + ) + + combined = Chem.RWMol() + mapping = AtomMapping() + # residue index -> (its N-side dummy idx, its C-side dummy idx) in `combined` + dummies: dict[int, tuple[int, int]] = {} + + for residue_index, token in enumerate(tokens): + monomer = get_monomer(token) + fragment = _fragment_mol(monomer.smiles) + offset = combined.GetNumAtoms() + combined.InsertMol(fragment) + + indices = [offset + a for a in range(fragment.GetNumAtoms())] + mapping.residue_atoms[residue_index] = indices + dummies[residue_index] = ( + offset + _find_dummy(fragment, ATTACH_N), + offset + _find_dummy(fragment, ATTACH_C), + ) + + # The backbone is N -> Cα -> C'(=O), so read them off the dummies' + # neighbours rather than by position, which side chains would shift. + n_dummy, c_dummy = dummies[residue_index] + nitrogen = next( + nbr.GetIdx() + for nbr in combined.GetAtomWithIdx(n_dummy).GetNeighbors() + if nbr.GetAtomicNum() == 7 + ) + carbonyl = next( + nbr.GetIdx() + for nbr in combined.GetAtomWithIdx(c_dummy).GetNeighbors() + if nbr.GetAtomicNum() == 6 + ) + alpha = next( + nbr.GetIdx() + for nbr in combined.GetAtomWithIdx(nitrogen).GetNeighbors() + if nbr.GetAtomicNum() == 6 and nbr.GetIdx() != n_dummy + and combined.GetBondBetweenAtoms(nbr.GetIdx(), carbonyl) is not None + ) + mapping.backbone_nitrogens[residue_index] = nitrogen + mapping.carbonyl_carbons[residue_index] = carbonyl + mapping.alpha_carbons[residue_index] = alpha + + # -- peptide bonds --------------------------------------------------- + for residue_index in range(len(tokens) - 1): + upstream_c = mapping.carbonyl_carbons[residue_index] + downstream_n = mapping.backbone_nitrogens[residue_index + 1] + combined.AddBond(upstream_c, downstream_n, Chem.BondType.SINGLE) + + # -- terminal caps --------------------------------------------------- + # Cap atoms belong to no residue of their own, so they are attributed to the + # terminal residue they modify. Leaving them untagged would make + # ``residue_atoms`` an incomplete cover of the molecule, which is exactly the + # kind of silent gap that lets an index-based check drift. + n_cap = N_TERMINAL_CAPS[n_terminus] + if n_cap is not None: + fragment = _fragment_mol(n_cap) + offset = combined.GetNumAtoms() + combined.InsertMol(fragment) + cap_atoms = [offset + a for a in range(fragment.GetNumAtoms())] + cap_dummy = offset + _find_dummy(fragment, ATTACH_N) + cap_anchor = next( + nbr.GetIdx() for nbr in combined.GetAtomWithIdx(cap_dummy).GetNeighbors() + ) + combined.AddBond(cap_anchor, mapping.backbone_nitrogens[0], Chem.BondType.SINGLE) + combined.RemoveAtom(cap_dummy) + mapping.residue_atoms[0].extend(a for a in cap_atoms if a != cap_dummy) + + c_cap = C_TERMINAL_CAPS[c_terminus] + if c_cap is not None: + fragment = _fragment_mol(c_cap) + offset = combined.GetNumAtoms() + combined.InsertMol(fragment) + cap_atoms = [offset + a for a in range(fragment.GetNumAtoms())] + cap_dummy = offset + _find_dummy(fragment, ATTACH_C) + cap_anchor = next( + nbr.GetIdx() for nbr in combined.GetAtomWithIdx(cap_dummy).GetNeighbors() + ) + combined.AddBond( + mapping.carbonyl_carbons[len(tokens) - 1], cap_anchor, Chem.BondType.SINGLE + ) + combined.RemoveAtom(cap_dummy) + mapping.residue_atoms[len(tokens) - 1].extend( + a for a in cap_atoms if a != cap_dummy + ) + + # -- strip the remaining attachment dummies -------------------------- + # Index arithmetic is deliberately avoided here. Cap attachment above + # already removed atoms, so a single "subtract the removals below me" pass + # over the original indices would be wrong — and was, silently, for the + # residues after the first cap. Instead every atom of interest is tagged + # with a durable property before any deletion, and the indices are re-read + # from the finished molecule afterwards. + for residue_index, indices in mapping.residue_atoms.items(): + for atom_index in indices: + atom = combined.GetAtomWithIdx(atom_index) + if atom.GetAtomicNum() != 0: + atom.SetIntProp("residue_index", residue_index) + for residue_index, atom_index in mapping.alpha_carbons.items(): + combined.GetAtomWithIdx(atom_index).SetIntProp("alpha_of", residue_index) + for residue_index, atom_index in mapping.backbone_nitrogens.items(): + combined.GetAtomWithIdx(atom_index).SetIntProp("amide_n_of", residue_index) + for residue_index, atom_index in mapping.carbonyl_carbons.items(): + combined.GetAtomWithIdx(atom_index).SetIntProp("carbonyl_of", residue_index) + + for index in sorted( + (atom.GetIdx() for atom in combined.GetAtoms() if atom.GetAtomicNum() == 0), + reverse=True, + ): + combined.RemoveAtom(index) + + return combined, _mapping_from_props(combined) + + +def _smiles_output_index(mol: Chem.Mol) -> dict[int, int]: + """Map working atom index -> position in the most recent SMILES output. + + ``MolToSmiles`` records its traversal as the private ``_smilesAtomOutputOrder`` + property, whose value at position ``k`` is the working index emitted there. + Inverting it gives what the mapping needs. The property is stored as a string + (``"[2,1,0,...,]"``) and is not surfaced by ``GetPropsAsDict``, so it has to + be read and parsed explicitly. + + Returns an empty dict when the property is absent, in which case the caller + leaves the mapping in working-molecule order rather than corrupting it. + """ + if not mol.HasProp("_smilesAtomOutputOrder"): + return {} + raw = mol.GetProp("_smilesAtomOutputOrder").strip() + inner = raw.strip("[]").rstrip(",") + if not inner: + return {} + try: + order = [int(part) for part in inner.split(",") if part.strip() != ""] + except ValueError: + return {} + return {old: new for new, old in enumerate(order)} + + +def _remap_mapping(mapping: AtomMapping, index_map: dict[int, int]) -> AtomMapping: + """Rewrite every recorded index through ``index_map``. + + Used to convert working-molecule indices into emitted-SMILES order, so the + published mapping is valid against a re-parse of the published string. + """ + + def convert(old: int) -> int: + return index_map.get(old, old) + + return AtomMapping( + residue_atoms={ + residue: sorted(convert(a) for a in atoms) + for residue, atoms in mapping.residue_atoms.items() + }, + alpha_carbons={k: convert(v) for k, v in mapping.alpha_carbons.items()}, + backbone_nitrogens={ + k: convert(v) for k, v in mapping.backbone_nitrogens.items() + }, + carbonyl_carbons={k: convert(v) for k, v in mapping.carbonyl_carbons.items()}, + crosslink_carbons=( + (convert(mapping.crosslink_carbons[0]), convert(mapping.crosslink_carbons[1])) + if mapping.crosslink_carbons + else None + ), + staple_atoms=[convert(a) for a in mapping.staple_atoms], + expelled_atoms=list(mapping.expelled_atoms), + ) + + +def _mapping_from_props(mol: Chem.Mol) -> AtomMapping: + """Rebuild the atom mapping by reading the durable tags off ``mol``. + + Reading the finished molecule is the only reliable way to get these indices: + any bookkeeping that tracks them through a sequence of ``RemoveAtom`` calls + has to model RDKit's renumbering exactly, and gets it wrong as soon as two + removals interleave with the recorded positions. + """ + mapping = AtomMapping() + for atom in mol.GetAtoms(): + index = atom.GetIdx() + if atom.HasProp("residue_index"): + mapping.residue_atoms.setdefault( + atom.GetIntProp("residue_index"), [] + ).append(index) + if atom.HasProp("alpha_of"): + mapping.alpha_carbons[atom.GetIntProp("alpha_of")] = index + if atom.HasProp("amide_n_of"): + mapping.backbone_nitrogens[atom.GetIntProp("amide_n_of")] = index + if atom.HasProp("carbonyl_of"): + mapping.carbonyl_carbons[atom.GetIntProp("carbonyl_of")] = index + return mapping + + +def _terminal_alkene_carbons( + mol: Chem.Mol, residue_atoms: list[int], alpha_carbon: int +) -> tuple[int, int]: + """Locate one anchor's tether alkene, as ``(inner_CH, terminal_CH2)``. + + Identified structurally rather than by index arithmetic: the terminal ``CH2`` + is the alkene carbon with exactly one heavy neighbour, and it must sit on the + side chain of this residue, reachable from Cα without crossing the backbone. + """ + candidates: list[tuple[int, int]] = [] + owned = set(residue_atoms) + for bond in mol.GetBonds(): + if bond.GetBondType() != Chem.BondType.DOUBLE: + continue + begin, end = bond.GetBeginAtom(), bond.GetEndAtom() + if begin.GetAtomicNum() != 6 or end.GetAtomicNum() != 6: + continue + if begin.GetIdx() not in owned or end.GetIdx() not in owned: + continue + for inner, terminal in ((begin, end), (end, begin)): + heavy = [n for n in terminal.GetNeighbors() if n.GetAtomicNum() > 1] + if len(heavy) == 1 and terminal.GetTotalNumHs() == 2: + candidates.append((inner.GetIdx(), terminal.GetIdx())) + + if not candidates: + raise SmilesBuildError( + BuildFailure.ANCHOR_NOT_OLEFINIC, + f"no terminal alkene found on the residue whose Cα is atom {alpha_carbon}", + ) + if len(candidates) > 1: + raise SmilesBuildError( + BuildFailure.ANCHOR_NOT_OLEFINIC, + f"ambiguous tether: {len(candidates)} terminal alkenes on the residue " + f"whose Cα is atom {alpha_carbon}", + ) + return candidates[0] + + +def _bridge_atoms(mol: Chem.Mol, start: int, end: int, blocked: set[int]) -> list[int]: + """Shortest path from ``start`` to ``end`` avoiding ``blocked`` atoms.""" + from collections import deque + + queue = deque([[start]]) + seen = {start} + while queue: + path = queue.popleft() + if path[-1] == end: + return path + for neighbour in mol.GetAtomWithIdx(path[-1]).GetNeighbors(): + index = neighbour.GetIdx() + if index in seen or index in blocked: + continue + seen.add(index) + queue.append(path + [index]) + return [] + + +# --------------------------------------------------------------------------- +# Public builder +# --------------------------------------------------------------------------- + + +def build_stapled_smiles( + sequence_tokens: list[str], + plan: StapleSmilesPlan, + n_terminus: str = N_TERM_FREE, + c_terminus: str = C_TERM_ACID, + olefin_geometry: OlefinGeometry | str = OlefinGeometry.E, + include_optional_topologies: bool = False, + strict: bool = True, +) -> StapledSmilesResult: + """Build the linear precursor and the RCM-stapled product. + + Args: + sequence_tokens: monomer tokens, e.g. ``["A", "S5", "L", "K", "A", "S5"]``. + plan: which positions to staple, and with which topology. + n_terminus: ``"free_amine"`` or ``"acetyl"``. + c_terminus: ``"free_acid"`` or ``"amide"``. + olefin_geometry: geometry of the product double bond. + include_optional_topologies: also allow ``R5-S8/i,i+7``. + strict: raise when validation fails, instead of returning an invalid + result with the reasons recorded. + + Returns: + A :class:`StapledSmilesResult` whose ``stapled_smiles`` is a genuinely + cyclised molecule. + + Raises: + SmilesBuildError: for an illegal request, or (when ``strict``) a product + that fails validation. + """ + tokens = [t.upper() for t in sequence_tokens] + geometry = OlefinGeometry(olefin_geometry) + _validate_request(tokens, plan, include_optional_topologies) + + i, j = plan.anchor_pair + linear_rw, mapping = _assemble_linear(tokens, n_terminus, c_terminus) + + linear_mol = linear_rw.GetMol() + try: + Chem.SanitizeMol(linear_mol) + except Exception as exc: + raise SmilesBuildError( + BuildFailure.SANITIZATION_FAILED, + f"linear precursor failed sanitization: {type(exc).__name__}: {exc}", + ) from exc + + inner_i, terminal_i = _terminal_alkene_carbons( + linear_mol, mapping.residue_atoms[i], mapping.alpha_carbons[i] + ) + inner_j, terminal_j = _terminal_alkene_carbons( + linear_mol, mapping.residue_atoms[j], mapping.alpha_carbons[j] + ) + + # -- ring-closing metathesis ----------------------------------------- + # Join the two inner alkene carbons, delete both terminal CH2 groups as the + # expelled ethene, and make the surviving bond a double bond. + product = Chem.RWMol(linear_mol) + product.AddBond(inner_i, inner_j, Chem.BondType.DOUBLE) + # Tag the crosslink carbons before deleting anything, for the same reason the + # linear assembly does: the deletions renumber the atoms above them. + product.GetAtomWithIdx(inner_i).SetIntProp("crosslink_side", 0) + product.GetAtomWithIdx(inner_j).SetIntProp("crosslink_side", 1) + for index in sorted((terminal_i, terminal_j), reverse=True): + product.RemoveAtom(index) + + product_mapping = _mapping_from_props(product) + crosslink: dict[int, int] = {} + for atom in product.GetAtoms(): + if atom.HasProp("crosslink_side"): + crosslink[atom.GetIntProp("crosslink_side")] = atom.GetIdx() + if set(crosslink) != {0, 1}: + raise SmilesBuildError( + BuildFailure.NO_CROSSLINK_BOND, + f"crosslink carbons were lost during ethene removal: found {crosslink}", + ) + product_mapping.crosslink_carbons = (crosslink[0], crosslink[1]) + product_mapping.expelled_atoms = sorted((terminal_i, terminal_j)) + + stapled_mol = product.GetMol() + try: + Chem.SanitizeMol(stapled_mol) + except Exception as exc: + raise SmilesBuildError( + BuildFailure.SANITIZATION_FAILED, + f"stapled product failed sanitization: {type(exc).__name__}: {exc}", + ) from exc + + # Re-read the mapping from the *sanitized* molecule. The tags survive both + # ``GetMol()`` and sanitization, and this is the molecule whose numbering the + # emitted SMILES will carry, so reading here is what keeps the published + # indices valid for a caller who re-parses the string. + product_mapping = _mapping_from_props(stapled_mol) + crosslink = {} + for atom in stapled_mol.GetAtoms(): + if atom.HasProp("crosslink_side"): + crosslink[atom.GetIntProp("crosslink_side")] = atom.GetIdx() + if set(crosslink) != {0, 1}: + raise SmilesBuildError( + BuildFailure.NO_CROSSLINK_BOND, + f"crosslink carbons were lost during sanitization: found {crosslink}", + ) + product_mapping.crosslink_carbons = (crosslink[0], crosslink[1]) + product_mapping.expelled_atoms = sorted((terminal_i, terminal_j)) + + # -- olefin geometry ------------------------------------------------- + cross_i, cross_j = product_mapping.crosslink_carbons + bond = stapled_mol.GetBondBetweenAtoms(cross_i, cross_j) + if bond is None: + raise SmilesBuildError( + BuildFailure.NO_CROSSLINK_BOND, + f"crosslink bond between atoms {cross_i} and {cross_j} is absent after build", + ) + if geometry is OlefinGeometry.UNSPECIFIED: + bond.SetStereo(Chem.BondStereo.STEREONONE) + else: + # Stereo atoms must be named for the parity to mean anything: pick the + # ring-side neighbour on each end, so E/Z refers to the macrocycle. + ref_i = next( + ( + n.GetIdx() + for n in stapled_mol.GetAtomWithIdx(cross_i).GetNeighbors() + if n.GetIdx() != cross_j + ), + None, + ) + ref_j = next( + ( + n.GetIdx() + for n in stapled_mol.GetAtomWithIdx(cross_j).GetNeighbors() + if n.GetIdx() != cross_i + ), + None, + ) + if ref_i is not None and ref_j is not None: + bond.SetStereoAtoms(ref_i, ref_j) + bond.SetStereo( + Chem.BondStereo.STEREOE + if geometry is OlefinGeometry.E + else Chem.BondStereo.STEREOZ + ) + # Setting the stereo descriptor alone is not enough for a double bond + # inside a ring: RDKit stores it on the bond but omits the ``/`` + # ``\`` markers when writing SMILES, so E and Z would serialise to + # the identical string. This call materialises the neighbour bond + # directions that the writer actually reads. + Chem.SetDoubleBondNeighborDirections(stapled_mol) + + if geometry is OlefinGeometry.UNSPECIFIED: + Chem.AssignStereochemistry(stapled_mol, cleanIt=True, force=True) + else: + # cleanIt=True would discard the ring-bond stereo just installed, so + # only assign what is missing. + Chem.AssignStereochemistry(stapled_mol, cleanIt=False, force=False) + + # -- staple bridge atoms, read off the graph ------------------------- + backbone_block = set(product_mapping.backbone_nitrogens.values()) | set( + product_mapping.carbonyl_carbons.values() + ) + bridge = _bridge_atoms( + stapled_mol, + product_mapping.alpha_carbons[i], + product_mapping.alpha_carbons[j], + blocked=backbone_block, + ) + product_mapping.staple_atoms = bridge + + # CIP labels and ring sizes are read while the mapping still holds working + # indices, since both queries address ``stapled_mol``. + anchor_cip: dict[int, str | None] = {} + rdCIPLabeler.AssignCIPLabels(stapled_mol) + for position in (i, j): + atom = stapled_mol.GetAtomWithIdx(product_mapping.alpha_carbons[position]) + anchor_cip[position] = ( + atom.GetProp("_CIPCode") if atom.HasProp("_CIPCode") else None + ) + + ring_sizes = [ + len(ring) + for ring in stapled_mol.GetRingInfo().AtomRings() + if cross_i in ring and cross_j in ring + ] + + linear_smiles = Chem.MolToSmiles(linear_mol, isomericSmiles=True) + stapled_smiles = Chem.MolToSmiles(stapled_mol, isomericSmiles=True) + + # Re-express the mapping in the *emitted SMILES* atom order. + # + # ``MolToSmiles`` chooses its own traversal root and order, so an index that + # is correct for ``stapled_mol`` points at an unrelated atom once a caller + # does ``MolFromSmiles(stapled_smiles)``. That silent mismatch is worse than + # useless for an atom-mapping API, so the published indices are converted + # into output order here using RDKit's own output-order record, and the + # tests assert them against a freshly re-parsed molecule. + output_index = _smiles_output_index(stapled_mol) + if output_index: + product_mapping = _remap_mapping(product_mapping, output_index) + + result = StapledSmilesResult( + sequence_tokens=tokens, + plan=plan, + linear_smiles=linear_smiles, + stapled_smiles=stapled_smiles, + linear_canonical=Chem.MolToSmiles(linear_mol, isomericSmiles=False), + stapled_canonical=Chem.MolToSmiles(stapled_mol, isomericSmiles=False), + atom_mapping=product_mapping, + n_terminus=n_terminus, + c_terminus=c_terminus, + olefin_geometry=geometry.value, + anchor_cip=anchor_cip, + linear_formula=rdMolDescriptors.CalcMolFormula(linear_mol), + stapled_formula=rdMolDescriptors.CalcMolFormula(stapled_mol), + linear_mw=float(Descriptors.MolWt(linear_mol)), + stapled_mw=float(Descriptors.MolWt(stapled_mol)), + macrocycle_size=min(ring_sizes) if ring_sizes else None, + expected_macrocycle_size=expected_macrocycle_size( + plan.spacing, tokens[i], tokens[j] + ), + staple_bridge_carbons=staple_carbon_count(tokens[i], tokens[j]), + ) + # Validated against the re-parsed product, not the working molecule: the + # published mapping is in emitted-SMILES order, and checking it against the + # string a caller would actually receive is the only check that proves the + # published indices are usable. + result.validation = validate_stapled_product(result) + + if strict and not result.is_valid: + raise SmilesBuildError( + BuildFailure(result.validation["first_failure"]), + f"stapled product failed validation: {result.validation['failures']}", + ) + return result + + +def build_from_sequence( + sequence: str, + anchor_pair: tuple[int, int] | None = None, + **kwargs: Any, +) -> StapledSmilesResult: + """Convenience wrapper: tokenize ``sequence`` and staple its two anchors. + + When ``anchor_pair`` is omitted the two anchor monomers present in the + sequence are used, which is what the StaPep audit needs. + """ + from staplebridge.hydrocarbon.tokenizer import tokenize_sequence + + tokens = tokenize_sequence(sequence) + if anchor_pair is None: + positions = [k for k, t in enumerate(tokens) if is_anchor_monomer(t)] + if len(positions) != REQUIRED_ANCHOR_COUNT: + raise SmilesBuildError( + BuildFailure.WRONG_ANCHOR_COUNT, + f"sequence {sequence!r} has {len(positions)} anchor monomers at " + f"{positions}; need exactly {REQUIRED_ANCHOR_COUNT}", + ) + anchor_pair = (positions[0], positions[1]) + plan = StapleSmilesPlan.from_tokens(tokens, anchor_pair) + return build_stapled_smiles(tokens, plan, **kwargs) + + +# --------------------------------------------------------------------------- +# Strict validation +# --------------------------------------------------------------------------- + + +def validate_stapled_product( + result: StapledSmilesResult, + linear_mol: Chem.Mol | None = None, + stapled_mol: Chem.Mol | None = None, +) -> dict[str, Any]: + """Strictly validate a built product, re-deriving the crosslink from the graph. + + Checks, each recorded by name: + + 1. both SMILES re-parse and sanitize; + 2. no valence errors; + 3. each is a single connected molecule; + 4. exactly two anchor residues, and the crosslink joins *their* tethers; + 5. the crosslink bond exists, is C-C, and is a double bond; + 6. the crosslink lies on a ring, and that ring has the arithmetically + expected size; + 7. no terminal ``CH2=`` remains (both tethers were consumed); + 8. the linear and stapled canonical SMILES differ. + + Check 6 is what makes a topology label insufficient: the ring is found in the + re-parsed product graph, so a molecule that was merely flagged as stapled has + no ring to find and fails here. + """ + failures: list[str] = [] + details: dict[str, Any] = {} + + reparsed_linear = Chem.MolFromSmiles(result.linear_smiles) + reparsed_stapled = Chem.MolFromSmiles(result.stapled_smiles) + details["linear_reparsed"] = reparsed_linear is not None + details["stapled_reparsed"] = reparsed_stapled is not None + if reparsed_linear is None or reparsed_stapled is None: + failures.append(BuildFailure.SANITIZATION_FAILED.value) + return _verdict(failures, details) + + for label, mol in (("linear", reparsed_linear), ("stapled", reparsed_stapled)): + problems = Chem.DetectChemistryProblems(mol) + if problems: + failures.append(BuildFailure.VALENCE_ERROR.value) + details[f"{label}_chemistry_problems"] = [p.Message() for p in problems] + + linear_fragments = len(Chem.GetMolFrags(reparsed_linear)) + stapled_fragments = len(Chem.GetMolFrags(reparsed_stapled)) + details["linear_fragment_count"] = linear_fragments + details["stapled_fragment_count"] = stapled_fragments + if linear_fragments != 1 or stapled_fragments != 1: + failures.append(BuildFailure.MULTIPLE_FRAGMENTS.value) + + # -- anchors --------------------------------------------------------- + anchor_positions = [ + k for k, t in enumerate(result.sequence_tokens) if is_anchor_monomer(t) + ] + details["anchor_positions"] = anchor_positions + if len(anchor_positions) != REQUIRED_ANCHOR_COUNT: + failures.append(BuildFailure.WRONG_ANCHOR_COUNT.value) + + # -- crosslink bond, in the working molecule ------------------------- + working = stapled_mol if stapled_mol is not None else reparsed_stapled + crosslink = result.atom_mapping.crosslink_carbons + details["crosslink_carbons"] = list(crosslink) if crosslink else None + if crosslink is None: + failures.append(BuildFailure.NO_CROSSLINK_BOND.value) + return _verdict(failures, details) + + cross_i, cross_j = crosslink + bond = working.GetBondBetweenAtoms(cross_i, cross_j) + if bond is None: + failures.append(BuildFailure.NO_CROSSLINK_BOND.value) + return _verdict(failures, details) + + both_carbon = ( + working.GetAtomWithIdx(cross_i).GetAtomicNum() == 6 + and working.GetAtomWithIdx(cross_j).GetAtomicNum() == 6 + ) + details["crosslink_is_carbon_carbon"] = both_carbon + details["crosslink_bond_type"] = str(bond.GetBondType()) + if not both_carbon: + failures.append(BuildFailure.NO_CROSSLINK_BOND.value) + if bond.GetBondType() != Chem.BondType.DOUBLE: + failures.append(BuildFailure.NO_CROSSLINK_BOND.value) + + # -- the crosslink must close a real macrocycle ----------------------- + details["crosslink_in_ring"] = bool(bond.IsInRing()) + if not bond.IsInRing(): + failures.append(BuildFailure.MACROCYCLE_NOT_CLOSED.value) + else: + rings = [ + len(ring) + for ring in working.GetRingInfo().AtomRings() + if cross_i in ring and cross_j in ring + ] + observed = min(rings) if rings else None + details["macrocycle_size"] = observed + details["expected_macrocycle_size"] = result.expected_macrocycle_size + if observed != result.expected_macrocycle_size: + failures.append(BuildFailure.UNEXPECTED_RING_SIZE.value) + + # -- both tethers consumed ------------------------------------------- + residual = working.GetSubstructMatches(Chem.MolFromSmarts("[CX3H2]=[CX3]")) + details["residual_terminal_alkenes"] = len(residual) + if residual: + failures.append(BuildFailure.RESIDUAL_TERMINAL_ALKENE.value) + + # -- the product is genuinely not the precursor ---------------------- + differs = result.linear_canonical != result.stapled_canonical + details["linear_differs_from_product"] = differs + if not differs: + failures.append(BuildFailure.LINEAR_EQUALS_PRODUCT.value) + + # Ethene (C2H4, 28.05) leaves, so the product must be lighter. + details["mw_loss_vs_linear"] = round(result.linear_mw - result.stapled_mw, 4) + return _verdict(failures, details) + + +def _verdict(failures: list[str], details: dict[str, Any]) -> dict[str, Any]: + """Package a validation outcome.""" + unique = sorted(set(failures)) + return { + "valid": not unique, + "failures": unique, + "first_failure": unique[0] if unique else BuildFailure.OK.value, + "n_checks_failed": len(unique), + **details, + } + + +def describe_builder() -> dict[str, Any]: + """Summary of builder capabilities, for reports.""" + return { + "supported_topologies": sorted(f"{p}/i,i+{s}" for p, s in SUPPORTED_TOPOLOGIES), + "optional_topologies": sorted(f"{p}/i,i+{s}" for p, s in OPTIONAL_TOPOLOGIES), + "n_terminal_caps": sorted(N_TERMINAL_CAPS), + "c_terminal_caps": sorted(C_TERMINAL_CAPS), + "olefin_geometries": [g.value for g in OlefinGeometry], + "rcm_expels": "ethene (C2H4)", + "validation_checks": [ + "sanitization", + "valence", + "single_fragment", + "exactly_two_anchors", + "crosslink_bond_is_CC_double", + "crosslink_closes_macrocycle_of_expected_size", + "no_residual_terminal_alkene", + "linear_differs_from_product", + ], + } diff --git a/staplebridge/hydrocarbon/terminal_energy.py b/staplebridge/hydrocarbon/terminal_energy.py new file mode 100644 index 0000000000000000000000000000000000000000..d2f9f87332d7639cc0021c5a38442478690c193b --- /dev/null +++ b/staplebridge/hydrocarbon/terminal_energy.py @@ -0,0 +1,227 @@ +"""Hydrocarbon terminal energy. + +``BridgeTrainer.terminal_energy`` is the lactam terminal energy and is **not +modified**. This module wraps it additively for the hydrocarbon branch: + + E_hydrocarbon = E_base + E_endpoint_prior + E_hydrocarbon_property + +``E_base`` is computed by delegating to the caller's existing terminal-energy +callable, so the geometry, edit-distance, property-penalty and cost terms behave +exactly as they do today. Hydrocarbon-specific endpoint-prior and strict SMILES property terms are both +optional. With neither enabled the total degrades to ``E_base`` bit-for-bit. +The property scorer is never imported or constructed by the lactam path. + +Because this wrapper is only constructed on the hydrocarbon path, the lactam +terminal energy is unchanged whether the prior is on or off. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, Protocol + +from staplebridge.chemistry.state import StapleState +from staplebridge.data.schemas import BuildingBlock, LeadExample +from staplebridge.hydrocarbon.actions import ( + FailureReason, + validate_hydrocarbon_staple, +) +from staplebridge.hydrocarbon.endpoint_prior import ( + EmpiricalHydrocarbonEndpointPrior, + EndpointPriorConfig, +) +from staplebridge.hydrocarbon.property_energy import ( + full_scoring_properties, + HydrocarbonPropertyEnergyConfig, + HydrocarbonPropertyScorer, + property_energy_terms, + required_energy_properties, + required_original_lead_properties, +) + + +class TerminalEnergyFn(Protocol): + """The signature of ``BridgeTrainer.terminal_energy``.""" + + def __call__( + self, z0: StapleState, zt: StapleState, lead: LeadExample + ) -> tuple[float, dict[str, Any]]: + ... + + +@dataclass +class HydrocarbonTerminalEnergyConfig: + """Config for the hydrocarbon terminal energy wrapper.""" + + #: Penalty when a terminal state is not a legal hydrocarbon staple. Applied + #: on top of whatever the base energy already charges for infeasibility. + invalid_topology_penalty: float = 10.0 + #: Also charge the penalty when the terminal state never got stapled. + penalize_unstapled: bool = True + + +class HydrocarbonTerminalEnergy: + """Additive hydrocarbon terminal energy. + + Args: + base_terminal_energy: the existing (lactam-style) terminal energy + callable, typically ``trainer.terminal_energy``. Delegated to + unchanged. + catalog: hydrocarbon blocks, used to validate the terminal topology. + endpoint_prior: the empirical prior. When ``None`` or disabled, this + class adds nothing to the base energy. + config: wrapper configuration. + """ + + def __init__( + self, + base_terminal_energy: TerminalEnergyFn | Callable[..., tuple[float, dict[str, Any]]], + catalog: list[BuildingBlock], + endpoint_prior: EmpiricalHydrocarbonEndpointPrior | None = None, + config: HydrocarbonTerminalEnergyConfig | None = None, + property_scorer: HydrocarbonPropertyScorer | None = None, + property_config: HydrocarbonPropertyEnergyConfig | None = None, + ) -> None: + self._base = base_terminal_energy + self.catalog = list(catalog) + self.catalog_index = {b.block_id: b for b in self.catalog} + self.endpoint_prior = endpoint_prior or EmpiricalHydrocarbonEndpointPrior() + self.cfg = config or HydrocarbonTerminalEnergyConfig() + self.property_scorer = property_scorer + self.property_cfg = property_config or HydrocarbonPropertyEnergyConfig() + + def __call__( + self, z0: StapleState, zt: StapleState, lead: LeadExample, *, energy_only: bool = False + ) -> tuple[float, dict[str, Any]]: + """Compute the hydrocarbon terminal energy for ``zt``. + + With ``energy_only=True`` the property scorer predicts only the + properties that reach the energy (see + :func:`required_energy_properties`), so the returned ``info`` carries + just those and the monitors are absent. The float energy is unchanged. + Callers that read ``info`` for logging, validation or ranking must + leave ``energy_only`` at its default. + """ + base_energy, info = self._base(z0, zt, lead) + energy = float(base_energy) + info = dict(info) + info["chemistry"] = "hydrocarbon" + info["base_terminal_energy"] = float(base_energy) + + block = self.catalog_index.get(zt.block_id) if zt.block_id else None + verdict = validate_hydrocarbon_staple( + zt.sequence_tokens, + None if zt.anchor_pair is None else tuple(zt.anchor_pair), + block, + self.catalog, + ) + info["hydrocarbon_topology_status"] = verdict.value + + topology_ok = verdict is FailureReason.OK + stapled = zt.topology == "stapled" + penalty = 0.0 + if not topology_ok or (self.cfg.penalize_unstapled and not stapled): + penalty = float(self.cfg.invalid_topology_penalty) + energy += penalty + info["hydrocarbon_invalid_topology_penalty"] = penalty + + breakdown = self.endpoint_prior.score_endpoint(zt, block) + energy += float(breakdown.total) + info.update(breakdown.as_dict()) + + info["hydrocarbon_endpoint_pass"] = bool(topology_ok and stapled) + + if self.property_scorer is not None and topology_ok and stapled: + if energy_only: + property_scores = self.property_scorer.score_energy_only( + zt, required_energy_properties(self.property_cfg) + ) + else: + # Preserve the historical full-scoring call exactly when the + # total switch is off; only the developability arm adds the + # solubility head to validation/ranking records. + if self.property_cfg.enable_developability_constraints: + property_scores = self.property_scorer.score( + zt, full_scoring_properties(self.property_cfg) + ) + else: + property_scores = self.property_scorer.score(zt) + if ( + self.property_cfg.enable_developability_constraints + or self.property_cfg.enable_halflife_preservation + or self.property_cfg.enable_joint_perm_halflife_support + ): + lead_key = ( + str(getattr(lead, "example_id")) + if getattr(lead, "example_id", None) is not None + else None + ) + original_scores = self.property_scorer.score_original_linear( + z0.sequence_tokens, + lead_key=lead_key, + properties=required_original_lead_properties(self.property_cfg), + ) + property_scores.update(original_scores) + info.update(property_scores) + property_energy, property_terms = property_energy_terms( + property_scores, self.property_cfg + ) + energy += float(property_energy) + info.update(property_terms) + else: + reason = ( + "disabled" + if self.property_scorer is None + else f"not_scored:{verdict.value}" + ) + info["hydrocarbon_property_status"] = reason + info["hydrocarbon_property_energy"] = 0.0 + info["hydrocarbon_toxicity_violation"] = False + + info["terminal_energy"] = float(energy) + return float(energy), info + + def describe(self) -> dict[str, Any]: + """Summary for logging and audits.""" + return { + "chemistry": "hydrocarbon", + "catalog_blocks": [b.block_id for b in self.catalog], + "invalid_topology_penalty": float(self.cfg.invalid_topology_penalty), + "endpoint_prior": self.endpoint_prior.describe(), + "property_energy": self.property_cfg.describe(), + } + + +def build_hydrocarbon_terminal_energy( + base_terminal_energy: TerminalEnergyFn, + catalog: list[BuildingBlock], + hydrocarbon_cfg: dict[str, Any] | None, + root: Any = None, + property_scorer: HydrocarbonPropertyScorer | None = None, +) -> HydrocarbonTerminalEnergy: + """Build the wrapper from a ``hydrocarbon`` config section. + + The prior is only constructed as enabled when the config says so, so a + config without an ``endpoint_prior`` block yields base-energy behaviour. + """ + section = dict(hydrocarbon_cfg or {}) + prior_cfg = EndpointPriorConfig.from_dict(section.get("endpoint_prior")) + prior = EmpiricalHydrocarbonEndpointPrior(prior_cfg, root=root) + + terminal_section = dict(section.get("terminal_energy") or {}) + cfg = HydrocarbonTerminalEnergyConfig() + for key, value in terminal_section.items(): + if hasattr(cfg, key): + current = getattr(cfg, key) + setattr(cfg, key, bool(value) if isinstance(current, bool) else float(value)) + + return HydrocarbonTerminalEnergy( + base_terminal_energy=base_terminal_energy, + catalog=catalog, + endpoint_prior=prior, + config=cfg, + property_scorer=property_scorer, + property_config=HydrocarbonPropertyEnergyConfig.from_dict( + terminal_section.get("property") + ), + ) diff --git a/staplebridge/hydrocarbon/tokenizer.py b/staplebridge/hydrocarbon/tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..ec348e31274046b728141cacda76482031debdf8 --- /dev/null +++ b/staplebridge/hydrocarbon/tokenizer.py @@ -0,0 +1,210 @@ +"""Multi-character monomer tokenizer for the hydrocarbon branch. + +The lactam path builds states with ``StapleState.from_sequence``, which is +``list(seq)`` - one character per residue. That is correct for K/D/E lactam and +is left untouched. Hydrocarbon anchors are two characters (``S5``, ``R8``), so +this module provides a *separate* tokenizing constructor. Lactam never calls it. + +Model-vocabulary projection +--------------------------- +``staplebridge.data.vocab.ALL_TOKENS`` has 24 entries and +:class:`~staplebridge.models.embeddings.TokenMLPEncoder` sizes its embedding as +``nn.Embedding(len(TOKEN_TO_ID), ...)``. Appending ``S5``/``R8`` to that list +would change the embedding matrix shape and break loading of every existing +checkpoint. So this module does **not** touch the vocab. Instead it projects +hydrocarbon monomers onto the ncAA tokens the vocab already carries: + + S3, S5, S8 -> "X" (X is already the StaPep alias for S5) + R3, R5, R8 -> "B" + Aib -> "X" + Nle -> "B" + +The projection is only applied where a tensor of token ids is needed. The +authoritative state keeps the true monomer tokens, so anchor typing, catalog +matching and the endpoint prior all see ``S5``/``R8`` exactly. +""" + +from __future__ import annotations + +from typing import Final + +from staplebridge.chemistry.state import StapleState +from staplebridge.data.vocab import TOKEN_TO_ID + +NATURAL_AA: Final[frozenset[str]] = frozenset("ACDEFGHIKLMNPQRSTVWY") + +#: Hydrocarbon anchor monomers recognised by the tokenizer. Longest-first so a +#: greedy scan consumes ``S5`` before it can mistake ``S`` for serine. +HYDROCARBON_ANCHOR_TOKENS: Final[tuple[str, ...]] = ( + "S3", + "S5", + "S8", + "R3", + "R5", + "R8", +) + +#: Non-anchor non-natural monomers the tokenizer accepts as whole segments. +OTHER_MONOMER_TOKENS: Final[tuple[str, ...]] = ("Aib", "AIB", "Nle", "NLE") + +#: Terminal modifications: no residue index, no contribution to length. +N_TERMINAL_MODS: Final[frozenset[str]] = frozenset({"AC"}) +C_TERMINAL_MODS: Final[frozenset[str]] = frozenset({"NH2"}) + +#: Projection onto tokens the existing model vocabulary already contains. See +#: the module docstring for why the vocab itself is not extended. +_MODEL_TOKEN_PROJECTION: Final[dict[str, str]] = { + "S3": "X", + "S5": "X", + "S8": "X", + "R3": "B", + "R5": "B", + "R8": "B", + "AIB": "X", + "NLE": "B", +} + +_MULTI_CHAR: Final[tuple[str, ...]] = tuple( + sorted( + {t.upper() for t in HYDROCARBON_ANCHOR_TOKENS + OTHER_MONOMER_TOKENS}, + key=len, + reverse=True, + ) +) + + +class HydrocarbonTokenizationError(ValueError): + """Raised when a sequence cannot be fully consumed into known monomers.""" + + +def is_anchor_token(token: str) -> bool: + """True when ``token`` is a hydrocarbon staple anchor monomer.""" + return token.upper() in {t.upper() for t in HYDROCARBON_ANCHOR_TOKENS} + + +def normalize_monomer(token: str) -> str: + """Canonicalise one monomer token (``s5`` -> ``S5``, ``Aib`` -> ``AIB``).""" + upper = token.upper() + if upper in {t.upper() for t in HYDROCARBON_ANCHOR_TOKENS}: + return upper + if upper in {"AIB", "NLE"}: + return upper + return upper + + +def tokenize_sequence(sequence: str) -> list[str]: + """Tokenize a hydrocarbon-style sequence into monomer tokens. + + Supports dash-delimited segments, undelimited runs and mixtures, plus + ``Ac-``/``-NH2`` terminal modifications (which are dropped from the residue + list because they carry no residue index). + + Args: + sequence: e.g. ``"TSFR8EYWALLS5"``, ``"Ac-ISF-R8-ELLDYY-S5-ESGS"``. + + Returns: + One canonical monomer token per residue. + + Raises: + HydrocarbonTokenizationError: if any part cannot be consumed. Nothing is + guessed. + """ + if sequence is None or not str(sequence).strip(): + raise HydrocarbonTokenizationError("empty sequence") + + text = str(sequence).strip().strip("-") + segments = [s for s in text.split("-") if s] + tokens: list[str] = [] + + for position, segment in enumerate(segments): + upper = segment.upper() + if upper in N_TERMINAL_MODS and position == 0: + continue + if upper in C_TERMINAL_MODS and position == len(segments) - 1: + continue + if upper in _MULTI_CHAR: + tokens.append(normalize_monomer(upper)) + continue + tokens.extend(_tokenize_run(segment)) + + if not tokens: + raise HydrocarbonTokenizationError( + f"sequence {sequence!r} contained no residues" + ) + return tokens + + +def _tokenize_run(run: str) -> list[str]: + """Longest-match scan of one undelimited run.""" + upper = run.upper() + tokens: list[str] = [] + index = 0 + while index < len(upper): + for candidate in _MULTI_CHAR: + if upper.startswith(candidate, index): + tokens.append(normalize_monomer(candidate)) + index += len(candidate) + break + else: + char = upper[index] + if char in NATURAL_AA: + tokens.append(char) + index += 1 + else: + raise HydrocarbonTokenizationError( + f"unrecognized character {char!r} at offset {index} of {run!r}" + ) + return tokens + + +def state_from_sequence(sequence: str, **kwargs: object) -> StapleState: + """Build a :class:`StapleState` whose tokens are hydrocarbon monomers. + + This is the hydrocarbon counterpart of ``StapleState.from_sequence``. It is + a separate function precisely so the lactam constructor keeps its exact + per-character behaviour. + """ + return StapleState(sequence_tokens=tokenize_sequence(sequence), **kwargs) + + +def anchor_positions(tokens: list[str]) -> list[int]: + """Residue indices of the hydrocarbon anchors, N-to-C.""" + return [i for i, t in enumerate(tokens) if is_anchor_token(t)] + + +def to_model_tokens(tokens: list[str]) -> list[str]: + """Project monomer tokens onto tokens present in the existing model vocab. + + Multi-character hydrocarbon monomers are mapped onto the ``X``/``B`` ncAA + tokens that ``staplebridge.data.vocab`` already defines, so the embedding + matrix keeps its original size and old checkpoints still load. Natural + residues pass through unchanged. + """ + projected: list[str] = [] + for token in tokens: + upper = token.upper() + if upper in _MODEL_TOKEN_PROJECTION: + projected.append(_MODEL_TOKEN_PROJECTION[upper]) + elif token in TOKEN_TO_ID: + projected.append(token) + elif upper in TOKEN_TO_ID: + projected.append(upper) + else: + projected.append("") + return projected + + +def to_display_sequence(tokens: list[str]) -> str: + """Human-readable sequence string, dash-separating multi-character monomers. + + ``"".join`` would render ``[..., "S5", ...]`` ambiguously against a real + ``S`` followed by a literal ``5``, so multi-character monomers are set off + with dashes. + """ + parts: list[str] = [] + for token in tokens: + if len(token) > 1: + parts.append(f"-{token}-") + else: + parts.append(token) + return "".join(parts).replace("--", "-").strip("-") diff --git a/staplebridge/integrations/__init__.py b/staplebridge/integrations/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5a500d206407f59deb5fe1f2b261fee5f89a782e --- /dev/null +++ b/staplebridge/integrations/__init__.py @@ -0,0 +1 @@ +"""External integration utilities.""" diff --git a/staplebridge/integrations/peptiverse.py b/staplebridge/integrations/peptiverse.py new file mode 100644 index 0000000000000000000000000000000000000000..2325ab2230fc97f69c21c4b849153162fee016f2 --- /dev/null +++ b/staplebridge/integrations/peptiverse.py @@ -0,0 +1,537 @@ +from __future__ import annotations + +import math +import os +import sys +import types +from pathlib import Path +from typing import Any, Sequence + +import torch + + +PROPERTY_ALIASES = { + "hemolysis": "hemolysis", + "nf": "nf", + "non_fouling": "nf", + "non-fouling": "nf", + "solubility": "solubility", + "toxicity": "toxicity", + "permeability": "permeability_penetrance", + "permeability_penetrance": "permeability_penetrance", + "pampa": "permeability_pampa", + "permeability_pampa": "permeability_pampa", + "caco2": "permeability_caco2", + "permeability_caco2": "permeability_caco2", + "half_life": "halflife", + "halflife": "halflife", + "binding_affinity": "binding_affinity", +} + + +def _to_finite_float(x: Any) -> float | None: + try: + v = float(x) + except (TypeError, ValueError): + return None + if math.isnan(v) or math.isinf(v): + return None + return v + + +def extract_numeric_score(raw: Any, property_key: str) -> float: + """Pull the canonical numeric score for ``property_key`` out of ``raw``. + + For ``binding_affinity`` we prefer ``raw["affinity"]`` then ``raw["score"]``. + For other properties we prefer ``raw["score"]``. Falls back to walking + other plausible keys if a top-level dict is given. Raises ``ValueError`` + if no finite float can be produced. When the priority keys do not yield a + value but the fallback dict-walk does, a ``warnings.warn`` is emitted so + silent mis-extraction is surfaced. + """ + import warnings as _warnings + + pkey = (property_key or "").lower() + candidates: list[Any] = [] + if isinstance(raw, dict): + if pkey == "binding_affinity": + candidates += [raw.get("affinity"), raw.get("score")] + else: + candidates += [raw.get("score"), raw.get(pkey)] + candidates += [raw.get("value"), raw.get("prediction"), raw.get("mean")] + elif isinstance(raw, (list, tuple)) and raw: + candidates.append(raw[0]) + else: + candidates.append(raw) + + for c in candidates: + v = _to_finite_float(c) + if v is not None: + return v + if isinstance(raw, dict): + finite_pairs = [(k, _to_finite_float(v)) for k, v in raw.items()] + finite_pairs = [(k, v) for k, v in finite_pairs if v is not None] + if finite_pairs: + if len(finite_pairs) > 1: + _warnings.warn( + f"PeptiVerse score extraction for property={property_key!r}: " + f"no priority key matched; falling back to dict walk over " + f"{[k for k, _ in finite_pairs]} -> picked first finite value " + f"{finite_pairs[0][1]} from key {finite_pairs[0][0]!r}. " + "This may be wrong — verify the predictor output schema.", + stacklevel=2, + ) + return finite_pairs[0][1] + raise ValueError( + f"Cannot extract finite numeric score for property={property_key!r} from raw={raw!r}" + ) + + +class PeptiVerseUnavailable(RuntimeError): + pass + + +class PeptiVerseWrapper: + """Local-import wrapper around ChatterjeeLab/PeptiVerse `inference.py`. + + The HF repo is expected to be cloned to disk at ``peptiverse_root``; we + insert that path into ``sys.path`` and import ``inference.PeptiVersePredictor`` + directly. No remote HTTP API is used. + + Parameters + ---------- + peptiverse_root : path to the local PeptiVerse repo clone. + manifest_path : best_models.txt path (defaults to ``peptiverse_root/best_models.txt``). + classifier_weight_root : weight root (defaults to ``peptiverse_root``). + device : ``"cpu"`` or ``"cuda"``. + strict : if True, raise on any init error. + uncertainty : default ``uncertainty`` flag passed to the underlying + predictor unless overridden per-call. + cache_enabled : memoize identical (property, mode, input, target, uncertainty) + calls to avoid re-running the underlying torch models. + """ + + def __init__( + self, + peptiverse_root: str | Path, + manifest_path: str | Path | None = None, + classifier_weight_root: str | Path | None = None, + device: str | None = None, + strict: bool = False, + uncertainty: bool = False, + cache_enabled: bool = True, + hf_cache_dir: str | Path | None = None, + esm_model_name_or_path: str | Path | None = None, + peptideclm_model_name_or_path: str | Path | None = None, + chemberta_model_name_or_path: str | Path | None = None, + offline: bool = False, + batch_size: int = 32, + ) -> None: + self.root = Path(peptiverse_root) if peptiverse_root else None + self.manifest_path = Path(manifest_path) if manifest_path else None + self.classifier_weight_root = ( + Path(classifier_weight_root) if classifier_weight_root else None + ) + self.device = device + self.strict = strict + self.uncertainty = bool(uncertainty) + self.cache_enabled = bool(cache_enabled) + self.hf_cache_dir = Path(hf_cache_dir) if hf_cache_dir else None + self.esm_model_name_or_path = ( + Path(esm_model_name_or_path) if esm_model_name_or_path else None + ) + self.peptideclm_model_name_or_path = ( + Path(peptideclm_model_name_or_path) + if peptideclm_model_name_or_path + else None + ) + self.chemberta_model_name_or_path = ( + Path(chemberta_model_name_or_path) + if chemberta_model_name_or_path + else None + ) + self.offline = bool(offline) + self.predictor = None + self.available = False + self.init_error: str | None = None + self._cache: dict[tuple, dict[str, Any]] = {} + self.cache_hits = 0 + self.cache_misses = 0 + self.batch_size = max(1, int(batch_size)) + self.batch_chunks = 0 + self.batch_oom_retries = 0 + self.batch_oom_giveups = 0 + self._init_predictor() + + @staticmethod + def _snapshot_from_cache(cache_root: Path, model_id: str) -> Path | None: + """Resolve a cached Hugging Face repo to a concrete local snapshot.""" + repo = cache_root / "hub" / f"models--{model_id.replace('/', '--')}" + snapshots = repo / "snapshots" + ref = repo / "refs" / "main" + if ref.is_file(): + revision = ref.read_text(encoding="utf-8").strip() + candidate = snapshots / revision + if candidate.is_dir(): + return candidate.resolve() + candidates = sorted(path for path in snapshots.glob("*") if path.is_dir()) + if len(candidates) == 1: + return candidates[0].resolve() + return None + + def _local_hf_models(self) -> dict[str, Path]: + configured = { + "esm_name": self.esm_model_name_or_path, + "clm_name": self.peptideclm_model_name_or_path, + "chemberta_name": self.chemberta_model_name_or_path, + } + model_ids = { + "esm_name": "facebook/esm2_t33_650M_UR50D", + "clm_name": "aaronfeller/PeptideCLM-23M-all", + "chemberta_name": "DeepChem/ChemBERTa-77M-MLM", + } + resolved: dict[str, Path] = {} + for argument, configured_path in configured.items(): + path = configured_path + if path is None and self.hf_cache_dir is not None: + path = self._snapshot_from_cache( + self.hf_cache_dir, model_ids[argument] + ) + if path is not None: + path = path.expanduser().resolve() + if not path.is_dir(): + raise FileNotFoundError( + f"local PeptiVerse HF model directory missing: {path}" + ) + resolved[argument] = path + return resolved + + def _init_predictor(self) -> None: + if self.hf_cache_dir is not None: + cache_root = self.hf_cache_dir.expanduser().resolve() + # Deliberately override inherited values. A stale user-level HF + # cache was the source of tokenizer lookup failures in offline + # runs; this wrapper's configured local model root is authoritative. + os.environ["HF_HOME"] = str(cache_root) + os.environ["HUGGINGFACE_HUB_CACHE"] = str(cache_root / "hub") + os.environ["TRANSFORMERS_CACHE"] = str(cache_root / "hub") + if self.offline: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + + if self.root is None or not self.root.exists(): + self.init_error = f"peptiverse_root does not exist: {self.root}" + if self.strict: + raise PeptiVerseUnavailable(self.init_error) + return + if not (self.root / "inference.py").exists(): + self.init_error = f"inference.py not found under {self.root}" + if self.strict: + raise PeptiVerseUnavailable(self.init_error) + return + + sys.path.insert(0, str(self.root)) + try: + # PeptiVerse imports ``tokenizer.my_tokenizers`` from its own + # source tree. Some environments also install an unrelated + # top-level ``tokenizer`` package; once that regular package is in + # sys.modules it shadows PeptiVerse's namespace directory even + # though the PeptiVerse root is first on sys.path. Bind the local + # namespace explicitly before importing inference. + local_tokenizer = self.root / "tokenizer" + if (local_tokenizer / "my_tokenizers.py").is_file(): + loaded = sys.modules.get("tokenizer") + loaded_paths = [ + str(Path(value).resolve()) + for value in (getattr(loaded, "__path__", None) or []) + ] + if str(local_tokenizer.resolve()) not in loaded_paths: + for name in list(sys.modules): + if name == "tokenizer" or name.startswith("tokenizer."): + del sys.modules[name] + namespace = types.ModuleType("tokenizer") + namespace.__path__ = [str(local_tokenizer)] # type: ignore[attr-defined] + namespace.__package__ = "tokenizer" + sys.modules["tokenizer"] = namespace + from inference import PeptiVersePredictor # type: ignore + + # Several PeptiVerse MAPIE calibration artifacts were serialized + # by a training script where this estimator lived in __main__. + # Re-export the identical local inference.py class there before + # joblib loads those artifacts. Point predictions do not use MAPIE + # when uncertainty=False, but resolving the class avoids noisy + # load failures and preserves optional uncertainty availability. + try: + from inference import PassthroughRegressor # type: ignore + + main_module = sys.modules.get("__main__") + if main_module is not None and not hasattr( + main_module, "PassthroughRegressor" + ): + setattr( + main_module, + "PassthroughRegressor", + PassthroughRegressor, + ) + except ImportError: + pass + + manifest = self.manifest_path or (self.root / "best_models.txt") + weight_root = self.classifier_weight_root or self.root + kwargs: dict[str, Any] = { + "manifest_path": str(manifest), + "classifier_weight_root": str(weight_root), + } + if self.device is not None: + kwargs["device"] = self.device + for argument, model_path in self._local_hf_models().items(): + kwargs[argument] = str(model_path) + self.predictor = PeptiVersePredictor(**kwargs) + self.available = True + except Exception as exc: # noqa: BLE001 - surface any import/init failure + self.init_error = f"{type(exc).__name__}: {exc}" + self.available = False + if self.strict: + raise PeptiVerseUnavailable(self.init_error) from exc + + @staticmethod + def normalize_property_name(key: str) -> str: + return PROPERTY_ALIASES.get(key.lower(), key) + + @staticmethod + def extract_numeric_score(raw: Any, property_key: str) -> float: + return extract_numeric_score(raw, property_key) + + def _ensure_available(self) -> None: + if not self.available or self.predictor is None: + raise PeptiVerseUnavailable( + self.init_error or "PeptiVerse predictor is not available" + ) + + def clear_cache(self) -> None: + self._cache.clear() + self.cache_hits = 0 + self.cache_misses = 0 + + def _cache_key( + self, property_key: str, mode: str, input_str: str, target_seq: str, uncertainty: bool + ) -> tuple: + return (property_key, mode, input_str, target_seq, bool(uncertainty)) + + def predict_property( + self, + property_key: str, + input_str: str, + mode: str = "wt", + uncertainty: bool | None = None, + ) -> dict[str, Any]: + """Run a single non-binding property head. + + Returns ``{"backend_used": "peptiverse", "raw": , "score": float, + "property": , "mode": , "cached": bool}``. + """ + self._ensure_available() + std_key = self.normalize_property_name(property_key) + unc = self.uncertainty if uncertainty is None else bool(uncertainty) + key = self._cache_key(std_key, mode, input_str, "", unc) + if self.cache_enabled and key in self._cache: + self.cache_hits += 1 + cached = dict(self._cache[key]) + cached["cached"] = True + return cached + + raw = self.predictor.predict_property( # type: ignore[union-attr] + std_key, mode, input_str, unc + ) + score = extract_numeric_score(raw, std_key) + out = { + "backend_used": "peptiverse", + "property": std_key, + "mode": mode, + "raw": raw, + "score": float(score), + "cached": False, + } + if self.cache_enabled: + self._cache[key] = {k: v for k, v in out.items() if k != "cached"} + self.cache_misses += 1 + return out + + def predict_property_batch( + self, + property_key: str, + input_strs: Sequence[str], + mode: str = "wt", + uncertainty: bool | None = None, + batch_size: int | None = None, + ) -> list[dict[str, Any]]: + """Run one non-binding property head over many inputs. + + Same per-item contract as :meth:`predict_property`, returned in input + order. Duplicates and cache hits are computed once. The speedup comes + from batching the *embedding* forward pass; the heads still see one + input at a time, so each score is produced by exactly the code path + :meth:`predict_property` would have used. + + Batching an embedder changes the padded sequence length, which perturbs + the pooled embedding at ~1e-7. That is why only the embeddings are + batched here and the results are cached under the same keys as the + scalar path: a later scalar call returns the batch-computed value, so a + run cannot mix the two conventions for the same input. Verify numeric + agreement for any new property head before relying on it. + + On CUDA OOM the offending chunk is retried at half the batch size, down + to scalar, so a long tail of large molecules degrades rather than fails. + """ + self._ensure_available() + std_key = self.normalize_property_name(property_key) + if std_key == "binding_affinity": + raise ValueError("use predict_binding_affinity for binding affinity") + unc = self.uncertainty if uncertainty is None else bool(uncertainty) + inputs = list(input_strs) + if not inputs: + return [] + + # Deduplicate while preserving first-seen order; only uncached, unique + # inputs reach the backend. + pending: list[str] = [] + seen: set[str] = set() + for value in inputs: + key = self._cache_key(std_key, mode, value, "", unc) + if self.cache_enabled and key in self._cache: + continue + if value in seen: + continue + seen.add(value) + pending.append(value) + + if pending: + self._embed_batch(std_key, mode, pending, unc, batch_size) + # Every input now either was cached or has a warm embedding, so this + # loop is the unmodified scalar path and defines the returned values. + return [ + self.predict_property(std_key, value, mode=mode, uncertainty=unc) + for value in inputs + ] + + def _embed_batch( + self, + std_key: str, + mode: str, + pending: Sequence[str], + uncertainty: bool, + batch_size: int | None, + ) -> None: + """Warm the underlying embedder's cache for ``pending`` in batches.""" + embedder, pooled_kind = self._embedder_for(std_key, mode) + if embedder is None: + return + size = int(batch_size or self.batch_size) + index = 0 + items = list(pending) + while index < len(items): + chunk = items[index : index + size] + try: + self._embed_chunk(embedder, pooled_kind, chunk) + except torch.cuda.OutOfMemoryError: + torch.cuda.empty_cache() + if len(chunk) == 1: + # Nothing left to split; let the scalar path surface it. + self.batch_oom_giveups += 1 + index += 1 + continue + size = max(1, len(chunk) // 2) + self.batch_oom_retries += 1 + continue + index += len(chunk) + self.batch_chunks += 1 + + def _embedder_for(self, std_key: str, mode: str) -> tuple[Any, str]: + """The embedder and pooling kind the head for ``std_key`` will request.""" + predictor = self.predictor + meta = getattr(predictor, "meta", {}).get((std_key, mode)) + if meta is None: + return None, "" + try: + embedder = predictor._get_embedder(meta["emb_tag"]) # type: ignore[union-attr] + except (KeyError, ValueError): + return None, "" + # torch_ckpt heads consume unpooled token embeddings; everything else + # consumes the masked mean pool. Mirrors PeptiVersePredictor._get_features. + return embedder, ("unpooled" if meta.get("kind") == "torch_ckpt" else "pooled") + + @staticmethod + def _embed_chunk(embedder: Any, pooled_kind: str, chunk: Sequence[str]) -> None: + """Forward ``chunk`` once and store per-item results in the embedder cache. + + Writes into the embedder's own ``_cache_pooled``/``_cache_unpooled`` so + the scalar ``pooled``/``unpooled`` calls that follow find them. Items + already cached are skipped. + """ + cache_attr = "_cache_pooled" if pooled_kind == "pooled" else "_cache_unpooled" + cache = getattr(embedder, cache_attr, None) + if cache is None: + return + stripped = [str(value).strip() for value in chunk] + fresh = [value for value in dict.fromkeys(stripped) if value not in cache] + if not fresh: + return + + tokenized = embedder._tokenize(fresh) + with torch.no_grad(): + hidden = embedder.model( + input_ids=tokenized["input_ids"], + attention_mask=tokenized["attention_mask"], + ).last_hidden_state + valid = embedder._valid_mask(tokenized["input_ids"], tokenized["attention_mask"]) + + if pooled_kind == "pooled": + weights = valid.unsqueeze(-1).float() + pooled = (hidden * weights).sum(dim=1) / weights.sum(dim=1).clamp(min=1e-9) + for position, value in enumerate(fresh): + cache[value] = pooled[position : position + 1] + return + + for position, value in enumerate(fresh): + row = valid[position] + features = hidden[position : position + 1, row, :] + mask = torch.ones( + (1, features.shape[1]), dtype=torch.bool, device=features.device + ) + cache[value] = (features, mask) + + def predict_binding_affinity( + self, + binder_str: str, + target_seq: str, + mode: str = "wt", + uncertainty: bool | None = None, + ) -> dict[str, Any]: + """Run the binding-affinity head against a specific protein target.""" + self._ensure_available() + if not target_seq: + raise ValueError("predict_binding_affinity requires target_seq") + unc = self.uncertainty if uncertainty is None else bool(uncertainty) + key = self._cache_key("binding_affinity", mode, binder_str, target_seq, unc) + if self.cache_enabled and key in self._cache: + self.cache_hits += 1 + cached = dict(self._cache[key]) + cached["cached"] = True + return cached + + raw = self.predictor.predict_binding_affinity( # type: ignore[union-attr] + mode, target_seq, binder_str, unc + ) + score = extract_numeric_score(raw, "binding_affinity") + out = { + "backend_used": "peptiverse", + "property": "binding_affinity", + "mode": mode, + "raw": raw, + "score": float(score), + "cached": False, + } + if self.cache_enabled: + self._cache[key] = {k: v for k, v in out.items() if k != "cached"} + self.cache_misses += 1 + return out diff --git a/staplebridge/oracles/__init__.py b/staplebridge/oracles/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f3b0d32f802e05f5e1927433a19f527fd5a43b92 --- /dev/null +++ b/staplebridge/oracles/__init__.py @@ -0,0 +1 @@ +"""""" diff --git a/staplebridge/oracles/anchor_prior.py b/staplebridge/oracles/anchor_prior.py new file mode 100644 index 0000000000000000000000000000000000000000..3636894a60bd538669a459eaf9361125a56b093f --- /dev/null +++ b/staplebridge/oracles/anchor_prior.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from typing import Any + +from staplebridge.oracles.base import AnchorPriorBase + + +class MockAnchorPrior(AnchorPriorBase): + """Anchor scorer. + + Without strong priors on what makes a "good" anchor, we just lightly + favor having an anchor at all and penalize touching protected positions. + Spacing preferences now live in the block.motif (e.g. STAPLE_LACTAM + accepts spacings 3 and 4), so this prior no longer hardcodes (4, 7). + """ + + def score_anchor(self, sequence: list[str], anchor_pair: tuple[int, int] | None, context: dict[str, Any] | None = None) -> float: + del sequence + if anchor_pair is None: + return -0.2 + i, j = anchor_pair + score = 0.5 # mild bias toward having an anchor + if context and context.get("protected_positions"): + protected = set(context["protected_positions"]) + if i in protected or j in protected: + score -= 1.0 + return score diff --git a/staplebridge/oracles/base.py b/staplebridge/oracles/base.py new file mode 100644 index 0000000000000000000000000000000000000000..ede69f891a73268dca8e798a8619bb7a5e0369a0 --- /dev/null +++ b/staplebridge/oracles/base.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from staplebridge.chemistry.state import StapleState +from staplebridge.data.schemas import BuildingBlock + + +class PeptidePriorBase(ABC): + @abstractmethod + def score_transition(self, old_state: StapleState, new_state: StapleState, context: dict[str, Any] | None = None) -> float: + pass + + def batch_score_transitions( + self, + old_state: StapleState, + new_states: list[StapleState], + context: dict[str, Any] | None = None, + ) -> list[float]: + """Score N candidate transitions from ``old_state`` at once. + + Default implementation just loops ``score_transition``; heavyweight + priors (ESM2) override this so they can share one model forward across + all candidates. Non-sequence-changing candidates are expected to + return exactly 0.0. + """ + return [ + self.score_transition(old_state, new, context) for new in new_states + ] + + def prewarm_requests( + self, pairs: list[tuple[StapleState, list[StapleState]]] + ) -> None: + """Prefetch model outputs for many (z, candidates) pairs at once. + + Default is a noop — heavy priors (ESM2) override this to run one + batched model forward covering every request across all pairs, so a + subsequent per-pair ``batch_score_transitions`` call becomes a pure + cache-lookup. + """ + del pairs + + +class AnchorPriorBase(ABC): + @abstractmethod + def score_anchor(self, sequence: list[str], anchor_pair: tuple[int, int] | None, context: dict[str, Any] | None = None) -> float: + pass + + +class BlockPriorBase(ABC): + @abstractmethod + def score_block( + self, + sequence: list[str], + anchor_pair: tuple[int, int] | None, + block: BuildingBlock | None, + context: dict[str, Any] | None = None, + ) -> float: + pass + + +class GeometryOracleBase(ABC): + @abstractmethod + def ctype( + self, + sequence: list[str], + anchor_pair: tuple[int, int] | None, + block: BuildingBlock | None, + *, + peptide_ca: list[tuple[float, float, float]] | None = None, + ) -> bool: + pass + + @abstractmethod + def cgeom( + self, + sequence: list[str], + anchor_pair: tuple[int, int] | None, + block: BuildingBlock | None, + *, + peptide_ca: list[tuple[float, float, float]] | None = None, + ) -> float: + pass diff --git a/staplebridge/oracles/block_prior.py b/staplebridge/oracles/block_prior.py new file mode 100644 index 0000000000000000000000000000000000000000..da1d34cab758546eb110a6c73043dbbb6e9eb4f9 --- /dev/null +++ b/staplebridge/oracles/block_prior.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from typing import Any + +from staplebridge.data.schemas import BuildingBlock +from staplebridge.oracles.base import BlockPriorBase + + +def _motif_satisfied(seq: list[str], i: int, j: int, motif: dict | None) -> bool: + if motif is None: + return True + spacings = motif.get("spacings") + if spacings != "any" and spacings is not None and (j - i) not in spacings: + return False + if motif.get("i_aa") and seq[i] not in motif["i_aa"]: + return False + if motif.get("j_aa") and seq[j] not in motif["j_aa"]: + return False + return True + + +class MockBlockPrior(BlockPriorBase): + def score_block( + self, + sequence: list[str], + anchor_pair: tuple[int, int] | None, + block: BuildingBlock | None, + context: dict[str, Any] | None = None, + ) -> float: + del context + if block is None: + return -0.3 + if anchor_pair is None: + return -1.0 + i, j = anchor_pair + score = 0.0 + score += 0.5 * block.synthetic_accessibility_score + score -= 0.2 * block.cost_score + if _motif_satisfied(sequence, i, j, block.motif): + score += 1.0 + else: + score -= 1.0 + return score diff --git a/staplebridge/oracles/catalog_block_prior.py b/staplebridge/oracles/catalog_block_prior.py new file mode 100644 index 0000000000000000000000000000000000000000..01e377fb5dd73c21a071aeb19e5d08a24b9224a5 --- /dev/null +++ b/staplebridge/oracles/catalog_block_prior.py @@ -0,0 +1,144 @@ +"""Catalog-informed block prior. + +Scores a building block conditional on the current sequence and anchor pair +using **only** catalog metadata — no PeptiVerse, no learned model. Signals: + + * spacing compatibility — is (j - i) in the block motif's spacings? + * residue compatibility — are (seq[i], seq[j]) in the block motif's + (i_aa, j_aa) sets? + * motif edit distance — how many substitutions to satisfy the motif? + * synthetic accessibility, SPPS, cost — pulled from BuildingBlock fields. + +Works with the existing ``BuildingBlock`` schema (motif dict) *and* legacy +schemas that expose ``allowed_anchor_spacings`` / ``compatible_residue_types`` +as top-level attributes; the schema has retired the latter, but callers may +inject them for ablations. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from staplebridge.data.schemas import BuildingBlock +from staplebridge.oracles.base import BlockPriorBase, GeometryOracleBase + + +@dataclass +class CatalogBlockPriorConfig: + weights: dict[str, float] = field(default_factory=lambda: { + "spacing": 2.0, + "residue_compatibility": 2.0, + "motif_edit_distance": 1.0, + "synthetic_accessibility": 0.5, + "spps": 0.5, + "cost": 0.2, + }) + max_motif_edits: int = 2 + no_block_penalty: float = 0.3 + no_anchor_penalty: float = 1.0 + + +class CatalogBlockPrior(BlockPriorBase): + """Score blocks purely from catalog + motif compatibility.""" + + def __init__( + self, + catalog: dict[str, BuildingBlock] | None = None, + geometry_oracle: GeometryOracleBase | None = None, + cfg: CatalogBlockPriorConfig | None = None, + ) -> None: + self.catalog = catalog or {} + self.geometry_oracle = geometry_oracle # currently unused; kept for + # symmetry with the anchor prior and for future extensions. + self.cfg = cfg or CatalogBlockPriorConfig() + + def _allowed_spacings(self, block: BuildingBlock) -> list[int] | None: + # Prefer legacy explicit field if present (ablations may set it). + allowed = getattr(block, "allowed_anchor_spacings", None) + if allowed: + return list(allowed) + motif = block.motif or {} + spacings = motif.get("spacings") + if spacings == "any": + return None # any spacing OK + return list(spacings) if spacings else None + + def _residue_ok(self, seq: list[str], i: int, j: int, block: BuildingBlock) -> tuple[bool, int]: + """Return (both_sides_ok, edit_count_needed_for_motif).""" + motif = block.motif or {} + i_aa = motif.get("i_aa") + j_aa = motif.get("j_aa") + legacy = getattr(block, "compatible_residue_types", None) + if not i_aa and isinstance(legacy, dict): + i_aa = legacy.get("i") + if not j_aa and isinstance(legacy, dict): + j_aa = legacy.get("j") + edits = 0 + i_ok = True + j_ok = True + if i_aa and 0 <= i < len(seq): + i_ok = seq[i] in i_aa + if not i_ok: + edits += 1 + if j_aa and 0 <= j < len(seq): + j_ok = seq[j] in j_aa + if not j_ok: + edits += 1 + return (i_ok and j_ok), edits + + def score_block( + self, + sequence: list[str], + anchor_pair: tuple[int, int] | None, + block: BuildingBlock | None, + context: dict[str, Any] | None = None, + ) -> float: + del context + if block is None: + return -self.cfg.no_block_penalty + if anchor_pair is None: + return -self.cfg.no_anchor_penalty + i, j = anchor_pair + if i > j: + i, j = j, i + w = self.cfg.weights + score = 0.0 + + # 1. Spacing + allowed = self._allowed_spacings(block) + if allowed is None: + score += 0.5 * float(w.get("spacing", 2.0)) + elif (j - i) in allowed: + score += float(w.get("spacing", 2.0)) + else: + # Distance-to-nearest-allowed spacing controls the size of the + # penalty (adjacent spacings hurt less than far-away ones). + gap = min(abs((j - i) - a) for a in allowed) + score -= float(w.get("spacing", 2.0)) * (0.5 + 0.5 * gap) + + # 2. Residue compatibility + motif edit distance + residues_ok, edits = self._residue_ok(sequence, i, j, block) + if residues_ok: + score += float(w.get("residue_compatibility", 2.0)) + else: + score -= float(w.get("residue_compatibility", 2.0)) * 0.5 + # motif_edit_distance term: 0 -> +weight, 1 -> +0.4*weight, 2 -> 0, + # >max -> -weight. + if edits == 0: + score += float(w.get("motif_edit_distance", 1.0)) + elif edits == 1: + score += 0.4 * float(w.get("motif_edit_distance", 1.0)) + elif edits == 2: + score += 0.0 + elif edits <= self.cfg.max_motif_edits: + score -= 0.5 * float(w.get("motif_edit_distance", 1.0)) + else: + score -= float(w.get("motif_edit_distance", 1.0)) + + # 3. Synthetic accessibility (larger better) / SPPS (larger better) / cost (smaller better) + score += float(w.get("synthetic_accessibility", 0.5)) * float(block.synthetic_accessibility_score) + score += float(w.get("spps", 0.5)) * float(block.spps_score) + score -= float(w.get("cost", 0.2)) * float(block.cost_score) + + return float(score) diff --git a/staplebridge/oracles/esm2_peptide_prior.py b/staplebridge/oracles/esm2_peptide_prior.py new file mode 100644 index 0000000000000000000000000000000000000000..ac5565f230f741a01d2e8fee8b8c9a2e4d7eae54 --- /dev/null +++ b/staplebridge/oracles/esm2_peptide_prior.py @@ -0,0 +1,565 @@ +"""ESM2-based peptide prior — batched, dedup-friendly. + +Scores a state transition by the change in ESM2 masked log-prob at the +sequence positions that changed. Non-sequence transitions (topology / anchor / +block only) return exactly ``0.0`` and cost nothing. + +Correctness note (motivating the batching design): + A substitution changes exactly one position ``p``. To score it we want + ``lp_new = log P(new_aa | seq_new masked at p)`` and + ``lp_old = log P(old_aa | seq_old masked at p)``. + Since ``seq_old`` and ``seq_new`` differ only at ``p`` and we mask ``p``, + the masked input is *identical* on both sides. One ESM2 forward over the + masked sequence yields the full 20-AA log-prob vector at ``p``, from + which we index both ``lp_old`` and ``lp_new``. + +That means: + * one ESM2 forward per unique ``(masked_seq_str, pos)`` covers many candidates + that only differ in the substituted amino acid at ``p``; + * ``batch_score_transitions`` collects unique requests across an entire + candidate neighborhood, deduplicates by ``(masked_seq, pos)``, and + pushes cache misses through **one** batched ESM2 forward pass; + * the cache stores the full 20-vector, so future single-position lookups + cost O(1) with no additional forwards. +""" + +from __future__ import annotations + +import math +import os +import sqlite3 +import struct +import threading +from collections import OrderedDict +from pathlib import Path +from typing import Any + +from staplebridge.chemistry.state import StapleState +from staplebridge.oracles.base import PeptidePriorBase +from staplebridge.utils.profiling import STAGE_TIMER + + +CANONICAL_AA = "ACDEFGHIKLMNPQRSTVWY" +_AA_TO_IDX = {a: i for i, a in enumerate(CANONICAL_AA)} +_DEFAULT_SURROGATE = {"X": "A", "B": "A"} +_VECTOR_LEN = len(CANONICAL_AA) # 20 +_VECTOR_STRUCT = struct.Struct(f"<{_VECTOR_LEN}f") # 80 bytes + + +class _SqliteVectorCache: + """Persistent cache mapping ``(masked_seq, pos) -> 20-AA log-prob vector``. + + Stores the vector as an 80-byte little-endian float32 blob. Failures never + propagate — the in-memory OrderedDict continues to satisfy hits. + """ + + def __init__(self, path: str | None, memory_size: int = 65536, commit_every: int = 1024) -> None: + self.memory_size = int(memory_size) + self._mem: "OrderedDict[tuple[str, int], list[float]]" = OrderedDict() + self._lock = threading.Lock() + self._db: sqlite3.Connection | None = None + self.path = path + # Batched-commit bookkeeping. Vectors are served from ``_mem`` the moment + # they are ``put`` (before any commit), so deferring the sqlite commit + # cannot change a single returned value or its byte content -- it only + # changes *when* rows are made durable. On a crash the at-most + # ``commit_every`` uncommitted vectors are simply recomputed next run. + self._commit_every = max(1, int(commit_every)) + self._pending = 0 + if path: + try: + p = Path(path) + p.parent.mkdir(parents=True, exist_ok=True) + self._db = sqlite3.connect(str(p), check_same_thread=False) + # WAL + NORMAL: one fsync per commit instead of the default + # journal's rewrite-and-fsync-per-commit. Combined with batched + # commits this removes the ~1 s/vector fsync that dominated the + # cold-cache rollout path. + self._db.execute("PRAGMA journal_mode=WAL") + self._db.execute("PRAGMA synchronous=NORMAL") + self._db.execute( + "CREATE TABLE IF NOT EXISTS vec_scores " + "(seq TEXT, pos INTEGER, vec BLOB, " + "PRIMARY KEY(seq, pos))" + ) + self._db.commit() + except Exception: + self._db = None + + def get(self, seq: str, pos: int) -> list[float] | None: + key = (seq, int(pos)) + with self._lock: + v = self._mem.get(key) + if v is not None: + self._mem.move_to_end(key) + return v + if self._db is not None: + try: + cur = self._db.execute( + "SELECT vec FROM vec_scores WHERE seq=? AND pos=?", + (seq, int(pos)), + ) + row = cur.fetchone() + if row is not None: + blob = row[0] + if isinstance(blob, (bytes, bytearray)) and len(blob) == _VECTOR_STRUCT.size: + vec = list(_VECTOR_STRUCT.unpack(bytes(blob))) + self._put_mem(key, vec) + return vec + except Exception: + pass + return None + + def put(self, seq: str, pos: int, vec: list[float]) -> None: + if len(vec) != _VECTOR_LEN: + return + key = (seq, int(pos)) + self._put_mem(key, list(vec)) + if self._db is not None: + try: + blob = _VECTOR_STRUCT.pack(*[float(x) for x in vec]) + self._db.execute( + "INSERT OR REPLACE INTO vec_scores(seq, pos, vec) VALUES (?, ?, ?)", + (seq, int(pos), blob), + ) + # Defer the commit; flush() (called once per forward batch by + # _get_vectors) makes the batch durable. A safety cap bounds the + # unflushed set if a single batch is unusually large. + self._pending += 1 + if self._pending >= self._commit_every: + self._db.commit() + self._pending = 0 + except Exception: + pass + + def flush(self) -> None: + """Commit any rows written since the last commit. Safe to call always.""" + if self._db is not None and self._pending: + try: + self._db.commit() + except Exception: + pass + self._pending = 0 + + def _put_mem(self, key: tuple[str, int], val: list[float]) -> None: + with self._lock: + self._mem[key] = val + self._mem.move_to_end(key) + while len(self._mem) > self.memory_size: + self._mem.popitem(last=False) + + +class ESM2DeltaPeptidePrior(PeptidePriorBase): + """Score transitions by masked-position log-prob under a frozen ESM2. + + Config surface (mirrors :func:`build_peptide_prior`): + model_name_or_path : str — local snapshot dir; can be omitted if a + live ``ESM2FrozenEncoder`` is passed. + cache_path : str | None — sqlite path for persistent caching. + Legacy scalar-cache DBs are ignored; + a new ``vec_scores`` table is used. + ncaa_policy : dict — {"mode": "canonical_surrogate", + "surrogates": {...}, + "unknown_penalty": 0.2} + offline : bool — force offline (default respects env vars). + device : str — "cpu" | "cuda" | "cuda:N". + temperature : float — softmax temperature (default 1.0). + max_batch_size : int — cap for a single ESM2 forward batch (safety + valve for very large neighborhoods). + """ + + def __init__( + self, + model_name_or_path: str | None = None, + *, + sequence_encoder: Any = None, + cache_path: str | None = None, + ncaa_policy: dict[str, Any] | None = None, + offline: bool = False, + device: str = "cpu", + temperature: float = 1.0, + max_batch_size: int = 128, + memory_cache_size: int = 65536, + strict_runtime: bool = False, + ) -> None: + self.model_name_or_path = model_name_or_path + self.device = device + self.temperature = max(float(temperature), 1e-6) + pol = dict(ncaa_policy or {}) + self.ncaa_mode = pol.get("mode", "canonical_surrogate") + self.surrogates = dict(_DEFAULT_SURROGATE) + self.surrogates.update(pol.get("surrogates") or {}) + self.unknown_penalty = float(pol.get("unknown_penalty", 0.2)) + self.offline = bool(offline) + self._external_encoder = sequence_encoder + self.max_batch_size = max(int(max_batch_size), 1) + self.strict_runtime = bool(strict_runtime) + + self._cache = _SqliteVectorCache(cache_path, memory_size=memory_cache_size) + self._model = None + self._tokenizer = None + self._vocab_ids: list[int] | None = None # canonical-order lookup + self._mask_id: int | None = None + self._init_lock = threading.Lock() + self._init_error: str | None = None + + # ------------------------------------------------------------------ + # Lazy model construction + # ------------------------------------------------------------------ + def _ensure_loaded(self) -> bool: + if self._model is not None or self._init_error: + return self._model is not None + with self._init_lock: + if self._model is not None: + return True + if self._init_error: + return False + if self.offline: + os.environ.setdefault("HF_HUB_OFFLINE", "1") + os.environ.setdefault("TRANSFORMERS_OFFLINE", "1") + try: + if self._external_encoder is not None: + self._tokenizer = self._external_encoder.tokenizer + from transformers import EsmForMaskedLM # type: ignore + + path = getattr(self._external_encoder, "model_path", None) or self.model_name_or_path + if path is None: + raise ValueError( + "ESM2DeltaPeptidePrior needs a model_name_or_path; " + "the passed encoder has no `model_path` attribute." + ) + self._model = EsmForMaskedLM.from_pretrained(path) + elif self.model_name_or_path: + from transformers import EsmForMaskedLM, EsmTokenizer # type: ignore + + self._tokenizer = EsmTokenizer.from_pretrained(self.model_name_or_path) + self._model = EsmForMaskedLM.from_pretrained(self.model_name_or_path) + else: + raise ValueError( + "ESM2DeltaPeptidePrior: neither sequence_encoder nor " + "model_name_or_path was provided." + ) + self._model.eval() + for p in self._model.parameters(): + p.requires_grad_(False) + try: + import torch # type: ignore + + if self.device.startswith("cuda") and torch.cuda.is_available(): + self._model.to(self.device) + else: + self._model.to("cpu") + except Exception: + pass + mask_tok = getattr(self._tokenizer, "mask_token", "") + self._mask_id = self._tokenizer.convert_tokens_to_ids(mask_tok) + self._vocab_ids = [ + int(self._tokenizer.convert_tokens_to_ids(aa)) + for aa in CANONICAL_AA + ] + return True + except Exception as exc: + self._init_error = f"{type(exc).__name__}: {exc}" + self._model = None + if self.strict_runtime: + raise RuntimeError( + f"strict ESM2 peptide prior unavailable: {self._init_error}" + ) from exc + return False + + def ensure_available(self) -> None: + """Eagerly verify the configured frozen model for strict runs.""" + if not self._ensure_loaded(): + raise RuntimeError(f"ESM2 peptide prior unavailable: {self._init_error or 'unknown error'}") + + # ------------------------------------------------------------------ + # Token surrogate + # ------------------------------------------------------------------ + def _canonicalize(self, seq: list[str]) -> tuple[str, float]: + out: list[str] = [] + penalty = 0.0 + for t in seq: + if isinstance(t, str) and len(t) == 1 and t.upper() in _AA_TO_IDX: + out.append(t.upper()) + continue + if self.ncaa_mode == "canonical_surrogate": + sub = self.surrogates.get(t) or self.surrogates.get(str(t).upper()) + if sub and sub in _AA_TO_IDX: + out.append(sub) + continue + out.append("A") + penalty += self.unknown_penalty + return "".join(out), penalty + + # ------------------------------------------------------------------ + # Batched forward — the whole point of this rewrite + # ------------------------------------------------------------------ + def _forward_batch( + self, requests: list[tuple[str, int]] + ) -> dict[tuple[str, int], list[float]]: + """Compute the 20-AA log-prob vector at ``pos`` under masked ``seq``. + + Returns a dict keyed by ``(seq, pos)``. Falls back to a zero-vector on + any error; the caller then treats it as a benign 0-delta contribution. + """ + if not requests or not self._ensure_loaded(): + return {} + try: + import torch # type: ignore + except Exception: + return {} + + results: dict[tuple[str, int], list[float]] = {} + # Chunk requests to bound peak activation memory. + for start in range(0, len(requests), self.max_batch_size): + chunk = requests[start : start + self.max_batch_size] + seqs = [seq for seq, _ in chunk] + try: + enc = self._tokenizer( + seqs, + return_tensors="pt", + add_special_tokens=True, + padding=True, + ) + input_ids = enc["input_ids"].clone() + attn = enc.get("attention_mask") + mask_positions: list[int] = [] + valid: list[bool] = [] + for i, (seq, pos) in enumerate(chunk): + mp = pos + 1 # shift + if mp <= 0 or mp >= input_ids.shape[1]: + mask_positions.append(0) + valid.append(False) + continue + input_ids[i, mp] = int(self._mask_id) # type: ignore[arg-type] + mask_positions.append(mp) + valid.append(True) + device = next(self._model.parameters()).device + input_ids = input_ids.to(device) + if attn is not None: + attn = attn.to(device) + with STAGE_TIMER.section("esm2_forward_time"): + with torch.no_grad(): + out = self._model(input_ids=input_ids, attention_mask=attn) + STAGE_TIMER.bump("esm2_forward_batches") + STAGE_TIMER.bump("esm2_forward_calls", len(chunk)) + # Gather per-row logits at each row's mask position. + logits = out.logits # [B, L, V] + idx = torch.arange(logits.shape[0], device=logits.device) + mp_tensor = torch.tensor( + mask_positions, device=logits.device, dtype=torch.long + ) + pos_logits = logits[idx, mp_tensor] / self.temperature # [B, V] + logprobs = torch.log_softmax(pos_logits, dim=-1) + # Column-select the 20 canonical-AA ids in one gather. + aa_ids = torch.tensor( + self._vocab_ids or [], device=logprobs.device, dtype=torch.long + ) + lp20 = logprobs.index_select(dim=-1, index=aa_ids) # [B, 20] + lp20_cpu = lp20.detach().to("cpu").tolist() + for i, (key, ok) in enumerate(zip(chunk, valid)): + if not ok: + results[key] = [0.0] * _VECTOR_LEN + else: + results[key] = list(lp20_cpu[i]) + except Exception as exc: + if self.strict_runtime: + raise RuntimeError(f"strict ESM2 peptide-prior forward failed: {exc}") from exc + # On any failure, degrade gracefully — mark this chunk as zeros. + for key in chunk: + results[key] = [0.0] * _VECTOR_LEN + return results + + def _get_vectors( + self, requests: list[tuple[str, int]] + ) -> dict[tuple[str, int], list[float]]: + """Return log-prob vectors for each ``(masked_seq, pos)`` request. + + Cache-hit vectors come from memory/sqlite; cache-miss requests are + deduped and pushed through a single batched ESM2 forward. + """ + vectors: dict[tuple[str, int], list[float]] = {} + misses: list[tuple[str, int]] = [] + seen_miss: set[tuple[str, int]] = set() + for key in requests: + if key in vectors: + continue + cached = self._cache.get(key[0], key[1]) + if cached is not None: + STAGE_TIMER.bump("esm2_cache_hit") + vectors[key] = cached + else: + if key not in seen_miss: + STAGE_TIMER.bump("esm2_cache_miss") + misses.append(key) + seen_miss.add(key) + if misses: + fresh = self._forward_batch(misses) + for key, vec in fresh.items(): + self._cache.put(key[0], key[1], vec) + vectors[key] = vec + # One commit per forward-batch group rather than one per vector. + self._cache.flush() + return vectors + + # ------------------------------------------------------------------ + # Scoring — single position substitution shortcut + # ------------------------------------------------------------------ + def _substitution_positions( + self, old_str: str, new_str: str + ) -> list[int] | None: + """Return positions where old and new differ; ``None`` if lengths don't match.""" + if len(old_str) != len(new_str): + return None + return [i for i in range(len(new_str)) if old_str[i] != new_str[i]] + + @staticmethod + def _mask_at(seq: str, pos: int) -> str: + # We keep the sequence string as-is; the actual mask token replaces the + # tokenized input in ``_forward_batch``. What matters is that all + # candidates masking the same (base_seq, pos) share the same key — + # base_seq is enough. The seq string carries every non-masked position + # verbatim. + return seq + + # ------------------------------------------------------------------ + # PeptidePriorBase — scalar API + # ------------------------------------------------------------------ + def score_transition( + self, + old_state: StapleState, + new_state: StapleState, + context: dict[str, Any] | None = None, + ) -> float: + return self.batch_score_transitions(old_state, [new_state], context)[0] + + def batch_score_transitions( + self, + old_state: StapleState, + new_states: list[StapleState], + context: dict[str, Any] | None = None, + ) -> list[float]: + del context + n = len(new_states) + if n == 0: + return [] + needed, plans, old_str = self._plan_and_collect(old_state, new_states) + vectors = self._get_vectors(needed) if needed else {} + return self._finalize_plans(plans, vectors) + + # ------------------------------------------------------------------ + # Public batching helpers (used by Decoder to share one ESM2 forward + # across many rollouts). + # ------------------------------------------------------------------ + + def prewarm_requests( + self, pairs: list[tuple[StapleState, list[StapleState]]] + ) -> None: + """Force one batched ESM2 forward covering all cache misses across + the given ``(z, candidates)`` pairs. Subsequent + ``batch_score_transitions`` calls with the same pairs will hit the + memory cache and do zero forwards. + """ + all_needed: list[tuple[str, int]] = [] + for z, cands in pairs: + if not cands: + continue + needed, _plans, _old = self._plan_and_collect(z, cands) + all_needed.extend(needed) + if all_needed: + self._get_vectors(all_needed) + + # ------------------------------------------------------------------ + # Internals shared with prewarm + # ------------------------------------------------------------------ + + def _plan_and_collect( + self, old_state: StapleState, new_states: list[StapleState] + ) -> tuple[list[tuple[str, int]], list[Any], str]: + old_tokens = list(old_state.sequence_tokens) + old_str, pen_old = self._canonicalize(old_tokens) + plans: list[Any] = [] + needed: list[tuple[str, int]] = [] + for new_state in new_states: + new_tokens = list(new_state.sequence_tokens) + if new_tokens == old_tokens: + plans.append(("zero",)) + continue + new_str, pen_new = self._canonicalize(new_tokens) + surrogate_penalty = -(pen_new - pen_old) + positions = self._substitution_positions(old_str, new_str) + if positions is None: + length_delta_pen = -self.unknown_penalty * abs( + len(new_str) - len(old_str) + ) + common = min(len(old_str), len(new_str)) + lc_positions = [i for i in range(common) if old_str[i] != new_str[i]] + for p in lc_positions: + needed.append((old_str, p)) + needed.append((new_str, p)) + plans.append( + ( + "length", + surrogate_penalty, + length_delta_pen, + lc_positions, + old_str, + new_str, + ) + ) + continue + for p in positions: + needed.append((old_str, p)) + plans.append(("subst", surrogate_penalty, positions, old_str, new_str)) + return needed, plans, old_str + + def _finalize_plans( + self, + plans: list[Any], + vectors: dict[tuple[str, int], list[float]], + ) -> list[float]: + def _lp(seq: str, pos: int, aa: str) -> float: + vec = vectors.get((seq, pos)) + if vec is None: + return 0.0 + i = _AA_TO_IDX.get(aa) + if i is None: + return 0.0 + return float(vec[i]) + + out: list[float] = [] + for plan in plans: + if plan[0] == "zero": + out.append(0.0) + continue + if plan[0] == "subst": + _, surrogate_penalty, positions, o_str, n_str = plan + if not positions: + out.append(float(surrogate_penalty)) + continue + delta = 0.0 + for p in positions: + lp_new = _lp(o_str, p, n_str[p]) + lp_old = _lp(o_str, p, o_str[p]) + delta += lp_new - lp_old + out.append(float(0.1 * delta + surrogate_penalty)) + continue + if plan[0] == "length": + _, surrogate_penalty, length_delta_pen, lc_positions, o_str, n_str = plan + delta = 0.0 + for p in lc_positions: + lp_new = _lp(n_str, p, n_str[p]) + lp_old = _lp(o_str, p, o_str[p]) + delta += lp_new - lp_old + out.append(float(delta + length_delta_pen + surrogate_penalty)) + continue + out.append(0.0) + return out + + @property + def is_available(self) -> bool: + return self._ensure_loaded() + + @property + def init_error(self) -> str | None: + return self._init_error diff --git a/staplebridge/oracles/heuristic_peptide_prior.py b/staplebridge/oracles/heuristic_peptide_prior.py new file mode 100644 index 0000000000000000000000000000000000000000..80927781411f95ebd38caa4b1a39a540ad99de44 --- /dev/null +++ b/staplebridge/oracles/heuristic_peptide_prior.py @@ -0,0 +1,56 @@ +"""Heuristic peptide prior — a slightly richer version of MockPeptidePrior. + +Kept as an ablation option (`peptide.backend: heuristic`). The default real +prior is :class:`ESM2DeltaPeptidePrior`; the mock/heuristic priors are +reserved for ablation studies and unit tests. +""" + +from __future__ import annotations + +from typing import Any + +from staplebridge.chemistry.state import StapleState +from staplebridge.data.vocab import AMINO_ACIDS, HYDROPHOBIC, RISKY_TOKENS +from staplebridge.oracles.base import PeptidePriorBase + + +class HeuristicPeptidePrior(PeptidePriorBase): + """Length + composition + risky-residue heuristic. + + Not a learned model, but discriminative enough to differentiate + substitutions that touch W/F/B (risky) from neutral ones. Only for + ablations; production runs must use ``esm2_delta``. + """ + + def __init__( + self, + target_length: int = 12, + length_weight: float = 1.0, + risky_penalty: float = 0.3, + hydrophobic_target: float = 0.45, + hydrophobic_weight: float = 0.4, + ) -> None: + self.target_length = int(target_length) + self.length_weight = float(length_weight) + self.risky_penalty = float(risky_penalty) + self.hydrophobic_target = float(hydrophobic_target) + self.hydrophobic_weight = float(hydrophobic_weight) + + def score_transition( + self, + old_state: StapleState, + new_state: StapleState, + context: dict[str, Any] | None = None, + ) -> float: + del old_state, context + seq = new_state.sequence_tokens + n = max(len(seq), 1) + valid = sum(1 for t in seq if t in AMINO_ACIDS) + length_pen = abs(len(seq) - self.target_length) / self.target_length + hydro = sum(1 for t in seq if t in HYDROPHOBIC) / n + risky = sum(1 for t in seq if t in RISKY_TOKENS) / n + composition_score = valid / n + length_score = -self.length_weight * length_pen + hydro_score = -self.hydrophobic_weight * abs(hydro - self.hydrophobic_target) + risky_score = -self.risky_penalty * risky + return float(composition_score + length_score + hydro_score + risky_score) diff --git a/staplebridge/oracles/motif_anchor_prior.py b/staplebridge/oracles/motif_anchor_prior.py new file mode 100644 index 0000000000000000000000000000000000000000..db077690bf9bd21e44af8b5f001d743188603974 --- /dev/null +++ b/staplebridge/oracles/motif_anchor_prior.py @@ -0,0 +1,190 @@ +"""Anchor priors — motif-support + geometry surrogate variant. + +The mock anchor prior in ``anchor_prior.py`` gives essentially uniform score +for any anchor, which lets the reference process wander into anchors that +require heavy edits. This prior scores an anchor pair by: + + 1. Whether ``j - i`` is inside ``allowed_spacings`` (default {3, 4}). + 2. How many residue substitutions would be needed to make the anchor + satisfy any block motif in the catalog (0 → strong bonus, + 1 → medium, 2 → small; > max_motif_edits → penalty). + 3. Whether the anchor overlaps ``context["protected_positions"]`` + (strong penalty when ``avoid_protected`` is enabled). + 4. An optional geometry surrogate: if a geometry oracle is available and + ``context["peptide_ca"]`` is provided, use ``-cgeom`` as an additive + term (larger is better). + +For debugging, callers can pass ``context["return_components"] = True`` and +retrieve the per-component breakdown from ``context["_components"]`` after the +call. This avoids changing the base-class signature. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from staplebridge.data.schemas import BuildingBlock +from staplebridge.oracles.base import AnchorPriorBase, GeometryOracleBase + + +@dataclass +class MotifSupportAnchorConfig: + allowed_spacings: tuple[int, ...] = (3, 4) + max_motif_edits: int = 2 + avoid_protected: bool = True + project_motif_for_geometry: bool = False + weights: dict[str, float] = field(default_factory=lambda: { + "valid_spacing": 1.0, + "existing_motif": 2.0, + "one_edit_motif": 1.2, + "two_edit_motif": 0.4, + "protected_penalty": 4.0, + "geometry_surrogate": 1.0, + }) + + +class MotifSupportAnchorPrior(AnchorPriorBase): + """Anchor scorer aware of catalog motifs, protected positions, and geometry.""" + + def __init__( + self, + catalog: dict[str, BuildingBlock] | None = None, + geometry_oracle: GeometryOracleBase | None = None, + cfg: MotifSupportAnchorConfig | None = None, + ) -> None: + self.catalog = catalog or {} + self.geometry_oracle = geometry_oracle + self.cfg = cfg or MotifSupportAnchorConfig() + + # ------------------------------------------------------------------ + # Motif support + # ------------------------------------------------------------------ + def _motif_edit_distance( + self, seq: list[str], i: int, j: int, motif: dict[str, Any] | None + ) -> int: + """Number of residue substitutions needed to satisfy ``motif`` at (i, j). + + Spacing mismatch counts as +infinity (motif is unreachable at this pair). + """ + if motif is None: + return 0 + spacings = motif.get("spacings") + if spacings and spacings != "any" and (j - i) not in spacings: + return 10**9 + edits = 0 + i_aa = motif.get("i_aa") + j_aa = motif.get("j_aa") + if i_aa and 0 <= i < len(seq) and seq[i] not in i_aa: + edits += 1 + if j_aa and 0 <= j < len(seq) and seq[j] not in j_aa: + edits += 1 + return edits + + def _best_motif_edit_distance( + self, seq: list[str], i: int, j: int + ) -> int: + if not self.catalog: + return 0 + best = min( + (self._motif_edit_distance(seq, i, j, b.motif) for b in self.catalog.values()), + default=0, + ) + return best + + # ------------------------------------------------------------------ + # AnchorPriorBase + # ------------------------------------------------------------------ + def score_anchor( + self, + sequence: list[str], + anchor_pair: tuple[int, int] | None, + context: dict[str, Any] | None = None, + ) -> float: + context = context or {} + components: dict[str, float] = {} + + if anchor_pair is None: + score = -0.2 + components["no_anchor"] = score + if context.get("return_components"): + context["_components"] = components + return score + + i, j = anchor_pair + if i > j: + i, j = j, i + w = self.cfg.weights + score = 0.0 + + # 1. spacing + spacing_ok = (j - i) in self.cfg.allowed_spacings + if spacing_ok: + components["valid_spacing"] = float(w.get("valid_spacing", 1.0)) + else: + components["valid_spacing"] = -float(w.get("valid_spacing", 1.0)) + score += components["valid_spacing"] + + # 2. motif support + edits = self._best_motif_edit_distance(sequence, i, j) + if edits == 0: + components["motif_support"] = float(w.get("existing_motif", 2.0)) + elif edits == 1: + components["motif_support"] = float(w.get("one_edit_motif", 1.2)) + elif edits == 2: + components["motif_support"] = float(w.get("two_edit_motif", 0.4)) + elif edits <= self.cfg.max_motif_edits: + components["motif_support"] = 0.0 + else: + components["motif_support"] = -float(w.get("existing_motif", 2.0)) + score += components["motif_support"] + + # 3. protected positions + protected_pen = 0.0 + if self.cfg.avoid_protected: + protected = set(context.get("protected_positions") or []) + if i in protected or j in protected: + protected_pen = -float(w.get("protected_penalty", 4.0)) + components["protected_penalty"] = protected_pen + score += protected_pen + + # 4. geometry surrogate + geom_score = 0.0 + if self.geometry_oracle is not None and w.get("geometry_surrogate", 0.0) > 0.0: + peptide_ca = context.get("peptide_ca") + block = None + if self.catalog: + # Prefer the block with a matching (i_aa, j_aa) motif; fall + # back to any block for a generic cgeom surrogate. + for b in self.catalog.values(): + if self._motif_edit_distance(sequence, i, j, b.motif) < 10**9: + block = b + break + if block is None: + block = next(iter(self.catalog.values())) + try: + geometry_sequence = sequence + if self.cfg.project_motif_for_geometry and block is not None: + # Project only the two motif tokens needed by this block so + # pre-edit hydrocarbon plans can use the PDB C-alpha window. + motif = block.motif or {} + i_tokens = motif.get("i_aa") or [] + j_tokens = motif.get("j_aa") or [] + if i_tokens and j_tokens: + geometry_sequence = list(sequence) + geometry_sequence[i] = i_tokens[0] + geometry_sequence[j] = j_tokens[0] + cgeom = float( + self.geometry_oracle.cgeom( + geometry_sequence, (i, j), block, peptide_ca=peptide_ca + ) + ) + geom_score = -float(w.get("geometry_surrogate", 1.0)) * cgeom + except Exception: + geom_score = 0.0 + components["geometry_surrogate"] = geom_score + score += geom_score + + if context.get("return_components"): + context["_components"] = components + return float(score) diff --git a/staplebridge/oracles/peptide_prior.py b/staplebridge/oracles/peptide_prior.py new file mode 100644 index 0000000000000000000000000000000000000000..bef15f3349d02ba0ff52bf1fc16de76d2261b9f0 --- /dev/null +++ b/staplebridge/oracles/peptide_prior.py @@ -0,0 +1,16 @@ +from __future__ import annotations + +from typing import Any + +from staplebridge.chemistry.state import StapleState +from staplebridge.data.vocab import AMINO_ACIDS +from staplebridge.oracles.base import PeptidePriorBase + + +class MockPeptidePrior(PeptidePriorBase): + def score_transition(self, old_state: StapleState, new_state: StapleState, context: dict[str, Any] | None = None) -> float: + del old_state, context + seq = new_state.sequence_tokens + valid = sum(1 for t in seq if t in AMINO_ACIDS) + length_pen = abs(len(seq) - 12) / 12 + return valid / max(len(seq), 1) - length_pen diff --git a/staplebridge/oracles/prior_factory.py b/staplebridge/oracles/prior_factory.py new file mode 100644 index 0000000000000000000000000000000000000000..7fb50914177c4b6616071691292b8a2fe79591fb --- /dev/null +++ b/staplebridge/oracles/prior_factory.py @@ -0,0 +1,238 @@ +"""Factories for peptide / anchor / block reference priors. + +The training script no longer instantiates ``MockPeptidePrior`` / +``MockAnchorPrior`` / ``MockBlockPrior`` directly. Instead it calls +:func:`build_reference_priors`, which reads the ``reference_priors`` block +from the training YAML and dispatches to the requested backend. + +Supported backends +------------------ +peptide : ``mock`` | ``heuristic`` | ``esm2_delta`` +anchor : ``mock`` | ``motif_support_geometry`` +block : ``mock`` | ``catalog_scored`` + +``strict_no_mock=true`` at the top level forces every backend to be a real +prior; a mock backend triggers :class:`ValueError` (surfaced by the trainer as +a ``SystemExit``). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from staplebridge.data.schemas import BuildingBlock +from staplebridge.oracles.anchor_prior import MockAnchorPrior +from staplebridge.oracles.base import ( + AnchorPriorBase, + BlockPriorBase, + GeometryOracleBase, + PeptidePriorBase, +) +from staplebridge.oracles.block_prior import MockBlockPrior +from staplebridge.oracles.catalog_block_prior import ( + CatalogBlockPrior, + CatalogBlockPriorConfig, +) +from staplebridge.oracles.esm2_peptide_prior import ESM2DeltaPeptidePrior +from staplebridge.oracles.heuristic_peptide_prior import HeuristicPeptidePrior +from staplebridge.oracles.motif_anchor_prior import ( + MotifSupportAnchorConfig, + MotifSupportAnchorPrior, +) +from staplebridge.oracles.peptide_prior import MockPeptidePrior + + +MOCK_PEPTIDE_BACKENDS = {"mock", "heuristic"} +MOCK_ANCHOR_BACKENDS = {"mock"} +MOCK_BLOCK_BACKENDS = {"mock"} + + +@dataclass +class PriorSelection: + peptide: PeptidePriorBase + anchor: AnchorPriorBase + block: BlockPriorBase + peptide_backend: str + anchor_backend: str + block_backend: str + strict_no_mock: bool = False + + @property + def any_mock(self) -> bool: + return ( + self.peptide_backend in MOCK_PEPTIDE_BACKENDS + or self.anchor_backend in MOCK_ANCHOR_BACKENDS + or self.block_backend in MOCK_BLOCK_BACKENDS + ) + + def to_manifest(self) -> dict[str, Any]: + return { + "peptide_backend": self.peptide_backend, + "anchor_backend": self.anchor_backend, + "block_backend": self.block_backend, + "strict_no_mock": bool(self.strict_no_mock), + "peptide_is_mock": self.peptide_backend in MOCK_PEPTIDE_BACKENDS, + "anchor_is_mock": self.anchor_backend in MOCK_ANCHOR_BACKENDS, + "block_is_mock": self.block_backend in MOCK_BLOCK_BACKENDS, + } + + +# --------------------------------------------------------------------------- +# Peptide prior +# --------------------------------------------------------------------------- +def build_peptide_prior( + cfg: dict[str, Any] | None, + sequence_encoder: Any = None, + device: str = "cpu", +) -> tuple[PeptidePriorBase, str]: + cfg = dict(cfg or {}) + backend = str(cfg.get("backend", "mock")).lower() + if backend == "mock": + return MockPeptidePrior(), "mock" + if backend == "heuristic": + return HeuristicPeptidePrior( + target_length=int(cfg.get("target_length", 12)), + length_weight=float(cfg.get("length_weight", 1.0)), + risky_penalty=float(cfg.get("risky_penalty", 0.3)), + hydrophobic_target=float(cfg.get("hydrophobic_target", 0.45)), + hydrophobic_weight=float(cfg.get("hydrophobic_weight", 0.4)), + ), "heuristic" + if backend == "esm2_delta": + prior = ESM2DeltaPeptidePrior( + model_name_or_path=cfg.get("model_name_or_path"), + sequence_encoder=sequence_encoder, + cache_path=cfg.get("cache_path"), + ncaa_policy=cfg.get("ncaa_policy"), + offline=bool(cfg.get("offline", False)), + device=str(cfg.get("device", device)), + temperature=float(cfg.get("temperature", 1.0)), + max_batch_size=int(cfg.get("max_batch_size", 128)), + memory_cache_size=int(cfg.get("memory_cache_size", 65536)), + strict_runtime=bool(cfg.get("strict_runtime", False)), + ) + return prior, "esm2_delta" + raise ValueError( + f"Unknown reference_priors.peptide.backend: {backend!r}. " + f"Expected one of: mock, heuristic, esm2_delta." + ) + + +# --------------------------------------------------------------------------- +# Anchor prior +# --------------------------------------------------------------------------- +def build_anchor_prior( + cfg: dict[str, Any] | None, + geometry_oracle: GeometryOracleBase | None = None, + catalog: dict[str, BuildingBlock] | None = None, +) -> tuple[AnchorPriorBase, str]: + cfg = dict(cfg or {}) + backend = str(cfg.get("backend", "mock")).lower() + if backend == "mock": + return MockAnchorPrior(), "mock" + if backend == "motif_support_geometry": + weights = dict(cfg.get("weights") or {}) + conf = MotifSupportAnchorConfig( + allowed_spacings=tuple(cfg.get("allowed_spacings") or (3, 4)), + max_motif_edits=int(cfg.get("max_motif_edits", 2)), + avoid_protected=bool(cfg.get("avoid_protected", True)), + project_motif_for_geometry=bool(cfg.get("project_motif_for_geometry", False)), + ) + # merge default weights with user overrides + conf.weights.update({k: float(v) for k, v in weights.items()}) + return MotifSupportAnchorPrior( + catalog=catalog, + geometry_oracle=geometry_oracle, + cfg=conf, + ), "motif_support_geometry" + raise ValueError( + f"Unknown reference_priors.anchor.backend: {backend!r}. " + f"Expected one of: mock, motif_support_geometry." + ) + + +# --------------------------------------------------------------------------- +# Block prior +# --------------------------------------------------------------------------- +def build_block_prior( + cfg: dict[str, Any] | None, + catalog: dict[str, BuildingBlock] | None = None, + geometry_oracle: GeometryOracleBase | None = None, +) -> tuple[BlockPriorBase, str]: + cfg = dict(cfg or {}) + backend = str(cfg.get("backend", "mock")).lower() + if backend == "mock": + return MockBlockPrior(), "mock" + if backend == "catalog_scored": + weights = dict(cfg.get("weights") or {}) + conf = CatalogBlockPriorConfig( + max_motif_edits=int(cfg.get("max_motif_edits", 2)), + no_block_penalty=float(cfg.get("no_block_penalty", 0.3)), + no_anchor_penalty=float(cfg.get("no_anchor_penalty", 1.0)), + ) + conf.weights.update({k: float(v) for k, v in weights.items()}) + return CatalogBlockPrior( + catalog=catalog, + geometry_oracle=geometry_oracle, + cfg=conf, + ), "catalog_scored" + raise ValueError( + f"Unknown reference_priors.block.backend: {backend!r}. " + f"Expected one of: mock, catalog_scored." + ) + + +# --------------------------------------------------------------------------- +# Bundle +# --------------------------------------------------------------------------- +def build_reference_priors( + cfg: dict[str, Any] | None, + catalog: dict[str, BuildingBlock] | None, + geometry_oracle: GeometryOracleBase | None, + sequence_encoder: Any = None, + device: str = "cpu", +) -> PriorSelection: + """Build all three priors from the ``reference_priors`` config section. + + ``cfg`` can be ``None`` (meaning "legacy: everything Mock"), a dict with + ``peptide/anchor/block`` sub-blocks, or a dict where those sub-blocks are + missing (they default to ``backend: mock``). + + ``strict_no_mock=true`` raises :class:`ValueError` if any resolved backend + is a mock/heuristic; the trainer surfaces this as :class:`SystemExit`. + """ + cfg = dict(cfg or {}) + strict = bool(cfg.get("strict_no_mock", False)) + peptide, pb = build_peptide_prior( + cfg.get("peptide"), sequence_encoder=sequence_encoder, device=device + ) + anchor, ab = build_anchor_prior( + cfg.get("anchor"), geometry_oracle=geometry_oracle, catalog=catalog + ) + block, bb = build_block_prior( + cfg.get("block"), catalog=catalog, geometry_oracle=geometry_oracle + ) + sel = PriorSelection( + peptide=peptide, + anchor=anchor, + block=block, + peptide_backend=pb, + anchor_backend=ab, + block_backend=bb, + strict_no_mock=strict, + ) + if strict and sel.any_mock: + offending = [] + if pb in MOCK_PEPTIDE_BACKENDS: + offending.append(f"peptide={pb}") + if ab in MOCK_ANCHOR_BACKENDS: + offending.append(f"anchor={ab}") + if bb in MOCK_BLOCK_BACKENDS: + offending.append(f"block={bb}") + raise ValueError( + "reference_priors.strict_no_mock=true but the following backends " + f"are still mock/heuristic: {', '.join(offending)}. Either switch " + "them to real backends (esm2_delta / motif_support_geometry / " + "catalog_scored) or set strict_no_mock=false for ablations." + ) + return sel diff --git a/staplebridge/reference/__init__.py b/staplebridge/reference/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f3b0d32f802e05f5e1927433a19f527fd5a43b92 --- /dev/null +++ b/staplebridge/reference/__init__.py @@ -0,0 +1 @@ +"""""" diff --git a/staplebridge/reference/energy.py b/staplebridge/reference/energy.py new file mode 100644 index 0000000000000000000000000000000000000000..793e9ea4354c7f05039e6f7df5f84d599f272a92 --- /dev/null +++ b/staplebridge/reference/energy.py @@ -0,0 +1,363 @@ +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 diff --git a/staplebridge/reference/kernel.py b/staplebridge/reference/kernel.py new file mode 100644 index 0000000000000000000000000000000000000000..9c478fd40f384476ea5fcd9ec0ee13c19b33ee94 --- /dev/null +++ b/staplebridge/reference/kernel.py @@ -0,0 +1,144 @@ +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] diff --git a/staplebridge/training/__init__.py b/staplebridge/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f3b0d32f802e05f5e1927433a19f527fd5a43b92 --- /dev/null +++ b/staplebridge/training/__init__.py @@ -0,0 +1 @@ +"""""" diff --git a/staplebridge/training/losses.py b/staplebridge/training/losses.py new file mode 100644 index 0000000000000000000000000000000000000000..825c95d53857b147d513ca4e53261ff48108a5f2 --- /dev/null +++ b/staplebridge/training/losses.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import torch + +from staplebridge.training.trajectory import WeightedTrajectory + + +def assign_lead_local_positive_weights( + trajectory_groups: list[list[WeightedTrajectory]], +) -> dict[str, float | int]: + """Assign lead-equal, within-lead Boltzmann weights to positive paths. + + Exact committed-plan completions are normal behavior-cloning positives. + Off-plan and unfinished paths remain explicit failures with zero positive + weight; teaching their chosen actions as positives would contradict strict + committed-plan inference. A separate failure-avoidance objective may be + added later without changing this positive-path definition. + """ + if not trajectory_groups: + return { + "n_leads": 0, + "n_positive": 0, + "n_failure": 0, + "n_leads_without_positive": 0, + "total_positive_weight": 0.0, + } + lead_scale = 1.0 / float(len(trajectory_groups)) + n_positive = 0 + n_failure = 0 + n_without_positive = 0 + for group in trajectory_groups: + for trajectory in group: + trajectory.weight = 0.0 + positives = [trajectory for trajectory in group if trajectory.is_positive] + n_positive += len(positives) + n_failure += len(group) - len(positives) + if not positives: + n_without_positive += 1 + continue + energies = torch.tensor( + [trajectory.terminal_energy for trajectory in positives], + dtype=torch.float64, + ) + weights = torch.softmax(-energies, dim=0).tolist() + for trajectory, weight in zip(positives, weights): + trajectory.weight = lead_scale * float(weight) + return { + "n_leads": len(trajectory_groups), + "n_positive": n_positive, + "n_failure": n_failure, + "n_leads_without_positive": n_without_positive, + "total_positive_weight": float( + sum(trajectory.weight for group in trajectory_groups for trajectory in group) + ), + } + + +def compute_weighted_path_loss(trajs: list[WeightedTrajectory], log_probs: list[torch.Tensor]) -> torch.Tensor: + device = log_probs[0].device if log_probs else torch.device("cpu") + weights = torch.tensor([t.weight for t in trajs], dtype=torch.float32, device=device) + if weights.numel() == 0: + return torch.tensor(0.0, requires_grad=True, device=device) + total = torch.tensor(0.0, dtype=torch.float32, device=device) + offset = 0 + for i, traj in enumerate(trajs): + n = len(traj.steps) + if n == 0: + continue + lp = torch.stack(log_probs[offset : offset + n]).sum() + total = total - weights[i] * lp + offset += n + return total diff --git a/staplebridge/training/main_loop.py b/staplebridge/training/main_loop.py new file mode 100644 index 0000000000000000000000000000000000000000..96b29b43b2ed6a32ae4e734acde44049bcda0b5b --- /dev/null +++ b/staplebridge/training/main_loop.py @@ -0,0 +1,953 @@ +"""StapleBridge training epoch and per-epoch validation. + +The Full Exact-SB formulation: + +* finite lead-specific feasible plan support H(x) via + ``property_free_hard_plan_support`` (chemistry, geometry, edit budget, + sequence identity, protected positions and exact committed-plan completion); +* ``q_ref(p|x)`` preserved on the legal plan space then conditioned on the hard + mask; +* exact finite-support teacher ``q*(p|x) ~ q_ref(p|x) exp(-beta E_T)``; +* amortized plan controller ``q_theta(p|x)`` masked and renormalised on the + same support; +* ``L_plan = KL(q* || q_theta)``; +* conditional execution policy trained on canonical-completing trajectories + with lead-local positive path weights. + +``validate_enabled`` computes the per-epoch validation metrics, including +``q_star_vs_q_theta_kl``, the checkpoint-selection metric. Post-hoc evaluation +reporting is not part of this training release. +""" +from __future__ import annotations + +import json +import random +import time +from collections import Counter +from typing import Any + +import numpy as np +import torch + +from staplebridge.chemistry.state import StapleState +from staplebridge.hydrocarbon.curriculum import build_hydrocarbon_demonstration_path +from staplebridge.hydrocarbon.exact_sb_cache import ( + energy_only_from_config, resolve_exact_sb_target, +) +from staplebridge.hydrocarbon.plan_control import ( + HydrocarbonPlanControlConfig, + best_committed_plan_trajectory, + HydrocarbonPlanHead, + completes_committed_plan, + controlled_plan_log_probabilities, + configured_plan_level_objective, + describe_plan, + empirical_log_probabilities, + hierarchical_plan_ranking_enabled, + mask_and_renormalize_plan_log_probabilities, + plan_entropy, + property_free_hard_plan_support, + sample_committed_plan_trajectory, + sample_distinct_plans, + select_plan_and_trajectory, +) +from staplebridge.hydrocarbon.plan_reference import enumerate_legal_plans +from staplebridge.hydrocarbon.plan_validation import ( + aggregate_q_star_diagnostics, + exact_sb_validation_enabled, + lead_q_star_diagnostics, +) +from staplebridge.hydrocarbon.property_energy import ( + HydrocarbonPropertyEnergyConfig, + HydrocarbonPropertyScorer, + required_original_lead_properties, +) +from staplebridge.hydrocarbon.tokenizer import tokenize_sequence +from staplebridge.models.control_kernel import state_to_features +from staplebridge.training.losses import ( + assign_lead_local_positive_weights, + compute_weighted_path_loss, +) +from staplebridge.training.records import ( + beam_decode, summarize, terminal_record, +) +from staplebridge.training.trajectory import TransitionStep, WeightedTrajectory +from staplebridge.utils.profiling import STAGE_TIMER + + +def plan_mode(plan: Any) -> str: + return f"{plan.ordered_pair}/i,i+{plan.spacing}" + + +def legal_plans(lead: Any, stack: dict[str, Any]) -> list[Any]: + return enumerate_legal_plans( + tokenize_sequence(lead.linear_sequence), + stack["catalog"], + protected_positions=lead.protected_positions, + filters=stack["sampler"].filters, + ) + + +def property_free_supported_plan_view( + lead: Any, + initial_state: StapleState, + stack: dict[str, Any], + config: dict[str, Any], +) -> dict[str, Any]: + """Build the one hard-constrained coarse-plan support used everywhere.""" + all_plans = legal_plans(lead, stack) + context = { + "protected_positions": lead.protected_positions, + "peptide_ca": (lead.target_context or {}).get("peptide_ca"), + } + all_reference_weights = stack["sampler"].plan_selection_weights( + initial_state, all_plans, context + ) + support_indices, verdicts = property_free_hard_plan_support( + initial_state=initial_state, + lead=lead, + plans=all_plans, + catalog=stack["catalog"], + catalog_index=stack["catalog_index"], + geometry=stack["geometry"], + config=config, + ) + return { + "all_plans": all_plans, + "all_reference_weights": all_reference_weights, + "support_indices": support_indices, + "plans": [all_plans[index] for index in support_indices], + # Preserve q_ref on the original legal space, then condition it on the + # hard mask. Do not recompute empirical within-mode counts afterward. + "reference_weights": [ + all_reference_weights[index] for index in support_indices + ], + "verdicts": verdicts, + "context": context, + } + + +def plan_key(plan: Any) -> tuple[Any, ...]: + return (plan.ordered_pair, plan.spacing, tuple(plan.anchor_pair), plan.block_id) + + +def encode_lead(policy: Any, kernel: Any, state: StapleState) -> torch.Tensor: + features = state_to_features(state, kernel.block_to_idx) + device = kernel.device + features = {key: value.to(device) for key, value in features.items()} + with torch.no_grad(): + return policy.encoder(**features).detach() + + +def selected_top1_rates(selected: list[dict[str, Any]]) -> dict[str, float]: + n = max(len(selected), 1) + return { + "top1_chemistry_valid_rate": sum(bool(row["chemistry_valid"]) for row in selected) / n, + "top1_stapled_rate": sum(bool(row["stapled"]) for row in selected) / n, + } + + +def train_enabled_epoch( + leads: list[Any], + config: dict[str, Any], + stack: dict[str, Any], + energy_fn: Any, + policy: Any, + kernel: Any, + optimizer: Any, + parameters: list[torch.nn.Parameter], + head: HydrocarbonPlanHead, + plan_cfg: HydrocarbonPlanControlConfig, + plan_rng: random.Random, + epoch: int, + exact_sb_cache: Any = None, +) -> tuple[list[dict[str, Any]], dict[str, Any]]: + train_cfg = config["training"] + horizon = int(train_cfg["horizon"]) + chunk_size = int(train_cfg["chunk_size"]) + rows_all: list[dict[str, Any]] = [] + losses: list[float] = [] + path_losses: list[float] = [] + plan_losses: list[float] = [] + entropies: list[float] = [] + coverage: list[float] = [] + disagreements: list[bool] = [] + exact_reverse_kls: list[float] = [] + exact_forward_objectives: list[float] = [] + exact_q_star_entropies: list[float] = [] + exact_plans_scored = 0 + joint_support_enabled = bool( + getattr( + getattr(energy_fn, "property_cfg", None), + "enable_joint_perm_halflife_support", + False, + ) + ) + joint_support_rows: list[dict[str, Any]] = [] + hard_mask_totals: Counter[str] = Counter() + exact_cache_sources: Counter[str] = Counter() + controlled_top_modes: Counter[str] = Counter() + weighting_totals: Counter[str] = Counter() + # --- profiling (opt-in via config profiling.chunk_timing) ------------ + # Wall-clock accounting only: every perf_counter below is read outside the + # computation it brackets and never feeds a tensor, and no RNG is drawn, so + # the trained model, loss and sampling stream are bit-for-bit identical + # whether profiling is on or off. + # + # CUDA attribution: the per-lead GPU stages (plan/qtheta, terminal/PeptiVerse, + # fine rollout's ESM2 forward) each END in a blocking device->host transfer + # (`.item()` / `.cpu()` / `.tolist()`) that already synchronises at the stage + # boundary, so their wall-clock is accurate without an added sync. The three + # chunk-level GPU stages (policy log-prob, backward, optimizer) have no such + # trailing transfer, so they get an explicit ``torch.cuda.synchronize()`` -- + # once per 32-lead chunk, never inside the per-lead or per-step loops. + chunk_timing = bool((config.get("profiling") or {}).get("chunk_timing", False)) + _cuda_sync = chunk_timing and torch.cuda.is_available() + + def _sync() -> None: + if _cuda_sync: + torch.cuda.synchronize() + + STAGE_KEYS = ( + "plan_qtheta", "qstar_target", "fine_rollout", "neighbor_gen", + "terminal_pv", "policy_logprob", "backward", "optimizer", + ) + stage_seconds: dict[str, float] = {key: 0.0 for key in STAGE_KEYS} + _profiling_scorer = getattr(energy_fn, "property_scorer", None) + _pv_wrapper = getattr(_profiling_scorer, "predictor", None) + _esm2_prev = STAGE_TIMER.snapshot() + epoch_started = time.perf_counter() + total_leads = len(leads) + + for chunk_start in range(0, len(leads), chunk_size): + chunk_leads = leads[chunk_start : chunk_start + chunk_size] + weighted_by_lead: list[list[WeightedTrajectory]] = [] + chunk_rows: list[dict[str, Any]] = [] + chunk_plan_losses: list[torch.Tensor] = [] + chunk_t: dict[str, float] = {key: 0.0 for key in STAGE_KEYS} + chunk_counts = { + "plans": 0, + "legal_plans_before_hard_mask": 0, + "hard_masked_plans": 0, + "leads_without_hard_support": 0, + "trajectories": 0, + "transitions": 0, + "neighbors": 0, + } + _cache_hits0 = exact_sb_cache.stats.hits if exact_sb_cache is not None else 0 + _cache_misses0 = exact_sb_cache.stats.misses if exact_sb_cache is not None else 0 + _pv_misses0 = _pv_wrapper.cache_misses if _pv_wrapper is not None else 0 + _pv_hits0 = _pv_wrapper.cache_hits if _pv_wrapper is not None else 0 + chunk_started = time.perf_counter() + for lead_index, lead in enumerate(chunk_leads, start=chunk_start): + lead_weighted: list[WeightedTrajectory] = [] + z0 = StapleState(sequence_tokens=tokenize_sequence(lead.linear_sequence)) + support_view = property_free_supported_plan_view( + lead, z0, stack, config + ) + all_plans = support_view["all_plans"] + support_indices = support_view["support_indices"] + plans = support_view["plans"] + ref_weights = support_view["reference_weights"] + context = support_view["context"] + chunk_counts["legal_plans_before_hard_mask"] += len(all_plans) + chunk_counts["plans"] += len(plans) + chunk_counts["hard_masked_plans"] += len(all_plans) - len(plans) + hard_mask_totals["legal_plans_before_hard_mask"] += len(all_plans) + hard_mask_totals["hard_supported_plans"] += len(plans) + hard_mask_totals["hard_masked_plans"] += len(all_plans) - len(plans) + if not plans: + chunk_counts["leads_without_hard_support"] += 1 + hard_mask_totals["leads_without_hard_support"] += 1 + if joint_support_enabled: + joint_support_rows.append( + { + "example_id": str(lead.example_id), + "n_hard_supported_plans": 0, + "joint_plan_count": 0, + "joint_nonempty": False, + "joint_fallback": False, + "q_star_support_size": 0, + "status": "no_hard_support", + } + ) + weighted_by_lead.append(lead_weighted) + continue + sampled = sample_distinct_plans( + plans, ref_weights, plan_cfg.plans_per_lead, plan_rng + ) + if not sampled: + raise RuntimeError(f"no sampled legal plan for {lead.example_id}") + key_to_index = {plan_key(plan): index for index, plan in enumerate(plans)} + _t = time.perf_counter() + lead_embedding = encode_lead(policy, kernel, z0) + all_controlled_logp = controlled_plan_log_probabilities( + head, + lead_embedding, + all_plans, + support_view["all_reference_weights"], + len(z0.sequence_tokens), + sequence_tokens=z0.sequence_tokens, + peptide_ca=context["peptide_ca"], + ) + masked_controlled_logp = mask_and_renormalize_plan_log_probabilities( + all_controlled_logp, support_indices + ) + controlled_logp = masked_controlled_logp[ + torch.tensor( + support_indices, + dtype=torch.long, + device=masked_controlled_logp.device, + ) + ] + chunk_t["plan_qtheta"] += time.perf_counter() - _t + joint_train_row: dict[str, Any] | None = None + if plan_cfg.exact_sb_objective: + reference_logp = empirical_log_probabilities( + ref_weights, controlled_logp.device + ) + # Deterministic half of the target (plan order, q_ref, terminal + # energies, log q*) may come from the persistent cache. q_theta + # is always the live `controlled_logp` computed above. + _t = time.perf_counter() + exact_energy_tensor, _, cache_info = resolve_exact_sb_target( + lead=lead, + plans=plans, + reference_log_probabilities=reference_logp, + beta=plan_cfg.exact_sb_beta, + energy_fn=energy_fn, + initial_state=z0, + build_terminal=lambda state, plan: build_hydrocarbon_demonstration_path( + state, plan, stack["catalog"] + )[-1], + cache=exact_sb_cache, + energy_only=energy_only_from_config(config), + ) + exact_cache_sources[str(cache_info["source"])] += 1 + target_support_mask = ( + torch.tensor( + cache_info["target_support_mask"], + dtype=torch.bool, + device=controlled_logp.device, + ) + if cache_info.get("target_support_mask") is not None + else None + ) + if joint_support_enabled: + joint_train_row = { + "example_id": str(lead.example_id), + "n_hard_supported_plans": len(plans), + "joint_plan_count": int(cache_info["joint_plan_count"]), + "joint_nonempty": bool(cache_info["joint_nonempty"]), + "joint_fallback": bool(cache_info["joint_fallback"]), + "q_star_support_size": int(cache_info["q_star_support_size"]), + "selection_semantics": "q_theta_top1_hard_plan", + "status": "scored", + } + joint_support_rows.append(joint_train_row) + chunk_t["qstar_target"] += time.perf_counter() - _t + _t = time.perf_counter() + exact_loss, forward_objective, log_q_star = configured_plan_level_objective( + controlled_logp, + [], + [], + plan_cfg.target_temperature, + exact_sb_objective=True, + reference_log_probabilities=reference_logp, + terminal_energies=exact_energy_tensor, + exact_sb_beta=plan_cfg.exact_sb_beta, + target_support_mask=target_support_mask, + ) + chunk_plan_losses.append(exact_loss) + exact_reverse_kls.append(float(exact_loss.detach().cpu().item())) + exact_forward_objectives.append( + float(forward_objective.detach().cpu().item()) + ) + q_star = log_q_star.exp() + if target_support_mask is None: + # Preserve historical false/fallback arithmetic exactly. + q_star_entropy = -(q_star * log_q_star).sum() + else: + q_star_entropy = -torch.where( + q_star > 0.0, + q_star * log_q_star, + torch.zeros_like(q_star), + ).sum() + exact_q_star_entropies.append( + float(q_star_entropy.detach().cpu().item()) + ) + exact_plans_scored += len(plans) + chunk_t["plan_qtheta"] += time.perf_counter() - _t + reference_top = int(np.argmax(ref_weights)) + controlled_top = int(torch.argmax(controlled_logp).item()) + if joint_train_row is not None: + joint_train_row.update( + { + "selected_plan_index": controlled_top, + "selected_plan": describe_plan(plans[controlled_top]), + "selected_satisfies_joint_condition": bool( + target_support_mask is not None + and target_support_mask[controlled_top].item() + ), + } + ) + disagreements.append(reference_top != controlled_top) + entropies.append(float(plan_entropy(controlled_logp).detach().cpu().item())) + coverage.append(len(sampled) / len(plans)) + controlled_top_modes[ + f"{plans[controlled_top].ordered_pair}/i,i+{plans[controlled_top].spacing}" + ] += 1 + + sampled_indices: list[int] = [] + penetrance_values: list[float] = [] + for trajectory_index, plan in enumerate(sampled): + _t = time.perf_counter() + trajectory = sample_committed_plan_trajectory( + stack["sampler"], + z0, + plan, + lead.protected_positions, + context, + horizon, + ) + chunk_t["fine_rollout"] += time.perf_counter() - _t + chunk_counts["trajectories"] += 1 + trajectory_context = {**context, "hydrocarbon_plan": plan} + steps: list[TransitionStep] = [] + for t, (state, next_state) in enumerate( + zip(trajectory.states[:-1], trajectory.states[1:]) + ): + _t = time.perf_counter() + candidates = stack["graph"].neighbors( + state, protected_positions=lead.protected_positions + ) + chunk_t["neighbor_gen"] += time.perf_counter() - _t + chunk_counts["neighbors"] += len(candidates) + chunk_counts["transitions"] += 1 + chosen = next( + (index for index, candidate in enumerate(candidates) if candidate == next_state), + None, + ) + if chosen is None: + raise RuntimeError("committed transition absent from graph") + steps.append(TransitionStep(state, next_state, candidates, chosen, t)) + terminal = trajectory.states[-1] + _t = time.perf_counter() + row = terminal_record( + z0, + terminal, + lead, + energy_fn, + stack, + epoch=epoch, + lead_index=lead_index, + trajectory_index=trajectory_index, + committed_plan=describe_plan(plan), + committed_plan_completed=trajectory.progress.plan_completed, + path_length=len(steps), + ) + chunk_t["terminal_pv"] += time.perf_counter() - _t + penetrance = row.get( + "hydrocarbon_permeability_penetrance_product_mean" + ) + penetrance_values.append(float(penetrance) if penetrance is not None else -1.0) + sampled_indices.append(key_to_index[plan_key(plan)]) + exact_completion = completes_committed_plan(terminal, plan) + failure_reason = ( + None + if exact_completion + else "off_plan_completion" + if terminal.topology == "stapled" + else "unfinished" + ) + row["trajectory_training_status"] = ( + "positive" if exact_completion else failure_reason + ) + lead_weighted.append( + WeightedTrajectory( + steps=steps, + terminal_state=terminal, + terminal_energy=float(row["terminal_energy"]), + context=trajectory_context, + is_positive=exact_completion, + failure_reason=failure_reason, + ) + ) + chunk_rows.append(row) + weighted_by_lead.append(lead_weighted) + if not plan_cfg.exact_sb_objective: + _t = time.perf_counter() + legacy_loss, _, _ = configured_plan_level_objective( + controlled_logp, + sampled_indices, + penetrance_values, + plan_cfg.target_temperature, + exact_sb_objective=False, + ) + chunk_plan_losses.append(legacy_loss) + chunk_t["plan_qtheta"] += time.perf_counter() - _t + + weight_diagnostics = assign_lead_local_positive_weights(weighted_by_lead) + for key in ("n_positive", "n_failure", "n_leads_without_positive"): + weighting_totals[key] += int(weight_diagnostics[key]) + weighted = [ + trajectory for group in weighted_by_lead for trajectory in group + ] + if not chunk_plan_losses: + rows_all.extend(chunk_rows) + continue + _sync() + _t = time.perf_counter() + log_probs = [ + kernel.log_prob_of( + step.state, + step.candidates, + step.chosen_idx, + step.t, + trajectory.context, + ) + for trajectory in weighted + for step in trajectory.steps + ] + path_loss = compute_weighted_path_loss(weighted, log_probs) + plan_loss = torch.stack(chunk_plan_losses).mean() + total_loss = path_loss + float(plan_cfg.loss_weight) * plan_loss + _sync() + chunk_t["policy_logprob"] += time.perf_counter() - _t + optimizer.zero_grad() + _t = time.perf_counter() + total_loss.backward() + _sync() + chunk_t["backward"] += time.perf_counter() - _t + _t = time.perf_counter() + torch.nn.utils.clip_grad_norm_( + parameters, float(train_cfg.get("grad_clip_norm", 1.0)) + ) + optimizer.step() + _sync() + chunk_t["optimizer"] += time.perf_counter() - _t + losses.append(float(total_loss.detach().cpu().item())) + path_losses.append(float(path_loss.detach().cpu().item())) + plan_losses.append(float(plan_loss.detach().cpu().item())) + rows_all.extend(chunk_rows) + + for key in STAGE_KEYS: + stage_seconds[key] += chunk_t[key] + if chunk_timing: + _esm2_now = STAGE_TIMER.snapshot() + esm2_diff = { + key: _esm2_now.get(key, 0) - _esm2_prev.get(key, 0) + for key in ( + "esm2_forward_time", "esm2_cache_hit", "esm2_cache_miss", + "esm2_forward_batches", "esm2_forward_calls", + ) + } + _esm2_prev = _esm2_now + leads_done = chunk_start + len(chunk_leads) + elapsed = time.perf_counter() - epoch_started + leads_per_min = leads_done / (elapsed / 60.0) if elapsed > 0 else 0.0 + eta_seconds = ( + (elapsed / leads_done) * (total_leads - leads_done) if leads_done else 0.0 + ) + record = { + "epoch": epoch + 1, + "chunk": chunk_start // chunk_size + 1, + "leads_done": leads_done, + "total_leads": total_leads, + "chunk_seconds": round(time.perf_counter() - chunk_started, 3), + "stage_seconds": {key: round(chunk_t[key], 4) for key in STAGE_KEYS}, + "esm2": { + key: (round(value, 4) if "time" in key else int(value)) + for key, value in esm2_diff.items() + }, + "counts": dict(chunk_counts), + "exact_sb_cache": { + "hits": (exact_sb_cache.stats.hits - _cache_hits0) + if exact_sb_cache is not None else 0, + "misses": (exact_sb_cache.stats.misses - _cache_misses0) + if exact_sb_cache is not None else 0, + }, + "peptiverse": { + "predictions": (_pv_wrapper.cache_misses - _pv_misses0) + if _pv_wrapper is not None else 0, + "cache_hits": (_pv_wrapper.cache_hits - _pv_hits0) + if _pv_wrapper is not None else 0, + }, + "leads_per_min": round(leads_per_min, 1), + "epoch_eta_min": round(eta_seconds / 60.0, 1), + } + print(f"[chunk timing] {json.dumps(record, ensure_ascii=False)}", flush=True) + + + metrics = { + "loss": float(np.mean(losses)), + "path_loss": float(np.mean(path_losses)), + "plan_loss": float(np.mean(plan_losses)), + "plan_entropy": float(np.mean(entropies)), + "plan_coverage": float(np.mean(coverage)), + "reference_vs_controlled_plan_disagreement": float(np.mean(disagreements)), + "controlled_plan_mode_mix": dict(controlled_top_modes), + "exact_sb_objective": bool(plan_cfg.exact_sb_objective), + "exact_sb_beta": float(plan_cfg.exact_sb_beta), + "plan_encoder_v2": bool(plan_cfg.plan_encoder_v2), + "legacy_target_temperature": float(plan_cfg.target_temperature), + "legacy_property_only_equivalent_beta": float( + 1.0 + / ( + plan_cfg.target_temperature + * float( + config["hydrocarbon"]["terminal_energy"]["property"][ + "penetrance_weight" + ] + ) + ) + ), + "q_star_vs_q_theta_kl": ( + float(np.mean(exact_reverse_kls)) if exact_reverse_kls else None + ), + "exact_sb_forward_objective": ( + float(np.mean(exact_forward_objectives)) + if exact_forward_objectives + else None + ), + "q_star_entropy": ( + float(np.mean(exact_q_star_entropies)) + if exact_q_star_entropies + else None + ), + "exact_sb_plans_scored": int(exact_plans_scored), + "exact_sb_target_sources": dict(exact_cache_sources), + "property_free_hard_plan_mask": True, + "hard_plan_mask_uses_peptiverse": False, + "hard_plan_mask_totals": dict(hard_mask_totals), + "trajectory_weighting": { + "normalization": "lead_local_positive_only", + **{key: int(value) for key, value in weighting_totals.items()}, + }, + "stage_seconds": {key: float(value) for key, value in stage_seconds.items()}, + } + if joint_support_enabled: + supported_joint_rows = [ + row for row in joint_support_rows if row["status"] == "scored" + ] + solution_leads = sum( + bool(row["joint_nonempty"]) for row in supported_joint_rows + ) + selected_joint = sum( + bool(row["selected_satisfies_joint_condition"]) + for row in supported_joint_rows + ) + metrics["joint_perm_halflife_support"] = { + "enabled": True, + "leads_total": len(joint_support_rows), + "leads_with_hard_support": len(supported_joint_rows), + "leads_without_hard_support": sum( + row["status"] == "no_hard_support" for row in joint_support_rows + ), + "leads_with_joint_solution": solution_leads, + "fallback_leads": sum( + bool(row["joint_fallback"]) for row in supported_joint_rows + ), + "joint_plan_count": sum( + int(row["joint_plan_count"]) for row in supported_joint_rows + ), + "q_star_support_size": sum( + int(row["q_star_support_size"]) for row in supported_joint_rows + ), + "selected_joint_solutions": selected_joint, + "joint_solution_recovery": ( + selected_joint / solution_leads if solution_leads else None + ), + "per_lead": joint_support_rows, + } + return rows_all, metrics + + +@torch.no_grad() +def validate_enabled( + leads: list[Any], + config: dict[str, Any], + stack: dict[str, Any], + energy_fn: Any, + policy: Any, + kernel: Any, + head: HydrocarbonPlanHead, + exact_sb_cache: Any = None, +) -> tuple[list[dict[str, Any]], dict[str, Any]]: + validation_cfg = config["validation"] + selected: list[dict[str, Any]] = [] + entropies: list[float] = [] + coverages: list[float] = [] + disagreements: list[bool] = [] + top1_plans: dict[str, str] = {} + mode_mix: Counter[str] = Counter() + # Diagnostics-only: opt-in exact-SB q* on the same property-free hard + # support used by q_theta. It never changes fine decoding or candidate + # ranking within the committed plan selected below. + exact_sb_diagnostics = exact_sb_validation_enabled(config) + exact_sb_beta = HydrocarbonPlanControlConfig.from_config(config).exact_sb_beta + q_star_rows: list[dict[str, Any]] = [] + q_star_cache_sources: Counter[str] = Counter() + joint_support_enabled = bool( + getattr( + getattr(energy_fn, "property_cfg", None), + "enable_joint_perm_halflife_support", + False, + ) + ) + joint_support_rows: list[dict[str, Any]] = [] + legal_plans_before_hard_mask = 0 + hard_supported_plans = 0 + leads_without_hard_support = 0 + # Wall-clock accounting only; see the note in train_enabled_epoch. + q_star_diagnostics_seconds = 0.0 + + for lead_index, lead in enumerate(leads): + z0 = StapleState(sequence_tokens=tokenize_sequence(lead.linear_sequence)) + support_view = property_free_supported_plan_view(lead, z0, stack, config) + all_plans = support_view["all_plans"] + support_indices = support_view["support_indices"] + plans = support_view["plans"] + ref_weights = support_view["reference_weights"] + context = support_view["context"] + legal_plans_before_hard_mask += len(all_plans) + hard_supported_plans += len(plans) + strict_support_at_lead = bool(plans) + if not plans: + leads_without_hard_support += 1 + if joint_support_enabled: + joint_support_rows.append( + { + "example_id": str(lead.example_id), + "n_hard_supported_plans": 0, + "joint_plan_count": 0, + "joint_nonempty": False, + "joint_fallback": False, + "q_star_support_size": 0, + "selected_satisfies_joint_condition": False, + "status": "no_hard_support", + } + ) + print( + f"[plan validation] leads={lead_index + 1}/{len(leads)} " + "status=no_strict_feasible_plan", + flush=True, + ) + continue + lead_embedding = encode_lead(policy, kernel, z0) + all_controlled_logp = controlled_plan_log_probabilities( + head, + lead_embedding, + all_plans, + support_view["all_reference_weights"], + len(z0.sequence_tokens), + sequence_tokens=z0.sequence_tokens, + peptide_ca=context["peptide_ca"], + ) + masked_controlled_logp = mask_and_renormalize_plan_log_probabilities( + all_controlled_logp, support_indices + ) + controlled_logp = masked_controlled_logp[ + torch.tensor( + support_indices, + dtype=torch.long, + device=masked_controlled_logp.device, + ) + ] + controlled_probabilities = controlled_logp.detach().exp().cpu() + plan_modes = [plan_mode(plan) for plan in plans] + q_theta_metrics = { + "q_theta_entropy": float(plan_entropy(controlled_logp).cpu().item()), + "q_theta_top1_mode": plan_modes[int(torch.argmax(controlled_logp).item())], + "q_theta_s5_s5_i4_probability_mass": float( + sum( + controlled_probabilities[index].item() + for index, mode in enumerate(plan_modes) + if mode == "S5-S5/i,i+4" + ) + ), + "q_theta_r8_s5_i7_probability_mass": float( + sum( + controlled_probabilities[index].item() + for index, mode in enumerate(plan_modes) + if mode == "R8-S5/i,i+7" + ) + ), + } + entropies.append(float(plan_entropy(controlled_logp).cpu().item())) + reference_top = int(np.argmax(ref_weights)) + controlled_top = int(torch.argmax(controlled_logp).item()) + disagreements.append(reference_top != controlled_top) + coverages.append(1.0) + + lead_q_star_row: dict[str, Any] | None = None + if exact_sb_diagnostics: + _diagnostics_started = time.perf_counter() + reference_logp = empirical_log_probabilities( + ref_weights, controlled_logp.device + ) + # Same deterministic half as training, same cache. q_theta below is + # always recomputed from the live head. + plan_energy_tensor, _, cache_info = resolve_exact_sb_target( + lead=lead, + plans=plans, + reference_log_probabilities=reference_logp, + beta=exact_sb_beta, + energy_fn=energy_fn, + initial_state=z0, + build_terminal=lambda state, plan: build_hydrocarbon_demonstration_path( + state, plan, stack["catalog"] + )[-1], + cache=exact_sb_cache, + energy_only=energy_only_from_config(config), + ) + q_star_cache_sources[str(cache_info["source"])] += 1 + target_support_mask = ( + torch.tensor( + cache_info["target_support_mask"], + dtype=torch.bool, + device=controlled_logp.device, + ) + if cache_info.get("target_support_mask") is not None + else None + ) + diagnostics = lead_q_star_diagnostics( + controlled_logp, + reference_logp, + plan_energy_tensor, + exact_sb_beta, + plan_modes=plan_modes, + target_support_mask=target_support_mask, + ) + lead_q_star_row = { + "example_id": str(lead.example_id), + "q_star_top1_plan": describe_plan( + plans[diagnostics["q_star_top1_index"]] + ), + "q_theta_top1_plan": describe_plan( + plans[diagnostics["q_theta_top1_index"]] + ), + "joint_perm_halflife_support_enabled": bool( + cache_info["joint_perm_halflife_support_enabled"] + ), + "joint_plan_count": cache_info["joint_plan_count"], + "joint_nonempty": cache_info["joint_nonempty"], + "joint_fallback": cache_info["joint_fallback"], + "q_star_support_size": int(cache_info["q_star_support_size"]), + **diagnostics, + } + q_star_rows.append(lead_q_star_row) + q_star_diagnostics_seconds += time.perf_counter() - _diagnostics_started + + hierarchical = hierarchical_plan_ranking_enabled(config) + decoded_by_plan: list[tuple[StapleState, float, int] | None] = [None] * len(plans) + decode_indices = [controlled_top] if hierarchical else range(len(plans)) + for plan_index in decode_indices: + plan = plans[plan_index] + decoded = beam_decode( + lead, + stack, + kernel, + int(validation_cfg["horizon"]), + int(validation_cfg["beam_size"]), + committed_plan=plan, + ) + if decoded: + decoded_by_plan[plan_index] = ( + best_committed_plan_trajectory(decoded, plan) if hierarchical else decoded[0] + ) + try: + joint, plan_index, terminal, path_logp, path_length = select_plan_and_trajectory( + controlled_logp.detach().cpu(), decoded_by_plan, hierarchical=hierarchical + ) + except RuntimeError as exc: + raise RuntimeError(f"no decoded plans for {lead.example_id}") from exc + selected_plan = plans[plan_index] + label = describe_plan(selected_plan) + top1_plans[lead.example_id] = label + mode_mix[ + f"{selected_plan.ordered_pair}/i,i+{selected_plan.spacing}" + ] += 1 + selected_record = terminal_record( + z0, + terminal, + lead, + energy_fn, + stack, + lead_index=lead_index, + committed_plan=label, + plan_log_probability=float(controlled_logp[plan_index].cpu().item()), + path_log_probability=path_logp, + joint_log_probability=joint, + path_length=path_length, + ) + selected.append(selected_record) + + print(f"[plan validation] leads={lead_index + 1}/{len(leads)}", flush=True) + + summary = summarize(selected, selected, len(leads)) + summary.update(selected_top1_rates(selected)) + summary.update( + { + "plan_control_enabled": True, + "hierarchical_plan_ranking": hierarchical_plan_ranking_enabled(config), + "plan_entropy": float(np.mean(entropies)) if entropies else None, + "plan_top1": top1_plans, + "plan_coverage": float(np.mean(coverages)) if coverages else 0.0, + "mode_mix": dict(mode_mix), + "reference_vs_controlled_plan_disagreement": ( + float(np.mean(disagreements)) if disagreements else None + ), + "property_free_hard_plan_mask": True, + "hard_plan_mask_uses_peptiverse": False, + "legal_plans_before_hard_mask": int(legal_plans_before_hard_mask), + "hard_supported_plans": int(hard_supported_plans), + "hard_masked_plans": int( + legal_plans_before_hard_mask - hard_supported_plans + ), + "leads_without_hard_support": int(leads_without_hard_support), + } + ) + summary["exact_sb_validation_diagnostics"] = exact_sb_diagnostics + if exact_sb_diagnostics: + # Additive only: no existing validation metric above is overwritten. + summary.update(aggregate_q_star_diagnostics(q_star_rows)) + summary["q_star_per_lead"] = q_star_rows + # Provenance only; not a metric and not used by any selection. + summary["q_star_target_sources"] = dict(q_star_cache_sources) + summary["q_star_diagnostics_seconds"] = float(q_star_diagnostics_seconds) + if joint_support_enabled: + supported_joint_rows = [ + row for row in joint_support_rows if row["status"] == "scored" + ] + solution_leads = sum( + bool(row["joint_nonempty"]) for row in supported_joint_rows + ) + selected_joint = sum( + bool(row["selected_satisfies_joint_condition"]) + for row in supported_joint_rows + ) + summary["joint_perm_halflife_support"] = { + "enabled": True, + "leads_total": len(joint_support_rows), + "leads_with_hard_support": len(supported_joint_rows), + "leads_without_hard_support": sum( + row["status"] == "no_hard_support" for row in joint_support_rows + ), + "leads_with_joint_solution": solution_leads, + "fallback_leads": sum( + bool(row["joint_fallback"]) for row in supported_joint_rows + ), + "joint_plan_count": sum( + int(row["joint_plan_count"]) for row in supported_joint_rows + ), + "q_star_support_size": sum( + int(row["q_star_support_size"]) for row in supported_joint_rows + ), + "selected_joint_solutions": selected_joint, + "joint_solution_recovery": ( + selected_joint / solution_leads if solution_leads else None + ), + "per_lead": joint_support_rows, + } + return selected, summary diff --git a/staplebridge/training/records.py b/staplebridge/training/records.py new file mode 100644 index 0000000000000000000000000000000000000000..fe1d73934f861eb3f3d0312e0837c648472009dd --- /dev/null +++ b/staplebridge/training/records.py @@ -0,0 +1,237 @@ +"""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") + diff --git a/staplebridge/training/stack.py b/staplebridge/training/stack.py new file mode 100644 index 0000000000000000000000000000000000000000..caf810e7ebf757aa9a79f4437eccd28920965cd0 --- /dev/null +++ b/staplebridge/training/stack.py @@ -0,0 +1,339 @@ +"""Construction of the StapleBridge training stack. + +Builds, in dependency order: the hydrocarbon staple catalog, the geometry +oracle, the transition graph, the reference priors (frozen ESM2 peptide prior, +anchor prior, block prior), the reference energy and kernel, the plan-aware +reference sampler, the terminal energy, the policy and value networks, the +controlled kernel, and the optimizer. + +The construction is numerically identical to the run that produced the released +checkpoint. +""" +from __future__ import annotations + +import random +from functools import wraps +from pathlib import Path +from typing import Any, Callable + +import numpy as np +import torch +import yaml + +from staplebridge.chemistry.edit_distance import weighted_edit_distance +from staplebridge.chemistry.state import StapleState +from staplebridge.data.dataset import load_leads +from staplebridge.data.schemas import LeadExample +from staplebridge.hydrocarbon.actions import HydrocarbonTransitionGraph +from staplebridge.hydrocarbon.catalog import hydrocarbon_catalog_from_config +from staplebridge.hydrocarbon.endpoint_prior import ( + EmpiricalHydrocarbonEndpointPrior, + EndpointPriorConfig, +) +from staplebridge.hydrocarbon.geometry import HydrocarbonGeometryOracle +from staplebridge.hydrocarbon.plan_reference import build_plan_aware_sampler +from staplebridge.hydrocarbon.property_energy import ( + HydrocarbonPropertyEnergyConfig, + HydrocarbonPropertyScorer, +) +from staplebridge.hydrocarbon.terminal_energy import ( + HydrocarbonTerminalEnergy, + HydrocarbonTerminalEnergyConfig, +) +from staplebridge.integrations.peptiverse import PeptiVerseWrapper +from staplebridge.models.control_kernel import ControlKernelConfig, ControlledKernel +from staplebridge.models.policy_net import PolicyNet +from staplebridge.models.value_net import ValueNet +from staplebridge.oracles.prior_factory import build_reference_priors +from staplebridge.reference.energy import ReferenceEnergy, ReferenceEnergyConfig +from staplebridge.reference.kernel import ReferenceKernel +from staplebridge.utils.paths import PACKAGE_ROOT + + +def install_embedding_cache(wrapper: PeptiVerseWrapper) -> dict[str, int]: + predictor = wrapper.predictor + if predictor is None: + raise RuntimeError("PeptiVerse predictor unavailable") + counts = {"hits": 0, "misses": 0} + for name in ("wt_embedder", "smiles_embedder", "chemberta_embedder"): + embedder = getattr(predictor, name) + for method_name in ("pooled", "unpooled"): + original: Callable[[str], Any] = getattr(embedder, method_name) + cache: dict[str, Any] = {} + + @wraps(original) + def cached(value: str, _original=original, _cache=cache): + if value in _cache: + counts["hits"] += 1 + return _cache[value] + counts["misses"] += 1 + result = _original(value) + _cache[value] = result + return result + + setattr(embedder, method_name, cached) + return counts + + +def build_stack(config: dict[str, Any], seed: int) -> dict[str, Any]: + hydro = dict(config.get("hydrocarbon") or {}) + catalog = hydrocarbon_catalog_from_config(hydro) + catalog_index = {block.block_id: block for block in catalog} + geometry = HydrocarbonGeometryOracle( + catalog, + sentinel_cgeom=float( + (hydro.get("geometry") or {}).get("sentinel_cgeom", 10.0) + ), + ) + graph = HydrocarbonTransitionGraph( + catalog, max_neighbors=int(config.get("max_neighbors", 128)) + ) + priors = build_reference_priors( + config.get("reference_priors"), + catalog=catalog_index, + geometry_oracle=geometry, + device=str((config.get("training") or {}).get("device", "cpu")), + ) + if (config.get("reference_priors") or {}).get("strict_no_mock"): + ensure_available = getattr(priors.peptide, "ensure_available", None) + if callable(ensure_available): + ensure_available() + reference_energy = ReferenceEnergy( + peptide_prior=priors.peptide, + anchor_prior=priors.anchor, + block_prior=priors.block, + geometry_oracle=geometry, + catalog_index=catalog_index, + config=ReferenceEnergyConfig(**dict(config.get("reference") or {})), + ) + base_kernel = ReferenceKernel( + reference_energy, group_normalize=True, substitution_downweight=0.25 + ) + sampler, plan_cfg = build_plan_aware_sampler( + graph, base_kernel, config, seed=seed, root=PACKAGE_ROOT + ) + + prior_cfg = EndpointPriorConfig.from_dict(hydro.get("endpoint_prior")) + if ( + float(prior_cfg.weight_pair) != 0.0 + or prior_cfg.use_length + or prior_cfg.use_relative_position + or prior_cfg.use_local_context + ): + raise RuntimeError( + "the main objective forbids empirical endpoint terms in terminal energy; " + "pair x spacing belongs only to plan-aware reference" + ) + # The empirical table is consumed by the plan-aware sampler above. Disable + # the terminal object completely so missing generated analysis artefacts do + # not trigger a reload and the same evidence cannot be counted twice. + prior_cfg.enabled = False + endpoint_prior = EmpiricalHydrocarbonEndpointPrior(prior_cfg, root=PACKAGE_ROOT) + return { + "hydro": hydro, + "catalog": catalog, + "catalog_index": catalog_index, + "geometry": geometry, + "graph": graph, + "sampler": sampler, + "plan_cfg": plan_cfg, + "endpoint_prior": endpoint_prior, + "reference_priors": priors, + "reference_priors_manifest": priors.to_manifest(), + } + + +def base_terminal_factory( + config: dict[str, Any], geometry: HydrocarbonGeometryOracle, catalog_index: dict +): + section = dict(config.get("terminal_energy") or {}) + lambda_close = float(section.get("lambda_close", 1.0)) + lambda_edit = float(section.get("lambda_edit", 0.2)) + lambda_cost = float(section.get("lambda_cost", 0.2)) + infeasible = float(section.get("infeasible_penalty", 20.0)) + + def base(z0: StapleState, zt: StapleState, lead: LeadExample): + block = catalog_index.get(zt.block_id) if zt.block_id else None + peptide_ca = (lead.target_context or {}).get("peptide_ca") + ctype = geometry.ctype( + zt.sequence_tokens, zt.anchor_pair, block, peptide_ca=peptide_ca + ) + cgeom = geometry.cgeom( + zt.sequence_tokens, zt.anchor_pair, block, peptide_ca=peptide_ca + ) + edit = weighted_edit_distance(zt, z0) + cost = block.cost_score if block else 1.0 + energy = lambda_close * cgeom + lambda_edit * edit + lambda_cost * cost + if not ctype or zt.topology != "stapled": + energy += infeasible + return float(energy), { + "cgeom": float(cgeom), + "edit": float(edit), + "ctype_ok": bool(ctype), + "cost": float(cost), + } + + return base + + +#: Config keys holding filesystem paths. Relative values are resolved against +#: PACKAGE_ROOT so the release runs correctly from any working directory; +#: absolute values are honoured untouched. This is a packaging concern only -- +#: in the official run every one of these was already an absolute path, so +#: resolution is a no-op there and no numerical behaviour depends on it. +_PATH_KEYS: tuple[tuple[str, ...], ...] = ( + ("data", "root"), + ("reference_priors", "peptide", "model_name_or_path"), + ("reference_priors", "peptide", "cache_path"), + ("hydrocarbon", "plan_control", "exact_sb_cache", "path"), + ("property_predictor", "peptiverse_root"), + ("property_predictor", "classifier_weight_root"), + ("property_predictor", "manifest_path"), + ("property_predictor", "hf_cache_dir"), + ("property_predictor", "esm_model_name_or_path"), + ("property_predictor", "peptideclm_model_name_or_path"), + ("property_predictor", "chemberta_model_name_or_path"), +) + + +def resolve_config_paths(config: dict[str, Any], root: Path | None = None) -> dict[str, Any]: + """Make every relative path in ``config`` absolute against ``root``.""" + root = root or PACKAGE_ROOT + for keys in _PATH_KEYS: + node = config + for key in keys[:-1]: + node = node.get(key) if isinstance(node, dict) else None + if not isinstance(node, dict): + break + else: + raw = node.get(keys[-1]) + if isinstance(raw, str) and raw and not Path(raw).is_absolute(): + node[keys[-1]] = str((root / raw).resolve()) + return config + + +def load_config(path: Path) -> dict[str, Any]: + with path.open("r", encoding="utf-8") as handle: + config = yaml.safe_load(handle) or {} + if not isinstance(config, dict): + raise ValueError("configuration root must be a mapping") + return resolve_config_paths(config) + + +def seed_everything(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + torch.use_deterministic_algorithms(False) + + +def select_leads(path: Path, n: int, max_length: int) -> list[LeadExample]: + leads = [lead for lead in load_leads(path) if len(lead.linear_sequence) <= max_length][:n] + if len(leads) != n: + raise RuntimeError(f"need exactly {n} eligible leads from {path}, found {len(leads)}") + return leads + + +def build_predictor(config: dict[str, Any]) -> tuple[PeptiVerseWrapper, dict[str, int]]: + cfg = dict(config["property_predictor"]) + hf_cache_dir = Path( + str(cfg.get("hf_cache_dir") or (PACKAGE_ROOT / "models/hf")) + ) + wrapper = PeptiVerseWrapper( + peptiverse_root=Path(cfg["peptiverse_root"]), + classifier_weight_root=Path(cfg["classifier_weight_root"]), + manifest_path=Path(cfg["manifest_path"]), + device=str(cfg["device"]), + strict=bool(cfg["strict"]), + uncertainty=bool(cfg.get("uncertainty", False)), + cache_enabled=bool(cfg.get("cache_enabled", True)), + hf_cache_dir=hf_cache_dir, + esm_model_name_or_path=cfg.get("esm_model_name_or_path"), + peptideclm_model_name_or_path=cfg.get( + "peptideclm_model_name_or_path" + ), + chemberta_model_name_or_path=cfg.get("chemberta_model_name_or_path"), + offline=bool(cfg.get("offline", True)), + batch_size=int(cfg.get("batch_size", 32)), + ) + if not wrapper.available or wrapper.predictor is None: + raise RuntimeError("strict PeptiVerse backend is unavailable; fallback is forbidden") + return wrapper, install_embedding_cache(wrapper) + + +def build_energy(config: dict[str, Any], stack: dict[str, Any], scorer: HydrocarbonPropertyScorer) -> HydrocarbonTerminalEnergy: + # The base terminal factory reads the generic terminal section. Map the + # unchanged current hydrocarbon training coefficients into that interface. + base_config = dict(config) + train_cfg = dict(config.get("training") or {}) + base_config["terminal_energy"] = { + "lambda_close": train_cfg.get("lambda_close", 1.0), + "lambda_edit": train_cfg.get("lambda_edit", 0.2), + "lambda_cost": train_cfg.get("lambda_cost", 0.2), + "infeasible_penalty": train_cfg.get("infeasible_penalty", 20.0), + } + base = base_terminal_factory(base_config, stack["geometry"], stack["catalog_index"]) + terminal = dict(stack["hydro"].get("terminal_energy") or {}) + prop_cfg = HydrocarbonPropertyEnergyConfig.from_dict(terminal.get("property")) + return HydrocarbonTerminalEnergy( + base, + stack["catalog"], + stack["endpoint_prior"], + HydrocarbonTerminalEnergyConfig( + invalid_topology_penalty=float(terminal.get("invalid_topology_penalty", 10.0)), + penalize_unstapled=bool(terminal.get("penalize_unstapled", True)), + ), + property_scorer=scorer, + property_config=prop_cfg, + ) + + +class PlanAwareKernelAdapter: + """Expose plan-conditioned probabilities through the generic kernel API.""" + + def __init__(self, plan_kernel: Any) -> None: + self.plan_kernel = plan_kernel + + def reference_logits( + self, + state: StapleState, + candidates: list[StapleState], + context: dict[str, Any] | None = None, + ) -> torch.Tensor: + context = dict(context or {}) + probabilities, _ = self.plan_kernel.plan_probs( + state, + candidates, + context.get("hydrocarbon_plan"), + context=context, + ) + return torch.log( + probabilities.clamp_min(torch.finfo(probabilities.dtype).tiny) + ) + + +def build_models(config: dict[str, Any], stack: dict[str, Any], device: torch.device): + model_cfg = dict(config.get("model") or {}) + train_cfg = dict(config.get("training") or {}) + emb_dim = int(model_cfg.get("emb_dim", 32)) + encoder_cfg = dict(config.get("sequence_encoder") or {}) or None + policy = PolicyNet(emb_dim=emb_dim, sequence_encoder_cfg=encoder_cfg).to(device) + value = ValueNet(emb_dim=emb_dim, sequence_encoder_cfg=encoder_cfg).to(device) + block_to_idx = {"": 0, **{block.block_id: index + 1 for index, block in enumerate(stack["catalog"])}} + kernel = ControlledKernel( + reference_kernel=PlanAwareKernelAdapter(stack["sampler"].kernel), + policy_net=policy, + value_net=value, + block_to_idx=block_to_idx, + cfg=ControlKernelConfig(mode=str(model_cfg.get("mode", "policy_tilt"))), + horizon=int(train_cfg["horizon"]), + ) + optimizer = torch.optim.Adam( + list(policy.parameters()) + list(value.parameters()), + lr=float(train_cfg.get("lr", 1e-3)), + ) + return policy, value, kernel, optimizer + + diff --git a/staplebridge/training/trajectory.py b/staplebridge/training/trajectory.py new file mode 100644 index 0000000000000000000000000000000000000000..f33748441dde302b0a92ac512fac779a2302f241 --- /dev/null +++ b/staplebridge/training/trajectory.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from staplebridge.chemistry.state import StapleState + + +@dataclass +class TransitionStep: + state: StapleState + next_state: StapleState + candidates: list[StapleState] + chosen_idx: int + t: int + + +@dataclass +class WeightedTrajectory: + steps: list[TransitionStep] + terminal_state: StapleState + terminal_energy: float + weight: float = 0.0 + context: dict[str, Any] = field(default_factory=dict) + # Hydrocarbon Exact-SB uses these fields to keep plan-conditioned positive + # paths separate from off-plan/unfinished failures. Defaults preserve all + # legacy callers. + is_positive: bool = True + failure_reason: str | None = None diff --git a/staplebridge/utils/__init__.py b/staplebridge/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f3b0d32f802e05f5e1927433a19f527fd5a43b92 --- /dev/null +++ b/staplebridge/utils/__init__.py @@ -0,0 +1 @@ +"""""" diff --git a/staplebridge/utils/paths.py b/staplebridge/utils/paths.py new file mode 100644 index 0000000000000000000000000000000000000000..96e46108d1ba6602d790e792f3efb6da99095389 --- /dev/null +++ b/staplebridge/utils/paths.py @@ -0,0 +1,19 @@ +"""Filesystem anchors for the released training package. + +``PACKAGE_ROOT`` is the release root (the directory containing ``staplebridge/``, +``configs/`` and ``scripts/``). It is used only to resolve *relative* paths given +in a config -- every absolute path in a config is honoured as-is -- so the +release can be checked out anywhere without editing code. +""" +from __future__ import annotations + +from pathlib import Path + +#: Release root: /../../.. == release/staplebridge_training/ +PACKAGE_ROOT = Path(__file__).resolve().parents[2] + + +def resolve_path(raw: str | Path, root: Path | None = None) -> Path: + """Absolute paths pass through; relative ones resolve against ``root``.""" + path = Path(raw) + return path if path.is_absolute() else (root or PACKAGE_ROOT) / path diff --git a/staplebridge/utils/profiling.py b/staplebridge/utils/profiling.py new file mode 100644 index 0000000000000000000000000000000000000000..cf876cc61f211fd04d99faff7934c13327d802cd --- /dev/null +++ b/staplebridge/utils/profiling.py @@ -0,0 +1,90 @@ +"""Lightweight per-stage timing accumulator. + +A single global ``TimingAccumulator`` collects wall-clock time and call counts +for named stages (e.g. ``reference_logits_time``, ``ESM2_prior_time``, +``PeptiVerse_time``, ``backward_time``, ...). Callers wrap their code in +``with STAGE_TIMER.section("stage_name"):`` and periodically call +``STAGE_TIMER.report_and_reset()`` to print/log a table. + +Cache hit/miss counters (``bump("esm2_cache_hit")``) live on the same object. + +The whole module is process-local and thread-safe enough for our single-process +training loop; no cross-process aggregation is attempted. +""" + +from __future__ import annotations + +import contextlib +import threading +import time +from collections import defaultdict + + +class TimingAccumulator: + def __init__(self) -> None: + self._lock = threading.Lock() + self._time: dict[str, float] = defaultdict(float) + self._calls: dict[str, int] = defaultdict(int) + self._counters: dict[str, int] = defaultdict(int) + + @contextlib.contextmanager + def section(self, name: str): + t0 = time.perf_counter() + try: + yield + finally: + dt = time.perf_counter() - t0 + with self._lock: + self._time[name] += dt + self._calls[name] += 1 + + def add(self, name: str, seconds: float) -> None: + with self._lock: + self._time[name] += float(seconds) + self._calls[name] += 1 + + def bump(self, name: str, amount: int = 1) -> None: + with self._lock: + self._counters[name] += int(amount) + + def snapshot(self) -> dict[str, float]: + with self._lock: + snap: dict[str, float] = {} + for k, v in self._time.items(): + snap[k] = float(v) + snap[f"{k}_calls"] = int(self._calls.get(k, 0)) + for k, v in self._counters.items(): + snap[k] = int(v) + return snap + + def reset(self) -> None: + with self._lock: + self._time.clear() + self._calls.clear() + self._counters.clear() + + def format_table(self, title: str = "timings") -> str: + with self._lock: + rows: list[tuple[str, float, int]] = [] + for k in sorted(self._time.keys()): + rows.append((k, float(self._time[k]), int(self._calls.get(k, 0)))) + counters = dict(self._counters) + lines = [f"[{title}]"] + for name, secs, calls in rows: + per = (secs / calls) if calls else 0.0 + lines.append( + f" {name:<28s} total={secs:>8.3f}s calls={calls:>8d} avg={per*1000:>8.3f}ms" + ) + if counters: + lines.append(" -- counters --") + for k in sorted(counters.keys()): + lines.append(f" {k:<28s} {counters[k]}") + return "\n".join(lines) + + def report_and_reset(self, title: str = "timings") -> str: + s = self.format_table(title) + self.reset() + return s + + +STAGE_TIMER = TimingAccumulator()