Skip to content

feat(synthetic): validate atom array hierarchy before gemmi conversion - #351

Open
DorisMai wants to merge 3 commits into
mainfrom
dm/validate-residue-spans
Open

feat(synthetic): validate atom array hierarchy before gemmi conversion#351
DorisMai wants to merge 3 commits into
mainfrom
dm/validate-residue-spans

Conversation

@DorisMai

@DorisMai DorisMai commented Jul 28, 2026

Copy link
Copy Markdown
Collaborator

Address Issue #340 , related to recently merged PR #322

atomarray_to_gemmi in 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_spans is 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 ids occupy one contiguous block each
  • (chain_id, res_id) residue keys occupy one contiguous block
  • atoms within a residue agree on res_name and hetero
  • (atom_name, altloc) is unique within a residue (gemmi's definition)

Tests: new TestGemmiHierarchyValidation covers each rejection case.

Summary by CodeRabbit

  • Bug Fixes

    • Strengthened structure conversion validation to catch malformed chain/residue/span data earlier.
    • Improved diagnostics for duplicate atoms, inconsistent per-residue fields, and noncontiguous residue or chain sequences.
    • Maintains correct behavior for alternate locations and prevents incorrect merging when residue IDs are shared across chain boundaries.
  • Tests

    • Expanded automated coverage for hierarchy validation, including edge cases for empty inputs, mixed residue attributes, alternate locations, and duplicate handling.
    • Refreshed test data setup to generate minimal, topology-controlled inputs.

@coderabbitai

coderabbitai Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

Changes

Synthetic hierarchy validation

Layer / File(s) Summary
Validated conversion pipeline
src/sampleworks/synthetic/synthetic_utils.py
atomarray_to_gemmi validates contiguous chain and residue spans, homogeneous residue fields, and unique (atom_name, altloc) pairs before constructing gemmi structures.
Hierarchy validation coverage
tests/synthetic/test_generate_synthetic_sf.py
Adds focused tests for valid and invalid spans, chain separation, alternate locations, empty inputs, and duplicate atoms; relocates the occupancy test without changing its behavior.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

Suggested reviewers: marcuscollins

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly summarizes the main change: adding hierarchy validation before Gemmi conversion.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches 💡 1
🛠️ Fix failing CI checks 💡
  • Fix failing CI checks
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch dm/validate-residue-spans

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

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

@DorisMai
DorisMai marked this pull request as ready for review July 28, 2026 21:04
@DorisMai
DorisMai requested review from a team, k-chrispens and marcuscollins as code owners July 28, 2026 21:04

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (2)
tests/synthetic/test_generate_synthetic_sf.py (1)

84-116: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a NumPy-style Parameters section to _residue_span_array.

The function takes 6 parameters with non-obvious default semantics (e.g. hetero=None vs altloc_id=None behave 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 win

Consider jaxtyping for the new shape-carrying np.ndarray params.

residue_boundary_mask (Lines 350, 497) and residue_span_start_idx (Line 420) are typed as bare np.ndarray with 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

📥 Commits

Reviewing files that changed from the base of the PR and between 089f89a and 05bff78.

📒 Files selected for processing (2)
  • src/sampleworks/synthetic/synthetic_utils.py
  • tests/synthetic/test_generate_synthetic_sf.py

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
tests/synthetic/test_generate_synthetic_sf.py (1)

93-100: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Use a NumPy-style docstring for _build_atom_array.

Document its parameters and AtomArray return value with Parameters and Returns sections.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 05bff78 and c50a13b.

📒 Files selected for processing (2)
  • src/sampleworks/synthetic/synthetic_utils.py
  • tests/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 marcuscollins left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Probably altlocs should be derived from the AtomArray directly as well?

altlocs,
),
level="atom",
identity="(chain_id, res_id, atom_name, altloc)",

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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]))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This honestly doesn't seem like a particularly useful test.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants