fix(rf3): squashed RF3 bug that was blocking running on models with unresolved residues#140
Merged
marcuscollins merged 4 commits intomainfrom Mar 9, 2026
Merged
Conversation
k-chrispens
requested review from
Copilot and
marcuscollins
and removed request for
Copilot
Contributor
📝 WalkthroughWalkthroughThis PR removes Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Tip Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs). Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests/integration/test_mismatch_integration.py`:
- Around line 638-640: The test duplicates the same assertion on
reward_inputs.input_coords.shape[-2], so replace the second duplicate with the
missing contract check (likely on reward_inputs.occupancies) to verify shapes
match expected_n_model; update the assertion to something like asserting
reward_inputs.occupancies.shape[-2] == real_pair_case.expected_n_model and keep
the existing checks on reward_inputs.elements.shape and
reward_inputs.input_coords.shape, referencing the reward_inputs object and
real_pair_case.expected_n_model to locate the change.
In `@tests/integration/test_pipeline_integration.py`:
- Around line 1076-1081: In the test docstring describing
process_structure_to_trajectory_input feeding into to_reward_inputs, replace the
multiplication sign '×' with a plain ASCII 'x' to satisfy Ruff's RUF002; update
the triple-quoted string that currently reads "wrapper × structure" to "wrapper
x structure" so the test's docstring no longer contains the non-ASCII
multiplication character.
- Around line 1113-1120: The assertion only checks atom counts but not the
validity of values; after computing reward_inputs =
processed.to_reward_inputs(device=device) and verifying n_state == n_reward, add
checks to ensure reward_inputs.elements and reward_inputs.occupancy contain
finite numbers (no NaN/Inf) and that occupancy values are within the expected
range (e.g., 0.0..1.0 or boolean 0/1 depending on the contract) so the test
fails if coordinates or occupancies are invalid; reference reward_inputs,
processed.to_reward_inputs, reward_inputs.elements, reward_inputs.occupancy,
state, n_state, and n_reward when locating where to insert these assertions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 29056aec-fb05-424b-bf50-ffad4008dc0c
📒 Files selected for processing (11)
src/sampleworks/core/rewards/protocol.pysrc/sampleworks/core/samplers/edm.pysrc/sampleworks/eval/generate_synthetic_density.pysrc/sampleworks/eval/grid_search_eval_utils.pysrc/sampleworks/eval/structure_utils.pysrc/sampleworks/models/rf3/wrapper.pysrc/sampleworks/utils/atom_array_utils.pysrc/sampleworks/utils/guidance_script_utils.pytests/integration/test_mismatch_integration.pytests/integration/test_pipeline_integration.pytests/utils/test_density_utils.py
Collaborator
marcuscollins
left a comment
marcuscollins
left a comment
There was a problem hiding this comment.
Just a couple questions to be sure of things. If you can answer them, I think it will be fine and I will approve.
| _save_trajectory( | ||
| trajectory, atom_array, output_dir, reward_param_mask, subdir_name, save_every | ||
| ) | ||
| _save_trajectory(trajectory, atom_array, output_dir, subdir_name, save_every) |
Collaborator
There was a problem hiding this comment.
The methods in this submodule get used for all models--did only RF3 have this issue that there might be NaN values and strange occupancies? None of the others needed the reward_param_mask?
Collaborator
Author
There was a problem hiding this comment.
Correct, specifically its model_atom_array was getting initialized incorrectly. Now that we have the reconciler, no need for reward_param_mask.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes a bug where RF3 models with experimentally unresolved residues would crash during guidance due to NaN coordinates and zero occupancy atoms propagating through the pipeline. The fix ensures the RF3 wrapper initializes unresolved atoms with noise-based coordinates (centered on the resolved atom centroid) and sets their occupancy to 1.0. As a result, the reward_param_mask / mask_like masking mechanism in RewardInputs is no longer needed and has been cleanly removed, simplifying the reward and trajectory saving logic. This also addresses issue #86 (NaN coordinates in output structures).
Changes:
- RF3 wrapper now replaces NaN coordinates with centroid+noise and sets occupancy=1.0 for all model atoms, ensuring clean atom arrays before downstream processing.
- Removed
reward_param_maskandmask_likefromRewardInputs, adding input validation (element/b_factor annotations, finite coords, valid occupancy) infrom_atom_arrayinstead. - Simplified trajectory saving (
_write_coords_into_array,save_everything) to require exact atom count match, and migrated remainingpyrightignore comments toty.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
src/sampleworks/core/rewards/protocol.py |
Removed reward_param_mask/mask_like fields; added input validation in from_atom_array; simplified tensor construction. |
src/sampleworks/models/rf3/wrapper.py |
Core fix: replace NaN coords with centroid+noise, set occupancy=1.0 for all model atoms. |
src/sampleworks/eval/structure_utils.py |
Removed mask-based input_coords filtering in to_reward_inputs; renamed variable; formatting cleanup. |
src/sampleworks/utils/guidance_script_utils.py |
Removed reward_param_mask from save/trajectory functions; simplified _write_coords_into_array. |
src/sampleworks/utils/atom_array_utils.py |
Removed pyright: ignore comments (migrated to ty). |
src/sampleworks/core/samplers/edm.py |
Comment update: pyright → ty. |
src/sampleworks/eval/generate_synthetic_density.py |
Comment update: pyright → ty. |
src/sampleworks/eval/grid_search_eval_utils.py |
Comment update: pyright → ty. |
tests/integration/test_reward_mask_mismatch.py |
New test file verifying model-only atoms participate in rewards and end-to-end pipeline works. |
tests/integration/test_pipeline_integration.py |
Removed reward_param_mask/mask_like from test constructions; added TestRealWrapperPreprocessing slow tests. |
tests/integration/test_mismatch_integration.py |
Removed RealPairExpectation tests and associated helpers; cleaned up reward assertions. |
tests/utils/test_density_utils.py |
Removed pyright: ignore comment. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
You can also share your feedback on Copilot code review. Take the survey.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/sampleworks/models/rf3/wrapper.py (1)
318-345: Assert the post-cleanupmodel_aainvariant.This block fixes the happy path, but it never validates that the final
model_aastill matchesatom_to_token_mapand contains only finite coordinates. Adding an explicit guard here would turn any upstream preprocessing drift into a clear failure instead of silently propagating a bad model atom array.♻️ Suggested guard
if len(model_aa) > num_atoms: model_aa = cast(AtomArray, model_aa[:num_atoms]) + if len(model_aa) != num_atoms: + raise ValueError( + "RF3 model atom array length mismatch: " + f"len(model_aa)={len(model_aa)}, num_atoms={num_atoms}" + ) # atomworks's add_missing_atoms adds unresolved atoms with # occupancy=0.0 and NaN coordinates when we get our atom array with # InferenceInput.from_atom_array. RF3 operates on these atoms (they're # in atom_to_token_map), so initialize their coordinates with noise and @@ # All atoms in model_aa are operated on by RF3 during diffusion. # Set occupancy to 1.0 regardless of what atomworks assigned (unresolved # atoms from add_missing_atoms get occupancy=0.0, but RF3 should use them) model_aa.set_annotation("occupancy", np.ones(len(model_aa), dtype=np.float32)) + if not np.isfinite(model_aa.coord).all(): + raise ValueError("RF3 model atom array still contains non-finite coordinates")As per coding guidelines, "Fail fast with clear error messages when invalid states are encountered."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/sampleworks/models/rf3/wrapper.py` around lines 318 - 345, Add a fail-fast guard after the cleanup block to validate the model_aa invariant: ensure len(model_aa) equals len(atom_to_token_map) (or that every index in atom_to_token_map is < len(model_aa)) and that all coordinates in model_aa.coord are finite (no NaN/Inf). If either check fails, raise a descriptive exception (including the lengths and a short diagnostics string) so upstream preprocessing drift is caught early; place this check immediately after the occupancy set via model_aa.set_annotation and before further use of model_aa/b_factor.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests/integration/test_reward_mask_mismatch.py`:
- Around line 237-240: The test currently only asserts shapes and finite losses
but not the numeric values of emitted coordinates; update the assertions to
verify numeric finiteness of the predicted coordinates by checking that all
entries in result.final_state are finite (e.g., use np.isfinite on
result.final_state.ravel() or result.final_state). Also keep the existing checks
on result.losses (result.losses -> valid_losses) but add an assert that all
coordinates in result.final_state are finite to prevent NaN-coordinate
regressions.
- Around line 255-260: The test currently fabricates `state = torch.randn(1,
case.n_model, 3)`, which masks bugs in `_preprocess()`; replace the synthetic
state with a real pipeline-produced state (e.g. use `processed.input_coords`
from running the pipeline preprocessing step or call the public pipeline method
that yields `processed`) and then compare its atom dimension against
`reward_inputs.elements.shape[-1]`; update the assertion to use `n_state =
processed.input_coords.shape[-2]` (or cast `processed.input_coords` to a tensor
as needed) and assert `n_state == n_reward` to ensure the test validates real
preprocessing output rather than `case.n_model`.
---
Nitpick comments:
In `@src/sampleworks/models/rf3/wrapper.py`:
- Around line 318-345: Add a fail-fast guard after the cleanup block to validate
the model_aa invariant: ensure len(model_aa) equals len(atom_to_token_map) (or
that every index in atom_to_token_map is < len(model_aa)) and that all
coordinates in model_aa.coord are finite (no NaN/Inf). If either check fails,
raise a descriptive exception (including the lengths and a short diagnostics
string) so upstream preprocessing drift is caught early; place this check
immediately after the occupancy set via model_aa.set_annotation and before
further use of model_aa/b_factor.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 6ea609b4-f03d-42eb-9d77-7a2021d47d64
📒 Files selected for processing (4)
src/sampleworks/models/rf3/wrapper.pytests/integration/test_mismatch_integration.pytests/integration/test_pipeline_integration.pytests/integration/test_reward_mask_mismatch.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/integration/test_pipeline_integration.py
k-chrispens
force-pushed
the
kmc/fix-rf3-nan-reward-mask
branch
from
71377be to
859566a
Compare
…_atom_array, causing indexing errors downstream
…eward must be computable on the whole model output. This is reasonable because the reconciler takes care of mismatches now, so we can now move the responsibility of ensuring the atoms are all valid for the model to the ModelWrapper itself.
k-chrispens
force-pushed
the
kmc/fix-rf3-nan-reward-mask
branch
from
859566a to
7c854b9
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/sampleworks/utils/guidance_script_utils.py (1)
108-115:⚠️ Potential issue | 🟠 MajorCall
.cpu()before.numpy()when saving trajectories on GPU.Lines 114 and 138 use
.detach().numpy(), which raisesRuntimeErroron CUDA tensors. The final-state path at line 308 already uses the correct pattern with.detach().cpu().numpy(). These trajectory save paths will cause GPU guidance runs to succeed during computation, then fail only during serialization.Proposed fix
- _write_coords_into_array(array_copy, coords.detach().numpy()) + _write_coords_into_array(array_copy, coords.detach().cpu().numpy())- _write_coords_into_array(array_copy, coords[0].detach().numpy()) + _write_coords_into_array(array_copy, coords[0].detach().cpu().numpy())🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/sampleworks/utils/guidance_script_utils.py` around lines 108 - 115, The trajectory-saving code uses coords.detach().numpy() which will raise on CUDA tensors; update the serialization paths that call _write_coords_into_array(...) to call coords.detach().cpu().numpy() instead (where trajectory is iterated and array_copy is prepared and passed to save_structure and _write_coords_into_array) so GPU tensors are moved to CPU before converting to numpy; mirror the same pattern used for the final-state save.
♻️ Duplicate comments (1)
tests/integration/test_mismatch_integration.py (1)
677-681:⚠️ Potential issue | 🟡 MinorReplace the duplicated shape assert with the missing
occupanciescheck.Line 681 repeats the
input_coordsassertion, so this test still never verifies thatreward_inputs.occupanciesstays aligned with model atom count after removingmask_like.As per coding guidelines, `tests/**/*.py` should "Verify test outputs match contracts: shapes, value ranges, mathematical properties".Proposed fix
assert reward_inputs.elements.shape[-1] == n_model assert reward_inputs.b_factors.shape[-1] == n_model assert reward_inputs.input_coords.shape[-2] == n_model - assert reward_inputs.input_coords.shape[-2] == n_model + assert reward_inputs.occupancies.shape[-1] == n_model🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integration/test_mismatch_integration.py` around lines 677 - 681, The fourth assert duplicates the input_coords shape check instead of verifying occupancies; replace the second "assert reward_inputs.input_coords.shape[-2] == n_model" with an assertion that reward_inputs.occupancies.shape[-1] == n_model so the test verifies occupancies remain aligned with n_model (use the same n_model variable and reward_inputs object as in the surrounding checks).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/sampleworks/core/rewards/protocol.py`:
- Around line 69-77: The current validation in RewardInputs.from_atom_array()
checks atom_array.coord (and occupancy/b_factor) but the reward pipeline uses
SampleworksProcessedStructure.input_coords as the canonical coordinates; move
NaN/occupancy coordinate validation out of RewardInputs.from_atom_array() and
instead perform it where the final coordinate path is assembled (i.e., validate
SampleworksProcessedStructure.input_coords before rewards are computed), or
alternatively require and document that wrappers producing model_atom_array must
sanitize model_atom_array.coord (and occupancy/b_factor) before calling
to_reward_inputs(); update checks to reference input_coords or add an explicit
precondition in the RewardInputs.from_atom_array() docstring mentioning
sanitized model_atom_array.
---
Outside diff comments:
In `@src/sampleworks/utils/guidance_script_utils.py`:
- Around line 108-115: The trajectory-saving code uses coords.detach().numpy()
which will raise on CUDA tensors; update the serialization paths that call
_write_coords_into_array(...) to call coords.detach().cpu().numpy() instead
(where trajectory is iterated and array_copy is prepared and passed to
save_structure and _write_coords_into_array) so GPU tensors are moved to CPU
before converting to numpy; mirror the same pattern used for the final-state
save.
---
Duplicate comments:
In `@tests/integration/test_mismatch_integration.py`:
- Around line 677-681: The fourth assert duplicates the input_coords shape check
instead of verifying occupancies; replace the second "assert
reward_inputs.input_coords.shape[-2] == n_model" with an assertion that
reward_inputs.occupancies.shape[-1] == n_model so the test verifies occupancies
remain aligned with n_model (use the same n_model variable and reward_inputs
object as in the surrounding checks).
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 77993f70-fbeb-4d6c-84e5-28e40962442e
📒 Files selected for processing (11)
src/sampleworks/core/rewards/protocol.pysrc/sampleworks/core/samplers/edm.pysrc/sampleworks/eval/generate_synthetic_density.pysrc/sampleworks/eval/grid_search_eval_utils.pysrc/sampleworks/eval/structure_utils.pysrc/sampleworks/models/rf3/wrapper.pysrc/sampleworks/utils/atom_array_utils.pysrc/sampleworks/utils/guidance_script_utils.pytests/integration/test_mismatch_integration.pytests/integration/test_pipeline_integration.pytests/utils/test_density_utils.py
🚧 Files skipped from review as they are similar to previous changes (6)
- src/sampleworks/eval/generate_synthetic_density.py
- src/sampleworks/eval/structure_utils.py
- src/sampleworks/models/rf3/wrapper.py
- tests/utils/test_density_utils.py
- src/sampleworks/eval/grid_search_eval_utils.py
- src/sampleworks/core/samplers/edm.py
Comment on lines
+69
to
+77
| # input validation: ensure atom_array has required annotations and valid values | ||
| if not hasattr(atom_array, "element"): | ||
| raise ValueError("Atom array must have 'element' annotation.") | ||
| if not hasattr(atom_array, "b_factor"): | ||
| raise ValueError("Atom array must have 'b_factor' annotation.") | ||
| if np.any(np.isnan(atom_array.coord)): | ||
| raise ValueError("Atom array contains NaN coordinates.") | ||
| if np.any((atom_array.occupancy <= 0) | (atom_array.occupancy > 1)): | ||
| raise ValueError("Atom array contains invalid occupancy values.") |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n=== to_reward_inputs implementation ===\n'
sed -n '1,220p' src/sampleworks/eval/structure_utils.py
printf '\n=== wrapper/model-atom-array sanitation sites ===\n'
rg -n -C3 'model_atom_array|occupanc|isnan|isfinite|filter_zero_occupancy|coord' \
src/sampleworks/models src/sampleworks/eval/structure_utils.pyRepository: diff-use/sampleworks
Length of output: 50376
🏁 Script executed:
#!/bin/bash
# Check Protenix wrapper for coordinate sanitization
rg -n -A10 'model_atom_array|model_aa\s*=|coord.*protenix' \
src/sampleworks/models/protenix/wrapper.py | head -100
# Also check if Protenix structure_processing has any sanitization
fd -e py protenix | xargs rg -l 'nan|isfinite|coord' | head -5Repository: diff-use/sampleworks
Length of output: 2209
Move the NaN/occupancy validation to the final coordinate path, or require all wrappers to sanitize model_atom_array beforehand.
The validation at lines 69-77 checks atom_array.coord before it's used in rewards, but RewardInputs.from_atom_array() uses this only for atom accounting. The actual reward coordinates come from SampleworksProcessedStructure.input_coords (line 97 of structure_utils.py), which are already clean and pre-aligned to the model space. Protenix's wrapper sets occupancy and b_factor on model_atom_array but does not sanitize its coordinates, so this validation will reject valid inputs where the template has placeholder NaNs but input_coords is finite. Either validate input_coords instead, or document that all wrappers must sanitize model_atom_array.coord before calling to_reward_inputs().
🧰 Tools
🪛 Ruff (0.15.4)
[warning] 71-71: Avoid specifying long messages outside the exception class
(TRY003)
[warning] 73-73: Avoid specifying long messages outside the exception class
(TRY003)
[warning] 75-75: Avoid specifying long messages outside the exception class
(TRY003)
[warning] 77-77: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/sampleworks/core/rewards/protocol.py` around lines 69 - 77, The current
validation in RewardInputs.from_atom_array() checks atom_array.coord (and
occupancy/b_factor) but the reward pipeline uses
SampleworksProcessedStructure.input_coords as the canonical coordinates; move
NaN/occupancy coordinate validation out of RewardInputs.from_atom_array() and
instead perform it where the final coordinate path is assembled (i.e., validate
SampleworksProcessedStructure.input_coords before rewards are computed), or
alternatively require and document that wrappers producing model_atom_array must
sanitize model_atom_array.coord (and occupancy/b_factor) before calling
to_reward_inputs(); update checks to reference input_coords or add an explicit
precondition in the RewardInputs.from_atom_array() docstring mentioning
sanitized model_atom_array.
marcuscollins
approved these changes
Mar 9, 2026
Collaborator
marcuscollins
left a comment
marcuscollins
left a comment
There was a problem hiding this comment.
Questions answered, thanks. Approved.
This was referenced Mar 20, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This also removed the reward_param_mask as it isn't needed anymore if we are computing common atoms and saving the atoms that the model is operating on instead of having to reconcile and save the original structure. This should close #86 as well.
Summary by CodeRabbit
Release Notes
New Features
Refactor