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/CODEBASE_GUIDE.md b/docs/CODEBASE_GUIDE.md new file mode 100644 index 00000000..4b46504e --- /dev/null +++ b/docs/CODEBASE_GUIDE.md @@ -0,0 +1,452 @@ +# Sampleworks Codebase Guide + +A function-level tour of `src/sampleworks/`, tracing how every part connects: what calls +what, and why. Read [AGENTS.md](../AGENTS.md) first for the design philosophy — this +document is the *call graph*, showing how that philosophy is wired together in practice. + +> **The one-sentence model.** Sampleworks treats a structure-prediction model (Boltz, +> Protenix, RF3) as a *prior* over realistic structures, then steers its diffusion +> sampling with the *gradient of an experimental-fit reward* (X-ray / cryo-EM density), +> drawing samples from the Bayesian posterior. Everything below is machinery for that one +> idea. + +--- + +## 1. The 30-second map + +``` +ENTRY POINTS (pyproject.toml [project.scripts]) + sampleworks-guidance → cli/guidance.py (one run, in-process) + sampleworks-runs → runs/cli.py (TOML-preset grid search; spawns one + subprocess per job, each re-entering the spine) + sampleworks-analysis → runs/analysis_cli.py (post-hoc eval; never enters the spine) + │ + │ the two guidance paths converge on ↓ + ▼ + run_guidance() in utils/guidance_script_utils.py ← THE SPINE + │ + ├─ get_reward_function_and_structure() → RealSpaceRewardFunction + atomworks structure + ├─ get_model_and_device() → Boltz/Protenix/RF3 wrapper on a device + ├─ AF3EDMSampler(EDMSamplerConfig) → the diffusion solver + ├─ {DataSpace,NoiseSpace,No}DPSScaler → per-step guidance rule + └─ {PureGuidance | FKSteering | LatentOptimization}.sample(structure, model, sampler, ...) + │ + └─ FOR each diffusion step: + sampler.step(coords, model, context, scaler) + ├─ model.step(noisy, t) → x̂₀ (denoised prediction) + ├─ reconciler.align(x̂₀, ref) → into experimental frame + └─ scaler.scale(x̂₀, context) → reward(x̂₀), ∇reward → guidance + ▼ + save_everything() → refined.cif, trajectory/*.cif, losses.txt, job_metadata.json +``` + +The five collaborators (`model`, `sampler`, `step_scaler`, `trajectory_scaler`, `reward`) +never import each other's concrete classes. They meet only through four **Protocols**, +which is what lets any model pair with any reward and any guidance strategy (the O(N+M) +promise in AGENTS.md). + +### 1.1 Directory orientation + +| Path under `src/sampleworks/` | What lives here | Guide section | +|-------------------------------|-----------------|---------------| +| `__init__.py` | Package init; the `NAN_CHECK` env toggle (`should_check_nans`) that turns NaN assertions on/off globally | §8 | +| `cli/` | `sampleworks-guidance` entry point | §10 | +| `runs/` | `sampleworks-runs` / `sampleworks-analysis` TOML-preset orchestrator | §10 | +| `core/samplers/` | Diffusion solver protocol + `AF3EDMSampler` | §4.2, §2 | +| `core/scalers/` | Step scalers (DPS) + trajectory scalers (PureGuidance/FKSteering/LatentOptimization) | §4.1, §4.3, §2 | +| `core/rewards/` | Reward protocol + `RealSpaceRewardFunction` | §5, §2 | +| `core/forward_models/` | Differentiable X-ray density calc + vendored qFit crystallography | §5 | +| `models/` | Model-wrapper protocols + Boltz/Protenix/RF3 wrappers + latent adapter | §6, §2 | +| `utils/` | The `run_guidance` spine, alignment, config, MSA, CIF/density helpers | §3, §7, §8 | +| `eval/` | Offline eval harness + synthetic-data generation + the sampling→reward bridge | §9, §4.1 | +| `metrics/` | LDDT / RMSD / sidechain quality metric framework | §9 | +| `data/` | Reference data — `protein_configs.csv` (per-protein map/selection configs) | §9 | + +Note the per-protein worker script `run_grid_search.py` lives at the **repository root** +(not under `src/`); the grid-search orchestrator spawns it as a subprocess (§10). + +--- + +## 2. The four protocols (the contracts everything speaks) + +These are the load-bearing interfaces. All are `@runtime_checkable` `typing.Protocol`s — +structural typing, no inheritance. + +| Protocol | File | Key methods | Implemented by | +|----------|------|-------------|----------------| +| **Model wrappers** | [models/protocol.py](../src/sampleworks/models/protocol.py) | `featurize()`, `step()`, `initialize_from_prior()` | Boltz1/2, Protenix, RF3 | +| **Samplers** | [core/samplers/protocol.py](../src/sampleworks/core/samplers/protocol.py) | `compute_schedule()`, `get_context_for_step()`, `step()` | `AF3EDMSampler` | +| **Scalers** | [core/scalers/protocol.py](../src/sampleworks/core/scalers/protocol.py) | `StepScalerProtocol.scale()` · `TrajectoryScalerProtocol.sample()` | DPS scalers · PureGuidance/FKSteering/LatentOptimization | +| **Rewards** | [core/rewards/protocol.py](../src/sampleworks/core/rewards/protocol.py) | `__call__()`, `precompute_unique_combinations()` | `RealSpaceRewardFunction` | + +### 2.1 The data objects that flow between them + +Understanding these five dataclasses is enough to read any part of the sampling loop. + +- **`GenerativeModelInput`** ([models/protocol.py:25](../src/sampleworks/models/protocol.py)) + — `{conditioning}`. Produced by `featurize()`. `conditioning` is a model-specific object + holding the cached trunk/pairformer output. Starting coordinates are not carried here; they + come from `initialize_from_prior(batch_size=...)` at sampling time. + +- **`StepParams`** ([core/samplers/protocol.py:44](../src/sampleworks/core/samplers/protocol.py)) + — the per-step context bundle. Carries `t, dt, noise_scale` (diffusion timing) plus + optional `reward`, `reward_inputs`, `reconciler`, `alignment_reference`, `metadata`. Built + fresh each step and *enriched* immutably via `.with_reward()`, `.with_reconciler()`, + `.with_metadata()`. This is how the reward and alignment reach the sampler without the + sampler knowing what they are. + +- **`SamplerStepOutput`** ([core/samplers/protocol.py:176](../src/sampleworks/core/samplers/protocol.py)) + — `{state, denoised, loss, log_proposal_correction}`. `log_proposal_correction` is the + log-ratio of base-to-guided proposal density, consumed only by FK-steering resampling. + +- **`RewardInputs`** ([core/rewards/protocol.py:17](../src/sampleworks/core/rewards/protocol.py)) + — `{elements, b_factors, occupancies, input_coords}`, pre-extracted once from the atom + array (`from_atom_array()`) so `scale()` doesn't re-extract every step. Validates: no NaN + coords/B-factors, occupancy in [0,1]. Tiles across ensemble/particle dims. + +- **`GuidanceOutput`** ([core/scalers/protocol.py:40](../src/sampleworks/core/scalers/protocol.py)) + — `{structure, final_state, trajectory, losses, metadata}`. What a trajectory scaler + returns; `metadata` carries `trajectory_denoised` and `model_atom_array`. + +--- + +## 3. The spine: `run_guidance()` step by step + +File: [utils/guidance_script_utils.py](../src/sampleworks/utils/guidance_script_utils.py). +Both guidance entry points reach `run_guidance()` — directly from `sampleworks-guidance`, or +once per spawned subprocess under `sampleworks-runs`. §10 walks both end to end. + +`run_guidance()` ([:427](../src/sampleworks/utils/guidance_script_utils.py)) is a thin +wrapper — it sets up per-job logging, times the run, catches exceptions into a `JobResult`, +and writes `job_metadata.json`. The real work is in `_run_guidance()` ([:481](../src/sampleworks/utils/guidance_script_utils.py)): + +1. **Load inputs & build reward** — `get_reward_function_and_structure()` ([:273](../src/sampleworks/utils/guidance_script_utils.py)): + - `resolve_mixed_hetatm_atom_altlocs()` fixes a CIF edge case (mixed ATOM/HETATM altlocs + that atomworks would misparse as an insertion), then `atomworks.parse(hydrogen_policy="remove")`. + - `XMap.fromfile(density, resolution)` loads the experimental map. + - `setup_scattering_params(em_mode)` builds the scattering-factor lookup table. + - Returns a **`RealSpaceRewardFunction`** + the structure dict. + +2. **Model-specific structure annotation** — dispatched on wrapper class name: + `annotate_structure_for_protenix` / `annotate_structure_for_rf3` / + `process_structure_for_boltz` (the last writes NPZ/manifest/MSA into the job's + `output_dir` so parallel grid jobs don't race). This stamps `ensemble_size`, + `recycling_steps`, etc. onto the structure dict for the wrapper's `featurize()` to read. + +3. **Build the sampler** — `AF3EDMSampler(EDMSamplerConfig(...))` ([:563](../src/sampleworks/utils/guidance_script_utils.py)). + Note `alignment_reverse_diffusion` defaults on only for Boltz. + +4. **Build the step scaler** — `dataspace` → `DataSpaceDPSScaler`, `noisespace` → + `NoiseSpaceDPSScaler`, `none` → `NoScalingScaler` ([:582](../src/sampleworks/utils/guidance_script_utils.py)). + +5. **Build the trajectory scaler & run** — `PureGuidance`, `FKSteering`, or `LatentOptimization`, + then call `.sample(structure, model_wrapper, sampler, step_scaler, reward_function[, num_particles])`. + `guidance_start` and `partial_diffusion_step` are converted from step counts to fractions + of `num_steps` here. `LatentOptimization` optimizes the model's cached trunk latents instead + of steering coordinates, so it ignores `step_scaler`. + +6. **Save** — `save_everything()` ([:320](../src/sampleworks/utils/guidance_script_utils.py)) + writes `refined.cif` (final ensemble coords written into the atom-array template, plus the + config injected as a `sampleworks` CIF category via `add_category_to_cif`), the denoised and + next-step trajectories (sub-sampled every 10 steps), and `losses.txt`. + +`run_guidance_job_queue()` ([:820](../src/sampleworks/utils/guidance_script_utils.py)) is +the batch entry: unpickle a list of `GuidanceConfig`, load the model **once**, loop +`run_guidance()` over jobs, emptying CUDA cache between them. + +--- + +## 4. The sampling loop in detail + +### 4.1 Trajectory scalers — the loop *around* the sampler + +All three live in `core/scalers/`. They implement `TrajectoryScalerProtocol.sample()`, whose job +is: featurize → sample prior → build reconciler → run the step loop → package +`GuidanceOutput`. + +**`PureGuidance.sample()`** ([core/scalers/pure_guidance.py:46](../src/sampleworks/core/scalers/pure_guidance.py)) — standard guided diffusion, no resampling: + +``` +features = model.featurize(structure) +coords = model.initialize_from_prior(ensemble_size, features) # Gaussian noise +processed = process_structure_to_trajectory_input(...) # → reconciler + reward_inputs +schedule = sampler.compute_schedule(num_steps) +for i in range(starting_step, num_steps): + context = sampler.get_context_for_step(i, schedule) # StepParams(t, dt, noise_scale) + if i >= guidance_start: + context = context.with_reward(reward, reward_inputs) # attach reward + context = context.with_reconciler(reconciler, alignment_reference) # attach alignment + out = sampler.step(coords, model, context, + scaler=step_scaler if guiding else None, + features=features) + coords = out.state # advance +``` + +**`FKSteering.sample()`** ([core/scalers/fk_steering.py:67](../src/sampleworks/core/scalers/fk_steering.py)) +— Feynman-Kač steering. Same skeleton but with **`num_particles`** whole ensembles evolving +in parallel and periodic **resampling** toward low loss: + +- `_run_step()` ([:225](../src/sampleworks/core/scalers/fk_steering.py)) — with no guidance, + runs all particles in one batched `sampler.step`; **with** guidance, loops per particle so + each gets its own gradient (required for correct FK weights). +- `_should_resample()` ([:352](../src/sampleworks/core/scalers/fk_steering.py)) — fires every + `resampling_interval` steps while noise > 0. +- `_resample_particles()` ([:367](../src/sampleworks/core/scalers/fk_steering.py)) — weights + `log_G = fk_lambda·(loss_prev − loss_curr) + log_proposal_correction`, then + `softmax → multinomial` to duplicate/drop whole ensembles. "Particles" are ensembles, not + single structures. +- Returns the single lowest-loss particle ensemble. + +**`LatentOptimization.sample()`** ([core/scalers/latent_optimization.py:123](../src/sampleworks/core/scalers/latent_optimization.py)) +— inference-time latent optimization (IT-opt). Instead of steering coordinates, it makes the +model's cached trunk latents (`s`, `z`) optimizable leaves, runs one or more full diffusion +rounds updating them against the reward, then samples the final ensemble with the latents +frozen. Coordinate guidance is not applied, so `step_scaler` is ignored. See +[IT_OPT_DESIGN.md](IT_OPT_DESIGN.md) for the algorithm and its mapping to the reference. + +All three call `process_structure_to_trajectory_input()` +([eval/structure_utils.py:106](../src/sampleworks/eval/structure_utils.py)), which cleans the +atom array, builds the **`AtomReconciler`**, tiles coordinates across the batch, and returns a +frozen `SampleworksProcessedStructure`. Its `.to_reward_inputs()` produces the `RewardInputs`. + +### 4.2 The sampler — `AF3EDMSampler.step()` + +File: [core/samplers/edm.py](../src/sampleworks/core/samplers/edm.py). This is the +Karras-EDM sampler (Karras et al. 2022, the Euler variant) as used in AlphaFold3 — +`step()` follows AF3 Supplementary Algorithm 18. `EDMSamplerConfig` defaults match the AF3 +parameterization *except* `gamma_min = 0.2` (AF3 uses 1.0). One `step()` +([:363](../src/sampleworks/core/samplers/edm.py)): + +1. `check_context()` — validate `t, dt, noise_scale` present. +2. Center coords; if `augmentation`, apply a random SO(3) rotation + translation + (`create_random_transform`). +3. Add stochastic noise: `noisy = augmented + eps·noise_scale`; set + `noisy.requires_grad_(scaler.requires_gradients)`. +4. **`x̂₀ = model_wrapper.step(noisy, t_hat, features)`** — the model denoises. *This is the + only call into the neural network per step.* +5. **Align** `x̂₀` into the experimental frame: if a reconciler is present, + `reconciler.align()`; else `align_to_reference_frame()`. Noise is carried into the aligned + frame via `transform_coords_and_noise_to_frame()` (rotation-only for the noise part, since + noise is translation-invariant). +6. Compute drift `delta = (noisy − x̂₀) / t_hat`. +7. If a scaler is passed, `_apply_scaler_guidance()` ([:301](../src/sampleworks/core/samplers/edm.py)): + calls `scaler.scale(x̂₀, context.with_metadata({"x_t": noisy}), model)` → guidance + direction + loss, scales by `guidance_strength()`, rotates the direction into the aligned + frame, optionally rescales to the diffusion magnitude, folds into `delta`, and computes the + `log_proposal_correction`. +8. Euler step: `next = noisy + step_scale·dt·delta`. +9. Return `SamplerStepOutput(state, denoised=x̂₀, loss, log_proposal_correction)`. + +The schedule (`EDMSchedule`, [:31](../src/sampleworks/core/samplers/edm.py)) is precomputed +once by `compute_schedule()` ([:210](../src/sampleworks/core/samplers/edm.py)) using the EDM +sigma schedule `σ = σ_data·(s_max^{1/p} + t·(s_min^{1/p} − s_max^{1/p}))^p` over +`num_steps+1` points; `gamma` (stochastic churn) is applied only where `σ > gamma_min`. +`get_context_for_step()` is an O(1) lookup that packages `t_hat, dt, eps_scale` into a +`StepParams`. + +### 4.3 Step scalers — the per-step guidance rule + +File: [core/scalers/step_scalers.py](../src/sampleworks/core/scalers/step_scalers.py). All +implement `StepScalerProtocol.scale(state, context, model) → (guidance_direction, loss)`. + +- **`DataSpaceDPSScaler`** ([:51](../src/sampleworks/core/scalers/step_scalers.py)) — enable + grad on `x̂₀`, compute `loss = reward(x̂₀, …)`, backprop → `∂loss/∂x̂₀`. Cheap; no backprop + through the model. `requires_gradients = False`. +- **`NoiseSpaceDPSScaler`** ([:100](../src/sampleworks/core/scalers/step_scalers.py)) — reads + `context.metadata["x_t"]` (the noisy state), computes `loss = reward(x̂₀)`, backprops + **through the model** to `∂loss/∂x_t`. More faithful DPS; `requires_gradients = True` (which + is what makes the sampler set `requires_grad` on the noisy state in step 3 above). +- **`NoScalingScaler`** ([:33](../src/sampleworks/core/scalers/step_scalers.py)) — returns + zeros (baseline/unguided). + +Both DPS variants optionally normalize the gradient. `guidance_strength()` returns the +per-step weight (`step_size`). + +--- + +## 5. The reward and forward model (coords → scalar → gradient) + +This is where "fit to experiment" becomes a differentiable number. +Files: [core/rewards/real_space_density.py](../src/sampleworks/core/rewards/real_space_density.py) +and [core/forward_models/xray/real_space_density.py](../src/sampleworks/core/forward_models/xray/real_space_density.py). + +``` +RealSpaceRewardFunction.__call__(coords, elements, b_factors, occupancies) [rewards/…:291] + │ + ├─ DifferentiableTransformer.forward(coords, …) [forward_models/…:501] + │ ├─ _compute_radial_densities() → per-atom radial density profiles + │ │ └─ torch.vmap over unique (element, b_factor) pairs + │ │ └─ GaussLegendreQuadrature ∫ scattering_integrand(s) ds + │ ├─ _compute_grid_coordinates() → Cartesian → fractional → grid + │ ├─ dilate points onto the grid → CUDA kernel (dilate_atom_centric) + │ │ or pure-torch (dilate_points_torch + scatter_add_) + │ └─ apply crystallographic symmetry (R,t per space-group op; F.grid_sample) + │ → density grid [batch, Dz, Dy, Dx] + ├─ .sum(0) → collapse batch → single grid + └─ self.loss(density, xmap.array) → L1/L2 vs. experimental map → SCALAR +``` + +Every operation is autograd-tracked, so `.backward()` yields `∂reward/∂coords` — the signal +the DPS scalers turn into guidance. Key details: + +- **`precompute_unique_combinations()`** ([rewards/…:228](../src/sampleworks/core/rewards/real_space_density.py)) + runs `torch.unique` *outside* vmap to avoid dynamic shapes; the results are passed in so the + vmapped radial integration sees only static shapes. +- **Crystallographic symmetry** is applied by the forward model, not the reward — the model + predicts a P1 asymmetric unit but the map is in the full crystal frame. Two equivalent + paths: CUDA (symmetry applied to atoms before dilation) or CPU (`XMap_torch.apply_symmetry` + after dilation via `grid_sample`). +- `setup_scattering_params(em_mode)` ([rewards/…:25](../src/sampleworks/core/rewards/real_space_density.py)) + picks X-ray Cromer-Mann (`ATOM_STRUCTURE_FACTORS`) vs. electron + (`ELECTRON_SCATTERING_FACTORS`) coefficients from the qFit dependency. +- Two extraction helpers sit alongside the reward: the module-level + `extract_density_inputs_from_atomarray` ([rewards/…:69](../src/sampleworks/core/rewards/real_space_density.py)) + and the method `RealSpaceRewardFunction.structure_to_reward_input` + ([rewards/…:258](../src/sampleworks/core/rewards/real_space_density.py)) — both turn an atom + array / structure dict into the element/B-factor/occupancy/coord tensors the reward consumes. + +**`real_space_density_deps/`** is vendored crystallography (from qFit): `spacegroups.py` +(space-group operators, 8k lines), `sf.py` (scattering factors + structure-factor calc), +`unitcell.py`, `volume.py` (`XMap`, CCP4 I/O), `transformer.py`, plus `utils/quadrature.py` +(Gauss-Legendre) and `ops/dilate_points_cuda.py` (the custom autograd + vmap CUDA kernel). + +--- + +## 6. Model wrappers (the pluggable priors) + +Files: [models/](../src/sampleworks/models/). Each wraps an external model behind +`FlowModelWrapper`. Common shape: + +- **`featurize(structure)`** — convert the atomworks dict into the model's native input, + run the expensive trunk/pairformer **once**, and cache it in a model-specific `Conditioning` + dataclass. Also loads the **model-space atom array** for reconciliation. +- **`step(x_t, t, features)`** — run only the diffusion/structure module on the cached + conditioning; return predicted clean coords. Called every diffusion step. +- **`initialize_from_prior(batch_size, features)`** — Gaussian noise at the model's atom count. + +| Wrapper | File | Notes | +|---------|------|-------| +| `Boltz1Wrapper` / `Boltz2Wrapper` | [boltz/wrapper.py](../src/sampleworks/models/boltz/wrapper.py) | Caches pairformer (`s_trunk, z_trunk`); Boltz2 also caches diffusion conditioning. Preprocessing writes NPZ/manifest/MSA. `process_structure_for_boltz` reconstructs the atom array from the processed NPZ. | +| `ProtenixWrapper` | [protenix/wrapper.py](../src/sampleworks/models/protenix/wrapper.py) | AF3 reimpl. `structure_processing.py` builds the Protenix JSON (entities, modifications, covalent bonds). Optional diffusion-shared-vars cache (`pair_z, p_lm, c_l`). | +| `RF3Wrapper` | [rf3/wrapper.py](../src/sampleworks/models/rf3/wrapper.py) | Baker AF3 replica. Uses an inference engine + trunk-with-recycling generator. Optional chiral-feature tracking/disabling (writes `chiral_grad_stats.json`). | + +The wrappers guarantee the atom array they hand back has valid coords/occupancy/B-factors, so +downstream `RewardInputs.from_atom_array()` never sees NaNs. Import failures are tolerated +([guidance_script_utils.py:51-75](../src/sampleworks/utils/guidance_script_utils.py)) because +Boltz/Protenix/RF3/Protpardelle have mutually incompatible dependencies and live in separate +pixi envs. + +--- + +## 7. Alignment & atom reconciliation (the SE(3) glue) + +Reward functions compare coordinates in the *fixed experimental frame*; models emit +arbitrary frames with possibly different atom sets. Two utilities bridge this — and the +**sampler**, not the reward, owns the timing (AGENTS.md "Alignment" pitfall). + +- **`AtomReconciler`** ([utils/atom_reconciler.py](../src/sampleworks/utils/atom_reconciler.py)) + — `from_arrays(model_array, struct_array)` normalizes atom IDs + (`chainidx_seqpos_atomname`, handling 0-based vs. author numbering), finds the common subset, + and returns index maps (or `.identity()` when they match). `align()` computes a rigid + transform on the common atoms and applies it to *all* model atoms. `struct_to_model()` maps + structure coords into model space differentiably. +- **`weighted_rigid_align_differentiable()`** ([utils/frame_transforms.py:332](../src/sampleworks/utils/frame_transforms.py)) + — weighted Kabsch/Procrustes via SVD (float32 for stability, reflection-corrected), + **preserving gradients** (unlike Boltz's detached version). `transform_coords_and_noise_to_frame()` + ([:560](../src/sampleworks/utils/frame_transforms.py)) applies full transform to coords but + rotation-only to noise. + +--- + +## 8. Utils (the support layer) + +Directory: [utils/](../src/sampleworks/utils/). Highlights beyond §3/§7: + +| File | Role | +|------|------| +| `guidance_script_arguments.py` | `GuidanceConfig` (all run params, `from_cli()` two-pass parse), `JobResult`, `_resolve_checkpoint()` (env → baked → ACTL → legacy), `validate_model_checkpoint()`. `as_dict()` remaps container↔host paths. | +| `guidance_constants.py` | Enums: `GuidanceType`, `StructurePredictor`, `StepScalers`, `TrajectoryScalers`, `Rewards`. | +| `cif_utils.py` | `resolve_mixed_hetatm_atom_altlocs()`, `add_category_to_cif()` (writes the `sampleworks` metadata block). | +| `atom_array_utils.py` | `make_normalized_atom_id()`, `filter_to_common_atoms()` (used by the reconciler). | +| `density_utils.py` | `compute_density_from_atomarray()`, `build_density_transformer()` — reusable forward-model wrappers used by synthetic-data generation and RSCC eval. | +| `msa.py` / `mmseqs2.py` | `MSAManager` — SHA3-keyed MSA cache, ColabFold/Protenix-server/mmseqs2 fetch, per-model formatting (CSV for Boltz, A3M for RF3). | +| `elements.py` | `element_to_scattering_idx()` — element symbol → scattering-table index. | +| `frame_transforms.py` | Rigid-transform algebra (forward/inverse/apply, random augmentation). | +| `torch_utils.py` | `try_gpu()` — pick the least-loaded GPU via `nvidia-smi`. | +| `imports.py` | `BOLTZ_/PROTENIX_/RF3_AVAILABLE` flags + `@require_*` test decorators. | +| `protein_input.py` | CSV parser for batch protein specs. | + +One package-level knob worth knowing: [`sampleworks/__init__.py`](../src/sampleworks/__init__.py) +reads the `NAN_CHECK` env var into `should_check_nans` (default on; set `NAN_CHECK=false`/`0` +to disable). [`torch_utils.py`](../src/sampleworks/utils/torch_utils.py) uses it to make +`assert_no_nans` either a real check or a no-op, gating expensive NaN assertions in hot paths. + +--- + +## 9. Metrics & eval + +**`metrics/`** — a pluggable metric framework used for validation/scoring: + +- `metric.py` — `Metric` ABC (`compute()`, `kwargs_to_compute_args`) + `MetricManager` + (tag-filtered batch computation). +- Concrete: `AllAtomLDDT` / `SelectedLDDT` ([lddt.py](../src/sampleworks/metrics/lddt.py)), + `AllAtomRMSD` ([rmsd.py](../src/sampleworks/metrics/rmsd.py), optional Kabsch), + `SidechainMetrics` ([sidechain_metrics.py](../src/sampleworks/metrics/sidechain_metrics.py), + topology/bond/clash checks), `ExtraInfo` (metadata pass-through). +- `metric_utils.py` — binning/masking helpers for LDDT/PAE-style scores. + +**`eval/`** — the offline evaluation and synthetic-data machinery: + +- `structure_utils.py` — **`SampleworksProcessedStructure` + `process_structure_to_trajectory_input()`** + (the bridge from §4.1 into `RewardInputs`), plus selection-string parsing and reference-structure loading. +- `eval_dataclasses.py` — `Trial`, `TrialList`, `ProteinConfig` (per-protein map/selection config, `from_csv()`). +- `grid_search_eval_utils.py` — `scan_grid_search_results()` walks a results tree of + `refined.cif` files into `Trial`s; `setup_evaluation_parameters()` + `parse_eval_args()` + standardize eval-script setup. +- `generate_synthetic_density.py` / `generate_synthetic_sf.py` — build synthetic maps/MTZs + from a structure using the forward models (via `compute_density_from_atomarray` / + `SFcalculator`); each has a CLI `main()` with single/batch modes. +- `metrics.py` — `rscc()` (real-space correlation coefficient). +- `occupancy_utils.py`, `synthetic_utils.py`, `constants.py` — altloc/occupancy handling for + synthetic ensembles. + +--- + +## 10. The two run paths, end to end + +**Direct (`sampleworks-guidance`)** — [cli/guidance.py](../src/sampleworks/cli/guidance.py): +`GuidanceConfig.from_cli(argv)` → `get_model_and_device()` → `run_guidance()` → exit code. One +process, one run. + +**Preset grid search (`sampleworks-runs`)** — [runs/](../src/sampleworks/runs/): +- `cli.py::main` → `run_cli()` with `EXPERIMENT_CLI_CONFIG` (presets in `experiments/`). +- `loader.py` reads a TOML preset (`_read_toml` → `_apply_overrides` (`--set a.b=c`) → + `_resolve_variables` (`${VAR}`) → `_build_preset`), producing `Preset`/`Job` dataclasses + ([schema.py](../src/sampleworks/runs/schema.py)). +- `runner.py::run` ([:619](../src/sampleworks/runs/runner.py)) resolves GPU assignments + (`_resolve_gpu_assignments` via `nvidia-smi`), builds one `JobInvocation` per job + (`_build_argv` picks the pixi-env Python and assembles the command), runs `pre_jobs` + sequentially (`_run_sequential`), then **spawns main jobs in parallel** (`_spawn` → + `subprocess.Popen` + `_tee` threads) and waits (`_wait_all`). +- The worker script defaults to `run_grid_search.py`, resolved by `_resolve_script_path` + from the repository root (`./run_grid_search.py`), with container/workspace fallbacks + (`/app/run_grid_search.py`, `/home/dev/workspace/run_grid_search.py`). A job may override it + via `Job.script`. +- Each spawned `run_grid_search.py` process ultimately calls the same `run_guidance()`. +- `analysis_cli.py` reuses `run_cli()` with `ANALYSIS_CLI_CONFIG` (presets in `analyses/`) for + post-hoc evaluation jobs. + +So the grid-search path fans out into many subprocesses that each re-enter the spine of §3. + +--- + +## 11. Where to start reading, by task + +- **"How does one guided step work?"** → `AF3EDMSampler.step()` (§4.2), then a DPS scaler (§4.3). +- **"How is the loop orchestrated?"** → `PureGuidance.sample()` / `FKSteering.sample()` (§4.1). +- **"How does experiment become a gradient?"** → `RealSpaceRewardFunction.__call__` + + `DifferentiableTransformer.forward` (§5). +- **"How do I add a model / reward / scaler?"** → satisfy the matching protocol in §2; see the + AGENTS.md "Adding New Components" section. +- **"How is a whole run wired?"** → `_run_guidance()` (§3) is the single best file to read. +- **"Why is alignment everywhere?"** → §7 (rewards live in the crystal frame; models don't). diff --git a/docs/IT_OPT_DESIGN.md b/docs/IT_OPT_DESIGN.md new file mode 100644 index 00000000..7e9a5a08 --- /dev/null +++ b/docs/IT_OPT_DESIGN.md @@ -0,0 +1,176 @@ +# IT-Opt — Design (read this first) + +Inference-time latent optimization (IT-opt) for sampleworks. This is the one document to read to +understand the feature: **what it is, the as-built algorithm, the components and where they live in +the code, per-model gradient readiness, and the design choices that make it correct.** + +Companions (read only if you need them): +- [IT_OPT_TESTING.md](IT_OPT_TESTING.md) — how to run, debug, and verify it, plus the open problems. +- [IT_OPT_REFERENCE_NOTES.md](developer_notes/IT_OPT_REFERENCE_NOTES.md) — deep dive on the external `it_opt/` + reference tree and its bug catalog (only relevant if you are re-porting from the reference). + +--- + +## 1. What it is + +A frozen structure predictor (Protenix / RF3 / Boltz) turns a sequence into cached post-trunk +latents — the **single** representation `s` and the **pair** representation `z` — and a diffusion +module decodes those latents into coordinates. IT-opt treats the model as a differentiable sampler +and **optimizes `s` and `z` themselves** (not the coordinates, not the weights) against an +experimental reward evaluated on the *denoised* structure. No model weights are trained. + +It is a `TrajectoryScalerProtocol`, a peer of `PureGuidance` and `FKSteering`, living in +[core/scalers/latent_optimization.py](../src/sampleworks/core/scalers/latent_optimization.py). It is +a faithful port of the reference `run_it_optimization`, with the reference's bugs fixed (§6). + +**v1 is latent-only.** Coordinate-space guidance is *not* applied — the attached step-scaler +returns a zero coordinate direction (§4, `_GradEnablingScaler`), so the only steering comes from the +evolving latents. + +## 2. The as-built loop + +``` +extract (s, z) from the trunk once -> clone into optimizable leaves (requires_grad=True) +for each outer round (fresh prior noise each round): + optimizer = Adam([s, z]) # ONE persistent Adam per round + for each diffusion step: + x̂₀ = differentiable denoise(x_t, s, z) # ONE forward, run under autograd + loss = reward(x̂₀) + anchor(s, z) + bond_geometry(x̂₀) + loss.backward(); clip s and z INDEPENDENTLY; optimizer.step() + x_t = step_output.state.detach() # advance; coordinate graph cut here +final clean sampling pass with the optimized latents -> saved ensemble +``` + +The gradient that matters is `∂reward(x̂₀)/∂latent` through **one** denoiser forward. Because the +coordinate state is detached every step (`coords = step_output.state.detach()`), the gradient never +flows through the coordinate recursion — it is a greedy, per-step latent gradient, **not** +backprop-through-sampling. (Unrolling the full sampler is intractable: the EDM sampler re-noises and +detaches each step, and a true unroll is O(steps·N²·d) memory.) + +## 3. Components and where they live + +| Component | File | What it is | +|---|---|---| +| `LatentOptimization` | [core/scalers/latent_optimization.py](../src/sampleworks/core/scalers/latent_optimization.py) | The scaler / entry point. Owns the loop, the ensemble, and the final sampling pass. | +| `LatentAnchor` | same file | `Σ wᵢ·mean((latentᵢ − baselineᵢ)²)` — mean-squared drift from the baseline latents. Regularizer, added to the loss directly (not a reward). | +| `_GradEnablingScaler` | same file | Minimal `StepScalerProtocol` that only turns autograd on (§4). | +| `BondGeometryReward` | [core/rewards/geometry.py](../src/sampleworks/core/rewards/geometry.py) | Bond-length + steric-clash hinges on `x̂₀`. Built only when `bond_length_weight > 0`. Bounded hinges (not the reference's exploding `exp(relu(...))`). | +| `AttrLatentIO` / `LatentIO` | [models/latent_adapter.py](../src/sampleworks/models/latent_adapter.py) | Reads/writes `s`/`z` by attribute name; writes via `dataclasses.replace` (honors `frozen=True, slots=True`). The only model-specific knowledge is a pair of attribute names. | +| CLI wiring | [utils/guidance_script_utils.py](../src/sampleworks/utils/guidance_script_utils.py) (`_run_guidance`), [utils/guidance_script_arguments.py](../src/sampleworks/utils/guidance_script_arguments.py) (`add_latent_opt_args`), [utils/guidance_constants.py](../src/sampleworks/utils/guidance_constants.py) (`GuidanceType.LATENT_OPT`) | Builds the scaler from CLI flags. | + +Private helpers, for orientation: `_leaf_latents` (promote s/z to leaves), `_optimize_one_round` +(one outer round), `_latent_adam_step` (score → backward → clip → step), `_sample_with_frozen_latents` +(final clean pass). + +### Constructor contract + +```python +LatentOptimization( + ensemble_size=1, num_steps=200, guidance_t_start=0.0, *, + outer_steps=1, learning_rate=0.05, max_grad_norm=1.0, + optimize_single=True, optimize_pair=True, + anchor_weight_single=0.0, anchor_weight_pair=0.0, + bond_length_weight=0.0, single_attr="s", pair_attr="z", +) +``` + +- `ensemble_size` structures are sampled in parallel and **share** the latents (per-member latents + are a documented follow-up). +- `guidance_t_start` is a fraction in `[0,1]`; stored as `guidance_start = int(guidance_t_start*num_steps)`. + Steps before it are plain frozen-latent diffusion. +- `outer_steps` resample rounds, fresh prior noise each. **Constructor default `1`; CLI default `2`.** +- `learning_rate` Adam LR; one persistent Adam is built **once per round**. +- `max_grad_norm` clips `s` and `z` **independently** (see §6 — a joint clip starves `s`). +- `single_attr`/`pair_attr` default to Boltz names; the wiring always passes model-resolved names. +- The `metadata` on the returned `GuidanceOutput` carries `"optimization_losses"` (per-round, + per-step data losses) and `"latent_drift"` (per-round relative L2 drift of each latent). It does + **not** emit the siblings' `"trajectory_denoised"` — treat scaler-specific keys as optional. + +## 4. The gradient gate + +The load-bearing mechanism. `AF3EDMSampler.step` reads `getattr(scaler, "requires_gradients", False)`; +if true it runs the denoiser under `torch.set_grad_enabled(True)` and returns a `denoised` (`x̂₀`) +that still carries a graph back to the latent leaves: + +```python +class _GradEnablingScaler: + requires_gradients = True + def scale(self, state, context, *, model=None): + return torch.zeros_like(state), torch.zeros(state.shape[0], device=state.device) +``` + +The zero direction means the trajectory advance is unguided; the flag is the *only* thing that turns +autograd on. `requires_gradients` is duck-typed (read via `getattr`, not declared on the protocol). +The sampler does **not** re-detach `state`/`denoised` — the scaler loop detaches between iterations. + +The denoised `x̂₀` is frame-aligned to the experimental reference by the sampler (via the +`AtomReconciler` + `alignment_reference`) *before* the reward sees it, so reward functions can assume +pre-aligned coordinates. + +## 5. Per-model gradient readiness + +The only model-specific concern is whether the cached latents can receive a gradient. The optimized +attributes are `s`/`z` for Boltz and `s_trunk`/`z_trunk` for Protenix/RF3 +(`DEFAULT_SINGLE_REP_ATTR` / `DEFAULT_PAIR_REP_ATTR` in `latent_adapter.py`). + +| Model | Status | Note | +|---|---|---| +| **Boltz1** | Clean | Trunk latents reach the diffusion module directly. | +| **RF3** | Clean | Same. | +| **Protenix** | Works, with two conditions | (a) `step()` keeps an injected leaf attached automatically via the `detach_unless_leaf` helper — a cached latent detaches under grad (avoids double-backward), but a leaf with `requires_grad=True` is kept attached. No manual edit needed. (b) For `z`, the diffusion shared-vars cache must be **off** (`pair_z`/`p_lm`/`c_l` are z-derived; a stale cache silently zeroes / corrupts the `z` gradient). `_run_guidance` disables it for `LATENT_OPT`. | +| **Boltz2** | Not yet | `z` (and part of `s`) are baked into a cached `diffusion_conditioning` at featurize time; needs the conditioning recomputed from the live latents. | + +All four models' conditioning is a `@dataclass(frozen=True, slots=True)`, swapped via +`dataclasses.replace` — hence `AttrLatentIO`. + +## 6. Design choices and invariants + +**Kept from the reference (the algorithm):** `s`/`z` as the optimized variables with `s_inputs` +frozen; one Adam step per diffusion step on the denoised `x̂₀`; latents persist across steps and +rounds; an outer resample loop; an anchor to the trunk baseline; a final clean sampling pass. + +**Fixed from the reference:** +- **Persistent Adam.** The reference rebuilt the optimizer *inside* the step loop, degenerating it to + `lr·sign(grad)` (signSGD) — "not actually running Adam," which also made its LR and grad-clip + inert. The port builds one Adam per round. +- **Independent per-latent clip.** `z`'s gradient is ~1e4× `s`'s, so a single joint `[s, z]` clip is + set by `z` and starves `s`. Clipping them separately decouples `s`'s step from `z`'s scale (Adam's + per-parameter normalization is what actually commensurates them). A joint clip was tried and + reverted 2026-07-20. +- **One forward per step.** The reference did a second `no_grad` forward to advance; the port advances + with `step_output.state.detach()` from the same differentiable step. +- Bounded geometry hinges instead of the reference's exploding `exp(relu(...))` clash term; no + hardcoded save-step. + +**Deliberate deviations:** the objective is a pluggable `RewardFunctionProtocol` (v1: +`RealSpaceRewardFunction`, density fit) rather than backbone-RMSD; the anchor is mean-squared rather +than the reference's Frobenius norm (shape-agnostic, so `w_s`/`w_z` are comparable); latents are +shared across the ensemble in v1. + +**Invariants worth stating:** +- **Optimization is not sampling.** An optimized `(s, z)` is a point estimate; reporting it as a + single state collapses the Boltzmann/population weighting. If populations matter, sample the latent + rather than optimizing it. Ensembles here come from fresh prior noise per member, not from a + posterior over latents. +- **On-manifold discipline.** Pushing `z` hard buys density fit with broken geometry; the anchor and + `BondGeometryReward` are the counter-pressure. Efficacy must be judged on held-out fit against + matched-compute baselines, not train-set loss (see [IT_OPT_TESTING.md](IT_OPT_TESTING.md)). + +## 7. Why `s`/`z` and not the MSA + +In the AF2/OpenFold tradition you could optimize the MSA representation `m`. The AF3-family models +(Protenix, Boltz) have **no persistent MSA latent** at the featurize→step boundary: the MSA module +writes only into `z`, and `s` is updated inside the Pairformer via pair-biased attention, seeded from +`s_inputs` rather than from an MSA row. So the post-trunk optimization levers are exactly `s_trunk` +(≈ the single rep) and `z_trunk` (≈ where the MSA information now lives). `m` is upstream of +featurize, rebuilt each recycle, and the trunk runs under `no_grad` — optimizing it would require +backprop through the trunk (ColabDesign/AfDesign territory) and is out of scope. + +## 8. Running it + +CLI: `sampleworks-guidance --model --guidance-type latent_opt …` with +`--which-latent {single,pair,both}` (default `pair`), `--learning-rate` (0.05), `--outer-steps` (2), +`--anchor-weight` (0.0), `--max-grad-norm` (1.0), `--bond-length-weight` (5e-5). Also reachable +programmatically through `run_guidance`. `_run_guidance` resolves the model name to the latent +attribute names and raises a clear `ValueError` for an unsupported model. See +[IT_OPT_TESTING.md](IT_OPT_TESTING.md) for the debug ladder and gradient check. diff --git a/docs/IT_OPT_TESTING.md b/docs/IT_OPT_TESTING.md new file mode 100644 index 00000000..f07f84fe --- /dev/null +++ b/docs/IT_OPT_TESTING.md @@ -0,0 +1,162 @@ +# IT-Opt — Testing, Verification, and Wiring + +How to run [`LatentOptimization`](../src/sampleworks/core/scalers/latent_optimization.py), debug it +when it misbehaves, and where it is wired into the pipeline. For the architecture, read +[IT_OPT_DESIGN.md](IT_OPT_DESIGN.md) first. + +Protenix is the primary test target (it is the model the reference algorithm was written for), so the +recipes below use it; Boltz1/RF3 need neither precondition in §1. + +--- + +## 1. Protenix preconditions (read first) + +**(a) The optimizable latent must reach the trunk.** Protenix's `step()` detaches the cached trunk +outputs when gradients are on (to avoid double-backward across the denoising steps) — correct for +coordinate guidance, but it would make a *latent* gradient **exactly zero**, silently. The +`detach_unless_leaf` helper in [models/protenix/wrapper.py](../src/sampleworks/models/protenix/wrapper.py) +handles both paths automatically: a cached latent detaches, but a latent IT-opt injected as an +optimizable leaf (`requires_grad=True`) is kept attached. No manual edit is needed — if the §2 Level-0 +check fails, confirm the leaf actually has `requires_grad=True` and that `single_attr`/`pair_attr` are +`s_trunk`/`z_trunk`. + +**(b) For `z` optimization, disable the diffusion shared-vars cache.** `pair_z`, `p_lm`, `c_l` are +cached from `z_trunk` at featurize time. Optimizing `z_trunk` with the cache on makes the diffusion +module read the **stale** `pair_z` — the `z` gradient is only partial and the forward is wrong. Build +the Protenix wrapper with `enable_diffusion_shared_vars_cache=False` (which `_run_guidance` does for +`LATENT_OPT`), or start single-only (`optimize_pair=False`), which the cache doesn't affect. + +> Recommended first Protenix run: `optimize_single=True, optimize_pair=False` with the cache at its +> default. It exercises the whole loop without the cache subtlety; then turn on `z` with the cache off. + +## 2. Debug ladder + +The runtime lives on the Astera pod; edit locally, sync, run in the Protenix env. Work up the levels +and stop at the first that fails. + +```bash +pixi run -e protenix-dev python your_it_opt_test.py # needs the checkpoint + a map/structure +pixi run -e protenix-dev python -m pytest tests/models/test_latent_adapter.py -q # CPU, no checkpoint +``` + +### Level 0 — Does the gradient reach the latent? (the #1 failure mode) + +Needs no reward; confirms precondition (a). Inject a `requires_grad` leaf as `z_trunk` (or +`s_trunk`), run one differentiable denoiser forward, and check the leaf received a gradient: + +```python +import torch +from sampleworks.utils.guidance_script_utils import get_model_and_device, get_reward_function_and_structure +from sampleworks.core.samplers.edm import AF3EDMSampler, EDMSamplerConfig +from sampleworks.models.latent_adapter import AttrLatentIO +from sampleworks.models.protocol import GenerativeModelInput + +device, model = get_model_and_device("cuda:0", "", "protenix") +reward, structure = get_reward_function_and_structure( + density="", device=device, em=False, loss_order=2, + resolution=1.8, structure_path="", +) + +feats = model.featurize(structure) # trunk pass +io = AttrLatentIO(single_attr="s_trunk", pair_attr="z_trunk") + +z0 = io.read_pair(feats.conditioning).detach() +z_leaf = z0.clone().requires_grad_(True) +cond = io.write_pair(feats.conditioning, z_leaf) +feats2 = GenerativeModelInput(conditioning=cond) # conditioning-only (x_init removed in #330) + +sampler = AF3EDMSampler(EDMSamplerConfig(device=str(device), augmentation=False)) +schedule = sampler.compute_schedule(num_steps=200) +t_hat = schedule.t_hat[100] +x = torch.as_tensor(model.initialize_from_prior(batch_size=1, features=feats2)) + +with torch.enable_grad(): + x0 = model.step(x, t_hat, features=feats2) + x0.sum().backward() + +print("z_leaf.grad is None:", z_leaf.grad is None) +print("z_leaf.grad abs-sum:", None if z_leaf.grad is None else z_leaf.grad.abs().sum().item()) +``` + +**Expect** `grad is None → False`, `abs-sum → > 0`. Repeat with `read_single`/`write_single` and +`s_trunk` to check `s`. + +> With the `z` cache **on**, this still shows a non-zero grad (through the direct `z_trunk` arg) but +> the forward used the stale `pair_z`. Level 0 confirms the leaf reaches the trunk, not cache +> correctness — that is why §1(b) matters for real `z` runs. + +### Level 1 — A tiny end-to-end run + +```python +from sampleworks.core.scalers.latent_optimization import LatentOptimization + +itopt = LatentOptimization( + ensemble_size=1, num_steps=20, outer_steps=1, # tiny + learning_rate=0.05, max_grad_norm=1.0, + optimize_single=True, optimize_pair=False, # single-only is cache-safe + anchor_weight_single=1.0, single_attr="s_trunk", pair_attr="z_trunk", +) +out = itopt.sample(structure, model, sampler, step_scaler=None, reward=reward) +opt = out.metadata["optimization_losses"][0] # per-step data losses, round 0 +print("first→last opt loss:", opt[0], "→", opt[-1]) +``` + +**Expect** the optimization loss trends **down** across the round and `out.final_state` has shape +`[ensemble_size, n_atoms, 3]`. + +### Level 2 — Scale up + +Raise `num_steps` (→200), `outer_steps` (→2–4), `ensemble_size`; turn on `z` (`optimize_pair=True` +**with the cache disabled**). Compare the final ensemble's density fit (RSCC) to an unguided baseline. + +## 3. Failure → diagnosis + +| Symptom | Likely cause | Fix | +|---|---|---| +| `latent.grad is None` / zero (Level 0) | Injected latent isn't a leaf reaching the trunk (§1a); or you optimized a latent the model doesn't route to. | Confirm `requires_grad=True`; use `single_attr="s_trunk"`, `pair_attr="z_trunk"` for Protenix. | +| `z` grad non-zero but structures don't respond | Stale `pair_z` cache (§1b). | `enable_diffusion_shared_vars_cache=False`, or optimize single-only first. | +| Optimization loss is flat | LR too small; anchor too strong (latents pinned); or gradient fully clipped. | Raise `learning_rate`; lower `anchor_weight_*`; raise `max_grad_norm`. Log the pre-clip grad norm. | +| Loss drops but structures degrade | Latent drifted off-manifold. | Raise `anchor_weight_*`; fewer steps/rounds; raise `bond_length_weight`. | +| NaNs | Augmentation/churn stochasticity, or bf16 latents. | `EDMSamplerConfig(augmentation=False)`; keep latents float32; lower LR. | +| `ValueError: no optimizable latent on the conditioning` | Wrong attr names (Boltz `"s"`/`"z"` on Protenix). | Use `s_trunk`/`z_trunk`. | +| `ValueError: State atom count != reward_inputs atom count` | Reward built from a different atom set than the model's. | Use the same structure/density the CLI takes. | + +Cheap instrumentation (one line each in `_latent_adam_step`): the pre-clip grad norm (is the clip +biting?) and the anchor's `mean((Δs)²)`, `mean((Δz)²)` (how far the latents drifted). + +## 4. CLI and wiring + +`GuidanceType.LATENT_OPT` routes to `LatentOptimization` in `_run_guidance` +([utils/guidance_script_utils.py](../src/sampleworks/utils/guidance_script_utils.py)): + +```bash +sampleworks-guidance --model protenix --guidance-type latent_opt … +``` + +The same seam resolves the model name to `single_attr`/`pair_attr` (raising a clear `ValueError` for +an unsupported model) and, for Protenix, disables the diffusion shared-vars cache so the `z` gradient +survives. The direct-script route in §2 is the better *debugging* surface — it isolates the scaler +from the grid-search / save machinery. + +**Production footprint** (each tagged with a greppable `IT-opt wiring` comment): +- [utils/guidance_constants.py](../src/sampleworks/utils/guidance_constants.py) — `GuidanceType.LATENT_OPT`. +- [utils/guidance_script_arguments.py](../src/sampleworks/utils/guidance_script_arguments.py) — + `add_latent_opt_args` (`--which-latent`, `--learning-rate`, `--outer-steps`, `--anchor-weight`, + `--max-grad-norm`, `--bond-length-weight`); the six names are also in `_DYNAMIC_ATTRS` so + `GuidanceConfig.from_cli` copies them onto the config. +- [utils/guidance_script_utils.py](../src/sampleworks/utils/guidance_script_utils.py) — the + `LATENT_OPT` dispatch, the Protenix cache-off decision, and the `save_trajectory` case. +- [core/scalers/latent_optimization.py](../src/sampleworks/core/scalers/latent_optimization.py) — + per-latent grad clips and the `latent_drift` diagnostic. + +## 5. Open problems + +1. **Density fit was measured with a local scorer**, not `scripts/eval/rscc_grid_search_script.py`, so + no number so far is repo-exact. A repo-exact run needs a depth-4 trial-dir tree + (`{PROTEIN}_native_occ/{model}_MD/{scaler}/ens{N}_gw{W}/refined.cif`); the generated ensembles are + flat, so symlinks suffice. +2. **Only native-occupancy rows were scored.** Enough to compare conditions on identical rows; + absolute fractions need a wider population. +3. **Five proteins are excluded.** 6RP1, 7Z0E, 4OLE, 8Z76, 2I6H raise "No common atoms found" + (chain/residue-naming mismatch). Fix with `scripts/patch_output_cif_files.py` (needs network for + `rcsb.fetch`) or sequence-based atom matching. diff --git a/src/sampleworks/core/rewards/geometry.py b/src/sampleworks/core/rewards/geometry.py new file mode 100644 index 00000000..bf1ff9a3 --- /dev/null +++ b/src/sampleworks/core/rewards/geometry.py @@ -0,0 +1,242 @@ +"""Coordinate-space bond-geometry regularizer for inference-time latent optimization (IT-opt). + +When the density reward is optimized aggressively, the latent update can distort the denoised +structure -- stretching covalent bonds and driving non-bonded atoms into steric clashes -- to buy a +little density fit. That is the overshoot observed on high-baseline targets (e.g. 1VME, where +z-optimization raised clashes while dropping RSCC). This module adds a differentiable penalty on the +denoised coordinates that the optimizer must trade against, so it cannot reach a good density score +through broken geometry. + +It is a faithful port of the reference ``BondLengthLossFunction`` +(https://github.com/sai-advaith/it_opt, ``protenix/src/losses/bond_length_loss_function.py``): +a bonded-pair length hinge plus a non-bonded steric-clash hinge. Both take the violation's positive +part; the bond term then raises it to ``bond_power`` (default 2) while the clash term stays linear, +matching the reference. The one deliberate divergence is how ``collision_loss`` reduces the +ensemble axis; see the comment there. It is meant +to be an additive term inside ``LatentOptimization``'s per-step loss, not a standalone objective, +and does nothing unless its weight is set. +""" + +from __future__ import annotations + +import gemmi +import numpy as np +import torch +from biotite.structure import AtomArray, connect_via_residue_names +from torch import Tensor + + +# ---- Reading the type hints in this file -------------------------------------- +# A ": type" after a name (or "-> type" after a function) is only a HINT -- it CLAIMS what a +# value should be, but nothing enforces it: pass the wrong type and Python still runs the code, +# and deleting every hint changes nothing. Hints are for humans (and optional checkers like ty). +# (In the table, "|" means "or".) +# +# with the hint plain Python what it claims (useless in runtime) +# element: str element should be a string +# atom_array: AtomArray atom_array should be a biotite AtomArray +# coords: Tensor coords should be a Tensor (~ a numpy array) +# bond_power: int = 2 bond_power = 2 should be an int +# bond_tolerance: float = 0.2 bond_tolerance = 0.2 should be a float +# device: torch.device | str device should be a torch.device or a str +# f(...) -> float f(...) f should return a float +# f(...) -> Tensor f(...) f should return a Tensor +# ------------------------------------------------------------------------------- + + +def _covalent_radius(element: str) -> float: + """Covalent radius of an element (Å). + + gemmi resolves an unrecognized symbol (e.g. the '?' some density inputs carry) to its unknown + element instead of raising, so this never fails -- it returns that element's 0.50 Å radius. + Such an atom gets a too-short ideal bond length, so a correct bond may be penalized, and a + shrunken clash sphere, so its clashes are under-reported. + + This function cannot flag either situation: an unresolvable symbol returns silently, and a + non-str element raises TypeError straight out of gemmi. If bond or clash numbers look wrong, + suspect the atom array's element symbols first. + """ + return float(gemmi.Element(element).covalent_r) + + +class BondGeometryReward: + """Penalize distorted bonds and steric clashes in the denoised structure. + + The topology (which atoms are bonded, their ideal bond lengths, the per-atom-pair collision + distances, and which pairs to score for clashes) is read once from ``atom_array`` at + construction, so each step is just cheap tensor ops on the coords. ``atom_array`` MUST be the + atom set whose ordering matches the coordinate tensor the reward sees -- the *model* atom array, + not the deposited structure -- or the bond indices will not line up with the coordinates. + + Despite the ``Reward`` name (the module's convention), this is a penalty: it is added to, + and minimized as part of, ``LatentOptimization``'s loss. + + Note: this does not satisfy ``RewardFunctionProtocol``. ``__call__`` takes only ``coords``, + not the protocol's per-atom ``elements``/``b_factors``/``occupancies``, so the planned + redefinition of the rewards interface will have to admit penalties of this shape. + """ + + def __init__( + self, + atom_array: AtomArray, + weight: float, + device: torch.device | str, + *, + bond_tolerance: float = 0.2, + clash_padding: float = 0.4, + bond_power: int = 2, + ): + """Build the bond topology and collision-distance matrix from ``atom_array`` once. + + ``weight`` scales the whole penalty; ``bond_tolerance`` (Å) is the slack a bond may deviate + from its ideal length before it is penalized; ``clash_padding`` (Å) is added to the sum of + covalent radii to set how close non-bonded atoms may approach before they count as clashing; + ``bond_power`` is the exponent on the bond-length hinge (2 = quadratic, per the reference). + """ + self.weight = weight + self.bond_tolerance = bond_tolerance + self.clash_padding = clash_padding + self.bond_power = bond_power + + # One covalent radius per atom, looked up once. Both penalties below are sums of two radii, + # so they index this instead of querying gemmi again per bond and per atom. + self._radii = self._build_radii(atom_array, device) + + # Each bond connects atom a to atom b and has an ideal length; these three arrays line up, + # one entry per bond. + self._bond_atom_a, self._bond_atom_b, self._bond_lengths = self._build_bonds( + atom_array, device + ) + self._collision_distances = self._build_collision_distances() + self._scored_pairs = self._build_scored_pairs(len(atom_array), device) + + # ================================ topology (built once) ================================ + + def _build_radii(self, atom_array: AtomArray, device) -> Tensor: + """Covalent radius of every atom, in atom-array order: [n_atoms]. + + A bond's ideal length is the sum of its two endpoints' radii, and a pair's collision + distance is the sum of that pair's radii, so every number both penalties need is a sum of + two entries of this tensor. Looking each atom up once here keeps the lookups to one pass. + """ + radii = [_covalent_radius(element) for element in atom_array.element] + return torch.tensor(radii, dtype=torch.float32, device=device) + + def _build_bonds(self, atom_array: AtomArray, device): + """Bonded-atom endpoints and ideal lengths, inferred from residue-name templates. + + biotite's ``connect_via_residue_names`` infers the intra- and inter-residue bonds from the + standard component templates and returns a ``BondList``. By default that list already holds + each bond exactly once with the lower atom index first, so ``as_array()`` gives us the edge + list directly -- its columns are atom i, atom j, bond type. (``get_all_bonds()`` would + instead return padded per-atom adjacency, which lists every bond twice.) A bond's ideal + length is the sum of the two atoms' covalent radii. + + Returns three aligned 1-D tensors, each of length n_bonds: the first endpoint of every bond, + the second endpoint, and the ideal length. + """ + bonds = connect_via_residue_names(atom_array).as_array() # [n_bonds, 3]: i, j, bond type + # as_array() is uint32, which torch refuses to convert, so cast to a signed integer. + pair_array = bonds[:, :2].astype(np.int64) + + atom_a = torch.tensor(pair_array[:, 0], dtype=torch.long, device=device) + atom_b = torch.tensor(pair_array[:, 1], dtype=torch.long, device=device) + + # Ideal length of a bond = sum of the two atoms' covalent radii. + bond_lengths = self._radii[atom_a] + self._radii[atom_b] + return atom_a, atom_b, bond_lengths + + def _build_collision_distances(self) -> Tensor: + """The [n_atoms, n_atoms] matrix of minimum non-clashing distances (covalent-radii sums). + + Entry (i, j) is the sum of atom i's and atom j's covalent radii; two non-bonded atoms closer + than this (plus ``clash_padding``) are overlapping. + """ + # radius_i + radius_j for every pair: a column plus a row vector broadcasts to [n, n]. + radius_column = self._radii.unsqueeze(1) # [n, 1] + radius_row = self._radii.unsqueeze(0) # [1, n] + return radius_column + radius_row + + def _build_scored_pairs(self, n_atoms: int, device) -> Tensor: + """Boolean [n_atoms, n_atoms] mask: which atom pairs count toward the clash penalty. + + We penalize every pair EXCEPT an atom with itself (distance 0) and directly-bonded atoms + (they should sit close together). Built once from the fixed bond topology; the mask + is symmetric, so each clashing pair is counted twice in the sum, matching the reference. + """ + scored = torch.ones((n_atoms, n_atoms), dtype=torch.bool, device=device) + # an atom never clashes with itself + scored.fill_diagonal_(False) + # directly-bonded atoms are supposed to sit close, so don't count them as clashes + scored[self._bond_atom_a, self._bond_atom_b] = False + scored[self._bond_atom_b, self._bond_atom_a] = False + return scored + + # ============================== the two penalty terms ============================== + + def bond_length_loss(self, coords: Tensor) -> Tensor: + """Hinge on bonded-pair length deviation beyond ``bond_tolerance``. + + ``coords`` is the denoised coordinate tensor, shape [ensemble_members, atoms, 3]. A bond is + free within ``bond_tolerance`` of its ideal length; past that the deviation is + raised to ``bond_power`` and summed over all bonds and ensemble members. Keeps the optimizer + from stretching or compressing covalent bonds to chase density. + """ + if self._bond_lengths.numel() == 0: + return coords.new_zeros(()) + + # Current length of every bond, for every ensemble member: [ensemble, n_bonds]. + pos_a = coords[:, self._bond_atom_a] + pos_b = coords[:, self._bond_atom_b] + lengths = (pos_a - pos_b).norm(dim=-1) + + # The deviation is unsigned, so a stretched and a compressed bond of equal error are + # penalized alike. The clamp below is the tolerance window -- deviations under it are free + # -- not a one-sided penalty like the one in collision_loss, which clamps a signed gap. + deviation = (lengths - self._bond_lengths).abs() + excess = (deviation - self.bond_tolerance).clamp(min=0) + return excess.pow(self.bond_power).sum() + + def collision_loss(self, coords: Tensor) -> Tensor: + """Hinge on non-bonded atoms closer than their collision distance + ``clash_padding``. + + ``coords`` is the denoised coordinate tensor, shape [ensemble, atoms, 3]. This is the + steric-clash term: per atom pair we take the worst member's overlap plus the ensemble mean, + then sum over the scored pairs. The reference used the worst member alone; the comment on + the reduction below compares the four options and says why we add the mean. Note the + pairwise-distance tensor is O(n_atoms**2), so memory scales with the square of the count. + """ + # Distance between every pair of atoms, for each ensemble member: [ensemble, n, n]. + # torch.cdist gives the same distances without building the [ensemble, n, n, 3] grid of + # coordinate differences a manual broadcast would (3x the distance matrix). Verified on the + # pinned torch (2.7): its gradient is finite and matches the broadcast form even for + # coincident atoms, so the historical cdist zero-distance NaN does not apply here. + distances = torch.cdist(coords, coords) + + # How far each pair is inside its allowed distance (positive = overlapping). We clamp per + # member, before reducing; the reference clamps after its max, which is equivalent there + # (relu commutes with max) but not once the mean below is added, where a comfortably + # separated member's negative gap would cancel another member's real overlap. + min_distance = self._collision_distances + self.clash_padding + overlap_per_member = (min_distance - distances).clamp(min=0) + + # Reduce the ensemble axis, then drop self- and bonded pairs. Four reductions were weighed; + # the value column is relative to max, the reference's choice, with E ensemble members: + # + # max 1x only the worst member gets gradient -- the rest are ignored + # mean 1/E .. 1x every member counted, but can weaken the penalty E-fold + # max + mean 1x .. 2x every member counted, never weaker than max, capped at 2x + # sum 1x .. Ex every member counted, but scales the penalty with E + # + # We take max + mean. The max keeps the worst member's full push, so the penalty can never + # come out weaker than it is today -- which matters because bond_length_weight is tuned to + # the smallest value that fixes clashes, and mean alone could drop below that floor. The + # mean then gives every other clashing member a share instead of nothing, which is what + # max alone failed to do. sum would do that too, but its value grows with ensemble size. + combined_overlap = overlap_per_member.max(dim=0).values + overlap_per_member.mean(dim=0) + scored_overlap = combined_overlap * self._scored_pairs # [n, n]; unscored pairs zeroed + return scored_overlap.sum() + + def __call__(self, coords: Tensor) -> Tensor: + """Weighted sum of the bond-length and collision penalties, added to the loss.""" + return self.weight * (self.bond_length_loss(coords) + self.collision_loss(coords)) diff --git a/src/sampleworks/core/scalers/latent_optimization.py b/src/sampleworks/core/scalers/latent_optimization.py new file mode 100644 index 00000000..ec2212df --- /dev/null +++ b/src/sampleworks/core/scalers/latent_optimization.py @@ -0,0 +1,566 @@ +"""Inference-time latent-space optimization (IT-opt). + +A Sampleworks port of the method in Maddipatla et al., "Inference-time optimization for +experiment-grounded protein ensemble generation" (ICML 2026, https://arxiv.org/abs/2602.24007). +Reference implementation: https://github.com/sai-advaith/it_opt -- cited paths such as +``run_it_optimization`` in ``protenix/it_optimization_manager.py`` are relative to that repo, and +name symbols rather than line numbers, which drift as that repo changes. + +It treats the frozen structure model as a differentiable sampler and optimizes its cached +post-trunk latents -- the single representation ``s`` and the pair representation ``z`` -- against +an experimental reward evaluated on the *denoised* structure. No model weights are updated. See +``docs/IT_OPT_DESIGN.md`` for the design and its mapping to the reference. + +The algorithm (kept faithful to the reference), with the reference's bugs fixed: + + extract (s, z) from the trunk once -> make them optimizable leaves + for each optimization round (fresh diffusion noise each round): + build ONE Adam over [s, z] # reference rebuilt it per step + for each diffusion step: + x0_hat = differentiable denoise(noisy_t, s, z) # one forward + loss = reward(x0_hat) + anchor(s, z) + loss.backward(); per-latent grad-clip (s, then z); adam.step() + advance the trajectory one step (latents frozen) + final clean sampling round with the optimized latents -> saved ensemble + +Terminology: a **round** is one full traversal of the ``num_steps`` diffusion schedule; a **step** +is one point on it, costing one denoiser forward. ``outer_steps`` counts the optimization rounds, +and one final sampling round follows them. "Pass" is used only in its usual sense of a forward +pass through the model, never as a synonym for either unit. + +Only Boltz1 and RF3 propagate gradients cleanly to the cached latents out of the +box. Protenix now does too -- its ``step()`` keeps an injected latent leaf attached +via ``detach_unless_leaf`` -- but still needs its diffusion cache disabled for z; +Boltz2 needs its diffusion conditioning recomputed. See ``docs/IT_OPT_TESTING.md``. + +v1 does ONLY latent optimization: coordinate-space guidance is not applied (the +attached scaler produces a zero coordinate direction), so the sole steering comes +from the evolving latents. +""" + +from __future__ import annotations + +import dataclasses +from collections.abc import Sequence + +import torch +from loguru import logger +from torch import Tensor +from tqdm import tqdm + +from sampleworks.core.rewards.geometry import BondGeometryReward +from sampleworks.core.rewards.protocol import RewardFunctionProtocol +from sampleworks.core.samplers.protocol import TrajectorySampler +from sampleworks.core.scalers.protocol import GuidanceOutput, StepScalerProtocol +from sampleworks.core.scalers.step_scalers import NoScalingScaler +from sampleworks.eval.structure_utils import process_structure_to_trajectory_input +from sampleworks.models.latent_adapter import AttrLatentIO +from sampleworks.models.protocol import FlowModelWrapper, GenerativeModelInput + + +# ---- Reading the type hints in this file -------------------------------------- +# A ": type" after a name (or "-> type" after a function) is only a HINT -- it CLAIMS what a +# value should be, but nothing enforces it: pass the wrong type and Python still runs the code, +# and deleting every hint changes nothing. Hints are for humans (and optional checkers like ty). +# (In the table, "|" means "or".) +# +# with the hint plain Python what it claims (useless in runtime) +# num_steps: int = 200 num_steps = 200 should be an int +# learning_rate: float = 0.05 learning_rate = 0.05 should be a float +# optimize_single: bool = True optimize_single = True should be True or False +# structure: dict structure should be a dict +# model: FlowModelWrapper model should be a FlowModelWrapper +# weights: Sequence[float] weights should be a list/tuple of floats +# latents: Sequence[Tensor] latents should be a list/tuple of Tensors +# losses: list[float] losses should be a list of floats +# denoised: Tensor | None denoised should be a Tensor, or None +# f(...) -> GuidanceOutput f(...) f should return a GuidanceOutput +# ------------------------------------------------------------------------------- + + +class _GradEnablingScaler(NoScalingScaler): + """``NoScalingScaler`` that also asks the sampler to run its denoiser under autograd. + + The zero-guidance behavior is inherited unchanged; the flag below is the only difference. + ``AF3EDMSampler`` runs its denoiser forward under autograd only when the attached scaler + exposes ``requires_gradients=True``, and it reads that with ``getattr(..., False)``. Plain + ``NoScalingScaler`` therefore leaves :attr:`SamplerStepOutput.denoised` detached, and IT-opt + would get no gradient back to the injected latent leaves. + + Keeping guidance at zero leaves the diffusion *advance* unguided, realizing "coordinate + guidance disabled" without editing the sampler. The reward is recomputed by the trajectory + scaler on ``denoised``, so there is exactly one denoiser forward and one reward evaluation + per step (the reference's Tier-2 single-forward optimization, for free). + """ + + requires_gradients = True + + +class LatentAnchor: + r"""On-manifold prior penalizing drift of the latents from their trunk baseline. + + ``penalty = Σ_i w_i · mean((latent_i − baseline_i)²)``. + + The reference ``AnchorLossFunction`` uses a per-sample Frobenius **norm** + ``‖latent − baseline‖`` (see ``anchor_loss_function.py``). We use the *mean + squared* deviation instead: it is shape-agnostic and normalizes for the very + different element counts of ``s`` and ``z`` (``z`` is O(tokens) larger), so the + weights ``w_s`` and ``w_z`` land on a comparable scale -- part of keeping the + two latent updates harmonized. Retune the weights when switching objectives. + + Planned move: this belongs under ``core.rewards`` with + ``BondGeometryReward``, so regularizers are collected and reusable. Deferred + because its input is *latents*, not coordinates, so it fits neither + ``RewardFunctionProtocol`` nor the coordinate-space shape of that package -- + the rewards interface has to admit latent-space terms first. + """ + + def __init__(self, weights: Sequence[float]): + self.weights = list(weights) + + def __call__(self, latents: Sequence[Tensor], baselines: Sequence[Tensor]) -> Tensor: + terms = [ + w * torch.mean((lat - base) ** 2) + for w, lat, base in zip(self.weights, latents, baselines) + ] + return torch.stack(terms).sum() + + +class LatentOptimization: + """Trajectory scaler that optimizes the model's ``s``/``z`` latents (IT-opt). + + Satisfies ``TrajectoryScalerProtocol`` and drops in alongside ``PureGuidance`` + / ``FKSteering``. + """ + + def __init__( + self, + ensemble_size: int = 1, + num_steps: int = 200, + guidance_t_start: float = 0.0, + *, + outer_steps: int = 1, + learning_rate: float = 0.05, + max_grad_norm: float = 1.0, + optimize_single: bool = True, + optimize_pair: bool = True, + anchor_weight_single: float = 0.0, + anchor_weight_pair: float = 0.0, + bond_length_weight: float = 5e-5, # smallest weight that fixes mean clash; see docstring + single_attr: str = "s", + pair_attr: str = "z", + ): + """Initialize the latent-optimization trajectory scaler. + + Parameters + ---------- + ensemble_size + Number of structures sampled in parallel (they share the latents). + num_steps + Number of diffusion steps in one round. + guidance_t_start + Fraction of ``num_steps`` after which to begin optimizing within each + round. Before it, steps are plain frozen-latent diffusion. Default 0 + (optimize from the first step, as the reference does). + outer_steps + Number of optimization rounds. Each traverses the whole schedule with fresh + prior noise, sharing the persisting latents (the reference's + ``outer_diffusion_steps``). Default 1 for a cheap first run. + learning_rate + Adam learning rate. A *real* persistent Adam is built once per round -- + unlike the reference, which rebuilt it every diffusion step and thereby + degenerated to signed gradient descent (its headline bug). + max_grad_norm + Per-latent gradient-clip threshold: each optimized latent (``s``, ``z``) + is clipped to this norm **independently** (matching the reference), so + ``s``'s update is not scaled down by ``z``'s much larger gradient. A single + joint clip over ``[s, z]`` does the opposite -- ``z`` dominates the shared + coefficient and starves ``s``. + optimize_single, optimize_pair + Which representations to optimize. Both by default. + anchor_weight_single, anchor_weight_pair + Weights of the on-manifold L2-to-baseline anchor per latent. + bond_length_weight + Weight of the coordinate-space bond-geometry penalty (bond-length + steric-clash + hinges; see ``BondGeometryReward``). Pure latent opt (0) improves density fit but + raises clashes; this penalty curbs that. Full-40 sweep -- mean/median clash, RSCC>=0.8 + (unguided baseline = 0.38 / 0.00 / 42%): + 0 -> 0.47 / 0.25 / 92% (density gained, but clashes rose) + 5e-5 -> 0.38 / 0.25 / 92% (mean clash back to baseline) [DEFAULT] + 1e-4 -> 0.39 / 0.25 / 92% + 1e-3 -> 0.35 / 0.00 / 98% (median clash 0; over-constrains, split 18->11%) + Default 5e-5: smallest weight that fixes the mean clash while keeping the density gain + and altloc diversity; use 1e-3 if you need median clash at 0 (accepts less spread). + single_attr, pair_attr + Conditioning attribute names for the single / pair representation + (``"s"``/``"z"`` for Boltz, ``"s_trunk"``/``"z_trunk"`` for + Protenix/RF3). + """ + logger.info( + f"Initialized LatentOptimization (IT-opt): outer_steps={outer_steps}, " + f"num_steps={num_steps}, lr={learning_rate}, " + f"optimize_single={optimize_single}, optimize_pair={optimize_pair}, " + f"bond_length_weight={bond_length_weight}." + ) + self.ensemble_size = ensemble_size + self.num_steps = num_steps + self.guidance_start = int(guidance_t_start * num_steps) + self.outer_steps = outer_steps + self.learning_rate = learning_rate + self.max_grad_norm = max_grad_norm + self.optimize_single = optimize_single + self.optimize_pair = optimize_pair + self.anchor_weight_single = anchor_weight_single + self.anchor_weight_pair = anchor_weight_pair + self.bond_length_weight = bond_length_weight + self.single_attr = single_attr + self.pair_attr = pair_attr + + def sample( + self, + structure: dict, + model: FlowModelWrapper, + sampler: TrajectorySampler, + step_scaler: StepScalerProtocol, + reward: RewardFunctionProtocol, + num_particles=1, + ) -> GuidanceOutput: + """Optimize the latents against ``reward``, then sample the final ensemble. + + ``step_scaler`` is ignored: v1 does latent optimization only (no + coordinate-space guidance). ``reward`` is evaluated on the + reconciler-aligned denoised prediction. + """ + io = AttrLatentIO( + single_attr=self.single_attr, + pair_attr=self.pair_attr if self.optimize_pair else None, + ) + + # --- extract (s, z) once and make them optimizable leaves --------------- + # no_grad avoids retaining a graph into the (frozen) trunk; we build leaves + # from detached copies. (reference: get_msa_features + clone/detach.) + with torch.no_grad(): + features = model.featurize(structure) + features, latents, baselines, anchor_weights = self._leaf_latents(features, io) + anchor = LatentAnchor(anchor_weights) + + # --- shared per-trajectory context (reconciler + reward inputs) --------- + coords = torch.as_tensor( + model.initialize_from_prior(batch_size=self.ensemble_size, features=features) + ) + processed = process_structure_to_trajectory_input( + structure=structure, + coords_from_prior=coords, + features=features, + ensemble_size=self.ensemble_size, + ) + reconciler = processed.reconciler.to(coords.device) + reward_inputs = processed.to_reward_inputs(device=coords.device) + schedule = sampler.compute_schedule(num_steps=self.num_steps) + grad_enabler = _GradEnablingScaler() + + # --- optional coordinate-space geometry penalty ------------------------- + # BondGeometryReward penalizes stretched bonds and steric clashes in the denoised structure, + # curbing the overshoot where an aggressive latent update trades geometry for density fit. + if self.bond_length_weight > 0: + # Prefer the model atom array; fall back to the input atom array if it is absent. + geometry_atom_array = processed.model_atom_array or processed.atom_array + bond_geometry = BondGeometryReward( + geometry_atom_array, self.bond_length_weight, coords.device + ) + else: + bond_geometry = None + + # --- optimize the latents (outer resample × inner diffusion steps) ------ + # These two lists are diagnostics only -- reported in the returned metadata, never read by + # the loop. One entry is appended per round. + optimization_losses = [] # per round: the per-step data losses + latent_drift = [] # per round: how far each latent moved from its baseline (see below) + for outer in tqdm(range(self.outer_steps)): + optimizer = torch.optim.Adam(latents, lr=self.learning_rate) # a fresh, persistent Adam + round_losses = self._optimize_one_round( + model=model, + sampler=sampler, + reward=reward, + features=features, + latents=latents, + baselines=baselines, + anchor=anchor, + optimizer=optimizer, + schedule=schedule, + reconciler=reconciler, + alignment_reference=processed.input_coords, + reward_inputs=reward_inputs, + grad_enabler=grad_enabler, + round_index=outer, + bond_geometry=bond_geometry, + ) + optimization_losses.append(round_losses) + # relative drift of each latent from its trunk baseline -- shows which latent + # actually moved (s vs z), which the loss trend alone cannot reveal. + latent_drift.append( + [ + float((lat.detach() - base).norm() / (base.norm() + 1e-12)) + for lat, base in zip(latents, baselines) + ] + ) + + # --- final clean sampling round with the optimized latents -------------- + final_coords, trajectory, losses = self._sample_with_frozen_latents( + model=model, + sampler=sampler, + reward=reward, + io=io, + features=features, + latents=latents, + schedule=schedule, + reconciler=reconciler, + alignment_reference=processed.input_coords, + reward_inputs=reward_inputs, + ) + + metadata = { + "optimization_losses": optimization_losses, + "latent_drift": latent_drift, + } + if reconciler.has_mismatch and processed.model_atom_array is not None: + metadata["model_atom_array"] = processed.model_atom_array + + return GuidanceOutput( + structure=structure, + final_state=final_coords, + trajectory=trajectory, + losses=losses, + metadata=metadata, + ) + + def _leaf_latents(self, features: GenerativeModelInput, io: AttrLatentIO): + """Replace ``s``/``z`` on the conditioning with fresh optimizable leaves. + + Returns the rewritten ``features`` plus parallel lists of leaves, their + detached baselines (anchor targets), and per-latent anchor weights. Each + leaf is a detached clone made ``requires_grad=True`` -- a true leaf severed + from any trunk graph, so Adam updates it directly (leaves persist and are + updated in place across rounds and steps). Shapes are preserved (whatever + the wrapper caches), so no assumption is made about a batch dimension. + """ + conditioning = features.conditioning + latents: list[Tensor] = [] + baselines: list[Tensor] = [] + anchor_weights: list[float] = [] + + # One row per optimizable latent: (optimize it?, its attribute on the conditioning, its + # anchor weight). We handle the single representation, then the pair. + specs = ( + (self.optimize_single, io.single_attr, self.anchor_weight_single), + (self.optimize_pair, io.pair_attr, self.anchor_weight_pair), + ) + for enabled, attr, anchor_weight in specs: + if not enabled: + continue + baseline = getattr(conditioning, attr, None) + if baseline is None: + # Enabled but absent means the configured attribute name is wrong for this model. + # Skipping it would optimize fewer latents than asked for and look like success. + raise ValueError( + f"LatentOptimization was asked to optimize {attr!r}, but the model's " + "conditioning does not expose it. Check the attribute names for this model." + ) + baseline = baseline.detach() + leaf = baseline.clone().requires_grad_(True) + # The conditioning is a frozen dataclass, so setattr would raise; replace() returns a + # copy with this one field swapped and every sidecar field left untouched. + conditioning = dataclasses.replace(conditioning, **{attr: leaf}) + latents.append(leaf) + baselines.append(baseline) + anchor_weights.append(anchor_weight) + + # A missing attribute already raised above, so reaching here means neither latent was + # enabled -- nothing to optimize, which is a misconfiguration rather than a valid run. + if not latents: + raise ValueError( + "LatentOptimization has no latent enabled " + f"(optimize_single={self.optimize_single}, optimize_pair={self.optimize_pair})." + ) + # Rebuild with the rewritten conditioning only. #330 removed x_init from + # GenerativeModelInput; initial coords now come from initialize_from_prior at sampling time. + features = GenerativeModelInput(conditioning=conditioning) + return features, latents, baselines, anchor_weights + + def _optimize_one_round( + self, + *, + model, + sampler, + reward, + features, + latents, + baselines, + anchor, + optimizer, + schedule, + reconciler, + alignment_reference, + reward_inputs, + grad_enabler, + round_index, + bond_geometry, + ) -> list[float]: + """One optimization round: a full traversal of the schedule that updates the latents. + + Fresh prior noise; the persisting latents are updated once per diffusion + step against the reward on that step's denoised prediction, then the + trajectory is advanced with the latents frozen. Returns the per-step data + losses. + """ + coords = torch.as_tensor( + model.initialize_from_prior(batch_size=self.ensemble_size, features=features) + ) + losses: list[float] = [] + steps = tqdm(range(self.num_steps), f"IT-opt round {round_index}") + for i in steps: + optimize = i >= self.guidance_start + context = sampler.get_context_for_step(i, schedule) + if optimize: + context = context.with_reward(reward, reward_inputs) + context = context.with_reconciler( + reconciler=reconciler, alignment_reference=alignment_reference + ) + + step_output = sampler.step( + state=coords, + model_wrapper=model, + context=context, + scaler=grad_enabler if optimize else None, + features=features, + ) + + if optimize: + data_loss = self._latent_adam_step( + denoised=step_output.denoised, + reward=reward, + reward_inputs=reward_inputs, + optimizer=optimizer, + anchor=anchor, + latents=latents, + baselines=baselines, + bond_geometry=bond_geometry, + ) + losses.append(data_loss) + steps.set_postfix(loss=data_loss) + + coords = step_output.state.detach() + return losses + + def _latent_adam_step( + self, + *, + denoised: Tensor | None, + reward: RewardFunctionProtocol, + reward_inputs, + optimizer: torch.optim.Optimizer, + anchor: LatentAnchor, + latents: Sequence[Tensor], + baselines: Sequence[Tensor], + bond_geometry: BondGeometryReward | None = None, + ) -> float: + """Score the denoised structure, backprop to the latents, take one Adam step. + + ``denoised`` is the sampler's aligned prediction and carries a graph back to + the latent leaves. When ``bond_geometry`` is supplied, its bond-length/clash + penalty is added to the loss alongside the anchor. Each latent is + gradient-clipped **separately** so ``s``'s step is not scaled down by ``z``'s + much larger gradient. Returns the data-only reward for logging. + """ + if denoised is None: + raise RuntimeError( + "Sampler returned no denoised prediction; the grad-enabling scaler " + "should have been attached this step so the latent gradient exists." + ) + data_loss = reward( + coordinates=denoised, + elements=reward_inputs.elements, + b_factors=reward_inputs.b_factors, + occupancies=reward_inputs.occupancies, + ) + loss = data_loss + anchor(latents, baselines) + # Add the coordinate-space geometry penalty (bond-length + clash) when it is enabled, so the + # latent update is pushed toward density fits that keep valid geometry; it backpropagates + # through the same denoised prediction to the latents. + if bond_geometry is not None: + loss = loss + bond_geometry(denoised) + + optimizer.zero_grad() + loss.backward() + # Clip each latent SEPARATELY (the reference makes two clip_grad_norm_ calls too), never + # as one joint [s, z] group. NB: with the density reward the gradients are tiny (grad-norm + # ~1e-4, far below a max_grad_norm of 1.0), so this clip is effectively inert here -- it + # bites only for rewards whose gradients exceed the threshold, where separate (not joint) + # clips keep s's step decoupled from z's much larger gradient scale. + for latent in latents: + torch.nn.utils.clip_grad_norm_(latent, self.max_grad_norm) + optimizer.step() + return float(data_loss.detach()) + + def _sample_with_frozen_latents( + self, + *, + model, + sampler, + reward, + io, + features, + latents, + schedule, + reconciler, + alignment_reference, + reward_inputs, + ): + """Final clean sampling round with the optimized latents held fixed. + + No optimization and no coordinate guidance (``scaler=None``). This produces + the ensemble that is returned/saved -- the reference's + ``run_diffusion_process_it_optimized``. The per-step reward (evaluated under + no-grad on the denoised prediction) is returned for logging. + """ + # Re-inject detached copies so the final round is purely a frozen sampler. + conditioning = features.conditioning + detached = [lat.detach() for lat in latents] + if self.optimize_single and io.read_single(conditioning) is not None: + conditioning = io.write_single(conditioning, detached.pop(0)) + if self.optimize_pair and io.read_pair(conditioning) is not None: + conditioning = io.write_pair(conditioning, detached.pop(0)) + frozen_features = GenerativeModelInput(conditioning=conditioning) # #330: conditioning-only + + coords = torch.as_tensor( + model.initialize_from_prior(batch_size=self.ensemble_size, features=frozen_features) + ) + trajectory: list[Tensor] = [] + losses: list[float | None] = [] + steps = tqdm(range(self.num_steps), "IT-opt final sampling") + for i in steps: + context = sampler.get_context_for_step(i, schedule).with_reconciler( + reconciler=reconciler, alignment_reference=alignment_reference + ) + step_output = sampler.step( + state=coords, + model_wrapper=model, + context=context, + scaler=None, + features=frozen_features, + ) + coords = step_output.state.detach() + trajectory.append(coords.clone().cpu()) + + if step_output.denoised is not None: + with torch.no_grad(): + loss = reward( + coordinates=step_output.denoised, + elements=reward_inputs.elements, + b_factors=reward_inputs.b_factors, + occupancies=reward_inputs.occupancies, + ) + losses.append(float(loss)) + else: + losses.append(None) + return coords, trajectory, losses diff --git a/src/sampleworks/models/latent_adapter.py b/src/sampleworks/models/latent_adapter.py new file mode 100644 index 00000000..9ebcf2b2 --- /dev/null +++ b/src/sampleworks/models/latent_adapter.py @@ -0,0 +1,144 @@ +"""Read and write a model's post-trunk latents (IT-opt plumbing). + +The latents are the *single* representation ``s`` and the *pair* representation ``z``. They live on +the model's *conditioning* -- the bundle of inputs ``featurize`` produces and hands to the ``step`` +(denoise) call, a (frozen) dataclass on which ``s`` and ``z`` are named attributes. This module lets +inference-time latent optimization (IT-opt) read those attributes and swap them out, so the scaler +in ``core/scalers/latent_optimization.py`` can optimize them against an experimental reward. See +``docs/IT_OPT_DESIGN.md``. + +What each piece is for +---------------------- +- :class:`AttrLatentIO` -- reads/writes ``s`` and ``z`` by attribute name. This is the piece the + IT-opt scaler actually uses: it reads the baseline latents once, holds its own optimizable copies, + and writes them back each step. +- :class:`LatentIO` -- the protocol (read/write contract) that :class:`AttrLatentIO` satisfies. + +Design goal +----------- +- **Minimal, model-agnostic.** The only model-specific knowledge is *which attribute of the + conditioning holds each representation* -- ``"s"``/``"z"`` for Boltz, ``"s_trunk"``/``"z_trunk"`` + for Protenix/RF3. No model package is imported here, so a 4th or 5th model is one more pair of + strings, not new code. +""" + +from __future__ import annotations + +import dataclasses +from typing import Protocol, runtime_checkable + +from torch import Tensor + + +# ---- Reading the type hints in this file -------------------------------------- +# A ": type" after a name (or "-> type" after a function) is only a HINT -- it CLAIMS what a +# value should be, but nothing enforces it: pass the wrong type and Python still runs the code, +# and deleting every hint changes nothing. Hints are for humans (and optional checkers like ty). +# (In the table, "|" means "or".) +# +# with the hint plain Python what it claims (useless in runtime) +# attr: str attr should be a string +# single: Tensor single should be a Tensor (~ a numpy array) +# pair_attr: str | None = None pair_attr = None should be a str, or None +# d: dict[str, str] = {} d = {} should be a dict, string -> string +# f(...) -> Tensor | None f(...) f should return a Tensor or None +# ------------------------------------------------------------------------------- + +# Convenience maps of the representation attribute name per model. 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", +} +DEFAULT_PAIR_REP_ATTR: dict[str, str] = { + "boltz1": "z", + "boltz2": "z", + "protenix": "z_trunk", + "rf3": "z_trunk", +} + + +# "conditioning" throughout this file is the model's conditioning object: a (frozen) dataclass +# carrying the latents ``s``/``z`` (and other cached state) as named attributes. +class AttrLatentIO: + """A general-purpose :class:`LatentIO` addressing each representation by attribute name. + + Works for any (dataclass) conditioning whose representations are stored in + named attributes. Writes use :func:`dataclasses.replace`, so non-tensor + sidecar state on the conditioning (e.g. RF3's chiral tracking arrays) is + preserved untouched. All three model wrappers use + ``@dataclass(frozen=True, slots=True)`` conditioning, which is exactly this + case. + + Parameters + ---------- + single_attr + Attribute holding the single representation, e.g. ``"s"`` (Boltz) or + ``"s_trunk"`` (Protenix/RF3). + pair_attr + Attribute holding the pair representation, e.g. ``"z"`` / ``"z_trunk"``. + When ``None`` (default), the pair accessors are no-ops -- single-rep-only + behaviour, matching the original scaffold. + """ + + def __init__(self, single_attr: str, pair_attr: str | None = None): + self.single_attr = single_attr + self.pair_attr = pair_attr + + def read_single(self, conditioning) -> Tensor: + # No default: a name that does not match this model's conditioning is a configuration + # error, and AttributeError says so more usefully than a silent None would. + return getattr(conditioning, self.single_attr) + + def write_single(self, conditioning, single: Tensor): + return dataclasses.replace(conditioning, **{self.single_attr: single}) # ty: ignore + + def read_pair(self, conditioning) -> Tensor | None: + # None here means "this io addresses no pair rep", which is a supported configuration -- + # unlike a missing attribute above. + if self.pair_attr is None: + return None + return getattr(conditioning, self.pair_attr) + + def write_pair(self, conditioning, pair: Tensor): + if self.pair_attr is None: + return conditioning + return dataclasses.replace(conditioning, **{self.pair_attr: pair}) # ty: ignore + + +# ============================ protocol (template / interface) ============================ +# A Protocol is just a checklist of methods a class must have -- think a C header, or an abstract +# base class, but *structural*: any class that already has these methods counts as a LatentIO +# without inheriting from anything. This definition does not run; it only documents the contract and +# lets the ty checker flag a mismatch. @runtime_checkable also lets isinstance(obj, LatentIO) work, +# which the unit tests use. The real implementation is above: AttrLatentIO is the LatentIO. + + +@runtime_checkable +class LatentIO(Protocol): + """Reads/writes the single and pair representations on a conditioning object. + + Implementations encapsulate the *only* model-specific knowledge in this + module: where each representation lives on the conditioning object. Pair-rep + methods return ``None`` / pass through unchanged when the implementation is + single-rep only. Each write returns a copy of the conditioning with one + representation replaced. + """ + + def read_single(self, conditioning) -> Tensor | None: + """Return the single representation tensor, or ``None`` if unavailable.""" + ... + + def write_single(self, conditioning, single: Tensor): + """Return a copy of ``conditioning`` with the single representation replaced.""" + ... + + def read_pair(self, conditioning) -> Tensor | None: + """Return the pair representation tensor, or ``None`` if unavailable.""" + ... + + def write_pair(self, conditioning, pair: Tensor): + """Return a copy of ``conditioning`` with the pair representation replaced.""" + ... diff --git a/src/sampleworks/models/protenix/wrapper.py b/src/sampleworks/models/protenix/wrapper.py index fa307a12..14b4a652 100644 --- a/src/sampleworks/models/protenix/wrapper.py +++ b/src/sampleworks/models/protenix/wrapper.py @@ -637,13 +637,24 @@ def step( t_tensor = match_batch(t_tensor, target_batch_size=x_t.shape[0]) - # When gradients are enabled, detach cached pairformer outputs so gradients - # only flow through the diffusion module (not back through the pairformer). - # The pairformer was computed with grad_needed=False, so its graph isn't retained. + # Detach cached pairformer outputs under grad so the many denoising steps that reuse them + # don't backprop through the trunk twice -- EXCEPT a latent that IT-opt has injected as an + # optimizable leaf (requires_grad=True), which is kept attached so its gradient survives. + # In the guidance path featurize() runs the trunk under no_grad, so every cached latent + # is a requires_grad=False constant and detaches -- the exact behavior before IT-opt. grad_needed = torch.is_grad_enabled() - s_inputs = cond.s_inputs.detach() if grad_needed else cond.s_inputs - s_trunk = cond.s_trunk.detach() if grad_needed else cond.s_trunk - z_trunk = cond.z_trunk.detach() if grad_needed else cond.z_trunk + + def detach_unless_leaf(latent: Tensor) -> Tensor: + """Detach a cached latent under grad, unless it is an optimizable IT-opt leaf.""" + return latent.detach() if grad_needed and not latent.requires_grad else latent + + s_inputs = detach_unless_leaf(cond.s_inputs) + s_trunk = detach_unless_leaf(cond.s_trunk) + z_trunk = detach_unless_leaf(cond.z_trunk) + # pair_z / p_lm / c_l are z-derived caches. When optimizing z_trunk, featurize with + # enable_diffusion_shared_vars_cache=False so these are None and the diffusion module + # recomputes them from the live z_trunk; otherwise the z gradient is only partial. + # See docs/IT_OPT_TESTING.md. pair_z = cond.pair_z.detach() if grad_needed and cond.pair_z is not None else cond.pair_z p_lm = cond.p_lm.detach() if grad_needed and cond.p_lm is not None else cond.p_lm c_l = cond.c_l.detach() if grad_needed and cond.c_l is not None else cond.c_l diff --git a/src/sampleworks/utils/guidance_constants.py b/src/sampleworks/utils/guidance_constants.py index f3be6965..6757b763 100644 --- a/src/sampleworks/utils/guidance_constants.py +++ b/src/sampleworks/utils/guidance_constants.py @@ -10,10 +10,12 @@ class GuidanceType(StrEnum): References: - Feynman-Kaç steering http://arxiv.org/abs/2501.06848 - Pure guidance (DPS) http://arxiv.org/abs/2209.14687 + - Latent optimization (IT-opt) https://arxiv.org/abs/2602.24007 """ FK_STEERING = "fk_steering" PURE_GUIDANCE = "pure_guidance" + LATENT_OPT = "latent_opt" class StructurePredictor(StrEnum): @@ -47,6 +49,7 @@ class TrajectoryScalers(StrEnum): PURE_GUIDANCE = "pure_guidance" FK_STEERING = "fk_steering" + LATENT_OPT = "latent_opt" class Rewards(StrEnum): diff --git a/src/sampleworks/utils/guidance_script_arguments.py b/src/sampleworks/utils/guidance_script_arguments.py index a934f00b..53e773e3 100644 --- a/src/sampleworks/utils/guidance_script_arguments.py +++ b/src/sampleworks/utils/guidance_script_arguments.py @@ -191,6 +191,14 @@ def validate_model_checkpoint( "num_gd_steps", "guidance_weight", "guidance_interval", + # latent optimization (IT-opt) -- must be listed here or from_cli() drops the parsed + # values and _run_guidance()'s getattr(args, ...) always sees the defaults (flags = no-ops). + "which_latent", + "learning_rate", + "outer_steps", + "anchor_weight", + "max_grad_norm", + "bond_length_weight", # model-specific "model_checkpoint", "method", @@ -646,9 +654,46 @@ def add_protpardelle_specific_args(parser: argparse.ArgumentParser | GuidanceCon "protpardelle": add_protpardelle_specific_args, } + +def add_latent_opt_args(parser: argparse.ArgumentParser | GuidanceConfig): + """Add CLI arguments specific to inference-time latent optimization (IT-opt).""" + parser.add_argument( + "--which-latent", + type=str, + default="pair", + choices=["single", "pair", "both"], + help="Which trunk latent(s) to optimize: single (s), pair (z), or both", + ) + parser.add_argument("--learning-rate", type=float, default=0.05, help="Adam learning rate") + parser.add_argument( + "--outer-steps", + type=int, + default=2, + help="Number of optimization rounds (fresh diffusion noise per round)", + ) + parser.add_argument( + "--anchor-weight", + type=float, + default=0.0, + help="On-manifold L2-to-baseline anchor weight for the optimized latent(s)", + ) + parser.add_argument( + "--max-grad-norm", type=float, default=1.0, help="Per-latent gradient-clip threshold" + ) + parser.add_argument( + "--bond-length-weight", + type=float, + default=5e-5, + help="Weight of the coordinate-space bond-geometry penalty (bond-length + steric-clash " + "hinges) added to the IT-opt loss. Default 5e-5 = smallest weight that fixes mean clash " + "while keeping density fit and diversity; use 1e-3 for median clash 0; 0 disables it", + ) + + _GUIDANCE_ARG_ADDERS: dict[str, Any] = { "pure_guidance": add_pure_guidance_args, "fk_steering": add_fk_steering_args, + "latent_opt": add_latent_opt_args, } diff --git a/src/sampleworks/utils/guidance_script_utils.py b/src/sampleworks/utils/guidance_script_utils.py index 9d1a5e61..f62f2351 100644 --- a/src/sampleworks/utils/guidance_script_utils.py +++ b/src/sampleworks/utils/guidance_script_utils.py @@ -24,6 +24,7 @@ ) from sampleworks.core.samplers.edm import AF3EDMSampler, EDMSamplerConfig from sampleworks.core.scalers.fk_steering import FKSteering +from sampleworks.core.scalers.latent_optimization import LatentOptimization # IT-opt wiring (added) from sampleworks.core.scalers.pure_guidance import PureGuidance from sampleworks.core.scalers.step_scalers import ( DataSpaceDPSScaler, @@ -85,7 +86,11 @@ def save_trajectory( save_every=10, ): """Dispatch trajectory serialization to the handler for the selected scaler.""" - if scaler_type == GuidanceType.PURE_GUIDANCE: + # IT-opt wiring (changed): this condition was `== GuidanceType.PURE_GUIDANCE`; we widened it to + # also accept LATENT_OPT. Latent optimization reuses the pure-guidance trajectory writer because + # its final sampling pass emits a trajectory with the same [ensemble, atoms, 3] layout that + # _save_trajectory already expects, so no separate writer is needed for it. + if scaler_type in (GuidanceType.PURE_GUIDANCE, GuidanceType.LATENT_OPT): _save_trajectory(trajectory, atom_array, output_dir, subdir_name, save_every) elif scaler_type == GuidanceType.FK_STEERING: _save_fk_steering_trajectory(trajectory, atom_array, output_dir, subdir_name, save_every) @@ -439,7 +444,7 @@ def run_guidance(args: GuidanceConfig, guidance_type: str, model_wrapper, device Result of the guidance run including status and timing. """ - log_path = getattr(args, "log_path", None) or os.path.join(args.output_dir, "run.log") + log_path = args.log_path or os.path.join(args.output_dir, "run.log") os.makedirs(os.path.dirname(log_path) or ".", exist_ok=True) # just in case log_path does not go to args.output_dir, make sure the latter exists @@ -500,16 +505,23 @@ def _run_guidance(args: GuidanceConfig, guidance_type: str, model_wrapper, devic if "Protenix" in wrapper_class_name: from sampleworks.models.protenix.wrapper import annotate_structure_for_protenix - structure = annotate_structure_for_protenix(structure, recycling_steps=recycling_steps) + structure = annotate_structure_for_protenix( + structure, + recycling_steps=recycling_steps, + # Disable diffusion shared-vars cache for LATENT_OPT so gradients can + # flow to z_trunk; cached tensors can otherwise become stale. + # Keep cache enabled for other guidance types. + enable_diffusion_shared_vars_cache=(guidance_type != GuidanceType.LATENT_OPT), + ) elif "RF3" in wrapper_class_name: from sampleworks.models.rf3.wrapper import annotate_structure_for_rf3 structure = annotate_structure_for_rf3( structure, recycling_steps=recycling_steps, - msa_path=getattr(args, "msa_path", None), - disable_chiral_features=getattr(args, "disable_chiral_features", False), - track_chiral_features=getattr(args, "track_chiral_features", False), + msa_path=args.msa_path, + disable_chiral_features=args.disable_chiral_features, + track_chiral_features=args.track_chiral_features, ) elif "Boltz" in wrapper_class_name: from sampleworks.models.boltz.wrapper import process_structure_for_boltz @@ -633,6 +645,76 @@ def _run_guidance(args: GuidanceConfig, guidance_type: str, model_wrapper, devic num_particles=args.num_particles, ) + refined_structure = result.structure + losses = result.losses if result.losses else [] + traj_denoised = result.metadata.get("trajectory_denoised", []) if result.metadata else [] + traj_next_step = list(result.trajectory) if result.trajectory else [] + + # ----- IT-opt wiring (added): the inference-time latent optimization branch (LATENT_OPT) ----- + # This branch belongs to _run_guidance() and is the entry point for latent optimization. Whereas + # pure guidance and FK steering steer the atomic coordinates, latent optimization instead + # optimizes the frozen model's cached trunk latents (the single representation s and/or the pair + # representation z) against the reward, then samples with those latents held fixed. We read the + # knobs off the config as the other branches do and hand the optimize-then-sample loop to + # LatentOptimization. + elif guidance_type == GuidanceType.LATENT_OPT: + logger.info("Initializing inference-time latent optimization (IT-opt)") + + # We import the representation-name maps inside this branch so the model-adapter dependency + # stays local to the only code that needs it. The attribute that stores each representation + # differs per model (for example "s_trunk"/"z_trunk" on Protenix and RF3 versus "s"/"z" on + # Boltz), so we look the names up by model instead of hard-coding them. + from sampleworks.models.latent_adapter import ( + DEFAULT_PAIR_REP_ATTR, + DEFAULT_SINGLE_REP_ATTR, + ) + + # LatentOptimization expects guidance start as a fraction of the schedule, but the config + # carries it as an integer step count, so we convert it here and default to optimizing from + # the first step. + guidance_t_start = args.guidance_start / num_steps if args.guidance_start > 0 else 0.0 + which_latent = args.which_latent # This is "single", "pair", or "both". + anchor_weight = args.anchor_weight + bond_length_weight = args.bond_length_weight + # GuidanceConfig exposes the model as `model_name` (not `args.model`); normalize to the + # lowercase key the DEFAULT_*_REP_ATTR maps use (as checkpoint resolution does below). + model_key = str(args.model_name).lower().replace("structurepredictor.", "") + try: + single_attr = DEFAULT_SINGLE_REP_ATTR[model_key] + pair_attr = DEFAULT_PAIR_REP_ATTR[model_key] + except KeyError as e: + raise ValueError( + "Latent optimization has no latent-attribute names registered for model " + f"{model_key!r}." + ) from e + + guidance = LatentOptimization( + ensemble_size=args.ensemble_size, + num_steps=num_steps, + guidance_t_start=guidance_t_start, + outer_steps=args.outer_steps, + learning_rate=args.learning_rate, + max_grad_norm=args.max_grad_norm, + optimize_single=which_latent in ("single", "both"), + optimize_pair=which_latent in ("pair", "both"), + single_attr=single_attr, + pair_attr=pair_attr, + anchor_weight_single=anchor_weight if which_latent in ("single", "both") else 0.0, + anchor_weight_pair=anchor_weight if which_latent in ("pair", "both") else 0.0, + bond_length_weight=bond_length_weight, + ) + + logger.info(f"Running latent optimization ({which_latent}) on model {model_key}") + # We still pass step_scaler so this call matches the signature the other guidance scalers + # use, but LatentOptimization ignores it -- v1 steers only through the latents. + result = guidance.sample( + structure=structure, + model=model_wrapper, + sampler=sampler, + step_scaler=step_scaler, + reward=reward_function, + ) + refined_structure = result.structure losses = result.losses if result.losses else [] traj_denoised = result.metadata.get("trajectory_denoised", []) if result.metadata else [] @@ -723,7 +805,7 @@ def get_job_result( runtime_seconds=round(end_time - start_time, 2), started_at=started_at.isoformat(), finished_at=ended_at.isoformat(), - log_path=getattr(args, "log_path", None) or os.path.join(args.output_dir, "run.log"), + log_path=args.log_path or os.path.join(args.output_dir, "run.log"), output_dir=args.output_dir, ) return result diff --git a/tests/models/test_latent_adapter.py b/tests/models/test_latent_adapter.py new file mode 100644 index 00000000..cb9a8d04 --- /dev/null +++ b/tests/models/test_latent_adapter.py @@ -0,0 +1,94 @@ +"""Tests for the latent read/write seam (``AttrLatentIO``) used by IT-opt. + +These run on CPU with no model checkpoints, using minimal mock conditioning +dataclasses whose representations are stored in named attributes. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +import torch +from sampleworks.models.latent_adapter import AttrLatentIO, LatentIO +from torch import Tensor + + +@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) + + +@dataclass(frozen=True) +class _PairConditioning: + """Conditioning carrying both single ``s`` and pair ``z`` representations.""" + + s: Tensor + z: Tensor + sidecar: str = "untouched" + + +@pytest.fixture +def s_tensor() -> Tensor: + torch.manual_seed(0) + return torch.randn(1, 8, 4) # [batch, tokens, d_s] + + +@pytest.fixture +def z_tensor() -> Tensor: + torch.manual_seed(1) + return torch.randn(1, 8, 8, 2) # [batch, tokens, tokens, d_z] + + +# --- 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_raises(s_tensor: Tensor): + """A name that matches no attribute is a configuration error, not an empty read.""" + io = AttrLatentIO(single_attr="does_not_exist") + with pytest.raises(AttributeError): + io.read_single(_SingleRepConditioning(s=s_tensor)) + + +# --- Pair (z) representation support ----------------------------------------- + + +def test_pair_io_none_by_default_is_noop(s_tensor: Tensor, z_tensor: Tensor): + """Single-arg AttrLatentIO leaves the pair rep unreachable (backward compatible).""" + io = AttrLatentIO("s") + cond = _PairConditioning(s=s_tensor, z=z_tensor) + assert io.read_pair(cond) is None + assert io.write_pair(cond, z_tensor * 9) is cond # unchanged + + +def test_pair_io_roundtrip(s_tensor: Tensor, z_tensor: Tensor): + io = AttrLatentIO(single_attr="s", pair_attr="z") + cond = _PairConditioning(s=s_tensor, z=z_tensor) + assert io.read_pair(cond) is z_tensor + new = io.write_pair(cond, z_tensor * 3) + torch.testing.assert_close(new.z, z_tensor * 3) + torch.testing.assert_close(new.s, s_tensor) # single untouched + assert new.sidecar == "untouched" diff --git a/tests/models/test_latent_optimization.py b/tests/models/test_latent_optimization.py new file mode 100644 index 00000000..3d5f258e --- /dev/null +++ b/tests/models/test_latent_optimization.py @@ -0,0 +1,139 @@ +"""Unit tests for the IT-opt scaler's core pieces. + +Covers the two genuinely-new, model-free parts of ``LatentOptimization``: + +- :class:`LatentAnchor` -- the on-manifold prior ``Σ w_i·mean((latent_i − baseline_i)²)``. +- :meth:`LatentOptimization._leaf_latents` -- turning the cached ``s``/``z`` on the + conditioning into fresh optimizable leaves (detach → clone → ``requires_grad``), + and raising loudly when the configured attribute names match no latent. + +CPU-only, no model checkpoints: a minimal frozen-dataclass conditioning stands in +for the real wrappers' ``@dataclass(frozen=True, slots=True)`` conditioning. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +import torch +from sampleworks.core.scalers.latent_optimization import LatentAnchor, LatentOptimization +from sampleworks.models.latent_adapter import AttrLatentIO +from sampleworks.models.protocol import GenerativeModelInput +from torch import Tensor + + +@dataclass(frozen=True) +class _Cond: + """Minimal conditioning carrying a single ``s`` and pair ``z`` (like the wrappers).""" + + s: Tensor + z: Tensor + + +@pytest.fixture +def features() -> GenerativeModelInput: + torch.manual_seed(0) + cond = _Cond(s=torch.randn(1, 6, 4), z=torch.randn(1, 6, 6, 2)) + return GenerativeModelInput(conditioning=cond) + + +def _make_opt(**overrides) -> LatentOptimization: + """Build a scaler for the leaf-logic tests. ``__init__`` only stores attrs + logs, + so no model/checkpoint is needed.""" + kwargs = dict( + optimize_single=True, + optimize_pair=True, + single_attr="s", + pair_attr="z", + anchor_weight_single=0.1, + anchor_weight_pair=0.2, + ) + kwargs.update(overrides) + return LatentOptimization(**kwargs) + + +# --- LatentAnchor ------------------------------------------------------------- + + +def test_anchor_is_zero_at_the_baseline(): + torch.manual_seed(0) + s, z = torch.randn(1, 6, 4), torch.randn(1, 6, 6, 2) + # latent == baseline -> no drift -> zero penalty, regardless of the weights + assert LatentAnchor([1.0, 2.0])([s, z], [s, z]).item() == pytest.approx(0.0) + + +def test_anchor_is_weighted_mean_squared_deviation(): + s0, z0 = torch.zeros(4), torch.zeros(4) + s, z = torch.full((4,), 2.0), torch.full((4,), 3.0) # mean sq dev = 4 and 9 + # 0.5 * 4 + 2.0 * 9 = 20 + assert LatentAnchor([0.5, 2.0])([s, z], [s0, z0]).item() == pytest.approx(20.0) + + +def test_anchor_gradient_reaches_the_latent(): + base = torch.zeros(4) + leaf = torch.full((4,), 2.0).requires_grad_(True) + LatentAnchor([1.0])([leaf], [base]).backward() + assert leaf.grad is not None + # d/dx of mean(x²) is 2x/n = 2*2/4 = 1 per element + torch.testing.assert_close(leaf.grad, torch.full((4,), 1.0)) + + +# --- LatentOptimization._leaf_latents ----------------------------------------- + + +def test_leaf_latents_makes_requires_grad_leaves(features): + opt = _make_opt() + io = AttrLatentIO("s", "z") + new_features, latents, baselines, weights = opt._leaf_latents(features, io) + + assert len(latents) == 2 + for leaf in latents: + assert leaf.requires_grad and leaf.is_leaf # a true leaf Adam can update directly + for leaf, base in zip(latents, baselines): + assert not base.requires_grad # the anchor target is detached + torch.testing.assert_close(leaf.detach(), base) # leaf starts exactly at the baseline + # the rewritten conditioning holds the SAME leaf objects, so the model reads + # the optimizable tensor rather than the original cached one + assert new_features.conditioning.s is latents[0] + assert new_features.conditioning.z is latents[1] + assert weights == [0.1, 0.2] + + +def test_leaf_latents_are_severed_from_the_original(features): + opt = _make_opt() + _, latents, baselines, _ = opt._leaf_latents(features, AttrLatentIO("s", "z")) + # a fresh clone/detach, not the original cached tensor (so optimizing it never + # writes back into, or backprops through, the frozen trunk output) + assert latents[0] is not features.conditioning.s + assert baselines[0] is not features.conditioning.s + + +def test_leaf_latents_single_only_respects_the_flag(features): + opt = _make_opt(optimize_pair=False) + _, latents, _, weights = opt._leaf_latents(features, AttrLatentIO("s", "z")) + assert len(latents) == 1 # pair skipped even though the io could address it + assert weights == [0.1] + + +def test_leaf_latents_raises_when_no_attribute_matches(features): + # wrong attribute names -> nothing to optimize -> a clear error, not a silent no-op + opt = _make_opt() + io = AttrLatentIO("does_not_exist", "also_missing") + with pytest.raises(ValueError, match="does not expose"): + opt._leaf_latents(features, io) + + +def test_leaf_latents_raises_when_only_one_attribute_matches(features): + # The dangerous case: one name is right, so the run would otherwise optimize half of what + # was asked for and still look successful. + opt = _make_opt() + io = AttrLatentIO("does_not_exist", "z") + with pytest.raises(ValueError, match="does not expose"): + opt._leaf_latents(features, io) + + +def test_leaf_latents_raises_when_nothing_is_enabled(features): + opt = _make_opt(optimize_single=False, optimize_pair=False) + with pytest.raises(ValueError, match="no latent enabled"): + opt._leaf_latents(features, AttrLatentIO("s", "z")) diff --git a/tests/utils/test_guidance_script_arguments.py b/tests/utils/test_guidance_script_arguments.py index 0b2f14e3..f2d8848a 100644 --- a/tests/utils/test_guidance_script_arguments.py +++ b/tests/utils/test_guidance_script_arguments.py @@ -376,3 +376,54 @@ def test_job_result_migrates_legacy_model_pickle() -> None: assert restored.model_name == "boltz2" assert "model" not in restored.__dict__ assert "model" not in restored.as_dict() + + +# ============================================================================ +# latent-opt CLI flag plumbing (regression: flags must reach the config) +# ============================================================================ + + +def test_latent_opt_cli_flags_reach_the_config(): + """Every --guidance-type latent_opt flag must land on the GuidanceConfig. + + Regression guard: these attribute names have to be listed in + ``_DYNAMIC_ATTRS`` or ``from_cli`` silently drops the parsed values and + ``_run_guidance``'s ``getattr(args, ...)`` falls back to the defaults, making + every flag a no-op. Passing *non-default* values here catches that. + """ + config = GuidanceConfig.from_cli( + argv=[ + "--protein", + "1vme", + "--structure", + "s.cif", + "--density", + "d.ccp4", + "--resolution", + "1.8", + "--output-dir", + "/tmp/out", + # IT-opt flags, all set to values that differ from their defaults + "--which-latent", + "single", # default "pair" + "--learning-rate", + "0.123", # default 0.05 + "--outer-steps", + "5", # default 2 + "--anchor-weight", + "0.7", # default 0.0 + "--max-grad-norm", + "2.0", # default 1.0 + "--bond-length-weight", + "0.001", # default 5e-5 + ], + model_name="protenix", + guidance_type="latent_opt", + ) + + assert config.which_latent == "single" + assert config.learning_rate == pytest.approx(0.123) + assert config.outer_steps == 5 + assert config.anchor_weight == pytest.approx(0.7) + assert config.max_grad_norm == pytest.approx(2.0) + assert config.bond_length_weight == pytest.approx(0.001)