feat(it-opt): give each ensemble member its own trunk latent by smallfishabc · Pull Request #366 · diff-use/sampleworks

smallfishabc · 2026-08-10T23:00:10Z

Why

IT-opt optimizes a single s/z latent shared by every ensemble member, so
members can only diverge through their prior noise. Under the density reward the
ensemble collapses toward one solution instead of spreading.

What changed

Each member now gets its own independent latent. _leaf_latents stacks
ensemble_size copies of the trunk baseline onto a leading batch dim, one
optimizable leaf per member. Since stock diffusion modules take an un-batched
conditioning, _PerMemberStepper slices that dim back apart and runs the
model's normal step once per member — same FLOPs as one batched call.

Stacked on fy/it-optimization (#313), which is merged in, so the diff here is
only the incremental work.

…ts from a brand new start and the loss function is penalized by the average of the ensemble.

The file was a zero-edit copy of the old eval/ path version. #329 moved it to tests/synthetic/, and its import of sampleworks.eval.generate_synthetic_sf no longer resolves, so collecting tests/eval/ failed at import.

…est-ensemble-sampling

Force-added, because .gitignore excludes it_opt_scratch/ so that the large run output trees (targets_out_*, patch_tree*, figures/, discard/) stay untracked. Only the runnable reproducibility set is committed here. This is the test + scoring workflow that IT_OPT_ARCHITECTURE_OVERVIEW.txt points at: the batch ensemble runner, the paper-metric RSCC/RMSD scorers and their simplified single-prediction counterparts, the SLURM driver, and the per-protein window selections every scorer reads. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Reconciles the branch without a force-push. The remote's two commits (4c441b1, 4ad6694) are patch-identical duplicates of e43f909 and 996e090: the branch was rebased onto d31f995 on 2026-07-29 and never force-pushed, so git reported a divergence that was only SHA-deep, never a content one. The merge's sole incoming change was tests/eval/test_generate_synthetic_sf.py, the stale pre-#329 duplicate that a585ee7 removed on purpose. That deletion is kept here: the file imports sampleworks.eval.generate_synthetic_sf and sampleworks.eval.synthetic_utils, but both modules now live under sampleworks.synthetic.*, so restoring it would break test collection outright. The canonical 350-line copy remains at tests/synthetic/test_generate_synthetic_sf.py. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Two capabilities, both needed to drive IT-opt from run_grid_search.py. latent_opt was already accepted by --scalers, but locked at the add_latent_opt_args defaults with no way to change them. The grid builds its GuidanceConfig directly and never calls from_cli(), so the _DYNAMIC_ATTRS copy loop that normally moves parsed CLI values onto a config is unreachable, and __post_init__ has already seeded every guidance-specific attribute with its argparse default. A new flag therefore looked accepted and did nothing. The LATENT_OPT branch in populate_config_for_guidance_type is what makes the flags take effect; the six tunables come from calling the shipped add_latent_opt_args adder rather than restating them, so names and defaults cannot drift from the sampleworks-guidance CLI. That branch deliberately does not set step_size or step_scaler_type. The adder does not define them, so leaving them unset lets _run_guidance fall through to the same values the scratch runner hit (step_size 0.01, noisespace). Setting them would quietly change what the recorded runs did. --jobs-per-gpu lets two jobs share a card, since a single job leaves the GPU idle while it featurizes and writes output. Two details make it correct: the worker index was doubling as the CUDA ordinal, so the device now wraps with i % len(gpus) -- without it, 4 GPUs x 2 asks for cuda:4..cuda:7 and dies at once; and the worker count is clamped to a whole multiple of the GPU count so every card carries the same load, rather than 5 workers on 4 GPUs leaving one card with two jobs and three with one. Capped at 2, which is the tested ceiling. The presets are one configuration differing only in GPU count, with settings taken from it_opt_scratch/slurm_ensemble_and_score.sh so runs are comparable to the recorded trees. ensemble-sizes and gradient-weights must stay single values: their defaults ("1 2 4 8", "0.01 0.1 0.2") would expand one job per protein into twelve. Verified on a 4x H100 pod at real settings (ensemble 8, 200 steps, outer_steps 2): 8 proteins over 8 workers, 2 per GPU across cuda:0-3, 8/8 SUCCESS in ~36 min. Existing behaviour is unchanged -- all 12 experiments/ presets still parse, _DYNAMIC_ATTRS is byte-identical, and the pure_guidance/fk_steering branches are untouched. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…RMSD scorers score_paper_rscc.py: add --maps-dir/--map-template to score non-uniform (0.5occ) maps; clean mixed ATOM/HETATM modified-residue altlocs from the reference (resolve_mixed_hetatm_atom_altlocs); fall back to filter_to_common_atoms(normalize_ids=True) in align_prediction_to_reference when strict matching finds no common atoms (chain P->A relabel). score_paper_rmsd.py: same reference cleaning; reference_conformers/prediction_lookup fall back to make_normalized_atom_id keys when the deposited keys are disjoint from the prediction, recovering chain-relabelled proteins. Verified on all 40: 33 working proteins byte-identical on both metrics; recovers 6NI5/6NI6 and the contiguous chain-relabel proteins (7Z0E partial due to sequence gaps). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…onformers Adds remap_altlocs_to_ab(cif_path). Some depositions label a residue's two conformers A/C or A/D rather than A/B, and tooling that assumes the two altlocs are literally named A and B then silently discards the second conformer (7Z0E, 7AVG). Per (chain, residue) this maps the alternate label onto the free B slot: an atom labelled A keeps A and the other non-blank label becomes B; if no A is present the two labels are assigned A/B in sorted order. Only positions carrying exactly two non-blank altlocs that are not already {A, B} are touched. Positions that are already A/B, hold a single or blank altloc, or carry three or more altlocs (ambiguous which two to keep) are left alone, so no metric moves for structures that were already A/B. A warning is logged per remapped position. Returns a temp CIF only when something changed, otherwise the original path unchanged. Known follow-ups: this shares its load / iterate-positions / write-temp-CIF scaffolding with resolve_mixed_hetatm_atom_altlocs in this module, and extracting that common core is deferred. No unit test yet.

align_prediction_to_reference paired atoms on (chain, res_id, atom_name) via filter_to_common_atoms, which never checks that the paired atoms are the same residue. Predictions are renumbered from 1 on chain 'A' while references keep the deposited numbering, so whenever the deposited range merely overlaps 1..N (3AZY is 8-258, 2YNT is 36-295) strict matching succeeded on a large but sequence-shifted set, pairing prediction residue i with reference residue i-offset. Kabsch then fit that shifted correspondence and reported a fold-scale error for structures that were actually correct. Adds two functions to score_paper_rscc: strict_frame_is_consistent(ref, pred) -- whether deposited identifiers may be used directly. Comparing residue names separates the two cases cleanly: a correct frame agrees on every pair, a shifted one only at the rate expected from residue-type frequency by chance. Accepts at >= MIN_RESIDUE_IDENTITY (0.95), else the caller falls back to normalized sequential ids. residue_identity(ref_common, pred_common) -- fraction of matched pairs whose res_name agrees, or None when the two matched sets differ in length. That happens because filter_to_common_atoms masks with np.isin, so a reference carrying two altlocs of one atom contributes both copies while a single-conformer prediction contributes one (6NI5, 6NI6). None keeps that distinct from "paired, but the residues disagree". align_prediction_to_reference now chooses normalization from that check rather than from whether strict matching raised, and raises when neither scheme pairs matching residues instead of silently fitting an unrelated construct. score_paper_rmsd adopts the same check in place of its raw_ref_keys.isdisjoint heuristic, so both scorers decide framing identically, and applies remap_altlocs_to_ab to the reference so the A/B conformer split below it keeps non-A/B alternates.

… runner The coord_guidance arm previously passed only step_scaler_type and step_size, which left it running a materially different recipe from the paper's: the paper normalizes the density gradient to the EDM denoising-update magnitude (gradient_normalization) and applies AF3 augmentation with a realign each step, and only under that normalization does step_size act as a fraction of the denoising step rather than a raw gradient multiplier. guidance_for_mode now forwards gradient_normalization and augmentation for coord_guidance, and build_config passes guidance_start through. The forwarding is mode-scoped, so the latent-opt arms are unaffected. Adds three flags, all defaulting to previous behavior so recorded runs are unchanged: --guidance-start (step at which guidance begins, -1 meaning from step 0; both coord DPS and IT-opt read it as a fraction of num_steps), --gradient-normalization, and --augmentation.

…est-ensemble-sampling # Conflicts: # src/sampleworks/core/scalers/latent_optimization.py

coderabbitai · 2026-08-10T23:00:21Z

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 08ec00fd-d519-4fff-ad91-183ee89ddb6b

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

Copilot

Pull request overview

This PR improves IT-opt ensemble diversity by giving each ensemble member its own optimizable post-trunk latent (s/z) instead of sharing one latent across all members, addressing ensemble collapse under density rewards. It also expands the surrounding IT-opt tooling (grid search worker scheduling, scoring scripts, and CIF pre-processing) to support multi-GPU runs and paper-metric evaluation.

Changes:

Update latent optimization to create per-member latent leaves and denoise via a per-member stepping adapter.
Make IT-opt CLI/grid argument propagation more robust and ensure Protenix intermediate dumps are rooted under the run output directory.
Add scratch workflow + scoring utilities (RSCC + min-altloc-RMSD) and CIF preprocessing to support paper-style evaluation.

Reviewed changes

Copilot reviewed 18 out of 18 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/rewards/test_geometry.py	Adds a small unit test for `_covalent_radius()` used by geometry regularization.
tests/models/test_latent_optimization.py	Updates tests to reflect per-member latent stacking semantics.
src/sampleworks/utils/guidance_script_utils.py	Roots Protenix per-sample JSON dumps under the job output directory.
src/sampleworks/utils/guidance_script_arguments.py	Deduplicates IT-opt tunable attribute propagation between CLI and grid-search config population.
src/sampleworks/utils/cif_utils.py	Adds `remap_altlocs_to_ab()` to normalize non-A/B altloc labels for A/B-only tooling.
src/sampleworks/core/scalers/latent_optimization.py	Implements per-member latent leaves and `_PerMemberStepper` to run unbatched model steps per ensemble member.
run_grid_search.py	Adds `--jobs-per-gpu` and corrects worker→GPU assignment when running multiple workers per GPU; reuses latent-opt arg adder for consistency.
it_opt_scratch/slurm_ensemble_and_score.sh	Adds a SLURM array workflow to generate ensembles and run scoring/aggregation.
it_opt_scratch/score_rscc_simplified.py	Self-contained RSCC scorer against a density map for one prediction.
it_opt_scratch/score_rmsd_simplified.py	Self-contained min-altloc-RMSD scorer for one prediction vs two-altloc reference.
it_opt_scratch/score_paper_simplified.py	Driver that joins both paper metrics into one CSV for a run tree.
it_opt_scratch/score_paper_rscc.py	Paper-style RSCC sweep driver over proteins/arms, matching published pipeline primitives.
it_opt_scratch/score_paper_rmsd.py	Paper-style min-altloc-RMSD sweep driver; uses CIF preprocessing to ensure A/B altloc semantics.
it_opt_scratch/run_targets.py	Batch runner to generate ensembles over target lists with optional multi-GPU fan-out.
it_opt_scratch/run_targets_simplified.py	Single-target runner wrapper around `run_guidance()` for the scratch workflow.
it_opt_scratch/paper_maxrmsd_selections.csv	Adds paper window selections used by scoring scripts.
experiments/it_opt_4gpu.toml	Preset for running IT-opt on 4 GPUs with `jobs-per-gpu` enabled.
experiments/it_opt_1gpu.toml	Preset for running IT-opt sequentially on 1 GPU.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+    def step(self, x_t: Tensor, t, *, features: GenerativeModelInput) -> Tensor:
+        """Loop the wrapped model's ``step`` over ensemble members; stack the per-member results."""
+        cond = features.conditioning
+        per_member: list[Tensor] = []
+        for i in range(self._ensemble_size):
+            # Slice only the OPTIMIZED latents (they carry the ensemble batch dim); a non-optimized
+            # latent stays the shared un-batched baseline already on ``cond``.
+            cond_i = cond
+            if self._optimize_single:
+                cond_i = self._io.write_single(cond_i, self._io.read_single(cond)[i])
+            if self._optimize_pair:
+                cond_i = self._io.write_pair(cond_i, self._io.read_pair(cond)[i])
+            t_i = t
+            if isinstance(t, Tensor) and t.ndim >= 1 and t.shape[0] == x_t.shape[0]:
+                t_i = t[i : i + 1]
+            features_i = GenerativeModelInput(conditioning=cond_i)
+            per_member.append(self._model.step(x_t[i : i + 1], t_i, features=features_i))
+        return torch.cat(per_member, dim=0)


The repo applies ruff and ty uniformly to every tracked Python file -- the ruff hooks carry no files/exclude filter and [tool.ty] has no exclude beyond a single eval script -- so tracking these experiment scripts subjected ~2,100 lines of run-once tooling to library-grade gates. That accounted for 20 of the 25 ruff errors on this branch and 28 of the 77 ty diagnostics, including the only hard error (list[LiteralString] vs list[str] in run_targets.py). The scripts are unchanged and stay on disk; only the index entries are dropped. Nothing under src/ imports them, so this removes no functionality -- the dependency runs scratch -> src, never the reverse. Committing the index change only, not the matching .gitignore rule: on the remote these paths simply no longer exist, so the ignore rule matters only to local working copies.

Measured against the base (fy/it-optimization) file by file, so only findings this branch introduces are addressed here; the rest are inherited and belong on #313. ruff (5 -> 0): - import sort in run_grid_search.py and tests/rewards/test_geometry.py, and an unused pytest import in the latter (ruff --fix) - rewrapped test_geometry.py's docstring, which ran to 111 columns and carried trailing whitespace, and added the missing final newline ty (1 -> 0 introduced): _PerMemberStepper subscripted read_pair() directly, and read_pair returns None when the io addresses no pair representation -- flagged as not-subscriptable, and the only ty finding this branch adds. sample() sets optimize_pair only together with a pair_attr, so the invariant holds through the public path, but it is established 140 lines away and nothing enforces it: an AttrLatentIO built directly with optimize_pair=True and pair_attr=None gives a None subscript deep in the denoise loop. Now raises where the assumption is made. Verified on the pod: ruff check and format clean across all 7 changed Python files, 40 tests pass, and the remaining ty counts (4/9/13 in latent_optimization, guidance_script_utils, test_latent_optimization) match the base exactly.

test_list_presets_returns_bundled_experiments and test_list_prints_all_experiment_presets both assert the exact set of bundled presets, so adding experiments/it_opt_1gpu.toml and it_opt_4gpu.toml necessarily failed them until the names were registered. Adds the two names in sorted position; no other behaviour changes. Verified on the pod: 45 tests pass across both files once the untracked local experiments/it_opt_8gpu.toml is set aside -- that file is not in git, so CI never sees it, but it will fail these two tests in any working tree that has it.

…ment set" This reverts commit e2393f4.

…est-ensemble-sampling

smallfishabc and others added 13 commits July 29, 2026 14:35

Change in md to remove paper related conclusions.

4c441b1

Change the optimization to N forward method so each optimization star…

4ad6694

…ts from a brand new start and the loss function is penalized by the average of the ensemble.

Change the optimization to N forward method so each optimization star…

996e090

…ts from a brand new start and the loss function is penalized by the average of the ensemble.

chore(tests): drop stale pre-#329 copy of the synthetic-sf test

a585ee7

The file was a zero-edit copy of the old eval/ path version. #329 moved it to tests/synthetic/, and its import of sampleworks.eval.generate_synthetic_sf no longer resolves, so collecting tests/eval/ failed at import.

Merge remote-tracking branch 'origin/fy/it-optimization' into fy/it-t…

d869af0

…est-ensemble-sampling

Merge remote-tracking branch 'origin/fy/it-optimization' into fy/it-t…

992fadf

…est-ensemble-sampling # Conflicts: # src/sampleworks/core/scalers/latent_optimization.py

Copilot AI lite review requested due to automatic review settings August 10, 2026 23:00

smallfishabc requested review from a team, Abdelsalam-Abbas, DorisMai, denis-zaitsev, k-chrispens, mag-astera and marcuscollins as code owners August 10, 2026 23:00

smallfishabc had a problem deploying to gpu-testing August 10, 2026 23:00 — with GitHub Actions Error

Copilot started reviewing on behalf of smallfishabc August 10, 2026 23:00 View session

Copilot AI reviewed Aug 10, 2026

View reviewed changes

smallfishabc mentioned this pull request Aug 11, 2026

feat(it-opt): per step latent space optimization #313

Open

smallfishabc had a problem deploying to gpu-testing August 11, 2026 00:46 — with GitHub Actions Error

smallfishabc had a problem deploying to gpu-testing August 13, 2026 20:18 — with GitHub Actions Error

Revert "test(runs): register the IT-opt presets in the bundled experi…

5c0b817

…ment set" This reverts commit e2393f4.

smallfishabc had a problem deploying to gpu-testing August 13, 2026 20:42 — with GitHub Actions Error

Merge remote-tracking branch 'origin/fy/it-optimization' into fy/it-t…

143db94

…est-ensemble-sampling

smallfishabc had a problem deploying to gpu-testing August 13, 2026 21:01 — with GitHub Actions Error

test(runs): register the IT-opt presets in the bundled set

c316ae8

smallfishabc requested a deployment to gpu-testing August 13, 2026 21:24 — with GitHub Actions Waiting

Navigation Menu

Uh oh!

feat(it-opt): give each ensemble member its own trunk latent - #366

feat(it-opt): give each ensemble member its own trunk latent#366
smallfishabc wants to merge 19 commits into
fy/it-optimizationfrom
fy/it-test-ensemble-sampling

smallfishabc commented Aug 10, 2026

Uh oh!

coderabbitai Bot commented Aug 10, 2026 •

edited

Loading

Review skipped

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

smallfishabc commented Aug 10, 2026

Why

What changed

Uh oh!

coderabbitai Bot commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

coderabbitai Bot commented Aug 10, 2026 •

edited

Loading