From 4c441b15b1c0d77c4c446403d4a930a20b8624eb Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Wed, 29 Jul 2026 14:35:05 -0700 Subject: [PATCH 01/15] Change in md to remove paper related conclusions. --- docs/IT_OPT_DESIGN.md | 2 +- docs/IT_OPT_TESTING.md | 66 +++++++----------------------------------- 2 files changed, 12 insertions(+), 56 deletions(-) diff --git a/docs/IT_OPT_DESIGN.md b/docs/IT_OPT_DESIGN.md index 5233f370..7e9a5a08 100644 --- a/docs/IT_OPT_DESIGN.md +++ b/docs/IT_OPT_DESIGN.md @@ -5,7 +5,7 @@ understand the feature: **what it is, the as-built algorithm, the components and 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 empirical results. +- [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). diff --git a/docs/IT_OPT_TESTING.md b/docs/IT_OPT_TESTING.md index 4f806893..f07f84fe 100644 --- a/docs/IT_OPT_TESTING.md +++ b/docs/IT_OPT_TESTING.md @@ -1,8 +1,8 @@ # IT-Opt — Testing, Verification, and Wiring How to run [`LatentOptimization`](../src/sampleworks/core/scalers/latent_optimization.py), debug it -when it misbehaves, what it produces on real targets, and where it is wired into the pipeline. For -the architecture, read [IT_OPT_DESIGN.md](IT_OPT_DESIGN.md) first. +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. @@ -149,58 +149,14 @@ from the grid-search / save machinery. - [core/scalers/latent_optimization.py](../src/sampleworks/core/scalers/latent_optimization.py) — per-latent grad clips and the `latent_drift` diagnostic. -## 5. Verification results - -Verified against the sampleworks-release paper's altloc metrics on the already-generated 40-protein -ensembles (native occupancy — no regeneration). Conditions: `baseline` (unguided), `coord_guidance` -(the paper's shipped guided method), and `z` (IT-opt) at bond-geometry weights 0 / 5e-5 / 1e-4 / 1e-3. - -| dimension (metric) | baseline | coord_guid | z (w=0) | z 5e-5 | -|---|---|---|---|---| -| density fit (RSCC ≥ 0.8) | 42% | inert | 92% | 92% | -| accuracy (nearer-altloc RMSD, med.) | 2.14 | inert | 1.02 | ~1.0 | -| reach both (RMSD max(A,B) med.; ≤2Å) | 2.18; 46% | inert | 1.3; 54% | ~1.3; 54% | -| diversity (ensembles that split A/B) | 7% | 8% | 18% | 22% | -| clean bimodal (clustering ≥ 0.5) | 10.6% | 10.6% | 8% | 9% | -| clashes (mean) | 0.38 | — | 0.47 | 0.38 | - -Bond-geometry weight sweep (mean / median clash / RSCC ≥ 0.8; unguided baseline = 0.38 / 0.00 / 42%): - -| weight | mean clash | median clash | RSCC ≥ 0.8 | diversity | -|---|---|---|---|---| -| 0 | 0.47 | 0.25 | 92% | 18% | -| **5e-5** | **0.38** | **0.25** | **92%** | **22%** (default) | -| 1e-4 | 0.39 | 0.25 | 92% | 17% | -| 1e-3 | 0.35 | 0.00 | 98% | 11% | - -**Interpretation.** IT-opt-`z` strongly improves density fit and accuracy (RSCC ≥ 0.8 goes 42% → 92%, -mirroring the paper's *guided* 45.4% → 96.0%; nearer-altloc RMSD halves). It modestly improves -diversity (~3× more ensembles reach both altlocs, 7% → 22%) but does **not** achieve clean bimodal -capture — the clustering silhouette stays flat and ~80% of ensembles still collapse to one -conformation, consistent with the paper's thesis. `coord_guidance` is inert on every metric here. The -default bond weight `5e-5` is the smallest that restores mean clash to baseline while keeping the full -density gain and the diversity; `1e-3` over-constrains and erodes diversity. Default -`bond_length_weight` is `5e-5` (`LatentOptimization.__init__` + `--bond-length-weight`); `0` disables -the penalty. - -## 6. Open problems - -1. **`coord_guidance` is inert vs. the paper.** The paper reports coordinate guidance taking RSCC ≥ - 0.8 from 45.4% to 96.0%; here it is indistinguishable from baseline on every metric, and it is the - *latent* optimization that reproduces the density jump. Likely a step-size/config difference (paper - optimal 0.1 for Protenix). Resolve with a `coord_guidance` guidance-strength sweep. -2. **RSCC uses a local scorer, not the repo-exact pipeline.** The numbers come from a homemade scorer, - not `scripts/eval/rscc_grid_search_script.py` (`process_group` → density transformer → Kabsch → - `extract_tight` at 2.0 Å → `rscc`). The homemade baseline (42%) matches the paper's (45.4%), which - calibrates it, but a repo-exact run is the final confirmation. Needs a depth-4 trial-dir tree +## 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. -3. **Absolute fractions sit above the paper's.** Baseline is 42% (RSCC ≥ 0.8) and 10.6% (clustering ≥ - 0.5) vs the paper's 45.4% and 1.6% — RSCC matches, clustering is ~7× higher. Likely row population: - we score native occupancy only (~85 rows), the paper aggregates the full 791-segment sweep. - Within-our-runs comparisons are sound; cross-to-paper *absolute* fractions are not until the - population is matched. -4. **Five proteins need patching.** 6RP1, 7Z0E, 4OLE, 8Z76, 2I6H raise "No common atoms found" (chain/ - residue-naming mismatch), so they drop from every aggregate. Resolve with - `scripts/patch_output_cif_files.py` (needs network for `rcsb.fetch`; the `~/.sampleworks/rcsb` - cache is empty) or sequence-based atom matching. +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. From 4ad6694b9e9f9226952ea4974733d79ab38fff2e Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Wed, 29 Jul 2026 14:38:22 -0700 Subject: [PATCH 02/15] Change the optimization to N forward method so each optimization starts from a brand new start and the loss function is penalized by the average of the ensemble. --- .../core/scalers/latent_optimization.py | 87 ++++++++- .../utils/guidance_script_utils.py | 4 + tests/eval/test_generate_synthetic_sf.py | 166 ++++++++++++++++++ tests/rewards/test_geometry.py | 17 ++ 4 files changed, 269 insertions(+), 5 deletions(-) create mode 100644 tests/eval/test_generate_synthetic_sf.py create mode 100644 tests/rewards/test_geometry.py diff --git a/src/sampleworks/core/scalers/latent_optimization.py b/src/sampleworks/core/scalers/latent_optimization.py index 56a9e4bd..d38e64fe 100644 --- a/src/sampleworks/core/scalers/latent_optimization.py +++ b/src/sampleworks/core/scalers/latent_optimization.py @@ -120,6 +120,62 @@ def __call__(self, latents: Sequence[Tensor], baselines: Sequence[Tensor]) -> Te return torch.stack(terms).sum() +class _PerMemberStepper: + """Denoise one ensemble member at a time, each with its OWN trunk latent. + + The it_opt "multiple-leaf" scheme gives every ensemble member its own latent, so ``s``/``z`` + carry a leading ``ensemble_size`` batch dim. Stock model diffusion modules instead take an + un-batched conditioning and broadcast it internally, so a batched latent either crashes + (Protenix) or mis-broadcasts (Boltz's ``multiplicity``). This adapter sidesteps that in a + model-agnostic way: for each member it slices that member's latent, runs the wrapped model's + normal un-batched ``step`` on that member alone, and stacks the results. Members stay + independent, so gradients stay per-member; the summed density reward couples them only through + the ensemble average. Cost is N forwards instead of one batched call, with an equivalent result. + + ``featurize`` / ``initialize_from_prior`` pass straight through to the wrapped model. + """ + + def __init__( + self, + model, + io: AttrLatentIO, + *, + optimize_single: bool, + optimize_pair: bool, + ensemble_size: int, + ): + self._model = model + self._io = io + self._optimize_single = optimize_single + self._optimize_pair = optimize_pair + self._ensemble_size = ensemble_size + + 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) + + def featurize(self, *args, **kwargs): + return self._model.featurize(*args, **kwargs) + + def initialize_from_prior(self, *args, **kwargs): + return self._model.initialize_from_prior(*args, **kwargs) + + class LatentOptimization: """Trajectory scaler that optimizes the model's ``s``/``z`` latents (IT-opt). @@ -253,6 +309,16 @@ def sample( schedule = sampler.compute_schedule(num_steps=self.num_steps) grad_enabler = _GradEnablingScaler() + # Denoise per member so each uses its own latent: stock model diffusion modules take an + # un-batched conditioning, so the batched per-member latents can't go through in one call. + stepper = _PerMemberStepper( + model, + io, + optimize_single=self.optimize_single, + optimize_pair=self.optimize_pair, + ensemble_size=self.ensemble_size, + ) + # --- 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 valid geometry for density @@ -275,7 +341,7 @@ def sample( for outer in 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, + model=stepper, sampler=sampler, reward=reward, features=features, @@ -303,7 +369,7 @@ def sample( # --- final clean sampling pass with the optimized latents --------------- final_coords, trajectory, losses = self._sample_with_frozen_latents( - model=model, + model=stepper, sampler=sampler, reward=reward, io=io, @@ -337,8 +403,12 @@ def _leaf_latents(self, features: GenerativeModelInput, io: AttrLatentIO): 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. + updated in place across rounds and steps). + + Each leaf gets a leading ``ensemble_size`` batch dimension -- one INDEPENDENT latent per + ensemble member, all cloned from the same trunk baseline (the it_opt "multiple-leaf" + scheme), so members can diverge rather than share one latent. The baseline kept for the + anchor stays un-batched and broadcasts across members. """ conditioning = features.conditioning latents: list[Tensor] = [] @@ -358,7 +428,14 @@ def _leaf_latents(self, features: GenerativeModelInput, io: AttrLatentIO): if baseline is None: continue baseline = baseline.detach() - leaf = baseline.clone().requires_grad_(True) + # it_opt "multiple-leaf" scheme: give each ensemble member its OWN latent. We stack + # ensemble_size independent copies of the trunk baseline into a leading batch dim, so + # each member gets its own gradient and can diverge, instead of collapsing onto one + # shared latent. (The reference does the same via batch-expand-then-clone.) Requires the + # diffusion module to accept a per-member (batched) conditioning -- verify with a + # batched gradcheck before relying on it. + member_copies = [baseline for _ in range(self.ensemble_size)] + leaf = torch.stack(member_copies).requires_grad_(True) conditioning = write(conditioning, leaf) latents.append(leaf) baselines.append(baseline) diff --git a/src/sampleworks/utils/guidance_script_utils.py b/src/sampleworks/utils/guidance_script_utils.py index 2e9313e0..de8717e9 100644 --- a/src/sampleworks/utils/guidance_script_utils.py +++ b/src/sampleworks/utils/guidance_script_utils.py @@ -507,6 +507,10 @@ def _run_guidance(args: GuidanceConfig, guidance_type: str, model_wrapper, devic structure = annotate_structure_for_protenix( structure, + # Root Protenix's per-sample input dump (protenix_input*.json) under the job's output + # dir; otherwise out_dir falls back to the input id, resolves against the CWD, and + # leaves a stray / folder there on every run. + out_dir=str(Path(args.output_dir) / "protenix_input"), recycling_steps=recycling_steps, # ensemble_size removed: #330 dropped it from the wrapper annotate signature; the # ensemble is now sized by the scaler via initialize_from_prior(). diff --git a/tests/eval/test_generate_synthetic_sf.py b/tests/eval/test_generate_synthetic_sf.py new file mode 100644 index 00000000..39534b37 --- /dev/null +++ b/tests/eval/test_generate_synthetic_sf.py @@ -0,0 +1,166 @@ +"""Tests for atomarray_to_gemmi in generate_synthetic_sf, using real 6b8x structure.""" + +import logging +from pathlib import Path + +import gemmi +import numpy as np +import pytest +import torch +from atomworks.io.transforms.atom_array import remove_waters +from biotite.structure import AtomArray +from sampleworks.eval.generate_synthetic_sf import atomarray_to_gemmi +from sampleworks.eval.synthetic_utils import assign_occupancies +from sampleworks.utils.atom_array_utils import ( + detect_altlocs, + keep_amino_acids, + keep_polymer, + load_structure_with_altlocs, + remove_hydrogens, +) +from SFC_Torch import SFcalculator +from SFC_Torch.io import PDBParser +from SFC_Torch.utils import assert_numpy + + +DMIN = 2.0 + + +@pytest.fixture(scope="module") +def stripped_gemmi(resources_dir: Path) -> gemmi.Structure: + """gemmi.Structure with hydrogens, ligands, and waters removed using gemmi methods.""" + s = gemmi.read_structure(str(resources_dir / "6b8x" / "6b8x_final.pdb")) + s.remove_hydrogens() + s.remove_ligands_and_waters() + return s + + +@pytest.fixture(scope="module") +def stripped_atom_array(resources_dir: Path) -> AtomArray: + """AtomArray with hydrogens, waters, and non-polymer/non-amino-acid atoms + removed using existing utils.""" + arr = load_structure_with_altlocs(resources_dir / "6b8x" / "6b8x_final.pdb") + arr = remove_hydrogens(arr) + arr = remove_waters(arr) + arr = keep_polymer(keep_amino_acids(arr)) + assert isinstance(arr, AtomArray) + return arr + + +@pytest.fixture(scope="module") +def gemmi_structure_from_atomarray( + stripped_atom_array, stripped_gemmi: gemmi.Structure +) -> gemmi.Structure: + """gemmi.Structure converted from the stripped AtomArray.""" + return atomarray_to_gemmi( + stripped_atom_array, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm + ) + + +def _compute_fprotein(gemmi_structure: gemmi.Structure, device: torch.device) -> np.ndarray: + """Compute |Fprotein| amplitudes from a gemmi structure via SFcalculator at ``DMIN`` + resolution. The final assert_numpy converts any tensor or list to a numpy array.""" + sfc = SFcalculator( + PDBParser(gemmi_structure), + mtzdata=None, + dmin=DMIN, + mode="xray", + anomalous=False, + set_experiment=False, + device=device, + ) + sfc.calc_fprotein() + return assert_numpy(sfc.Fprotein_asu) + + +class TestAtomArrayToGemmi: + """Tests for atomarray_to_gemmi using the 6b8x structure.""" + + def test_cell_matches_pdb(self, gemmi_structure_from_atomarray, stripped_gemmi): + """Unit cell parameters are preserved through the biotite→gemmi conversion.""" + result = gemmi_structure_from_atomarray.cell + expected = stripped_gemmi.cell + assert result.a == pytest.approx(expected.a) + assert result.b == pytest.approx(expected.b) + assert result.c == pytest.approx(expected.c) + assert result.alpha == pytest.approx(expected.alpha) + assert result.beta == pytest.approx(expected.beta) + assert result.gamma == pytest.approx(expected.gamma) + + def test_space_group_matches_pdb(self, gemmi_structure_from_atomarray, stripped_gemmi): + """Space group is preserved through the biotite→gemmi conversion.""" + assert gemmi_structure_from_atomarray.spacegroup_hm == stripped_gemmi.spacegroup_hm + + def test_atoms_match_pdb(self, gemmi_structure_from_atomarray, stripped_gemmi): + """Atom names and positions match the original PDB, in the same order.""" + # atom order is preserved: biotite keeps PDB file order, array2hier reconstructs it + parser_from_atomarray = PDBParser(gemmi_structure_from_atomarray) + parser_from_gemmi = PDBParser(stripped_gemmi) + assert np.array_equal(parser_from_atomarray.atom_name, parser_from_gemmi.atom_name) + np.testing.assert_allclose( + parser_from_atomarray.atom_pos, parser_from_gemmi.atom_pos, atol=1e-3 + ) + + def test_occupancy_change_is_applied(self, stripped_atom_array, stripped_gemmi): + """Custom occupancy values are correctly written to each altloc group.""" + occ_values = [0.2, 0.8, 0.0] + altloc_info = detect_altlocs(stripped_atom_array) + arr = assign_occupancies(stripped_atom_array, altloc_info, "custom", occ_values) + parser = PDBParser( + atomarray_to_gemmi(arr, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm) + ) + for altloc, expected in zip(altloc_info.altloc_ids, occ_values): + assert np.allclose(parser.atom_occ[altloc_info.atom_masks[altloc]], expected) + + def test_fprotein_matches_direct_gemmi( + self, gemmi_structure_from_atomarray, stripped_gemmi, device + ): + """Fprotein amplitudes from the converted structure match those from + the original gemmi structure.""" + f_atomarray = _compute_fprotein(gemmi_structure_from_atomarray, device) + f_direct = _compute_fprotein(stripped_gemmi, device) + np.testing.assert_allclose(np.abs(f_atomarray), np.abs(f_direct), atol=1e-3) + + def test_occupancy_warns_on_extra_values(self, stripped_atom_array, caplog): + """A warning is logged when more occupancy values are provided than there are altlocs.""" + altloc_info = detect_altlocs(stripped_atom_array) + with caplog.at_level(logging.WARNING): + assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.2, 0.8, 0.0, 0.0]) + assert "Extra values will be ignored" in caplog.text + + def test_occupancy_warns_on_missing_values(self, stripped_atom_array, caplog): + """A warning is logged when fewer occupancy values are provided than there are altlocs.""" + altloc_info = detect_altlocs(stripped_atom_array) + with caplog.at_level(logging.WARNING): + assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.5, 0.5]) + assert "Missing values are automatically set to 0" in caplog.text + + def test_occupancy_raises_on_out_of_range(self, stripped_atom_array): + """ValueError is raised when an occupancy value is outside [0.0, 1.0].""" + altloc_info = detect_altlocs(stripped_atom_array) + with pytest.raises(ValueError, match="out of range"): + assign_occupancies(stripped_atom_array, altloc_info, "custom", [1.5, 0.0, 0.0]) + + def test_occupancy_raises_on_bad_sum(self, stripped_atom_array): + """ValueError is raised when occupancy values do not sum to 1.0.""" + altloc_info = detect_altlocs(stripped_atom_array) + with pytest.raises(ValueError, match="sum to 1.0"): + assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.3, 0.3, 0.3]) + + def test_fprotein_changes_with_occupancy(self, stripped_atom_array, stripped_gemmi, device): + """Fprotein amplitudes differ when occupancies changes from uniform to custom values.""" + altloc_info = detect_altlocs(stripped_atom_array) + + arr_uniform = assign_occupancies(stripped_atom_array, altloc_info, "uniform") + f_uniform = _compute_fprotein( + atomarray_to_gemmi(arr_uniform, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm), + device, + ) + + arr_custom = assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.2, 0.8, 0.0]) + f_custom = _compute_fprotein( + atomarray_to_gemmi(arr_custom, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm), + device, + ) + + assert not np.allclose(np.abs(f_uniform), np.abs(f_custom), atol=1e-3) diff --git a/tests/rewards/test_geometry.py b/tests/rewards/test_geometry.py new file mode 100644 index 00000000..1b510505 --- /dev/null +++ b/tests/rewards/test_geometry.py @@ -0,0 +1,17 @@ +"""A test case for testing the geometry reward + +This test case is designed to verify the reward function that evaluates the geometry of a molecular structure. +It checks whether the reward function correctly computes the reward +based on the provided atomic coordinates and other relevant parameters. + + +""" + +import pytest +import gemmi +from sampleworks.core.rewards.geometry import _covalent_radius + +def test_known_element_matches_gemmi(): + r = _covalent_radius("C") + assert r > 0, "Covalent radius for Carbon should be greater than 0" + assert r == gemmi.Element("C").covalent_r \ No newline at end of file From 996e09064d7d948376bce7ba0eaa34ca6bfc2b3c Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Wed, 29 Jul 2026 14:38:22 -0700 Subject: [PATCH 03/15] Change the optimization to N forward method so each optimization starts from a brand new start and the loss function is penalized by the average of the ensemble. --- .../core/scalers/latent_optimization.py | 87 ++++++++- .../utils/guidance_script_utils.py | 4 + tests/eval/test_generate_synthetic_sf.py | 166 ++++++++++++++++++ tests/rewards/test_geometry.py | 17 ++ 4 files changed, 269 insertions(+), 5 deletions(-) create mode 100644 tests/eval/test_generate_synthetic_sf.py create mode 100644 tests/rewards/test_geometry.py diff --git a/src/sampleworks/core/scalers/latent_optimization.py b/src/sampleworks/core/scalers/latent_optimization.py index 56a9e4bd..d38e64fe 100644 --- a/src/sampleworks/core/scalers/latent_optimization.py +++ b/src/sampleworks/core/scalers/latent_optimization.py @@ -120,6 +120,62 @@ def __call__(self, latents: Sequence[Tensor], baselines: Sequence[Tensor]) -> Te return torch.stack(terms).sum() +class _PerMemberStepper: + """Denoise one ensemble member at a time, each with its OWN trunk latent. + + The it_opt "multiple-leaf" scheme gives every ensemble member its own latent, so ``s``/``z`` + carry a leading ``ensemble_size`` batch dim. Stock model diffusion modules instead take an + un-batched conditioning and broadcast it internally, so a batched latent either crashes + (Protenix) or mis-broadcasts (Boltz's ``multiplicity``). This adapter sidesteps that in a + model-agnostic way: for each member it slices that member's latent, runs the wrapped model's + normal un-batched ``step`` on that member alone, and stacks the results. Members stay + independent, so gradients stay per-member; the summed density reward couples them only through + the ensemble average. Cost is N forwards instead of one batched call, with an equivalent result. + + ``featurize`` / ``initialize_from_prior`` pass straight through to the wrapped model. + """ + + def __init__( + self, + model, + io: AttrLatentIO, + *, + optimize_single: bool, + optimize_pair: bool, + ensemble_size: int, + ): + self._model = model + self._io = io + self._optimize_single = optimize_single + self._optimize_pair = optimize_pair + self._ensemble_size = ensemble_size + + 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) + + def featurize(self, *args, **kwargs): + return self._model.featurize(*args, **kwargs) + + def initialize_from_prior(self, *args, **kwargs): + return self._model.initialize_from_prior(*args, **kwargs) + + class LatentOptimization: """Trajectory scaler that optimizes the model's ``s``/``z`` latents (IT-opt). @@ -253,6 +309,16 @@ def sample( schedule = sampler.compute_schedule(num_steps=self.num_steps) grad_enabler = _GradEnablingScaler() + # Denoise per member so each uses its own latent: stock model diffusion modules take an + # un-batched conditioning, so the batched per-member latents can't go through in one call. + stepper = _PerMemberStepper( + model, + io, + optimize_single=self.optimize_single, + optimize_pair=self.optimize_pair, + ensemble_size=self.ensemble_size, + ) + # --- 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 valid geometry for density @@ -275,7 +341,7 @@ def sample( for outer in 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, + model=stepper, sampler=sampler, reward=reward, features=features, @@ -303,7 +369,7 @@ def sample( # --- final clean sampling pass with the optimized latents --------------- final_coords, trajectory, losses = self._sample_with_frozen_latents( - model=model, + model=stepper, sampler=sampler, reward=reward, io=io, @@ -337,8 +403,12 @@ def _leaf_latents(self, features: GenerativeModelInput, io: AttrLatentIO): 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. + updated in place across rounds and steps). + + Each leaf gets a leading ``ensemble_size`` batch dimension -- one INDEPENDENT latent per + ensemble member, all cloned from the same trunk baseline (the it_opt "multiple-leaf" + scheme), so members can diverge rather than share one latent. The baseline kept for the + anchor stays un-batched and broadcasts across members. """ conditioning = features.conditioning latents: list[Tensor] = [] @@ -358,7 +428,14 @@ def _leaf_latents(self, features: GenerativeModelInput, io: AttrLatentIO): if baseline is None: continue baseline = baseline.detach() - leaf = baseline.clone().requires_grad_(True) + # it_opt "multiple-leaf" scheme: give each ensemble member its OWN latent. We stack + # ensemble_size independent copies of the trunk baseline into a leading batch dim, so + # each member gets its own gradient and can diverge, instead of collapsing onto one + # shared latent. (The reference does the same via batch-expand-then-clone.) Requires the + # diffusion module to accept a per-member (batched) conditioning -- verify with a + # batched gradcheck before relying on it. + member_copies = [baseline for _ in range(self.ensemble_size)] + leaf = torch.stack(member_copies).requires_grad_(True) conditioning = write(conditioning, leaf) latents.append(leaf) baselines.append(baseline) diff --git a/src/sampleworks/utils/guidance_script_utils.py b/src/sampleworks/utils/guidance_script_utils.py index 2e9313e0..de8717e9 100644 --- a/src/sampleworks/utils/guidance_script_utils.py +++ b/src/sampleworks/utils/guidance_script_utils.py @@ -507,6 +507,10 @@ def _run_guidance(args: GuidanceConfig, guidance_type: str, model_wrapper, devic structure = annotate_structure_for_protenix( structure, + # Root Protenix's per-sample input dump (protenix_input*.json) under the job's output + # dir; otherwise out_dir falls back to the input id, resolves against the CWD, and + # leaves a stray / folder there on every run. + out_dir=str(Path(args.output_dir) / "protenix_input"), recycling_steps=recycling_steps, # ensemble_size removed: #330 dropped it from the wrapper annotate signature; the # ensemble is now sized by the scaler via initialize_from_prior(). diff --git a/tests/eval/test_generate_synthetic_sf.py b/tests/eval/test_generate_synthetic_sf.py new file mode 100644 index 00000000..39534b37 --- /dev/null +++ b/tests/eval/test_generate_synthetic_sf.py @@ -0,0 +1,166 @@ +"""Tests for atomarray_to_gemmi in generate_synthetic_sf, using real 6b8x structure.""" + +import logging +from pathlib import Path + +import gemmi +import numpy as np +import pytest +import torch +from atomworks.io.transforms.atom_array import remove_waters +from biotite.structure import AtomArray +from sampleworks.eval.generate_synthetic_sf import atomarray_to_gemmi +from sampleworks.eval.synthetic_utils import assign_occupancies +from sampleworks.utils.atom_array_utils import ( + detect_altlocs, + keep_amino_acids, + keep_polymer, + load_structure_with_altlocs, + remove_hydrogens, +) +from SFC_Torch import SFcalculator +from SFC_Torch.io import PDBParser +from SFC_Torch.utils import assert_numpy + + +DMIN = 2.0 + + +@pytest.fixture(scope="module") +def stripped_gemmi(resources_dir: Path) -> gemmi.Structure: + """gemmi.Structure with hydrogens, ligands, and waters removed using gemmi methods.""" + s = gemmi.read_structure(str(resources_dir / "6b8x" / "6b8x_final.pdb")) + s.remove_hydrogens() + s.remove_ligands_and_waters() + return s + + +@pytest.fixture(scope="module") +def stripped_atom_array(resources_dir: Path) -> AtomArray: + """AtomArray with hydrogens, waters, and non-polymer/non-amino-acid atoms + removed using existing utils.""" + arr = load_structure_with_altlocs(resources_dir / "6b8x" / "6b8x_final.pdb") + arr = remove_hydrogens(arr) + arr = remove_waters(arr) + arr = keep_polymer(keep_amino_acids(arr)) + assert isinstance(arr, AtomArray) + return arr + + +@pytest.fixture(scope="module") +def gemmi_structure_from_atomarray( + stripped_atom_array, stripped_gemmi: gemmi.Structure +) -> gemmi.Structure: + """gemmi.Structure converted from the stripped AtomArray.""" + return atomarray_to_gemmi( + stripped_atom_array, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm + ) + + +def _compute_fprotein(gemmi_structure: gemmi.Structure, device: torch.device) -> np.ndarray: + """Compute |Fprotein| amplitudes from a gemmi structure via SFcalculator at ``DMIN`` + resolution. The final assert_numpy converts any tensor or list to a numpy array.""" + sfc = SFcalculator( + PDBParser(gemmi_structure), + mtzdata=None, + dmin=DMIN, + mode="xray", + anomalous=False, + set_experiment=False, + device=device, + ) + sfc.calc_fprotein() + return assert_numpy(sfc.Fprotein_asu) + + +class TestAtomArrayToGemmi: + """Tests for atomarray_to_gemmi using the 6b8x structure.""" + + def test_cell_matches_pdb(self, gemmi_structure_from_atomarray, stripped_gemmi): + """Unit cell parameters are preserved through the biotite→gemmi conversion.""" + result = gemmi_structure_from_atomarray.cell + expected = stripped_gemmi.cell + assert result.a == pytest.approx(expected.a) + assert result.b == pytest.approx(expected.b) + assert result.c == pytest.approx(expected.c) + assert result.alpha == pytest.approx(expected.alpha) + assert result.beta == pytest.approx(expected.beta) + assert result.gamma == pytest.approx(expected.gamma) + + def test_space_group_matches_pdb(self, gemmi_structure_from_atomarray, stripped_gemmi): + """Space group is preserved through the biotite→gemmi conversion.""" + assert gemmi_structure_from_atomarray.spacegroup_hm == stripped_gemmi.spacegroup_hm + + def test_atoms_match_pdb(self, gemmi_structure_from_atomarray, stripped_gemmi): + """Atom names and positions match the original PDB, in the same order.""" + # atom order is preserved: biotite keeps PDB file order, array2hier reconstructs it + parser_from_atomarray = PDBParser(gemmi_structure_from_atomarray) + parser_from_gemmi = PDBParser(stripped_gemmi) + assert np.array_equal(parser_from_atomarray.atom_name, parser_from_gemmi.atom_name) + np.testing.assert_allclose( + parser_from_atomarray.atom_pos, parser_from_gemmi.atom_pos, atol=1e-3 + ) + + def test_occupancy_change_is_applied(self, stripped_atom_array, stripped_gemmi): + """Custom occupancy values are correctly written to each altloc group.""" + occ_values = [0.2, 0.8, 0.0] + altloc_info = detect_altlocs(stripped_atom_array) + arr = assign_occupancies(stripped_atom_array, altloc_info, "custom", occ_values) + parser = PDBParser( + atomarray_to_gemmi(arr, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm) + ) + for altloc, expected in zip(altloc_info.altloc_ids, occ_values): + assert np.allclose(parser.atom_occ[altloc_info.atom_masks[altloc]], expected) + + def test_fprotein_matches_direct_gemmi( + self, gemmi_structure_from_atomarray, stripped_gemmi, device + ): + """Fprotein amplitudes from the converted structure match those from + the original gemmi structure.""" + f_atomarray = _compute_fprotein(gemmi_structure_from_atomarray, device) + f_direct = _compute_fprotein(stripped_gemmi, device) + np.testing.assert_allclose(np.abs(f_atomarray), np.abs(f_direct), atol=1e-3) + + def test_occupancy_warns_on_extra_values(self, stripped_atom_array, caplog): + """A warning is logged when more occupancy values are provided than there are altlocs.""" + altloc_info = detect_altlocs(stripped_atom_array) + with caplog.at_level(logging.WARNING): + assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.2, 0.8, 0.0, 0.0]) + assert "Extra values will be ignored" in caplog.text + + def test_occupancy_warns_on_missing_values(self, stripped_atom_array, caplog): + """A warning is logged when fewer occupancy values are provided than there are altlocs.""" + altloc_info = detect_altlocs(stripped_atom_array) + with caplog.at_level(logging.WARNING): + assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.5, 0.5]) + assert "Missing values are automatically set to 0" in caplog.text + + def test_occupancy_raises_on_out_of_range(self, stripped_atom_array): + """ValueError is raised when an occupancy value is outside [0.0, 1.0].""" + altloc_info = detect_altlocs(stripped_atom_array) + with pytest.raises(ValueError, match="out of range"): + assign_occupancies(stripped_atom_array, altloc_info, "custom", [1.5, 0.0, 0.0]) + + def test_occupancy_raises_on_bad_sum(self, stripped_atom_array): + """ValueError is raised when occupancy values do not sum to 1.0.""" + altloc_info = detect_altlocs(stripped_atom_array) + with pytest.raises(ValueError, match="sum to 1.0"): + assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.3, 0.3, 0.3]) + + def test_fprotein_changes_with_occupancy(self, stripped_atom_array, stripped_gemmi, device): + """Fprotein amplitudes differ when occupancies changes from uniform to custom values.""" + altloc_info = detect_altlocs(stripped_atom_array) + + arr_uniform = assign_occupancies(stripped_atom_array, altloc_info, "uniform") + f_uniform = _compute_fprotein( + atomarray_to_gemmi(arr_uniform, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm), + device, + ) + + arr_custom = assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.2, 0.8, 0.0]) + f_custom = _compute_fprotein( + atomarray_to_gemmi(arr_custom, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm), + device, + ) + + assert not np.allclose(np.abs(f_uniform), np.abs(f_custom), atol=1e-3) diff --git a/tests/rewards/test_geometry.py b/tests/rewards/test_geometry.py new file mode 100644 index 00000000..1b510505 --- /dev/null +++ b/tests/rewards/test_geometry.py @@ -0,0 +1,17 @@ +"""A test case for testing the geometry reward + +This test case is designed to verify the reward function that evaluates the geometry of a molecular structure. +It checks whether the reward function correctly computes the reward +based on the provided atomic coordinates and other relevant parameters. + + +""" + +import pytest +import gemmi +from sampleworks.core.rewards.geometry import _covalent_radius + +def test_known_element_matches_gemmi(): + r = _covalent_radius("C") + assert r > 0, "Covalent radius for Carbon should be greater than 0" + assert r == gemmi.Element("C").covalent_r \ No newline at end of file From a585ee7262e49ba41c52bd52992f52698e5a7658 Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Wed, 29 Jul 2026 15:14:12 -0700 Subject: [PATCH 04/15] chore(tests): drop stale pre-#329 copy of the synthetic-sf test 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. --- tests/eval/test_generate_synthetic_sf.py | 166 ----------------------- 1 file changed, 166 deletions(-) delete mode 100644 tests/eval/test_generate_synthetic_sf.py diff --git a/tests/eval/test_generate_synthetic_sf.py b/tests/eval/test_generate_synthetic_sf.py deleted file mode 100644 index 39534b37..00000000 --- a/tests/eval/test_generate_synthetic_sf.py +++ /dev/null @@ -1,166 +0,0 @@ -"""Tests for atomarray_to_gemmi in generate_synthetic_sf, using real 6b8x structure.""" - -import logging -from pathlib import Path - -import gemmi -import numpy as np -import pytest -import torch -from atomworks.io.transforms.atom_array import remove_waters -from biotite.structure import AtomArray -from sampleworks.eval.generate_synthetic_sf import atomarray_to_gemmi -from sampleworks.eval.synthetic_utils import assign_occupancies -from sampleworks.utils.atom_array_utils import ( - detect_altlocs, - keep_amino_acids, - keep_polymer, - load_structure_with_altlocs, - remove_hydrogens, -) -from SFC_Torch import SFcalculator -from SFC_Torch.io import PDBParser -from SFC_Torch.utils import assert_numpy - - -DMIN = 2.0 - - -@pytest.fixture(scope="module") -def stripped_gemmi(resources_dir: Path) -> gemmi.Structure: - """gemmi.Structure with hydrogens, ligands, and waters removed using gemmi methods.""" - s = gemmi.read_structure(str(resources_dir / "6b8x" / "6b8x_final.pdb")) - s.remove_hydrogens() - s.remove_ligands_and_waters() - return s - - -@pytest.fixture(scope="module") -def stripped_atom_array(resources_dir: Path) -> AtomArray: - """AtomArray with hydrogens, waters, and non-polymer/non-amino-acid atoms - removed using existing utils.""" - arr = load_structure_with_altlocs(resources_dir / "6b8x" / "6b8x_final.pdb") - arr = remove_hydrogens(arr) - arr = remove_waters(arr) - arr = keep_polymer(keep_amino_acids(arr)) - assert isinstance(arr, AtomArray) - return arr - - -@pytest.fixture(scope="module") -def gemmi_structure_from_atomarray( - stripped_atom_array, stripped_gemmi: gemmi.Structure -) -> gemmi.Structure: - """gemmi.Structure converted from the stripped AtomArray.""" - return atomarray_to_gemmi( - stripped_atom_array, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm - ) - - -def _compute_fprotein(gemmi_structure: gemmi.Structure, device: torch.device) -> np.ndarray: - """Compute |Fprotein| amplitudes from a gemmi structure via SFcalculator at ``DMIN`` - resolution. The final assert_numpy converts any tensor or list to a numpy array.""" - sfc = SFcalculator( - PDBParser(gemmi_structure), - mtzdata=None, - dmin=DMIN, - mode="xray", - anomalous=False, - set_experiment=False, - device=device, - ) - sfc.calc_fprotein() - return assert_numpy(sfc.Fprotein_asu) - - -class TestAtomArrayToGemmi: - """Tests for atomarray_to_gemmi using the 6b8x structure.""" - - def test_cell_matches_pdb(self, gemmi_structure_from_atomarray, stripped_gemmi): - """Unit cell parameters are preserved through the biotite→gemmi conversion.""" - result = gemmi_structure_from_atomarray.cell - expected = stripped_gemmi.cell - assert result.a == pytest.approx(expected.a) - assert result.b == pytest.approx(expected.b) - assert result.c == pytest.approx(expected.c) - assert result.alpha == pytest.approx(expected.alpha) - assert result.beta == pytest.approx(expected.beta) - assert result.gamma == pytest.approx(expected.gamma) - - def test_space_group_matches_pdb(self, gemmi_structure_from_atomarray, stripped_gemmi): - """Space group is preserved through the biotite→gemmi conversion.""" - assert gemmi_structure_from_atomarray.spacegroup_hm == stripped_gemmi.spacegroup_hm - - def test_atoms_match_pdb(self, gemmi_structure_from_atomarray, stripped_gemmi): - """Atom names and positions match the original PDB, in the same order.""" - # atom order is preserved: biotite keeps PDB file order, array2hier reconstructs it - parser_from_atomarray = PDBParser(gemmi_structure_from_atomarray) - parser_from_gemmi = PDBParser(stripped_gemmi) - assert np.array_equal(parser_from_atomarray.atom_name, parser_from_gemmi.atom_name) - np.testing.assert_allclose( - parser_from_atomarray.atom_pos, parser_from_gemmi.atom_pos, atol=1e-3 - ) - - def test_occupancy_change_is_applied(self, stripped_atom_array, stripped_gemmi): - """Custom occupancy values are correctly written to each altloc group.""" - occ_values = [0.2, 0.8, 0.0] - altloc_info = detect_altlocs(stripped_atom_array) - arr = assign_occupancies(stripped_atom_array, altloc_info, "custom", occ_values) - parser = PDBParser( - atomarray_to_gemmi(arr, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm) - ) - for altloc, expected in zip(altloc_info.altloc_ids, occ_values): - assert np.allclose(parser.atom_occ[altloc_info.atom_masks[altloc]], expected) - - def test_fprotein_matches_direct_gemmi( - self, gemmi_structure_from_atomarray, stripped_gemmi, device - ): - """Fprotein amplitudes from the converted structure match those from - the original gemmi structure.""" - f_atomarray = _compute_fprotein(gemmi_structure_from_atomarray, device) - f_direct = _compute_fprotein(stripped_gemmi, device) - np.testing.assert_allclose(np.abs(f_atomarray), np.abs(f_direct), atol=1e-3) - - def test_occupancy_warns_on_extra_values(self, stripped_atom_array, caplog): - """A warning is logged when more occupancy values are provided than there are altlocs.""" - altloc_info = detect_altlocs(stripped_atom_array) - with caplog.at_level(logging.WARNING): - assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.2, 0.8, 0.0, 0.0]) - assert "Extra values will be ignored" in caplog.text - - def test_occupancy_warns_on_missing_values(self, stripped_atom_array, caplog): - """A warning is logged when fewer occupancy values are provided than there are altlocs.""" - altloc_info = detect_altlocs(stripped_atom_array) - with caplog.at_level(logging.WARNING): - assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.5, 0.5]) - assert "Missing values are automatically set to 0" in caplog.text - - def test_occupancy_raises_on_out_of_range(self, stripped_atom_array): - """ValueError is raised when an occupancy value is outside [0.0, 1.0].""" - altloc_info = detect_altlocs(stripped_atom_array) - with pytest.raises(ValueError, match="out of range"): - assign_occupancies(stripped_atom_array, altloc_info, "custom", [1.5, 0.0, 0.0]) - - def test_occupancy_raises_on_bad_sum(self, stripped_atom_array): - """ValueError is raised when occupancy values do not sum to 1.0.""" - altloc_info = detect_altlocs(stripped_atom_array) - with pytest.raises(ValueError, match="sum to 1.0"): - assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.3, 0.3, 0.3]) - - def test_fprotein_changes_with_occupancy(self, stripped_atom_array, stripped_gemmi, device): - """Fprotein amplitudes differ when occupancies changes from uniform to custom values.""" - altloc_info = detect_altlocs(stripped_atom_array) - - arr_uniform = assign_occupancies(stripped_atom_array, altloc_info, "uniform") - f_uniform = _compute_fprotein( - atomarray_to_gemmi(arr_uniform, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm), - device, - ) - - arr_custom = assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.2, 0.8, 0.0]) - f_custom = _compute_fprotein( - atomarray_to_gemmi(arr_custom, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm), - device, - ) - - assert not np.allclose(np.abs(f_uniform), np.abs(f_custom), atol=1e-3) From 4cf854de8d1049c2773c13647662118e8027caeb Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Tue, 4 Aug 2026 11:17:00 -0700 Subject: [PATCH 05/15] chore(it-opt): track the ensemble runner and scoring scripts 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) --- it_opt_scratch/paper_maxrmsd_selections.csv | 41 ++ it_opt_scratch/run_targets.py | 452 ++++++++++++++++++++ it_opt_scratch/run_targets_simplified.py | 186 ++++++++ it_opt_scratch/score_paper_rmsd.py | 242 +++++++++++ it_opt_scratch/score_paper_rscc.py | 299 +++++++++++++ it_opt_scratch/score_paper_simplified.py | 154 +++++++ it_opt_scratch/score_rmsd_simplified.py | 259 +++++++++++ it_opt_scratch/score_rscc_simplified.py | 240 +++++++++++ it_opt_scratch/slurm_ensemble_and_score.sh | 148 +++++++ 9 files changed, 2021 insertions(+) create mode 100644 it_opt_scratch/paper_maxrmsd_selections.csv create mode 100644 it_opt_scratch/run_targets.py create mode 100644 it_opt_scratch/run_targets_simplified.py create mode 100644 it_opt_scratch/score_paper_rmsd.py create mode 100644 it_opt_scratch/score_paper_rscc.py create mode 100644 it_opt_scratch/score_paper_simplified.py create mode 100644 it_opt_scratch/score_rmsd_simplified.py create mode 100644 it_opt_scratch/score_rscc_simplified.py create mode 100755 it_opt_scratch/slurm_ensemble_and_score.sh diff --git a/it_opt_scratch/paper_maxrmsd_selections.csv b/it_opt_scratch/paper_maxrmsd_selections.csv new file mode 100644 index 00000000..e0138d2c --- /dev/null +++ b/it_opt_scratch/paper_maxrmsd_selections.csv @@ -0,0 +1,41 @@ +protein,selection +1VME,chain A and resi 1-1;chain A and resi 164-164;chain A and resi 207-207;chain A and resi 218-218;chain A and resi 245-245;chain A and resi 269-269;chain A and resi 318-320;chain A and resi 326-326;chain A and resi 373-373;chain A and resi 67-67;chain A and resi 70-70 +2A26,chain A and resi 11-11;chain A and resi 16-18;chain A and resi 26-26;chain A and resi 30-30;chain A and resi 45-45 +2BWD,chain A and resi 106-106;chain A and resi 141-141;chain A and resi 16-16;chain A and resi 168-168;chain A and resi 197-199;chain A and resi 211-211;chain A and resi 22-22;chain A and resi 239-239;chain A and resi 276-276;chain A and resi 29-29;chain A and resi 302-302;chain A and resi 313-313;chain A and resi 321-321;chain A and resi 323-323;chain A and resi 325-325;chain A and resi 33-33;chain A and resi 330-330;chain A and resi 340-340;chain A and resi 36-36;chain A and resi 46-46;chain A and resi 50-50;chain A and resi 78-78;chain A and resi 92-92;chain A and resi 94-94;chain A and resi 98-98 +2E7Z,chain A and resi 105-105;chain A and resi 141-141;chain A and resi 195-197;chain A and resi 23-23;chain A and resi 277-279;chain A and resi 339-339;chain A and resi 344-344;chain A and resi 376-376;chain A and resi 446-446;chain A and resi 508-509;chain A and resi 549-551;chain A and resi 633-635;chain A and resi 664-666;chain A and resi 81-81;chain A and resi 93-93 +2I6H,chain B and resi 105-105;chain B and resi 108-108;chain B and resi 11-11;chain B and resi 131-131;chain B and resi 146-146;chain B and resi 153-155;chain B and resi 159-159;chain B and resi 165-165;chain B and resi 169-171;chain B and resi 23-23;chain B and resi 26-26;chain B and resi 33-33;chain B and resi 56-56 +2IGT,chain A and resi 108-109;chain A and resi 121-121;chain A and resi 128-128;chain A and resi 131-131;chain A and resi 181-181;chain A and resi 185-185;chain A and resi 19-19;chain A and resi 200-200;chain A and resi 215-215;chain A and resi 225-225;chain A and resi 235-235;chain A and resi 254-256;chain A and resi 261-261;chain A and resi 290-292;chain A and resi 53-53;chain A and resi 58-59;chain A and resi 62-62;chain A and resi 77-77;chain A and resi 94-94;chain A and resi 99-99 +2P97,chain A and resi 100-100;chain A and resi 112-112;chain A and resi 129-129;chain A and resi 14-14;chain A and resi 141-141;chain A and resi 147-147;chain A and resi 155-155;chain A and resi 166-166;chain A and resi 18-20;chain A and resi 180-180;chain A and resi 186-186;chain A and resi 2-2;chain A and resi 35-35;chain A and resi 56-57;chain A and resi 67-67;chain A and resi 97-97 +2QCV,chain A and resi 111-111;chain A and resi 144-144;chain A and resi 198-198;chain A and resi 2-4;chain A and resi 215-215;chain A and resi 261-261;chain A and resi 282-282;chain A and resi 284-284;chain A and resi 39-39;chain A and resi 54-54;chain A and resi 77-77;chain A and resi 85-85 +2XLW,chain A and resi 10-10;chain A and resi 106-106;chain A and resi 111-111;chain A and resi 115-115;chain A and resi 121-122;chain A and resi 18-20;chain A and resi 30-30;chain A and resi 37-37;chain A and resi 52-52;chain A and resi 60-62;chain A and resi 66-66;chain A and resi 69-71;chain A and resi 81-83;chain A and resi 87-89;chain A and resi 92-92;chain A and resi 98-99 +2YL0,chain A and resi 10-10;chain A and resi 102-104;chain A and resi 106-106;chain A and resi 115-115;chain A and resi 19-19;chain A and resi 30-30;chain A and resi 37-37;chain A and resi 40-40;chain A and resi 46-46;chain A and resi 52-52;chain A and resi 60-62;chain A and resi 66-66;chain A and resi 68-70;chain A and resi 74-76;chain A and resi 81-81;chain A and resi 88-89;chain A and resi 92-92;chain A and resi 94-94 +3DT2,chain A and resi 103-103;chain A and resi 109-111;chain A and resi 11-11;chain A and resi 120-120;chain A and resi 151-151;chain A and resi 153-154;chain A and resi 17-17;chain A and resi 191-191;chain A and resi 209-209;chain A and resi 21-21;chain A and resi 274-274;chain A and resi 276-276;chain A and resi 296-296;chain A and resi 320-320;chain A and resi 36-36;chain A and resi 39-39;chain A and resi 390-390;chain A and resi 41-42;chain A and resi 422-422;chain A and resi 50-51;chain A and resi 502-502;chain A and resi 510-510;chain A and resi 524-524;chain A and resi 545-545;chain A and resi 55-55;chain A and resi 578-578;chain A and resi 583-583;chain A and resi 587-587;chain A and resi 597-597;chain A and resi 613-613;chain A and resi 70-71 +3HVV,chain A and resi 58-60;chain A and resi 93-93 +3HYN,chain A and resi 121-121;chain A and resi 127-127;chain A and resi 133-134;chain A and resi 141-141;chain A and resi 149-149;chain A and resi 151-151;chain A and resi 160-160;chain A and resi 166-166;chain A and resi 182-182;chain A and resi 19-21;chain A and resi 31-31;chain A and resi 38-38;chain A and resi 45-45;chain A and resi 60-62;chain A and resi 69-69;chain A and resi 71-71;chain A and resi 78-78;chain A and resi 91-93;chain A and resi 97-97 +3I7M,chain A and resi 112-112;chain A and resi 130-130;chain A and resi 133-133;chain A and resi 15-15;chain A and resi 29-29;chain A and resi 38-39;chain A and resi 60-60;chain A and resi 68-68;chain A and resi 70-70;chain A and resi 78-79;chain A and resi 8-9;chain A and resi 85-85;chain A and resi 88-88;chain A and resi 96-96 +3IAC,chain A and resi 111-111;chain A and resi 186-186;chain A and resi 200-200;chain A and resi 212-212;chain A and resi 287-287;chain A and resi 317-317;chain A and resi 334-334;chain A and resi 55-57;chain A and resi 84-84 +3L4P,chain A and resi 11-11;chain A and resi 124-124;chain A and resi 13-13;chain A and resi 132-132;chain A and resi 134-134;chain A and resi 166-166;chain A and resi 171-173;chain A and resi 18-18;chain A and resi 182-182;chain A and resi 2-2;chain A and resi 204-206;chain A and resi 216-216;chain A and resi 233-233;chain A and resi 240-240;chain A and resi 261-261;chain A and resi 271-273;chain A and resi 303-305;chain A and resi 318-320;chain A and resi 334-334;chain A and resi 352-352;chain A and resi 372-372;chain A and resi 382-382;chain A and resi 396-396;chain A and resi 466-466;chain A and resi 468-468;chain A and resi 520-520;chain A and resi 55-55;chain A and resi 553-553;chain A and resi 557-557;chain A and resi 559-559;chain A and resi 579-579;chain A and resi 58-58;chain A and resi 587-587;chain A and resi 594-595;chain A and resi 597-597;chain A and resi 604-605;chain A and resi 661-661;chain A and resi 680-680;chain A and resi 722-724;chain A and resi 733-733;chain A and resi 736-737;chain A and resi 742-742;chain A and resi 748-750;chain A and resi 761-761;chain A and resi 783-783;chain A and resi 789-789;chain A and resi 82-84;chain A and resi 826-826;chain A and resi 834-834;chain A and resi 842-843;chain A and resi 847-847;chain A and resi 889-889 +3T94,chain A and resi 116-118;chain A and resi 132-133;chain A and resi 152-152;chain A and resi 204-205;chain A and resi 211-211;chain A and resi 229-229;chain A and resi 232-232;chain A and resi 42-42;chain A and resi 91-91 +4GMU,chain A and resi 120-120;chain A and resi 123-123;chain A and resi 14-14;chain A and resi 190-191;chain A and resi 224-224;chain A and resi 226-227;chain A and resi 235-235;chain A and resi 244-244;chain A and resi 253-253;chain A and resi 312-312;chain A and resi 350-350;chain A and resi 39-39;chain A and resi 403-403;chain A and resi 41-42;chain A and resi 442-442;chain A and resi 475-476;chain A and resi 50-51;chain A and resi 503-503;chain A and resi 543-543;chain A and resi 545-545;chain A and resi 56-56;chain A and resi 578-578;chain A and resi 582-582;chain A and resi 587-587;chain A and resi 606-607;chain A and resi 609-609;chain A and resi 613-613;chain A and resi 619-619;chain A and resi 70-70;chain A and resi 85-87;chain A and resi 97-97 +4OLE,chain B and resi 395-395;chain B and resi 424-426 +5A71,chain A and resi 105-105;chain A and resi 139-139;chain A and resi 159-159;chain A and resi 182-182;chain A and resi 194-195;chain A and resi 214-215;chain A and resi 281-283;chain A and resi 298-298;chain A and resi 74-74;chain A and resi 76-76 +5I09,chain A and resi 15-15;chain A and resi 166-166;chain A and resi 17-17;chain A and resi 170-170;chain A and resi 174-174;chain A and resi 188-188;chain A and resi 21-21;chain A and resi 213-213;chain A and resi 281-281;chain A and resi 303-303;chain A and resi 31-31;chain A and resi 325-325;chain A and resi 34-34;chain A and resi 351-351;chain A and resi 57-57;chain A and resi 69-71;chain A and resi 91-91 +5IMV,chain A and resi 106-106;chain A and resi 113-113;chain A and resi 119-119;chain A and resi 123-123;chain A and resi 127-128;chain A and resi 131-131;chain A and resi 133-133;chain A and resi 155-155;chain A and resi 24-24;chain A and resi 59-59;chain A and resi 6-6;chain A and resi 65-67;chain A and resi 74-74;chain A and resi 82-82;chain A and resi 84-84 +5MC8,chain A and resi 103-104;chain A and resi 107-107;chain A and resi 115-115;chain A and resi 128-129;chain A and resi 131-131;chain A and resi 136-136;chain A and resi 148-148;chain A and resi 149-149;chain A and resi 166-166;chain A and resi 168-168;chain A and resi 172-172;chain A and resi 176-176;chain A and resi 179-179;chain A and resi 181-181;chain A and resi 194-194;chain A and resi 196-196;chain A and resi 198-199;chain A and resi 202-202;chain A and resi 206-206;chain A and resi 210-210;chain A and resi 214-214;chain A and resi 223-223;chain A and resi 235-235;chain A and resi 238-238;chain A and resi 240-241;chain A and resi 246-246;chain A and resi 269-269;chain A and resi 275-275;chain A and resi 285-285;chain A and resi 293-293;chain A and resi 296-296;chain A and resi 300-300;chain A and resi 317-317;chain A and resi 324-324;chain A and resi 346-346;chain A and resi 357-359;chain A and resi 36-36;chain A and resi 366-366;chain A and resi 377-378;chain A and resi 43-43;chain A and resi 51-53;chain A and resi 79-80;chain A and resi 93-93 +5MHX,chain A and resi 131-131;chain A and resi 162-164;chain A and resi 17-17;chain A and resi 179-179;chain A and resi 191-191;chain A and resi 198-198;chain A and resi 201-201;chain A and resi 219-220;chain A and resi 251-252;chain A and resi 265-265;chain A and resi 368-369;chain A and resi 377-377;chain A and resi 380-381;chain A and resi 409-409;chain A and resi 433-433;chain A and resi 443-443;chain A and resi 46-47;chain A and resi 465-465;chain A and resi 476-476;chain A and resi 57-57 +5RBR,chain A and resi 11-13;chain A and resi 149-149;chain A and resi 154-154;chain A and resi 164-164;chain A and resi 174-174;chain A and resi 182-182;chain A and resi 206-206;chain A and resi 215-215;chain A and resi 223-224;chain A and resi 240-240;chain A and resi 246-246;chain A and resi 249-251;chain A and resi 268-268;chain A and resi 276-276;chain A and resi 279-281;chain A and resi 289-289;chain A and resi 291-291;chain A and resi 298-300;chain A and resi 311-311;chain A and resi 325-325;chain A and resi 329-329;chain A and resi 68-68;chain A and resi 71-71;chain A and resi 9-9 +5SOP,chain A and resi 10-12;chain A and resi 100-102;chain A and resi 113-113;chain A and resi 131-133;chain A and resi 138-140;chain A and resi 157-159;chain A and resi 20-22;chain A and resi 28-30;chain A and resi 38-38;chain A and resi 4-6;chain A and resi 42-42;chain A and resi 47-49;chain A and resi 87-88 +6B8X,chain A and resi 105-105;chain A and resi 115-115;chain A and resi 121-121;chain A and resi 147-148;chain A and resi 150-152;chain A and resi 157-157;chain A and resi 159-159;chain A and resi 176-176;chain A and resi 181-183;chain A and resi 196-197;chain A and resi 207-207;chain A and resi 216-216;chain A and resi 221-221;chain A and resi 224-224;chain A and resi 226-226;chain A and resi 233-233;chain A and resi 235-235;chain A and resi 239-241;chain A and resi 24-24;chain A and resi 246-246;chain A and resi 252-253;chain A and resi 26-26;chain A and resi 265-265;chain A and resi 269-269;chain A and resi 272-272;chain A and resi 276-276;chain A and resi 278-280;chain A and resi 32-32;chain A and resi 36-36;chain A and resi 72-72;chain A and resi 78-80;chain A and resi 92-92 +6DUR,chain A and resi 216-216;chain A and resi 269-269;chain A and resi 29-31;chain A and resi 359-361;chain A and resi 389-391;chain A and resi 397-399;chain A and resi 445-447;chain A and resi 66-66 +6NI5,chain A and resi 101-101;chain A and resi 114-114;chain A and resi 116-117;chain A and resi 12-12;chain A and resi 124-124;chain A and resi 135-135;chain A and resi 138-139;chain A and resi 15-15;chain A and resi 151-153;chain A and resi 191-191;chain A and resi 193-193;chain A and resi 201-202;chain A and resi 204-206;chain A and resi 209-210;chain A and resi 21-21;chain A and resi 213-215;chain A and resi 219-219;chain A and resi 24-24;chain A and resi 30-32;chain A and resi 35-35;chain A and resi 39-39;chain A and resi 4-4;chain A and resi 43-43;chain A and resi 52-52;chain A and resi 61-62;chain A and resi 64-64;chain A and resi 66-66;chain A and resi 74-74;chain A and resi 86-86;chain A and resi 9-9;chain A and resi 92-94 +6NI6,chain A and resi 101-101;chain A and resi 114-114;chain A and resi 116-117;chain A and resi 12-12;chain A and resi 124-124;chain A and resi 135-135;chain A and resi 138-139;chain A and resi 15-15;chain A and resi 151-153;chain A and resi 174-174;chain A and resi 178-179;chain A and resi 181-181;chain A and resi 191-191;chain A and resi 193-193;chain A and resi 201-202;chain A and resi 204-206;chain A and resi 209-210;chain A and resi 21-21;chain A and resi 213-215;chain A and resi 219-219;chain A and resi 30-32;chain A and resi 35-35;chain A and resi 39-39;chain A and resi 4-4;chain A and resi 52-52;chain A and resi 61-62;chain A and resi 64-64;chain A and resi 66-66;chain A and resi 74-74;chain A and resi 86-86;chain A and resi 92-94 +6RP1,chain C and resi 10-10;chain C and resi 185-185;chain C and resi 26-26;chain C and resi 287-287;chain C and resi 290-290;chain C and resi 326-328;chain C and resi 360-360;chain C and resi 418-418;chain C and resi 431-431;chain C and resi 465-465;chain C and resi 475-475;chain C and resi 5-5;chain C and resi 513-513;chain C and resi 522-522;chain C and resi 536-536;chain C and resi 64-64 +6YVM,chain A and resi 224-225;chain A and resi 234-235;chain A and resi 241-241;chain A and resi 245-245;chain A and resi 248-248;chain A and resi 251-251;chain A and resi 261-262;chain A and resi 286-288;chain A and resi 295-297;chain A and resi 303-303;chain A and resi 308-308;chain A and resi 323-323;chain A and resi 327-327;chain A and resi 330-330;chain A and resi 335-335;chain A and resi 338-338;chain A and resi 343-343;chain A and resi 352-352;chain A and resi 355-356;chain A and resi 359-359;chain A and resi 361-362;chain A and resi 365-365;chain A and resi 371-371;chain A and resi 374-374;chain A and resi 391-393;chain A and resi 395-395;chain A and resi 404-404;chain A and resi 411-413;chain A and resi 427-427;chain A and resi 438-438;chain A and resi 444-444;chain A and resi 454-454;chain A and resi 461-461;chain A and resi 465-465;chain A and resi 476-476 +7AVG,chain A and resi 1-1;chain A and resi 102-104;chain A and resi 106-106;chain A and resi 109-109;chain A and resi 112-114;chain A and resi 116-118;chain A and resi 123-123;chain A and resi 126-126;chain A and resi 13-13;chain A and resi 15-16;chain A and resi 19-19;chain A and resi 24-24;chain A and resi 37-37;chain A and resi 39-39;chain A and resi 43-44;chain A and resi 47-49;chain A and resi 5-5;chain A and resi 52-52;chain A and resi 65-65;chain A and resi 67-67;chain A and resi 7-7;chain A and resi 81-81;chain A and resi 85-85;chain A and resi 86-87;chain A and resi 89-89 +7FRD,chain A and resi 100-102;chain A and resi 105-105;chain A and resi 107-107;chain A and resi 134-136;chain A and resi 157-159;chain A and resi 162-163;chain A and resi 169-169;chain A and resi 24-26;chain A and resi 28-29;chain A and resi 3-5;chain A and resi 31-31;chain A and resi 38-38;chain A and resi 47-49;chain A and resi 58-58;chain A and resi 87-88;chain A and resi 9-11;chain A and resi 90-90 +7P6M,chain A and resi 102-104;chain A and resi 108-109;chain A and resi 111-111;chain A and resi 113-113;chain A and resi 116-117;chain A and resi 119-120;chain A and resi 123-124;chain A and resi 15-15;chain A and resi 21-21;chain A and resi 39-39;chain A and resi 44-45;chain A and resi 47-48;chain A and resi 5-5;chain A and resi 52-52;chain A and resi 63-63;chain A and resi 65-65;chain A and resi 78-79;chain A and resi 81-81;chain A and resi 85-87;chain A and resi 89-89 +7Z0E,chain P and resi 100-100;chain P and resi 104-104;chain P and resi 11-13;chain P and resi 115-115;chain P and resi 122-124;chain P and resi 131-131;chain P and resi 132-132;chain P and resi 133-133;chain P and resi 137-139;chain P and resi 145-145;chain P and resi 15-15;chain P and resi 18-20;chain P and resi 181-183;chain P and resi 193-194;chain P and resi 195-197;chain P and resi 202-202;chain P and resi 204-204;chain P and resi 206-208;chain P and resi 209-209;chain P and resi 214-216;chain P and resi 22-22;chain P and resi 224-224;chain P and resi 226-226;chain P and resi 23-24;chain P and resi 27-27;chain P and resi 28-28;chain P and resi 32-32;chain P and resi 45-47;chain P and resi 48-48;chain P and resi 54-56;chain P and resi 59-59;chain P and resi 60-60;chain P and resi 64-65;chain P and resi 7-7;chain P and resi 82-84;chain P and resi 93-94;chain P and resi 96-97 +8VPW,chain A and resi 101-101;chain A and resi 117-117;chain A and resi 12-12;chain A and resi 124-124;chain A and resi 126-126;chain A and resi 138-139;chain A and resi 142-142;chain A and resi 147-147;chain A and resi 149-149;chain A and resi 15-15;chain A and resi 151-153;chain A and resi 162-162;chain A and resi 165-165;chain A and resi 201-202;chain A and resi 204-204;chain A and resi 210-210;chain A and resi 214-215;chain A and resi 31-31;chain A and resi 35-35;chain A and resi 42-42;chain A and resi 49-49;chain A and resi 51-52;chain A and resi 67-67;chain A and resi 72-72;chain A and resi 74-74;chain A and resi 95-97 +8VQ1,chain A and resi 101-101;chain A and resi 117-117;chain A and resi 12-12;chain A and resi 124-124;chain A and resi 138-139;chain A and resi 142-142;chain A and resi 147-147;chain A and resi 149-149;chain A and resi 15-15;chain A and resi 151-153;chain A and resi 162-162;chain A and resi 165-165;chain A and resi 201-202;chain A and resi 204-204;chain A and resi 210-210;chain A and resi 214-215;chain A and resi 35-35;chain A and resi 42-42;chain A and resi 49-49;chain A and resi 52-52;chain A and resi 67-67;chain A and resi 72-72;chain A and resi 74-74;chain A and resi 95-97 +8Z76,chain C and resi 148-148;chain C and resi 162-162;chain C and resi 230-230;chain C and resi 248-248;chain C and resi 253-253;chain C and resi 256-258;chain C and resi 470-470;chain C and resi 82-82 +9BN8,chain A and resi 113-113;chain A and resi 151-151;chain A and resi 203-203;chain A and resi 208-208;chain A and resi 218-218;chain A and resi 226-228;chain A and resi 238-238;chain A and resi 251-251;chain A and resi 315-315;chain A and resi 327-327;chain A and resi 33-33;chain A and resi 336-337;chain A and resi 364-364;chain A and resi 375-375;chain A and resi 394-394;chain A and resi 4-4;chain A and resi 428-428;chain A and resi 431-431;chain A and resi 438-438 diff --git a/it_opt_scratch/run_targets.py b/it_opt_scratch/run_targets.py new file mode 100644 index 00000000..c36d2351 --- /dev/null +++ b/it_opt_scratch/run_targets.py @@ -0,0 +1,452 @@ +"""Batch conformational-ensemble generation over a list of protein targets. + +This drives the sampleworks guidance pipeline through its public entry point (`run_guidance`): the +model is loaded ONCE, a `GuidanceConfig` is built per run, and the pipeline itself does +featurize -> reward -> scaler -> sample -> save. Output for each (target, mode) lands under +// (refined.cif + trajectory + losses). A per-run failure is recorded and the batch +continues; a batch_summary.json is written at the end. + +Where the inputs come from (three forms, checked in this order): + * CSV --targets x.csv columns: name,structure,density,resolution + (out dir = /) + * JSON --targets x.json list of {name, density, structure, resolution, out_dir} + * IDs --proteins 2YL0,5I09 / --proteins-file ids.txt + paths are built from --structure-dir/--structure-template and + --density-dir/--density-template, so you only list PDB IDs. Templates take + {pdb} (as written), {PDB} (upper) and {pdb_low} (lower). + * none the built-in DEFAULT_TARGETS (1vme) + +Where the outputs go: + /// refined.cif trajectory/ losses.txt job_metadata.json run.log + --summary path batch_summary.json (per-shard copies alongside it when sharded) + +Modes (--modes, comma-separated, or 'all') -> guidance type: + baseline unguided sampling (pure_guidance, no step scaler) + s_only IT-opt, which_latent=single + z_only IT-opt, which_latent=pair + s_plus_z IT-opt, which_latent=both + coord_guidance shipped coordinate DPS (pure_guidance + noise-space step scaler) + +Run on the pod from the repo root -- one process, explicit CSV: + pixi run -e protenix-dev python it_opt_scratch/run_targets.py \ + --targets /home/dev/test_data/proteins.csv --output-base it_opt_scratch/targets_out \ + --modes baseline,z_only,coord_guidance --ensemble-size 4 --num-steps 200 --outer-steps 2 + +Run on 4 GPUs with 2 proteins in flight per GPU (8 worker processes), IDs + directories: + pixi run -e protenix-dev python -u it_opt_scratch/run_targets.py \ + --proteins-file it_opt_scratch/regen11.txt \ + --structure-dir /home/dev/test_data/processed \ + --density-dir it_opt_scratch/targets_out_11_regen_0.5occA_0.5occB/density_maps \ + --density-template '{PDB}_0.5occA_0.5occB_1.00A.ccp4' \ + --name-template '{PDB}_0.5occA_0.5occB' --resolution 1.0 \ + --output-base it_opt_scratch/targets_out_11_regen_0.5occA_0.5occB \ + --modes s_plus_z --ensemble-size 8 --bond-length-weight 5e-5 \ + --gpus 0,1,2,3 --jobs-per-gpu 2 + +Add --dry-run to either command to print the resolved plan (paths, existence, shard +assignment) without loading the model. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import os +import subprocess +import sys +import traceback +from pathlib import Path + +import torch + +from sampleworks.utils.guidance_constants import GuidanceType, StructurePredictor +from sampleworks.utils.guidance_script_arguments import GuidanceConfig +from sampleworks.utils.guidance_script_utils import get_model_and_device, run_guidance + +REPO = Path(__file__).resolve().parents[1] +ALL_MODES = ["baseline", "s_only", "z_only", "s_plus_z", "coord_guidance"] + +DEFAULT_TARGETS = [ + { + "name": "1vme", + "density": "tests/resources/1vme/1vme_final_carved_edited_0.5occA_0.5occB_1.80A.ccp4", + "structure": "tests/resources/1vme/1vme_final_carved_edited_0.5occA_0.5occB.cif", + "resolution": 1.8, + "out_dir": "it_opt_scratch/targets_out/1vme", + }, +] + + +# ============================ the batch (top-down) ============================ + +def main() -> None: + args = parse_args() + targets = collect_targets(args) + modes = ALL_MODES if args.modes == "all" else [m for m in args.modes.split(",") if m] + + if args.dry_run: # resolve and report the plan without touching a GPU + print_plan(targets, modes, args) + return + + # A parent process (--gpus given, no --shard-index) only fans out; the children do the work. + if args.gpus and args.shard_index is None: + launch_workers(targets, args) + return + + if args.num_shards > 1: # this process is one worker: take its slice, round-robin + targets = targets[args.shard_index :: args.num_shards] + + print(f"targets={[t['name'] for t in targets]} modes={modes} model={args.model}\n" + f"ensemble_size={args.ensemble_size} num_steps={args.num_steps} " + f"outer_steps={args.outer_steps} lr={args.lr} anchor={args.anchor}") + + device, model = get_model_and_device(args.device, args.checkpoint, StructurePredictor(args.model)) + + summary: list[dict] = [] + for target in targets: + print(f"\n######### TARGET: {target['name']} #########") + for mode in modes: + summary.append(generate_one(target, mode, model, device, args)) + + write_summary(summary, shard_summary_path(resolve(args.summary), args.shard_index)) + + +def generate_one(target: dict, mode: str, model, device, args) -> dict: + """Generate + save one ensemble for (target, mode) via run_guidance; return a summary record.""" + name = target["name"] + out_dir = resolve(target["out_dir"]) / mode + try: + guidance_type, extras = guidance_for_mode(mode, args) + config = build_config(target, guidance_type, out_dir, args) + for key, value in extras.items(): # mode-specific args the arg-adders would otherwise set + setattr(config, key, value) + out_dir.mkdir(parents=True, exist_ok=True) + torch.manual_seed(args.seed) # same start noise across modes -> comparable ensembles + job = run_guidance(config, guidance_type, model, device) + status = getattr(job, "status", "unknown") + print(f"[{name}/{mode}] {status} -> {out_dir}/refined.cif") + return {"target": name, "mode": mode, "status": "OK" if status == "success" else "FAILED", + "out_dir": str(out_dir)} + except Exception: + print(f"[{name}/{mode}] FAILED:") + traceback.print_exc() + return {"target": name, "mode": mode, "status": "FAILED"} + + +# ============================== target inputs =============================== +# Added for the ID-and-directory input mode. The CSV/JSON contract below is unchanged -- +# collect_targets() only falls through to the ID mode when --targets is not given. + +def collect_targets(args) -> list[dict]: + """Build the target list from --targets, or from --proteins/--proteins-file, or the default. + + Every target is a dict with name, structure, density, resolution, out_dir -- the same + shape run_guidance has always consumed. Missing input files abort the run here rather + than partway through the batch.""" + output_base = resolve(args.output_base) + if args.targets: + targets = load_targets(args.targets, output_base) + elif args.proteins or args.proteins_file: + targets = targets_from_ids(read_ids(args.proteins, args.proteins_file), output_base, args) + else: + targets = DEFAULT_TARGETS + + missing = [ + f"{t['name']}: {role}={t[role]}" + for t in targets + for role in ("structure", "density") + if not resolve(t[role]).exists() + ] + if missing and not args.allow_missing: + raise SystemExit( + "input files not found (pass --allow-missing to stage a run anyway):\n " + + "\n ".join(missing) + ) + return targets + + +def read_ids(inline: str | None, path: str | None) -> list[str]: + """PDB IDs from a comma-separated --proteins value and/or a --proteins-file (one per + line, '#' comments and blank lines skipped).""" + ids = [i.strip() for i in (inline or "").split(",") if i.strip()] + if path: + for line in resolve(path).read_text().splitlines(): + line = line.split("#")[0].strip() + if line: + ids.append(line) + return ids + + +def targets_from_ids(ids: list[str], output_base: Path, args) -> list[dict]: + """Expand PDB IDs into targets using the structure/density directory templates.""" + return [ + { + "name": fill(args.name_template, pdb), + "structure": str(resolve(args.structure_dir) / fill(args.structure_template, pdb)), + "density": str(resolve(args.density_dir) / fill(args.density_template, pdb)), + "resolution": args.resolution, + "out_dir": str(output_base / fill(args.name_template, pdb)), + } + for pdb in ids + ] + + +def fill(template: str, pdb: str) -> str: + """Substitute one PDB ID into a path/name template, in whichever case it needs.""" + return template.format(pdb=pdb, PDB=pdb.upper(), pdb_low=pdb.lower()) + + +# ================================= fan-out ================================== +# Added so one command can drive several GPUs: the parent re-runs this same script once per +# worker with --shard-index/--num-shards, then merges the per-shard summaries. + +def launch_workers(targets: list[dict], args) -> None: + """Run one child process per (GPU x --jobs-per-gpu) slot, wait, and merge summaries.""" + gpus = [g.strip() for g in args.gpus.split(",") if g.strip()] + n_workers = min(len(gpus) * args.jobs_per_gpu, len(targets)) # no empty workers + log_dir = resolve(args.output_base) / "shards" + log_dir.mkdir(parents=True, exist_ok=True) + summary_path = resolve(args.summary) + + env = dict(os.environ) | thread_env(n_workers) + print(f"fan-out: {len(targets)} targets over {n_workers} workers on GPUs {gpus} " + f"({args.jobs_per_gpu}/GPU), threads/worker={env['OMP_NUM_THREADS']}") + + children = [] + for worker in range(n_workers): + device = f"cuda:{gpus[worker % len(gpus)]}" + log_path = log_dir / f"shard_{worker}.log" + command = [sys.executable, "-u", str(Path(__file__).resolve()), *worker_argv(sys.argv[1:]), + "--shard-index", str(worker), "--num-shards", str(n_workers), + "--device", device, "--summary", str(summary_path)] + print(f" worker {worker} -> {device}, {len(targets[worker::n_workers])} targets, " + f"log {log_path}") + with log_path.open("w") as log: + children.append( + subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT, env=env) + ) + + codes = [child.wait() for child in children] + print(f"workers exited with {codes}") + merge_summaries(summary_path, n_workers) + + +def worker_argv(argv: list[str]) -> list[str]: + """This run's flags with the ones the parent sets per worker removed, in both the + '--flag value' and '--flag=value' spellings, so children never re-fan-out.""" + parent_only = ("--gpus", "--jobs-per-gpu", "--device", "--summary") + kept, skip_value = [], False + for token in argv: + if skip_value: + skip_value = False + elif token in parent_only: + skip_value = True + elif not any(token.startswith(flag + "=") for flag in parent_only): + kept.append(token) + return kept + + +def thread_env(n_workers: int) -> dict[str, str]: + """Per-worker BLAS/OMP thread counts, divided from the cgroup CPU limit. + + On a HAMi vGPU profile os.cpu_count() reports the whole node, not the cgroup quota, so + the thread pools oversubscribe and the run stalls -- that is what hung the first ens8 + attempt. Read /sys/fs/cgroup/cpu.max when it is there and fall back otherwise. + + An OMP_NUM_THREADS already in the environment wins. This function only knows its own + worker count, so two fan-outs sharing a pod (e.g. one arm on GPUs 0-2, another on GPU 3) + would each claim the whole quota; setting the variable by hand is how you split it.""" + if os.environ.get("OMP_NUM_THREADS"): + threads = os.environ["OMP_NUM_THREADS"] + return {"OMP_NUM_THREADS": threads, "MKL_NUM_THREADS": threads, + "OPENBLAS_NUM_THREADS": threads} + + cores = os.cpu_count() or n_workers + try: + quota, period = Path("/sys/fs/cgroup/cpu.max").read_text().split() + if quota != "max": + cores = float(quota) / float(period) + except (OSError, ValueError): + pass + threads = str(max(1, int(cores // n_workers))) + return {"OMP_NUM_THREADS": threads, "MKL_NUM_THREADS": threads, "OPENBLAS_NUM_THREADS": threads} + + +def shard_summary_path(path: Path, shard_index: int | None) -> Path: + """batch_summary.json -> batch_summary_shard2.json, so workers never overwrite each other.""" + if shard_index is None: + return path + return path.with_name(f"{path.stem}_shard{shard_index}{path.suffix}") + + +def merge_summaries(summary_path: Path, n_workers: int) -> None: + """Concatenate the per-shard summaries into the single --summary file.""" + merged: list[dict] = [] + for worker in range(n_workers): + shard_path = shard_summary_path(summary_path, worker) + if shard_path.exists(): + merged.extend(json.loads(shard_path.read_text())) + else: + print(f"WARNING: no summary from worker {worker} ({shard_path}) -- check its log") + write_summary(merged, summary_path) + + +def print_plan(targets: list[dict], modes: list[str], args) -> None: + """--dry-run: show what would run, where it reads from, and where it would be written.""" + n_workers = ( + min(len([g for g in args.gpus.split(",") if g.strip()]) * args.jobs_per_gpu, len(targets)) + if args.gpus else 1 + ) + print(f"PLAN: {len(targets)} targets x {len(modes)} modes = {len(targets) * len(modes)} runs, " + f"{n_workers} worker(s), model={args.model} ensemble_size={args.ensemble_size}\n" + f"modes={modes} summary -> {resolve(args.summary)}") + for index, target in enumerate(targets): + print(f"\n[{index}] {target['name']} worker={index % n_workers} " + f"res={target['resolution']}\n" + f" structure {mark(target['structure'])} {target['structure']}\n" + f" density {mark(target['density'])} {target['density']}\n" + f" out {resolve(target['out_dir'])}//refined.cif") + + +def mark(path: str) -> str: + """'ok' / 'MISSING' tag for one input path in the --dry-run plan.""" + return "ok " if resolve(path).exists() else "MISSING" + + +# ================================= plumbing ================================= + +def guidance_for_mode(mode: str, args) -> tuple[GuidanceType, dict]: + """Map a mode to its GuidanceType and the extra GuidanceConfig attributes it needs.""" + if mode == "baseline": + return GuidanceType.PURE_GUIDANCE, {"step_scaler_type": "none"} + if mode == "coord_guidance": + return GuidanceType.PURE_GUIDANCE, {"step_scaler_type": "noisespace", "step_size": args.step_size} + if mode in ("s_only", "z_only", "s_plus_z"): + which = {"s_only": "single", "z_only": "pair", "s_plus_z": "both"}[mode] + return GuidanceType.LATENT_OPT, { + "which_latent": which, + "learning_rate": args.lr, + "outer_steps": args.outer_steps, + "anchor_weight": args.anchor, + "max_grad_norm": 1.0, + "bond_length_weight": args.bond_length_weight, + } + raise ValueError(f"unknown mode {mode!r}; pick from {ALL_MODES}") + + +def build_config(target: dict, guidance_type: GuidanceType, out_dir: Path, args) -> GuidanceConfig: + config = GuidanceConfig( + protein=target["name"], + structure=str(resolve(target["structure"])), + density=str(resolve(target["density"])), + model_name=args.model, # GuidanceConfig renamed this field from `model` (merge from main) + guidance_type=guidance_type, + log_path=str(out_dir / "run.log"), + output_dir=str(out_dir), + resolution=float(target["resolution"]), + num_diffusion_steps=args.num_steps, + align_to_input=True, + ) + config.ensemble_size = args.ensemble_size # set dynamically (not a declared GuidanceConfig field) + return config + + +def parse_args() -> argparse.Namespace: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--targets", help=".csv (name,structure,density,resolution) or .json list of dicts") + ap.add_argument("--output-base", default="it_opt_scratch/targets_out", + help="base output dir for CSV targets: ///") + ap.add_argument("--modes", default="baseline,z_only,coord_guidance", + help=f"comma-separated subset of {ALL_MODES}, or 'all'") + + ids = ap.add_argument_group("target list from PDB IDs (used when --targets is not given)") + ids.add_argument("--proteins", help="comma-separated PDB IDs, e.g. 2YL0,5I09,5MC8") + ids.add_argument("--proteins-file", dest="proteins_file", + help="file of PDB IDs, one per line ('#' comments allowed)") + ids.add_argument("--structure-dir", dest="structure_dir", + default="/home/dev/test_data/processed", + help="root holding the input structures") + ids.add_argument("--structure-template", dest="structure_template", + default="{PDB}/{PDB}_single_001_density_input.cif", + help="path under --structure-dir; takes {pdb}, {PDB}, {pdb_low}") + ids.add_argument("--density-dir", dest="density_dir", default="", + help="directory holding the .ccp4 maps") + ids.add_argument("--density-template", dest="density_template", + default="{PDB}_0.5occA_0.5occB_1.00A.ccp4", + help="filename under --density-dir; takes {pdb}, {PDB}, {pdb_low}") + ids.add_argument("--name-template", dest="name_template", default="{PDB}", + help="target (and output dir) name, e.g. '{PDB}_0.5occA_0.5occB'") + ids.add_argument("--resolution", type=float, default=1.0, + help="resolution for ID-built targets") + ids.add_argument("--allow-missing", dest="allow_missing", action="store_true", + help="do not abort when a structure/density file is absent (e.g. staging " + "a run locally for pod paths)") + + fan = ap.add_argument_group("multi-GPU fan-out") + fan.add_argument("--gpus", help="comma-separated GPU indices, e.g. 0,1,2,3; one child " + "process per GPU x --jobs-per-gpu, then summaries are merged") + fan.add_argument("--jobs-per-gpu", dest="jobs_per_gpu", type=int, default=1, + help="concurrent runs per GPU; 2 needs ~2x the weights resident, " + "so check VRAM") + fan.add_argument("--shard-index", dest="shard_index", type=int, default=None, + help="set by the parent on each worker; targets[shard_index::num_shards]") + fan.add_argument("--num-shards", dest="num_shards", type=int, default=1, + help="set by the parent on each worker") + ap.add_argument("--dry-run", dest="dry_run", action="store_true", + help="print the resolved plan (paths, existence, shards) and exit") + + ap.add_argument("--model", default="protenix", choices=[m.value for m in StructurePredictor]) + ap.add_argument("--ensemble-size", dest="ensemble_size", type=int, default=4) + ap.add_argument("--num-steps", dest="num_steps", type=int, default=200) + ap.add_argument("--outer-steps", dest="outer_steps", type=int, default=2) + ap.add_argument("--lr", type=float, default=0.05) + ap.add_argument("--anchor", type=float, default=0.0, help="on-manifold anchor weight (IT-opt)") + ap.add_argument("--bond-length-weight", dest="bond_length_weight", type=float, default=0.0, + help="coordinate-space bond-geometry penalty weight (IT-opt); 0 disables") + ap.add_argument("--step-size", dest="step_size", type=float, default=0.1, help="coord-guidance DPS step") + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--device", default="cuda:0") + ap.add_argument("--checkpoint", default=None) + ap.add_argument("--summary", default=None, + help="batch summary JSON; defaults to /batch_summary.json") + + args = ap.parse_args() + if args.summary is None: # keep the summary inside the run tree it describes + args.summary = str(resolve(args.output_base) / "batch_summary.json") + return args + + +def resolve(path: str) -> Path: + """Absolute path as-is; relative path is taken relative to the repo root.""" + p = Path(path) + return p if p.is_absolute() else REPO / p + + +def load_targets(path: str, output_base: Path) -> list[dict]: + """Load targets from a .csv (name,structure,density,resolution) or a .json list of dicts. + For CSV rows, out_dir defaults to /.""" + path = Path(path) + if path.suffix.lower() == ".csv": + return [ + { + "name": row["name"], + "structure": row["structure"], + "density": row["density"], + "resolution": float(row["resolution"]), + "out_dir": str(output_base / row["name"]), + } + for row in csv.DictReader(path.read_text().splitlines()) + ] + return json.loads(path.read_text()) + + +def write_summary(summary: list[dict], path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(summary, indent=2)) + n_ok = sum(1 for s in summary if s.get("status") == "OK") + print(f"\n===== DONE: {n_ok}/{len(summary)} runs OK. Summary -> {path} =====") + + +if __name__ == "__main__": + main() diff --git a/it_opt_scratch/run_targets_simplified.py b/it_opt_scratch/run_targets_simplified.py new file mode 100644 index 00000000..5481efd3 --- /dev/null +++ b/it_opt_scratch/run_targets_simplified.py @@ -0,0 +1,186 @@ +"""Generate one conformational ensemble: one structure + one density map -> one output dir. + +This is the refinement step of a longer workflow, so it expects a model that is already sitting +in the map's frame. Starting from a sequence and a map, the usual route is: + + 1. predict sequence -> coordinates (Protenix inference, no map involved) + 2. place put that model in the map's frame (molecular replacement, e.g. Phaser; not + part of this repo) + 3. refine this script + +A deposited structure or an existing MR solution is already placed, so it starts at step 3 -- +which is what the runs in it_opt_scratch/ have done. + +That is also what --structure is for. It carries the sequence and composition (there is no +separate sequence input in this pipeline), the atom identity used for reconciliation, and the +reference frame: the density reward scores coordinates in the map's frame and aligns against +the input structure rather than the map, so a model in an unrelated frame will not score +meaningfully. + +Otherwise, read this as a plain input/output tool -- one run per invocation, no state: + + IN --structure one .cif / .pdb + --density one .ccp4 map + --resolution that map's resolution, in Angstrom + --mode which guidance arm to run (list below) + (everything else is a knob with a default; see --help) + + OUT --output-dir refined.cif trajectory/ losses.txt job_metadata.json run.log + + EXIT the pipeline's own exit code -- 0 on success, non-zero on failure, so a caller + looping over many targets can tell which ones need rerunning. + +The generation itself is not implemented here: this builds a GuidanceConfig and hands it to +the shipped `run_guidance()`, which does featurize -> reward -> scaler -> sample -> save. + +Modes: + baseline unguided sampling (pure_guidance, no step scaler) + s_only IT-opt, which_latent=single + z_only IT-opt, which_latent=pair + s_plus_z IT-opt, which_latent=both + coord_guidance shipped coordinate DPS (pure_guidance + noise-space step scaler) + +One run: + pixi run -e protenix-dev python it_opt_scratch/run_targets_simplified.py \ + --structure /home/dev/test_data/processed/2YL0/2YL0_single_001_density_input.cif \ + --density density_maps/2YL0_0.5occA_0.5occB_1.00A.ccp4 \ + --resolution 1.0 --mode s_plus_z --ensemble-size 8 --bond-length-weight 5e-5 \ + --output-dir out/2YL0_s_plus_z_ens8 --device cuda:0 + +A whole CSV, one after another on one GPU (add --skip-existing to resume a stopped sweep): + tail -n +2 targets.csv | while IFS=, read -r name structure density resolution; do + pixi run -e protenix-dev python it_opt_scratch/run_targets_simplified.py \ + --structure "$structure" --density "$density" --resolution "$resolution" \ + --mode s_plus_z --ensemble-size 8 --bond-length-weight 5e-5 \ + --output-dir "out/$name" --device cuda:0 || echo "FAILED: $name" + done + +Spreading a CSV across several GPUs is deliberately NOT this script's job -- that belongs to a +separate driver that calls this one. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import torch + +from sampleworks.utils.guidance_constants import GuidanceType, StructurePredictor +from sampleworks.utils.guidance_script_arguments import GuidanceConfig +from sampleworks.utils.guidance_script_utils import get_model_and_device, run_guidance + +MODES = ("baseline", "coord_guidance", "s_only", "z_only", "s_plus_z") +WHICH_LATENT = {"s_only": "single", "z_only": "pair", "s_plus_z": "both"} + + +def main(): + """Read the inputs, run one ensemble generation, return the pipeline's exit code.""" + args = parse_args() + + structure = Path(args.structure).expanduser() + density = Path(args.density).expanduser() + out_dir = Path(args.output_dir).expanduser() + + if not structure.exists(): + sys.exit(f"structure not found: {structure}") + if not density.exists(): + sys.exit(f"density not found: {density}") + + if args.skip_existing and (out_dir / "refined.cif").exists(): + print(f"[skip] already done: {out_dir}/refined.cif") + return 0 + out_dir.mkdir(parents=True, exist_ok=True) + + config = build_config(args, structure, density, out_dir) + device, model = get_model_and_device( + args.device, args.checkpoint, StructurePredictor(args.model) + ) + torch.manual_seed(args.seed) # fixed seed -> the same start noise across modes + + job = run_guidance(config, config.guidance_type, model, device) + print(f"[{args.mode}] {job.status} -> {out_dir}/refined.cif") + return job.exit_code + + +def build_config(args, structure: Path, density: Path, out_dir: Path) -> GuidanceConfig: + """Everything this run will do, in one place: the inputs, the knobs, and the mode. + + The mode-specific fields are assigned directly rather than declared on GuidanceConfig, + which is how the shipped per-guidance-type arg-adders set them too. + """ + # baseline and coord_guidance run the shipped pure_guidance path; the rest optimize latents. + if args.mode in ("baseline", "coord_guidance"): + guidance_type = GuidanceType.PURE_GUIDANCE + else: + guidance_type = GuidanceType.LATENT_OPT + + config = GuidanceConfig( + protein=args.name or structure.stem, + structure=str(structure), + density=str(density), + resolution=args.resolution, + output_dir=str(out_dir), + log_path=str(out_dir / "run.log"), + model_name=args.model, + guidance_type=guidance_type, + num_diffusion_steps=args.num_steps, + align_to_input=True, + ) + config.ensemble_size = args.ensemble_size + + if args.mode == "baseline": # unguided: no step scaler at all + config.step_scaler_type = "none" + elif args.mode == "coord_guidance": # gradients applied in coordinate space + config.step_scaler_type = "noisespace" + config.step_size = args.step_size + else: # s_only / z_only / s_plus_z: optimize the latents instead of the coordinates + config.which_latent = WHICH_LATENT[args.mode] + config.learning_rate = args.lr + config.outer_steps = args.outer_steps + config.anchor_weight = args.anchor + config.max_grad_norm = args.max_grad_norm + config.bond_length_weight = args.bond_length_weight + + return config + + +def parse_args(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + inputs = ap.add_argument_group("in / out") + inputs.add_argument("--structure", required=True, help="input .cif / .pdb") + inputs.add_argument("--density", required=True, help="input .ccp4 map") + inputs.add_argument("--resolution", type=float, required=True, help="map resolution, Angstrom") + inputs.add_argument("--output-dir", dest="output_dir", required=True, + help="everything this run produces is written here") + inputs.add_argument("--mode", required=True, choices=MODES) + inputs.add_argument("--name", help="label recorded in job_metadata (default: structure stem)") + inputs.add_argument("--skip-existing", dest="skip_existing", action="store_true", + help="succeed without running if refined.cif is already there") + + run = ap.add_argument_group("sampling") + run.add_argument("--ensemble-size", dest="ensemble_size", type=int, default=4) + run.add_argument("--num-steps", dest="num_steps", type=int, default=200) + run.add_argument("--seed", type=int, default=0) + run.add_argument("--model", default="protenix", choices=[m.value for m in StructurePredictor]) + run.add_argument("--device", default="cuda:0") + run.add_argument("--checkpoint", default=None) + + it = ap.add_argument_group("IT-opt modes only (s_only / z_only / s_plus_z)") + it.add_argument("--lr", type=float, default=0.05) + it.add_argument("--outer-steps", dest="outer_steps", type=int, default=2) + it.add_argument("--anchor", type=float, default=0.0, help="on-manifold anchor weight") + it.add_argument("--bond-length-weight", dest="bond_length_weight", type=float, default=0.0, + help="coordinate-space bond-geometry penalty; 0 disables") + it.add_argument("--max-grad-norm", dest="max_grad_norm", type=float, default=1.0) + + ap.add_argument("--step-size", dest="step_size", type=float, default=0.1, + help="coord_guidance mode only: DPS step size") + return ap.parse_args() + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/it_opt_scratch/score_paper_rmsd.py b/it_opt_scratch/score_paper_rmsd.py new file mode 100644 index 00000000..f9369133 --- /dev/null +++ b/it_opt_scratch/score_paper_rmsd.py @@ -0,0 +1,242 @@ +"""Score IT-opt arms with the paper's min-altloc-RMSD metric (Figure 3E). + +The paper computes, per altloc selection, the *minimum* RMSD over the predicted ensemble to +altloc A and to altloc B of the deposited reference, separately. Its purpose is diagnostic: an +ensemble can raise RSCC by fitting one altloc well, and only the pair (min-to-A, min-to-B) +reveals whether both conformations were reached. + +Scope matches the paper's RSCC protocol -- the same 3-residue max-RMSD subsegments -- and the +prediction is placed by the same global uniform-weight Kabsch used in ``score_paper_rscc.py``. +That is deliberate: aligning on the subsegment itself would let a wrong local conformation be +rotated into agreement. + +Per-altloc conformers are read with gemmi, which exposes the altloc character directly. +``load_any`` drops the altloc annotation entirely and ``parse`` keeps only the first altloc, +so neither can separate the two conformers. Atoms with a blank altloc are shared and belong to +both conformers. + +Comparator: ``occ_sweep_results/min_altloc_rmsd_results.csv`` (protenix, 0.5/0.5), already +computed by the published pipeline over the same selections. + +Usage +----- + python it_opt_scratch/score_paper_rmsd.py \ + --runs-dir it_opt_scratch/patch_tree --dir-template "{protein}" \ + --target-filename refined-patched.cif \ + --inputs-dir /home/dev/test_data \ + --selections-csv it_opt_scratch/paper_maxrmsd_selections.csv \ + --arms baseline coord_guidance s_only s_plus_z z_only \ + --out it_opt_scratch/patch_tree/itopt_paper_rmsd.csv +""" + +from __future__ import annotations + +import argparse +import csv +import re +import traceback +from pathlib import Path + +import gemmi +import numpy as np +import pandas as pd +import torch +from atomworks.io.utils.io_utils import load_any +from loguru import logger +from sampleworks.utils.atom_array_utils import remove_atoms_with_any_nan_coords + +# Same directory as this script, which is sys.path[0] when run as `python it_opt_scratch/...`. +from score_paper_rscc import align_prediction_to_reference, read_selections + +SELECTION_RE = re.compile(r"^chain\s+(\S+)\s+and\s+resi\s+(-?\d+)\s*-\s*(-?\d+)$") + + +def parse_selection(selection: str) -> tuple[str, range]: + """Split ``chain A and resi 12-14`` into ``("A", range(12, 15))``.""" + m = SELECTION_RE.match(selection.strip()) + if m is None: + raise ValueError(f"unparseable selection: {selection!r}") + chain, lo, hi = m.group(1), int(m.group(2)), int(m.group(3)) + return chain, range(lo, hi + 1) + + +def reference_conformers( + ref_path: Path, selections: list[str] +) -> dict[str, dict[str, dict[tuple[str, int, str], np.ndarray]]]: + """Per selection, the altloc-A and altloc-B conformers keyed by (chain, res_id, atom_name). + + Atoms with a blank altloc are shared between conformers, so they appear in both. + """ + st = gemmi.read_structure(str(ref_path)) + by_res: dict[tuple[str, int], list] = {} + for chain in st[0]: + for res in chain: + by_res.setdefault((chain.name, res.seqid.num), []).extend( + (a.name, a.altloc, np.array([a.pos.x, a.pos.y, a.pos.z])) for a in res + ) + + out: dict[str, dict[str, dict[tuple[str, int, str], np.ndarray]]] = {} + for sel in selections: + chain, residues = parse_selection(sel) + conformers: dict[str, dict[tuple[str, int, str], np.ndarray]] = {"A": {}, "B": {}} + for res_id in residues: + for name, altloc, xyz in by_res.get((chain, res_id), []): + alt = altloc.strip() + targets = ("A", "B") if alt == "" else (alt,) + for t in targets: + if t in conformers: + conformers[t][(chain, res_id, name)] = xyz + out[sel] = conformers + return out + + +def prediction_lookup(atom_array) -> dict[tuple[str, int, str], np.ndarray]: + """Map (chain, res_id, atom_name) -> per-model coordinates, shape [n_models, 3].""" + coords = atom_array.coord + if coords.ndim == 2: # single model -> add the model axis + coords = coords[None] + return { + (str(c), int(r), str(n)): coords[:, i] + for i, (c, r, n) in enumerate( + zip(atom_array.chain_id, atom_array.res_id, atom_array.atom_name, strict=True) + ) + } + + +def min_rmsd_over_ensemble( + conformer: dict[tuple[str, int, str], np.ndarray], + pred: dict[tuple[str, int, str], np.ndarray], +) -> tuple[float, int]: + """Minimum over ensemble members of RMSD to one reference conformer, plus atoms matched.""" + keys = [k for k in conformer if k in pred] + if not keys: + return float("nan"), 0 + ref = np.stack([conformer[k] for k in keys]) # [n_atoms, 3] + prd = np.stack([pred[k] for k in keys], axis=1) # [n_models, n_atoms, 3] + per_model = np.sqrt(((prd - ref[None]) ** 2).sum(-1).mean(-1)) # [n_models] + return float(per_model.min()), len(keys) + + +def score_protein( + protein: str, + selections: list[str], + runs_dir: Path, + inputs_dir: Path, + arms: list[str], + target_filename: str, + dir_template: str, +) -> list[dict]: + rows: list[dict] = [] + ref_path = inputs_dir / "processed" / protein / f"{protein}_single_001_density_input.cif" + + def fail(arm: str, err: str) -> None: + for sel in selections: + rows.append( + { + "protein": protein, + "arm": arm, + "selection": sel, + "min_rmsd_to_A": np.nan, + "min_rmsd_to_B": np.nan, + "n_atoms_A": 0, + "n_atoms_B": 0, + "error": err, + } + ) + + try: + conformers = reference_conformers(ref_path, selections) + # Alignment target: the same array score_paper_rscc.py aligns against, so both metrics + # place the prediction identically. + ref_atom_array = remove_atoms_with_any_nan_coords(load_any(str(ref_path))) + except Exception as e: # noqa: BLE001 + logger.error(f"{protein}: setup failed: {e}\n{traceback.format_exc()}") + for arm in arms: + fail(arm, f"setup: {e}") + return rows + + for arm in arms: + cif = runs_dir / dir_template.format(protein=protein) / arm / target_filename + if not cif.exists(): + fail(arm, f"{target_filename} missing") + continue + try: + aa = remove_atoms_with_any_nan_coords(load_any(str(cif))) + aa = align_prediction_to_reference(ref_atom_array, aa) + pred = prediction_lookup(aa) + except Exception as e: # noqa: BLE001 + logger.error(f"{protein}/{arm}: {e}\n{traceback.format_exc()}") + fail(arm, str(e)) + continue + + for sel in selections: + rmsd_a, n_a = min_rmsd_over_ensemble(conformers[sel]["A"], pred) + rmsd_b, n_b = min_rmsd_over_ensemble(conformers[sel]["B"], pred) + rows.append( + { + "protein": protein, + "arm": arm, + "selection": sel, + "min_rmsd_to_A": rmsd_a, + "min_rmsd_to_B": rmsd_b, + "n_atoms_A": n_a, + "n_atoms_B": n_b, + "error": None if n_a and n_b else "no matching atoms", + } + ) + return rows + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--runs-dir", type=Path, required=True) + p.add_argument("--inputs-dir", type=Path, required=True) + p.add_argument("--selections-csv", type=Path, required=True) + p.add_argument("--arms", nargs="+", default=["baseline", "s_only", "s_plus_z", "z_only"]) + p.add_argument("--proteins", nargs="+", default=None) + p.add_argument("--dir-template", default="{protein}_native_occ") + p.add_argument("--target-filename", default="refined.cif") + p.add_argument("--out", type=Path, required=True) + args = p.parse_args() + + selections = read_selections(args.selections_csv) + if args.proteins: + wanted = {x.upper() for x in args.proteins} + selections = {k: v for k, v in selections.items() if k in wanted} + logger.info( + f"{len(selections)} proteins, {sum(len(v) for v in selections.values())} selections, " + f"arms={args.arms}" + ) + + rows: list[dict] = [] + for i, (protein, sels) in enumerate(sorted(selections.items()), 1): + logger.info(f"[{i}/{len(selections)}] {protein} ({len(sels)} selections)") + rows.extend( + score_protein( + protein, sels, args.runs_dir, args.inputs_dir, args.arms, + args.target_filename, args.dir_template, + ) + ) + pd.DataFrame(rows).to_csv(args.out, index=False) + + df = pd.DataFrame(rows) + df.to_csv(args.out, index=False) + ok = df[df.min_rmsd_to_A.notna() & df.min_rmsd_to_B.notna()] + logger.info(f"wrote {args.out}: {len(df)} rows, {len(ok)} scored") + if not ok.empty: + nearer = ok[["min_rmsd_to_A", "min_rmsd_to_B"]].min(axis=1) + worse = ok[["min_rmsd_to_A", "min_rmsd_to_B"]].max(axis=1) + summary = pd.DataFrame( + { + "n": ok.groupby("arm").size(), + "med_nearer": nearer.groupby(ok.arm).median(), + "med_max": worse.groupby(ok.arm).median(), + "max_le_2A": worse.le(2.0).groupby(ok.arm).mean(), + "max_le_1A": worse.le(1.0).groupby(ok.arm).mean(), + } + ) + logger.info(f"\n{summary.to_string()}") + + +if __name__ == "__main__": + main() diff --git a/it_opt_scratch/score_paper_rscc.py b/it_opt_scratch/score_paper_rscc.py new file mode 100644 index 00000000..bfd67823 --- /dev/null +++ b/it_opt_scratch/score_paper_rscc.py @@ -0,0 +1,299 @@ +"""Score IT-opt arms with the paper's RSCC protocol (791 max-RMSD subsegments). + +The published scorer is ``scripts/eval/rscc_grid_search_script.py``. This driver reuses its +exact primitives -- the differentiable density forward model, the global uniform-weight Kabsch +alignment, and ``extract_tight`` at 2.0 A -- so the numbers are directly comparable to the +occupancy-sweep results. What it does *not* reuse is the Trial/ProteinConfig directory scanner, +because the IT-opt output tree is ``{PROTEIN}_native_occ/{arm}/refined.cif`` (depth 3) rather +than the ``{PROTEIN}_{occ}/{model}/{scaler}/ens{N}_gw{W}/`` depth-4 grid-search layout, and +``native_occ`` does not parse as an occupancy key. + +Selections come from the segmentation CSV (one row per protein, semicolon-joined), which encodes +the paper's rule: the contiguous 3-residue window maximising altloc A-B RMSD, or the whole +segment when it is 3 residues or shorter. + +The RSCC mask is built from reference coordinates including *all* altlocs, matching +``get_reference_structure_coords``, which unions the altloc-A-only and altloc-B-only coordinate +sets for exactly this purpose. + +Usage +----- + pixi run -e analysis python it_opt_scratch/score_paper_rscc.py \ + --runs-dir it_opt_scratch/targets_out_40 \ + --inputs-dir /home/dev/test_data \ + --selections-csv it_opt_scratch/paper_maxrmsd_selections.csv \ + --arms baseline s_only s_plus_z z_only \ + --out it_opt_scratch/patch_tree/itopt_paper_rscc.csv +""" + +from __future__ import annotations + +import argparse +import copy +import csv +import traceback +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +from atomworks.io.utils.io_utils import load_any +from loguru import logger +from sampleworks.eval.constants import DEFAULT_SELECTION_PADDING +from sampleworks.eval.metrics import rscc +from sampleworks.eval.structure_utils import extract_selection_coordinates +from sampleworks.utils.atom_array_utils import ( + filter_to_common_atoms, + remove_atoms_with_any_nan_coords, +) +from sampleworks.utils.density_utils import build_density_transformer, run_density_transformer +from sampleworks.utils.frame_transforms import ( + apply_forward_transform, + weighted_rigid_align_differentiable, +) +from sampleworks.utils.framework_utils import match_batch + +from sampleworks.core.forward_models.xray.real_space_density_deps.qfit.volume import XMap + +RESOLUTION = 1.0 + + +def read_selections(csv_path: Path) -> dict[str, list[str]]: + """Read the per-protein segmentation CSV into {PROTEIN: [selection, ...]}.""" + out: dict[str, list[str]] = {} + with open(csv_path) as fh: + for row in csv.DictReader(fh): + sels = [s.strip() for s in row["selection"].split(";") if s.strip()] + if sels: + out[row["protein"].strip().upper()] = sels + return out + + +def align_prediction_to_reference(ref_atom_array, pred_atom_array): + """Global uniform-weight Kabsch of prediction onto reference, applied to every predicted atom. + + Mirrors rscc_grid_search_script.py: the transform is fitted on the atoms common to both + structures, then applied to the whole predicted array. Alignment is deliberately global -- + fitting on the 3-residue window itself would let a wrong local conformation be rotated into + apparent agreement. + """ + try: + ref_common, pred_common = filter_to_common_atoms(ref_atom_array, pred_atom_array) + except RuntimeError: + # The prediction relabels every chain to 'A' and renumbers residues from 1, while the + # reference keeps the deposited chain id and numbering (e.g. chain 'P', res 5-234). Strict + # (chain,res,name) matching then finds nothing. Fall back to sequential per-chain matching, + # which realigns these (otherwise identical) structures. This only runs when strict matching + # raises, so the proteins that already align keep their exact matched-atom set unchanged. + ref_common, pred_common = filter_to_common_atoms( + ref_atom_array, pred_atom_array, normalize_ids=True + ) + ref_t = torch.from_numpy(ref_common.coord).float() + pred_t = torch.from_numpy(pred_common.coord).float() + ref_t = match_batch(ref_t, pred_t.shape[0]) + if ref_t.ndim != 3 or ref_t.shape[1] != pred_t.shape[1]: + raise ValueError(f"shape mismatch: ref {tuple(ref_t.shape)} vs pred {tuple(pred_t.shape)}") + + n_atoms = ref_t.shape[1] + _, transform = weighted_rigid_align_differentiable( + true_coords=pred_t, + pred_coords=ref_t, + weights=torch.ones(1, n_atoms), + mask=torch.ones(1, n_atoms), + return_transforms=True, + allow_gradients=False, + ) + moved = apply_forward_transform( + torch.from_numpy(pred_atom_array.coord), transform, rotation_only=False + ) + pred_atom_array.coord = moved.numpy() + return pred_atom_array + + +def ensure_density_annotations(atom_array): + """Add the occupancy / b_factor annotations the density forward model requires.""" + n = atom_array.coord.shape[-2] + for name, value in (("occupancy", 1.0), ("b_factor", 20.0)): + if name not in atom_array.get_annotation_categories(): + atom_array.set_annotation(name, np.full(n, value)) + return atom_array + + +def score_protein( + protein: str, + selections: list[str], + runs_dir: Path, + inputs_dir: Path, + arms: list[str], + device: torch.device, + target_filename: str = "refined.cif", + dir_template: str = "{protein}_native_occ", + maps_dir: Path | None = None, + map_template: str = "{protein}_uniform_1.00A.ccp4", +) -> list[dict]: + """Compute per-(arm, selection) RSCC for one protein.""" + rows: list[dict] = [] + maps_base = maps_dir if maps_dir is not None else inputs_dir / "density_maps" + map_path = maps_base / map_template.format(protein=protein) + ref_path = inputs_dir / "processed" / protein / f"{protein}_single_001_density_input.cif" + + def fail(arm: str, err: str) -> None: + for sel in selections: + rows.append( + { + "protein": protein, + "arm": arm, + "selection": sel, + "rscc": np.nan, + "error": err, + "base_map_path": str(map_path), + } + ) + + try: + base_xmap = XMap.fromfile(str(map_path), resolution=RESOLUTION).canonical_unit_cell() + transformer, _ = build_density_transformer(base_xmap, em_mode=False, device=device) + # Mask coords keep every altloc: load_any preserves them, whereas parse() drops all but + # the first. The mask must cover both conformations, as the published scorer does by + # unioning the A-only and B-only reference structures. + ref_all_altlocs = load_any(str(ref_path)) + sel_coords = {} + for sel in selections: + try: + coords = extract_selection_coordinates(ref_all_altlocs, sel) + except Exception as e: # noqa: BLE001 - selection syntax varies per protein + logger.warning(f"{protein}: selection {sel!r} failed: {e}") + continue + if len(coords) and np.isfinite(coords).all(): + sel_coords[sel] = coords + else: + logger.warning(f"{protein}: selection {sel!r} empty or non-finite") + # Alignment target. load_any, not parse: parse() reconciles against the + # pdbx_poly_seq_scheme that the patcher inherits from the deposited RCSB entry, and with + # atomworks 2.1.1 and no CCD mirror that turns ~93% of atoms into NaN placeholders. + # load_any reads atom_site as written, and is the reader the patcher itself uses. + ref_atom_array = remove_atoms_with_any_nan_coords(load_any(str(ref_path))) + except Exception as e: # noqa: BLE001 - per-protein setup failure should not kill the sweep + logger.error(f"{protein}: setup failed: {e}\n{traceback.format_exc()}") + for arm in arms: + fail(arm, f"setup: {e}") + return rows + + base_cache: dict[str, np.ndarray] = {} + for arm in arms: + cif = runs_dir / dir_template.format(protein=protein) / arm / target_filename + if not cif.exists(): + fail(arm, f"{target_filename} missing") + continue + try: + aa = remove_atoms_with_any_nan_coords(load_any(str(cif))) + # load_any does not populate these annotations even when the CIF carries the columns; + # the forward model requires both. The written values are occupancy 1.0 and B 20.0, + # which is also what parse() supplied to the published scorer. + aa = ensure_density_annotations(aa) + aa = align_prediction_to_reference(ref_atom_array, aa) + # One density from the whole ensemble: RSCC scores the multi-conformer model, not + # individual members. + computed = run_density_transformer(transformer, aa) + computed_xmap = copy.copy(base_xmap) + computed_xmap.array = computed.cpu().numpy() + if computed_xmap.array.shape != base_xmap.array.shape: + raise ValueError( + f"density shape {computed_xmap.array.shape} != base {base_xmap.array.shape}" + ) + except Exception as e: # noqa: BLE001 + logger.error(f"{protein}/{arm}: {e}\n{traceback.format_exc()}") + fail(arm, str(e)) + continue + + for sel, coords in sel_coords.items(): + row = { + "protein": protein, + "arm": arm, + "selection": sel, + "base_map_path": str(map_path), + "error": None, + } + try: + extracted_base = base_cache.get(sel) + if extracted_base is None: + _, extracted_base = base_xmap.extract_tight( + coords, padding=DEFAULT_SELECTION_PADDING + ) + if extracted_base is None or extracted_base.shape[0] == 0: + raise ValueError("extracted base map empty") + base_cache[sel] = extracted_base + _, extracted_computed = computed_xmap.extract_tight( + coords, padding=DEFAULT_SELECTION_PADDING + ) + if extracted_computed is None or extracted_computed.shape[0] == 0: + raise ValueError("extracted computed map empty") + row["rscc"] = rscc(extracted_base, extracted_computed) + except Exception as e: # noqa: BLE001 + row["error"] = str(e) + row["rscc"] = np.nan + rows.append(row) + return rows + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--runs-dir", type=Path, required=True) + p.add_argument("--inputs-dir", type=Path, required=True) + p.add_argument("--selections-csv", type=Path, required=True) + p.add_argument("--arms", nargs="+", default=["baseline", "s_only", "s_plus_z", "z_only"]) + p.add_argument("--proteins", nargs="+", default=None, help="Subset; default all in the CSV.") + p.add_argument("--dir-template", default="{protein}_native_occ", + help="Per-protein dir name under --runs-dir; '{protein}' is substituted.") + p.add_argument("--target-filename", default="refined.cif", + help="CIF to score in each arm dir; use refined-patched.cif after patching.") + p.add_argument("--maps-dir", type=Path, default=None, + help="Dir holding the density maps; default /density_maps.") + p.add_argument("--map-template", default="{protein}_uniform_1.00A.ccp4", + help="Map filename template; '{protein}' is substituted. Use " + "'{protein}_0.5occA_0.5occB_1.00A.ccp4' for the 0.5/0.5 occupancy maps.") + p.add_argument("--out", type=Path, required=True) + args = p.parse_args() + + selections = read_selections(args.selections_csv) + if args.proteins: + wanted = {x.upper() for x in args.proteins} + selections = {k: v for k, v in selections.items() if k in wanted} + total = sum(len(v) for v in selections.values()) + logger.info(f"{len(selections)} proteins, {total} selections, arms={args.arms}") + + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + logger.info(f"device={device}") + + rows: list[dict] = [] + for i, (protein, sels) in enumerate(sorted(selections.items()), 1): + logger.info(f"[{i}/{len(selections)}] {protein} ({len(sels)} selections)") + rows.extend( + score_protein( + protein, sels, args.runs_dir, args.inputs_dir, args.arms, device, + target_filename=args.target_filename, + dir_template=args.dir_template, + maps_dir=args.maps_dir, + map_template=args.map_template, + ) + ) + pd.DataFrame(rows).to_csv(args.out, index=False) # checkpoint after each protein + + df = pd.DataFrame(rows) + df.to_csv(args.out, index=False) + logger.info(f"wrote {args.out}: {len(df)} rows") + + ok = df[df["rscc"].notna()] + logger.info(f"scored {len(ok)}/{len(df)} rows") + if not ok.empty: + summary = ok.groupby("arm")["rscc"].agg( + n="size", + median="median", + frac_ge_08=lambda s: (s >= 0.8).mean(), + frac_ge_09=lambda s: (s >= 0.9).mean(), + ) + logger.info(f"\n{summary.to_string()}") + + +if __name__ == "__main__": + main() diff --git a/it_opt_scratch/score_paper_simplified.py b/it_opt_scratch/score_paper_simplified.py new file mode 100644 index 00000000..5d69c6a3 --- /dev/null +++ b/it_opt_scratch/score_paper_simplified.py @@ -0,0 +1,154 @@ +"""Score one run tree with both paper metrics, in one command and one CSV. + +This is a thin driver, not a new metric. RSCC comes from score_paper_rscc.py and min-altloc-RMSD +from score_paper_rmsd.py, both called unchanged -- so the numbers are identical to running those +two scripts separately. It exists so you remember one command instead of two with eight matching +flags, and get one table instead of two you have to join by hand. + + IN --runs-dir tree of // + --dir-template per-protein dir name, e.g. '{protein}_0.5occA_0.5occB' + --arms which arm sub-dirs to score + --inputs-dir holds processed/{PROTEIN}/{PROTEIN}_single_001_density_input.cif + --maps-dir holds the density maps (default /density_maps) + --map-template map filename, e.g. '{protein}_0.5occA_0.5occB_1.00A.ccp4' + --selections-csv the paper's per-protein 3-residue max-RMSD windows + + OUT --out one row per (protein, arm, selection): + rscc, min_rmsd_to_A, min_rmsd_to_B, n_atoms_A, n_atoms_B, + rscc_error, rmsd_error, base_map_path + rewritten after every protein, so a long sweep is resumable-by-eye + +Run it on the 11-protein regen tree: + pixi run -e analysis python it_opt_scratch/score_paper_simplified.py \ + --runs-dir it_opt_scratch/targets_out_11_regen_0.5occA_0.5occB_ens8 \ + --dir-template '{protein}_0.5occA_0.5occB' \ + --maps-dir it_opt_scratch/targets_out_11_regen_0.5occA_0.5occB/density_maps \ + --map-template '{protein}_0.5occA_0.5occB_1.00A.ccp4' \ + --inputs-dir /home/dev/test_data \ + --selections-csv it_opt_scratch/paper_maxrmsd_selections.csv \ + --arms s_plus_z \ + --out it_opt_scratch/regen11_scores.csv + +Each prediction is loaded and aligned once per metric rather than once in total. That is a +deliberate trade: calling the two scorers as they are keeps this file honest about the numbers, +and scoring is seconds per arm against minutes per generation run. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import pandas as pd +import torch +from loguru import logger + +# Same directory as this script, which is sys.path[0] when run as `python it_opt_scratch/...`. +import score_paper_rmsd +import score_paper_rscc + +KEY = ["protein", "arm", "selection"] + + +def main() -> None: + args = parse_args() + + selections = score_paper_rscc.read_selections(args.selections_csv) + if args.proteins: + wanted = {p.upper() for p in args.proteins} + selections = {k: v for k, v in selections.items() if k in wanted} + + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + logger.info( + f"{len(selections)} proteins, {sum(len(v) for v in selections.values())} selections, " + f"arms={args.arms}, device={device}" + ) + + scored = [] + for i, (protein, sels) in enumerate(sorted(selections.items()), 1): + logger.info(f"[{i}/{len(selections)}] {protein} ({len(sels)} selections)") + scored.append(score_one_protein(protein, sels, args, device)) + pd.concat(scored).to_csv(args.out, index=False) # checkpoint after every protein + + table = pd.concat(scored) + table.to_csv(args.out, index=False) + logger.info(f"wrote {args.out}: {len(table)} rows") + report(table) + + +def score_one_protein(protein: str, sels: list[str], args, device) -> pd.DataFrame: + """Both metrics for one protein, joined on (protein, arm, selection). + + An outer join because the two scorers can disagree on which selections are scoreable: the + RSCC side drops a selection whose reference coordinates are empty or non-finite, while the + RMSD side still emits a row for it. + """ + rscc_rows = score_paper_rscc.score_protein( + protein, sels, args.runs_dir, args.inputs_dir, args.arms, device, + target_filename=args.target_filename, + dir_template=args.dir_template, + maps_dir=args.maps_dir, + map_template=args.map_template, + ) + rmsd_rows = score_paper_rmsd.score_protein( + protein, sels, args.runs_dir, args.inputs_dir, args.arms, + args.target_filename, args.dir_template, + ) + rscc_table = pd.DataFrame(rscc_rows).rename(columns={"error": "rscc_error"}) + rmsd_table = pd.DataFrame(rmsd_rows).rename(columns={"error": "rmsd_error"}) + return pd.merge(rscc_table, rmsd_table, on=KEY, how="outer") + + +def report(table: pd.DataFrame) -> None: + """Per-arm summary, in the same terms the two original scorers print.""" + rscc_ok = table[table["rscc"].notna()] + if not rscc_ok.empty: + summary = rscc_ok.groupby("arm")["rscc"].agg( + n="size", + median="median", + frac_ge_08=lambda s: (s >= 0.8).mean(), + frac_ge_09=lambda s: (s >= 0.9).mean(), + ) + logger.info(f"RSCC\n{summary.to_string()}") + + both = ["min_rmsd_to_A", "min_rmsd_to_B"] + rmsd_ok = table[table[both].notna().all(axis=1)] + if not rmsd_ok.empty: + nearer = rmsd_ok[both].min(axis=1) # the altloc the ensemble reached + worse = rmsd_ok[both].max(axis=1) # the one it had to also reach to score well + summary = pd.DataFrame( + { + "n": rmsd_ok.groupby("arm").size(), + "med_nearer": nearer.groupby(rmsd_ok.arm).median(), + "med_max": worse.groupby(rmsd_ok.arm).median(), + "max_le_2A": worse.le(2.0).groupby(rmsd_ok.arm).mean(), + "max_le_1A": worse.le(1.0).groupby(rmsd_ok.arm).mean(), + } + ) + logger.info(f"min-altloc-RMSD\n{summary.to_string()}") + + +def parse_args() -> argparse.Namespace: + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--runs-dir", type=Path, required=True) + ap.add_argument("--inputs-dir", type=Path, required=True) + ap.add_argument("--selections-csv", type=Path, required=True) + ap.add_argument("--out", type=Path, required=True) + ap.add_argument("--arms", nargs="+", default=["baseline", "s_only", "s_plus_z", "z_only"]) + ap.add_argument("--proteins", nargs="+", default=None, help="subset; default all in the CSV") + ap.add_argument("--dir-template", default="{protein}_native_occ", + help="per-protein dir under --runs-dir; '{protein}' is substituted") + ap.add_argument("--target-filename", default="refined.cif", + help="CIF to score in each arm dir; refined-patched.cif after patching") + ap.add_argument("--maps-dir", type=Path, default=None, + help="dir holding the maps; default /density_maps") + ap.add_argument("--map-template", default="{protein}_uniform_1.00A.ccp4", + help="map filename; use '{protein}_0.5occA_0.5occB_1.00A.ccp4' for the " + "paper's 0.5/0.5 occupancy maps") + return ap.parse_args() + + +if __name__ == "__main__": + main() diff --git a/it_opt_scratch/score_rmsd_simplified.py b/it_opt_scratch/score_rmsd_simplified.py new file mode 100644 index 00000000..25a56c2e --- /dev/null +++ b/it_opt_scratch/score_rmsd_simplified.py @@ -0,0 +1,259 @@ +"""Min-altloc-RMSD for one prediction against a two-conformer reference. Self-contained. + +Answers: did the ensemble actually reach BOTH conformations the crystal shows, or just one? + +For each of the paper's 3-residue windows, this reports two numbers -- the closest any ensemble +member gets to altloc A, and the closest any member gets to altloc B, measured separately. The +pair is the point. RSCC can look good when an ensemble fits one conformer well and ignores the +other; only the pair reveals whether both were found. + + IN --prediction refined.cif from a run (multi-model = the ensemble) + --reference the deposited .cif, which must contain altloc A and B + --selections-csv the paper's per-protein windows + --protein which row of that CSV to use, e.g. 2YL0 + + OUT --out one row per window: + min_rmsd_to_A closest member to conformer A, Angstrom + min_rmsd_to_B closest member to conformer B, Angstrom + n_atoms_A/_B atoms actually compared (0 = nothing matched) + + EXIT 0 if windows scored, 1 if none did. + +How to read it: both numbers small means the ensemble captured both conformers. One small and +one large means it collapsed onto a single conformation -- so max(A, B) is the honest per-window +summary, and the median of that across windows is the honest per-arm summary. + +Two details that decide the numbers: + +* Altlocs are read with gemmi, which exposes the altloc character directly. `load_any` drops the + annotation and `parse` keeps only the first altloc, so neither can separate the conformers. + Atoms with a blank altloc are shared and belong to both. +* The prediction is placed by a GLOBAL uniform-weight Kabsch onto the reference, never fitted on + the window itself -- fitting locally would let a wrong conformation be rotated into agreement. + +Run it: + pixi run -e analysis python it_opt_scratch/score_rmsd_simplified.py \ + --prediction out/2YL0/refined-patched.cif \ + --reference /home/dev/test_data/processed/2YL0/2YL0_single_001_density_input.cif \ + --selections-csv it_opt_scratch/paper_maxrmsd_selections.csv \ + --protein 2YL0 --out 2YL0_rmsd.csv + +Use the patched CIF (scripts/patch_output_cif_files.py) when you have one. Atoms are matched by +(chain, residue number, atom name), and for some proteins the raw prediction relabels chains and +renumbers from 1 while the reference keeps its deposited numbering -- then nothing matches and +every window comes back with n_atoms 0. This script stops with that message rather than writing +a file full of blanks. +""" + +from __future__ import annotations + +import argparse +import csv +import re +import sys +from pathlib import Path + +import gemmi +import numpy as np +import pandas as pd +import torch +from atomworks.io.utils.io_utils import load_any +from loguru import logger +from sampleworks.utils.atom_array_utils import ( + filter_to_common_atoms, + remove_atoms_with_any_nan_coords, +) +from sampleworks.utils.frame_transforms import ( + apply_forward_transform, + weighted_rigid_align_differentiable, +) +from sampleworks.utils.framework_utils import match_batch + +SELECTION_RE = re.compile(r"^chain\s+(\S+)\s+and\s+resi\s+(-?\d+)\s*-\s*(-?\d+)$") + + +def main() -> int: + args = parse_args() + + windows = read_windows(args.selections_csv, args.protein) + if not windows: + sys.exit(f"no selections for {args.protein} in {args.selections_csv}") + + # The two reference conformers per window, straight from the altloc characters. + conformers = reference_conformers(args.reference, windows) + + prediction = remove_atoms_with_any_nan_coords(load_any(str(args.prediction))) + alignment_target = remove_atoms_with_any_nan_coords(load_any(str(args.reference))) + prediction = align_to_reference(alignment_target, prediction) + predicted_atoms = prediction_lookup(prediction) + + rows = [] + for window in windows: + rmsd_a, n_a = min_rmsd_over_ensemble(conformers[window]["A"], predicted_atoms) + rmsd_b, n_b = min_rmsd_over_ensemble(conformers[window]["B"], predicted_atoms) + rows.append( + { + "protein": args.protein, + "selection": window, + "min_rmsd_to_A": rmsd_a, + "min_rmsd_to_B": rmsd_b, + "n_atoms_A": n_a, + "n_atoms_B": n_b, + } + ) + + table = pd.DataFrame(rows) + table.to_csv(args.out, index=False) + + scored = table[table[["n_atoms_A", "n_atoms_B"]].gt(0).all(axis=1)] + if scored.empty: + logger.error( + "no atoms matched between prediction and reference. Atoms are matched by " + "(chain, residue number, atom name) -- the prediction is probably renumbered. " + "Patch it first: scripts/patch_output_cif_files.py" + ) + return 1 + + worse = scored[["min_rmsd_to_A", "min_rmsd_to_B"]].max(axis=1) + logger.info(f"wrote {args.out}: {len(scored)}/{len(table)} windows scored") + logger.info( + f"median max(A,B) {worse.median():.3f} A " + f"windows with both conformers within 1 A: {worse.le(1.0).mean():.0%}" + ) + return 0 + + +def read_windows(csv_path: Path, protein: str) -> list[str]: + """The paper's windows for one protein: the semicolon-joined `selection` column of its row.""" + with open(csv_path) as fh: + for row in csv.DictReader(fh): + if row["protein"].strip().upper() == protein.upper(): + return [s.strip() for s in row["selection"].split(";") if s.strip()] + return [] + + +def parse_selection(selection: str) -> tuple[str, range]: + """Split ``chain A and resi 12-14`` into ``("A", range(12, 15))``.""" + match = SELECTION_RE.match(selection.strip()) + if match is None: + raise ValueError(f"unparseable selection: {selection!r}") + chain, low, high = match.group(1), int(match.group(2)), int(match.group(3)) + return chain, range(low, high + 1) + + +def reference_conformers(reference_path: Path, windows: list[str]) -> dict: + """Per window, the altloc-A and altloc-B conformers keyed by (chain, residue, atom name). + + Read with gemmi because it exposes the altloc character per atom. An atom with a blank + altloc is shared between the two conformers, so it is placed in both. + """ + structure = gemmi.read_structure(str(reference_path)) + atoms_by_residue: dict[tuple[str, int], list] = {} + for chain in structure[0]: + for residue in chain: + atoms_by_residue.setdefault((chain.name, residue.seqid.num), []).extend( + (a.name, a.altloc, np.array([a.pos.x, a.pos.y, a.pos.z])) for a in residue + ) + + out = {} + for window in windows: + chain_name, residues = parse_selection(window) + conformers = {"A": {}, "B": {}} + for residue_id in residues: + for atom_name, altloc, xyz in atoms_by_residue.get((chain_name, residue_id), []): + shared = altloc.strip() == "" + for label in ("A", "B") if shared else (altloc.strip(),): + if label in conformers: + conformers[label][(chain_name, residue_id, atom_name)] = xyz + out[window] = conformers + return out + + +def prediction_lookup(atom_array) -> dict: + """Map (chain, residue, atom name) -> that atom's coordinates in every ensemble member. + + Values have shape [n_models, 3]; a single-model file is given a leading axis of 1. + """ + coords = atom_array.coord + if coords.ndim == 2: + coords = coords[None] + return { + (str(chain), int(residue), str(name)): coords[:, i] + for i, (chain, residue, name) in enumerate( + zip(atom_array.chain_id, atom_array.res_id, atom_array.atom_name, strict=True) + ) + } + + +def min_rmsd_over_ensemble(conformer: dict, predicted_atoms: dict) -> tuple[float, int]: + """Closest any single ensemble member gets to one reference conformer. + + The minimum is over members, not an average: the question is whether ANY member found this + conformation, not whether the ensemble is centred on it. + """ + shared_keys = [key for key in conformer if key in predicted_atoms] + if not shared_keys: + return float("nan"), 0 + reference = np.stack([conformer[key] for key in shared_keys]) # [n_atoms, 3] + predicted = np.stack([predicted_atoms[key] for key in shared_keys], axis=1) # [n_models, n, 3] + per_member = np.sqrt(((predicted - reference[None]) ** 2).sum(-1).mean(-1)) # [n_models] + return float(per_member.min()), len(shared_keys) + + +def align_to_reference(reference, prediction): + """Global uniform-weight Kabsch of prediction onto reference, applied to every atom. + + Alignment is deliberately GLOBAL: fitting on the 3-residue window itself would let a wrong + local conformation be rotated into apparent agreement. + """ + try: + reference_common, prediction_common = filter_to_common_atoms(reference, prediction) + except RuntimeError: + # Some predictions relabel every chain to 'A' and renumber residues from 1 while the + # reference keeps its deposited chain id and numbering, so strict (chain, res, name) + # matching finds nothing. Sequential per-chain matching realigns these otherwise + # identical structures. Only runs when strict matching raises. + reference_common, prediction_common = filter_to_common_atoms( + reference, prediction, normalize_ids=True + ) + + reference_coords = torch.from_numpy(reference_common.coord).float() + prediction_coords = torch.from_numpy(prediction_common.coord).float() + reference_coords = match_batch(reference_coords, prediction_coords.shape[0]) + if reference_coords.ndim != 3 or reference_coords.shape[1] != prediction_coords.shape[1]: + raise ValueError( + f"shape mismatch: reference {tuple(reference_coords.shape)} " + f"vs prediction {tuple(prediction_coords.shape)}" + ) + + n_atoms = reference_coords.shape[1] + _, transform = weighted_rigid_align_differentiable( + true_coords=prediction_coords, + pred_coords=reference_coords, + weights=torch.ones(1, n_atoms), + mask=torch.ones(1, n_atoms), + return_transforms=True, + allow_gradients=False, + ) + moved = apply_forward_transform( + torch.from_numpy(prediction.coord), transform, rotation_only=False + ) + prediction.coord = moved.numpy() + return prediction + + +def parse_args() -> argparse.Namespace: + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--prediction", type=Path, required=True, help="refined.cif to score") + ap.add_argument("--reference", type=Path, required=True, + help="deposited .cif containing altlocs A and B") + ap.add_argument("--selections-csv", type=Path, required=True, help="the paper's windows") + ap.add_argument("--protein", required=True, help="which protein's windows, e.g. 2YL0") + ap.add_argument("--out", type=Path, required=True, help="output CSV") + return ap.parse_args() + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/it_opt_scratch/score_rscc_simplified.py b/it_opt_scratch/score_rscc_simplified.py new file mode 100644 index 00000000..0b93179e --- /dev/null +++ b/it_opt_scratch/score_rscc_simplified.py @@ -0,0 +1,240 @@ +"""RSCC for one prediction against one density map. Self-contained: just run it. + +Scores how well a predicted ensemble explains observed density, over the paper's 3-residue +windows. One prediction per invocation, so a sweep is a plain loop in the caller. + + IN --prediction refined.cif from a run (an ensemble is expected; see below) + --reference the deposited .cif -- the alignment target, and the source of the + coordinates that define each window + --map the observed .ccp4 + --resolution that map's resolution, in Angstrom + --selections-csv the paper's per-protein windows (semicolon-joined column) + --protein which row of that CSV to use, e.g. 2YL0 + + OUT --out one row per window: protein, selection, rscc, error + + EXIT 0 if every window scored, 1 if any errored. + +The whole ensemble becomes a SINGLE calculated density, compared to the map once. That is the +point of the metric: RSCC scores the multi-conformer model, not its members individually, which +is how an ensemble can explain density that no single conformer can. + +Run it: + pixi run -e analysis python it_opt_scratch/score_rscc_simplified.py \ + --prediction out/2YL0/refined.cif \ + --reference /home/dev/test_data/processed/2YL0/2YL0_single_001_density_input.cif \ + --map density_maps/2YL0_0.5occA_0.5occB_1.00A.ccp4 \ + --resolution 1.0 \ + --selections-csv it_opt_scratch/paper_maxrmsd_selections.csv \ + --protein 2YL0 --out 2YL0_rscc.csv + +Every number-determining step is written out in this file on purpose -- the alignment, the +occupancy/B values, the window cropping. They match scripts/eval/rscc_grid_search_script.py and +score_paper_rscc.py; changing any of them changes the metric, so the comments say why each is +the way it is. +""" + +from __future__ import annotations + +import argparse +import copy +import csv +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +from atomworks.io.utils.io_utils import load_any +from loguru import logger +from sampleworks.eval.constants import DEFAULT_SELECTION_PADDING +from sampleworks.eval.metrics import rscc +from sampleworks.eval.structure_utils import extract_selection_coordinates +from sampleworks.utils.atom_array_utils import ( + filter_to_common_atoms, + remove_atoms_with_any_nan_coords, +) +from sampleworks.utils.density_utils import build_density_transformer, run_density_transformer +from sampleworks.utils.frame_transforms import ( + apply_forward_transform, + weighted_rigid_align_differentiable, +) +from sampleworks.utils.framework_utils import match_batch + +from sampleworks.core.forward_models.xray.real_space_density_deps.qfit.volume import XMap + + +def main() -> int: + args = parse_args() + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + + windows = read_windows(args.selections_csv, args.protein) + if not windows: + sys.exit(f"no selections for {args.protein} in {args.selections_csv}") + + require_unit_cell(args.prediction) + + # The observed map, plus the forward model that turns coordinates into a calculated one. + observed = XMap.fromfile(str(args.map), resolution=args.resolution).canonical_unit_cell() + transformer, _ = build_density_transformer(observed, em_mode=False, device=device) + + # Two reads of the reference, for two different jobs: + # - windows are defined over ALL altlocs, so the mask covers both conformations; + # - alignment needs NaN-free coordinates. + # load_any, not parse: parse() reconciles against pdbx_poly_seq_scheme and, without a CCD + # mirror, turns most atoms into NaN placeholders. load_any reads atom_site as written. + reference_all_altlocs = load_any(str(args.reference)) + alignment_target = remove_atoms_with_any_nan_coords(load_any(str(args.reference))) + + window_coords = {} + for selection in windows: + coords = extract_selection_coordinates(reference_all_altlocs, selection) + if len(coords) and np.isfinite(coords).all(): + window_coords[selection] = coords + else: + logger.warning(f"{args.protein}: window {selection!r} empty or non-finite, skipping") + + prediction = remove_atoms_with_any_nan_coords(load_any(str(args.prediction))) + prediction = add_density_annotations(prediction) + prediction = align_to_reference(alignment_target, prediction) + + calculated = copy.copy(observed) + calculated.array = run_density_transformer(transformer, prediction).cpu().numpy() + if calculated.array.shape != observed.array.shape: + sys.exit(f"density shape {calculated.array.shape} != observed {observed.array.shape}") + + rows = [ + score_window(args.protein, selection, coords, observed, calculated) + for selection, coords in window_coords.items() + ] + + table = pd.DataFrame(rows) + table.to_csv(args.out, index=False) + scored = int(table["rscc"].notna().sum()) + logger.info(f"wrote {args.out}: {scored}/{len(table)} windows scored") + if scored: + logger.info(f"median rscc {table['rscc'].median():.4f}") + return 0 if scored == len(table) else 1 + + +def require_unit_cell(cif_path: Path) -> None: + """Stop early if the CIF has no crystallographic header. + + X-ray density is calculated in the full crystal frame, so the forward model needs the unit + cell and space group. A freshly generated refined.cif carries neither -- the header is added + afterwards by scripts/patch_output_cif_files.py, which fetches it from the PDB entry. Without + this check, scoring an unpatched file yields numbers that look reasonable and are not. + """ + text = cif_path.read_text() + if "_cell." not in text: + sys.exit( + f"{cif_path} has no unit cell -- it looks like an unpatched refined.cif.\n" + "Add the crystallographic header first:\n" + " python scripts/patch_output_cif_files.py --input-dir " + "--grid-search-input-dir --rcsb-pattern '/([0-9][A-Za-z0-9]{3})'\n" + "then score the resulting refined-patched.cif." + ) + + +def read_windows(csv_path: Path, protein: str) -> list[str]: + """The paper's windows for one protein: the semicolon-joined `selection` column of its row.""" + with open(csv_path) as fh: + for row in csv.DictReader(fh): + if row["protein"].strip().upper() == protein.upper(): + return [s.strip() for s in row["selection"].split(";") if s.strip()] + return [] + + +def align_to_reference(reference, prediction): + """Global uniform-weight Kabsch of prediction onto reference, applied to every atom. + + Alignment is deliberately GLOBAL: fitting on the 3-residue window itself would let a wrong + local conformation be rotated into apparent agreement. The transform is fitted on the atoms + common to both structures, then applied to the whole predicted array. + """ + try: + reference_common, prediction_common = filter_to_common_atoms(reference, prediction) + except RuntimeError: + # The prediction relabels every chain to 'A' and renumbers residues from 1, while the + # reference keeps the deposited chain id and numbering (e.g. chain 'P', res 5-234). + # Strict (chain, res, name) matching then finds nothing, so fall back to sequential + # per-chain matching. This only runs when strict matching raises, so proteins that + # already align keep their exact matched-atom set. + reference_common, prediction_common = filter_to_common_atoms( + reference, prediction, normalize_ids=True + ) + + reference_coords = torch.from_numpy(reference_common.coord).float() + prediction_coords = torch.from_numpy(prediction_common.coord).float() + reference_coords = match_batch(reference_coords, prediction_coords.shape[0]) + if reference_coords.ndim != 3 or reference_coords.shape[1] != prediction_coords.shape[1]: + raise ValueError( + f"shape mismatch: reference {tuple(reference_coords.shape)} " + f"vs prediction {tuple(prediction_coords.shape)}" + ) + + n_atoms = reference_coords.shape[1] + _, transform = weighted_rigid_align_differentiable( + true_coords=prediction_coords, + pred_coords=reference_coords, + weights=torch.ones(1, n_atoms), + mask=torch.ones(1, n_atoms), + return_transforms=True, + allow_gradients=False, + ) + moved = apply_forward_transform( + torch.from_numpy(prediction.coord), transform, rotation_only=False + ) + prediction.coord = moved.numpy() + return prediction + + +def add_density_annotations(atom_array): + """Add the occupancy / b_factor annotations the density forward model requires. + + load_any does not populate these even when the CIF carries the columns. The values are + occupancy 1.0 and B 20.0, which is what parse() supplied to the published scorer. + """ + n_atoms = atom_array.coord.shape[-2] + for name, value in (("occupancy", 1.0), ("b_factor", 20.0)): + if name not in atom_array.get_annotation_categories(): + atom_array.set_annotation(name, np.full(n_atoms, value)) + return atom_array + + +def score_window(protein: str, selection: str, coords, observed, calculated) -> dict: + """RSCC over one window: crop both maps to it, then correlate. + + Cropping is `extract_tight` around the window's reference coordinates, so observed and + calculated are compared on exactly the same voxels. + """ + row = {"protein": protein, "selection": selection, "rscc": None, "error": None} + try: + _, observed_crop = observed.extract_tight(coords, padding=DEFAULT_SELECTION_PADDING) + _, calculated_crop = calculated.extract_tight(coords, padding=DEFAULT_SELECTION_PADDING) + if observed_crop is None or observed_crop.shape[0] == 0: + raise ValueError("observed crop empty") + if calculated_crop is None or calculated_crop.shape[0] == 0: + raise ValueError("calculated crop empty") + row["rscc"] = rscc(observed_crop, calculated_crop) + except Exception as err: # noqa: BLE001 - one bad window should not lose the others + row["error"] = str(err) + return row + + +def parse_args() -> argparse.Namespace: + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--prediction", type=Path, required=True, help="refined.cif to score") + ap.add_argument("--reference", type=Path, required=True, help="deposited .cif") + ap.add_argument("--map", type=Path, required=True, help="observed .ccp4") + ap.add_argument("--resolution", type=float, default=1.0, help="map resolution, Angstrom") + ap.add_argument("--selections-csv", type=Path, required=True, help="the paper's windows") + ap.add_argument("--protein", required=True, help="which protein's windows, e.g. 2YL0") + ap.add_argument("--out", type=Path, required=True, help="output CSV") + return ap.parse_args() + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/it_opt_scratch/slurm_ensemble_and_score.sh b/it_opt_scratch/slurm_ensemble_and_score.sh new file mode 100755 index 00000000..843db970 --- /dev/null +++ b/it_opt_scratch/slurm_ensemble_and_score.sh @@ -0,0 +1,148 @@ +#!/usr/bin/env bash +# Generate an ensemble for one target and score it, as a SLURM array task. +# +# IN a CSV of targets: name,structure,density,resolution (one row per target) +# OUT //refined-patched.cif the ensemble +# //rscc.csv per-window RSCC +# //rmsd.csv per-window min-altloc-RMSD +# /rscc_all.csv all targets, after `aggregate` +# /rmsd_all.csv all targets, after `aggregate` +# +# Each array task handles exactly one row, so a failed target fails only its own task and can be +# requeued on its own. Nothing is appended to a shared file until `aggregate`, so tasks never +# race each other. +# +# Stages (argument 1, default `all`): +# generate GPU. sample the ensemble, then add the crystallographic header +# score CPU (GPU optional). RSCC + min-altloc-RMSD into two per-target CSVs +# all both, in one task +# aggregate concatenate every per-target CSV into the two final ones (run once, at the end) +# +# Submit as two dependent arrays so the CPU stage does not sit on a GPU: +# +# N=$(( $(wc -l < targets.csv) - 1 )) +# gen=$(sbatch --parsable --array=1-$N --gres=gpu:1 --cpus-per-task=6 --mem=64G \ +# --time=2:00:00 it_opt_scratch/slurm_ensemble_and_score.sh generate) +# scr=$(sbatch --parsable --array=1-$N --dependency=aftercorr:$gen --cpus-per-task=4 --mem=32G \ +# --time=1:00:00 it_opt_scratch/slurm_ensemble_and_score.sh score) +# sbatch --dependency=afterany:$scr --cpus-per-task=1 --mem=4G \ +# it_opt_scratch/slurm_ensemble_and_score.sh aggregate +# +# Or in one array (simpler, wastes the GPU during scoring): +# sbatch --array=1-$N --gres=gpu:1 it_opt_scratch/slurm_ensemble_and_score.sh all +# +# Without SLURM it runs row 1 unless you set TASK_ID: +# TASK_ID=3 bash it_opt_scratch/slurm_ensemble_and_score.sh all +# +# BEFORE SUBMITTING: the header-patching step downloads the PDB entry from RCSB. Compute nodes +# on many clusters have no outbound network. Warm the cache on a login node first by running the +# `generate` stage for every target there, or pre-populate ~/.sampleworks/rcsb -- otherwise every +# task fails at step 2. See the note in check_prerequisites below. + +set -euo pipefail + +# ------------------------------- settings ------------------------------- +# Override any of these by exporting them before sbatch, e.g. `MODE=z_only sbatch ...`. +REPO="${REPO:-/home/dev/workspace}" +TARGETS="${TARGETS:-$REPO/it_opt_scratch/targets.csv}" +OUTPUT_ROOT="${OUTPUT_ROOT:-$REPO/it_opt_scratch/slurm_out}" +SELECTIONS="${SELECTIONS:-$REPO/it_opt_scratch/paper_maxrmsd_selections.csv}" +PROCESSED_DIR="${PROCESSED_DIR:-/home/dev/test_data/processed}" + +MODE="${MODE:-s_plus_z}" +ENSEMBLE_SIZE="${ENSEMBLE_SIZE:-8}" +BOND_LENGTH_WEIGHT="${BOND_LENGTH_WEIGHT:-5e-5}" +NUM_STEPS="${NUM_STEPS:-200}" + +GEN_ENV="${GEN_ENV:-protenix-dev}" # has protenix + torch +ANALYSIS_ENV="${ANALYSIS_ENV:-analysis}" # has gemmi + the density tooling +# ------------------------------------------------------------------------ + +STAGE="${1:-all}" + +# aggregate only reads OUTPUT_ROOT, so it runs anywhere -- no repo, no pixi, no GPU. +if [[ "$STAGE" == "aggregate" ]]; then + # Keep the header from the first file, skip it in the rest. + for metric in rscc rmsd; do + out="$OUTPUT_ROOT/${metric}_all.csv" + first=1 + : > "$out" + for f in "$OUTPUT_ROOT"/*/"${metric}.csv"; do + [[ -e "$f" ]] || continue + if [[ $first == 1 ]]; then cat "$f"; first=0; else tail -n +2 "$f"; fi + done >> "$out" + echo "[aggregate] $(( $(wc -l < "$out") - 1 )) rows -> $out" + done + exit 0 +fi + +cd "$REPO" + +# ---- which target is this task? Row 1 of the CSV is the header, so add one. ---- +TASK_ID="${SLURM_ARRAY_TASK_ID:-${TASK_ID:-1}}" +row=$(( TASK_ID + 1 )) +line=$(sed -n "${row}p" "$TARGETS") +[[ -n "$line" ]] || { echo "no row $row in $TARGETS" >&2; exit 1; } +IFS=, read -r NAME STRUCTURE DENSITY RESOLUTION <<< "$line" + +# The window list and the PDB header lookup are keyed by the bare 4-character PDB id, while the +# target name usually carries a suffix describing the map (e.g. 2YL0_0.5occA_0.5occB). +PDB="${NAME%%_*}" +OUT_DIR="$OUTPUT_ROOT/$NAME" +mkdir -p "$OUT_DIR" + +echo "[task $TASK_ID] $NAME (pdb $PDB) stage=$STAGE mode=$MODE ens=$ENSEMBLE_SIZE" +echo " structure $STRUCTURE" +echo " density $DENSITY @ ${RESOLUTION} A" +echo " out $OUT_DIR" + +if [[ "$STAGE" == "generate" || "$STAGE" == "all" ]]; then + echo "[1/4] sampling the ensemble" + pixi run -e "$GEN_ENV" python -u it_opt_scratch/run_targets_simplified.py \ + --structure "$STRUCTURE" \ + --density "$DENSITY" \ + --resolution "$RESOLUTION" \ + --mode "$MODE" \ + --ensemble-size "$ENSEMBLE_SIZE" \ + --num-steps "$NUM_STEPS" \ + --bond-length-weight "$BOND_LENGTH_WEIGHT" \ + --name "$NAME" \ + --skip-existing \ + --output-dir "$OUT_DIR" + + # RSCC needs the unit cell and space group, which sampling does not write. This fetches them + # from the PDB entry (network!) and writes refined-patched.cif alongside refined.cif. + echo "[2/4] adding the crystallographic header" + pixi run -e "$ANALYSIS_ENV" python scripts/patch_output_cif_files.py \ + --input-dir "$OUT_DIR" \ + --depth 1 \ + --cif-pattern refined.cif \ + --rcsb-pattern "($PDB)" \ + --grid-search-input-dir "$PROCESSED_DIR" +fi + +if [[ "$STAGE" == "score" || "$STAGE" == "all" ]]; then + SCORED_CIF="$OUT_DIR/refined-patched.cif" + [[ -f "$SCORED_CIF" ]] || { echo "missing $SCORED_CIF -- did the generate stage finish?" >&2; exit 1; } + REFERENCE="$PROCESSED_DIR/$PDB/${PDB}_single_001_density_input.cif" + + echo "[3/4] RSCC" + pixi run -e "$ANALYSIS_ENV" python it_opt_scratch/score_rscc_simplified.py \ + --prediction "$SCORED_CIF" \ + --reference "$REFERENCE" \ + --map "$DENSITY" \ + --resolution "$RESOLUTION" \ + --selections-csv "$SELECTIONS" \ + --protein "$PDB" \ + --out "$OUT_DIR/rscc.csv" + + echo "[4/4] min-altloc-RMSD" + pixi run -e "$ANALYSIS_ENV" python it_opt_scratch/score_rmsd_simplified.py \ + --prediction "$SCORED_CIF" \ + --reference "$REFERENCE" \ + --selections-csv "$SELECTIONS" \ + --protein "$PDB" \ + --out "$OUT_DIR/rmsd.csv" +fi + +echo "[task $TASK_ID] $NAME done" From db905bfca202563dfa3bb66d705e0bd3d01c3836 Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Tue, 4 Aug 2026 12:01:44 -0700 Subject: [PATCH 06/15] feat(grid-search): run latent_opt (IT-opt) with 2 jobs per GPU 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) --- experiments/it_opt_1gpu.toml | 41 +++++++++++++++++ experiments/it_opt_4gpu.toml | 44 ++++++++++++++++++ run_grid_search.py | 45 ++++++++++++++++--- .../utils/guidance_script_arguments.py | 37 ++++++++++++--- 4 files changed, 156 insertions(+), 11 deletions(-) create mode 100644 experiments/it_opt_1gpu.toml create mode 100644 experiments/it_opt_4gpu.toml diff --git a/experiments/it_opt_1gpu.toml b/experiments/it_opt_1gpu.toml new file mode 100644 index 00000000..3722d33e --- /dev/null +++ b/experiments/it_opt_1gpu.toml @@ -0,0 +1,41 @@ +description = "IT-opt s_plus_z on 1 GPU: all proteins run sequentially in one worker." + +# Run with: +# sampleworks-runs --preset it_opt_1gpu +# +# One job per protein, one worker, so 40 proteins is 40 jobs run back to back. The sibling +# it_opt_4gpu.toml is this same configuration with gpu_count = 4, which splits those jobs over +# four workers. +# +# The settings below are the defaults of it_opt_scratch/slurm_ensemble_and_score.sh, so a run +# here is comparable to the recorded trees. To sample a different mode, change which-latent +# (single = s, pair = z, both = s+z); for the unguided or coordinate-DPS references use +# scalers = "pure_guidance" with step-scaler-type = "none" or "noisespace" instead. + +[defaults] +DATA_DIR = "/data/inputs" +RESULTS_DIR = "/data/results/it_opt_1gpu" +PROTEINS_CSV = "${DATA_DIR}/proteins.csv" + +# ensemble-sizes and gradient-weights MUST stay single values. Their defaults are "1 2 4 8" and +# "0.01 0.1 0.2", which would expand one job per protein into twelve. gradient-weights only +# feeds the step scaler, which latent_opt does not use, so its value is arbitrary. +# align-to-input matches the scratch runner, which set it unconditionally; the grid defaults it +# off. bond-length-weight is named explicitly because the scratch runner defaulted it to 0.0. +[shared_args] +proteins = "${PROTEINS_CSV}" +model = "protenix" +scalers = "latent_opt" +which-latent = "both" +ensemble-sizes = "8" +gradient-weights = "0.0" +num-diffusion-steps = 200 +bond-length-weight = 5e-5 +align-to-input = true + +[[jobs]] +name = "s_plus_z" +env = "protenix" +gpu_count = 1 +output_subdir = "s_plus_z" +args = {} diff --git a/experiments/it_opt_4gpu.toml b/experiments/it_opt_4gpu.toml new file mode 100644 index 00000000..d8694685 --- /dev/null +++ b/experiments/it_opt_4gpu.toml @@ -0,0 +1,44 @@ +description = "IT-opt s_plus_z on 4 GPUs: proteins split across four workers, one per GPU." + +# Run with: +# sampleworks-runs --preset it_opt_4gpu +# +# Identical to it_opt_1gpu.toml except for gpu_count and jobs-per-gpu. run_grid_search sets +# max_workers = len(gpus) * jobs-per-gpu and gives worker i the slice jobs[i::max_workers], so +# 4 GPUs x 2 gives 8 workers and 40 proteins become 5 jobs each rather than 40 in a row. +# +# jobs-per-gpu = 2 matches the recorded run (--jobs-per-gpu 2 on 4 GPUs). Two jobs share a card, +# each holding its own copy of the weights: a measured single job peaked near 11 GiB of an 80 GiB +# H100, so two fit with wide margin. Raise it further only after checking VRAM on the largest +# protein in the set, since cost scales with atom count. +# +# Settings are the defaults of it_opt_scratch/slurm_ensemble_and_score.sh; see the sibling +# preset for how to switch mode. + +[defaults] +DATA_DIR = "/data/inputs" +RESULTS_DIR = "/data/results/it_opt_4gpu" +PROTEINS_CSV = "${DATA_DIR}/proteins.csv" + +# ensemble-sizes and gradient-weights MUST stay single values. Their defaults are "1 2 4 8" and +# "0.01 0.1 0.2", which would expand one job per protein into twelve. gradient-weights only +# feeds the step scaler, which latent_opt does not use, so its value is arbitrary. +# align-to-input matches the scratch runner, which set it unconditionally; the grid defaults it +# off. bond-length-weight is named explicitly because the scratch runner defaulted it to 0.0. +[shared_args] +proteins = "${PROTEINS_CSV}" +model = "protenix" +scalers = "latent_opt" +which-latent = "both" +ensemble-sizes = "8" +gradient-weights = "0.0" +num-diffusion-steps = 200 +bond-length-weight = 5e-5 +align-to-input = true + +[[jobs]] +name = "s_plus_z" +env = "protenix" +gpu_count = 4 +output_subdir = "s_plus_z" +args = { jobs-per-gpu = 2 } diff --git a/run_grid_search.py b/run_grid_search.py index b3336417..058aa6a8 100755 --- a/run_grid_search.py +++ b/run_grid_search.py @@ -19,7 +19,12 @@ from loguru import logger as log from sampleworks.utils.guidance_constants import GuidanceType, StructurePredictor -from sampleworks.utils.guidance_script_arguments import GuidanceConfig, JobConfig, JobResult +from sampleworks.utils.guidance_script_arguments import ( + GuidanceConfig, + JobConfig, + JobResult, + add_latent_opt_args, +) from sampleworks.utils.protein_input import ProteinInput @@ -223,12 +228,29 @@ def run_grid_search( successful = 0 failed = 0 - max_workers = len(gpus) - log.info(f"Running {len(jobs)} jobs with {max_workers} parallel workers") + # Workers per GPU (--jobs-per-gpu) lets two jobs share a card, which is worthwhile because a + # single job leaves the GPU idle while it featurizes and writes output. Two constraints shape + # the worker count: + # never more workers than jobs -- a surplus worker gets an empty queue, and the + # worker_job_queues[worker_num][0] lookup below then raises IndexError before any model + # loads. Easy to hit once jobs_per_gpu multiplies the count. + # keep it a whole multiple of the GPU count, so every card carries the same load. Clamping + # to the job count alone leaves e.g. 5 workers on 4 GPUs: one card runs two jobs + # concurrently while the other three run one and then idle. + max_workers = min(len(gpus) * args.jobs_per_gpu, len(jobs)) + if max_workers > len(gpus): + max_workers -= max_workers % len(gpus) + gpus_used = min(max_workers, len(gpus)) + log.info( + f"Running {len(jobs)} jobs with {max_workers} parallel workers " + f"({gpus_used} GPU(s) x {max_workers // gpus_used} jobs/GPU)" + ) - # Divide the job among the workers: + # Divide the job among the workers. The worker index is no longer the GPU ordinal, so the + # device wraps with i % len(gpus): workers 0..7 across 4 GPUs pair up as 0,1,2,3,0,1,2,3. + # Passing i directly would ask for cuda:4..cuda:7 on a 4-GPU box and fail immediately. worker_job_queues = [ - [build_args_for_process_pool(j, args, i) for j in jobs[i::max_workers]] + [build_args_for_process_pool(j, args, i % len(gpus)) for j in jobs[i::max_workers]] for i in range(max_workers) ] # we'll pickle each job queue separately and then execute each job queue in a separate process @@ -706,6 +728,11 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--augmentation", action="store_true", help="Enable augmentation") parser.add_argument("--align-to-input", action="store_true", help="Align to input structure") + # Latent optimization (IT-opt) arguments, used when --scalers includes latent_opt. + # Registered from the shipped adder rather than restated here so that flag names, types and + # defaults stay identical to the sampleworks-guidance CLI. + add_latent_opt_args(parser) + # RF3-specific arguments parser.add_argument( "--disable-chiral-features", @@ -730,6 +757,14 @@ def parse_args() -> argparse.Namespace: default="auto", help="Max parallel jobs (default: auto = number of GPUs)", ) + parser.add_argument( + "--jobs-per-gpu", + type=int, + default=1, + choices=(1, 2), + help="Concurrent jobs per GPU (default: 1). Each holds its own copy of the model weights, " + "so 2 needs roughly twice the VRAM; 2 is the tested ceiling", + ) parser.add_argument( "--dry-run", action="store_true", diff --git a/src/sampleworks/utils/guidance_script_arguments.py b/src/sampleworks/utils/guidance_script_arguments.py index 53e773e3..bccc75ea 100644 --- a/src/sampleworks/utils/guidance_script_arguments.py +++ b/src/sampleworks/utils/guidance_script_arguments.py @@ -178,6 +178,18 @@ def validate_model_checkpoint( return str(checkpoint_path) +# Inference-time latent optimization (IT-opt) tunables. Named once here so that from_cli() +# and populate_config_for_guidance_type(), the two independent paths that copy CLI values onto +# a config, cannot drift apart as flags are added. +_LATENT_OPT_ATTRS = ( + "which_latent", + "learning_rate", + "outer_steps", + "anchor_weight", + "max_grad_norm", + "bond_length_weight", +) + # Attributes set dynamically by add_*_args helpers that should be copied # from a parsed argparse.Namespace onto a GuidanceConfig instance. _DYNAMIC_ATTRS = [ @@ -193,12 +205,7 @@ def validate_model_checkpoint( "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", + *_LATENT_OPT_ATTRS, # model-specific "model_checkpoint", "method", @@ -404,6 +411,24 @@ def populate_config_for_guidance_type(self, job: JobConfig, args: argparse.Names self.fk_lambda = args.fk_lambda self.fk_resampling_interval = args.fk_resampling_interval self.ensemble_size = job.ensemble_size + elif job.scaler == GuidanceType.LATENT_OPT: + # IT-opt tunables are scalar flags rather than grid axes, so they come straight off + # the driver's namespace. This copy is what makes them work at all: the grid builds + # its GuidanceConfig directly and never calls from_cli(), so the _DYNAMIC_ATTRS loop + # that normally moves parsed values onto a config is unreachable from here. Without + # it every --which-latent / --learning-rate / --outer-steps flag is a silent no-op, + # left at the add_latent_opt_args default that __post_init__ already seeded. + for attr in _LATENT_OPT_ATTRS: + value = getattr(args, attr, None) + if value is not None: + setattr(self, attr, value) + # _run_guidance reads ensemble_size directly rather than via getattr, so it must + # be set here or a latent_opt job raises AttributeError. + self.ensemble_size = job.ensemble_size + # Deliberately no step_size / step_scaler_type here. add_latent_opt_args does not + # define them, so leaving them unset makes _run_guidance fall through to the same + # defaults the scratch runner hit (step_size 0.01, noisespace). Setting them would + # silently change what the recorded runs did. else: self.step_size = job.gradient_weight self.step_scaler_type = args.step_scaler_type From 907faa82e5e569808a572a97123ff67864b4be59 Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Wed, 5 Aug 2026 15:52:11 -0700 Subject: [PATCH 07/15] fix(it-opt): handle chain-relabel & CYS/CSO references in paper RSCC/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 --- it_opt_scratch/score_paper_rmsd.py | 71 +++++++++++++++++++++++++----- it_opt_scratch/score_paper_rscc.py | 5 +++ 2 files changed, 65 insertions(+), 11 deletions(-) diff --git a/it_opt_scratch/score_paper_rmsd.py b/it_opt_scratch/score_paper_rmsd.py index f9369133..89cea77b 100644 --- a/it_opt_scratch/score_paper_rmsd.py +++ b/it_opt_scratch/score_paper_rmsd.py @@ -43,7 +43,12 @@ import torch from atomworks.io.utils.io_utils import load_any from loguru import logger -from sampleworks.utils.atom_array_utils import remove_atoms_with_any_nan_coords +from sampleworks.utils.atom_array_utils import ( + make_atom_id, + make_normalized_atom_id, + remove_atoms_with_any_nan_coords, +) +from sampleworks.utils.cif_utils import resolve_mixed_hetatm_atom_altlocs # Same directory as this script, which is sys.path[0] when run as `python it_opt_scratch/...`. from score_paper_rscc import align_prediction_to_reference, read_selections @@ -61,13 +66,26 @@ def parse_selection(selection: str) -> tuple[str, range]: def reference_conformers( - ref_path: Path, selections: list[str] -) -> dict[str, dict[str, dict[tuple[str, int, str], np.ndarray]]]: - """Per selection, the altloc-A and altloc-B conformers keyed by (chain, res_id, atom_name). + ref_path: Path, selections: list[str], normalize: bool = False +) -> dict[str, dict[str, dict]]: + """Per selection, the altloc-A and altloc-B conformers keyed by atom id. Atoms with a blank altloc are shared between conformers, so they appear in both. + + ``normalize=False`` keys by deposited ``(chain, res_id, atom_name)`` -- the default. + ``normalize=True`` keys by ``make_normalized_atom_id`` (the same normalization + ``filter_to_common_atoms(normalize_ids=True)`` uses for RSCC), a fallback for a relabelled + reference. gemmi reads coordinates because ``load_any`` drops the altloc character; the biotite + load only builds the deposited-key -> normalized-key map. """ st = gemmi.read_structure(str(ref_path)) + raw_to_norm = None + if normalize: + ref_bio = load_any(str(ref_path), altloc="all") + raw_to_norm = dict( + zip(make_atom_id(ref_bio), make_normalized_atom_id(ref_bio), strict=True) + ) + by_res: dict[tuple[str, int], list] = {} for chain in st[0]: for res in chain: @@ -75,26 +93,41 @@ def reference_conformers( (a.name, a.altloc, np.array([a.pos.x, a.pos.y, a.pos.z])) for a in res ) - out: dict[str, dict[str, dict[tuple[str, int, str], np.ndarray]]] = {} + out: dict[str, dict[str, dict]] = {} for sel in selections: chain, residues = parse_selection(sel) - conformers: dict[str, dict[tuple[str, int, str], np.ndarray]] = {"A": {}, "B": {}} + conformers: dict[str, dict] = {"A": {}, "B": {}} for res_id in residues: for name, altloc, xyz in by_res.get((chain, res_id), []): + if normalize: + key = raw_to_norm.get(f"{chain}_{res_id}_{name}") + if key is None: + continue # atom absent from the biotite load + else: + key = (chain, res_id, name) alt = altloc.strip() targets = ("A", "B") if alt == "" else (alt,) for t in targets: if t in conformers: - conformers[t][(chain, res_id, name)] = xyz + conformers[t][key] = xyz out[sel] = conformers return out -def prediction_lookup(atom_array) -> dict[tuple[str, int, str], np.ndarray]: - """Map (chain, res_id, atom_name) -> per-model coordinates, shape [n_models, 3].""" +def prediction_lookup(atom_array, normalize: bool = False) -> dict: + """Map atom id -> per-model coordinates, shape [n_models, 3]. + + ``normalize=False`` keys by raw ``(chain, res_id, atom_name)`` -- the default, so a prediction + whose chain/numbering already match the reference scores exactly as before. ``normalize=True`` + keys by ``make_normalized_atom_id`` (sequential per-chain numbering + chain-index), used only as + a fallback for a relabelled/renumbered prediction (Protenix chain 'A' vs deposited 'P'). + """ coords = atom_array.coord if coords.ndim == 2: # single model -> add the model axis coords = coords[None] + if normalize: + ids = make_normalized_atom_id(atom_array) + return {ids[i]: coords[:, i] for i in range(len(ids))} return { (str(c), int(r), str(n)): coords[:, i] for i, (c, r, n) in enumerate( @@ -128,6 +161,9 @@ def score_protein( ) -> list[dict]: rows: list[dict] = [] ref_path = inputs_dir / "processed" / protein / f"{protein}_single_001_density_input.cif" + # Match generation + the RSCC scorer: collapse mixed ATOM/HETATM modified-residue positions + # (e.g. CYS+CSO) so the reference has the same atoms as a prediction from the cleaned CIF. + ref_path = resolve_mixed_hetatm_atom_altlocs(ref_path) def fail(arm: str, err: str) -> None: for sel in selections: @@ -145,7 +181,7 @@ def fail(arm: str, err: str) -> None: ) try: - conformers = reference_conformers(ref_path, selections) + conformers_raw = reference_conformers(ref_path, selections, normalize=False) # Alignment target: the same array score_paper_rscc.py aligns against, so both metrics # place the prediction identically. ref_atom_array = remove_atoms_with_any_nan_coords(load_any(str(ref_path))) @@ -155,6 +191,9 @@ def fail(arm: str, err: str) -> None: fail(arm, f"setup: {e}") return rows + raw_ref_keys = {k for c in conformers_raw.values() for d in (c["A"], c["B"]) for k in d} + conformers_norm = None # built lazily, only if a prediction needs the fallback + for arm in arms: cif = runs_dir / dir_template.format(protein=protein) / arm / target_filename if not cif.exists(): @@ -163,7 +202,17 @@ def fail(arm: str, err: str) -> None: try: aa = remove_atoms_with_any_nan_coords(load_any(str(cif))) aa = align_prediction_to_reference(ref_atom_array, aa) - pred = prediction_lookup(aa) + pred = prediction_lookup(aa, normalize=False) + # Prefer exact (chain,res,name) matching so a protein whose frame already agrees keeps + # its exact score. Only when the prediction shares NO atoms with the reference (chain + # relabelled, e.g. deposited 'P' vs Protenix 'A') fall back to normalized keys. + if raw_ref_keys.isdisjoint(pred): + if conformers_norm is None: + conformers_norm = reference_conformers(ref_path, selections, normalize=True) + conformers = conformers_norm + pred = prediction_lookup(aa, normalize=True) + else: + conformers = conformers_raw except Exception as e: # noqa: BLE001 logger.error(f"{protein}/{arm}: {e}\n{traceback.format_exc()}") fail(arm, str(e)) diff --git a/it_opt_scratch/score_paper_rscc.py b/it_opt_scratch/score_paper_rscc.py index bfd67823..846822f0 100644 --- a/it_opt_scratch/score_paper_rscc.py +++ b/it_opt_scratch/score_paper_rscc.py @@ -46,6 +46,7 @@ filter_to_common_atoms, remove_atoms_with_any_nan_coords, ) +from sampleworks.utils.cif_utils import resolve_mixed_hetatm_atom_altlocs from sampleworks.utils.density_utils import build_density_transformer, run_density_transformer from sampleworks.utils.frame_transforms import ( apply_forward_transform, @@ -136,6 +137,10 @@ def score_protein( maps_base = maps_dir if maps_dir is not None else inputs_dir / "density_maps" map_path = maps_base / map_template.format(protein=protein) ref_path = inputs_dir / "processed" / protein / f"{protein}_single_001_density_input.cif" + # Match generation: collapse modified-residue positions (mixed ATOM/HETATM, different resname, + # e.g. CYS+CSO) that atomworks would otherwise duplicate into an extra residue, so the reference + # carries the same atoms as a prediction generated from the cleaned CIF. No-op otherwise. + ref_path = resolve_mixed_hetatm_atom_altlocs(ref_path) def fail(arm: str, err: str) -> None: for sel in selections: From 3c7bc01efe6eab5897388c3b2d9528015ccfa3be Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Mon, 10 Aug 2026 15:21:14 -0700 Subject: [PATCH 08/15] feat(utils): relabel non-A/B altlocs so A/B-only tooling keeps both conformers 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. --- src/sampleworks/utils/cif_utils.py | 74 ++++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) diff --git a/src/sampleworks/utils/cif_utils.py b/src/sampleworks/utils/cif_utils.py index 82989b9a..1e4b96e5 100644 --- a/src/sampleworks/utils/cif_utils.py +++ b/src/sampleworks/utils/cif_utils.py @@ -12,6 +12,7 @@ from loguru import logger from sampleworks.utils.atom_array_utils import ( + BLANK_ALTLOC_IDS, find_all_altloc_ids, save_structure_to_cif, select_altloc, @@ -247,6 +248,79 @@ def resolve_mixed_hetatm_atom_altlocs(cif_path: Path | str) -> Path: return tmp_path +def remap_altlocs_to_ab(cif_path: Path | str) -> Path: + """Relabel a residue's non-A/B alternate altloc(s) to fill the A/B slots. + + Some depositions label a residue's two conformers ``A``/``C`` (or ``A``/``D``, etc.) instead of + ``A``/``B``. Tools that assume the two altlocs are literally named A and B -- e.g. a + min-RMSD-to-altloc-A/altloc-B scorer -- then silently drop the second conformer. This is an + optional pre-processing step that, per residue, maps the alternate label onto the free ``B`` + slot: an atom with altloc ``A`` stays ``A`` and the other non-blank altloc becomes ``B``; if + ``A`` is absent, the two labels are assigned to ``A``/``B`` in sorted order. + + Only residues with **exactly two** non-blank altlocs that are **not already** ``{A, B}`` are + touched. Residues that are already A/B, have a single or blank altloc, or carry three or more + altlocs (ambiguous which two to keep) are left unchanged, so the metric's numbers do not move + for structures that were already A/B. + + A warning is logged for every remapped ``(chain, residue)`` position. + + Parameters + ---------- + cif_path + Path to the input CIF file. + + Returns + ------- + Path + Path to a remapped temporary CIF file if any position was changed, or the original + ``cif_path`` unchanged if none were. + """ + cif_path = Path(cif_path) + atom_array = load_any(cif_path, altloc="all", extra_fields=["occupancy", "b_factor"]) + if isinstance(atom_array, AtomArrayStack): + atom_array = atom_array[0] + if not hasattr(atom_array, "altloc_id"): + return cif_path # no altloc annotation -> nothing to remap + + altloc = atom_array.altloc_id.copy() # per-atom altloc characters, mutated in place below + chain_id = atom_array.chain_id + res_id = atom_array.res_id + remapped = 0 + + for chain in np.unique(chain_id): + for rid in np.unique(res_id[chain_id == chain]): + pos = (chain_id == chain) & (res_id == rid) + labels = sorted(set(altloc[pos].tolist()) - BLANK_ALTLOC_IDS) + if len(labels) != 2 or set(labels) == {"A", "B"}: + continue + if "A" in labels: + mapping = {next(x for x in labels if x != "A"): "B"} + else: + mapping = {labels[0]: "A", labels[1]: "B"} + for src, dst in mapping.items(): + sel = pos & (altloc == src) + altloc[sel] = dst + remapped += int(sel.sum()) + logger.warning( + f"Chain {chain}, residue {rid}: remapped altlocs {labels} -> A/B so the alternate " + "conformer is not dropped by A/B-only tooling." + ) + + if remapped == 0: + return cif_path + + atom_array.set_annotation("altloc_id", altloc) + with tempfile.NamedTemporaryFile( + mode="w", suffix=".cif", prefix="sampleworks_altloc_ab_", delete=False + ) as tmp_file: + tmp_path = Path(tmp_file.name) + + save_structure_to_cif(atom_array, tmp_path) + logger.info(f"Wrote altloc-remapped CIF to temporary file: {tmp_path}") + return tmp_path + + def add_category_to_cif( ciffile: CIFFile, data: dict[str, Any], From 7087cd95ca72db8158cc2a86fd8e5d9c9a045265 Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Mon, 10 Aug 2026 15:21:30 -0700 Subject: [PATCH 09/15] fix(it-opt): pair scorer atoms by verified residue frame, not id overlap 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. --- it_opt_scratch/score_paper_rmsd.py | 27 +++++++---- it_opt_scratch/score_paper_rscc.py | 78 +++++++++++++++++++++++++----- 2 files changed, 84 insertions(+), 21 deletions(-) diff --git a/it_opt_scratch/score_paper_rmsd.py b/it_opt_scratch/score_paper_rmsd.py index 89cea77b..60aee1b6 100644 --- a/it_opt_scratch/score_paper_rmsd.py +++ b/it_opt_scratch/score_paper_rmsd.py @@ -48,10 +48,14 @@ make_normalized_atom_id, remove_atoms_with_any_nan_coords, ) -from sampleworks.utils.cif_utils import resolve_mixed_hetatm_atom_altlocs +from sampleworks.utils.cif_utils import remap_altlocs_to_ab, resolve_mixed_hetatm_atom_altlocs # Same directory as this script, which is sys.path[0] when run as `python it_opt_scratch/...`. -from score_paper_rscc import align_prediction_to_reference, read_selections +from score_paper_rscc import ( + align_prediction_to_reference, + read_selections, + strict_frame_is_consistent, +) SELECTION_RE = re.compile(r"^chain\s+(\S+)\s+and\s+resi\s+(-?\d+)\s*-\s*(-?\d+)$") @@ -164,6 +168,9 @@ def score_protein( # Match generation + the RSCC scorer: collapse mixed ATOM/HETATM modified-residue positions # (e.g. CYS+CSO) so the reference has the same atoms as a prediction from the cleaned CIF. ref_path = resolve_mixed_hetatm_atom_altlocs(ref_path) + # Relabel non-A/B alternate altlocs (e.g. A/C) to A/B so the A/B-only conformer split below + # does not drop the second conformer (7Z0E, 7AVG). No-op for already-A/B structures. + ref_path = remap_altlocs_to_ab(ref_path) def fail(arm: str, err: str) -> None: for sel in selections: @@ -191,7 +198,6 @@ def fail(arm: str, err: str) -> None: fail(arm, f"setup: {e}") return rows - raw_ref_keys = {k for c in conformers_raw.values() for d in (c["A"], c["B"]) for k in d} conformers_norm = None # built lazily, only if a prediction needs the fallback for arm in arms: @@ -202,17 +208,18 @@ def fail(arm: str, err: str) -> None: try: aa = remove_atoms_with_any_nan_coords(load_any(str(cif))) aa = align_prediction_to_reference(ref_atom_array, aa) - pred = prediction_lookup(aa, normalize=False) - # Prefer exact (chain,res,name) matching so a protein whose frame already agrees keeps - # its exact score. Only when the prediction shares NO atoms with the reference (chain - # relabelled, e.g. deposited 'P' vs Protenix 'A') fall back to normalized keys. - if raw_ref_keys.isdisjoint(pred): + # Key on deposited (chain,res,name) only when those identifiers demonstrably pair the + # same residues. Sharing *some* keys is not enough: a prediction renumbered from 1 + # overlaps a deposited range like 8-258 on most of its residues while denoting + # different ones, which silently compared residue i against residue i-offset. + if strict_frame_is_consistent(ref_atom_array, aa): + conformers = conformers_raw + pred = prediction_lookup(aa, normalize=False) + else: if conformers_norm is None: conformers_norm = reference_conformers(ref_path, selections, normalize=True) conformers = conformers_norm pred = prediction_lookup(aa, normalize=True) - else: - conformers = conformers_raw except Exception as e: # noqa: BLE001 logger.error(f"{protein}/{arm}: {e}\n{traceback.format_exc()}") fail(arm, str(e)) diff --git a/it_opt_scratch/score_paper_rscc.py b/it_opt_scratch/score_paper_rscc.py index 846822f0..e8f551de 100644 --- a/it_opt_scratch/score_paper_rscc.py +++ b/it_opt_scratch/score_paper_rscc.py @@ -70,6 +70,57 @@ def read_selections(csv_path: Path) -> dict[str, list[str]]: return out +MIN_RESIDUE_IDENTITY = 0.95 # matched pairs below this share a numbering frame only by accident + + +def strict_frame_is_consistent(ref_atom_array, pred_atom_array) -> bool: + """Whether prediction and reference can be paired on deposited (chain, res_id, atom_name). + + ``filter_to_common_atoms`` pairs atoms by identifier alone, never checking 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 succeeds on a large but sequence-shifted set, + pairing prediction residue i with reference residue i-offset. Kabsch then fits that shifted + correspondence and reports a fold-scale error for a structure that is actually correct. + + Comparing residue names separates the two cases cleanly: a correct frame agrees on every + pair, a shifted one agrees only at the ~5% rate expected from 20 residue types by chance. + + Parameters + ---------- + ref_atom_array, pred_atom_array + Reference and predicted structures, as loaded by ``load_any``. + + Returns + ------- + bool + True when strict matching pairs chemically identical residues and so may be used + directly; False when the caller should fall back to normalized (sequential) ids. + """ + try: + ref_common, pred_common = filter_to_common_atoms(ref_atom_array, pred_atom_array) + except RuntimeError: + return False # no shared identifiers at all, e.g. deposited chain 'B' vs Protenix 'A' + identity = residue_identity(ref_common, pred_common) + return identity is not None and identity >= MIN_RESIDUE_IDENTITY + + +def residue_identity(ref_common, pred_common) -> float | None: + """Fraction of matched pairs whose residue names agree, or None if they cannot be paired. + + ``filter_to_common_atoms`` masks with ``np.isin``, so a reference carrying two altlocs of one + atom contributes both copies of a duplicated (chain, res_id, atom_name) key while a + single-conformer prediction contributes one. The two matched sets then differ in length and + there is no atom-for-atom correspondence to score at all (6NI5 and 6NI6 are the known cases). + Returning None keeps that distinct from "paired, but the residues disagree". + """ + ref_names = np.asarray(ref_common.res_name) + pred_names = np.asarray(pred_common.res_name) + if ref_names.shape != pred_names.shape: + return None + return float((ref_names == pred_names).mean()) + + def align_prediction_to_reference(ref_atom_array, pred_atom_array): """Global uniform-weight Kabsch of prediction onto reference, applied to every predicted atom. @@ -78,17 +129,22 @@ def align_prediction_to_reference(ref_atom_array, pred_atom_array): fitting on the 3-residue window itself would let a wrong local conformation be rotated into apparent agreement. """ - try: - ref_common, pred_common = filter_to_common_atoms(ref_atom_array, pred_atom_array) - except RuntimeError: - # The prediction relabels every chain to 'A' and renumbers residues from 1, while the - # reference keeps the deposited chain id and numbering (e.g. chain 'P', res 5-234). Strict - # (chain,res,name) matching then finds nothing. Fall back to sequential per-chain matching, - # which realigns these (otherwise identical) structures. This only runs when strict matching - # raises, so the proteins that already align keep their exact matched-atom set unchanged. - ref_common, pred_common = filter_to_common_atoms( - ref_atom_array, pred_atom_array, normalize_ids=True - ) + # Sequential per-chain matching is used unless the deposited identifiers demonstrably pair the + # same residues, so a partial numbering overlap can no longer produce a silently shifted fit. + normalize = not strict_frame_is_consistent(ref_atom_array, pred_atom_array) + ref_common, pred_common = filter_to_common_atoms( + ref_atom_array, pred_atom_array, normalize_ids=normalize + ) + if normalize: + identity = residue_identity(ref_common, pred_common) + # None means the matched sets are different sizes; the explicit shape check below reports + # that with the actual atom counts, which is the more useful message for it. + if identity is not None and identity < MIN_RESIDUE_IDENTITY: + raise ValueError( + f"neither deposited nor sequential atom ids pair matching residues " + f"(sequential identity {identity:.2f}); the reference and the prediction are " + f"not the same construct" + ) ref_t = torch.from_numpy(ref_common.coord).float() pred_t = torch.from_numpy(pred_common.coord).float() ref_t = match_batch(ref_t, pred_t.shape[0]) From 4d959a49d532695ff796fc2a3d38e91f8ecfe25a Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Mon, 10 Aug 2026 15:22:06 -0700 Subject: [PATCH 10/15] feat(it-opt): forward the paper's coordinate-guidance recipe from the 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. --- it_opt_scratch/run_targets.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/it_opt_scratch/run_targets.py b/it_opt_scratch/run_targets.py index c36d2351..eb49395a 100644 --- a/it_opt_scratch/run_targets.py +++ b/it_opt_scratch/run_targets.py @@ -321,7 +321,16 @@ def guidance_for_mode(mode: str, args) -> tuple[GuidanceType, dict]: if mode == "baseline": return GuidanceType.PURE_GUIDANCE, {"step_scaler_type": "none"} if mode == "coord_guidance": - return GuidanceType.PURE_GUIDANCE, {"step_scaler_type": "noisespace", "step_size": args.step_size} + # The paper's coordinate guidance normalizes the density gradient to the EDM denoising-update + # magnitude (gradient_normalization) and applies AF3 augmentation + realign each step; only + # then does step_size act as a fraction of the denoising step. These are mode-scoped so the + # latent-opt arms are unaffected. + return GuidanceType.PURE_GUIDANCE, { + "step_scaler_type": "noisespace", + "step_size": args.step_size, + "gradient_normalization": args.gradient_normalization, + "augmentation": args.augmentation, + } if mode in ("s_only", "z_only", "s_plus_z"): which = {"s_only": "single", "z_only": "pair", "s_plus_z": "both"}[mode] return GuidanceType.LATENT_OPT, { @@ -346,6 +355,7 @@ def build_config(target: dict, guidance_type: GuidanceType, out_dir: Path, args) output_dir=str(out_dir), resolution=float(target["resolution"]), num_diffusion_steps=args.num_steps, + guidance_start=args.guidance_start, # -1 -> guide from step 0; e.g. 120 -> last low-noise steps align_to_input=True, ) config.ensemble_size = args.ensemble_size # set dynamically (not a declared GuidanceConfig field) @@ -399,12 +409,22 @@ def parse_args() -> argparse.Namespace: ap.add_argument("--model", default="protenix", choices=[m.value for m in StructurePredictor]) ap.add_argument("--ensemble-size", dest="ensemble_size", type=int, default=4) ap.add_argument("--num-steps", dest="num_steps", type=int, default=200) + ap.add_argument("--guidance-start", dest="guidance_start", type=int, default=-1, + help="step at which guidance begins (coord DPS and IT-opt both use it as a " + "fraction of num_steps); -1 means from step 0. e.g. 120 of 200 guides " + "only the last low-noise steps, matching the tuned coordinate-guidance recipe") ap.add_argument("--outer-steps", dest="outer_steps", type=int, default=2) ap.add_argument("--lr", type=float, default=0.05) ap.add_argument("--anchor", type=float, default=0.0, help="on-manifold anchor weight (IT-opt)") ap.add_argument("--bond-length-weight", dest="bond_length_weight", type=float, default=0.0, help="coordinate-space bond-geometry penalty weight (IT-opt); 0 disables") ap.add_argument("--step-size", dest="step_size", type=float, default=0.1, help="coord-guidance DPS step") + ap.add_argument("--gradient-normalization", dest="gradient_normalization", action="store_true", + help="coord guidance: normalize the density gradient to the denoising-update " + "magnitude before scaling by --step-size (the paper's recipe; makes 0.1 mean " + "10%% of the denoising step rather than 0.1x the raw gradient)") + ap.add_argument("--augmentation", action="store_true", + help="coord guidance: apply AF3 random augmentation + realign each step (paper recipe)") ap.add_argument("--seed", type=int, default=0) ap.add_argument("--device", default="cuda:0") ap.add_argument("--checkpoint", default=None) From f725c919f4a9723445989a03e0a442c75212ddad Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Mon, 10 Aug 2026 17:26:47 -0700 Subject: [PATCH 11/15] chore(it-opt): untrack the experiment scratch tooling 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. --- it_opt_scratch/paper_maxrmsd_selections.csv | 41 -- it_opt_scratch/run_targets.py | 472 -------------------- it_opt_scratch/run_targets_simplified.py | 186 -------- it_opt_scratch/score_paper_rmsd.py | 298 ------------ it_opt_scratch/score_paper_rscc.py | 360 --------------- it_opt_scratch/score_paper_simplified.py | 154 ------- it_opt_scratch/score_rmsd_simplified.py | 259 ----------- it_opt_scratch/score_rscc_simplified.py | 240 ---------- it_opt_scratch/slurm_ensemble_and_score.sh | 148 ------ 9 files changed, 2158 deletions(-) delete mode 100644 it_opt_scratch/paper_maxrmsd_selections.csv delete mode 100644 it_opt_scratch/run_targets.py delete mode 100644 it_opt_scratch/run_targets_simplified.py delete mode 100644 it_opt_scratch/score_paper_rmsd.py delete mode 100644 it_opt_scratch/score_paper_rscc.py delete mode 100644 it_opt_scratch/score_paper_simplified.py delete mode 100644 it_opt_scratch/score_rmsd_simplified.py delete mode 100644 it_opt_scratch/score_rscc_simplified.py delete mode 100755 it_opt_scratch/slurm_ensemble_and_score.sh diff --git a/it_opt_scratch/paper_maxrmsd_selections.csv b/it_opt_scratch/paper_maxrmsd_selections.csv deleted file mode 100644 index e0138d2c..00000000 --- a/it_opt_scratch/paper_maxrmsd_selections.csv +++ /dev/null @@ -1,41 +0,0 @@ -protein,selection -1VME,chain A and resi 1-1;chain A and resi 164-164;chain A and resi 207-207;chain A and resi 218-218;chain A and resi 245-245;chain A and resi 269-269;chain A and resi 318-320;chain A and resi 326-326;chain A and resi 373-373;chain A and resi 67-67;chain A and resi 70-70 -2A26,chain A and resi 11-11;chain A and resi 16-18;chain A and resi 26-26;chain A and resi 30-30;chain A and resi 45-45 -2BWD,chain A and resi 106-106;chain A and resi 141-141;chain A and resi 16-16;chain A and resi 168-168;chain A and resi 197-199;chain A and resi 211-211;chain A and resi 22-22;chain A and resi 239-239;chain A and resi 276-276;chain A and resi 29-29;chain A and resi 302-302;chain A and resi 313-313;chain A and resi 321-321;chain A and resi 323-323;chain A and resi 325-325;chain A and resi 33-33;chain A and resi 330-330;chain A and resi 340-340;chain A and resi 36-36;chain A and resi 46-46;chain A and resi 50-50;chain A and resi 78-78;chain A and resi 92-92;chain A and resi 94-94;chain A and resi 98-98 -2E7Z,chain A and resi 105-105;chain A and resi 141-141;chain A and resi 195-197;chain A and resi 23-23;chain A and resi 277-279;chain A and resi 339-339;chain A and resi 344-344;chain A and resi 376-376;chain A and resi 446-446;chain A and resi 508-509;chain A and resi 549-551;chain A and resi 633-635;chain A and resi 664-666;chain A and resi 81-81;chain A and resi 93-93 -2I6H,chain B and resi 105-105;chain B and resi 108-108;chain B and resi 11-11;chain B and resi 131-131;chain B and resi 146-146;chain B and resi 153-155;chain B and resi 159-159;chain B and resi 165-165;chain B and resi 169-171;chain B and resi 23-23;chain B and resi 26-26;chain B and resi 33-33;chain B and resi 56-56 -2IGT,chain A and resi 108-109;chain A and resi 121-121;chain A and resi 128-128;chain A and resi 131-131;chain A and resi 181-181;chain A and resi 185-185;chain A and resi 19-19;chain A and resi 200-200;chain A and resi 215-215;chain A and resi 225-225;chain A and resi 235-235;chain A and resi 254-256;chain A and resi 261-261;chain A and resi 290-292;chain A and resi 53-53;chain A and resi 58-59;chain A and resi 62-62;chain A and resi 77-77;chain A and resi 94-94;chain A and resi 99-99 -2P97,chain A and resi 100-100;chain A and resi 112-112;chain A and resi 129-129;chain A and resi 14-14;chain A and resi 141-141;chain A and resi 147-147;chain A and resi 155-155;chain A and resi 166-166;chain A and resi 18-20;chain A and resi 180-180;chain A and resi 186-186;chain A and resi 2-2;chain A and resi 35-35;chain A and resi 56-57;chain A and resi 67-67;chain A and resi 97-97 -2QCV,chain A and resi 111-111;chain A and resi 144-144;chain A and resi 198-198;chain A and resi 2-4;chain A and resi 215-215;chain A and resi 261-261;chain A and resi 282-282;chain A and resi 284-284;chain A and resi 39-39;chain A and resi 54-54;chain A and resi 77-77;chain A and resi 85-85 -2XLW,chain A and resi 10-10;chain A and resi 106-106;chain A and resi 111-111;chain A and resi 115-115;chain A and resi 121-122;chain A and resi 18-20;chain A and resi 30-30;chain A and resi 37-37;chain A and resi 52-52;chain A and resi 60-62;chain A and resi 66-66;chain A and resi 69-71;chain A and resi 81-83;chain A and resi 87-89;chain A and resi 92-92;chain A and resi 98-99 -2YL0,chain A and resi 10-10;chain A and resi 102-104;chain A and resi 106-106;chain A and resi 115-115;chain A and resi 19-19;chain A and resi 30-30;chain A and resi 37-37;chain A and resi 40-40;chain A and resi 46-46;chain A and resi 52-52;chain A and resi 60-62;chain A and resi 66-66;chain A and resi 68-70;chain A and resi 74-76;chain A and resi 81-81;chain A and resi 88-89;chain A and resi 92-92;chain A and resi 94-94 -3DT2,chain A and resi 103-103;chain A and resi 109-111;chain A and resi 11-11;chain A and resi 120-120;chain A and resi 151-151;chain A and resi 153-154;chain A and resi 17-17;chain A and resi 191-191;chain A and resi 209-209;chain A and resi 21-21;chain A and resi 274-274;chain A and resi 276-276;chain A and resi 296-296;chain A and resi 320-320;chain A and resi 36-36;chain A and resi 39-39;chain A and resi 390-390;chain A and resi 41-42;chain A and resi 422-422;chain A and resi 50-51;chain A and resi 502-502;chain A and resi 510-510;chain A and resi 524-524;chain A and resi 545-545;chain A and resi 55-55;chain A and resi 578-578;chain A and resi 583-583;chain A and resi 587-587;chain A and resi 597-597;chain A and resi 613-613;chain A and resi 70-71 -3HVV,chain A and resi 58-60;chain A and resi 93-93 -3HYN,chain A and resi 121-121;chain A and resi 127-127;chain A and resi 133-134;chain A and resi 141-141;chain A and resi 149-149;chain A and resi 151-151;chain A and resi 160-160;chain A and resi 166-166;chain A and resi 182-182;chain A and resi 19-21;chain A and resi 31-31;chain A and resi 38-38;chain A and resi 45-45;chain A and resi 60-62;chain A and resi 69-69;chain A and resi 71-71;chain A and resi 78-78;chain A and resi 91-93;chain A and resi 97-97 -3I7M,chain A and resi 112-112;chain A and resi 130-130;chain A and resi 133-133;chain A and resi 15-15;chain A and resi 29-29;chain A and resi 38-39;chain A and resi 60-60;chain A and resi 68-68;chain A and resi 70-70;chain A and resi 78-79;chain A and resi 8-9;chain A and resi 85-85;chain A and resi 88-88;chain A and resi 96-96 -3IAC,chain A and resi 111-111;chain A and resi 186-186;chain A and resi 200-200;chain A and resi 212-212;chain A and resi 287-287;chain A and resi 317-317;chain A and resi 334-334;chain A and resi 55-57;chain A and resi 84-84 -3L4P,chain A and resi 11-11;chain A and resi 124-124;chain A and resi 13-13;chain A and resi 132-132;chain A and resi 134-134;chain A and resi 166-166;chain A and resi 171-173;chain A and resi 18-18;chain A and resi 182-182;chain A and resi 2-2;chain A and resi 204-206;chain A and resi 216-216;chain A and resi 233-233;chain A and resi 240-240;chain A and resi 261-261;chain A and resi 271-273;chain A and resi 303-305;chain A and resi 318-320;chain A and resi 334-334;chain A and resi 352-352;chain A and resi 372-372;chain A and resi 382-382;chain A and resi 396-396;chain A and resi 466-466;chain A and resi 468-468;chain A and resi 520-520;chain A and resi 55-55;chain A and resi 553-553;chain A and resi 557-557;chain A and resi 559-559;chain A and resi 579-579;chain A and resi 58-58;chain A and resi 587-587;chain A and resi 594-595;chain A and resi 597-597;chain A and resi 604-605;chain A and resi 661-661;chain A and resi 680-680;chain A and resi 722-724;chain A and resi 733-733;chain A and resi 736-737;chain A and resi 742-742;chain A and resi 748-750;chain A and resi 761-761;chain A and resi 783-783;chain A and resi 789-789;chain A and resi 82-84;chain A and resi 826-826;chain A and resi 834-834;chain A and resi 842-843;chain A and resi 847-847;chain A and resi 889-889 -3T94,chain A and resi 116-118;chain A and resi 132-133;chain A and resi 152-152;chain A and resi 204-205;chain A and resi 211-211;chain A and resi 229-229;chain A and resi 232-232;chain A and resi 42-42;chain A and resi 91-91 -4GMU,chain A and resi 120-120;chain A and resi 123-123;chain A and resi 14-14;chain A and resi 190-191;chain A and resi 224-224;chain A and resi 226-227;chain A and resi 235-235;chain A and resi 244-244;chain A and resi 253-253;chain A and resi 312-312;chain A and resi 350-350;chain A and resi 39-39;chain A and resi 403-403;chain A and resi 41-42;chain A and resi 442-442;chain A and resi 475-476;chain A and resi 50-51;chain A and resi 503-503;chain A and resi 543-543;chain A and resi 545-545;chain A and resi 56-56;chain A and resi 578-578;chain A and resi 582-582;chain A and resi 587-587;chain A and resi 606-607;chain A and resi 609-609;chain A and resi 613-613;chain A and resi 619-619;chain A and resi 70-70;chain A and resi 85-87;chain A and resi 97-97 -4OLE,chain B and resi 395-395;chain B and resi 424-426 -5A71,chain A and resi 105-105;chain A and resi 139-139;chain A and resi 159-159;chain A and resi 182-182;chain A and resi 194-195;chain A and resi 214-215;chain A and resi 281-283;chain A and resi 298-298;chain A and resi 74-74;chain A and resi 76-76 -5I09,chain A and resi 15-15;chain A and resi 166-166;chain A and resi 17-17;chain A and resi 170-170;chain A and resi 174-174;chain A and resi 188-188;chain A and resi 21-21;chain A and resi 213-213;chain A and resi 281-281;chain A and resi 303-303;chain A and resi 31-31;chain A and resi 325-325;chain A and resi 34-34;chain A and resi 351-351;chain A and resi 57-57;chain A and resi 69-71;chain A and resi 91-91 -5IMV,chain A and resi 106-106;chain A and resi 113-113;chain A and resi 119-119;chain A and resi 123-123;chain A and resi 127-128;chain A and resi 131-131;chain A and resi 133-133;chain A and resi 155-155;chain A and resi 24-24;chain A and resi 59-59;chain A and resi 6-6;chain A and resi 65-67;chain A and resi 74-74;chain A and resi 82-82;chain A and resi 84-84 -5MC8,chain A and resi 103-104;chain A and resi 107-107;chain A and resi 115-115;chain A and resi 128-129;chain A and resi 131-131;chain A and resi 136-136;chain A and resi 148-148;chain A and resi 149-149;chain A and resi 166-166;chain A and resi 168-168;chain A and resi 172-172;chain A and resi 176-176;chain A and resi 179-179;chain A and resi 181-181;chain A and resi 194-194;chain A and resi 196-196;chain A and resi 198-199;chain A and resi 202-202;chain A and resi 206-206;chain A and resi 210-210;chain A and resi 214-214;chain A and resi 223-223;chain A and resi 235-235;chain A and resi 238-238;chain A and resi 240-241;chain A and resi 246-246;chain A and resi 269-269;chain A and resi 275-275;chain A and resi 285-285;chain A and resi 293-293;chain A and resi 296-296;chain A and resi 300-300;chain A and resi 317-317;chain A and resi 324-324;chain A and resi 346-346;chain A and resi 357-359;chain A and resi 36-36;chain A and resi 366-366;chain A and resi 377-378;chain A and resi 43-43;chain A and resi 51-53;chain A and resi 79-80;chain A and resi 93-93 -5MHX,chain A and resi 131-131;chain A and resi 162-164;chain A and resi 17-17;chain A and resi 179-179;chain A and resi 191-191;chain A and resi 198-198;chain A and resi 201-201;chain A and resi 219-220;chain A and resi 251-252;chain A and resi 265-265;chain A and resi 368-369;chain A and resi 377-377;chain A and resi 380-381;chain A and resi 409-409;chain A and resi 433-433;chain A and resi 443-443;chain A and resi 46-47;chain A and resi 465-465;chain A and resi 476-476;chain A and resi 57-57 -5RBR,chain A and resi 11-13;chain A and resi 149-149;chain A and resi 154-154;chain A and resi 164-164;chain A and resi 174-174;chain A and resi 182-182;chain A and resi 206-206;chain A and resi 215-215;chain A and resi 223-224;chain A and resi 240-240;chain A and resi 246-246;chain A and resi 249-251;chain A and resi 268-268;chain A and resi 276-276;chain A and resi 279-281;chain A and resi 289-289;chain A and resi 291-291;chain A and resi 298-300;chain A and resi 311-311;chain A and resi 325-325;chain A and resi 329-329;chain A and resi 68-68;chain A and resi 71-71;chain A and resi 9-9 -5SOP,chain A and resi 10-12;chain A and resi 100-102;chain A and resi 113-113;chain A and resi 131-133;chain A and resi 138-140;chain A and resi 157-159;chain A and resi 20-22;chain A and resi 28-30;chain A and resi 38-38;chain A and resi 4-6;chain A and resi 42-42;chain A and resi 47-49;chain A and resi 87-88 -6B8X,chain A and resi 105-105;chain A and resi 115-115;chain A and resi 121-121;chain A and resi 147-148;chain A and resi 150-152;chain A and resi 157-157;chain A and resi 159-159;chain A and resi 176-176;chain A and resi 181-183;chain A and resi 196-197;chain A and resi 207-207;chain A and resi 216-216;chain A and resi 221-221;chain A and resi 224-224;chain A and resi 226-226;chain A and resi 233-233;chain A and resi 235-235;chain A and resi 239-241;chain A and resi 24-24;chain A and resi 246-246;chain A and resi 252-253;chain A and resi 26-26;chain A and resi 265-265;chain A and resi 269-269;chain A and resi 272-272;chain A and resi 276-276;chain A and resi 278-280;chain A and resi 32-32;chain A and resi 36-36;chain A and resi 72-72;chain A and resi 78-80;chain A and resi 92-92 -6DUR,chain A and resi 216-216;chain A and resi 269-269;chain A and resi 29-31;chain A and resi 359-361;chain A and resi 389-391;chain A and resi 397-399;chain A and resi 445-447;chain A and resi 66-66 -6NI5,chain A and resi 101-101;chain A and resi 114-114;chain A and resi 116-117;chain A and resi 12-12;chain A and resi 124-124;chain A and resi 135-135;chain A and resi 138-139;chain A and resi 15-15;chain A and resi 151-153;chain A and resi 191-191;chain A and resi 193-193;chain A and resi 201-202;chain A and resi 204-206;chain A and resi 209-210;chain A and resi 21-21;chain A and resi 213-215;chain A and resi 219-219;chain A and resi 24-24;chain A and resi 30-32;chain A and resi 35-35;chain A and resi 39-39;chain A and resi 4-4;chain A and resi 43-43;chain A and resi 52-52;chain A and resi 61-62;chain A and resi 64-64;chain A and resi 66-66;chain A and resi 74-74;chain A and resi 86-86;chain A and resi 9-9;chain A and resi 92-94 -6NI6,chain A and resi 101-101;chain A and resi 114-114;chain A and resi 116-117;chain A and resi 12-12;chain A and resi 124-124;chain A and resi 135-135;chain A and resi 138-139;chain A and resi 15-15;chain A and resi 151-153;chain A and resi 174-174;chain A and resi 178-179;chain A and resi 181-181;chain A and resi 191-191;chain A and resi 193-193;chain A and resi 201-202;chain A and resi 204-206;chain A and resi 209-210;chain A and resi 21-21;chain A and resi 213-215;chain A and resi 219-219;chain A and resi 30-32;chain A and resi 35-35;chain A and resi 39-39;chain A and resi 4-4;chain A and resi 52-52;chain A and resi 61-62;chain A and resi 64-64;chain A and resi 66-66;chain A and resi 74-74;chain A and resi 86-86;chain A and resi 92-94 -6RP1,chain C and resi 10-10;chain C and resi 185-185;chain C and resi 26-26;chain C and resi 287-287;chain C and resi 290-290;chain C and resi 326-328;chain C and resi 360-360;chain C and resi 418-418;chain C and resi 431-431;chain C and resi 465-465;chain C and resi 475-475;chain C and resi 5-5;chain C and resi 513-513;chain C and resi 522-522;chain C and resi 536-536;chain C and resi 64-64 -6YVM,chain A and resi 224-225;chain A and resi 234-235;chain A and resi 241-241;chain A and resi 245-245;chain A and resi 248-248;chain A and resi 251-251;chain A and resi 261-262;chain A and resi 286-288;chain A and resi 295-297;chain A and resi 303-303;chain A and resi 308-308;chain A and resi 323-323;chain A and resi 327-327;chain A and resi 330-330;chain A and resi 335-335;chain A and resi 338-338;chain A and resi 343-343;chain A and resi 352-352;chain A and resi 355-356;chain A and resi 359-359;chain A and resi 361-362;chain A and resi 365-365;chain A and resi 371-371;chain A and resi 374-374;chain A and resi 391-393;chain A and resi 395-395;chain A and resi 404-404;chain A and resi 411-413;chain A and resi 427-427;chain A and resi 438-438;chain A and resi 444-444;chain A and resi 454-454;chain A and resi 461-461;chain A and resi 465-465;chain A and resi 476-476 -7AVG,chain A and resi 1-1;chain A and resi 102-104;chain A and resi 106-106;chain A and resi 109-109;chain A and resi 112-114;chain A and resi 116-118;chain A and resi 123-123;chain A and resi 126-126;chain A and resi 13-13;chain A and resi 15-16;chain A and resi 19-19;chain A and resi 24-24;chain A and resi 37-37;chain A and resi 39-39;chain A and resi 43-44;chain A and resi 47-49;chain A and resi 5-5;chain A and resi 52-52;chain A and resi 65-65;chain A and resi 67-67;chain A and resi 7-7;chain A and resi 81-81;chain A and resi 85-85;chain A and resi 86-87;chain A and resi 89-89 -7FRD,chain A and resi 100-102;chain A and resi 105-105;chain A and resi 107-107;chain A and resi 134-136;chain A and resi 157-159;chain A and resi 162-163;chain A and resi 169-169;chain A and resi 24-26;chain A and resi 28-29;chain A and resi 3-5;chain A and resi 31-31;chain A and resi 38-38;chain A and resi 47-49;chain A and resi 58-58;chain A and resi 87-88;chain A and resi 9-11;chain A and resi 90-90 -7P6M,chain A and resi 102-104;chain A and resi 108-109;chain A and resi 111-111;chain A and resi 113-113;chain A and resi 116-117;chain A and resi 119-120;chain A and resi 123-124;chain A and resi 15-15;chain A and resi 21-21;chain A and resi 39-39;chain A and resi 44-45;chain A and resi 47-48;chain A and resi 5-5;chain A and resi 52-52;chain A and resi 63-63;chain A and resi 65-65;chain A and resi 78-79;chain A and resi 81-81;chain A and resi 85-87;chain A and resi 89-89 -7Z0E,chain P and resi 100-100;chain P and resi 104-104;chain P and resi 11-13;chain P and resi 115-115;chain P and resi 122-124;chain P and resi 131-131;chain P and resi 132-132;chain P and resi 133-133;chain P and resi 137-139;chain P and resi 145-145;chain P and resi 15-15;chain P and resi 18-20;chain P and resi 181-183;chain P and resi 193-194;chain P and resi 195-197;chain P and resi 202-202;chain P and resi 204-204;chain P and resi 206-208;chain P and resi 209-209;chain P and resi 214-216;chain P and resi 22-22;chain P and resi 224-224;chain P and resi 226-226;chain P and resi 23-24;chain P and resi 27-27;chain P and resi 28-28;chain P and resi 32-32;chain P and resi 45-47;chain P and resi 48-48;chain P and resi 54-56;chain P and resi 59-59;chain P and resi 60-60;chain P and resi 64-65;chain P and resi 7-7;chain P and resi 82-84;chain P and resi 93-94;chain P and resi 96-97 -8VPW,chain A and resi 101-101;chain A and resi 117-117;chain A and resi 12-12;chain A and resi 124-124;chain A and resi 126-126;chain A and resi 138-139;chain A and resi 142-142;chain A and resi 147-147;chain A and resi 149-149;chain A and resi 15-15;chain A and resi 151-153;chain A and resi 162-162;chain A and resi 165-165;chain A and resi 201-202;chain A and resi 204-204;chain A and resi 210-210;chain A and resi 214-215;chain A and resi 31-31;chain A and resi 35-35;chain A and resi 42-42;chain A and resi 49-49;chain A and resi 51-52;chain A and resi 67-67;chain A and resi 72-72;chain A and resi 74-74;chain A and resi 95-97 -8VQ1,chain A and resi 101-101;chain A and resi 117-117;chain A and resi 12-12;chain A and resi 124-124;chain A and resi 138-139;chain A and resi 142-142;chain A and resi 147-147;chain A and resi 149-149;chain A and resi 15-15;chain A and resi 151-153;chain A and resi 162-162;chain A and resi 165-165;chain A and resi 201-202;chain A and resi 204-204;chain A and resi 210-210;chain A and resi 214-215;chain A and resi 35-35;chain A and resi 42-42;chain A and resi 49-49;chain A and resi 52-52;chain A and resi 67-67;chain A and resi 72-72;chain A and resi 74-74;chain A and resi 95-97 -8Z76,chain C and resi 148-148;chain C and resi 162-162;chain C and resi 230-230;chain C and resi 248-248;chain C and resi 253-253;chain C and resi 256-258;chain C and resi 470-470;chain C and resi 82-82 -9BN8,chain A and resi 113-113;chain A and resi 151-151;chain A and resi 203-203;chain A and resi 208-208;chain A and resi 218-218;chain A and resi 226-228;chain A and resi 238-238;chain A and resi 251-251;chain A and resi 315-315;chain A and resi 327-327;chain A and resi 33-33;chain A and resi 336-337;chain A and resi 364-364;chain A and resi 375-375;chain A and resi 394-394;chain A and resi 4-4;chain A and resi 428-428;chain A and resi 431-431;chain A and resi 438-438 diff --git a/it_opt_scratch/run_targets.py b/it_opt_scratch/run_targets.py deleted file mode 100644 index eb49395a..00000000 --- a/it_opt_scratch/run_targets.py +++ /dev/null @@ -1,472 +0,0 @@ -"""Batch conformational-ensemble generation over a list of protein targets. - -This drives the sampleworks guidance pipeline through its public entry point (`run_guidance`): the -model is loaded ONCE, a `GuidanceConfig` is built per run, and the pipeline itself does -featurize -> reward -> scaler -> sample -> save. Output for each (target, mode) lands under -// (refined.cif + trajectory + losses). A per-run failure is recorded and the batch -continues; a batch_summary.json is written at the end. - -Where the inputs come from (three forms, checked in this order): - * CSV --targets x.csv columns: name,structure,density,resolution - (out dir = /) - * JSON --targets x.json list of {name, density, structure, resolution, out_dir} - * IDs --proteins 2YL0,5I09 / --proteins-file ids.txt - paths are built from --structure-dir/--structure-template and - --density-dir/--density-template, so you only list PDB IDs. Templates take - {pdb} (as written), {PDB} (upper) and {pdb_low} (lower). - * none the built-in DEFAULT_TARGETS (1vme) - -Where the outputs go: - /// refined.cif trajectory/ losses.txt job_metadata.json run.log - --summary path batch_summary.json (per-shard copies alongside it when sharded) - -Modes (--modes, comma-separated, or 'all') -> guidance type: - baseline unguided sampling (pure_guidance, no step scaler) - s_only IT-opt, which_latent=single - z_only IT-opt, which_latent=pair - s_plus_z IT-opt, which_latent=both - coord_guidance shipped coordinate DPS (pure_guidance + noise-space step scaler) - -Run on the pod from the repo root -- one process, explicit CSV: - pixi run -e protenix-dev python it_opt_scratch/run_targets.py \ - --targets /home/dev/test_data/proteins.csv --output-base it_opt_scratch/targets_out \ - --modes baseline,z_only,coord_guidance --ensemble-size 4 --num-steps 200 --outer-steps 2 - -Run on 4 GPUs with 2 proteins in flight per GPU (8 worker processes), IDs + directories: - pixi run -e protenix-dev python -u it_opt_scratch/run_targets.py \ - --proteins-file it_opt_scratch/regen11.txt \ - --structure-dir /home/dev/test_data/processed \ - --density-dir it_opt_scratch/targets_out_11_regen_0.5occA_0.5occB/density_maps \ - --density-template '{PDB}_0.5occA_0.5occB_1.00A.ccp4' \ - --name-template '{PDB}_0.5occA_0.5occB' --resolution 1.0 \ - --output-base it_opt_scratch/targets_out_11_regen_0.5occA_0.5occB \ - --modes s_plus_z --ensemble-size 8 --bond-length-weight 5e-5 \ - --gpus 0,1,2,3 --jobs-per-gpu 2 - -Add --dry-run to either command to print the resolved plan (paths, existence, shard -assignment) without loading the model. -""" - -from __future__ import annotations - -import argparse -import csv -import json -import os -import subprocess -import sys -import traceback -from pathlib import Path - -import torch - -from sampleworks.utils.guidance_constants import GuidanceType, StructurePredictor -from sampleworks.utils.guidance_script_arguments import GuidanceConfig -from sampleworks.utils.guidance_script_utils import get_model_and_device, run_guidance - -REPO = Path(__file__).resolve().parents[1] -ALL_MODES = ["baseline", "s_only", "z_only", "s_plus_z", "coord_guidance"] - -DEFAULT_TARGETS = [ - { - "name": "1vme", - "density": "tests/resources/1vme/1vme_final_carved_edited_0.5occA_0.5occB_1.80A.ccp4", - "structure": "tests/resources/1vme/1vme_final_carved_edited_0.5occA_0.5occB.cif", - "resolution": 1.8, - "out_dir": "it_opt_scratch/targets_out/1vme", - }, -] - - -# ============================ the batch (top-down) ============================ - -def main() -> None: - args = parse_args() - targets = collect_targets(args) - modes = ALL_MODES if args.modes == "all" else [m for m in args.modes.split(",") if m] - - if args.dry_run: # resolve and report the plan without touching a GPU - print_plan(targets, modes, args) - return - - # A parent process (--gpus given, no --shard-index) only fans out; the children do the work. - if args.gpus and args.shard_index is None: - launch_workers(targets, args) - return - - if args.num_shards > 1: # this process is one worker: take its slice, round-robin - targets = targets[args.shard_index :: args.num_shards] - - print(f"targets={[t['name'] for t in targets]} modes={modes} model={args.model}\n" - f"ensemble_size={args.ensemble_size} num_steps={args.num_steps} " - f"outer_steps={args.outer_steps} lr={args.lr} anchor={args.anchor}") - - device, model = get_model_and_device(args.device, args.checkpoint, StructurePredictor(args.model)) - - summary: list[dict] = [] - for target in targets: - print(f"\n######### TARGET: {target['name']} #########") - for mode in modes: - summary.append(generate_one(target, mode, model, device, args)) - - write_summary(summary, shard_summary_path(resolve(args.summary), args.shard_index)) - - -def generate_one(target: dict, mode: str, model, device, args) -> dict: - """Generate + save one ensemble for (target, mode) via run_guidance; return a summary record.""" - name = target["name"] - out_dir = resolve(target["out_dir"]) / mode - try: - guidance_type, extras = guidance_for_mode(mode, args) - config = build_config(target, guidance_type, out_dir, args) - for key, value in extras.items(): # mode-specific args the arg-adders would otherwise set - setattr(config, key, value) - out_dir.mkdir(parents=True, exist_ok=True) - torch.manual_seed(args.seed) # same start noise across modes -> comparable ensembles - job = run_guidance(config, guidance_type, model, device) - status = getattr(job, "status", "unknown") - print(f"[{name}/{mode}] {status} -> {out_dir}/refined.cif") - return {"target": name, "mode": mode, "status": "OK" if status == "success" else "FAILED", - "out_dir": str(out_dir)} - except Exception: - print(f"[{name}/{mode}] FAILED:") - traceback.print_exc() - return {"target": name, "mode": mode, "status": "FAILED"} - - -# ============================== target inputs =============================== -# Added for the ID-and-directory input mode. The CSV/JSON contract below is unchanged -- -# collect_targets() only falls through to the ID mode when --targets is not given. - -def collect_targets(args) -> list[dict]: - """Build the target list from --targets, or from --proteins/--proteins-file, or the default. - - Every target is a dict with name, structure, density, resolution, out_dir -- the same - shape run_guidance has always consumed. Missing input files abort the run here rather - than partway through the batch.""" - output_base = resolve(args.output_base) - if args.targets: - targets = load_targets(args.targets, output_base) - elif args.proteins or args.proteins_file: - targets = targets_from_ids(read_ids(args.proteins, args.proteins_file), output_base, args) - else: - targets = DEFAULT_TARGETS - - missing = [ - f"{t['name']}: {role}={t[role]}" - for t in targets - for role in ("structure", "density") - if not resolve(t[role]).exists() - ] - if missing and not args.allow_missing: - raise SystemExit( - "input files not found (pass --allow-missing to stage a run anyway):\n " - + "\n ".join(missing) - ) - return targets - - -def read_ids(inline: str | None, path: str | None) -> list[str]: - """PDB IDs from a comma-separated --proteins value and/or a --proteins-file (one per - line, '#' comments and blank lines skipped).""" - ids = [i.strip() for i in (inline or "").split(",") if i.strip()] - if path: - for line in resolve(path).read_text().splitlines(): - line = line.split("#")[0].strip() - if line: - ids.append(line) - return ids - - -def targets_from_ids(ids: list[str], output_base: Path, args) -> list[dict]: - """Expand PDB IDs into targets using the structure/density directory templates.""" - return [ - { - "name": fill(args.name_template, pdb), - "structure": str(resolve(args.structure_dir) / fill(args.structure_template, pdb)), - "density": str(resolve(args.density_dir) / fill(args.density_template, pdb)), - "resolution": args.resolution, - "out_dir": str(output_base / fill(args.name_template, pdb)), - } - for pdb in ids - ] - - -def fill(template: str, pdb: str) -> str: - """Substitute one PDB ID into a path/name template, in whichever case it needs.""" - return template.format(pdb=pdb, PDB=pdb.upper(), pdb_low=pdb.lower()) - - -# ================================= fan-out ================================== -# Added so one command can drive several GPUs: the parent re-runs this same script once per -# worker with --shard-index/--num-shards, then merges the per-shard summaries. - -def launch_workers(targets: list[dict], args) -> None: - """Run one child process per (GPU x --jobs-per-gpu) slot, wait, and merge summaries.""" - gpus = [g.strip() for g in args.gpus.split(",") if g.strip()] - n_workers = min(len(gpus) * args.jobs_per_gpu, len(targets)) # no empty workers - log_dir = resolve(args.output_base) / "shards" - log_dir.mkdir(parents=True, exist_ok=True) - summary_path = resolve(args.summary) - - env = dict(os.environ) | thread_env(n_workers) - print(f"fan-out: {len(targets)} targets over {n_workers} workers on GPUs {gpus} " - f"({args.jobs_per_gpu}/GPU), threads/worker={env['OMP_NUM_THREADS']}") - - children = [] - for worker in range(n_workers): - device = f"cuda:{gpus[worker % len(gpus)]}" - log_path = log_dir / f"shard_{worker}.log" - command = [sys.executable, "-u", str(Path(__file__).resolve()), *worker_argv(sys.argv[1:]), - "--shard-index", str(worker), "--num-shards", str(n_workers), - "--device", device, "--summary", str(summary_path)] - print(f" worker {worker} -> {device}, {len(targets[worker::n_workers])} targets, " - f"log {log_path}") - with log_path.open("w") as log: - children.append( - subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT, env=env) - ) - - codes = [child.wait() for child in children] - print(f"workers exited with {codes}") - merge_summaries(summary_path, n_workers) - - -def worker_argv(argv: list[str]) -> list[str]: - """This run's flags with the ones the parent sets per worker removed, in both the - '--flag value' and '--flag=value' spellings, so children never re-fan-out.""" - parent_only = ("--gpus", "--jobs-per-gpu", "--device", "--summary") - kept, skip_value = [], False - for token in argv: - if skip_value: - skip_value = False - elif token in parent_only: - skip_value = True - elif not any(token.startswith(flag + "=") for flag in parent_only): - kept.append(token) - return kept - - -def thread_env(n_workers: int) -> dict[str, str]: - """Per-worker BLAS/OMP thread counts, divided from the cgroup CPU limit. - - On a HAMi vGPU profile os.cpu_count() reports the whole node, not the cgroup quota, so - the thread pools oversubscribe and the run stalls -- that is what hung the first ens8 - attempt. Read /sys/fs/cgroup/cpu.max when it is there and fall back otherwise. - - An OMP_NUM_THREADS already in the environment wins. This function only knows its own - worker count, so two fan-outs sharing a pod (e.g. one arm on GPUs 0-2, another on GPU 3) - would each claim the whole quota; setting the variable by hand is how you split it.""" - if os.environ.get("OMP_NUM_THREADS"): - threads = os.environ["OMP_NUM_THREADS"] - return {"OMP_NUM_THREADS": threads, "MKL_NUM_THREADS": threads, - "OPENBLAS_NUM_THREADS": threads} - - cores = os.cpu_count() or n_workers - try: - quota, period = Path("/sys/fs/cgroup/cpu.max").read_text().split() - if quota != "max": - cores = float(quota) / float(period) - except (OSError, ValueError): - pass - threads = str(max(1, int(cores // n_workers))) - return {"OMP_NUM_THREADS": threads, "MKL_NUM_THREADS": threads, "OPENBLAS_NUM_THREADS": threads} - - -def shard_summary_path(path: Path, shard_index: int | None) -> Path: - """batch_summary.json -> batch_summary_shard2.json, so workers never overwrite each other.""" - if shard_index is None: - return path - return path.with_name(f"{path.stem}_shard{shard_index}{path.suffix}") - - -def merge_summaries(summary_path: Path, n_workers: int) -> None: - """Concatenate the per-shard summaries into the single --summary file.""" - merged: list[dict] = [] - for worker in range(n_workers): - shard_path = shard_summary_path(summary_path, worker) - if shard_path.exists(): - merged.extend(json.loads(shard_path.read_text())) - else: - print(f"WARNING: no summary from worker {worker} ({shard_path}) -- check its log") - write_summary(merged, summary_path) - - -def print_plan(targets: list[dict], modes: list[str], args) -> None: - """--dry-run: show what would run, where it reads from, and where it would be written.""" - n_workers = ( - min(len([g for g in args.gpus.split(",") if g.strip()]) * args.jobs_per_gpu, len(targets)) - if args.gpus else 1 - ) - print(f"PLAN: {len(targets)} targets x {len(modes)} modes = {len(targets) * len(modes)} runs, " - f"{n_workers} worker(s), model={args.model} ensemble_size={args.ensemble_size}\n" - f"modes={modes} summary -> {resolve(args.summary)}") - for index, target in enumerate(targets): - print(f"\n[{index}] {target['name']} worker={index % n_workers} " - f"res={target['resolution']}\n" - f" structure {mark(target['structure'])} {target['structure']}\n" - f" density {mark(target['density'])} {target['density']}\n" - f" out {resolve(target['out_dir'])}//refined.cif") - - -def mark(path: str) -> str: - """'ok' / 'MISSING' tag for one input path in the --dry-run plan.""" - return "ok " if resolve(path).exists() else "MISSING" - - -# ================================= plumbing ================================= - -def guidance_for_mode(mode: str, args) -> tuple[GuidanceType, dict]: - """Map a mode to its GuidanceType and the extra GuidanceConfig attributes it needs.""" - if mode == "baseline": - return GuidanceType.PURE_GUIDANCE, {"step_scaler_type": "none"} - if mode == "coord_guidance": - # The paper's coordinate guidance normalizes the density gradient to the EDM denoising-update - # magnitude (gradient_normalization) and applies AF3 augmentation + realign each step; only - # then does step_size act as a fraction of the denoising step. These are mode-scoped so the - # latent-opt arms are unaffected. - return GuidanceType.PURE_GUIDANCE, { - "step_scaler_type": "noisespace", - "step_size": args.step_size, - "gradient_normalization": args.gradient_normalization, - "augmentation": args.augmentation, - } - if mode in ("s_only", "z_only", "s_plus_z"): - which = {"s_only": "single", "z_only": "pair", "s_plus_z": "both"}[mode] - return GuidanceType.LATENT_OPT, { - "which_latent": which, - "learning_rate": args.lr, - "outer_steps": args.outer_steps, - "anchor_weight": args.anchor, - "max_grad_norm": 1.0, - "bond_length_weight": args.bond_length_weight, - } - raise ValueError(f"unknown mode {mode!r}; pick from {ALL_MODES}") - - -def build_config(target: dict, guidance_type: GuidanceType, out_dir: Path, args) -> GuidanceConfig: - config = GuidanceConfig( - protein=target["name"], - structure=str(resolve(target["structure"])), - density=str(resolve(target["density"])), - model_name=args.model, # GuidanceConfig renamed this field from `model` (merge from main) - guidance_type=guidance_type, - log_path=str(out_dir / "run.log"), - output_dir=str(out_dir), - resolution=float(target["resolution"]), - num_diffusion_steps=args.num_steps, - guidance_start=args.guidance_start, # -1 -> guide from step 0; e.g. 120 -> last low-noise steps - align_to_input=True, - ) - config.ensemble_size = args.ensemble_size # set dynamically (not a declared GuidanceConfig field) - return config - - -def parse_args() -> argparse.Namespace: - ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - ap.add_argument("--targets", help=".csv (name,structure,density,resolution) or .json list of dicts") - ap.add_argument("--output-base", default="it_opt_scratch/targets_out", - help="base output dir for CSV targets: ///") - ap.add_argument("--modes", default="baseline,z_only,coord_guidance", - help=f"comma-separated subset of {ALL_MODES}, or 'all'") - - ids = ap.add_argument_group("target list from PDB IDs (used when --targets is not given)") - ids.add_argument("--proteins", help="comma-separated PDB IDs, e.g. 2YL0,5I09,5MC8") - ids.add_argument("--proteins-file", dest="proteins_file", - help="file of PDB IDs, one per line ('#' comments allowed)") - ids.add_argument("--structure-dir", dest="structure_dir", - default="/home/dev/test_data/processed", - help="root holding the input structures") - ids.add_argument("--structure-template", dest="structure_template", - default="{PDB}/{PDB}_single_001_density_input.cif", - help="path under --structure-dir; takes {pdb}, {PDB}, {pdb_low}") - ids.add_argument("--density-dir", dest="density_dir", default="", - help="directory holding the .ccp4 maps") - ids.add_argument("--density-template", dest="density_template", - default="{PDB}_0.5occA_0.5occB_1.00A.ccp4", - help="filename under --density-dir; takes {pdb}, {PDB}, {pdb_low}") - ids.add_argument("--name-template", dest="name_template", default="{PDB}", - help="target (and output dir) name, e.g. '{PDB}_0.5occA_0.5occB'") - ids.add_argument("--resolution", type=float, default=1.0, - help="resolution for ID-built targets") - ids.add_argument("--allow-missing", dest="allow_missing", action="store_true", - help="do not abort when a structure/density file is absent (e.g. staging " - "a run locally for pod paths)") - - fan = ap.add_argument_group("multi-GPU fan-out") - fan.add_argument("--gpus", help="comma-separated GPU indices, e.g. 0,1,2,3; one child " - "process per GPU x --jobs-per-gpu, then summaries are merged") - fan.add_argument("--jobs-per-gpu", dest="jobs_per_gpu", type=int, default=1, - help="concurrent runs per GPU; 2 needs ~2x the weights resident, " - "so check VRAM") - fan.add_argument("--shard-index", dest="shard_index", type=int, default=None, - help="set by the parent on each worker; targets[shard_index::num_shards]") - fan.add_argument("--num-shards", dest="num_shards", type=int, default=1, - help="set by the parent on each worker") - ap.add_argument("--dry-run", dest="dry_run", action="store_true", - help="print the resolved plan (paths, existence, shards) and exit") - - ap.add_argument("--model", default="protenix", choices=[m.value for m in StructurePredictor]) - ap.add_argument("--ensemble-size", dest="ensemble_size", type=int, default=4) - ap.add_argument("--num-steps", dest="num_steps", type=int, default=200) - ap.add_argument("--guidance-start", dest="guidance_start", type=int, default=-1, - help="step at which guidance begins (coord DPS and IT-opt both use it as a " - "fraction of num_steps); -1 means from step 0. e.g. 120 of 200 guides " - "only the last low-noise steps, matching the tuned coordinate-guidance recipe") - ap.add_argument("--outer-steps", dest="outer_steps", type=int, default=2) - ap.add_argument("--lr", type=float, default=0.05) - ap.add_argument("--anchor", type=float, default=0.0, help="on-manifold anchor weight (IT-opt)") - ap.add_argument("--bond-length-weight", dest="bond_length_weight", type=float, default=0.0, - help="coordinate-space bond-geometry penalty weight (IT-opt); 0 disables") - ap.add_argument("--step-size", dest="step_size", type=float, default=0.1, help="coord-guidance DPS step") - ap.add_argument("--gradient-normalization", dest="gradient_normalization", action="store_true", - help="coord guidance: normalize the density gradient to the denoising-update " - "magnitude before scaling by --step-size (the paper's recipe; makes 0.1 mean " - "10%% of the denoising step rather than 0.1x the raw gradient)") - ap.add_argument("--augmentation", action="store_true", - help="coord guidance: apply AF3 random augmentation + realign each step (paper recipe)") - ap.add_argument("--seed", type=int, default=0) - ap.add_argument("--device", default="cuda:0") - ap.add_argument("--checkpoint", default=None) - ap.add_argument("--summary", default=None, - help="batch summary JSON; defaults to /batch_summary.json") - - args = ap.parse_args() - if args.summary is None: # keep the summary inside the run tree it describes - args.summary = str(resolve(args.output_base) / "batch_summary.json") - return args - - -def resolve(path: str) -> Path: - """Absolute path as-is; relative path is taken relative to the repo root.""" - p = Path(path) - return p if p.is_absolute() else REPO / p - - -def load_targets(path: str, output_base: Path) -> list[dict]: - """Load targets from a .csv (name,structure,density,resolution) or a .json list of dicts. - For CSV rows, out_dir defaults to /.""" - path = Path(path) - if path.suffix.lower() == ".csv": - return [ - { - "name": row["name"], - "structure": row["structure"], - "density": row["density"], - "resolution": float(row["resolution"]), - "out_dir": str(output_base / row["name"]), - } - for row in csv.DictReader(path.read_text().splitlines()) - ] - return json.loads(path.read_text()) - - -def write_summary(summary: list[dict], path: Path) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(summary, indent=2)) - n_ok = sum(1 for s in summary if s.get("status") == "OK") - print(f"\n===== DONE: {n_ok}/{len(summary)} runs OK. Summary -> {path} =====") - - -if __name__ == "__main__": - main() diff --git a/it_opt_scratch/run_targets_simplified.py b/it_opt_scratch/run_targets_simplified.py deleted file mode 100644 index 5481efd3..00000000 --- a/it_opt_scratch/run_targets_simplified.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Generate one conformational ensemble: one structure + one density map -> one output dir. - -This is the refinement step of a longer workflow, so it expects a model that is already sitting -in the map's frame. Starting from a sequence and a map, the usual route is: - - 1. predict sequence -> coordinates (Protenix inference, no map involved) - 2. place put that model in the map's frame (molecular replacement, e.g. Phaser; not - part of this repo) - 3. refine this script - -A deposited structure or an existing MR solution is already placed, so it starts at step 3 -- -which is what the runs in it_opt_scratch/ have done. - -That is also what --structure is for. It carries the sequence and composition (there is no -separate sequence input in this pipeline), the atom identity used for reconciliation, and the -reference frame: the density reward scores coordinates in the map's frame and aligns against -the input structure rather than the map, so a model in an unrelated frame will not score -meaningfully. - -Otherwise, read this as a plain input/output tool -- one run per invocation, no state: - - IN --structure one .cif / .pdb - --density one .ccp4 map - --resolution that map's resolution, in Angstrom - --mode which guidance arm to run (list below) - (everything else is a knob with a default; see --help) - - OUT --output-dir refined.cif trajectory/ losses.txt job_metadata.json run.log - - EXIT the pipeline's own exit code -- 0 on success, non-zero on failure, so a caller - looping over many targets can tell which ones need rerunning. - -The generation itself is not implemented here: this builds a GuidanceConfig and hands it to -the shipped `run_guidance()`, which does featurize -> reward -> scaler -> sample -> save. - -Modes: - baseline unguided sampling (pure_guidance, no step scaler) - s_only IT-opt, which_latent=single - z_only IT-opt, which_latent=pair - s_plus_z IT-opt, which_latent=both - coord_guidance shipped coordinate DPS (pure_guidance + noise-space step scaler) - -One run: - pixi run -e protenix-dev python it_opt_scratch/run_targets_simplified.py \ - --structure /home/dev/test_data/processed/2YL0/2YL0_single_001_density_input.cif \ - --density density_maps/2YL0_0.5occA_0.5occB_1.00A.ccp4 \ - --resolution 1.0 --mode s_plus_z --ensemble-size 8 --bond-length-weight 5e-5 \ - --output-dir out/2YL0_s_plus_z_ens8 --device cuda:0 - -A whole CSV, one after another on one GPU (add --skip-existing to resume a stopped sweep): - tail -n +2 targets.csv | while IFS=, read -r name structure density resolution; do - pixi run -e protenix-dev python it_opt_scratch/run_targets_simplified.py \ - --structure "$structure" --density "$density" --resolution "$resolution" \ - --mode s_plus_z --ensemble-size 8 --bond-length-weight 5e-5 \ - --output-dir "out/$name" --device cuda:0 || echo "FAILED: $name" - done - -Spreading a CSV across several GPUs is deliberately NOT this script's job -- that belongs to a -separate driver that calls this one. -""" - -from __future__ import annotations - -import argparse -import sys -from pathlib import Path - -import torch - -from sampleworks.utils.guidance_constants import GuidanceType, StructurePredictor -from sampleworks.utils.guidance_script_arguments import GuidanceConfig -from sampleworks.utils.guidance_script_utils import get_model_and_device, run_guidance - -MODES = ("baseline", "coord_guidance", "s_only", "z_only", "s_plus_z") -WHICH_LATENT = {"s_only": "single", "z_only": "pair", "s_plus_z": "both"} - - -def main(): - """Read the inputs, run one ensemble generation, return the pipeline's exit code.""" - args = parse_args() - - structure = Path(args.structure).expanduser() - density = Path(args.density).expanduser() - out_dir = Path(args.output_dir).expanduser() - - if not structure.exists(): - sys.exit(f"structure not found: {structure}") - if not density.exists(): - sys.exit(f"density not found: {density}") - - if args.skip_existing and (out_dir / "refined.cif").exists(): - print(f"[skip] already done: {out_dir}/refined.cif") - return 0 - out_dir.mkdir(parents=True, exist_ok=True) - - config = build_config(args, structure, density, out_dir) - device, model = get_model_and_device( - args.device, args.checkpoint, StructurePredictor(args.model) - ) - torch.manual_seed(args.seed) # fixed seed -> the same start noise across modes - - job = run_guidance(config, config.guidance_type, model, device) - print(f"[{args.mode}] {job.status} -> {out_dir}/refined.cif") - return job.exit_code - - -def build_config(args, structure: Path, density: Path, out_dir: Path) -> GuidanceConfig: - """Everything this run will do, in one place: the inputs, the knobs, and the mode. - - The mode-specific fields are assigned directly rather than declared on GuidanceConfig, - which is how the shipped per-guidance-type arg-adders set them too. - """ - # baseline and coord_guidance run the shipped pure_guidance path; the rest optimize latents. - if args.mode in ("baseline", "coord_guidance"): - guidance_type = GuidanceType.PURE_GUIDANCE - else: - guidance_type = GuidanceType.LATENT_OPT - - config = GuidanceConfig( - protein=args.name or structure.stem, - structure=str(structure), - density=str(density), - resolution=args.resolution, - output_dir=str(out_dir), - log_path=str(out_dir / "run.log"), - model_name=args.model, - guidance_type=guidance_type, - num_diffusion_steps=args.num_steps, - align_to_input=True, - ) - config.ensemble_size = args.ensemble_size - - if args.mode == "baseline": # unguided: no step scaler at all - config.step_scaler_type = "none" - elif args.mode == "coord_guidance": # gradients applied in coordinate space - config.step_scaler_type = "noisespace" - config.step_size = args.step_size - else: # s_only / z_only / s_plus_z: optimize the latents instead of the coordinates - config.which_latent = WHICH_LATENT[args.mode] - config.learning_rate = args.lr - config.outer_steps = args.outer_steps - config.anchor_weight = args.anchor - config.max_grad_norm = args.max_grad_norm - config.bond_length_weight = args.bond_length_weight - - return config - - -def parse_args(): - ap = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - inputs = ap.add_argument_group("in / out") - inputs.add_argument("--structure", required=True, help="input .cif / .pdb") - inputs.add_argument("--density", required=True, help="input .ccp4 map") - inputs.add_argument("--resolution", type=float, required=True, help="map resolution, Angstrom") - inputs.add_argument("--output-dir", dest="output_dir", required=True, - help="everything this run produces is written here") - inputs.add_argument("--mode", required=True, choices=MODES) - inputs.add_argument("--name", help="label recorded in job_metadata (default: structure stem)") - inputs.add_argument("--skip-existing", dest="skip_existing", action="store_true", - help="succeed without running if refined.cif is already there") - - run = ap.add_argument_group("sampling") - run.add_argument("--ensemble-size", dest="ensemble_size", type=int, default=4) - run.add_argument("--num-steps", dest="num_steps", type=int, default=200) - run.add_argument("--seed", type=int, default=0) - run.add_argument("--model", default="protenix", choices=[m.value for m in StructurePredictor]) - run.add_argument("--device", default="cuda:0") - run.add_argument("--checkpoint", default=None) - - it = ap.add_argument_group("IT-opt modes only (s_only / z_only / s_plus_z)") - it.add_argument("--lr", type=float, default=0.05) - it.add_argument("--outer-steps", dest="outer_steps", type=int, default=2) - it.add_argument("--anchor", type=float, default=0.0, help="on-manifold anchor weight") - it.add_argument("--bond-length-weight", dest="bond_length_weight", type=float, default=0.0, - help="coordinate-space bond-geometry penalty; 0 disables") - it.add_argument("--max-grad-norm", dest="max_grad_norm", type=float, default=1.0) - - ap.add_argument("--step-size", dest="step_size", type=float, default=0.1, - help="coord_guidance mode only: DPS step size") - return ap.parse_args() - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/it_opt_scratch/score_paper_rmsd.py b/it_opt_scratch/score_paper_rmsd.py deleted file mode 100644 index 60aee1b6..00000000 --- a/it_opt_scratch/score_paper_rmsd.py +++ /dev/null @@ -1,298 +0,0 @@ -"""Score IT-opt arms with the paper's min-altloc-RMSD metric (Figure 3E). - -The paper computes, per altloc selection, the *minimum* RMSD over the predicted ensemble to -altloc A and to altloc B of the deposited reference, separately. Its purpose is diagnostic: an -ensemble can raise RSCC by fitting one altloc well, and only the pair (min-to-A, min-to-B) -reveals whether both conformations were reached. - -Scope matches the paper's RSCC protocol -- the same 3-residue max-RMSD subsegments -- and the -prediction is placed by the same global uniform-weight Kabsch used in ``score_paper_rscc.py``. -That is deliberate: aligning on the subsegment itself would let a wrong local conformation be -rotated into agreement. - -Per-altloc conformers are read with gemmi, which exposes the altloc character directly. -``load_any`` drops the altloc annotation entirely and ``parse`` keeps only the first altloc, -so neither can separate the two conformers. Atoms with a blank altloc are shared and belong to -both conformers. - -Comparator: ``occ_sweep_results/min_altloc_rmsd_results.csv`` (protenix, 0.5/0.5), already -computed by the published pipeline over the same selections. - -Usage ------ - python it_opt_scratch/score_paper_rmsd.py \ - --runs-dir it_opt_scratch/patch_tree --dir-template "{protein}" \ - --target-filename refined-patched.cif \ - --inputs-dir /home/dev/test_data \ - --selections-csv it_opt_scratch/paper_maxrmsd_selections.csv \ - --arms baseline coord_guidance s_only s_plus_z z_only \ - --out it_opt_scratch/patch_tree/itopt_paper_rmsd.csv -""" - -from __future__ import annotations - -import argparse -import csv -import re -import traceback -from pathlib import Path - -import gemmi -import numpy as np -import pandas as pd -import torch -from atomworks.io.utils.io_utils import load_any -from loguru import logger -from sampleworks.utils.atom_array_utils import ( - make_atom_id, - make_normalized_atom_id, - remove_atoms_with_any_nan_coords, -) -from sampleworks.utils.cif_utils import remap_altlocs_to_ab, resolve_mixed_hetatm_atom_altlocs - -# Same directory as this script, which is sys.path[0] when run as `python it_opt_scratch/...`. -from score_paper_rscc import ( - align_prediction_to_reference, - read_selections, - strict_frame_is_consistent, -) - -SELECTION_RE = re.compile(r"^chain\s+(\S+)\s+and\s+resi\s+(-?\d+)\s*-\s*(-?\d+)$") - - -def parse_selection(selection: str) -> tuple[str, range]: - """Split ``chain A and resi 12-14`` into ``("A", range(12, 15))``.""" - m = SELECTION_RE.match(selection.strip()) - if m is None: - raise ValueError(f"unparseable selection: {selection!r}") - chain, lo, hi = m.group(1), int(m.group(2)), int(m.group(3)) - return chain, range(lo, hi + 1) - - -def reference_conformers( - ref_path: Path, selections: list[str], normalize: bool = False -) -> dict[str, dict[str, dict]]: - """Per selection, the altloc-A and altloc-B conformers keyed by atom id. - - Atoms with a blank altloc are shared between conformers, so they appear in both. - - ``normalize=False`` keys by deposited ``(chain, res_id, atom_name)`` -- the default. - ``normalize=True`` keys by ``make_normalized_atom_id`` (the same normalization - ``filter_to_common_atoms(normalize_ids=True)`` uses for RSCC), a fallback for a relabelled - reference. gemmi reads coordinates because ``load_any`` drops the altloc character; the biotite - load only builds the deposited-key -> normalized-key map. - """ - st = gemmi.read_structure(str(ref_path)) - raw_to_norm = None - if normalize: - ref_bio = load_any(str(ref_path), altloc="all") - raw_to_norm = dict( - zip(make_atom_id(ref_bio), make_normalized_atom_id(ref_bio), strict=True) - ) - - by_res: dict[tuple[str, int], list] = {} - for chain in st[0]: - for res in chain: - by_res.setdefault((chain.name, res.seqid.num), []).extend( - (a.name, a.altloc, np.array([a.pos.x, a.pos.y, a.pos.z])) for a in res - ) - - out: dict[str, dict[str, dict]] = {} - for sel in selections: - chain, residues = parse_selection(sel) - conformers: dict[str, dict] = {"A": {}, "B": {}} - for res_id in residues: - for name, altloc, xyz in by_res.get((chain, res_id), []): - if normalize: - key = raw_to_norm.get(f"{chain}_{res_id}_{name}") - if key is None: - continue # atom absent from the biotite load - else: - key = (chain, res_id, name) - alt = altloc.strip() - targets = ("A", "B") if alt == "" else (alt,) - for t in targets: - if t in conformers: - conformers[t][key] = xyz - out[sel] = conformers - return out - - -def prediction_lookup(atom_array, normalize: bool = False) -> dict: - """Map atom id -> per-model coordinates, shape [n_models, 3]. - - ``normalize=False`` keys by raw ``(chain, res_id, atom_name)`` -- the default, so a prediction - whose chain/numbering already match the reference scores exactly as before. ``normalize=True`` - keys by ``make_normalized_atom_id`` (sequential per-chain numbering + chain-index), used only as - a fallback for a relabelled/renumbered prediction (Protenix chain 'A' vs deposited 'P'). - """ - coords = atom_array.coord - if coords.ndim == 2: # single model -> add the model axis - coords = coords[None] - if normalize: - ids = make_normalized_atom_id(atom_array) - return {ids[i]: coords[:, i] for i in range(len(ids))} - return { - (str(c), int(r), str(n)): coords[:, i] - for i, (c, r, n) in enumerate( - zip(atom_array.chain_id, atom_array.res_id, atom_array.atom_name, strict=True) - ) - } - - -def min_rmsd_over_ensemble( - conformer: dict[tuple[str, int, str], np.ndarray], - pred: dict[tuple[str, int, str], np.ndarray], -) -> tuple[float, int]: - """Minimum over ensemble members of RMSD to one reference conformer, plus atoms matched.""" - keys = [k for k in conformer if k in pred] - if not keys: - return float("nan"), 0 - ref = np.stack([conformer[k] for k in keys]) # [n_atoms, 3] - prd = np.stack([pred[k] for k in keys], axis=1) # [n_models, n_atoms, 3] - per_model = np.sqrt(((prd - ref[None]) ** 2).sum(-1).mean(-1)) # [n_models] - return float(per_model.min()), len(keys) - - -def score_protein( - protein: str, - selections: list[str], - runs_dir: Path, - inputs_dir: Path, - arms: list[str], - target_filename: str, - dir_template: str, -) -> list[dict]: - rows: list[dict] = [] - ref_path = inputs_dir / "processed" / protein / f"{protein}_single_001_density_input.cif" - # Match generation + the RSCC scorer: collapse mixed ATOM/HETATM modified-residue positions - # (e.g. CYS+CSO) so the reference has the same atoms as a prediction from the cleaned CIF. - ref_path = resolve_mixed_hetatm_atom_altlocs(ref_path) - # Relabel non-A/B alternate altlocs (e.g. A/C) to A/B so the A/B-only conformer split below - # does not drop the second conformer (7Z0E, 7AVG). No-op for already-A/B structures. - ref_path = remap_altlocs_to_ab(ref_path) - - def fail(arm: str, err: str) -> None: - for sel in selections: - rows.append( - { - "protein": protein, - "arm": arm, - "selection": sel, - "min_rmsd_to_A": np.nan, - "min_rmsd_to_B": np.nan, - "n_atoms_A": 0, - "n_atoms_B": 0, - "error": err, - } - ) - - try: - conformers_raw = reference_conformers(ref_path, selections, normalize=False) - # Alignment target: the same array score_paper_rscc.py aligns against, so both metrics - # place the prediction identically. - ref_atom_array = remove_atoms_with_any_nan_coords(load_any(str(ref_path))) - except Exception as e: # noqa: BLE001 - logger.error(f"{protein}: setup failed: {e}\n{traceback.format_exc()}") - for arm in arms: - fail(arm, f"setup: {e}") - return rows - - conformers_norm = None # built lazily, only if a prediction needs the fallback - - for arm in arms: - cif = runs_dir / dir_template.format(protein=protein) / arm / target_filename - if not cif.exists(): - fail(arm, f"{target_filename} missing") - continue - try: - aa = remove_atoms_with_any_nan_coords(load_any(str(cif))) - aa = align_prediction_to_reference(ref_atom_array, aa) - # Key on deposited (chain,res,name) only when those identifiers demonstrably pair the - # same residues. Sharing *some* keys is not enough: a prediction renumbered from 1 - # overlaps a deposited range like 8-258 on most of its residues while denoting - # different ones, which silently compared residue i against residue i-offset. - if strict_frame_is_consistent(ref_atom_array, aa): - conformers = conformers_raw - pred = prediction_lookup(aa, normalize=False) - else: - if conformers_norm is None: - conformers_norm = reference_conformers(ref_path, selections, normalize=True) - conformers = conformers_norm - pred = prediction_lookup(aa, normalize=True) - except Exception as e: # noqa: BLE001 - logger.error(f"{protein}/{arm}: {e}\n{traceback.format_exc()}") - fail(arm, str(e)) - continue - - for sel in selections: - rmsd_a, n_a = min_rmsd_over_ensemble(conformers[sel]["A"], pred) - rmsd_b, n_b = min_rmsd_over_ensemble(conformers[sel]["B"], pred) - rows.append( - { - "protein": protein, - "arm": arm, - "selection": sel, - "min_rmsd_to_A": rmsd_a, - "min_rmsd_to_B": rmsd_b, - "n_atoms_A": n_a, - "n_atoms_B": n_b, - "error": None if n_a and n_b else "no matching atoms", - } - ) - return rows - - -def main() -> None: - p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--runs-dir", type=Path, required=True) - p.add_argument("--inputs-dir", type=Path, required=True) - p.add_argument("--selections-csv", type=Path, required=True) - p.add_argument("--arms", nargs="+", default=["baseline", "s_only", "s_plus_z", "z_only"]) - p.add_argument("--proteins", nargs="+", default=None) - p.add_argument("--dir-template", default="{protein}_native_occ") - p.add_argument("--target-filename", default="refined.cif") - p.add_argument("--out", type=Path, required=True) - args = p.parse_args() - - selections = read_selections(args.selections_csv) - if args.proteins: - wanted = {x.upper() for x in args.proteins} - selections = {k: v for k, v in selections.items() if k in wanted} - logger.info( - f"{len(selections)} proteins, {sum(len(v) for v in selections.values())} selections, " - f"arms={args.arms}" - ) - - rows: list[dict] = [] - for i, (protein, sels) in enumerate(sorted(selections.items()), 1): - logger.info(f"[{i}/{len(selections)}] {protein} ({len(sels)} selections)") - rows.extend( - score_protein( - protein, sels, args.runs_dir, args.inputs_dir, args.arms, - args.target_filename, args.dir_template, - ) - ) - pd.DataFrame(rows).to_csv(args.out, index=False) - - df = pd.DataFrame(rows) - df.to_csv(args.out, index=False) - ok = df[df.min_rmsd_to_A.notna() & df.min_rmsd_to_B.notna()] - logger.info(f"wrote {args.out}: {len(df)} rows, {len(ok)} scored") - if not ok.empty: - nearer = ok[["min_rmsd_to_A", "min_rmsd_to_B"]].min(axis=1) - worse = ok[["min_rmsd_to_A", "min_rmsd_to_B"]].max(axis=1) - summary = pd.DataFrame( - { - "n": ok.groupby("arm").size(), - "med_nearer": nearer.groupby(ok.arm).median(), - "med_max": worse.groupby(ok.arm).median(), - "max_le_2A": worse.le(2.0).groupby(ok.arm).mean(), - "max_le_1A": worse.le(1.0).groupby(ok.arm).mean(), - } - ) - logger.info(f"\n{summary.to_string()}") - - -if __name__ == "__main__": - main() diff --git a/it_opt_scratch/score_paper_rscc.py b/it_opt_scratch/score_paper_rscc.py deleted file mode 100644 index e8f551de..00000000 --- a/it_opt_scratch/score_paper_rscc.py +++ /dev/null @@ -1,360 +0,0 @@ -"""Score IT-opt arms with the paper's RSCC protocol (791 max-RMSD subsegments). - -The published scorer is ``scripts/eval/rscc_grid_search_script.py``. This driver reuses its -exact primitives -- the differentiable density forward model, the global uniform-weight Kabsch -alignment, and ``extract_tight`` at 2.0 A -- so the numbers are directly comparable to the -occupancy-sweep results. What it does *not* reuse is the Trial/ProteinConfig directory scanner, -because the IT-opt output tree is ``{PROTEIN}_native_occ/{arm}/refined.cif`` (depth 3) rather -than the ``{PROTEIN}_{occ}/{model}/{scaler}/ens{N}_gw{W}/`` depth-4 grid-search layout, and -``native_occ`` does not parse as an occupancy key. - -Selections come from the segmentation CSV (one row per protein, semicolon-joined), which encodes -the paper's rule: the contiguous 3-residue window maximising altloc A-B RMSD, or the whole -segment when it is 3 residues or shorter. - -The RSCC mask is built from reference coordinates including *all* altlocs, matching -``get_reference_structure_coords``, which unions the altloc-A-only and altloc-B-only coordinate -sets for exactly this purpose. - -Usage ------ - pixi run -e analysis python it_opt_scratch/score_paper_rscc.py \ - --runs-dir it_opt_scratch/targets_out_40 \ - --inputs-dir /home/dev/test_data \ - --selections-csv it_opt_scratch/paper_maxrmsd_selections.csv \ - --arms baseline s_only s_plus_z z_only \ - --out it_opt_scratch/patch_tree/itopt_paper_rscc.csv -""" - -from __future__ import annotations - -import argparse -import copy -import csv -import traceback -from pathlib import Path - -import numpy as np -import pandas as pd -import torch -from atomworks.io.utils.io_utils import load_any -from loguru import logger -from sampleworks.eval.constants import DEFAULT_SELECTION_PADDING -from sampleworks.eval.metrics import rscc -from sampleworks.eval.structure_utils import extract_selection_coordinates -from sampleworks.utils.atom_array_utils import ( - filter_to_common_atoms, - remove_atoms_with_any_nan_coords, -) -from sampleworks.utils.cif_utils import resolve_mixed_hetatm_atom_altlocs -from sampleworks.utils.density_utils import build_density_transformer, run_density_transformer -from sampleworks.utils.frame_transforms import ( - apply_forward_transform, - weighted_rigid_align_differentiable, -) -from sampleworks.utils.framework_utils import match_batch - -from sampleworks.core.forward_models.xray.real_space_density_deps.qfit.volume import XMap - -RESOLUTION = 1.0 - - -def read_selections(csv_path: Path) -> dict[str, list[str]]: - """Read the per-protein segmentation CSV into {PROTEIN: [selection, ...]}.""" - out: dict[str, list[str]] = {} - with open(csv_path) as fh: - for row in csv.DictReader(fh): - sels = [s.strip() for s in row["selection"].split(";") if s.strip()] - if sels: - out[row["protein"].strip().upper()] = sels - return out - - -MIN_RESIDUE_IDENTITY = 0.95 # matched pairs below this share a numbering frame only by accident - - -def strict_frame_is_consistent(ref_atom_array, pred_atom_array) -> bool: - """Whether prediction and reference can be paired on deposited (chain, res_id, atom_name). - - ``filter_to_common_atoms`` pairs atoms by identifier alone, never checking 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 succeeds on a large but sequence-shifted set, - pairing prediction residue i with reference residue i-offset. Kabsch then fits that shifted - correspondence and reports a fold-scale error for a structure that is actually correct. - - Comparing residue names separates the two cases cleanly: a correct frame agrees on every - pair, a shifted one agrees only at the ~5% rate expected from 20 residue types by chance. - - Parameters - ---------- - ref_atom_array, pred_atom_array - Reference and predicted structures, as loaded by ``load_any``. - - Returns - ------- - bool - True when strict matching pairs chemically identical residues and so may be used - directly; False when the caller should fall back to normalized (sequential) ids. - """ - try: - ref_common, pred_common = filter_to_common_atoms(ref_atom_array, pred_atom_array) - except RuntimeError: - return False # no shared identifiers at all, e.g. deposited chain 'B' vs Protenix 'A' - identity = residue_identity(ref_common, pred_common) - return identity is not None and identity >= MIN_RESIDUE_IDENTITY - - -def residue_identity(ref_common, pred_common) -> float | None: - """Fraction of matched pairs whose residue names agree, or None if they cannot be paired. - - ``filter_to_common_atoms`` masks with ``np.isin``, so a reference carrying two altlocs of one - atom contributes both copies of a duplicated (chain, res_id, atom_name) key while a - single-conformer prediction contributes one. The two matched sets then differ in length and - there is no atom-for-atom correspondence to score at all (6NI5 and 6NI6 are the known cases). - Returning None keeps that distinct from "paired, but the residues disagree". - """ - ref_names = np.asarray(ref_common.res_name) - pred_names = np.asarray(pred_common.res_name) - if ref_names.shape != pred_names.shape: - return None - return float((ref_names == pred_names).mean()) - - -def align_prediction_to_reference(ref_atom_array, pred_atom_array): - """Global uniform-weight Kabsch of prediction onto reference, applied to every predicted atom. - - Mirrors rscc_grid_search_script.py: the transform is fitted on the atoms common to both - structures, then applied to the whole predicted array. Alignment is deliberately global -- - fitting on the 3-residue window itself would let a wrong local conformation be rotated into - apparent agreement. - """ - # Sequential per-chain matching is used unless the deposited identifiers demonstrably pair the - # same residues, so a partial numbering overlap can no longer produce a silently shifted fit. - normalize = not strict_frame_is_consistent(ref_atom_array, pred_atom_array) - ref_common, pred_common = filter_to_common_atoms( - ref_atom_array, pred_atom_array, normalize_ids=normalize - ) - if normalize: - identity = residue_identity(ref_common, pred_common) - # None means the matched sets are different sizes; the explicit shape check below reports - # that with the actual atom counts, which is the more useful message for it. - if identity is not None and identity < MIN_RESIDUE_IDENTITY: - raise ValueError( - f"neither deposited nor sequential atom ids pair matching residues " - f"(sequential identity {identity:.2f}); the reference and the prediction are " - f"not the same construct" - ) - ref_t = torch.from_numpy(ref_common.coord).float() - pred_t = torch.from_numpy(pred_common.coord).float() - ref_t = match_batch(ref_t, pred_t.shape[0]) - if ref_t.ndim != 3 or ref_t.shape[1] != pred_t.shape[1]: - raise ValueError(f"shape mismatch: ref {tuple(ref_t.shape)} vs pred {tuple(pred_t.shape)}") - - n_atoms = ref_t.shape[1] - _, transform = weighted_rigid_align_differentiable( - true_coords=pred_t, - pred_coords=ref_t, - weights=torch.ones(1, n_atoms), - mask=torch.ones(1, n_atoms), - return_transforms=True, - allow_gradients=False, - ) - moved = apply_forward_transform( - torch.from_numpy(pred_atom_array.coord), transform, rotation_only=False - ) - pred_atom_array.coord = moved.numpy() - return pred_atom_array - - -def ensure_density_annotations(atom_array): - """Add the occupancy / b_factor annotations the density forward model requires.""" - n = atom_array.coord.shape[-2] - for name, value in (("occupancy", 1.0), ("b_factor", 20.0)): - if name not in atom_array.get_annotation_categories(): - atom_array.set_annotation(name, np.full(n, value)) - return atom_array - - -def score_protein( - protein: str, - selections: list[str], - runs_dir: Path, - inputs_dir: Path, - arms: list[str], - device: torch.device, - target_filename: str = "refined.cif", - dir_template: str = "{protein}_native_occ", - maps_dir: Path | None = None, - map_template: str = "{protein}_uniform_1.00A.ccp4", -) -> list[dict]: - """Compute per-(arm, selection) RSCC for one protein.""" - rows: list[dict] = [] - maps_base = maps_dir if maps_dir is not None else inputs_dir / "density_maps" - map_path = maps_base / map_template.format(protein=protein) - ref_path = inputs_dir / "processed" / protein / f"{protein}_single_001_density_input.cif" - # Match generation: collapse modified-residue positions (mixed ATOM/HETATM, different resname, - # e.g. CYS+CSO) that atomworks would otherwise duplicate into an extra residue, so the reference - # carries the same atoms as a prediction generated from the cleaned CIF. No-op otherwise. - ref_path = resolve_mixed_hetatm_atom_altlocs(ref_path) - - def fail(arm: str, err: str) -> None: - for sel in selections: - rows.append( - { - "protein": protein, - "arm": arm, - "selection": sel, - "rscc": np.nan, - "error": err, - "base_map_path": str(map_path), - } - ) - - try: - base_xmap = XMap.fromfile(str(map_path), resolution=RESOLUTION).canonical_unit_cell() - transformer, _ = build_density_transformer(base_xmap, em_mode=False, device=device) - # Mask coords keep every altloc: load_any preserves them, whereas parse() drops all but - # the first. The mask must cover both conformations, as the published scorer does by - # unioning the A-only and B-only reference structures. - ref_all_altlocs = load_any(str(ref_path)) - sel_coords = {} - for sel in selections: - try: - coords = extract_selection_coordinates(ref_all_altlocs, sel) - except Exception as e: # noqa: BLE001 - selection syntax varies per protein - logger.warning(f"{protein}: selection {sel!r} failed: {e}") - continue - if len(coords) and np.isfinite(coords).all(): - sel_coords[sel] = coords - else: - logger.warning(f"{protein}: selection {sel!r} empty or non-finite") - # Alignment target. load_any, not parse: parse() reconciles against the - # pdbx_poly_seq_scheme that the patcher inherits from the deposited RCSB entry, and with - # atomworks 2.1.1 and no CCD mirror that turns ~93% of atoms into NaN placeholders. - # load_any reads atom_site as written, and is the reader the patcher itself uses. - ref_atom_array = remove_atoms_with_any_nan_coords(load_any(str(ref_path))) - except Exception as e: # noqa: BLE001 - per-protein setup failure should not kill the sweep - logger.error(f"{protein}: setup failed: {e}\n{traceback.format_exc()}") - for arm in arms: - fail(arm, f"setup: {e}") - return rows - - base_cache: dict[str, np.ndarray] = {} - for arm in arms: - cif = runs_dir / dir_template.format(protein=protein) / arm / target_filename - if not cif.exists(): - fail(arm, f"{target_filename} missing") - continue - try: - aa = remove_atoms_with_any_nan_coords(load_any(str(cif))) - # load_any does not populate these annotations even when the CIF carries the columns; - # the forward model requires both. The written values are occupancy 1.0 and B 20.0, - # which is also what parse() supplied to the published scorer. - aa = ensure_density_annotations(aa) - aa = align_prediction_to_reference(ref_atom_array, aa) - # One density from the whole ensemble: RSCC scores the multi-conformer model, not - # individual members. - computed = run_density_transformer(transformer, aa) - computed_xmap = copy.copy(base_xmap) - computed_xmap.array = computed.cpu().numpy() - if computed_xmap.array.shape != base_xmap.array.shape: - raise ValueError( - f"density shape {computed_xmap.array.shape} != base {base_xmap.array.shape}" - ) - except Exception as e: # noqa: BLE001 - logger.error(f"{protein}/{arm}: {e}\n{traceback.format_exc()}") - fail(arm, str(e)) - continue - - for sel, coords in sel_coords.items(): - row = { - "protein": protein, - "arm": arm, - "selection": sel, - "base_map_path": str(map_path), - "error": None, - } - try: - extracted_base = base_cache.get(sel) - if extracted_base is None: - _, extracted_base = base_xmap.extract_tight( - coords, padding=DEFAULT_SELECTION_PADDING - ) - if extracted_base is None or extracted_base.shape[0] == 0: - raise ValueError("extracted base map empty") - base_cache[sel] = extracted_base - _, extracted_computed = computed_xmap.extract_tight( - coords, padding=DEFAULT_SELECTION_PADDING - ) - if extracted_computed is None or extracted_computed.shape[0] == 0: - raise ValueError("extracted computed map empty") - row["rscc"] = rscc(extracted_base, extracted_computed) - except Exception as e: # noqa: BLE001 - row["error"] = str(e) - row["rscc"] = np.nan - rows.append(row) - return rows - - -def main() -> None: - p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--runs-dir", type=Path, required=True) - p.add_argument("--inputs-dir", type=Path, required=True) - p.add_argument("--selections-csv", type=Path, required=True) - p.add_argument("--arms", nargs="+", default=["baseline", "s_only", "s_plus_z", "z_only"]) - p.add_argument("--proteins", nargs="+", default=None, help="Subset; default all in the CSV.") - p.add_argument("--dir-template", default="{protein}_native_occ", - help="Per-protein dir name under --runs-dir; '{protein}' is substituted.") - p.add_argument("--target-filename", default="refined.cif", - help="CIF to score in each arm dir; use refined-patched.cif after patching.") - p.add_argument("--maps-dir", type=Path, default=None, - help="Dir holding the density maps; default /density_maps.") - p.add_argument("--map-template", default="{protein}_uniform_1.00A.ccp4", - help="Map filename template; '{protein}' is substituted. Use " - "'{protein}_0.5occA_0.5occB_1.00A.ccp4' for the 0.5/0.5 occupancy maps.") - p.add_argument("--out", type=Path, required=True) - args = p.parse_args() - - selections = read_selections(args.selections_csv) - if args.proteins: - wanted = {x.upper() for x in args.proteins} - selections = {k: v for k, v in selections.items() if k in wanted} - total = sum(len(v) for v in selections.values()) - logger.info(f"{len(selections)} proteins, {total} selections, arms={args.arms}") - - device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - logger.info(f"device={device}") - - rows: list[dict] = [] - for i, (protein, sels) in enumerate(sorted(selections.items()), 1): - logger.info(f"[{i}/{len(selections)}] {protein} ({len(sels)} selections)") - rows.extend( - score_protein( - protein, sels, args.runs_dir, args.inputs_dir, args.arms, device, - target_filename=args.target_filename, - dir_template=args.dir_template, - maps_dir=args.maps_dir, - map_template=args.map_template, - ) - ) - pd.DataFrame(rows).to_csv(args.out, index=False) # checkpoint after each protein - - df = pd.DataFrame(rows) - df.to_csv(args.out, index=False) - logger.info(f"wrote {args.out}: {len(df)} rows") - - ok = df[df["rscc"].notna()] - logger.info(f"scored {len(ok)}/{len(df)} rows") - if not ok.empty: - summary = ok.groupby("arm")["rscc"].agg( - n="size", - median="median", - frac_ge_08=lambda s: (s >= 0.8).mean(), - frac_ge_09=lambda s: (s >= 0.9).mean(), - ) - logger.info(f"\n{summary.to_string()}") - - -if __name__ == "__main__": - main() diff --git a/it_opt_scratch/score_paper_simplified.py b/it_opt_scratch/score_paper_simplified.py deleted file mode 100644 index 5d69c6a3..00000000 --- a/it_opt_scratch/score_paper_simplified.py +++ /dev/null @@ -1,154 +0,0 @@ -"""Score one run tree with both paper metrics, in one command and one CSV. - -This is a thin driver, not a new metric. RSCC comes from score_paper_rscc.py and min-altloc-RMSD -from score_paper_rmsd.py, both called unchanged -- so the numbers are identical to running those -two scripts separately. It exists so you remember one command instead of two with eight matching -flags, and get one table instead of two you have to join by hand. - - IN --runs-dir tree of // - --dir-template per-protein dir name, e.g. '{protein}_0.5occA_0.5occB' - --arms which arm sub-dirs to score - --inputs-dir holds processed/{PROTEIN}/{PROTEIN}_single_001_density_input.cif - --maps-dir holds the density maps (default /density_maps) - --map-template map filename, e.g. '{protein}_0.5occA_0.5occB_1.00A.ccp4' - --selections-csv the paper's per-protein 3-residue max-RMSD windows - - OUT --out one row per (protein, arm, selection): - rscc, min_rmsd_to_A, min_rmsd_to_B, n_atoms_A, n_atoms_B, - rscc_error, rmsd_error, base_map_path - rewritten after every protein, so a long sweep is resumable-by-eye - -Run it on the 11-protein regen tree: - pixi run -e analysis python it_opt_scratch/score_paper_simplified.py \ - --runs-dir it_opt_scratch/targets_out_11_regen_0.5occA_0.5occB_ens8 \ - --dir-template '{protein}_0.5occA_0.5occB' \ - --maps-dir it_opt_scratch/targets_out_11_regen_0.5occA_0.5occB/density_maps \ - --map-template '{protein}_0.5occA_0.5occB_1.00A.ccp4' \ - --inputs-dir /home/dev/test_data \ - --selections-csv it_opt_scratch/paper_maxrmsd_selections.csv \ - --arms s_plus_z \ - --out it_opt_scratch/regen11_scores.csv - -Each prediction is loaded and aligned once per metric rather than once in total. That is a -deliberate trade: calling the two scorers as they are keeps this file honest about the numbers, -and scoring is seconds per arm against minutes per generation run. -""" - -from __future__ import annotations - -import argparse -from pathlib import Path - -import pandas as pd -import torch -from loguru import logger - -# Same directory as this script, which is sys.path[0] when run as `python it_opt_scratch/...`. -import score_paper_rmsd -import score_paper_rscc - -KEY = ["protein", "arm", "selection"] - - -def main() -> None: - args = parse_args() - - selections = score_paper_rscc.read_selections(args.selections_csv) - if args.proteins: - wanted = {p.upper() for p in args.proteins} - selections = {k: v for k, v in selections.items() if k in wanted} - - device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - logger.info( - f"{len(selections)} proteins, {sum(len(v) for v in selections.values())} selections, " - f"arms={args.arms}, device={device}" - ) - - scored = [] - for i, (protein, sels) in enumerate(sorted(selections.items()), 1): - logger.info(f"[{i}/{len(selections)}] {protein} ({len(sels)} selections)") - scored.append(score_one_protein(protein, sels, args, device)) - pd.concat(scored).to_csv(args.out, index=False) # checkpoint after every protein - - table = pd.concat(scored) - table.to_csv(args.out, index=False) - logger.info(f"wrote {args.out}: {len(table)} rows") - report(table) - - -def score_one_protein(protein: str, sels: list[str], args, device) -> pd.DataFrame: - """Both metrics for one protein, joined on (protein, arm, selection). - - An outer join because the two scorers can disagree on which selections are scoreable: the - RSCC side drops a selection whose reference coordinates are empty or non-finite, while the - RMSD side still emits a row for it. - """ - rscc_rows = score_paper_rscc.score_protein( - protein, sels, args.runs_dir, args.inputs_dir, args.arms, device, - target_filename=args.target_filename, - dir_template=args.dir_template, - maps_dir=args.maps_dir, - map_template=args.map_template, - ) - rmsd_rows = score_paper_rmsd.score_protein( - protein, sels, args.runs_dir, args.inputs_dir, args.arms, - args.target_filename, args.dir_template, - ) - rscc_table = pd.DataFrame(rscc_rows).rename(columns={"error": "rscc_error"}) - rmsd_table = pd.DataFrame(rmsd_rows).rename(columns={"error": "rmsd_error"}) - return pd.merge(rscc_table, rmsd_table, on=KEY, how="outer") - - -def report(table: pd.DataFrame) -> None: - """Per-arm summary, in the same terms the two original scorers print.""" - rscc_ok = table[table["rscc"].notna()] - if not rscc_ok.empty: - summary = rscc_ok.groupby("arm")["rscc"].agg( - n="size", - median="median", - frac_ge_08=lambda s: (s >= 0.8).mean(), - frac_ge_09=lambda s: (s >= 0.9).mean(), - ) - logger.info(f"RSCC\n{summary.to_string()}") - - both = ["min_rmsd_to_A", "min_rmsd_to_B"] - rmsd_ok = table[table[both].notna().all(axis=1)] - if not rmsd_ok.empty: - nearer = rmsd_ok[both].min(axis=1) # the altloc the ensemble reached - worse = rmsd_ok[both].max(axis=1) # the one it had to also reach to score well - summary = pd.DataFrame( - { - "n": rmsd_ok.groupby("arm").size(), - "med_nearer": nearer.groupby(rmsd_ok.arm).median(), - "med_max": worse.groupby(rmsd_ok.arm).median(), - "max_le_2A": worse.le(2.0).groupby(rmsd_ok.arm).mean(), - "max_le_1A": worse.le(1.0).groupby(rmsd_ok.arm).mean(), - } - ) - logger.info(f"min-altloc-RMSD\n{summary.to_string()}") - - -def parse_args() -> argparse.Namespace: - ap = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - ap.add_argument("--runs-dir", type=Path, required=True) - ap.add_argument("--inputs-dir", type=Path, required=True) - ap.add_argument("--selections-csv", type=Path, required=True) - ap.add_argument("--out", type=Path, required=True) - ap.add_argument("--arms", nargs="+", default=["baseline", "s_only", "s_plus_z", "z_only"]) - ap.add_argument("--proteins", nargs="+", default=None, help="subset; default all in the CSV") - ap.add_argument("--dir-template", default="{protein}_native_occ", - help="per-protein dir under --runs-dir; '{protein}' is substituted") - ap.add_argument("--target-filename", default="refined.cif", - help="CIF to score in each arm dir; refined-patched.cif after patching") - ap.add_argument("--maps-dir", type=Path, default=None, - help="dir holding the maps; default /density_maps") - ap.add_argument("--map-template", default="{protein}_uniform_1.00A.ccp4", - help="map filename; use '{protein}_0.5occA_0.5occB_1.00A.ccp4' for the " - "paper's 0.5/0.5 occupancy maps") - return ap.parse_args() - - -if __name__ == "__main__": - main() diff --git a/it_opt_scratch/score_rmsd_simplified.py b/it_opt_scratch/score_rmsd_simplified.py deleted file mode 100644 index 25a56c2e..00000000 --- a/it_opt_scratch/score_rmsd_simplified.py +++ /dev/null @@ -1,259 +0,0 @@ -"""Min-altloc-RMSD for one prediction against a two-conformer reference. Self-contained. - -Answers: did the ensemble actually reach BOTH conformations the crystal shows, or just one? - -For each of the paper's 3-residue windows, this reports two numbers -- the closest any ensemble -member gets to altloc A, and the closest any member gets to altloc B, measured separately. The -pair is the point. RSCC can look good when an ensemble fits one conformer well and ignores the -other; only the pair reveals whether both were found. - - IN --prediction refined.cif from a run (multi-model = the ensemble) - --reference the deposited .cif, which must contain altloc A and B - --selections-csv the paper's per-protein windows - --protein which row of that CSV to use, e.g. 2YL0 - - OUT --out one row per window: - min_rmsd_to_A closest member to conformer A, Angstrom - min_rmsd_to_B closest member to conformer B, Angstrom - n_atoms_A/_B atoms actually compared (0 = nothing matched) - - EXIT 0 if windows scored, 1 if none did. - -How to read it: both numbers small means the ensemble captured both conformers. One small and -one large means it collapsed onto a single conformation -- so max(A, B) is the honest per-window -summary, and the median of that across windows is the honest per-arm summary. - -Two details that decide the numbers: - -* Altlocs are read with gemmi, which exposes the altloc character directly. `load_any` drops the - annotation and `parse` keeps only the first altloc, so neither can separate the conformers. - Atoms with a blank altloc are shared and belong to both. -* The prediction is placed by a GLOBAL uniform-weight Kabsch onto the reference, never fitted on - the window itself -- fitting locally would let a wrong conformation be rotated into agreement. - -Run it: - pixi run -e analysis python it_opt_scratch/score_rmsd_simplified.py \ - --prediction out/2YL0/refined-patched.cif \ - --reference /home/dev/test_data/processed/2YL0/2YL0_single_001_density_input.cif \ - --selections-csv it_opt_scratch/paper_maxrmsd_selections.csv \ - --protein 2YL0 --out 2YL0_rmsd.csv - -Use the patched CIF (scripts/patch_output_cif_files.py) when you have one. Atoms are matched by -(chain, residue number, atom name), and for some proteins the raw prediction relabels chains and -renumbers from 1 while the reference keeps its deposited numbering -- then nothing matches and -every window comes back with n_atoms 0. This script stops with that message rather than writing -a file full of blanks. -""" - -from __future__ import annotations - -import argparse -import csv -import re -import sys -from pathlib import Path - -import gemmi -import numpy as np -import pandas as pd -import torch -from atomworks.io.utils.io_utils import load_any -from loguru import logger -from sampleworks.utils.atom_array_utils import ( - filter_to_common_atoms, - remove_atoms_with_any_nan_coords, -) -from sampleworks.utils.frame_transforms import ( - apply_forward_transform, - weighted_rigid_align_differentiable, -) -from sampleworks.utils.framework_utils import match_batch - -SELECTION_RE = re.compile(r"^chain\s+(\S+)\s+and\s+resi\s+(-?\d+)\s*-\s*(-?\d+)$") - - -def main() -> int: - args = parse_args() - - windows = read_windows(args.selections_csv, args.protein) - if not windows: - sys.exit(f"no selections for {args.protein} in {args.selections_csv}") - - # The two reference conformers per window, straight from the altloc characters. - conformers = reference_conformers(args.reference, windows) - - prediction = remove_atoms_with_any_nan_coords(load_any(str(args.prediction))) - alignment_target = remove_atoms_with_any_nan_coords(load_any(str(args.reference))) - prediction = align_to_reference(alignment_target, prediction) - predicted_atoms = prediction_lookup(prediction) - - rows = [] - for window in windows: - rmsd_a, n_a = min_rmsd_over_ensemble(conformers[window]["A"], predicted_atoms) - rmsd_b, n_b = min_rmsd_over_ensemble(conformers[window]["B"], predicted_atoms) - rows.append( - { - "protein": args.protein, - "selection": window, - "min_rmsd_to_A": rmsd_a, - "min_rmsd_to_B": rmsd_b, - "n_atoms_A": n_a, - "n_atoms_B": n_b, - } - ) - - table = pd.DataFrame(rows) - table.to_csv(args.out, index=False) - - scored = table[table[["n_atoms_A", "n_atoms_B"]].gt(0).all(axis=1)] - if scored.empty: - logger.error( - "no atoms matched between prediction and reference. Atoms are matched by " - "(chain, residue number, atom name) -- the prediction is probably renumbered. " - "Patch it first: scripts/patch_output_cif_files.py" - ) - return 1 - - worse = scored[["min_rmsd_to_A", "min_rmsd_to_B"]].max(axis=1) - logger.info(f"wrote {args.out}: {len(scored)}/{len(table)} windows scored") - logger.info( - f"median max(A,B) {worse.median():.3f} A " - f"windows with both conformers within 1 A: {worse.le(1.0).mean():.0%}" - ) - return 0 - - -def read_windows(csv_path: Path, protein: str) -> list[str]: - """The paper's windows for one protein: the semicolon-joined `selection` column of its row.""" - with open(csv_path) as fh: - for row in csv.DictReader(fh): - if row["protein"].strip().upper() == protein.upper(): - return [s.strip() for s in row["selection"].split(";") if s.strip()] - return [] - - -def parse_selection(selection: str) -> tuple[str, range]: - """Split ``chain A and resi 12-14`` into ``("A", range(12, 15))``.""" - match = SELECTION_RE.match(selection.strip()) - if match is None: - raise ValueError(f"unparseable selection: {selection!r}") - chain, low, high = match.group(1), int(match.group(2)), int(match.group(3)) - return chain, range(low, high + 1) - - -def reference_conformers(reference_path: Path, windows: list[str]) -> dict: - """Per window, the altloc-A and altloc-B conformers keyed by (chain, residue, atom name). - - Read with gemmi because it exposes the altloc character per atom. An atom with a blank - altloc is shared between the two conformers, so it is placed in both. - """ - structure = gemmi.read_structure(str(reference_path)) - atoms_by_residue: dict[tuple[str, int], list] = {} - for chain in structure[0]: - for residue in chain: - atoms_by_residue.setdefault((chain.name, residue.seqid.num), []).extend( - (a.name, a.altloc, np.array([a.pos.x, a.pos.y, a.pos.z])) for a in residue - ) - - out = {} - for window in windows: - chain_name, residues = parse_selection(window) - conformers = {"A": {}, "B": {}} - for residue_id in residues: - for atom_name, altloc, xyz in atoms_by_residue.get((chain_name, residue_id), []): - shared = altloc.strip() == "" - for label in ("A", "B") if shared else (altloc.strip(),): - if label in conformers: - conformers[label][(chain_name, residue_id, atom_name)] = xyz - out[window] = conformers - return out - - -def prediction_lookup(atom_array) -> dict: - """Map (chain, residue, atom name) -> that atom's coordinates in every ensemble member. - - Values have shape [n_models, 3]; a single-model file is given a leading axis of 1. - """ - coords = atom_array.coord - if coords.ndim == 2: - coords = coords[None] - return { - (str(chain), int(residue), str(name)): coords[:, i] - for i, (chain, residue, name) in enumerate( - zip(atom_array.chain_id, atom_array.res_id, atom_array.atom_name, strict=True) - ) - } - - -def min_rmsd_over_ensemble(conformer: dict, predicted_atoms: dict) -> tuple[float, int]: - """Closest any single ensemble member gets to one reference conformer. - - The minimum is over members, not an average: the question is whether ANY member found this - conformation, not whether the ensemble is centred on it. - """ - shared_keys = [key for key in conformer if key in predicted_atoms] - if not shared_keys: - return float("nan"), 0 - reference = np.stack([conformer[key] for key in shared_keys]) # [n_atoms, 3] - predicted = np.stack([predicted_atoms[key] for key in shared_keys], axis=1) # [n_models, n, 3] - per_member = np.sqrt(((predicted - reference[None]) ** 2).sum(-1).mean(-1)) # [n_models] - return float(per_member.min()), len(shared_keys) - - -def align_to_reference(reference, prediction): - """Global uniform-weight Kabsch of prediction onto reference, applied to every atom. - - Alignment is deliberately GLOBAL: fitting on the 3-residue window itself would let a wrong - local conformation be rotated into apparent agreement. - """ - try: - reference_common, prediction_common = filter_to_common_atoms(reference, prediction) - except RuntimeError: - # Some predictions relabel every chain to 'A' and renumber residues from 1 while the - # reference keeps its deposited chain id and numbering, so strict (chain, res, name) - # matching finds nothing. Sequential per-chain matching realigns these otherwise - # identical structures. Only runs when strict matching raises. - reference_common, prediction_common = filter_to_common_atoms( - reference, prediction, normalize_ids=True - ) - - reference_coords = torch.from_numpy(reference_common.coord).float() - prediction_coords = torch.from_numpy(prediction_common.coord).float() - reference_coords = match_batch(reference_coords, prediction_coords.shape[0]) - if reference_coords.ndim != 3 or reference_coords.shape[1] != prediction_coords.shape[1]: - raise ValueError( - f"shape mismatch: reference {tuple(reference_coords.shape)} " - f"vs prediction {tuple(prediction_coords.shape)}" - ) - - n_atoms = reference_coords.shape[1] - _, transform = weighted_rigid_align_differentiable( - true_coords=prediction_coords, - pred_coords=reference_coords, - weights=torch.ones(1, n_atoms), - mask=torch.ones(1, n_atoms), - return_transforms=True, - allow_gradients=False, - ) - moved = apply_forward_transform( - torch.from_numpy(prediction.coord), transform, rotation_only=False - ) - prediction.coord = moved.numpy() - return prediction - - -def parse_args() -> argparse.Namespace: - ap = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - ap.add_argument("--prediction", type=Path, required=True, help="refined.cif to score") - ap.add_argument("--reference", type=Path, required=True, - help="deposited .cif containing altlocs A and B") - ap.add_argument("--selections-csv", type=Path, required=True, help="the paper's windows") - ap.add_argument("--protein", required=True, help="which protein's windows, e.g. 2YL0") - ap.add_argument("--out", type=Path, required=True, help="output CSV") - return ap.parse_args() - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/it_opt_scratch/score_rscc_simplified.py b/it_opt_scratch/score_rscc_simplified.py deleted file mode 100644 index 0b93179e..00000000 --- a/it_opt_scratch/score_rscc_simplified.py +++ /dev/null @@ -1,240 +0,0 @@ -"""RSCC for one prediction against one density map. Self-contained: just run it. - -Scores how well a predicted ensemble explains observed density, over the paper's 3-residue -windows. One prediction per invocation, so a sweep is a plain loop in the caller. - - IN --prediction refined.cif from a run (an ensemble is expected; see below) - --reference the deposited .cif -- the alignment target, and the source of the - coordinates that define each window - --map the observed .ccp4 - --resolution that map's resolution, in Angstrom - --selections-csv the paper's per-protein windows (semicolon-joined column) - --protein which row of that CSV to use, e.g. 2YL0 - - OUT --out one row per window: protein, selection, rscc, error - - EXIT 0 if every window scored, 1 if any errored. - -The whole ensemble becomes a SINGLE calculated density, compared to the map once. That is the -point of the metric: RSCC scores the multi-conformer model, not its members individually, which -is how an ensemble can explain density that no single conformer can. - -Run it: - pixi run -e analysis python it_opt_scratch/score_rscc_simplified.py \ - --prediction out/2YL0/refined.cif \ - --reference /home/dev/test_data/processed/2YL0/2YL0_single_001_density_input.cif \ - --map density_maps/2YL0_0.5occA_0.5occB_1.00A.ccp4 \ - --resolution 1.0 \ - --selections-csv it_opt_scratch/paper_maxrmsd_selections.csv \ - --protein 2YL0 --out 2YL0_rscc.csv - -Every number-determining step is written out in this file on purpose -- the alignment, the -occupancy/B values, the window cropping. They match scripts/eval/rscc_grid_search_script.py and -score_paper_rscc.py; changing any of them changes the metric, so the comments say why each is -the way it is. -""" - -from __future__ import annotations - -import argparse -import copy -import csv -import sys -from pathlib import Path - -import numpy as np -import pandas as pd -import torch -from atomworks.io.utils.io_utils import load_any -from loguru import logger -from sampleworks.eval.constants import DEFAULT_SELECTION_PADDING -from sampleworks.eval.metrics import rscc -from sampleworks.eval.structure_utils import extract_selection_coordinates -from sampleworks.utils.atom_array_utils import ( - filter_to_common_atoms, - remove_atoms_with_any_nan_coords, -) -from sampleworks.utils.density_utils import build_density_transformer, run_density_transformer -from sampleworks.utils.frame_transforms import ( - apply_forward_transform, - weighted_rigid_align_differentiable, -) -from sampleworks.utils.framework_utils import match_batch - -from sampleworks.core.forward_models.xray.real_space_density_deps.qfit.volume import XMap - - -def main() -> int: - args = parse_args() - device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - - windows = read_windows(args.selections_csv, args.protein) - if not windows: - sys.exit(f"no selections for {args.protein} in {args.selections_csv}") - - require_unit_cell(args.prediction) - - # The observed map, plus the forward model that turns coordinates into a calculated one. - observed = XMap.fromfile(str(args.map), resolution=args.resolution).canonical_unit_cell() - transformer, _ = build_density_transformer(observed, em_mode=False, device=device) - - # Two reads of the reference, for two different jobs: - # - windows are defined over ALL altlocs, so the mask covers both conformations; - # - alignment needs NaN-free coordinates. - # load_any, not parse: parse() reconciles against pdbx_poly_seq_scheme and, without a CCD - # mirror, turns most atoms into NaN placeholders. load_any reads atom_site as written. - reference_all_altlocs = load_any(str(args.reference)) - alignment_target = remove_atoms_with_any_nan_coords(load_any(str(args.reference))) - - window_coords = {} - for selection in windows: - coords = extract_selection_coordinates(reference_all_altlocs, selection) - if len(coords) and np.isfinite(coords).all(): - window_coords[selection] = coords - else: - logger.warning(f"{args.protein}: window {selection!r} empty or non-finite, skipping") - - prediction = remove_atoms_with_any_nan_coords(load_any(str(args.prediction))) - prediction = add_density_annotations(prediction) - prediction = align_to_reference(alignment_target, prediction) - - calculated = copy.copy(observed) - calculated.array = run_density_transformer(transformer, prediction).cpu().numpy() - if calculated.array.shape != observed.array.shape: - sys.exit(f"density shape {calculated.array.shape} != observed {observed.array.shape}") - - rows = [ - score_window(args.protein, selection, coords, observed, calculated) - for selection, coords in window_coords.items() - ] - - table = pd.DataFrame(rows) - table.to_csv(args.out, index=False) - scored = int(table["rscc"].notna().sum()) - logger.info(f"wrote {args.out}: {scored}/{len(table)} windows scored") - if scored: - logger.info(f"median rscc {table['rscc'].median():.4f}") - return 0 if scored == len(table) else 1 - - -def require_unit_cell(cif_path: Path) -> None: - """Stop early if the CIF has no crystallographic header. - - X-ray density is calculated in the full crystal frame, so the forward model needs the unit - cell and space group. A freshly generated refined.cif carries neither -- the header is added - afterwards by scripts/patch_output_cif_files.py, which fetches it from the PDB entry. Without - this check, scoring an unpatched file yields numbers that look reasonable and are not. - """ - text = cif_path.read_text() - if "_cell." not in text: - sys.exit( - f"{cif_path} has no unit cell -- it looks like an unpatched refined.cif.\n" - "Add the crystallographic header first:\n" - " python scripts/patch_output_cif_files.py --input-dir " - "--grid-search-input-dir --rcsb-pattern '/([0-9][A-Za-z0-9]{3})'\n" - "then score the resulting refined-patched.cif." - ) - - -def read_windows(csv_path: Path, protein: str) -> list[str]: - """The paper's windows for one protein: the semicolon-joined `selection` column of its row.""" - with open(csv_path) as fh: - for row in csv.DictReader(fh): - if row["protein"].strip().upper() == protein.upper(): - return [s.strip() for s in row["selection"].split(";") if s.strip()] - return [] - - -def align_to_reference(reference, prediction): - """Global uniform-weight Kabsch of prediction onto reference, applied to every atom. - - Alignment is deliberately GLOBAL: fitting on the 3-residue window itself would let a wrong - local conformation be rotated into apparent agreement. The transform is fitted on the atoms - common to both structures, then applied to the whole predicted array. - """ - try: - reference_common, prediction_common = filter_to_common_atoms(reference, prediction) - except RuntimeError: - # The prediction relabels every chain to 'A' and renumbers residues from 1, while the - # reference keeps the deposited chain id and numbering (e.g. chain 'P', res 5-234). - # Strict (chain, res, name) matching then finds nothing, so fall back to sequential - # per-chain matching. This only runs when strict matching raises, so proteins that - # already align keep their exact matched-atom set. - reference_common, prediction_common = filter_to_common_atoms( - reference, prediction, normalize_ids=True - ) - - reference_coords = torch.from_numpy(reference_common.coord).float() - prediction_coords = torch.from_numpy(prediction_common.coord).float() - reference_coords = match_batch(reference_coords, prediction_coords.shape[0]) - if reference_coords.ndim != 3 or reference_coords.shape[1] != prediction_coords.shape[1]: - raise ValueError( - f"shape mismatch: reference {tuple(reference_coords.shape)} " - f"vs prediction {tuple(prediction_coords.shape)}" - ) - - n_atoms = reference_coords.shape[1] - _, transform = weighted_rigid_align_differentiable( - true_coords=prediction_coords, - pred_coords=reference_coords, - weights=torch.ones(1, n_atoms), - mask=torch.ones(1, n_atoms), - return_transforms=True, - allow_gradients=False, - ) - moved = apply_forward_transform( - torch.from_numpy(prediction.coord), transform, rotation_only=False - ) - prediction.coord = moved.numpy() - return prediction - - -def add_density_annotations(atom_array): - """Add the occupancy / b_factor annotations the density forward model requires. - - load_any does not populate these even when the CIF carries the columns. The values are - occupancy 1.0 and B 20.0, which is what parse() supplied to the published scorer. - """ - n_atoms = atom_array.coord.shape[-2] - for name, value in (("occupancy", 1.0), ("b_factor", 20.0)): - if name not in atom_array.get_annotation_categories(): - atom_array.set_annotation(name, np.full(n_atoms, value)) - return atom_array - - -def score_window(protein: str, selection: str, coords, observed, calculated) -> dict: - """RSCC over one window: crop both maps to it, then correlate. - - Cropping is `extract_tight` around the window's reference coordinates, so observed and - calculated are compared on exactly the same voxels. - """ - row = {"protein": protein, "selection": selection, "rscc": None, "error": None} - try: - _, observed_crop = observed.extract_tight(coords, padding=DEFAULT_SELECTION_PADDING) - _, calculated_crop = calculated.extract_tight(coords, padding=DEFAULT_SELECTION_PADDING) - if observed_crop is None or observed_crop.shape[0] == 0: - raise ValueError("observed crop empty") - if calculated_crop is None or calculated_crop.shape[0] == 0: - raise ValueError("calculated crop empty") - row["rscc"] = rscc(observed_crop, calculated_crop) - except Exception as err: # noqa: BLE001 - one bad window should not lose the others - row["error"] = str(err) - return row - - -def parse_args() -> argparse.Namespace: - ap = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - ap.add_argument("--prediction", type=Path, required=True, help="refined.cif to score") - ap.add_argument("--reference", type=Path, required=True, help="deposited .cif") - ap.add_argument("--map", type=Path, required=True, help="observed .ccp4") - ap.add_argument("--resolution", type=float, default=1.0, help="map resolution, Angstrom") - ap.add_argument("--selections-csv", type=Path, required=True, help="the paper's windows") - ap.add_argument("--protein", required=True, help="which protein's windows, e.g. 2YL0") - ap.add_argument("--out", type=Path, required=True, help="output CSV") - return ap.parse_args() - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/it_opt_scratch/slurm_ensemble_and_score.sh b/it_opt_scratch/slurm_ensemble_and_score.sh deleted file mode 100755 index 843db970..00000000 --- a/it_opt_scratch/slurm_ensemble_and_score.sh +++ /dev/null @@ -1,148 +0,0 @@ -#!/usr/bin/env bash -# Generate an ensemble for one target and score it, as a SLURM array task. -# -# IN a CSV of targets: name,structure,density,resolution (one row per target) -# OUT //refined-patched.cif the ensemble -# //rscc.csv per-window RSCC -# //rmsd.csv per-window min-altloc-RMSD -# /rscc_all.csv all targets, after `aggregate` -# /rmsd_all.csv all targets, after `aggregate` -# -# Each array task handles exactly one row, so a failed target fails only its own task and can be -# requeued on its own. Nothing is appended to a shared file until `aggregate`, so tasks never -# race each other. -# -# Stages (argument 1, default `all`): -# generate GPU. sample the ensemble, then add the crystallographic header -# score CPU (GPU optional). RSCC + min-altloc-RMSD into two per-target CSVs -# all both, in one task -# aggregate concatenate every per-target CSV into the two final ones (run once, at the end) -# -# Submit as two dependent arrays so the CPU stage does not sit on a GPU: -# -# N=$(( $(wc -l < targets.csv) - 1 )) -# gen=$(sbatch --parsable --array=1-$N --gres=gpu:1 --cpus-per-task=6 --mem=64G \ -# --time=2:00:00 it_opt_scratch/slurm_ensemble_and_score.sh generate) -# scr=$(sbatch --parsable --array=1-$N --dependency=aftercorr:$gen --cpus-per-task=4 --mem=32G \ -# --time=1:00:00 it_opt_scratch/slurm_ensemble_and_score.sh score) -# sbatch --dependency=afterany:$scr --cpus-per-task=1 --mem=4G \ -# it_opt_scratch/slurm_ensemble_and_score.sh aggregate -# -# Or in one array (simpler, wastes the GPU during scoring): -# sbatch --array=1-$N --gres=gpu:1 it_opt_scratch/slurm_ensemble_and_score.sh all -# -# Without SLURM it runs row 1 unless you set TASK_ID: -# TASK_ID=3 bash it_opt_scratch/slurm_ensemble_and_score.sh all -# -# BEFORE SUBMITTING: the header-patching step downloads the PDB entry from RCSB. Compute nodes -# on many clusters have no outbound network. Warm the cache on a login node first by running the -# `generate` stage for every target there, or pre-populate ~/.sampleworks/rcsb -- otherwise every -# task fails at step 2. See the note in check_prerequisites below. - -set -euo pipefail - -# ------------------------------- settings ------------------------------- -# Override any of these by exporting them before sbatch, e.g. `MODE=z_only sbatch ...`. -REPO="${REPO:-/home/dev/workspace}" -TARGETS="${TARGETS:-$REPO/it_opt_scratch/targets.csv}" -OUTPUT_ROOT="${OUTPUT_ROOT:-$REPO/it_opt_scratch/slurm_out}" -SELECTIONS="${SELECTIONS:-$REPO/it_opt_scratch/paper_maxrmsd_selections.csv}" -PROCESSED_DIR="${PROCESSED_DIR:-/home/dev/test_data/processed}" - -MODE="${MODE:-s_plus_z}" -ENSEMBLE_SIZE="${ENSEMBLE_SIZE:-8}" -BOND_LENGTH_WEIGHT="${BOND_LENGTH_WEIGHT:-5e-5}" -NUM_STEPS="${NUM_STEPS:-200}" - -GEN_ENV="${GEN_ENV:-protenix-dev}" # has protenix + torch -ANALYSIS_ENV="${ANALYSIS_ENV:-analysis}" # has gemmi + the density tooling -# ------------------------------------------------------------------------ - -STAGE="${1:-all}" - -# aggregate only reads OUTPUT_ROOT, so it runs anywhere -- no repo, no pixi, no GPU. -if [[ "$STAGE" == "aggregate" ]]; then - # Keep the header from the first file, skip it in the rest. - for metric in rscc rmsd; do - out="$OUTPUT_ROOT/${metric}_all.csv" - first=1 - : > "$out" - for f in "$OUTPUT_ROOT"/*/"${metric}.csv"; do - [[ -e "$f" ]] || continue - if [[ $first == 1 ]]; then cat "$f"; first=0; else tail -n +2 "$f"; fi - done >> "$out" - echo "[aggregate] $(( $(wc -l < "$out") - 1 )) rows -> $out" - done - exit 0 -fi - -cd "$REPO" - -# ---- which target is this task? Row 1 of the CSV is the header, so add one. ---- -TASK_ID="${SLURM_ARRAY_TASK_ID:-${TASK_ID:-1}}" -row=$(( TASK_ID + 1 )) -line=$(sed -n "${row}p" "$TARGETS") -[[ -n "$line" ]] || { echo "no row $row in $TARGETS" >&2; exit 1; } -IFS=, read -r NAME STRUCTURE DENSITY RESOLUTION <<< "$line" - -# The window list and the PDB header lookup are keyed by the bare 4-character PDB id, while the -# target name usually carries a suffix describing the map (e.g. 2YL0_0.5occA_0.5occB). -PDB="${NAME%%_*}" -OUT_DIR="$OUTPUT_ROOT/$NAME" -mkdir -p "$OUT_DIR" - -echo "[task $TASK_ID] $NAME (pdb $PDB) stage=$STAGE mode=$MODE ens=$ENSEMBLE_SIZE" -echo " structure $STRUCTURE" -echo " density $DENSITY @ ${RESOLUTION} A" -echo " out $OUT_DIR" - -if [[ "$STAGE" == "generate" || "$STAGE" == "all" ]]; then - echo "[1/4] sampling the ensemble" - pixi run -e "$GEN_ENV" python -u it_opt_scratch/run_targets_simplified.py \ - --structure "$STRUCTURE" \ - --density "$DENSITY" \ - --resolution "$RESOLUTION" \ - --mode "$MODE" \ - --ensemble-size "$ENSEMBLE_SIZE" \ - --num-steps "$NUM_STEPS" \ - --bond-length-weight "$BOND_LENGTH_WEIGHT" \ - --name "$NAME" \ - --skip-existing \ - --output-dir "$OUT_DIR" - - # RSCC needs the unit cell and space group, which sampling does not write. This fetches them - # from the PDB entry (network!) and writes refined-patched.cif alongside refined.cif. - echo "[2/4] adding the crystallographic header" - pixi run -e "$ANALYSIS_ENV" python scripts/patch_output_cif_files.py \ - --input-dir "$OUT_DIR" \ - --depth 1 \ - --cif-pattern refined.cif \ - --rcsb-pattern "($PDB)" \ - --grid-search-input-dir "$PROCESSED_DIR" -fi - -if [[ "$STAGE" == "score" || "$STAGE" == "all" ]]; then - SCORED_CIF="$OUT_DIR/refined-patched.cif" - [[ -f "$SCORED_CIF" ]] || { echo "missing $SCORED_CIF -- did the generate stage finish?" >&2; exit 1; } - REFERENCE="$PROCESSED_DIR/$PDB/${PDB}_single_001_density_input.cif" - - echo "[3/4] RSCC" - pixi run -e "$ANALYSIS_ENV" python it_opt_scratch/score_rscc_simplified.py \ - --prediction "$SCORED_CIF" \ - --reference "$REFERENCE" \ - --map "$DENSITY" \ - --resolution "$RESOLUTION" \ - --selections-csv "$SELECTIONS" \ - --protein "$PDB" \ - --out "$OUT_DIR/rscc.csv" - - echo "[4/4] min-altloc-RMSD" - pixi run -e "$ANALYSIS_ENV" python it_opt_scratch/score_rmsd_simplified.py \ - --prediction "$SCORED_CIF" \ - --reference "$REFERENCE" \ - --selections-csv "$SELECTIONS" \ - --protein "$PDB" \ - --out "$OUT_DIR/rmsd.csv" -fi - -echo "[task $TASK_ID] $NAME done" From dea8f7fa93ff25b085f4b45e3aa754ef9f57b6a9 Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Mon, 10 Aug 2026 17:45:52 -0700 Subject: [PATCH 12/15] fix(it-opt): clear this branch's own lint and type findings 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. --- run_grid_search.py | 2 +- src/sampleworks/core/scalers/latent_optimization.py | 10 +++++++++- tests/rewards/test_geometry.py | 10 ++++------ 3 files changed, 14 insertions(+), 8 deletions(-) diff --git a/run_grid_search.py b/run_grid_search.py index 058aa6a8..bf193cee 100755 --- a/run_grid_search.py +++ b/run_grid_search.py @@ -20,10 +20,10 @@ from loguru import logger as log from sampleworks.utils.guidance_constants import GuidanceType, StructurePredictor from sampleworks.utils.guidance_script_arguments import ( + add_latent_opt_args, GuidanceConfig, JobConfig, JobResult, - add_latent_opt_args, ) from sampleworks.utils.protein_input import ProteinInput diff --git a/src/sampleworks/core/scalers/latent_optimization.py b/src/sampleworks/core/scalers/latent_optimization.py index 39ca118b..32de5227 100644 --- a/src/sampleworks/core/scalers/latent_optimization.py +++ b/src/sampleworks/core/scalers/latent_optimization.py @@ -167,7 +167,15 @@ def step(self, x_t: Tensor, t, *, features: GenerativeModelInput) -> Tensor: 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]) + # read_pair returns None when the io addresses no pair rep. sample() only sets + # optimize_pair together with a pair_attr, so this is a misconfigured io reaching + # us directly; say so rather than failing on a None subscript mid-denoise. + pair = self._io.read_pair(cond) + if pair is None: + raise ValueError( + "optimize_pair is set, but the io addresses no pair representation." + ) + cond_i = self._io.write_pair(cond_i, pair[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] diff --git a/tests/rewards/test_geometry.py b/tests/rewards/test_geometry.py index 1b510505..ca2e0048 100644 --- a/tests/rewards/test_geometry.py +++ b/tests/rewards/test_geometry.py @@ -1,17 +1,15 @@ """A test case for testing the geometry reward -This test case is designed to verify the reward function that evaluates the geometry of a molecular structure. -It checks whether the reward function correctly computes the reward +This test case is designed to verify the reward function that evaluates the geometry of a +molecular structure. It checks whether the reward function correctly computes the reward based on the provided atomic coordinates and other relevant parameters. - - """ -import pytest import gemmi from sampleworks.core.rewards.geometry import _covalent_radius + def test_known_element_matches_gemmi(): r = _covalent_radius("C") assert r > 0, "Covalent radius for Carbon should be greater than 0" - assert r == gemmi.Element("C").covalent_r \ No newline at end of file + assert r == gemmi.Element("C").covalent_r From e2393f4e23af5f926d1115dc89e4f6862f474423 Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Thu, 13 Aug 2026 13:18:09 -0700 Subject: [PATCH 13/15] test(runs): register the IT-opt presets in the bundled experiment set 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. --- tests/runs/test_cli.py | 2 ++ tests/runs/test_loader.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/tests/runs/test_cli.py b/tests/runs/test_cli.py index a20f141c..52c2e65e 100644 --- a/tests/runs/test_cli.py +++ b/tests/runs/test_cli.py @@ -20,6 +20,8 @@ def test_list_prints_all_experiment_presets(capsys: pytest.CaptureFixture[str]) "boltz2_md", "boltz2_xrd", "full_8gpu", + "it_opt_1gpu", + "it_opt_4gpu", "protenix", "protenix_dual", "rf3", diff --git a/tests/runs/test_loader.py b/tests/runs/test_loader.py index 0d3dd672..3e5e8eb5 100644 --- a/tests/runs/test_loader.py +++ b/tests/runs/test_loader.py @@ -15,6 +15,8 @@ "boltz2_md", "boltz2_xrd", "full_8gpu", + "it_opt_1gpu", + "it_opt_4gpu", "protenix", "protenix_dual", "rf3", From 5c0b8170787fceaa4075509cb2122d1869929c36 Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Thu, 13 Aug 2026 13:42:41 -0700 Subject: [PATCH 14/15] Revert "test(runs): register the IT-opt presets in the bundled experiment set" This reverts commit e2393f4e23af5f926d1115dc89e4f6862f474423. --- tests/runs/test_cli.py | 2 -- tests/runs/test_loader.py | 2 -- 2 files changed, 4 deletions(-) diff --git a/tests/runs/test_cli.py b/tests/runs/test_cli.py index 52c2e65e..a20f141c 100644 --- a/tests/runs/test_cli.py +++ b/tests/runs/test_cli.py @@ -20,8 +20,6 @@ def test_list_prints_all_experiment_presets(capsys: pytest.CaptureFixture[str]) "boltz2_md", "boltz2_xrd", "full_8gpu", - "it_opt_1gpu", - "it_opt_4gpu", "protenix", "protenix_dual", "rf3", diff --git a/tests/runs/test_loader.py b/tests/runs/test_loader.py index 3e5e8eb5..0d3dd672 100644 --- a/tests/runs/test_loader.py +++ b/tests/runs/test_loader.py @@ -15,8 +15,6 @@ "boltz2_md", "boltz2_xrd", "full_8gpu", - "it_opt_1gpu", - "it_opt_4gpu", "protenix", "protenix_dual", "rf3", From c316ae85633cf149b0e3c06ba057d1ac40716be7 Mon Sep 17 00:00:00 2001 From: Feng Yu Date: Thu, 13 Aug 2026 14:23:44 -0700 Subject: [PATCH 15/15] test(runs): register the IT-opt presets in the bundled set --- tests/runs/conftest.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/runs/conftest.py b/tests/runs/conftest.py index 78204e7a..c0a492ff 100644 --- a/tests/runs/conftest.py +++ b/tests/runs/conftest.py @@ -41,6 +41,8 @@ def force_pixi_argv(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: "boltz2_md", "boltz2_xrd", "full_8gpu", + "it_opt_1gpu", + "it_opt_4gpu", "protenix", "protenix_dual", "protpardelle",