feat(rewards): add reciprocal-space reward (Fprotein from SFC)#272
feat(rewards): add reciprocal-space reward (Fprotein from SFC)#272DorisMai wants to merge 14 commits into
Conversation
📝 WalkthroughWalkthroughAdds reciprocal-space structure-factor reward scoring, shared AtomArray-to-gemmi conversion, multi-label synthetic MTZ output, and expanded reward test coverage for contract, SF-specific, and real-space behavior. ChangesStructure Factor Reward and Synthetic SF Pipeline
Reward Validation and Tests
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 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: 7
🧹 Nitpick comments (2)
src/sampleworks/eval/generate_synthetic_sf.py (1)
137-145: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse a NumPy-style docstring for this new helper.
The summary is useful, but this new function is missing the required
Parameters/Returnssections. As per coding guidelines, “Always include NumPy-style docstrings for every function and class.”Proposed docstring update
- """Build a one-amplitude rs.DataSet with labelled F / SIGF / PHIF columns. - - ``sfc.prepare_dataset`` returns an amplitude column and a phase column (degrees) - for the given ``structure_factor_column`` attribute. We auto-detect those by MTZ - dtype (rather than assuming the unexposed ``FMODEL`` / ``PHIFMODEL`` names), - rename them to ``F{label}`` / ``PHIF{label}``, and synthesize a ``SIGF{label}`` - column so several structure-factor sets (e.g. protein and total) can coexist in - one MTZ. - """ + """Build a one-amplitude dataset with labelled F / SIGF / PHIF columns. + + Parameters + ---------- + sfc + Structure-factor calculator containing the requested ASU amplitudes. + label + Output column label suffix, e.g. ``protein`` or ``total``. + structure_factor_column + SFcalculator attribute passed to ``prepare_dataset``. + miller_index_column + SFcalculator attribute containing Miller indices. + sigma_f_scale + Scale factor used to synthesize dummy SIGF values from amplitudes. + + Returns + ------- + rs.DataSet + Dataset containing ``F{label}``, ``SIGF{label}``, and ``PHIF{label}``. + """🤖 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/eval/generate_synthetic_sf.py` around lines 137 - 145, The new helper in generate_synthetic_sf.py has a summary docstring but is missing the required NumPy-style structure. Update the docstring for the helper that builds the rs.DataSet to use NumPy format with explicit Parameters and Returns sections, documenting each input and the returned dataset/columns clearly; keep the existing behavior unchanged and ensure the docstring matches the function name and its role in auto-detecting and renaming the structure-factor columns.Source: Coding guidelines
tests/rewards/test_reward_function_contract.py (1)
47-66: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFreeze
RewardCasebefore sharing it across tests.This bundle is passed around as shared state, so leaving the dataclass mutable makes accidental test-side mutation hard to spot.
@dataclass(frozen=True)matches the repo's immutable-state convention and still works withbatch(). As per coding guidelines, "Use frozen dataclasses with functional updates for immutable state management."Suggested change
-@dataclass +@dataclass(frozen=True) class RewardCase:🤖 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_reward_function_contract.py` around lines 47 - 66, Make RewardCase immutable by marking the RewardCase dataclass as frozen so shared test state cannot be mutated accidentally. Update the RewardCase definition to use a frozen dataclass while keeping batch() unchanged, since it only reads fields and still works with immutable instances. Use the existing RewardCase symbol in tests/rewards/test_reward_function_contract.py to locate the class and apply the repo’s immutable-state convention.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 152-163: Validate the batch_partition argument eagerly in the
structure_factor constructor before assigning it on self, and raise a clear
ValueError when it is zero or negative. Add the check in the initializer where
device, mtzfile, and batch_partition are set so invalid OOM-tuning input is
rejected immediately instead of failing later in calc_fprotein_batch().
- Around line 69-78: The auto-detection in structure_factor.py is ambiguous
because it independently chooses the first amplitude and sigma columns, which
can silently pair the wrong MTZ labels in multi-label datasets. Update the
selection logic in the structure-factor loading path to either require explicit
expcolumns when more than one StructureFactorAmplitudeDtype or
StandardDeviationDtype candidate exists, or ensure the chosen sigma column is
matched to the selected amplitude label in the same ds.select_mtzdtype/return
path.
In `@src/sampleworks/eval/generate_synthetic_sf.py`:
- Around line 543-547: The help text in generate_synthetic_sf should not escape
the braces around label because it is a plain string, not an f-string. Update
the help text in the argument definition near the existing bulk solvent option
so users see F{label}/SIGF{label}/PHIF{label} in --help, and verify the same
wording is used consistently wherever that option text is defined.
In `@tests/eval/test_generate_synthetic_sf.py`:
- Around line 135-139: The round-trip test in test_generate_synthetic_sf is only
checking that each non-blank altloc label exists, so it can miss cases where
some atoms lose their altloc assignment. Update the assertion near
find_all_altloc_ids/loaded to verify multiplicity as well, either by comparing
per-label counts for altloc_id or by comparing the full altloc_id annotation
when order is stable. Keep the check black-box by asserting observable
annotation behavior rather than implementation details.
In `@tests/rewards/reward_input_helpers.py`:
- Around line 7-13: The shared helper build_scattering_indices() still has a
prose-only docstring, so update it to the repository’s NumPy-style format. Add
the standard sections for parameters, returns, and any relevant notes/details so
the contract is explicit for the reward tests that reuse it, while keeping the
behavior unchanged.
In `@tests/rewards/test_reward_function_contract.py`:
- Around line 248-270: The debug-only `test_gradient_descent_loss_trace` in
`test_reward_function_contract.py` should be removed from collected tests
because it only prints loss values and never asserts anything. Delete this
`test_...` method from `TestRewardFunctionContract`, or if you want to keep the
trace, rename it to a non-test helper so pytest won’t run it; make sure no
`print`-only optimization loop remains in the test suite.
- Around line 39-44: The module-level test marker in the reward contract suite
should include both GPU and slow tagging. Update the existing pytestmark
assignment in the test module so the suite remains GPU-marked via
pytest.mark.gpu while also adding pytest.mark.slow, keeping the change localized
to the module-level marker used by these reward contract tests.
---
Nitpick comments:
In `@src/sampleworks/eval/generate_synthetic_sf.py`:
- Around line 137-145: The new helper in generate_synthetic_sf.py has a summary
docstring but is missing the required NumPy-style structure. Update the
docstring for the helper that builds the rs.DataSet to use NumPy format with
explicit Parameters and Returns sections, documenting each input and the
returned dataset/columns clearly; keep the existing behavior unchanged and
ensure the docstring matches the function name and its role in auto-detecting
and renaming the structure-factor columns.
In `@tests/rewards/test_reward_function_contract.py`:
- Around line 47-66: Make RewardCase immutable by marking the RewardCase
dataclass as frozen so shared test state cannot be mutated accidentally. Update
the RewardCase definition to use a frozen dataclass while keeping batch()
unchanged, since it only reads fields and still works with immutable instances.
Use the existing RewardCase symbol in
tests/rewards/test_reward_function_contract.py to locate the class and apply the
repo’s immutable-state convention.
🪄 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
Run ID: b9ae1d00-b4d9-46d0-87da-383baa1d031f
📒 Files selected for processing (11)
src/sampleworks/core/rewards/structure_factor.pysrc/sampleworks/eval/generate_synthetic_sf.pysrc/sampleworks/eval/synthetic_utils.pytests/conftest.pytests/eval/test_generate_synthetic_sf.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/reward_input_helpers.pytests/rewards/test_real_space_density_reward.pytests/rewards/test_reward_function_contract.pytests/rewards/test_structure_factor_reward.py
DorisMai
force-pushed
the
dm/add-sfc-reward
branch
from
b3f83d6 to
de306a5
Compare
DorisMai
requested review from
a team,
Abdelsalam-Abbas,
k-chrispens and
mag-astera
as code owners
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/eval/synthetic_utils.py`:
- Around line 265-276: Residue grouping in the AtomArray-to-Gemmi conversion
currently ignores insertion codes, so residues with the same res_id but
different ins_code get merged. Update the residue identity used in the
synthetic_utils flow (the key comparison around current_res / prev_key) to
include atom_array.ins_code, and propagate that value into gemmi.SeqId when
creating each residue so author insertion codes are preserved on write.
🪄 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
Run ID: 17750ecd-c9a6-4df7-8992-9af5dd711124
📒 Files selected for processing (4)
src/sampleworks/eval/synthetic_utils.pytests/eval/test_generate_synthetic_sf.pytests/resources/1vme/1vme_final_crystalframe_0.5occA_0.5occB_1.80A.ciftests/rewards/test_structure_factor_reward.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/eval/test_generate_synthetic_sf.py
- tests/rewards/test_structure_factor_reward.py
| # New residue on any (chain_id, res_id) change. Keying on the pair (not res_id | ||
| # alone) keeps residues that share a res_id across a chain boundary separate. | ||
| if key != prev_key: | ||
| if prev_key is not None: | ||
| assert current_res is not None # prev_key set implies a residue exists | ||
| current_chain.add_residue(current_res) | ||
| current_res = gemmi.Residue() | ||
| current_res.name = str(atom_array.res_name[i]) | ||
| current_res.seqid = gemmi.SeqId(str(res_id)) # author numbering | ||
| current_res.label_seq = res_id # label numbering -> label_seq_id on write | ||
| current_res.subchain = chain_id # label_asym_id == chain name | ||
| prev_key = key |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check whether ins_code is present/used in the atom arrays flowing into atomarray_to_gemmi
rg -nP '\bins_code\b' -C2 --type=py
# Confirm biotite AtomArray exposes ins_code in this environment's version
rg -nP 'atomworks|biotite' pyproject.tomlRepository: diff-use/sampleworks
Length of output: 279
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the target file and locate the converter/callers.
ast-grep outline src/sampleworks/eval/synthetic_utils.py --view expanded || true
rg -n "atomarray_to_gemmi|ins_code|label_seq|res_id|SeqId" src/sampleworks -C 3
# Inspect the relevant slice of the converter.
sed -n '220,320p' src/sampleworks/eval/synthetic_utils.py
# Find where the atom arrays come from and whether any pipeline preserves insertion codes.
rg -n "get_annotation_categories|ins_code|AtomArray|AtomArrayStack|load_structure|to_atom_array|atom_array" src -C 2Repository: diff-use/sampleworks
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Narrow search for actual insertion-code handling in the repo.
rg -n "\bins_code\b|ins[-_ ]?code|insertion code" src tests .
# Inspect the exact converter slice with line numbers.
sed -n '250,290p' src/sampleworks/eval/synthetic_utils.py | cat -n
# Inspect the atom-array ID helpers, since they define the matching semantics used upstream.
sed -n '580,760p' src/sampleworks/utils/atom_array_utils.py | cat -n
# Check whether the code assumes res_id alone anywhere else for residue identity.
rg -n "\bres_id\b.*\bres_id\b|make_atom_id|make_normalized_atom_id|filter_to_common_atoms" src/sampleworks/utils/atom_array_utils.py src/sampleworks/eval/synthetic_utils.py -C 2Repository: diff-use/sampleworks
Length of output: 21337
🌐 Web query:
biotite AtomArray ins_code annotation PDBx PDB insertion code AtomArray load_any
💡 Result:
In Biotite, the insertion code is a mandatory annotation within an AtomArray, stored in the 'ins_code' array [1]. When loading structures from PDBx/mmCIF files, Biotite maps the PDBx field _atom_site.pdbx_PDB_ins_code to this 'ins_code' annotation [2]. Regarding your specific queries: 1. 'ins_code' annotation: Every AtomArray object in Biotite stores fundamental structural attributes as NumPy arrays, including 'ins_code' (residue insertion code), 'chain_id', 'res_id', 'res_name', 'atom_name', and 'element' [1][3]. These are essential components of the molecular structure representation in Biotite [1]. 2. PDBx/PDB integration: When using PDBx (mmCIF) files, the Biotite parser automatically attempts to extract the insertion code [2]. Note that historical issues have been reported where files lacking the _atom_site.pdbx_PDB_ins_code field caused errors, though this has been addressed in recent updates to better handle missing values [4]. 3. load_any: Biotite does not have a single monolithic 'load_any' function; instead, structure loading is handled through format-specific modules in biotite.structure.io, such as pdb.get_structure or pdbx.get_structure [5][6][7]. To load a structure and preserve attributes like insertion codes, you should use the get_structure method provided by the appropriate IO module (e.g., pdbx for PDBx/mmCIF files) [5][7]. These methods automatically populate the standard annotations, including 'ins_code', if the information is available in the source file [5][2].
Citations:
- 1: https://www.biotite-python.org/latest/tutorial/structure/atoms.html
- 2: https://github.com/biotite-dev/biotite/blob/master/src/biotite/structure/io/pdbx/convert.py
- 3: https://www.biotite-python.org/latest/tutorial/structure/index.html
- 4: Failure to parse mmCIF files lacking _atom_site.pdbx_PDB_ins_code biotite-dev/biotite#869
- 5: https://www.biotite-python.org/latest/apidoc/biotite.structure.io.pdbx.get_structure.html
- 6: https://www.biotite-python.org/latest/apidoc/biotite.structure.io.pdb.get_structure.html
- 7: https://www.biotite-python.org/latest/tutorial/structure/io.html
Include ins_code in residue identity src/sampleworks/eval/synthetic_utils.py:265-276
Biotite exposes ins_code on AtomArray, and these CIF inputs already carry _atom_site.pdbx_PDB_ins_code, so residues that share a res_id but differ by insertion code will be merged here and written back without the author insertion code. Fold ins_code into the residue boundary key and gemmi.SeqId.
🤖 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/eval/synthetic_utils.py` around lines 265 - 276, Residue
grouping in the AtomArray-to-Gemmi conversion currently ignores insertion codes,
so residues with the same res_id but different ins_code get merged. Update the
residue identity used in the synthetic_utils flow (the key comparison around
current_res / prev_key) to include atom_array.ins_code, and propagate that value
into gemmi.SeqId when creating each residue so author insertion codes are
preserved on write.
There was a problem hiding this comment.
Currently the entire repo does not use ins_code, and key on (chain_id, res_id) or (chain_id, res_id, atom_name). I am not sure when or how often we encounter ins_code. I did notice that Vratin uses (chain_id, res_id, ins_code) in WaterFlow. @marcuscollins Do you think this is worth opening an issue to fix repo-wide or low priority / assume no ins code handling? either way I don't think it should block this PR.
There was a problem hiding this comment.
@DorisMai Agreed this shouldn't block this PR — merging residues that share res_id but differ by ins_code is a pre-existing, repo-wide convention here rather than something introduced by this change. Since it's a broader design question (repo-wide key convention vs. one-off fix), it does seem worth tracking separately so it doesn't get lost, especially given the inconsistency you noticed with WaterFlow's (chain_id, res_id, ins_code) keying.
If you and @marcuscollins agree, I'm happy to open a follow-up GitHub issue summarizing this discussion (current repo convention, the WaterFlow inconsistency, and the specific spot in atomarray_to_gemmi as a starting point) — just let me know and I'll file it.
✏️ Learnings added
Learnt from: DorisMai
Repo: diff-use/sampleworks PR: 272
File: src/sampleworks/eval/synthetic_utils.py:265-276
Timestamp: 2026-07-06T17:02:04.734Z
Learning: In the diff-use/sampleworks repo, residue identity throughout the codebase (e.g., in `src/sampleworks/eval/synthetic_utils.py`'s `atomarray_to_gemmi`) is generally keyed on (chain_id, res_id) or (chain_id, res_id, atom_name), without incorporating `ins_code`. This is an intentional, repo-wide convention, not a bug specific to one function. Note: WaterFlow code (referenced by DorisMai/Vratin) uses (chain_id, res_id, ins_code) instead, which is an inconsistency that may warrant a future repo-wide follow-up issue, but should not block individual PRs.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
🧠 Learnings used
Learnt from: k-chrispens
Repo: diff-use/sampleworks PR: 73
File: src/sampleworks/core/samplers/protocol.py:1-18
Timestamp: 2026-02-02T19:13:51.594Z
Learning: In Python files across the repo, do not suggest or add inline or file-level suppressions for Ruff F722 (jaxtyping dimension strings). Instead rely on a global ignore configured in pyproject.toml (as done by the sampleworks project). Verify that F722 is included in the global ignore, and avoid recommending per-file # ruff: noqa: F722 directives unless explicitly documented as an exception.
Learnt from: marcuscollins
Repo: diff-use/sampleworks PR: 132
File: src/sampleworks/utils/guidance_script_utils.py:585-586
Timestamp: 2026-03-05T16:30:40.514Z
Learning: In Python code, if enums subclassing StrEnum are used (e.g., GuidanceType, StructurePredictor), they serialize to plain strings with json.dump and pickle without special handling. Do not flag these as non-serializable in code reviews. Treat StrEnum values as strings for JSON serialization and ensure tests cover that behavior; no extra pickle handling needed.
There was a problem hiding this comment.
@CodeRabbit go ahead and create the issue. Please indicate in the issue that we need to research when ins_code is used and why. That should inform our decision on what to do.
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
repo="diff-use/sampleworks"
body=$(cat <<'EOF'
## Summary
`atomarray_to_gemmi` in `src/sampleworks/eval/synthetic_utils.py` (and residue-identity logic elsewhere in the repo) keys residues on `(chain_id, res_id)` or `(chain_id, res_id, atom_name)`, without incorporating `ins_code`. Biotite `AtomArray` objects do expose an `ins_code` annotation, and CIF inputs can carry `_atom_site.pdbx_PDB_ins_code`, so residues that share a `res_id` but differ by insertion code could currently be merged/misidentified.
Separately, WaterFlow code keys residue identity on `(chain_id, res_id, ins_code)`, which is inconsistent with the convention used in this repo.
## Why this needs research first
Before deciding whether to change the repo-wide convention (or leave it as-is), we need to understand:
- How often, and in what real-world structures/datasets, `ins_code` is actually populated (non-blank) in inputs processed by this codebase.
- Whether any current pipelines (synthetic SF/density generation, reward functions, structure loading utilities) could produce silently incorrect results due to ignoring `ins_code`.
- Why WaterFlow chose to include `ins_code` in its identity key, and whether that reflects a real correctness requirement or just defensive coding.
## Affected areas
- `src/sampleworks/eval/synthetic_utils.py` (`atomarray_to_gemmi`)
- `src/sampleworks/utils/atom_array_utils.py` (residue/atom identity helpers)
- Any other residue-identity keying in the repo that currently omits `ins_code`
## Acceptance criteria
- [ ] Document findings on when/why `ins_code` is populated and relevant in practice (e.g., which structures/datasets), and whether it's already handled by upstream loading/filtering steps.
- [ ] Decide whether to adopt `(chain_id, res_id, ins_code)` repo-wide, keep the current convention, or handle it case-by-case.
- [ ] If a change is warranted, list the specific functions/files to update.
## References
- PR: https://github.com/diff-use/sampleworks/pull/272
- Discussion: https://github.com/diff-use/sampleworks/pull/272#discussion_r3530470906
- Requested by: `@marcuscollins`
EOF
)
gh issue create \
--repo "$repo" \
--title "Research: when/why is ins_code needed for residue identity?" \
--assignee "marcuscollins" \
--body "$body"Length of output: 2503
@marcuscollins Done — I've opened #306 to track this, framed as a research task on when/why ins_code is used and why, which should inform whether we adopt (chain_id, res_id, ins_code) repo-wide, keep the current convention, or handle it case-by-case. It includes backlinks to this PR and discussion, the affected files, and acceptance criteria.
(ᵔᴥᵔ)
🧠 Learnings used
Learnt from: DorisMai
Repo: diff-use/sampleworks PR: 272
File: src/sampleworks/eval/synthetic_utils.py:265-276
Timestamp: 2026-07-06T17:02:04.734Z
Learning: In diff-use/sampleworks, treat residue identity as following the repo convention: key residue-related data primarily by (chain_id, res_id) and, when needed, (chain_id, res_id, atom_name)—do not require ins_code to be included in these identity keys. For example, functions like atomarray_to_gemmi in src/sampleworks/eval/synthetic_utils.py should follow this approach. Note that some external/related code (e.g., WaterFlow) uses (chain_id, res_id, ins_code); this inconsistency may justify a future repo-wide follow-up, but it should not be used to block individual PRs that adhere to the existing convention.
52f093b to
56e9fd5
Compare
…fixing in prepare()
…y rewriting array2hier
DorisMai
force-pushed
the
dm/add-sfc-reward
branch
from
56e9fd5 to
c3f689f
Compare
marcuscollins
left a comment
marcuscollins
left a comment
There was a problem hiding this comment.
Lots of comments to look at. The biggest concern I have is about picking the first type-matching column for amplitudes, etc...
I'm also confused about how some of the changes in generate_synthetic_sf.py relate to the changes in structure_factor.py. They appear disconnected, but maybe I'm missing something. Can you explain?
| # New residue on any (chain_id, res_id) change. Keying on the pair (not res_id | ||
| # alone) keeps residues that share a res_id across a chain boundary separate. | ||
| if key != prev_key: | ||
| if prev_key is not None: | ||
| assert current_res is not None # prev_key set implies a residue exists | ||
| current_chain.add_residue(current_res) | ||
| current_res = gemmi.Residue() | ||
| current_res.name = str(atom_array.res_name[i]) | ||
| current_res.seqid = gemmi.SeqId(str(res_id)) # author numbering | ||
| current_res.label_seq = res_id # label numbering -> label_seq_id on write | ||
| current_res.subchain = chain_id # label_asym_id == chain name | ||
| prev_key = key |
There was a problem hiding this comment.
@CodeRabbit go ahead and create the issue. Please indicate in the issue that we need to research when ins_code is used and why. That should inform our decision on what to do.
| # Ftotal uses the default (unrefined) scales kiso=1, kmask=0.35, small uaniso, | ||
| # matching generate_synthetic_sf's Ftotal. They don't depend on coordinates, so | ||
| # set them once here. _set_scales only needs atom_pos_frac (set at construction) | ||
| # and n_bins (set by the experiment init) for dtype/device and bin count. |
There was a problem hiding this comment.
what is the "experiment init"?
There was a problem hiding this comment.
SFC's internal method (SFcalculator.init_mtz()) called in its own __init__
| if self.bulk_solvent != "off": | ||
| self.sfc._set_scales(requires_grad=False) | ||
|
|
||
| # |Eo| are computed in SFC's experiment init (inside a try/except). |
There was a problem hiding this comment.
remind me, what is Eo?
There was a problem hiding this comment.
observed amplitudes normalized in each resolution bin
| unit_cell=self.unit_cell, | ||
| space_group=self.space_group, | ||
| ) | ||
| self.sfc = SFcalculator(pdbmodel=PDBParser(gemmi_structure), **self._sfc_kwargs) |
There was a problem hiding this comment.
so does self.sfc contain the actual amplitudes and sigmas? I don't see anywhere here that explicitly stores the data in the mtz file. I would call out somewhere where that data lives to make the code easier to understand.
There was a problem hiding this comment.
updated the docstring of prepare to describe the sfc object better
| # Multi-conformer combination: complex sum over the ensemble [batch, n_hkl]. | ||
| # occ = 1/E (set per call) makes the summed |F| the multi-conformer total. | ||
| Fprotein_batch = self.sfc.calc_fprotein_batch( | ||
| coordinates, Return=True, PARTITION=self.batch_partition |
There was a problem hiding this comment.
omg we need to talk to Minhuan about his Python style. capital letters kwarg names!?
My previous PR implemented atomarray_to_gemmi() that correctly feed into SFC, and optionally saved a cif file for my visual inspection. However, that cif file doesn't correctly round trip back to atomarray. I noticed this issue/needed to correct it in order to generate the tests for the structure factor reward function, which happens in the middle of this PR, hence the diffs look a bit messy. |
2 participants
What changed
StructureFactorRewardFunction(core/rewards/structure_factor.py) via SFC. Two-phase construction (__init__config +prepare(atom_array)) as SFC requires knowing topology.|Ftotal| = |Fprotein| + |Fsolvent|. The solvent part can have different modes: off (Fprotein only), combined (one mask for averaged protein), or one mask per conformation. This requires SFC to be the latest version that fixes the solvent scaling issue.normalize_amplitudefor testing andbatch_partitionin case of OOMeval/generate_synthetic_sf.pyto generate test data with both Fprotein and Ftotal in the same mtz for debugging/development .eval/synthetic_utils.py. Tests added intests/eval/test_generate_synthetic_sf.py.test_real_space_density_reward.pytotest_reward_function_contract.py. Real space specific tests that are unused (e.g. vmap related) remain untouched.test_structure_factor_reward.pyand synthetic test data (1vme cif and mtz) generated by:Next steps
Summary by CodeRabbit
Summary of changes
New Features
off,combined,per_conformer).Bug Fixes
Tests