diff --git a/.gitignore b/.gitignore index f678b659..d0ed7983 100644 --- a/.gitignore +++ b/.gitignore @@ -233,3 +233,7 @@ initial_dataset_40*/ !src/sampleworks/data/protein_configs.csv .idea + +# local tooling / build artifacts +.continue/ +pyproject.toml.pixi.bak diff --git a/docs/latent_adapter/generate_solution_diagram.py b/docs/latent_adapter/generate_solution_diagram.py new file mode 100644 index 00000000..73ce8e36 --- /dev/null +++ b/docs/latent_adapter/generate_solution_diagram.py @@ -0,0 +1,418 @@ +#!/usr/bin/env python3 +"""Generate a data-flow diagram for the latent-adapter injection solution. + +Mirrors the style of the existing X-ray guidance pipeline figure: light-gray +boxes for components, white plaintext boxes for function calls / file:line +annotations, dotted edges for the reward/target feed. + +Color legend +------------ + green : NEW code (models/latent_adapter.py) -- Step 1 + orange : WIRING (one-line change in an existing file) + gray : UNCHANGED existing pipeline (zero edits) + yellow : the injection seam (inside LatentAdaptedWrapper.featurize) + blue/dashed : Step 2 (future) -- swap the injector, add training + +Usage +----- + pip install graphviz # python binding + # also needs the Graphviz system package providing `dot`: + # macOS: brew install graphviz + # debian: apt-get install graphviz + python generate_solution_diagram.py # -> latent_adapter_solution.png + .gv + +If the `graphviz` python package or `dot` is unavailable, the script still +writes the .gv source; render it yourself with: + dot -Tpng latent_adapter_solution.gv -o latent_adapter_solution.png +""" + +from __future__ import annotations + +import sys + +OUT = "latent_adapter_solution" + +# --- styles ------------------------------------------------------------------ +NEW = dict(style="filled", fillcolor="#d5f5e3", color="#1e8449", shape="box") +WIRING = dict(style="filled", fillcolor="#fdebd0", color="#b9770e", shape="box") +UNCHANGED = dict(style="filled", fillcolor="#eeeeee", color="#888888", shape="box") +SEAM = dict(style="filled", fillcolor="#fcf3cf", color="#b7950b", shape="box") +FUTURE = dict(style="dashed,filled", fillcolor="#d6eaf8", color="#2471a3", shape="box") +NOTE = dict(shape="plaintext", fontsize="9") # white annotation boxes +CLUSTER_ATTRS = dict(style="rounded", color="#bbbbbb", fontsize="11", labeljust="l") + + +def build(include_step2: bool = True, out: str = OUT, formats=("png", "svg")): + try: + from graphviz import Digraph + except ImportError: + _emit_raw_dot() + return + + g = Digraph(out) + # splines="spline" (default) so edge labels render; "ortho" drops them. + g.attr(rankdir="TB", fontname="Helvetica", splines="spline", nodesep="0.45", ranksep="0.6") + g.attr("node", fontname="Helvetica", fontsize="10", margin="0.12,0.08") + g.attr("edge", fontname="Helvetica", fontsize="9", color="#555555") + + # ---------------------------------------------------------------- inputs + with g.subgraph(name="cluster_inputs") as c: + c.attr(label="INPUTS (reused, unchanged)", **CLUSTER_ATTRS) + c.node("struct_in", "args.structure (.pdb / .cif)\natomworks.parse(...)", **UNCHANGED) + c.node("struct", "structure: dict\n(asym_unit, chain_info, __config)", **UNCHANGED) + c.node("density_in", "args.density + resolution\nXMap.fromfile(...)", **UNCHANGED) + c.node( + "reward", + "reward_function = RealSpaceRewardFunction\nholds target density + loss (L1/MSE)\n" + "rewards/real_space_density.py", + **UNCHANGED, + ) + c.node("struct_in_note", "get_reward_function_and_structure()\nguidance_script_utils.py:225", **NOTE) + c.edge("struct_in", "struct") + c.edge("density_in", "reward") + c.edge("struct_in_note", "reward", style="invis") + + # ----------------------------------------------------------- wiring layer + with g.subgraph(name="cluster_wiring") as c: + c.attr(label="WIRING (one-line change)", **CLUSTER_ATTRS) + c.node( + "wire", + "get_model_and_device()\nguidance_script_utils.py:165\n\n" + "model = LatentAdaptedWrapper(\n" + " BoltzWrapper(...), # or Protenix / RF3\n" + ' AttrLatentIO("s"), # "s_trunk" for Protenix/RF3\n' + " AffineInjector()) # Step 1", + **WIRING, + ) + + # -------------------------------------------------- NEW adapter (Step 1) + with g.subgraph(name="cluster_adapter") as c: + c.attr(label="NEW models/latent_adapter.py (Step 1)", **CLUSTER_ATTRS) + c.node( + "wrapper", + "LatentAdaptedWrapper[C]\n(decorator; satisfies FlowModelWrapper)\n" + "fields: inner, latent_io, injector, training_adapter", + **NEW, + ) + c.node( + "io", + 'AttrLatentIO(single_attr)\nread_single / write_single\n' + "the ONLY model-specific knowledge\n" + '(one string: "s" | "s_trunk")', + **NEW, + ) + c.node("injector", "AffineInjector (pseudo-MLP)\ns′ = k·s + b\nk=1,b=0 → identity", **NEW) + c.node("protocols", "Protocols\nLatentIO / LatentInjector\n(Step-2 swap points)", **NEW) + c.edge("wrapper", "io", style="dashed", arrowhead="none", label="uses") + c.edge("wrapper", "injector", style="dashed", arrowhead="none", label="uses") + c.edge("protocols", "injector", style="invis") + + # --------------------------------------------- the injection seam (flow) + with g.subgraph(name="cluster_seam") as c: + c.attr(label="THE SEAM LatentAdaptedWrapper.featurize()", **CLUSTER_ATTRS) + c.node("f1", "1. feats = inner.featurize(structure)\n→ GenerativeModelInput[C]\n(conditioning: s, z, …)", **SEAM) + c.node("f2", "2. s = latent_io.read_single(cond)", **SEAM) + c.node("f3", "3. s′ = injector(s) # k·s + b", **SEAM) + c.node( + "f4", + "4. if not training_adapter:\n s′ = s′.detach()\n" + "# gradient isolation: guidance hits coords only", + **SEAM, + ) + c.node("f5", "5. cond′ = latent_io.write_single(cond, s′)\n# dataclasses.replace → sidecar state preserved", **SEAM) + c.node("f6", "6. return GenerativeModelInput(x_init, cond′)", **SEAM) + for a, b in [("f1", "f2"), ("f2", "f3"), ("f3", "f4"), ("f4", "f5"), ("f5", "f6")]: + c.edge(a, b) + + # ------------------------------------------- downstream (UNCHANGED) + with g.subgraph(name="cluster_down") as c: + c.attr(label="DOWNSTREAM (UNCHANGED — zero edits)", **CLUSTER_ATTRS) + c.node( + "pg", + "PureGuidance.sample()\npure_guidance.py:74\nfeatures = model.featurize(structure) # once\n" + "for i in range(num_steps): sampler.step(… features=features)", + **UNCHANGED, + ) + c.node( + "edm", + "AF3EDMSampler.step()\nedm.py:301-487\nnoisy_state.requires_grad_(True) # coords only", + **UNCHANGED, + ) + c.node( + "mstep", + "model.step(x_t, t, features=features)\nboltz/wrapper.py:874\ncond = features.conditioning ← reads cond′\n" + "(Protenix detaches cond when grad_needed)", + **UNCHANGED, + ) + c.node("guid", "reward.backward() on COORDS\n→ guidance (B,N,3) → guided Euler delta", **UNCHANGED) + c.edge("pg", "edm") + c.edge("edm", "mstep") + c.edge("mstep", "guid") + c.edge("guid", "edm", label="next step", constraint="false", style="dotted") + + # ------------------------------------------------ Step 2 (future) + if include_step2: + with g.subgraph(name="cluster_step2") as c: + c.attr(label="STEP 2 (future — same seam, swap pieces)", **CLUSTER_ATTRS) + c.node("mlp", "MLPInjector / FiLMInjector\nzero-init last layer → identity at start", **FUTURE) + c.node("dens", "DensityInjector\nforward(s, e) e = reward features\n(AlphaSAXS-style conditioning)", **FUTURE) + c.node( + "latentguid", + "LatentGuidance scaler\ntraining_adapter=True\none-step-denoiser proxy\noptimize injector.parameters()", + **FUTURE, + ) + c.edge("mlp", "dens", style="invis") + c.edge("dens", "latentguid", style="invis") + + # ------------------------------------------------ verification ladder + g.node( + "ladder", + "VERIFICATION LADDER\n" + "1 transform correct .............. ✅ test (k≠1)\n" + "2 zero-impact at identity ........ ✅ test (k=1,b=0)\n" + "3 gradient isolation ............. ✅ test (detach/attach)\n" + "4 real model consumes cond′ ....... ⛳ cluster smoke run\n" + "5 injection improves data fit .... \U0001f51c Step 2", + shape="note", + style="filled", + fillcolor="#fef9e7", + fontsize="9", + ) + + # ---------------------------------------------------------- cross edges + g.edge("struct", "wire") + g.edge("wire", "wrapper", label="constructs") + g.edge("wrapper", "f1", label="featurize()") + g.edge("io", "f2", style="dotted", arrowhead="none") + g.edge("injector", "f3", style="dotted", arrowhead="none") + g.edge("f6", "pg", label="GenerativeModelInput[C]\n(conditioning = cond′)") + g.edge("reward", "guid", style="dotted", label="target density") + # Step 2 swaps + if include_step2: + g.edge("injector", "mlp", style="dashed", color="#2471a3", label="swap") + g.edge("reward", "dens", style="dashed", color="#2471a3", label="feeds e") + g.edge("latentguid", "wrapper", style="dashed", color="#2471a3", label="flips training_adapter") + g.edge("guid", "ladder", style="invis") + + written = [] + for fmt in formats: + g.format = fmt + written.append(g.render(out, cleanup=True)) + print("Wrote " + ", ".join(written)) + + +def build_downstream_impact(out: str = OUT + "_downstream", formats=("png", "svg")): + """Upstream is a BLACK BOX; expand what a swapped s' does downstream. + + Every downstream consumer is tagged by impact category: + green INTENDED -- the effect we want + gray INVARIANT -- provably unaffected (derives from atom arrays / x_init, not s) + orange INDIRECT -- same mechanism, but now operating on a shifted trajectory + red RISK -- can break; needs a guard + """ + try: + from graphviz import Digraph + except ImportError: + _emit_raw_dot() + return + + INTENDED = dict(style="filled", fillcolor="#d5f5e3", color="#1e8449", shape="box") + INVARIANT = dict(style="filled", fillcolor="#eeeeee", color="#888888", shape="box") + INDIRECT = dict(style="filled", fillcolor="#fdebd0", color="#b9770e", shape="box") + RISK = dict(style="filled,bold", fillcolor="#f5b7b1", color="#c0392b", shape="box") + + g = Digraph(out) + g.attr(rankdir="TB", fontname="Helvetica", splines="spline", nodesep="0.4", ranksep="0.55") + g.attr("node", fontname="Helvetica", fontsize="10", margin="0.12,0.08") + g.attr("edge", fontname="Helvetica", fontsize="9", color="#555555") + + # --- black box upstream + the single injection point -------------------- + g.node( + "blackbox", + "UPSTREAM (BLACK BOX — assumed correct)\n" + "inner.featurize(structure) → encoder / trunk\n" + "→ conditioning with single representation s [B, tokens, d_s]", + shape="box3d", + style="filled", + fillcolor="#2c3e50", + fontcolor="white", + ) + g.node( + "inject", + "INJECTION (the only change)\n" + "s′ = k·s + b (identical shape / dtype / device)\n" + "detach() in sampling mode", + style="filled", + fillcolor="#d5f5e3", + color="#1e8449", + shape="box", + penwidth="2", + ) + g.edge("blackbox", "inject", label="s") + + # --- INTENDED ------------------------------------------------------------ + with g.subgraph(name="cluster_intended") as c: + c.attr(label="INTENDED EFFECT", **CLUSTER_ATTRS) + c.node("step", "model.step(): diffusion forward\nuses s_trunk=cond.s′ (boltz/wrapper.py:874)\n→ x̂₀ changes", **INTENDED) + c.node("reuse", "cached once, reused across all 200 steps\n→ a CONSTANT bias that compounds over the trajectory", **INTENDED) + c.node("xout", "sampled coordinates / ensemble shift\n(the steering we want)", **INTENDED) + c.edge("step", "xout") + c.edge("reuse", "step", style="dotted") + + # --- INVARIANT (safe) ---------------------------------------------------- + with g.subgraph(name="cluster_inv") as c: + c.attr(label="PRESERVED INVARIANTS (provably unaffected by s′)", **CLUSTER_ATTRS) + c.node("shapes", "tensor shapes / atom counts / x_init\n(s is [B,tokens,d], not coords → no shape change)", **INVARIANT) + c.node("recon", "process_structure_to_trajectory_input\nreconciler · model_atom_array · RewardInputs\n(derive from atom arrays, NOT s)", **INVARIANT) + c.node("bf", "b_factors / occupancies / elements\n(from atom arrays)", **INVARIANT) + c.node("det", "determinism / seeding\n(injection adds no randomness)", **INVARIANT) + c.node("out_shape", "GuidanceOutput / trajectory / save_everything\nSAME shapes, different values", **INVARIANT) + + # --- INDIRECT ------------------------------------------------------------ + with g.subgraph(name="cluster_indirect") as c: + c.attr(label="INDIRECT EFFECTS (mechanism unchanged, inputs shifted)", **CLUSTER_ATTRS) + c.node("reward", "reward on COORDS, grad on coords only\n(cond detached) → structurally identical,\nbut evaluated on shifted x̂₀", **INDIRECT) + c.node("guid", "guidance / guided Euler delta\n(edm.py) → different direction", **INDIRECT) + c.node("fk", "FK steering: resampling weights\n→ particle distribution reweighted", **INDIRECT) + c.edge("reward", "guid") + + # --- RISK ---------------------------------------------------------------- + with g.subgraph(name="cluster_risk") as c: + c.attr(label="RISKS / FAILURE MODES (add a guard)", **CLUSTER_ATTRS) + c.node("nan", "OOD s′ (large k or b) → unstable x̂₀ / NaN\npropagates to reward → whole trajectory\nGUARD: clamp |Δ|, assert_no_nans on step output", **RISK) + c.node("leak", "training_adapter=True during sampling\n→ autograd leak (Boltz step does NOT detach)\nGUARD: detach-in-sampling (already in wrapper)", **RISK) + c.node("collapse", "uniform bias across ensemble dim\n→ diversity collapse (loses flexibility)\nGUARD: per-member / noise-conditioned injection", **RISK) + + # --- propagation edges from injection ------------------------------------ + g.edge("inject", "step", label="s′ via features.conditioning", penwidth="2") + g.edge("inject", "reuse", style="dotted") + g.edge("inject", "shapes", style="dashed", arrowhead="empty", label="no effect") + g.edge("inject", "recon", style="dashed", arrowhead="empty", label="no effect") + g.edge("xout", "reward", label="shifted coords") + g.edge("guid", "step", label="next step", constraint="false", style="dotted") + g.edge("inject", "nan", color="#c0392b") + g.edge("inject", "leak", color="#c0392b", style="dashed") + g.edge("step", "collapse", color="#c0392b", style="dotted") + g.edge("xout", "out_shape", style="dashed", arrowhead="empty") + + # --- legend -------------------------------------------------------------- + g.node( + "legend", + "IMPACT LEGEND\n" + "green INTENDED — desired steering\n" + "gray INVARIANT — provably safe (no s dependence)\n" + "orange INDIRECT — shifted inputs, same mechanism\n" + "red RISK — needs a guard", + shape="note", + style="filled", + fillcolor="#fef9e7", + fontsize="9", + ) + + written = [] + for fmt in formats: + g.format = fmt + written.append(g.render(out, cleanup=True)) + print("Wrote " + ", ".join(written)) + + +def build_final_recommendation(out: str = OUT + "_final", formats=("png", "svg")): + """Finalized design: a DECOUPLED latent-optimization pre-pass + FROZEN sampler. + + See latent_space_optimization.md §4. Targets the post-trunk single rep + (s_trunk); the diffusion sampler is unchanged. ``m`` is never touched. + """ + try: + from graphviz import Digraph + except ImportError: + _emit_raw_dot() + return + + NEW = dict(style="filled", fillcolor="#d5f5e3", color="#1e8449", shape="box") + FROZEN = dict(style="filled", fillcolor="#eeeeee", color="#888888", shape="box") + CHOICE = dict(style="filled", fillcolor="#fcf3cf", color="#b7950b", shape="box") + + g = Digraph(out) + g.attr(rankdir="TB", fontname="Helvetica", splines="spline", nodesep="0.45", ranksep="0.6") + g.attr("node", fontname="Helvetica", fontsize="10", margin="0.12,0.08") + g.attr("edge", fontname="Helvetica", fontsize="9", color="#555555") + + with g.subgraph(name="cluster_pre") as c: + c.attr(label="DECOUPLED LATENT-OPT PRE-PASS (NEW — runs once)", **CLUSTER_ATTRS) + c.node("F", "inner.featurize(structure)\ntrunk runs ONCE → caches s_trunk, z_trunk", **NEW) + c.node("D", "DeltaInjector\ndelta = nn.Parameter(zeros) → identity at init\ns' = s_trunk + delta", **NEW) + c.node("OBJ", "objective(s') (no sampler unroll)", **CHOICE) + c.node("A", "(a) aux head h(s') → observable\nvs experimental target\nsampler NOT called", **NEW) + c.node("B", "(b) ONE model.step(x_t,t,feats) → x̂₀\nreuse real_space_density reward\nsingle denoiser call, not the loop", **NEW) + c.node("UPD", "loss.backward() → grad on delta ONLY\nopt.step() (repeat opt_steps)\nregularize ‖delta‖", **NEW) + c.node("WR", "feats' = write_single(feats,\n (s_trunk+delta).detach())", **NEW) + c.edge("F", "D"); c.edge("D", "OBJ") + c.edge("OBJ", "A"); c.edge("OBJ", "B") + c.edge("A", "UPD"); c.edge("B", "UPD") + c.edge("UPD", "WR", label="converged") + + with g.subgraph(name="cluster_samp") as c: + c.attr(label="FROZEN SAMPLER (UNCHANGED — zero edits)", **CLUSTER_ATTRS) + c.node("PG", "PureGuidance.sample()\nfeatures = feats' (fixed across steps)", **FROZEN) + c.node("EDM", "AF3EDMSampler.step()\nnoisy_state.requires_grad_(True) # COORDS only", **FROZEN) + c.node("MS", "model.step(x_t, t, features=feats')\nreads optimized s_trunk", **FROZEN) + c.node("G", "reward.backward() on coords\n→ guided Euler delta", **FROZEN) + c.edge("PG", "EDM"); c.edge("EDM", "MS"); c.edge("MS", "G") + c.edge("G", "EDM", label="next step", constraint="false", style="dotted") + + g.edge("WR", "PG", label="optimized features", penwidth="2") + + g.node( + "note", + "EASE (post-trunk latent): RF3 / Boltz1 = clean;\n" + "Protenix = bypass step() detach + recompute pair_z;\n" + "Boltz2 = recompute diffusion_conditioning.\n" + "m is NEVER touched — its info is already in s_trunk / z_trunk.", + shape="note", style="filled", fillcolor="#fef9e7", fontsize="9", + ) + g.edge("WR", "note", style="invis") + + written = [] + for fmt in formats: + g.format = fmt + written.append(g.render(out, cleanup=True)) + print("Wrote " + ", ".join(written)) + + +def _emit_raw_dot(): + """Fallback: graphviz python package missing -> write a .gv the user can render.""" + print("graphviz python package not found; writing raw DOT instead.", file=sys.stderr) + print(f"Render with: dot -Tpng {OUT}.gv -o {OUT}.png", file=sys.stderr) + # Minimal raw-DOT mirror of the structure above. + dot = '''digraph latent_adapter_solution { + rankdir=TB; node [shape=box, style=filled, fontname=Helvetica]; + struct [label="structure: dict", fillcolor="#eeeeee"]; + reward [label="RealSpaceRewardFunction\\n(target density)", fillcolor="#eeeeee"]; + wire [label="get_model_and_device()\\nmodel = LatentAdaptedWrapper(\\n BoltzWrapper(...), AttrLatentIO(\\"s\\"), AffineInjector())", fillcolor="#fdebd0"]; + wrapper [label="LatentAdaptedWrapper[C]\\n(satisfies FlowModelWrapper)", fillcolor="#d5f5e3"]; + io [label="AttrLatentIO(\\"s\\"|\\"s_trunk\\")\\nonly model-specific knowledge", fillcolor="#d5f5e3"]; + injector[label="AffineInjector s'=k*s+b\\nk=1,b=0 -> identity", fillcolor="#d5f5e3"]; + seam [label="featurize: read s -> inject -> detach(if sampling) -> write s'", fillcolor="#fcf3cf"]; + pg [label="PureGuidance.sample() / AF3EDMSampler.step()\\nUNCHANGED", fillcolor="#eeeeee"]; + mstep [label="model.step(): cond = features.conditioning (reads s')", fillcolor="#eeeeee"]; + step2 [label="STEP 2: MLP/FiLM/Density injector + LatentGuidance training", fillcolor="#d6eaf8", style="dashed,filled"]; + struct -> wire -> wrapper -> seam -> pg -> mstep; + wrapper -> io [style=dashed,arrowhead=none]; wrapper -> injector [style=dashed,arrowhead=none]; + reward -> mstep [style=dotted,label="target density"]; + injector -> step2 [style=dashed,label="swap"]; +} +''' + with open(f"{OUT}.gv", "w") as fh: + fh.write(dot) + print(f"Wrote {OUT}.gv") + + +if __name__ == "__main__": + # Full map (PNG + SVG) + build(include_step2=True, out=OUT, formats=("png", "svg")) + # Compact current-state view: Step 1 only (PNG + SVG) + build(include_step2=False, out=OUT + "_step1", formats=("png", "svg")) + # Downstream-impact view: upstream as black box, downstream expanded + build_downstream_impact(out=OUT + "_downstream", formats=("png", "svg")) + # Finalized decoupled latent-optimization pre-pass + frozen sampler + build_final_recommendation(out=OUT + "_final", formats=("png", "svg")) diff --git a/docs/latent_adapter/implementation_roadmap.md b/docs/latent_adapter/implementation_roadmap.md new file mode 100644 index 00000000..01444ba6 --- /dev/null +++ b/docs/latent_adapter/implementation_roadmap.md @@ -0,0 +1,382 @@ +# Implementation Roadmap — Legibility Refactor → Latent-Opt Injection → Validation + +> Companion to **`latent_space_optimization.md`** (the verified *what/where* of the +> latent pre-pass: targets, `DeltaInjector`, per-model cache handling) and +> `latent_adapter.py`. +> +> This doc is the **order of operations and the scientific guardrails** around that +> design: +> 1. the legibility refactor that creates a clean insertion seam, +> 2. the latent-opt injection (deferring mechanism details to the companion doc), +> 3. the per-model migration, +> 4. the validation falsifier and the scientific-validity contract that say *when a +> result may be believed*. +> +> Derived from a design review of `guidance_script_utils.py`, `edm.py`, +> `pure_guidance.py`, and `boltz/wrapper.py`. Confidence levels are stated inline. + +--- + +## 0. The frame: intervention depth + +Every way experimental data can change what the model generates sits on one axis — +**where the intervention acts** — running from "after the prior" to "inside the +prior." Capability to *relocate support* increases with depth; so does machinery and +risk. + +| Depth | Mechanism | Relocates support? | Cost / machinery | Status in repo | +|---|---|---|---|---| +| 1 | Output selection (best-of-N / importance-weighted) | No — support-limited | Huge compute; unbiased; any reward | **Not implemented** (should be the baseline) | +| 2 | Coordinate steering (DPS / FK) | No — support-limited | Biased; cheap; needs differentiable reward | **What the code does today** | +| 3 | Latent remapping (`s_trunk` / `z_trunk`) | **Yes** | Per-model; needs manifold regularizer; inference-opt *or* learned | **This work** (see companion doc) | +| 4 | Weight fine-tuning (reward alignment) | Yes, globally | Most expensive; risks forgetting the physics prior | Future | + +Two established conclusions that motivate moving to depth 3 (≈85% confidence, +pending held-out evidence): + +- Depths 1–2 are **support-limited**: they steer/select within what the frozen + trunk already generates. Coordinate guidance (depth 2) is, within support, a + *biased variance-reduction trick* over best-of-N selection (depth 1). +- Only depth ≥3 changes **capability**, because only it edits the *information* in + `s`/`z`. Coordinate guidance cannot inject information absent from `s_trunk`. + +⇒ For the scientifically interesting regime (under-sampled / alternate states), the +work is at depth 3. This roadmap gets there safely. + +--- + +## 1. Phase 0 — Legibility refactor (PREREQUISITE, no behavior change) + +**Why first:** `guidance_script_utils.py` is ~710 lines, of which ~38% is pure +save/serialize/metadata and the scientific flow (`_run_guidance`) does not begin +until **line 427**. Injecting latent-opt into this file means surgery in the middle +of plumbing. Extracting the plumbing first turns the insertion into a one-line +change in an obvious place. This is also a standing code-style preference: keep the +scientific narrative front-and-center, mechanism in its own module. + +### Steps + +1. **Create `runs/outputs.py`** and move (verbatim, no logic change): + - `save_everything` (currently L272) + - `save_trajectory`, `_save_trajectory`, `_save_fk_steering_trajectory` (L66–150) + - `save_losses` (L151), `_write_coords_into_array` (L83) + - `_write_job_metadata` (L604), `get_job_result` (L639), `epoch_seconds` (L634) + +2. **Create `runs/factory.py`** (or, better, fold into the wrappers later) and move: + - `get_model_and_device` (L165) + - the model-specific structure-prep dispatch currently inside `_run_guidance` + (`process_structure_for_boltz` / `annotate_structure_for_*`, the + `"Boltz" in wrapper_class_name` branches). This is the string-dispatch smell + flagged by repo issue #192; relocating it is the first step toward killing it. + +3. **Rewrite `_run_guidance` as a linear top-down narrative** in the surviving file + (rename the file to `runs/guidance.py` if desired). Target shape (~5 readable + steps, science visible on the first screen): + + ```python + def _run_guidance(args, guidance_type, model_wrapper, device): + reward, structure = load_inputs(args, device) # inputs + structure = prepare_for_model(model_wrapper, structure, args) # factory + sampler = build_sampler(args, device, model_wrapper) + scaler = build_step_scaler(args) + guidance = build_guidance(guidance_type, args) + result = guidance.sample(structure, model_wrapper, sampler, scaler, reward) + save_everything(args, result, guidance_type) # outputs + ``` + +4. **Order rule:** entry/orchestration functions at the **top** of each file, + helpers below or imported. (`_run_guidance` at L427-after-100-lines-of-savers is + the anti-pattern.) + +5. **Add an optional `features=None` param** to `PureGuidance.sample` and + `FKSteering.sample`: if provided, skip the internal `model.featurize(structure)` + (currently `pure_guidance.py` L76). Backward-compatible; this is the hook Phase 1 + uses to pass an optimized latent in. + +### Acceptance criteria +- Behavior identical: the only diff is code **relocated**, not changed. Verify by + re-running the 1VME quick-start and diffing outputs. +- `_run_guidance` reads in one screen; "what does this do scientifically" is + answerable from the top of the file. +- No new imports cross the science↔plumbing boundary except the explicit + `from .outputs import save_everything`. + +--- + +## 2. Phase 1 — Latent-opt pre-pass (mechanism: see companion doc §4) + +The latent design itself (targets, `DeltaInjector`, objective choices, per-model +cache handling) is finalized in **`latent_space_optimization.md` §3–4**. This +section only records *how it lands in the refactored code* and *which slice to build +first*. + +### Seam (now obvious after Phase 0) +```python +features = model.featurize(structure) +if args.latent_opt: + features = run_latent_optimization(model, features, reward, reward_inputs, cfg) # NEW pre-pass +result = guidance.sample(structure, model=model, ..., features=features) # unchanged sampler +``` +The pre-pass is **decoupled**: trunk runs once, `delta` is optimized against a +differentiable objective, the optimized latent is detached and handed to the +**unchanged** sampler. Depth-3 (relocate) then composes with optional depth-2 +(local steer) on top. + +### Start slice (corrects an earlier suggestion) +- **Begin with RF3 or Boltz1, `s_trunk`, `DeltaInjector`, objective (a) aux-head** + — per companion doc §3/§4.4, these cache nothing extra, so the optimized latent + flows straight through `step()` with **zero model surgery**. +- **Boltz2 is the hardest case, not the first:** `s_trunk` is baked into the cached + `diffusion_conditioning` (`boltz/wrapper.py` ~L805–826), so a boundary injection + reaches only the `s_trunk` arg, not `q/c/biases`. It requires **recomputing + `diffusion_conditioning`** from the perturbed `s,z`. Do it after the easy models + prove signal. +- **Protenix is medium:** `step()` detaches conditioning latents; bypass that for + the leaf and recompute `pair_z`/`p_lm`/`c_l` if perturbing `z`. + +### Decisions to lock before coding +1. **Target:** `s_trunk` first; add `z_trunk` only when the MSA-information lever is + explicitly needed. +2. **Objective:** (a) aux-head (sampler never called during opt) preferred; (b) + single `model.step` `x̂₀` call as a cheap differentiable decoder if a + coordinate-space density reward is wanted. Neither unrolls the sampler. +3. **Regularizer:** `‖delta‖` penalty + a hard clamp on `‖delta‖` (manifold guard; + see §5). + +--- + +## 3. Phase 2 — Per-model capability (later; the "per-model idea") + +Promote the pre-pass to an **optional capability Protocol** rather than a free +function with model branches: + +```python +@runtime_checkable +class LatentSteerable(Protocol): + def optimize_latent(self, features, reward, reward_inputs, cfg) -> GenerativeModelInput: ... +``` + +- Each wrapper implements its own `s/z → conditioning` rule (Boltz2 recomputes + `diffusion_conditioning`; RF3/Boltz1 pass through; Protenix recomputes its caches). + This is the model-specific part that *cannot* be shared — and the read of + `boltz/wrapper.py` proves it: the relationship between `(s,z)` and what the + denoiser consumes differs per model. +- The **sampler/decoder stays shared and unchanged** (the only thing that should be + shared — it is comparability infrastructure, not the intervention). +- A driver checks `isinstance(model, LatentSteerable)` and uses it if present; models + without the capability simply don't offer latent-opt. No lowest-common-denominator + abstraction, no string dispatch. + +--- + +## 4. Phase 3 — Learnable / amortized remapping (future) + +Resolves "static steering → learnable process." Requires three additions the current +architecture lacks (today the graph is severed every step at `edm.py:420`, and there +is no trainable parameter, optimizer, or training entry point): + +1. **A graph-preserving / truncated-BPTT sampler mode** — if the single-step `x̂₀` + decode proves too biased for the objective gradient. +2. **A `fit`/trainer** (optimizer + dataset loop) distinct from `sample`, calling the + same wrapper capability and backprop-ing into its latent parameters. +3. **An amortized adapter** trained across many `(structure, data)` pairs that + *predicts* the remap, vs. per-target inference-time optimization. + +Defer until per-target latent-opt (Phase 1) shows it beats the baselines in §5. + +--- + +## 5. Validation — the falsifier (MUST precede any claim) + +No architecture matters until you can distinguish a **real recovered ensemble** from +**forward-model overfitting**. Build the judge before the thing it judges. + +- **Held-out signal:** held-out reflections (R-free), and/or blind recovery of a + *known* alternate conformation withheld from the reward. +- **Matched-compute baselines, all scored on held-out data:** + 1. best-of-N **importance-weighted** selection (depth 1), + 2. coordinate guidance (depth 2, current), + 3. latent-opt (depth 3, this work). +- **The decisive test:** *does latent-opt beat matched-compute best-of-N on held-out + fit?* If depth-2 guidance does **not** beat depth-1 selection at equal + forward-pass budget, the guidance layer is adding nothing but search. If depth-3 + does not beat both, the representation remap is not earning its complexity. + +Report compute as forward-pass count, not wall-clock, so the comparison is fair. + +--- + +## 6. Scientific-validity contract (caveats as enforced requirements) + +These are not warnings to remember; encode them so they cannot be silently skipped. + +1. **Forward-model overfitting (the real hazard).** Unregularized optimization + against a differentiable density term will produce non-physical structures that + fit the map. *Defense:* `‖delta‖` regularizer **and** held-out R-free (train-fit + improves while held-out degrades = overfitting detected). The goal is to leave the + prior's *default mode*, **not** its *valid-structure manifold*. +2. **Optimization ≠ sampling.** `argmin_delta L(decode)` is a point estimate; it + collapses to a mode and destroys Boltzmann/population weighting. If relative + populations matter, *sample* the latent (retain noise) rather than optimize to a + point — or report the result as a single state, not an ensemble. +3. **Cross-model comparison is outcome-based, not operation-based.** Latent-opt is a + *different operation* per model (different `s/z` shapes and `→conditioning` + rules). Comparability comes only from the held-out outcome metric (§5), never from + claiming the operations are equivalent. Do **not** assert operational equivalence + across models. +4. **Manifold guard.** Clamp `‖delta‖` to avoid OOD / NaN `x̂₀`. Carry the guards + from the downstream-impact diagram (companion doc §4.4). + +--- + +## 7. Sequencing summary + +``` +Phase 0 Legibility refactor (no behavior change) ← do first; unblocks the rest +Phase 1 Latent-opt pre-pass, RF3/Boltz1, s_trunk, aux-head ← companion doc has the mechanism + ▸ build §5 baselines + held-out harness IN PARALLEL with Phase 1 +Phase 2 Promote to LatentSteerable capability (per-model) +Phase 3 Differentiable-sampler mode + trainer (amortized) ← only if Phase 1 beats baselines +``` + +The legibility refactor is the cheapest, highest-leverage step and a hard +prerequisite. The validation harness (§5) gates everything downstream: it is the +falsifier that tells you whether depth-3 is worth building at all. + +--- + +## 8. Newcomer execution order (feature-first, refactor-earned) + +> **This inverts §1–§7's ordering on purpose.** §0–§7 are the *architecturally* +> correct sequence (refactor → feature). As a **new team member**, the *socially* +> correct sequence is the opposite: ship the additive feature first, earn standing +> and context, and treat the shared-code refactor (Phase 0) as something you earn the +> right to propose — not your opening move (Chesterton's fence: don't move a fence +> before you know why it's there). Each step below has a concrete code anchor and its +> blast radius. + +### Step 1 — Land latent-opt as a self-contained additive module *(blast radius: ~zero)* +- **Anchor (new code):** `src/sampleworks/models/latent_adapter.py` (already present, + untracked). Mechanism is finalized in `latent_space_optimization.md` §4.1 + (`DeltaInjector`, lines 126–133) and §4.3 (decoupled pre-pass, lines 151–161). +- **Anchor (start model — RF3, easiest):** `RF3Wrapper` + (`src/sampleworks/models/rf3/wrapper.py:211`). `step` at `:549` reads + `S_trunk_I=cond.s_trunk` directly (`:598`) — clean passthrough, **no cache to + recompute**. `s_trunk` is a first-class field (`:51`). Add an additive + `optimize_latent(...)` method here; guard so other models are untouched. +- **Alt start (Boltz1, also easy):** `Boltz1Wrapper` + (`src/sampleworks/models/boltz/wrapper.py:967`); `step` at `:1202` reads + `s_trunk=cond.s` (`:1260`), no extra cache. +- **Do NOT start with Boltz2** — `step` at `:846` consumes a *cached* + `diffusion_conditioning` baked at `:820`; the perturbation is partially ignored + unless recomputed. Hardest case, last. + +### Step 1 (extended) — Protenix slice: two blockers, one additive fix +Protenix is the only model with **both** failure modes (verified against +`src/sampleworks/models/protenix/wrapper.py`). When you extend past RF3/Boltz1: + +- **Blocker 1 — `step` detaches the latent under grad** (`protenix/wrapper.py:672–681`): + a deliberate DPS optimization (grad flows only to coordinates, not back through the + no-grad pairformer). Fatal here — a `requires_grad` `s_trunk` leaf is detached at + `:677` on the first op. *RF3/Boltz don't do this* (RF3 passes the leaf straight + through at `rf3/wrapper.py:598`). +- **Blocker 2 — `z` lives in derived caches** (`:600–624`): `pair_z`/`p_lm`/`c_l` are + precomputed from `z` via `prepare_cache`; `z_trunk` reaches the diffusion module + *mostly through them*, not the `z_trunk` arg. Direct analog of Boltz2's + `diffusion_conditioning` baking. +- **Mechanical detail:** `ProtenixConditioning` is + `@dataclass(frozen=True, slots=True)` (`:43`) — write with + `dataclasses.replace(cond, s_trunk=...)`, not in place (Boltz's is mutable). + +**`s_trunk`-only is still the easy first slice on Protenix:** `s_trunk` goes +*directly* to the diffusion module (`:688`) and is **not** in the `pair_z` cache, so +only Blocker 1 applies — **no cache recompute needed**. + +**Additive fix (newcomer-safe — no edit to the shared `step`):** in your decoupled +pre-pass, call the diffusion module directly, mirroring `step` (`:683–693`) *minus* +the detach: +```python +x0 = self.model.diffusion_module.forward( + x_noisy=x_t, t_hat_noise_level=t, + input_feature_dict=cond.features, + s_inputs=cond.s_inputs.detach(), # keep frozen + s_trunk=s0 + delta, # ← un-detached leaf (the variable) + z_trunk=cond.z_trunk.detach(), + pair_z=cond.pair_z, p_lm=cond.p_lm, c_l=cond.c_l, # s-only: caches stay valid +) +``` +This bypasses the detach for exactly the tensor you optimize, touches no shared code, +and is Protenix's equivalent of RF3's "just pass the leaf." + +**`z_trunk` later (two options):** +1. Recompute caches in your decode: + `pair_z = self.model.diffusion_module.diffusion_conditioning.prepare_cache(relp, z0+dz, ...)`, + then `p_lm/c_l = ...atom_attention_encoder.prepare_cache(..., pair_z, ...)` so grad + flows `z → pair_z → diffusion` (Boltz2-style recompute). +2. **Hypothesis to test (~70%):** set `enable_diffusion_shared_vars_cache=False` + (`ProtenixConfig` at `:84` → drives `:600`); caches become `None` (`:626`) and — *if* + `diffusion_module.forward` recomputes conditioning from `z_trunk` when they're + `None` — the `z_trunk` gradient flows automatically (per-call recompute cost). + Verify the `None` path actually recomputes before relying on it. + +**Per-model start order:** RF3/Boltz1 (pass the leaf) → Boltz2 (recompute +`diffusion_conditioning`) → Protenix (bypass detach via direct `forward`; recompute +caches for `z`). + +### Step 2 — The one acceptable edit to shared code *(blast radius: one optional param × 2 methods)* +- **Anchor:** `PureGuidance.sample` at `src/sampleworks/core/scalers/pure_guidance.py:46`; + it calls `features = model.featurize(structure)` at `:74`. Add `features=None`; if + provided, skip the re-featurize. Mirror in `FKSteering.sample` + (`src/sampleworks/core/scalers/fk_steering.py:67`, featurize at `:99`). +- **Framing for review:** "I need to hand the sampler a pre-optimized latent" — + backward-compatible, obviously feature-justified. If you want *zero* shared edits in + PR #1, wrap/subclass instead and propose this param afterward. + +### Step 3 — Build the validation harness *(blast radius: new files only)* +- **Anchor (reward to reuse):** `src/sampleworks/core/rewards/real_space_density.py`. +- **Anchor (baseline + held-out):** add a new script under `scripts/eval/` mirroring + the sweep pattern in `run_grid_search.py`; compare best-of-N IW selection vs. + coordinate guidance vs. latent-opt on held-out R-free. Align with existing + `scripts/eval/EVALUATION.md` and `src/sampleworks/eval/`. +- **Why now:** additive, touches no one's architecture, earns *scientific* + credibility — the §5 falsifier that justifies everything later. + +### Step 4 — Boy-scout only your own path *(blast radius: a local helper in a file you already touch)* +- **Anchor:** the `diffusion_conditioning` call inside `_pairformer_pass` + (`src/sampleworks/models/boltz/wrapper.py:820`, within the method at `:747`). When + you reach the Boltz2 slice, factor out a `_compute_diffusion_conditioning(s, z, ...)` + helper that **both** `featurize` and your `optimize_latent` call. Scoped to exactly + the code your feature forces you to touch — no crusade. + +### Step 5 — Keep a private friction log *(blast radius: none — observe, don't edit)* +Record each pain point with its concrete cost; this converts intuition into evidence. +- **String dispatch:** `src/sampleworks/utils/guidance_script_utils.py:450–480` + (the `if "Protenix"/"RF3"/"Boltz" in wrapper_class_name` chain). +- **Buried entry point:** `_run_guidance` at `guidance_script_utils.py:427`. +- **Scattered config contract:** the `getattr(args, ...)` defaults throughout + `guidance_script_utils.py` and `src/sampleworks/utils/guidance_script_arguments.py`. + +### Step 6 — Raise improvements as questions, aligned to debt the team already owns +- **Anchor:** the existing TODO at `guidance_script_utils.py:443` + (`# See https://github.com/diff-use/sampleworks/issues/192 ...`). The team *already + agrees* the dispatch is debt — align with #192 rather than your own opinion. +- **Anchor (the welcomed extension):** propose your per-model latent work as an + optional capability on the Protocol at `src/sampleworks/models/protocol.py:99` + (`FlowModelWrapper`). The README explicitly invites "new ModelWrappers" and + "differentiable modules for new data modalities" — your Phase 2 `LatentSteerable` + is exactly that, so it lands as *welcomed addition*, not *unsolicited refactor*. +- **Social anchor:** the author (`pyproject.toml` → Karson). Ask before touching + shared code: "I'm adding latent-opt; is there history behind the saving/dispatch + before I'd consider factoring any out?" + +### Step 7 — The Phase 0 legibility refactor: earned, deferred +- **Anchor:** the ~270 plumbing lines in `guidance_script_utils.py` (savers at + `:66–163`, `save_everything` `:272`, `_write_job_metadata` `:604`, `get_job_result` + `:639`) → `runs/outputs.py`. +- **Precondition:** standing earned (Steps 1–3 shipped) + a populated friction log + (Step 5) + author buy-in (Step 6). Then land it in small, behavior-preserving + slices — never as a week-one rewrite. + +**One-line strategy:** *additive feature + validation first (Steps 1–3, your +mandate, ~zero blast radius) → observe and align (Steps 4–6) → propose the shared +refactor only once earned (Step 7).* diff --git a/docs/latent_adapter/latent_adapter_solution.png b/docs/latent_adapter/latent_adapter_solution.png new file mode 100644 index 00000000..9b44622f Binary files /dev/null and b/docs/latent_adapter/latent_adapter_solution.png differ diff --git a/docs/latent_adapter/latent_adapter_solution.svg b/docs/latent_adapter/latent_adapter_solution.svg new file mode 100644 index 00000000..07e08e09 --- /dev/null +++ b/docs/latent_adapter/latent_adapter_solution.svg @@ -0,0 +1,383 @@ + + + + + + +latent_adapter_solution + + +cluster_inputs + +INPUTS  (reused, unchanged) + + +cluster_wiring + +WIRING  (one-line change) + + +cluster_adapter + +NEW  models/latent_adapter.py  (Step 1) + + +cluster_seam + +THE SEAM  LatentAdaptedWrapper.featurize() + + +cluster_down + +DOWNSTREAM  (UNCHANGED — zero edits) + + +cluster_step2 + +STEP 2  (future — same seam, swap pieces) + + + +struct_in + +args.structure (.pdb / .cif) +atomworks.parse(...) + + + +struct + +structure: dict +(asym_unit, chain_info, _<model>_config) + + + +struct_in->struct + + + + + +wire + +get_model_and_device() +guidance_script_utils.py:165 +model = LatentAdaptedWrapper( +    BoltzWrapper(...),   # or Protenix / RF3 +    AttrLatentIO("s"),   # "s_trunk" for Protenix/RF3 +    AffineInjector())    # Step 1 + + + +struct->wire + + + + + +density_in + +args.density + resolution +XMap.fromfile(...) + + + +reward + +reward_function = RealSpaceRewardFunction +holds target density + loss (L1/MSE) +rewards/real_space_density.py + + + +density_in->reward + + + + + +guid + +reward.backward() on COORDS +→ guidance (B,N,3) → guided Euler delta + + + +reward->guid + + +target density + + + +dens + +DensityInjector +forward(s, e)   e = reward features +(AlphaSAXS-style conditioning) + + + +reward->dens + + +feeds e + + + +struct_in_note +get_reward_function_and_structure() +guidance_script_utils.py:225 + + + + +wrapper + +LatentAdaptedWrapper[C] +(decorator; satisfies FlowModelWrapper) +fields: inner, latent_io, injector, training_adapter + + + +wire->wrapper + + +constructs + + + +io + +AttrLatentIO(single_attr) +read_single / write_single +the ONLY model-specific knowledge +(one string: "s" | "s_trunk") + + + +wrapper->io + +uses + + + +injector + +AffineInjector  (pseudo-MLP) +s′ = k·s + b +k=1,b=0 → identity + + + +wrapper->injector + +uses + + + +f1 + +1. feats = inner.featurize(structure) +→ GenerativeModelInput[C] +(conditioning: s, z, …) + + + +wrapper->f1 + + +featurize() + + + +f2 + +2. s = latent_io.read_single(cond) + + + +io->f2 + + + + +f3 + +3. s′ = injector(s)   # k·s + b + + + +injector->f3 + + + + +mlp + +MLPInjector / FiLMInjector +zero-init last layer → identity at start + + + +injector->mlp + + +swap + + + +protocols + +Protocols +LatentIO  /  LatentInjector +(Step-2 swap points) + + + + +f1->f2 + + + + + +f2->f3 + + + + + +f4 + +4. if not training_adapter: +       s′ = s′.detach() +# gradient isolation: guidance hits coords only + + + +f3->f4 + + + + + +f5 + +5. cond′ = latent_io.write_single(cond, s′) +# dataclasses.replace → sidecar state preserved + + + +f4->f5 + + + + + +f6 + +6. return GenerativeModelInput(x_init, cond′) + + + +f5->f6 + + + + + +pg + +PureGuidance.sample() +pure_guidance.py:74 +features = model.featurize(structure)  # once +for i in range(num_steps): sampler.step(… features=features) + + + +f6->pg + + +GenerativeModelInput[C] +(conditioning = cond′) + + + +edm + +AF3EDMSampler.step() +edm.py:301-487 +noisy_state.requires_grad_(True)  # coords only + + + +pg->edm + + + + + +mstep + +model.step(x_t, t, features=features) +boltz/wrapper.py:874 +cond = features.conditioning  ← reads cond′ +(Protenix detaches cond when grad_needed) + + + +edm->mstep + + + + + +mstep->guid + + + + + +guid->edm + + +next step + + + +ladder + + + +VERIFICATION LADDER +1 transform correct .............. ✅ test (k≠1) +2 zero-impact at identity ........ ✅ test (k=1,b=0) +3 gradient isolation ............. ✅ test (detach/attach) +4 real model consumes cond′ ....... ⛳ cluster smoke run +5 injection improves data fit .... 🔜 Step 2 + + + + + +latentguid + +LatentGuidance scaler +training_adapter=True +one-step-denoiser proxy +optimize injector.parameters() + + + + +latentguid->wrapper + + +flips training_adapter + + + diff --git a/docs/latent_adapter/latent_adapter_solution_downstream.png b/docs/latent_adapter/latent_adapter_solution_downstream.png new file mode 100644 index 00000000..36edf12e Binary files /dev/null and b/docs/latent_adapter/latent_adapter_solution_downstream.png differ diff --git a/docs/latent_adapter/latent_adapter_solution_downstream.svg b/docs/latent_adapter/latent_adapter_solution_downstream.svg new file mode 100644 index 00000000..c2cbf5f0 --- /dev/null +++ b/docs/latent_adapter/latent_adapter_solution_downstream.svg @@ -0,0 +1,258 @@ + + + + + + +latent_adapter_solution_downstream + + +cluster_intended + +INTENDED EFFECT + + +cluster_inv + +PRESERVED INVARIANTS  (provably unaffected by s′) + + +cluster_indirect + +INDIRECT EFFECTS  (mechanism unchanged, inputs shifted) + + +cluster_risk + +RISKS / FAILURE MODES  (add a guard) + + + +blackbox + + + + +UPSTREAM  (BLACK BOX — assumed correct) +inner.featurize(structure) → encoder / trunk +→ conditioning with single representation  s   [B, tokens, d_s] + + + +inject + +INJECTION  (the only change) +s′ = k·s + b   (identical shape / dtype / device) +detach() in sampling mode + + + +blackbox->inject + + +s + + + +step + +model.step(): diffusion forward +uses s_trunk=cond.s′  (boltz/wrapper.py:874) +→ x̂₀ changes + + + +inject->step + + +s′ via features.conditioning + + + +reuse + +cached once, reused across all 200 steps +→ a CONSTANT bias that compounds over the trajectory + + + +inject->reuse + + + + + +shapes + +tensor shapes / atom counts / x_init +(s is [B,tokens,d], not coords → no shape change) + + + +inject->shapes + + +no effect + + + +recon + +process_structure_to_trajectory_input +reconciler · model_atom_array · RewardInputs +(derive from atom arrays, NOT s) + + + +inject->recon + + +no effect + + + +nan + +OOD s′ (large k or b) → unstable x̂₀ / NaN +propagates to reward → whole trajectory +GUARD: clamp |Δ|, assert_no_nans on step output + + + +inject->nan + + + + + +leak + +training_adapter=True during sampling +→ autograd leak (Boltz step does NOT detach) +GUARD: detach-in-sampling (already in wrapper) + + + +inject->leak + + + + + +xout + +sampled coordinates / ensemble shift +(the steering we want) + + + +step->xout + + + + + +collapse + +uniform bias across ensemble dim +→ diversity collapse (loses flexibility) +GUARD: per-member / noise-conditioned injection + + + +step->collapse + + + + + +reuse->step + + + + + +out_shape + +GuidanceOutput / trajectory / save_everything +SAME shapes, different values + + + +xout->out_shape + + + + + +reward + +reward on COORDS, grad on coords only +(cond detached) → structurally identical, +but evaluated on shifted x̂₀ + + + +xout->reward + + +shifted coords + + + +bf + +b_factors / occupancies / elements +(from atom arrays) + + + +det + +determinism / seeding +(injection adds no randomness) + + + +guid + +guidance / guided Euler delta +(edm.py) → different direction + + + +reward->guid + + + + + +guid->step + + +next step + + + +fk + +FK steering: resampling weights +→ particle distribution reweighted + + + +legend + + + +IMPACT LEGEND +green  INTENDED — desired steering +gray   INVARIANT — provably safe (no s dependence) +orange INDIRECT — shifted inputs, same mechanism +red    RISK — needs a guard + + + diff --git a/docs/latent_adapter/latent_adapter_solution_step1.png b/docs/latent_adapter/latent_adapter_solution_step1.png new file mode 100644 index 00000000..7238668a Binary files /dev/null and b/docs/latent_adapter/latent_adapter_solution_step1.png differ diff --git a/docs/latent_adapter/latent_adapter_solution_step1.svg b/docs/latent_adapter/latent_adapter_solution_step1.svg new file mode 100644 index 00000000..9b921e0a --- /dev/null +++ b/docs/latent_adapter/latent_adapter_solution_step1.svg @@ -0,0 +1,331 @@ + + + + + + +latent_adapter_solution_step1 + + +cluster_inputs + +INPUTS  (reused, unchanged) + + +cluster_wiring + +WIRING  (one-line change) + + +cluster_adapter + +NEW  models/latent_adapter.py  (Step 1) + + +cluster_seam + +THE SEAM  LatentAdaptedWrapper.featurize() + + +cluster_down + +DOWNSTREAM  (UNCHANGED — zero edits) + + + +struct_in + +args.structure (.pdb / .cif) +atomworks.parse(...) + + + +struct + +structure: dict +(asym_unit, chain_info, _<model>_config) + + + +struct_in->struct + + + + + +wire + +get_model_and_device() +guidance_script_utils.py:165 +model = LatentAdaptedWrapper( +    BoltzWrapper(...),   # or Protenix / RF3 +    AttrLatentIO("s"),   # "s_trunk" for Protenix/RF3 +    AffineInjector())    # Step 1 + + + +struct->wire + + + + + +density_in + +args.density + resolution +XMap.fromfile(...) + + + +reward + +reward_function = RealSpaceRewardFunction +holds target density + loss (L1/MSE) +rewards/real_space_density.py + + + +density_in->reward + + + + + +guid + +reward.backward() on COORDS +→ guidance (B,N,3) → guided Euler delta + + + +reward->guid + + +target density + + + +struct_in_note +get_reward_function_and_structure() +guidance_script_utils.py:225 + + + + +wrapper + +LatentAdaptedWrapper[C] +(decorator; satisfies FlowModelWrapper) +fields: inner, latent_io, injector, training_adapter + + + +wire->wrapper + + +constructs + + + +io + +AttrLatentIO(single_attr) +read_single / write_single +the ONLY model-specific knowledge +(one string: "s" | "s_trunk") + + + +wrapper->io + +uses + + + +injector + +AffineInjector  (pseudo-MLP) +s′ = k·s + b +k=1,b=0 → identity + + + +wrapper->injector + +uses + + + +f1 + +1. feats = inner.featurize(structure) +→ GenerativeModelInput[C] +(conditioning: s, z, …) + + + +wrapper->f1 + + +featurize() + + + +f2 + +2. s = latent_io.read_single(cond) + + + +io->f2 + + + + +f3 + +3. s′ = injector(s)   # k·s + b + + + +injector->f3 + + + + +protocols + +Protocols +LatentIO  /  LatentInjector +(Step-2 swap points) + + + + +f1->f2 + + + + + +f2->f3 + + + + + +f4 + +4. if not training_adapter: +       s′ = s′.detach() +# gradient isolation: guidance hits coords only + + + +f3->f4 + + + + + +f5 + +5. cond′ = latent_io.write_single(cond, s′) +# dataclasses.replace → sidecar state preserved + + + +f4->f5 + + + + + +f6 + +6. return GenerativeModelInput(x_init, cond′) + + + +f5->f6 + + + + + +pg + +PureGuidance.sample() +pure_guidance.py:74 +features = model.featurize(structure)  # once +for i in range(num_steps): sampler.step(… features=features) + + + +f6->pg + + +GenerativeModelInput[C] +(conditioning = cond′) + + + +edm + +AF3EDMSampler.step() +edm.py:301-487 +noisy_state.requires_grad_(True)  # coords only + + + +pg->edm + + + + + +mstep + +model.step(x_t, t, features=features) +boltz/wrapper.py:874 +cond = features.conditioning  ← reads cond′ +(Protenix detaches cond when grad_needed) + + + +edm->mstep + + + + + +mstep->guid + + + + + +guid->edm + + +next step + + + +ladder + + + +VERIFICATION LADDER +1 transform correct .............. ✅ test (k≠1) +2 zero-impact at identity ........ ✅ test (k=1,b=0) +3 gradient isolation ............. ✅ test (detach/attach) +4 real model consumes cond′ ....... ⛳ cluster smoke run +5 injection improves data fit .... 🔜 Step 2 + + + + diff --git a/docs/latent_adapter/latent_space_optimization.md b/docs/latent_adapter/latent_space_optimization.md new file mode 100644 index 00000000..4c9b7e2d --- /dev/null +++ b/docs/latent_adapter/latent_space_optimization.md @@ -0,0 +1,236 @@ +# Latent-space optimization across AF3-style wrappers + +> Consolidated analysis + final recommendation for doing **AF2/OpenFold2-style +> latent-space optimization** (single representation and/or MSA representation) +> on the RF3, Protenix, and Boltz wrappers in `sampleworks`. +> +> Companion to `latent_adapter.py` and the solution diagrams in this folder. +> All upstream claims below were verified by reading the original source +> (OpenFold `evoformer.py`, Protenix `pairformer.py`/`protenix.py`, Boltz +> `trunkv2.py`/`pairformer.py`), not summaries. + +--- + +## 1. The core architectural fact + +In **AF2 / OpenFold2** the single representation and the MSA are *the same +persistent latent*: + +```python +# openfold/model/evoformer.py:984 / :1062 +s = self.linear(m[..., 0, :, :]) # single rep = Linear of MSA row 0 +return m, z, s # m and z carried through all 48 blocks +``` + +In the **AF3 family (RF3 / Protenix / Boltz)** this is no longer true. Two +properties hold in every one of them: + +1. **The MSA module writes only to `z`, never to `s`.** + - Protenix `MSAModule.forward(...) -> z`; the internal `msa_sample` is + discarded (`pairformer.py:916`). Its only coupling to the rest of the + network is `z = z + OuterProductMean(m)` (`pairformer.py:660`). + - Boltz `MSAModule.forward(...) -> z` (`trunkv2.py:669`); `MSALayer` does + `z = z + outer_product_mean(m)` (`trunkv2.py:751`). +2. **The single representation is updated inside the Pairformer, independent of + the MSA**, via pair-biased attention: + - Protenix: `s = s + attention_pair_bias(a=s, z=z); s = s + single_transition(s)` + (`pairformer.py:217-223`). + - Boltz: `s = s + attention(s, z); s = s + transition_s(s)` + (`pairformer.py:107-111`). + - It is seeded from `s_inputs` (which folds in the MSA `profile` + + `deletion_mean` summary), **not** from MSA row 0. + +So `m` is a **transient, mid-trunk** tensor: rebuilt fresh from features every +recycle, used to enrich `z`, then dropped. **MSA reaches the single rep only +indirectly, through `z`.** + +### The reframe that matters + +`return z` is **information transfer, not destruction.** Latent-space +optimization optimizes *information*, and after the MSA module the MSA +information lives in `z` (and then propagates into `s` via the Pairformer). +Therefore the optimization targets are the **post-trunk latents** the wrapper +already caches: + +| AF2 / OpenFold2 lever | Equivalent here | +|---|---| +| optimize single rep `s` | optimize **`s_trunk`** (post-Pairformer single) | +| optimize MSA `m` | optimize **`z_trunk`** (where MSA info is encoded) | +| persistent, separable | entangled in `z`/`s`, exposed only post-trunk | + +`m` itself is the **wrong target**: it sits *upstream* of the +`featurize`→`step` boundary, so touching it forces a trunk re-run; it is rebuilt +each recycle (so an injected value is overwritten); and the trunk runs under +`no_grad`. Reserve direct `m` optimization for on-manifold steering *through* +the MSA nonlinearity (AfDesign / ColabDesign-style backprop-through-trunk with +`N_cycle=1`) — not for this work. + +--- + +## 2. How the pipeline consumes the latent + +``` +PureGuidance.sample() # pure_guidance.py + features = model.featurize(structure) # ← runs the trunk ONCE; caches s_trunk/z_trunk + for i in range(num_steps): + AF3EDMSampler.step(state, model, ctx, features=features) + noisy_state.requires_grad_(True) # edm.py:420 ← COORDS are the grad var + x_hat_0 = model.step(noisy_state, t, features)# edm.py:426 ← latent passed as constant + guidance = reward.backward() on coords # guidance flows through coordinates only +``` + +Two consequences: + +- The latent (`features.conditioning`) is **fixed across all sampling steps**. +- The existing guidance is **coordinate-space** (classifier/DPS-style); it never + differentiates w.r.t. the latent. + +⇒ **Latent optimization must be a separate, decoupled pass that produces an +optimized `features`, after which the sampler runs unchanged.** This is exactly +the seam `latent_adapter.py` already occupies. + +--- + +## 3. Which model is easiest (AF2-style, post-trunk latent) + +| Model | `s_trunk` (single) | `z_trunk` (MSA-info) | direct `m` | Main blocker | **Verdict** | +|---|---|---|---|---|---| +| **RF3** | ✅ clean — `step()` reads it directly, no detach, no extra cache | ✅ clean | ❌ opaque generator trunk, `no_grad`, bf16 | bf16 / EMA-shadow; opaque trunk for `m` | **Easiest** | +| **Boltz1** | ✅ clean — no precomputed cache, no detach | ✅ clean | ⚠️ loop inline (hookable) but `m` rebuilt each recycle | none for post-trunk | **Easiest** | +| **Protenix** | ⚠️ `step()` **detaches** conditioning latents when grad on (`protenix/wrapper.py:676-681`) — bypass for your leaf | ⚠️ + recompute cached `pair_z`/`p_lm`/`c_l` (`:600-624`) | ❌ opaque `get_pairformer_output`, `no_grad` | detach in `step()` + z-derived caches | **Medium** | +| **Boltz2** | ⚠️ **partial** — `s_trunk` baked into cached `diffusion_conditioning` (`boltz/wrapper.py:819-826`); a boundary injection only reaches the `s_trunk` arg, not `q/c/biases` | ❌ must recompute `diffusion_conditioning` (depends on `s` **and** `z`) | ⚠️ loop inline but same rebuild issue | `diffusion_conditioning` cache | **Hardest** | + +`DEFAULT_SINGLE_REP_ATTR` (in `latent_adapter.py`): Boltz=`s`, Protenix/RF3=`s_trunk`. + +The split between "easy" and "hard" is entirely about **speed caches**: Protenix +and Boltz2 hoist diffusion conditioning out of the trunk and cache it so +`step()` is cheap. Those caches must be bypassed (detach) or recomputed for an +optimized latent to actually reach the denoiser. RF3 and Boltz1 cache nothing +extra, so the optimized latent flows straight through. + +--- + +## 4. Final recommendation + +**Optimize the cached post-trunk latent in a decoupled pre-pass; keep diffusion +a frozen sampler.** Concretely: + +### 4.1 Target and mechanism +- **Target:** `s_trunk` first (single rep; reuses the existing adapter seam and + propagates cleanly on RF3/Boltz1). Add `z_trunk` only when the explicit + MSA-information lever is needed. +- **Mechanism:** a **direct additive latent perturbation** — the most direct + form of latent-space optimization: + + ```python + class DeltaInjector(nn.Module): + """Direct latent-space optimization: a free additive perturbation.""" + def __init__(self, shape): + super().__init__() + self.delta = nn.Parameter(torch.zeros(shape)) # identity at init + def forward(self, x): + return x + self.delta + ``` + `delta` starts at zero (preserves the zero-impact guarantee), is the + optimization variable, and is regularized (`‖delta‖`) to stay on-manifold. + Use `training_adapter=True` so it is **not** detached during the optimization + pass. + +### 4.2 Objective (diffusion stays a sampler) +Pick by available signal, neither unrolls the sampler: +- **(a) Auxiliary-head / AlphaSAXS objective (preferred):** a small + differentiable head `h(latent) → predicted observable`, optimized against the + experimental target. The sampler is **never** called during optimization; it + runs once at the end to decode the optimized latent. +- **(b) Single-denoiser-call objective:** call `model.step(x_t, t, features)` + **once** (one `x̂₀` evaluation, not the loop) as a cheap differentiable + decoder, reuse the existing reward (`rewards/real_space_density.py`), backprop + to `delta`. Minimal coupling that still yields coordinate-space rewards. + +### 4.3 Integration (zero edits to the sampler) +``` +LatentOptimizer (NEW, decoupled pre-pass) + feats = inner.featurize(structure) # trunk once + delta = DeltaInjector(s_trunk.shape) # leaf params + for _ in range(opt_steps): + s' = s_trunk + delta + loss = objective(s') # (a) head or (b) one step() call + loss.backward(); opt.step() # grad → delta only + feats' = write_single(feats, (s_trunk+delta).detach()) + return feats' # hand to PureGuidance.sample() unchanged +``` + +### 4.4 Start here +- **RF3 or Boltz1**, `s_trunk`, `DeltaInjector`, objective (a). Zero model + surgery, no trunk re-run, frozen sampler, clean gradient. +- Then `z_trunk` for the MSA-information lever. +- **Guards** (carry over from the downstream-impact diagram): clamp `‖delta‖` to + avoid OOD/NaN `x̂₀`; for **Protenix** skip the `step()` detach for the leaf and + recompute `pair_z`/`p_lm`/`c_l` if perturbing `z`; for **Boltz2** recompute + `diffusion_conditioning` after injection or the perturbation is partially + ignored. + +--- + +## 5. Existing implementations / prior art +- **Backprop-through-a-folding-model to optimize the MSA/sequence input** is what + **ColabDesign / AfDesign** does for AF2 — the precedent for true `m`/input + optimization, and it confirms the full-trunk-backprop cost profile. +- **In this repo:** the only latent-optimization scaffold is the single-rep + injection in `latent_adapter.py`. No `m` path exists; no wrapper exposes `m`. +- Diffusion-model guidance in this space is otherwise coordinate-space + (the existing `PureGuidance` / `AF3EDMSampler` path). + +--- + +## 6. Visualization — the finalized decoupled design + +```mermaid +flowchart TB + subgraph PRE["DECOUPLED LATENT-OPT PRE-PASS (NEW — runs once, before sampling)"] + F["inner.featurize(structure)
trunk runs ONCE → caches s_trunk, z_trunk"] + D["DeltaInjector: delta = nn.Parameter(zeros)
s' = s_trunk + delta"] + OBJ{"objective(s')"} + A["(a) aux head h(s') → observable
vs experimental target
sampler NOT called"] + B["(b) ONE model.step(x_t,t,feats) → x̂₀
reuse real_space_density reward
single denoiser call, not the loop"] + UPD["loss.backward() → grad on delta only
opt.step() (repeat opt_steps)"] + WR["feats' = write_single(feats, (s_trunk+delta).detach())"] + F --> D --> OBJ + OBJ --> A --> UPD + OBJ --> B --> UPD + UPD -->|converged| WR + end + + subgraph SAMP["FROZEN SAMPLER (UNCHANGED — zero edits)"] + PG["PureGuidance.sample()
features = feats' (fixed across steps)"] + EDM["AF3EDMSampler.step()
noisy_state.requires_grad_(True) # COORDS only"] + MS["model.step(x_t, t, features=feats')
reads optimized s_trunk"] + G["reward.backward() on coords → guided Euler delta"] + PG --> EDM --> MS --> G + G -.next step.-> EDM + end + + WR ==>|optimized features| PG + + classDef new fill:#d5f5e3,stroke:#1e8449; + classDef frozen fill:#eeeeee,stroke:#888888; + class F,D,A,B,UPD,WR new; + class PG,EDM,MS,G frozen; +``` + +**Read:** the trunk runs once; `delta` (the additive latent perturbation) is the +only optimization variable; the objective avoids unrolling the sampler; the +optimized latent is detached and handed to the **unchanged** sampler. `m` never +appears — its information is already inside `s_trunk`/`z_trunk`. + +## 7. Diagrams (Graphviz) +A `build_final_recommendation()` was added to `generate_solution_diagram.py` +(render with `pip install graphviz` + system `dot`). See also the existing +figures in this folder: +- `latent_adapter_solution[_step1].svg` — the injection seam & pipeline. +- `latent_adapter_solution_downstream.svg` — what a swapped `s'` affects + downstream (intended / invariant / indirect / risk). + +The **Step 2** nodes in that diagram (MLP/FiLM/Density injector, `LatentGuidance` +with `training_adapter=True`, one-step-denoiser proxy) are precisely the +`DeltaInjector` + decoupled-objective design finalized in §4. diff --git a/src/sampleworks/models/latent_adapter.py b/src/sampleworks/models/latent_adapter.py new file mode 100644 index 00000000..62a57935 --- /dev/null +++ b/src/sampleworks/models/latent_adapter.py @@ -0,0 +1,191 @@ +"""Latent-space injection adapter for experimentally-guided sampling. + +This module inserts a small, swappable transform on a model's *single +representation* (the post-trunk latent ``s``) at the conditioning boundary +between :meth:`featurize` and :meth:`step`. It is the Step-1 scaffold for an +AlphaSAXS-style latent injection (see project notes): the experimental signal +will eventually drive this transform, but here the transform is a deliberately +trivial affine ``k * s + b`` ("pseudo MLP") so the plumbing and tests can be +validated before any real module is introduced. + +Design goals +------------ +- **Minimal change.** Everything lives behind :class:`LatentAdaptedWrapper`, + which itself satisfies the ``FlowModelWrapper`` protocol. The samplers, + scalers, eval, and their tests consume the protocol, so nothing downstream + changes. With an identity transform (``k=1, b=0``) the wrapper is + behaviourally identical to the model it wraps. +- **Heterogeneous across models.** The only model-specific knowledge is *which + attribute of the conditioning object holds the single representation*. That is + a single string (``"s"`` for Boltz, ``"s_trunk"`` for Protenix/RF3), supplied + via :class:`AttrLatentIO`. No model package is imported here, so a 4th or 5th + model is one more string, not new code. +- **Gradient isolation.** During sampling the injected latent is detached so the + existing coordinate-only guidance gradient (and all downstream tests) behave + exactly as before. Training the transform happens in a separate pass (Step 2). +""" + +from __future__ import annotations + +import dataclasses +from typing import Generic, Protocol, runtime_checkable, TypeVar + +import torch +from torch import nn, Tensor + +from sampleworks.models.protocol import GenerativeModelInput + +C = TypeVar("C") + +# Convenience map of the single-representation attribute name per model. +# This is documentation/config only -- the adapter never imports these models. +DEFAULT_SINGLE_REP_ATTR: dict[str, str] = { + "boltz1": "s", + "boltz2": "s", + "protenix": "s_trunk", + "rf3": "s_trunk", +} + + +@runtime_checkable +class LatentIO(Protocol[C]): + """Reads/writes the single representation on a model's conditioning object. + + Implementations encapsulate the *only* model-specific knowledge in this + module: where the single representation lives on conditioning ``C``. + """ + + def read_single(self, conditioning: C) -> Tensor | None: + """Return the single representation tensor, or ``None`` if unavailable.""" + ... + + def write_single(self, conditioning: C, single: Tensor) -> C: + """Return a copy of ``conditioning`` with the single representation replaced.""" + ... + + +class AttrLatentIO: + """Generic :class:`LatentIO` that addresses the single rep by attribute name. + + Works for any (dataclass) conditioning whose single representation is stored + in a single attribute. ``write_single`` uses :func:`dataclasses.replace`, so + non-tensor sidecar state on the conditioning (e.g. RF3's chiral tracking + arrays) is preserved untouched. + + Parameters + ---------- + single_attr + Name of the attribute holding the single representation, e.g. ``"s"`` + for Boltz or ``"s_trunk"`` for Protenix/RF3. + """ + + def __init__(self, single_attr: str): + self.single_attr = single_attr + + def read_single(self, conditioning: C) -> Tensor | None: + if conditioning is None: + return None + return getattr(conditioning, self.single_attr, None) + + def write_single(self, conditioning: C, single: Tensor) -> C: + if dataclasses.is_dataclass(conditioning) and not isinstance(conditioning, type): + return dataclasses.replace(conditioning, **{self.single_attr: single}) # ty: ignore + # Fallback for non-dataclass conditioning: mutate a shallow copy. + import copy + + new_cond = copy.copy(conditioning) + setattr(new_cond, self.single_attr, single) + return new_cond + + +@runtime_checkable +class LatentInjector(Protocol): + """A transform applied to the single representation. Swappable in Step 2.""" + + def __call__(self, single: Tensor) -> Tensor: + """Map a single representation to a modified single representation.""" + ... + + +class AffineInjector(nn.Module): + """Step-1 "pseudo MLP": an elementwise affine transform ``k * s + b``. + + With the defaults (``k=1, b=0``) this is the identity, which makes + :class:`LatentAdaptedWrapper` behaviourally identical to the wrapped model -- + the basis of the zero-impact guarantee. ``k`` and ``b`` are learnable + :class:`~torch.nn.Parameter` s so that Step 2 can either optimize them + directly or replace this module wholesale with an MLP/FiLM head. + + Parameters + ---------- + k_init, b_init + Initial scale and shift. Defaults give the identity transform. + """ + + def __init__(self, k_init: float = 1.0, b_init: float = 0.0): + super().__init__() + self.k = nn.Parameter(torch.tensor(float(k_init))) + self.b = nn.Parameter(torch.tensor(float(b_init))) + + def forward(self, single: Tensor) -> Tensor: + return self.k * single + self.b + + +class LatentAdaptedWrapper(Generic[C]): # noqa: UP046 (Python 3.11 compatibility) + """Decorator over any ``FlowModelWrapper`` that injects into the latent. + + Satisfies the ``FlowModelWrapper`` protocol itself, so it is a drop-in + replacement for the model it wraps. Only :meth:`featurize` is augmented; + ``step`` and ``initialize_from_prior`` delegate verbatim. + + Parameters + ---------- + inner + The wrapped model wrapper (Boltz/Protenix/RF3 or a mock). + latent_io + Accessor for the single representation on ``inner``'s conditioning. + injector + Transform applied to the single representation. Defaults to an identity + :class:`AffineInjector`. + training_adapter + When ``False`` (default, i.e. sampling), the injected latent is detached + so downstream guidance gradients and behaviour are unchanged. Set + ``True`` only inside a dedicated optimization pass. + """ + + def __init__( + self, + inner, + latent_io: LatentIO[C], + injector: LatentInjector | None = None, + *, + training_adapter: bool = False, + ): + self.inner = inner + self.latent_io = latent_io + self.injector = injector if injector is not None else AffineInjector() + self.training_adapter = training_adapter + + def featurize(self, structure: dict) -> GenerativeModelInput[C]: + """Featurize via ``inner``, then inject into the single representation.""" + feats = self.inner.featurize(structure) + single = self.latent_io.read_single(feats.conditioning) + if single is None: + return feats # nothing to inject into; pass through unchanged + + injected = self.injector(single) + if not self.training_adapter: + # Honour the framework's constant-conditioning contract during + # sampling: detach so guidance backprop reaches coords only. + injected = injected.detach() + + new_cond = self.latent_io.write_single(feats.conditioning, injected) + return GenerativeModelInput(x_init=feats.x_init, conditioning=new_cond) + + def step(self, *args, **kwargs): + """Delegate denoising to the wrapped model.""" + return self.inner.step(*args, **kwargs) + + def initialize_from_prior(self, *args, **kwargs): + """Delegate prior initialization to the wrapped model.""" + return self.inner.initialize_from_prior(*args, **kwargs) diff --git a/tests/models/test_latent_adapter.py b/tests/models/test_latent_adapter.py new file mode 100644 index 00000000..7cb01986 --- /dev/null +++ b/tests/models/test_latent_adapter.py @@ -0,0 +1,156 @@ +"""Tests for the Step-1 latent injection adapter (pseudo-MLP ``k*s + b``). + +These run on CPU with no model checkpoints, using a minimal mock wrapper whose +conditioning carries a single-representation tensor. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import pytest +import torch +from torch import Tensor + +from sampleworks.models.latent_adapter import ( + AffineInjector, + AttrLatentIO, + LatentAdaptedWrapper, + LatentInjector, + LatentIO, +) +from sampleworks.models.protocol import FlowModelWrapper, GenerativeModelInput + + +@dataclass(frozen=True) +class _SingleRepConditioning: + """Minimal conditioning carrying a single representation ``s`` plus sidecar state.""" + + s: Tensor + sidecar: str = "untouched" # stands in for non-tensor state (e.g. RF3 chiral arrays) + + +class _SingleRepWrapper: + """Mock FlowModelWrapper returning a deterministic single representation.""" + + def __init__(self, s: Tensor): + self._s = s + + def featurize(self, structure: dict, **kwargs: Any) -> GenerativeModelInput[_SingleRepConditioning]: + return GenerativeModelInput( + x_init=torch.zeros(1, self._s.shape[-2], 3), + conditioning=_SingleRepConditioning(s=self._s.clone()), + ) + + def step(self, x_t, t, *, features=None): + return torch.zeros_like(x_t) + + def initialize_from_prior(self, batch_size, features=None, *, shape=None): + return torch.randn(batch_size, self._s.shape[-2], 3) + + +@pytest.fixture +def s_tensor() -> Tensor: + torch.manual_seed(0) + return torch.randn(1, 8, 4) # [batch, tokens, d_s] + + +# --- AffineInjector: the pseudo-MLP ------------------------------------------ + + +def test_affine_injector_identity_at_init(s_tensor: Tensor): + """Default k=1, b=0 must be the identity.""" + out = AffineInjector()(s_tensor) + torch.testing.assert_close(out, s_tensor) + + +def test_affine_injector_applies_kx_plus_b(s_tensor: Tensor): + out = AffineInjector(k_init=2.0, b_init=3.0)(s_tensor) + torch.testing.assert_close(out, 2.0 * s_tensor + 3.0) + + +def test_affine_injector_satisfies_injector_protocol(): + assert isinstance(AffineInjector(), LatentInjector) + + +# --- AttrLatentIO: the only model-specific knowledge ------------------------- + + +def test_attr_latent_io_roundtrip(s_tensor: Tensor): + io = AttrLatentIO(single_attr="s") + cond = _SingleRepConditioning(s=s_tensor) + assert io.read_single(cond) is s_tensor + new = io.write_single(cond, s_tensor * 5) + torch.testing.assert_close(new.s, s_tensor * 5) + + +def test_attr_latent_io_preserves_sidecar_state(s_tensor: Tensor): + """replace() must leave non-tensor state (e.g. RF3 chiral arrays) intact.""" + io = AttrLatentIO(single_attr="s") + cond = _SingleRepConditioning(s=s_tensor, sidecar="keep-me") + new = io.write_single(cond, s_tensor + 1) + assert new.sidecar == "keep-me" + + +def test_attr_latent_io_satisfies_protocol(): + assert isinstance(AttrLatentIO("s"), LatentIO) + + +def test_attr_latent_io_missing_attr_returns_none(s_tensor: Tensor): + io = AttrLatentIO(single_attr="does_not_exist") + assert io.read_single(_SingleRepConditioning(s=s_tensor)) is None + + +# --- LatentAdaptedWrapper: the decorator ------------------------------------- + + +def test_adapted_wrapper_satisfies_flow_protocol(s_tensor: Tensor): + wrapped = LatentAdaptedWrapper(_SingleRepWrapper(s_tensor), AttrLatentIO("s")) + assert isinstance(wrapped, FlowModelWrapper) + + +def test_identity_injection_is_zero_impact(s_tensor: Tensor): + """k=1,b=0 => featurized single rep is unchanged: the zero-impact guarantee.""" + inner = _SingleRepWrapper(s_tensor) + wrapped = LatentAdaptedWrapper(inner, AttrLatentIO("s"), AffineInjector()) + base = inner.featurize({}) + out = wrapped.featurize({}) + torch.testing.assert_close(out.conditioning.s, base.conditioning.s) + + +def test_affine_injection_is_applied(s_tensor: Tensor): + inner = _SingleRepWrapper(s_tensor) + wrapped = LatentAdaptedWrapper(inner, AttrLatentIO("s"), AffineInjector(2.0, 3.0)) + out = wrapped.featurize({}) + torch.testing.assert_close(out.conditioning.s, 2.0 * inner.featurize({}).conditioning.s + 3.0) + + +def test_sampling_mode_detaches_latent(s_tensor: Tensor): + """In sampling mode the injected latent is detached -> guidance hits coords only.""" + wrapped = LatentAdaptedWrapper( + _SingleRepWrapper(s_tensor), AttrLatentIO("s"), AffineInjector(), training_adapter=False + ) + out = wrapped.featurize({}) + assert out.conditioning.s.requires_grad is False + + +def test_training_mode_keeps_graph(s_tensor: Tensor): + """In training mode the latent stays attached so the injector can be optimized.""" + wrapped = LatentAdaptedWrapper( + _SingleRepWrapper(s_tensor), AttrLatentIO("s"), AffineInjector(), training_adapter=True + ) + out = wrapped.featurize({}) + assert out.conditioning.s.requires_grad is True + # gradient flows back to the injector's k and b + out.conditioning.s.sum().backward() + assert wrapped.injector.k.grad is not None + assert wrapped.injector.b.grad is not None + + +def test_step_and_prior_delegate(s_tensor: Tensor): + inner = _SingleRepWrapper(s_tensor) + wrapped = LatentAdaptedWrapper(inner, AttrLatentIO("s")) + x = torch.randn(2, 8, 3) + torch.testing.assert_close(wrapped.step(x, torch.tensor(0.5), features=None), torch.zeros_like(x)) + assert wrapped.initialize_from_prior(2, shape=(8, 3)).shape == (2, 8, 3)