feat(rewards): add reciprocal-space reward (Fprotein from SFC)

[DorisMai]

What changed

• Added StructureFactorRewardFunction (core/rewards/structure_factor.py) via SFC. Two-phase construction (__init__ config + prepare(atom_array)) as SFC requires knowing topology.
  • ensembles combine as a complex sum.: v1 fits |Fprotein|
  • restricted by SFC to occupancy and B-factor shared across batch
  • exposes normalize_amplitude for testing and batch_partition in case of OOM
• Patched eval/generate_synthetic_sf.py to generate test data with both Fprotein and Ftotal in the same mtz for debugging/development and support the round trip of cif --> atomarray --> gemmi --> cif --> atomarray.
• Refactored and added reward tests.
  • Extracted common functions from test_real_space_density_reward.py to test_reward_function_contract.py. Real space specific tests that are unused (e.g. vmap related) remain untouched.
  • Added test_structure_factor_reward.py and synthetic test data (1vme cif and mtz) generated from eval/generate_synthetic_sf.py.

Next steps before PR ready to review

• Visual inspection on a few cases for validation
• Add unscaled |Ftotal| (should be trivial)

Next steps

• Wire the reward into a guidance step scaler for initial 40 experiment
• Add scaling optimization

[coderabbitai]

Important

Review skipped

Draft detected.

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

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 3b4df947-d9f5-4f67-86ba-f6db92a10217

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

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch dm/add-sfc-reward

---

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

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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