Add pose_stack_from_backbone_coords: in-memory PoseStack construction from N/CA/C/O

[ssiddhantsharma]

Motivation

pose_stack_from_pdb requires writing backbone coordinates to a temporary PDB
file and reading it back — expensive I/O for workloads (gradient-based
hallucination, ΔG scoring loops) that already have coordinates in memory as
tensors.

Changes

New: tmol/io/pose_stack_from_backbone_coords.py

Builds a PoseStack directly from (N, CA, C, O) backbone coordinates and
integer residue type indices, with no file I/O. Follows the same adapter
pattern as pose_stack_from_openfold and pose_stack_from_rosettafold2.
Missing side-chain leaf atoms are reconstructed from ideal internal coordinates
geometry by the existing pose_stack_from_canonical_form pipeline.

Public API (exported from tmol):

Symbol	Description
pose_stack_from_backbone_coords(coords, res_types, chain_id, device)	Main entry point
canonical_form_from_backbone_coords(...)	Stable serializable intermediate
canonical_ordering_for_backbone_coords()	Memoized canonical ordering
packed_block_types_for_backbone_coords(device)	Memoized packed block types

Conventions:

• Atom order: N=0, CA=1, C=2, O=3 — matches AlphaFold2 / mosaic
• Default aa_order="ARNDCQEGHILKMFPSTWYV" (AlphaFold2 / mosaic standard; overridable)
• res_types=-1 marks padding positions
• Single-pose (L, 4, 3) inputs are auto-unsqueezed; batched (n_poses, L, 4, 3) also accepted

Cherry-pick: PoseStack.split and PDBInfo.split (from PR #379)

Adds PoseStack.split(index) and PDBInfo.split(index) from
aleaverfay/split_pose_stack (#379).
These are included here because they compose naturally with the new adapter:
build complex, binder, and target as a 3-pose batch with
PoseStackBuilder.from_poses, score once, and index per-pose energies —
rather than three separate pose_stack_from_pdb + score calls.

Usage

import torch, tmol

# coords: (L, 4, 3) — N/CA/C/O per residue
# res_types: (L,)   — indices into "ARNDCQEGHILKMFPSTWYV"
# chain_id: (L,)    — integer chain identifiers (int32)

device = torch.device("cuda")
ps = tmol.pose_stack_from_backbone_coords(
    torch.from_numpy(coords).float(),
    torch.from_numpy(res_types).long(),
    torch.from_numpy(chain_id).int(),
    device,
)
sfxn = tmol.beta2016_score_function(device)
energy = sfxn.render_whole_pose_scoring_module(ps)(ps.coords).sum()

Tests

tmol/tests/io/test_pose_stack_from_backbone_coords.py

Test	What it checks
test_pose_stack_from_backbone_coords_shape	n_poses and max_n_blocks match residue count (ubiquitin, 76 res)
test_pose_stack_from_backbone_coords_device	returned PoseStack.coords is on the requested device
test_pose_stack_from_backbone_coords_single_pose_squeeze	3-D input (L, 4, 3) is auto-unsqueezed to a 1-pose batch
test_backbone_coords_side_chains_built	total atom count > 4×n_res, confirming ideal side chains were added
test_pose_stack_from_backbone_coords_two_chains	splitting a single chain into two produces correct max_n_blocks
test_memoization_of_backbone_coords_paramdb	_paramdb_for_backbone_coords() returns the same object on repeated calls
test_memoization_of_canonical_ordering	canonical_ordering_for_backbone_coords() is memoized
test_memoization_of_packed_block_types	packed_block_types_for_backbone_coords(device) is memoized
test_device_of_packed_block_types	PackedBlockTypes is on the correct device
test_padding_residues_ignored	res_types=-1 padding positions are excluded; 3 valid + 2 padding → max_n_blocks=3