feat(rewards): add reciprocal-space reward (Fprotein from SFC) #272
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -4,12 +4,14 @@ | |
| import traceback | ||
| from pathlib import Path | ||
|
|
||
| import gemmi | ||
| from atomworks.io.transforms.atom_array import remove_waters | ||
| from biotite.structure import AtomArray | ||
| from loguru import logger | ||
| from sampleworks.eval.structure_utils import apply_selection | ||
| from sampleworks.utils.atom_array_utils import ( | ||
| AltlocInfo, | ||
| BLANK_ALTLOC_IDS, | ||
| detect_altlocs, | ||
| keep_amino_acids, | ||
| keep_polymer, | ||
|
|
@@ -194,3 +196,106 @@ def load_structure_for_synthetic_reward( | |
| raise ValueError(f"Invalid occupancy mode '{occupancy_mode}'") | ||
|
|
||
| return atom_array | ||
|
|
||
|
|
||
| def atomarray_to_gemmi( | ||
| atom_array: AtomArray, | ||
| unit_cell: gemmi.UnitCell | None = None, | ||
| space_group: str | None = None, | ||
| ) -> gemmi.Structure: | ||
| """Convert a biotite AtomArray to a gemmi.Structure for SFcalculator. | ||
|
|
||
| Anisotropic B-factors are set to zero since biotite does not store them. | ||
| Blank altloc labels are converted from biotite's '' to gemmi's '\\x00'. If | ||
| the atom array has no ``altloc_id`` annotation (e.g. arrays reconstructed by | ||
| a model wrapper), all altlocs default to blank. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| atom_array | ||
| Input structure with occupancy and b_factor annotations | ||
| unit_cell | ||
| Crystallographic unit cell for the structure. If None, gemmi defaults | ||
| to (1.0, 1.0, 1.0, 90.0, 90.0, 90.0) in units of Angstroms and degrees. | ||
| space_group | ||
| Space group (in Hermann-Mauguin string format) for the structure. If | ||
| empty or invalid, SFcalculator defaults to P1. | ||
|
|
||
| Returns | ||
| ------- | ||
| gemmi.Structure | ||
| Structure ready to be wrapped by SFC_Torch.io.PDBParser | ||
|
|
||
| Notes | ||
| ----- | ||
| Built directly rather than via ``SFC_Torch.io.array2hier``: that helper packs atom | ||
| identity into a hyphen-delimited ``cra_name`` string, sets only the author seqid, and | ||
| keys residue boundaries on ``res_id`` alone (merging residues that share a res_id across | ||
| a chain boundary). Building the hierarchy here from the AtomArray annotations avoids the | ||
| string round-trip and lets us set the author seqid, ``label_seq`` (so ``label_seq_id`` | ||
| survives a write -> atomworks reload; otherwise res_id collapses to -1), and ``subchain`` | ||
| (a <=1-char ``label_asym_id`` that SFcalculator's PDB-header step accepts) at construction. | ||
| """ | ||
| n = len(atom_array) | ||
|
DorisMai marked this conversation as resolved.
|
||
| # altloc_id is not a mandatory biotite annotation; default to blank when absent. | ||
| # gemmi uses '\x00' for blank altloc. | ||
| if "altloc_id" in atom_array.get_annotation_categories(): | ||
| atom_altloc = ["\x00" if a in BLANK_ALTLOC_IDS else a for a in atom_array.altloc_id] | ||
| else: | ||
| atom_altloc = ["\x00"] * n | ||
|
|
||
| model = gemmi.Model("1") # numeric name -> valid mmCIF pdbx_PDB_model_num | ||
| current_chain: gemmi.Chain | None = None | ||
| current_res: gemmi.Residue | None = None | ||
| prev_key: tuple[str, int] | None = None # (chain_id, res_id) of the previous atom | ||
|
|
||
| for i in range(n): | ||
|
DorisMai marked this conversation as resolved.
|
||
| chain_id = str(atom_array.chain_id[i]) | ||
| res_id = int(atom_array.res_id[i]) | ||
| key = (chain_id, res_id) | ||
|
|
||
| if current_chain is None or chain_id != current_chain.name: | ||
|
DorisMai marked this conversation as resolved.
|
||
| if current_chain is not None: | ||
| assert current_res is not None # a chain always holds >=1 residue here | ||
| current_chain.add_residue(current_res) | ||
| model.add_chain(current_chain) | ||
| current_chain = gemmi.Chain(chain_id) | ||
| prev_key = None # force a fresh residue for the new chain | ||
|
|
||
| # 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 | ||
|
Comment on lines
+265
to
+276
Contributor
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:
💡 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:
Include 🤖 Prompt for AI Agents
Collaborator
Author
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.
Contributor
There was a problem hiding this comment.
If you and ✏️ Learnings added
🧠 Learnings used
Collaborator
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.
Contributor
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
(ᵔᴥᵔ) 🧠 Learnings used |
||
|
|
||
| assert current_res is not None # created on the first atom / every chain switch | ||
|
DorisMai marked this conversation as resolved.
|
||
| atom = gemmi.Atom() | ||
| atom.name = str(atom_array.atom_name[i]) | ||
| atom.element = gemmi.Element(str(atom_array.element[i])) | ||
| atom.pos = gemmi.Position(*atom_array.coord[i]) # coord[i]: (3,) float | ||
| atom.b_iso = float(atom_array.b_factor[i]) | ||
| atom.aniso = gemmi.SMat33f(0, 0, 0, 0, 0, 0) # biotite stores no anisotropic B | ||
| atom.occ = float(atom_array.occupancy[i]) | ||
| atom.altloc = atom_altloc[i] | ||
| current_res.add_atom(atom) | ||
|
|
||
| if current_chain is not None: # flush the trailing residue + chain | ||
| assert current_res is not None # non-empty input -> a residue was built | ||
| current_chain.add_residue(current_res) | ||
| model.add_chain(current_chain) | ||
|
|
||
| structure = gemmi.Structure() | ||
| structure.add_model(model) | ||
| structure.setup_entities() # SFcalculator/PDBParser expects entities assigned | ||
| if unit_cell is not None: | ||
| structure.cell = unit_cell | ||
| if space_group is not None: | ||
| structure.spacegroup_hm = space_group | ||
| return structure | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I've been confused in the past about what the standard names of these are. Are there other common conventions? I feel like there were a couple. Should we address that?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Each refinement program has there own names used, so I don't think there is a standard (which is why auto-detect by column type). Usually F* is for amplitude, SIGF* is for amplitude uncertainty, PHIF* is for phase, and that's the convention I am following here.