feat(synthetic): validate atom array hierarchy before gemmi conversion - #351
feat(synthetic): validate atom array hierarchy before gemmi conversion#351DorisMai wants to merge 3 commits into
Conversation
📝 WalkthroughWalkthroughChangesSynthetic hierarchy validation
Estimated code review effort: 4 (Complex) | ~45 minutes 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 |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/synthetic/test_generate_synthetic_sf.py (1)
84-116: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a NumPy-style Parameters section to
_residue_span_array.The function takes 6 parameters with non-obvious default semantics (e.g.
hetero=Nonevsaltloc_id=Nonebehave differently — one defaults to all-False, the other leaves the annotation unset). A proper Parameters/Returns section would make this contract explicit rather than relying on prose notes.As per coding guidelines,
**/*.py: "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/synthetic/test_generate_synthetic_sf.py` around lines 84 - 116, Expand the `_residue_span_array` docstring with a NumPy-style Parameters section documenting all six arguments, including the distinct `hetero=None` and `altloc_id=None` behaviors, and add a Returns section describing the constructed `AtomArray`.Source: Coding guidelines
src/sampleworks/synthetic/synthetic_utils.py (1)
350-384: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider jaxtyping for the new shape-carrying
np.ndarrayparams.
residue_boundary_mask(Lines 350, 497) andresidue_span_start_idx(Line 420) are typed as barenp.ndarraywith shape documented only via inline comments (e.g.# (n_atoms - 1,) bool). Using jaxtyping annotations (e.g.Bool[np.ndarray, "n_atoms-1"]) would make these shapes machine-checkable and consistent with the guideline for this path.♻️ Example annotation for one signature
+from jaxtyping import Bool + def _check_residue_fields_homogeneous( - atom_array: AtomArray, residue_boundary_mask: np.ndarray + atom_array: AtomArray, residue_boundary_mask: Bool[np.ndarray, "n_atoms-1"] ) -> None:As per coding guidelines,
src/sampleworks/**/*.py: "Use jaxtyping annotations to document array and tensor shapes where applicable."Also applies to: 420-439, 460-507
🤖 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/synthetic/synthetic_utils.py` around lines 350 - 384, Update the shape-carrying array parameters in _check_residue_fields_homogeneous and the related residue-span helpers to use jaxtyping annotations, including Bool for residue_boundary_mask and an appropriate integer-shaped annotation for residue_span_start_idx. Add or reuse the required jaxtyping imports, encode the documented dimensions in the annotations, and remove redundant inline shape comments while preserving behavior.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.
Nitpick comments:
In `@src/sampleworks/synthetic/synthetic_utils.py`:
- Around line 350-384: Update the shape-carrying array parameters in
_check_residue_fields_homogeneous and the related residue-span helpers to use
jaxtyping annotations, including Bool for residue_boundary_mask and an
appropriate integer-shaped annotation for residue_span_start_idx. Add or reuse
the required jaxtyping imports, encode the documented dimensions in the
annotations, and remove redundant inline shape comments while preserving
behavior.
In `@tests/synthetic/test_generate_synthetic_sf.py`:
- Around line 84-116: Expand the `_residue_span_array` docstring with a
NumPy-style Parameters section documenting all six arguments, including the
distinct `hetero=None` and `altloc_id=None` behaviors, and add a Returns section
describing the constructed `AtomArray`.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 7c267525-982e-4ad1-aae1-0e65b5aa6915
📒 Files selected for processing (2)
src/sampleworks/synthetic/synthetic_utils.pytests/synthetic/test_generate_synthetic_sf.py
05bff78 to
34e735d
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/synthetic/test_generate_synthetic_sf.py (1)
93-100: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse a NumPy-style docstring for
_build_atom_array.Document its parameters and
AtomArrayreturn value withParametersandReturnssections.Proposed docstring shape
- """Build a minimal AtomArray for residue-span validation tests. + """Build a minimal AtomArray for residue-span validation tests. - Only the fields a test varies need to be passed; coords are distinct per atom and - element/b_factor/occupancy are uniform, since none of them participate in residue - grouping or span validation. ``hetero`` defaults to biotite's all-False, and - ``altloc_id`` is left unset (not blank) when omitted, exercising the - missing-annotation path in ``_resolve_altlocs_for_gemmi``. + Parameters + ---------- + chain_id, res_id, res_name, atom_name + Per-atom hierarchy annotations. + hetero + Optional per-atom hetero flags. + altloc_id + Optional per-atom alternate-location identifiers. + + Returns + ------- + AtomArray + Minimal array with deterministic coordinates and required annotations. """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/synthetic/test_generate_synthetic_sf.py` around lines 93 - 100, Update the _build_atom_array docstring to use NumPy-style sections: add a Parameters section documenting each argument and a Returns section describing the returned AtomArray, while preserving the existing behavioral notes.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.
Nitpick comments:
In `@tests/synthetic/test_generate_synthetic_sf.py`:
- Around line 93-100: Update the _build_atom_array docstring to use NumPy-style
sections: add a Parameters section documenting each argument and a Returns
section describing the returned AtomArray, while preserving the existing
behavioral notes.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: dc429e78-3dce-498a-91bf-1aa77402da97
📒 Files selected for processing (2)
src/sampleworks/synthetic/synthetic_utils.pytests/synthetic/test_generate_synthetic_sf.py
🚧 Files skipped from review as they are similar to previous changes (1)
- src/sampleworks/synthetic/synthetic_utils.py
marcuscollins
left a comment
There was a problem hiding this comment.
I think there are some improvements you could make here (see comments), so I'd encourage you to make those before merging, but I will approve and leave changes up to you.
| Structure whose atoms are grouped into residues. Atoms of a residue are | ||
| assumed contiguous (true for arrays loaded in file order). | ||
| Structure to check. | ||
| residue_boundary_mask : np.ndarray |
There was a problem hiding this comment.
Can't this come from the AtomArray itself? Is there a reason to calculate it separately outside this function? If it is computed inside this function, you guard against someone passing an incorrect boundary mask.
| atom_array.chain_id.tolist(), | ||
| atom_array.res_id.tolist(), | ||
| atom_array.atom_name.tolist(), | ||
| altlocs, |
There was a problem hiding this comment.
Probably altlocs should be derived from the AtomArray directly as well?
| altlocs, | ||
| ), | ||
| level="atom", | ||
| identity="(chain_id, res_id, atom_name, altloc)", |
There was a problem hiding this comment.
It might be good to derive the iterator passed to keys from these identity names, something like
identities = ("chain_id", "res_id", "atom_name", "altloc")
zipped = zip(atom_array.get(k).tolist() for k in identities
_check_keys_unique(zipped, level="atom", identities = str(identities))that way if anything ever changes you only have to change it one place.
Not likely to change, so I wouldn't block this PR on this change, but something to use in the future at least.
There was a problem hiding this comment.
You could even make another intermediate method, check_atom_array_keys say, that accepted the AtomArray, level, and identities, that you could re-use in several places here.
| ---------- | ||
| atom_array : AtomArray | ||
| Structure to check. | ||
| residue_span_start_idx : np.ndarray |
There was a problem hiding this comment.
Again I think I'd get this from the AtomArray internally to avoid mistakes.
| # (n_atoms - 1,) bool | ||
| residue_boundary_mask = (chain_id[1:] != chain_id[:-1]) | (res_id[1:] != res_id[:-1]) | ||
| # (n_residues,) int | ||
| residue_span_start_idx = np.flatnonzero(np.concatenate([[True], residue_boundary_mask])) |
There was a problem hiding this comment.
Okay, so I guess the point of creating these here, rather than inside the check* methods, is that residue_boundary_mask gets used in more than one place. If that's the case, it would make sense to make two separate methods _get_residue_bndy_mask and _get_residue_span_starts, so that you can individually unit test them.
| with pytest.raises(ValueError, match="empty AtomArray"): | ||
| atomarray_to_gemmi(AtomArray(0)) | ||
|
|
||
| def test_single_atom_array_converts(self): |
There was a problem hiding this comment.
This honestly doesn't seem like a particularly useful test.
Address Issue #340 , related to recently merged PR #322
atomarray_to_gemmiin generating synthetic SF data assumed contiguous chains/residues and read per-residue fields(
res_name,hetero, seqid, subchain) from each residue's first atom. For defensive programming, this PR added validation checks during the atomarray --> gemmi conversion againsgt malformed inputs and added tests on the validation checks. The behavior on generating synthetic data from valid files should be unchanged.Specifcally,
_prepare_residue_spansis a new function that first validates these assumptions before yielding the residue spans in atom array to build the gemmi Structure. Validation checks include:(chain_id, res_id)residue keys occupy one contiguous blockres_nameandhetero(atom_name, altloc)is unique within a residue (gemmi's definition)Tests: new
TestGemmiHierarchyValidationcovers each rejection case.Summary by CodeRabbit
Bug Fixes
Tests