feat(rewards): add differentiable structure-factor reward - #324
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughChangesStructure-factor reward
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant RewardCaller
participant StructureFactorRewardFunction
participant SFcalculator
RewardCaller->>StructureFactorRewardFunction: provide batched coordinates, B-factors, and occupancies
StructureFactorRewardFunction->>SFcalculator: compute ensemble protein amplitudes
StructureFactorRewardFunction->>SFcalculator: compute optional solvent-adjusted amplitudes
StructureFactorRewardFunction-->>RewardCaller: return masked amplitude loss
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
📝 Generate docstrings
🧪 Generate unit tests (beta)
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}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
src/sampleworks/core/rewards/structure_factor.py (1)
121-121: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a class-level docstring.
StructureFactorRewardFunctionhas only an__init__docstring, so class-level API documentation is absent. As per coding guidelines, “Add NumPy-style docstrings to every function and class.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/sampleworks/core/rewards/structure_factor.py` at line 121, Add a NumPy-style class-level docstring directly inside StructureFactorRewardFunction, documenting its purpose and relevant attributes or behavior; retain the existing __init__ docstring and avoid changing implementation logic.Source: Coding guidelines
tests/rewards/conftest.py (1)
63-106: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the new fixtures with NumPy-style docstrings.
The four added fixtures rely on comments or have no docstring. Add concise
Parameters/Returnsdocumentation. As per coding guidelines, “Add NumPy-style docstrings to every function and class.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/rewards/conftest.py` around lines 63 - 106, Add concise NumPy-style docstrings to the fixtures mtz_path_1vme, structure_1vme_sf, test_coordinates_1vme_sf, and reward_function_1vme_sf, documenting their parameters and return values. Replace or supplement the existing inline comments as needed while preserving the fixture behavior.Source: Coding guidelines
tests/rewards/test_structure_factor_reward.py (1)
342-370: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert free-set exclusion through the public reward call.
The test reads
sfc.free_flagand_reflection_mask, coupling it to internals. Use a public custom loss that returns the received reflection count, then assert the excluded-free reward receives fewer reflections.As per coding guidelines, “Write black-box tests that verify public behavior and contracts … rather than implementation details.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/rewards/test_structure_factor_reward.py` around lines 342 - 370, Update test_exclude_free_set_drops_reflections to verify exclusion through the public reward-call interface instead of inspecting sfc.free_flag or _reflection_mask. Define a public custom loss that records or returns the number of received reflections, pass it to both prepared rewards, invoke each reward, and assert that the exclude_free_reflections=True result receives fewer reflections than the unrestricted result.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/sampleworks/core/rewards/structure_factor.py`:
- Around line 221-223: Update the constructor validation immediately before
assigning self.batch_partition to require batch_partition to be an int but not a
bool, while preserving the existing positive-value check and error behavior for
invalid values. Ensure non-integer positive values are rejected before reaching
PARTITION.
- Around line 357-362: Update the reflection-mask construction near
self._reflection_mask to validate mask_np before converting it to a tensor.
Raise a clear ValueError when mask_np.any() is false, while preserving the
existing outlier and optional free-reflection filtering for non-empty masks.
- Around line 241-242: Update the sfcalculator_kwargs handling in the
surrounding reward initialization to reject or remove reward-owned keys mtzdata,
pdbmodel, expcolumns, device, and set_experiment before updating
self._sfc_kwargs. Preserve only safe extension options such as n_bins and
free-flag configuration, and ensure reserved arguments cannot override
reward-controlled values or reach the constructor twice.
---
Nitpick comments:
In `@src/sampleworks/core/rewards/structure_factor.py`:
- Line 121: Add a NumPy-style class-level docstring directly inside
StructureFactorRewardFunction, documenting its purpose and relevant attributes
or behavior; retain the existing __init__ docstring and avoid changing
implementation logic.
In `@tests/rewards/conftest.py`:
- Around line 63-106: Add concise NumPy-style docstrings to the fixtures
mtz_path_1vme, structure_1vme_sf, test_coordinates_1vme_sf, and
reward_function_1vme_sf, documenting their parameters and return values. Replace
or supplement the existing inline comments as needed while preserving the
fixture behavior.
In `@tests/rewards/test_structure_factor_reward.py`:
- Around line 342-370: Update test_exclude_free_set_drops_reflections to verify
exclusion through the public reward-call interface instead of inspecting
sfc.free_flag or _reflection_mask. Define a public custom loss that records or
returns the number of received reflections, pass it to both prepared rewards,
invoke each reward, and assert that the exclude_free_reflections=True result
receives fewer reflections than the unrestricted result.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 9e1ae148-0d34-43f4-92a5-09a9a8383486
📒 Files selected for processing (6)
src/sampleworks/core/rewards/structure_factor.pytests/resources/1vme/1vme_final_crystalframe_0.5occA_0.5occB_1.80A.ciftests/resources/1vme/1vme_final_crystalframe_0.5occA_0.5occB_1.80A.mtztests/rewards/conftest.pytests/rewards/test_reward_function_contract.pytests/rewards/test_structure_factor_reward.py
marcuscollins
left a comment
marcuscollins
left a comment
There was a problem hiding this comment.
There are a couple testing things I think should be added/changed, and a bunch of smaller issues that should be easy to address. Very close to ready.
There was a problem hiding this comment.
I would rather not add binary files to a git repo. The problem is that a binary file can't have tracked changes, and so if you change it, you keep both the old and the new files in the history forever. This leads to the repo growing out of control if you change the binary files often enough.
In any case for unit testing, it is usually better to create files like this "on the fly" as needed for testing (that would include the corresponding CIF too). Integration testing (which we aren't really doing right now) could use a real file like this, but then we'd store it somewhere else, not in the repo.
There was a problem hiding this comment.
the fixtures and the resources files were both mirroring the density reward ones. I am changing the files for SF to be generated on the fly. Would you want an issue opened about moving the density ccp4 files out as well or keep it as legacy code/file?
| self.sfc._set_scales(requires_grad=False) | ||
|
|
||
| # |Eo| are computed in SFC's experiment init (inside a try/except). | ||
| if self.normalize_amplitude and getattr(self.sfc, "Eo", None) is None: |
There was a problem hiding this comment.
Should we set the SFC args/kwargs in accordance with self.normalize_amplitude instead, and make sure Eo are computed? Or is that already the case and the calculation might have failed?
There was a problem hiding this comment.
there is no args/kwargs in SFC to control this, Eo is always computed when an mtz file is provided. If Eo computation failed, outlier mask is just all False, so I will modify the check here such that it logs a warning even if we don't use Eo (self.normalize_amplitude = False)
| if self.bulk_solvent == "per_conformer": | ||
| # calc_fsolvent_batch masks each conformer (from Fprotein_asu_batch, set by | ||
| # calc_fprotein_batch); the mean applies the 1/E weight -> <mask(rho)>. | ||
| Fmask_HKL_batch = self.sfc.calc_fsolvent_batch( |
There was a problem hiding this comment.
I haven't looked at the tests yet, so maybe you've got this: Make sure to check that this method gets called when self.bulk_solvent == "per_conformer". You can do that by making a mock for sfc.calc_fsolvent_batch (e.g. using MagicMock) and then the mock should keep track of how many times it is called.
There was a problem hiding this comment.
There is a test (test_ftotal_modes_diverge_for_distinct_conformer_ensemble) that check the resulted Ftotal differs depending on the bulk_solvent mode being "per_conformer" or "combined". I added a MagicMock component to that test to also count the correct dispatch function call, which I think helps diagnose if there is incidental agreement between <mask(protein_density)> and mask(<protein_density>). I will caveat that I've never heard of MagicMock/Mock before so I am not sure if I used it the best way.
There was a problem hiding this comment.
🧹 Nitpick comments (3)
tests/rewards/test_structure_factor_reward.py (1)
460-470: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocstring says "committed MTZ", but the source MTZ is generated per session.
mtz_path_1vmecomes from the synthetic generator fixture, not a committed file — worth correcting to avoid misleading future readers.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/rewards/test_structure_factor_reward.py` around lines 460 - 470, Update the docstring of the mtz_all_free_1vme fixture to describe the generated source MTZ accurately, replacing the reference to a committed MTZ while preserving the documented flag column name and test-set convention.tests/rewards/conftest.py (1)
95-141: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffFixture couples to a private generator and its filename convention.
_process_single_rowis private and returnsNone, so the fixture must reconstruct{stem}_sf_input.cif/{stem}_{resolution:.2f}A.mtzby hand — a silent break if the generator's naming changes. Since the comment already notes this, consider making the generator return the written paths (or accept explicit output names) so the fixture can consume a public contract.As per coding guidelines, tests should "verify public behavior and contracts" rather than implementation details.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/rewards/conftest.py` around lines 95 - 141, The fixture currently depends on private _process_single_row behavior and hard-coded output filenames. Update the generator’s public contract to return the generated CIF and MTZ paths (or accept explicit output names), then have the fixture consume that returned contract instead of reconstructing names from source_cif and resolution. Preserve the existing generation parameters and missing-output failure behavior.Source: Coding guidelines
src/sampleworks/core/rewards/structure_factor.py (1)
387-392: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse
sfc.init_scales(requires_grad=False)instead of_set_scales.
generate_synthetic_sf.pyusesinit_scalesfor the same default scale initialization, while_set_scalesis private SFC_Torch/SFcalculattorch internal API. If this code intentionally skipsinit_scales’s extra setup, add a short justification; otherwise call the public API consistently.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/sampleworks/core/rewards/structure_factor.py` around lines 387 - 392, Replace the private self.sfc._set_scales call in the bulk-solvent initialization block with the public self.sfc.init_scales(requires_grad=False) API, matching generate_synthetic_sf.py’s default scale initialization. If the extra setup performed by init_scales must be skipped, document that exception briefly instead.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@src/sampleworks/core/rewards/structure_factor.py`:
- Around line 387-392: Replace the private self.sfc._set_scales call in the
bulk-solvent initialization block with the public
self.sfc.init_scales(requires_grad=False) API, matching
generate_synthetic_sf.py’s default scale initialization. If the extra setup
performed by init_scales must be skipped, document that exception briefly
instead.
In `@tests/rewards/conftest.py`:
- Around line 95-141: The fixture currently depends on private
_process_single_row behavior and hard-coded output filenames. Update the
generator’s public contract to return the generated CIF and MTZ paths (or accept
explicit output names), then have the fixture consume that returned contract
instead of reconstructing names from source_cif and resolution. Preserve the
existing generation parameters and missing-output failure behavior.
In `@tests/rewards/test_structure_factor_reward.py`:
- Around line 460-470: Update the docstring of the mtz_all_free_1vme fixture to
describe the generated source MTZ accurately, replacing the reference to a
committed MTZ while preserving the documented flag column name and test-set
convention.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 35b347ea-e816-4f7d-a83d-16344a8fa70c
📒 Files selected for processing (4)
src/sampleworks/core/rewards/structure_factor.pytests/rewards/conftest.pytests/rewards/test_reward_function_contract.pytests/rewards/test_structure_factor_reward.py
marcuscollins
left a comment
marcuscollins
left a comment
There was a problem hiding this comment.
Thanks for the changes. I left a few more comments (look for unresolved comments) that you can think about and decide what to do, but after you look at those, feel free to merge.
| if self.bulk_solvent == "per_conformer": | ||
| # calc_fsolvent_batch masks each conformer (from Fprotein_asu_batch, set by | ||
| # calc_fprotein_batch); the mean applies the 1/E weight -> <mask(rho)>. | ||
| Fmask_HKL_batch = self.sfc.calc_fsolvent_batch( |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
src/sampleworks/core/rewards/structure_factor.py (1)
142-146: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTypos in the new docstring paragraph.
Niavely→Naively,overrride→override, andthere are not arguments→there are no arguments.✏️ Proposed fix
- the structure by the MTZ's whenever they disagree. Niavely overriding the cell - and/or the space group of MTZ can be dangerous because the reflections' indexing - could become wrong, so there are not arguments here to overrride them. + the structure by the MTZ's whenever they disagree. Naively overriding the cell + and/or the space group of MTZ can be dangerous because the reflections' indexing + could become wrong, so there are no arguments here to override them.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/sampleworks/core/rewards/structure_factor.py` around lines 142 - 146, Correct the typos in the docstring paragraph near SFcalculator metadata handling: change “Niavely” to “Naively,” “overrride” to “override,” and “there are not arguments” to “there are no arguments,” without altering the meaning.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@src/sampleworks/core/rewards/structure_factor.py`:
- Around line 142-146: Correct the typos in the docstring paragraph near
SFcalculator metadata handling: change “Niavely” to “Naively,” “overrride” to
“override,” and “there are not arguments” to “there are no arguments,” without
altering the meaning.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: dedbf114-d1b8-480e-b1ea-38b96c3965ca
📒 Files selected for processing (2)
src/sampleworks/core/rewards/structure_factor.pytests/rewards/test_structure_factor_reward.py
"E" collided with the normalized structure-factor amplitude |E| that `normalize_amplitude` scores, so the ensemble/batch dimension is now spelled `batch_size` throughout the docstrings and comments. Docs only; no behavior change. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2 participants
The older PR #272 is broken into 3 smaller PRs for (1) improve synthetic data generation, (2) refactor reward tests, and (3) add structure factor reward. This PR addresses (3).
Changes relative to main branch
StructureFactorRewardFunctionChanges relative to the previous PR #272
Most changes were addressing reviewer comments, except for MTZ parsing and inferring column types
_resolve_mtz_metadataand cached;mtzdatamoved out of_sfc_kwargsintoprepare()to reduce the number of file reading. Dropped_detect_mtz_metadata._resolve_expcolumnsrewritten to take the parsedrs.Dataset. Auto-detect mtz column now requires exactly one amplitude + one sigma column — a multi-set MTZ (Fprotein+Ftotal) raises instead of silently picking the first with a warning._build_reflection_maskto also validate the built mask retains a reasonable number of reflectionsSummary by CodeRabbit
New Features
Tests