Backbone-only side-chain completion: scorable PoseStack from N/CA/C/O (follow-up to #465) - #380

The gap

Every structure-prediction adapter on master requires the side-chain heavy
atoms to already be there. pose_stack_from_canonical_form builds missing
leaf atoms but rejects blocks missing non-leaf atoms
(fail_on_missing_nonleaf_atoms=not return_block_has_missing_atoms), so
pose_stack_from_openfold, pose_stack_from_atomworks and
pose_stack_from_atom37_and_biotite all reject backbone-only input. Only
pose_stack_from_biotite completes side chains, and it needs a caller-supplied
AtomArray.

Backbone generators — RFdiffusion, mosaic, partial-diffusion and hallucination
loops — emit exactly N/CA/C/O and an integer sequence, and nothing else. Today
the only way to score one is to write a PDB and read it back.

This PR adds that one missing step: bare (L, 4, 3) backbone plus residue-type
indices in, scorable PoseStack out.

API

ps = tmol.pose_stack_from_backbone_coords(
    coords,      # (L, 4, 3) or (n_poses, L, 4, 3) — N, CA, C, O
    res_types,   # indices into "ARNDCQEGHILKMFPSTWYV"; -1 is padding
    chain_id,
    device,
    sidechain_completion="pack",   # or "none"
)

sidechain_completion is a required decision, spelled out rather than implied:

• "pack" (default) runs build_missing_sidechains — a full
  score-function-driven Dunbrack + OptH packing job. The docstring says so at
  the call site: expensive, dependent on beta2016 and the Dunbrack library, and
  the rebuilt atoms are not a differentiable function of the input.
• "none" does the coordinate conversion only. Absent side chains stay
  NaN, the result is not scorable, and the conversion is cheap and fully
  differentiable — for callers that place side chains themselves or want to
  batch the packing.

Under both modes the supplied backbone is returned bit-for-bit and stays on
the autograd tape: after packing, every finite input atom is routed back to the
canonical tensor through the same _restore_canonical_input_coords mechanism
pose_stack_from_atom37_and_biotite uses.

Two things the new tests turned up

Backbone-only input leaves a NaN HA on every non-GLY residue. HA's
internal-coordinate frame is CA / N / CB, so build_missing_leaf_atoms — which
runs before packing — cannot place it when CB does not exist yet, and the
packer only rebuilds heavy atoms and never revisits it. On ubiquitin that was 70
NaN atoms in the 70 blocks flagged as incomplete. "pack" therefore rebuilds the
pose once after packing, when the heavy atoms exist; that second pass costs
~0.00 s against ~0.70 s for the packing itself and deliberately does not pass
return_block_has_missing_atoms, so anything still absent raises rather than
silently returning NaN. This is not specific to this adapter — any backbone-only
input reaching pose_stack_from_canonical_form hits it.

Side-chain placement is not reproducible from a torch seed on CPU. The CPU
simulated annealer draws from libc rand()
(tmol/pack/compiled/compiled.cpu.cpp, lines 29 / 109 / 220), which
torch.manual_seed does not control; repeated packs of the same input differ by
up to ~10 Å. test_default_pose_builds_reuse_packing_setup passes only because
a complete 1UBQ needs no Dunbrack packing and takes the OptH-only path.

That reaches topology too, not just coordinates: the HIS tautomer is resolved
from the packed side chain, so ubiquitin's His68 comes back as HIS on some
builds and HIS_D on others. The test here therefore asserts what is actually
guaranteed — setup reuse, atom layout, and a bit-identical backbone — and says
why in the docstring, rather than asserting a reproducibility the library does
not currently provide. Happy to tighten it if you would rather seed the
annealer.

Reuse — no new format-specific caches

The chemistry is shared verbatim with the atomworks adapter, whose protein
subset is already exactly what a backbone needs:

• _paramdb_for_atomworks / _restype_set_for_atomworks /
  canonical_ordering_for_atomworks / packed_block_types_for_atomworks
• _get_aw_2_tmol_mappings for the atom and residue-type tables

Two facts make that reuse exact, and both are asserted in the tests rather than
assumed: AlphaFold2's "ARNDCQEGHILKMFPSTWYV" is alphabetical by three-letter
code and therefore equals ATOMWORKS_NAME3S[1:21], and N/CA/C/O occupy atom37
slots (0, 1, 2, 4) for all 20 canonical amino acids. Padding maps to the
atomworks <M> token, which the existing tables already send to tmol restype
-1 with no real atoms.

The packing setup is held in a memoized PoseBuildContext, so the score
function and Dunbrack sampler are built once per device across repeated calls,
not once per structure.

Tests

tmol/tests/io/test_pose_stack_from_backbone_coords.py

Changes from the previous revision of this PR

• Rebased onto current master; the PoseStack.split / PDBInfo.split
  cherry-pick is dropped now that Add code to split out a single Pose from a PoseStack #379 is merged (it was the source of the
  conflicts).
• Dropped _paramdb_for_backbone_coords, canonical_ordering_for_backbone_coords
  and packed_block_types_for_backbone_coords in favour of the atomworks ones,
  and dropped the module's own backbone atom-name table in favour of
  _get_aw_2_tmol_mappings. Public surface is down from five names to two.
• Side-chain completion is now an explicit sidechain_completion argument
  instead of an unconditional hidden step.
• The supplied backbone is now restored after packing and stays differentiable.
• Tests rewritten around backbone preservation, side-chain completeness,
  device equivalence and setup reuse rather than shape and device alone. Those
  tests are what caught the NaN HA, which the previous revision shipped.

Two notes on points raised in review that I did not change, for
transparency: CanonicalForm.coords is Tensor[torch.float32] by contract, so
the float32 cast is the framework's, not this adapter's; and the protein-only
restype subset is the same one pose_stack_from_atomworks and
pose_stack_from_openfold enforce. Happy to follow if either is generalized
repo-wide.

Finally — if you'd rather this live inside _pose_stack_from_atomworks.py as a
sibling of pose_stack_from_atomworks, or would rather relax
canonical_form_from_atomworks to accept its own <M> mask token as padding
(which would reduce this to a thin wrapper), say the word and I'll move it.