Fix fragments with interleaved atom indices (#47)

[joshkamm]

Description

Fixes #47: PrimitiveInternalCoordinates crashed with a cryptic, message-less RuntimeError when an input geometry had two disconnected molecular fragments whose atom indices interleave (e.g. fragment A = atoms {0,2}, fragment B = {1,3}) instead of falling into separate contiguous ranges. One real-world trigger is a catalyst whose ligands bind/dissociate over a reaction mechanism, or (the case that surfaced this again) a mobile-proton H-bonded pair where fragment membership depends on which atom currently holds the proton.

Root cause

PrimitiveInternalCoordinates represented each connected-component fragment as a numeric window (min(atom_idx), max(atom_idx)+1) instead of its real atom-index set (frag.L() in topology.py, already available everywhere it's needed). This window is only a valid stand-in for "the atoms in this fragment" when those indices happen to be contiguous. With interleaved fragments, two fragments' windows overlap, and several call sites that assumed windows exactly partition [0, natoms) broke in different ways:

• get_hybrid_indices (the reported crash site): re-expanded each fragment's (min, max) window back into a range() and concatenated across fragments to find "leftover" hybrid (pure-Cartesian) atoms. With overlapping windows this produces duplicate indices; the second .remove() of an already-claimed index raised ValueError, silently caught by a bare except: and re-raised as a bare raise RuntimeError — the exact cryptic error this issue reports.
• reorderPrimsByFrag — reached whenever options['connect'] is False, which is the default, and live during SE-GSM growth (se_gsm.py). This turned out to be worse than "misassigns primitives to the wrong block": its primitive-matching loop iterates over all primitives once per block with no early exit, so a primitive whose atoms fall inside two overlapping windows gets appended to self.Internals twice, silently, with no exception at all.
• wilsonB (called every optimization step to build the Wilson B-matrix): I initially assumed the fix here could be "pass the full xyz array instead of a block slice for interleaved blocks," but verified that's wrong — block_matrix.full_matrix assembles the final matrix via scipy.linalg.block_diag(*matlist), which assumes each block owns an exclusive, correctly-positioned column range. Feeding it a wrong-sized block would misalign columns even though the per-primitive derivative math itself would be numerically fine.

Fix

Rather than patching each symptom, the fix restores the invariant everything else already assumes: block_info windows must always be true, non-overlapping, gap-free partitions of [0, natoms).

• Added _merge_block_windows, a small pure function that merges overlapping (sa, ea, sp, ep, kind) windows (standard sorted-interval-merge) into windows that are each exactly the true contiguous atom range they cover. This is correct because every raw window already comes from a disjoint atom set (connected components), so the union of all raw windows has no gaps — merging any overlapping chain therefore can't accidentally include atoms that don't belong. For ordinary, non-interleaved geometries, merged == raw exactly, so there's zero behavior or performance change in the common case.
• Wired this into both newMakePrimitives and reorderPrimsByFrag, replacing the raw block_info with the merged partition.
• Fixed reorderPrimsByFrag's primitive-duplication bug independently by matching primitives to a fragment via real membership (frag.L()) instead of the numeric window — this alone fixes duplication regardless of the window-merge work.
• Fixed get_hybrid_indices to build its atom list directly from each fragment's real frag.L() instead of re-expanding a (min,max) window, and replaced the bare RuntimeError with a descriptive one for the one case that should still be an error (an atom genuinely claimed by two fragments — a real topology inconsistency, not an interleaving artifact).
• block_info tuples grew a 5th field (kind, 'reg'/'hyb') so append_prim_to_block's hybrid-block detection (previously inferred from "exactly 3 primitives," which breaks once merged hybrid blocks can hold more) becomes an explicit tag instead of a heuristic.
• newMakePrimitives's addcart branch had the same contiguous-range assumption (for i in range(info[0], info[1])) when adding per-atom Cartesian primitives; fixed by mirroring the adjacent addtr branch, which already correctly iterates the fragment's real node list.

Because block_info windows are guaranteed correct partitions after this fix, wilsonB, second_derivatives, calcCg, GMatrix, and GInverse_SVD needed no changes at all — their block-local Cartesian slicing was never the actual problem, only the bookkeeping that decided block boundaries.

Two additional bugs found while verifying against real geometries

Static analysis of the plan wasn't enough — running the fix against an actual interleaved geometry surfaced two more issues in slots.py's Rotator class (used by RotationA/B/C, the translation/rotation primitives TRIC uses by default), both only reachable once a block can span more than one fragment's Translation/Rotation primitives:

1. Rotator.derivative selected xsel = xyz directly instead of xyz[relative_a, :] — there's even a # [relative_a, :] comment showing the correct slice was written and then disabled. This was silently correct before merging was possible, because a block always held exactly one Rotator's atoms (xyz == self.a's atoms already). Once two fragments share a merged block, this fed the whole block's atoms into a computation that expected only this primitive's own 2+ atoms, raising a shape-mismatch error.
2. Rotator.derivative's cache-hit path returned self.stored_deriv[relative_a] instead of self.stored_deriv, a shape inconsistent with the freshly-computed path — again silently correct only when relative_a happened to span the entire block. The maintainer had already left a comment flagging this exact cache as unreliable ("stored_der does not currently work in block-matrix formulism"), which is a good sign this was a real, previously-known-suspicious latent bug rather than something new.

Both are one-line-ish fixes with comments explaining why the old code was silently correct before and why the new code is needed now.

Todos

• Fix get_hybrid_indices (crash site)
• Fix reorderPrimsByFrag's primitive-duplication bug
• Fix newMakePrimitives's addcart contiguous-range bug
• Add _merge_block_windows and wire it into both primitive-building methods
• Fix append_prim_to_block's hybrid-block heuristic and add_union_primitives's tuple unpacking for the new 5-tuple block_info
• Fix the two Rotator.derivative bugs found during verification
• Add regression tests (pyGSM/tests/test_interleaved_fragments.py): crash regression, no-duplicate-primitives, block_info partition invariant, 3-fragment chained interleaving, and a permutation-equivalence check (same molecule built contiguous vs. interleaved, wilsonB and primitive values compared after permuting back)
• Verify no regression on ordinary (non-interleaved) geometries: byte-identical wilsonB output before/after on a synthetic case, and DelocalizedInternalCoordinates construction + reorderPrimitives() confirmed working unchanged on a real molecule (diels_alder.xyz)

Questions

• Is master...claude/fragment-indices-error-clarity-rd5s93 the branch naming/base you'd want, or should this target a different base branch?
• second_derivatives/calcCg (which the window-merge fix also covers) don't appear to have a live caller in the current GSM/optimizer driver flow — I didn't add direct tests for them beyond the fact that they share the now-fixed block_info invariant. Flag if you know of a call path I should exercise directly.

Status

• Ready to go — all new tests pass locally (pytest pyGSM/tests/test_interleaved_fragments.py), and existing lightweight tests (test_pygsm.py, plus manual DelocalizedInternalCoordinates/reorderPrimitives() construction on diels_alder.xyz) show no regression. I could not run test_basic_mecp.py's full optimization test in this environment (no xtb binary available), so CI running that would be good extra confidence.

---

Generated by Claude Code

[joshkamm]

Real-scale validation from a downstream project

Two Athena runs of asymmetric-aminoallylation's List GSM ensemble (185-atom reactant complexes, 2 non-covalently-bonded fragments whose atom-index membership genuinely varies per conformer — a mobile N–H···N proton) tested this fix at real scale, beyond this PR's own unit tests. Full detail in asymmetric-aminoallylation#3:

• Pre-fix (driver-side permute/unpermute workaround, not this PR): 165/331 conformers converged; 158 failures were the exact get_hybrid_indices crash this PR fixes, just relocated to the DE-GSM stage (a second call site the driver-level workaround didn't cover).
• With this PR (workaround removed entirely, atoms passed through in original interleaved order): 249/331 converged, zero recurrence of the crash across any of the 331 conformers, at any stage.

The remaining 82 failures are unrelated (24h walltime, an unrelated SE-GSM growth-stall bug, early xTB rejections, Hessian LinAlgError, etc. — breakdown in the linked issue).

This also plausibly exercises the Rotator.derivative fix, not just get_hybrid_indices: this system's 2-fragment split overlaps by exactly one atom, so for most conformers the window-merge logic collapses both fragments' Translation/Rotation primitives into a single block_info block — the specific scenario that bug required to manifest.

Not exercised by this validation (gaps worth knowing about, not blockers): 3+ fragment interleaving (this system is always exactly 2 fragments — only your own synthetic 3-fragment test covers that), the MECI/SEAM/TS-SEAM/BEALES_CG code paths (this driver only uses standard SE-GSM/DE-GSM), and non-xTB levels of theory.

---

Generated by Claude Code