diff --git a/docs/api/io.rst b/docs/api/io.rst index 0a13a0f5d..639fb0c23 100644 --- a/docs/api/io.rst +++ b/docs/api/io.rst @@ -8,7 +8,10 @@ coordinates with general Biotite topology, including nucleic acids and ligands, use :func:`tmol.io.pose_stack_from_atom37_and_biotite`. Repeated diffusion, guidance, and search workloads should bind their fixed topology once with :func:`tmol.io.prepare_pose_stack_from_atom37`; its returned callable accepts -each coordinate batch and optimizes hydrogens by default. +each coordinate batch and optimizes hydrogens by default. For backbone-only +input -- N/CA/C/O with no side chains, as produced by backbone generators -- +use :func:`tmol.io.pose_stack_from_backbone_coords`, which completes the +missing side chains with the packer. .. automodule:: tmol.io :members: diff --git a/tmol/__init__.py b/tmol/__init__.py index 65e556637..c0177cc67 100644 --- a/tmol/__init__.py +++ b/tmol/__init__.py @@ -51,6 +51,7 @@ def include_paths(): "atom_records_from_pose_stack", "beta2016_score_function", "build_kinforest_network", + "canonical_form_from_backbone_coords", "canonical_form_from_openfold", "canonical_form_from_pdb", "canonical_form_from_rosettafold2", @@ -67,6 +68,7 @@ def include_paths(): "one2three", "packed_block_types_for_openfold", "packed_block_types_for_rosettafold2", + "pose_stack_from_backbone_coords", "pose_stack_from_canonical_form", "pose_stack_from_openfold", "pose_stack_from_pdb", @@ -93,6 +95,7 @@ def include_paths(): ( "CanonicalOrdering", "atom_records_from_pose_stack", + "canonical_form_from_backbone_coords", "canonical_form_from_openfold", "canonical_form_from_pdb", "canonical_form_from_rosettafold2", @@ -103,6 +106,7 @@ def include_paths(): "extended_pose_stack_from_sequences", "packed_block_types_for_openfold", "packed_block_types_for_rosettafold2", + "pose_stack_from_backbone_coords", "pose_stack_from_canonical_form", "pose_stack_from_openfold", "pose_stack_from_pdb", diff --git a/tmol/io/__init__.py b/tmol/io/__init__.py index 4d2d73248..609ce8878 100644 --- a/tmol/io/__init__.py +++ b/tmol/io/__init__.py @@ -93,6 +93,10 @@ packed_block_types_for_rosettafold2, _paramdb_for_rosettafold2, ) +from ._pose_stack_from_backbone_coords import ( # noqa: F401 + pose_stack_from_backbone_coords, + canonical_form_from_backbone_coords, +) from ._write_pose_stack_pdb import ( # noqa: F401 write_pose_stack_pdb, atom_records_from_pose_stack, @@ -112,6 +116,7 @@ "biotite_from_pose_stack", "build_context_from_biotite", "canonical_form_from_atomworks", + "canonical_form_from_backbone_coords", "canonical_form_from_biotite", "canonical_form_from_pdb", "canonical_form_from_pose_stack", @@ -131,6 +136,7 @@ "pose_stack_from_atomworks", "prepare_pose_stack_from_atom37", "pose_stack_from_atom37_and_biotite", + "pose_stack_from_backbone_coords", "pose_stack_from_biotite", "pose_stack_from_openfold", "pose_stack_from_pdb", diff --git a/tmol/io/_pose_stack_from_backbone_coords.py b/tmol/io/_pose_stack_from_backbone_coords.py new file mode 100644 index 000000000..e2b214df5 --- /dev/null +++ b/tmol/io/_pose_stack_from_backbone_coords.py @@ -0,0 +1,276 @@ +"""Build a PoseStack from backbone-only N/CA/C/O coordinates. + +tmol builds missing leaf atoms but rejects blocks missing non-leaf atoms, so +the other adapters cannot take a bare backbone. This one completes the absent +side chains with the packer. Chemistry is shared with the atomworks adapter. +""" + +from typing import Literal + +import toolz +import torch + +from tmol.chemical import one2three +from tmol.io import CanonicalForm +from tmol.io._build_context import PoseBuildContext +from tmol.io._pose_stack_from_atomworks import ( + ATOMWORKS_NAME3S, + _get_aw_2_tmol_mappings, + _paramdb_for_atomworks, + _restype_set_for_atomworks, + canonical_ordering_for_atomworks, + packed_block_types_for_atomworks, +) +from tmol.pose import PoseStack + +# AlphaFold2/mosaic ordering; alphabetical by three-letter code and so equal to +# ATOMWORKS_NAME3S[1:21]. +_DEFAULT_AA_ORDER = "ARNDCQEGHILKMFPSTWYV" + +# N, CA, C, O in the atom37 layout. The same four slots for all 20 amino acids, +# which is what lets a bare (L, 4, 3) tensor reuse the atomworks tables. +_BACKBONE_ATOM37_SLOTS = (0, 1, 2, 4) + + +def pose_stack_from_backbone_coords( + coords: torch.Tensor, + res_types: torch.Tensor, + chain_id: torch.Tensor, + device: torch.device, + *, + aa_order: str = _DEFAULT_AA_ORDER, + sidechain_completion: Literal["pack", "none"] = "pack", + no_optH: bool = False, + **kwargs, +) -> PoseStack: + """Build a PoseStack from backbone N/CA/C/O coordinates, with no file I/O. + + The input has no side-chain heavy atoms, so ``sidechain_completion`` says + how to supply them. ``"pack"`` runs build_missing_sidechains -- a full + score-function-driven Dunbrack/OptH job, far more expensive than the + conversion, whose output is not a differentiable function of the input. + ``"none"`` converts only, leaving absent side chains NaN and the pose + unscorable. Either way the supplied backbone is returned unchanged and + stays on the autograd tape. + + Args: + coords: ``(max_n_res, 4, 3)`` or ``(n_poses, max_n_res, 4, 3)`` in + N, CA, C, O order. Non-finite entries mark absent atoms. Cast to + float32, as CanonicalForm requires. + res_types: indices into ``aa_order``; ``-1`` marks padding. + chain_id: chain identifiers, shaped like ``res_types``. Residues in a + chain must be consecutive. + device: device for the returned PoseStack. + aa_order: one-letter codes defining the ``res_types`` mapping. + no_optH: leave rebuilt hydrogens at ideal positions when packing. + kwargs: passed through to pose_stack_from_canonical_form. + + Example:: + + ps = tmol.pose_stack_from_backbone_coords( + coords, res_types, chain_id, torch.device("cuda") + ) + sfxn = tmol.beta2016_score_function(ps.packed_block_types.device) + energy = sfxn.render_whole_pose_scoring_module(ps)(ps.coords).sum() + """ + from tmol.io import ( + canonical_form_from_pose_stack, + pose_stack_from_canonical_form, + ) + + if sidechain_completion not in ("pack", "none"): + raise ValueError( + f"sidechain_completion must be 'pack' or 'none'; " + f"got {sidechain_completion!r}" + ) + + if coords.dim() == 3: + coords = coords.unsqueeze(0) + res_types = res_types.unsqueeze(0) + chain_id = chain_id.unsqueeze(0) + + cf = canonical_form_from_backbone_coords( + coords.to(device), res_types.to(device), chain_id.to(device), aa_order + ) + context = _build_context_for_backbone_coords(device) + co = context.canonical_ordering + pbt = context.packed_block_types + + packing = sidechain_completion == "pack" + wants_missing_mask = bool(kwargs.pop("return_block_has_missing_atoms", False)) + wants_atom_mapping = bool(kwargs.pop("return_atom_mapping", False)) + + pose_stack, opt_return_vals = pose_stack_from_canonical_form( + co, + pbt, + *cf, + return_block_has_missing_atoms=True, + return_atom_mapping=wants_atom_mapping or packing, + **kwargs, + ) + block_has_missing_atoms = opt_return_vals.pop("block_has_missing_atoms") + has_missing_atoms = block_has_missing_atoms is not None and bool( + torch.any(block_has_missing_atoms) + ) + + if packing and block_has_missing_atoms is not None: + from tmol.io._pose_stack_from_biotite import _restore_canonical_input_coords + from tmol.pack import build_missing_sidechains + + if has_missing_atoms or not no_optH: + pose_stack = build_missing_sidechains( + pose_stack, + ( + context._packing_score_function + if has_missing_atoms + else context._opth_score_function + ), + context._dunbrack_sampler, + block_has_missing_atoms, + no_optH=no_optH, + has_missing_atoms=has_missing_atoms, + ) + + if has_missing_atoms: + # HA is built from CA/N/CB, and leaf building runs before + # packing, so a bare backbone leaves it unplaced. Rebuild now + # that the heavy atoms exist. No missing-atom flag here: a gap + # that survives packing should raise, not come back NaN. + pose_stack, opt_return_vals = pose_stack_from_canonical_form( + co, + pbt, + *canonical_form_from_pose_stack(co, pose_stack), + return_atom_mapping=True, + **kwargs, + ) + + # Packing works on values, not the autograd graph; route the input + # atoms back to the canonical tensor. + pose_stack = _restore_canonical_input_coords( + pose_stack, + cf.coords, + opt_return_vals["can_atom_mapping"], + opt_return_vals["ps_atom_mapping"], + ) + + if wants_missing_mask: + opt_return_vals["block_has_missing_atoms"] = block_has_missing_atoms + if not wants_atom_mapping: + opt_return_vals.pop("can_atom_mapping", None) + opt_return_vals.pop("ps_atom_mapping", None) + + if opt_return_vals: + return pose_stack, opt_return_vals + return pose_stack + + +def canonical_form_from_backbone_coords( + coords: torch.Tensor, + res_types: torch.Tensor, + chain_id: torch.Tensor, + aa_order: str = _DEFAULT_AA_ORDER, +) -> CanonicalForm: + """Build a CanonicalForm from backbone N/CA/C/O coordinates. + + Every non-backbone atom is NaN, so the result needs + ``return_block_has_missing_atoms=True`` to reach a PoseStack. Its residue + type indices refer to canonical_ordering_for_atomworks, which is what to + rebuild with after a round-trip through disk. + """ + assert coords.dim() == 4, "coords must be 4D (n_poses, max_n_res, 4, 3)" + assert coords.shape[2:] == (4, 3), "atom dimension must be 4: N, CA, C, O" + assert res_types.shape == coords.shape[:2], "res_types must be (n_poses, max_n_res)" + assert chain_id.shape == coords.shape[:2], "chain_id must be (n_poses, max_n_res)" + assert coords.device == res_types.device + assert coords.device == chain_id.device + + device = coords.device + n_poses, max_n_res = coords.shape[:2] + + co = canonical_ordering_for_atomworks() + aw2t_rtmap, aw2t_atmap, aw_at_is_real = _get_aw_2_tmol_mappings(device) + + # Slice the atom tables down to the backbone up front, so nothing of size + # (n_poses, max_n_res, 37) is materialized. + slots = torch.tensor(_BACKBONE_ATOM37_SLOTS, dtype=torch.int64, device=device) + bb_atmap = aw2t_atmap[:, slots] + bb_is_real = aw_at_is_real[:, slots] + + res_types = res_types.to(torch.int64) + padding = res_types < 0 + tokens = _atomworks_tokens_for_aa_order(aa_order, device) + if bool(torch.any(res_types[~padding] >= len(tokens))): + bad = res_types[~padding & (res_types >= len(tokens))].unique() + raise ValueError( + f"res_types must be in range [0, {len(tokens) - 1}] or -1 for " + f"padding; got out-of-range values: {bad.tolist()}" + ) + # Padding becomes the atomworks "" token, which the mappings already + # send to restype -1 with no real atoms. + aw_tokens = torch.where(padding, 0, tokens[res_types.clamp(min=0)]) + + tmol_restypes = aw2t_rtmap[aw_tokens] + atom_mapping = bb_atmap[aw_tokens] + at_is_real = bb_is_real[aw_tokens] + + n_bb_ats = len(_BACKBONE_ATOM37_SLOTS) + pose_ind = ( + torch.arange(n_poses, dtype=torch.int64, device=device) + .reshape(-1, 1, 1) + .expand(n_poses, max_n_res, n_bb_ats) + ) + res_ind = ( + torch.arange(max_n_res, dtype=torch.int64, device=device) + .reshape(1, -1, 1) + .expand(n_poses, max_n_res, n_bb_ats) + ) + + tmol_coords = torch.full( + (n_poses, max_n_res, co.max_n_canonical_atoms, 3), + float("nan"), + dtype=torch.float32, + device=device, + ) + tmol_coords[ + pose_ind[at_is_real], + res_ind[at_is_real], + atom_mapping[at_is_real], + ] = coords.to(torch.float32)[at_is_real] + + return CanonicalForm( + chain_id=chain_id.to(torch.int32), + res_types=tmol_restypes.to(torch.int32), + coords=tmol_coords, + res_labels=None, + residue_insertion_codes=None, + chain_labels=None, + atom_occupancy=None, + atom_b_factor=None, + disulfides=None, + res_not_connected=None, + ) + + +@toolz.functoolz.memoize +def _atomworks_tokens_for_aa_order(aa_order: str, device: torch.device): + """Map positions in aa_order onto atomworks protein token indices.""" + if len(set(aa_order)) != len(aa_order): + raise ValueError(f"aa_order must not repeat a one-letter code: {aa_order!r}") + tokens = [] + for aa in aa_order: + name3 = one2three(aa) + if name3 not in ATOMWORKS_NAME3S: + raise ValueError(f"aa_order entry {aa!r} ({name3}) is not a canonical AA") + tokens.append(ATOMWORKS_NAME3S.index(name3)) + return torch.tensor(tokens, dtype=torch.int64, device=device) + + +@toolz.functoolz.memoize +def _build_context_for_backbone_coords(device: torch.device) -> PoseBuildContext: + """Build context for the canonical amino acids, shared across calls.""" + return PoseBuildContext( + canonical_ordering=canonical_ordering_for_atomworks(), + packed_block_types=packed_block_types_for_atomworks(device), + parameter_database=_paramdb_for_atomworks(), + restype_set=_restype_set_for_atomworks(), + ) diff --git a/tmol/tests/io/test_pose_stack_from_backbone_coords.py b/tmol/tests/io/test_pose_stack_from_backbone_coords.py new file mode 100644 index 000000000..bda61e53a --- /dev/null +++ b/tmol/tests/io/test_pose_stack_from_backbone_coords.py @@ -0,0 +1,327 @@ +"""Tests for backbone-only (N/CA/C/O) side-chain completion.""" + +import pytest +import torch + +from tmol.chemical import three2one +from tmol.io import ( + canonical_form_from_backbone_coords, + canonical_form_from_pdb, + canonical_form_from_pose_stack, + canonical_ordering_for_atomworks, + packed_block_types_for_atomworks, + pose_stack_from_backbone_coords, +) +from tmol.io._pose_stack_from_atomworks import _paramdb_for_atomworks +from tmol.tests._torch import requires_cuda +from tmol.io._pose_stack_from_backbone_coords import ( + _BACKBONE_ATOM37_SLOTS, + _DEFAULT_AA_ORDER, + _atomworks_tokens_for_aa_order, + _build_context_for_backbone_coords, +) + +_BACKBONE_ATOM_NAMES = ("N", "CA", "C", "O") +_CANONICALIZE_NAME3 = {"HIS_D": "HIS", "HIS_POS": "HIS", "CYD": "CYS"} + + +def _backbone_from_pdb(pdb, device): + """Extract (coords, res_types, chain_id) for a PDB via the canonical form.""" + co = canonical_ordering_for_atomworks() + cf = canonical_form_from_pdb(co, pdb, device) + + n_poses, max_n_res = cf.res_types.shape[:2] + coords = torch.full( + (n_poses, max_n_res, 4, 3), float("nan"), dtype=torch.float32, device=device + ) + res_types = torch.full((n_poses, max_n_res), -1, dtype=torch.int64, device=device) + + for p in range(n_poses): + for r in range(max_n_res): + rt_ind = int(cf.res_types[p, r]) + if rt_ind < 0: + continue + name3 = co.restype_io_equiv_classes[rt_ind] + at_map = co.restypes_atom_index_mapping[name3] + for slot, at_name in enumerate(_BACKBONE_ATOM_NAMES): + if at_name in at_map: + coords[p, r, slot] = cf.coords[p, r, at_map[at_name]] + one = three2one(_CANONICALIZE_NAME3.get(name3, name3)) + res_types[p, r] = _DEFAULT_AA_ORDER.index(one) + + return coords, res_types, cf.chain_id.to(torch.int64) + + +def test_build_context_reuses_atomworks_chemistry(torch_device): + context = _build_context_for_backbone_coords(torch_device) + + assert context.canonical_ordering is canonical_ordering_for_atomworks() + assert context.packed_block_types is packed_block_types_for_atomworks(torch_device) + assert context.parameter_database is _paramdb_for_atomworks() + + +def test_default_aa_order_matches_atomworks_protein_tokens(): + """AF2 one-letter order is alphabetical by name3, so it is a +1 shift.""" + tokens = _atomworks_tokens_for_aa_order(_DEFAULT_AA_ORDER, torch.device("cpu")) + assert torch.equal(tokens, torch.arange(1, 21, dtype=torch.int64)) + + +def test_backbone_slots_are_uniform_across_restypes(): + """The (L, 4, 3) reuse of the atomworks tables depends on this.""" + from tmol.io._pose_stack_from_atomworks import ( + ATOMWORKS_ATOM37_NAMES, + ATOMWORKS_NAME3S, + ) + + for name3 in ATOMWORKS_NAME3S[1:21]: + row = ATOMWORKS_ATOM37_NAMES[name3] + slots = tuple(row.index(a) for a in _BACKBONE_ATOM_NAMES) + assert slots == _BACKBONE_ATOM37_SLOTS + + +def test_aa_order_rejects_non_canonical_and_repeats(): + with pytest.raises(ValueError, match="not a canonical AA"): + _atomworks_tokens_for_aa_order("ARNDCQEGHILKMFPSTWYX", torch.device("cpu")) + with pytest.raises(ValueError, match="must not repeat"): + _atomworks_tokens_for_aa_order("AARNDCQEGHILKMFPSTWY", torch.device("cpu")) + + +def test_packing_preserves_the_supplied_backbone(ubq_pdb, torch_device): + coords, res_types, chain_id = _backbone_from_pdb(ubq_pdb, torch_device) + + torch.manual_seed(0) + pose_stack = pose_stack_from_backbone_coords( + coords, res_types, chain_id, torch_device + ) + + co = canonical_ordering_for_atomworks() + round_tripped = canonical_form_from_pose_stack(co, pose_stack) + + n_res = int((res_types[0] >= 0).sum()) + for r in range(n_res): + name3 = co.restype_io_equiv_classes[int(round_tripped.res_types[0, r])] + at_map = co.restypes_atom_index_mapping[name3] + for slot, at_name in enumerate(_BACKBONE_ATOM_NAMES): + if at_name not in at_map: + continue + torch.testing.assert_close( + round_tripped.coords[0, r, at_map[at_name]], + coords[0, r, slot], + rtol=0.0, + atol=0.0, + ) + + +def test_sidechains_are_built_and_no_nans_remain(ubq_pdb, torch_device): + coords, res_types, chain_id = _backbone_from_pdb(ubq_pdb, torch_device) + n_res = int((res_types[0] >= 0).sum()) + + torch.manual_seed(0) + packed = pose_stack_from_backbone_coords(coords, res_types, chain_id, torch_device) + assert not torch.any(torch.isnan(packed.coords[packed.real_atoms])) + assert int(packed.real_atoms.sum()) > 4 * n_res + + +def test_gradients_reach_the_input_backbone_through_packing(ubq_pdb, torch_device): + coords, res_types, chain_id = _backbone_from_pdb(ubq_pdb, torch_device) + coords = coords[:, :8].clone() + res_types = res_types[:, :8].clone() + chain_id = chain_id[:, :8].clone() + coords.requires_grad_(True) + + torch.manual_seed(0) + pose_stack = pose_stack_from_backbone_coords( + coords, res_types, chain_id, torch_device + ) + + torch.nan_to_num(pose_stack.coords).sum().backward() + assert coords.grad is not None + finite = torch.isfinite(coords.detach()).all(dim=-1) + assert torch.count_nonzero(coords.grad[finite]) > 0 + + +def test_completion_none_leaves_sidechains_absent(ubq_pdb, torch_device): + coords, res_types, chain_id = _backbone_from_pdb(ubq_pdb, torch_device) + + unpacked = pose_stack_from_backbone_coords( + coords, res_types, chain_id, torch_device, sidechain_completion="none" + ) + torch.manual_seed(0) + packed = pose_stack_from_backbone_coords( + coords, res_types, chain_id, torch_device, sidechain_completion="pack" + ) + + real = unpacked.real_atoms + assert torch.any(torch.isnan(unpacked.coords[real])) + assert not torch.any(torch.isnan(packed.coords[packed.real_atoms])) + + +def test_completion_none_is_differentiable(ubq_pdb, torch_device): + coords, res_types, chain_id = _backbone_from_pdb(ubq_pdb, torch_device) + coords = coords[:, :8].clone().requires_grad_(True) + + pose_stack = pose_stack_from_backbone_coords( + coords, + res_types[:, :8], + chain_id[:, :8], + torch_device, + sidechain_completion="none", + ) + torch.nan_to_num(pose_stack.coords).sum().backward() + assert coords.grad is not None + assert torch.count_nonzero(coords.grad) > 0 + + +def test_invalid_completion_policy_is_rejected(ubq_pdb, torch_device): + coords, res_types, chain_id = _backbone_from_pdb(ubq_pdb, torch_device) + with pytest.raises(ValueError, match="sidechain_completion"): + pose_stack_from_backbone_coords( + coords, res_types, chain_id, torch_device, sidechain_completion="ideal" + ) + + +def test_repeated_builds_are_stable_and_reuse_setup(ubq_pdb, torch_device): + """Packed side chains are deliberately not compared: the CPU annealer draws + from libc rand() (pack/compiled/compiled.cpu.cpp), which torch.manual_seed + does not control. That reaches topology too, since the HIS tautomer is + resolved from the packed side chain. + """ + coords, res_types, chain_id = _backbone_from_pdb(ubq_pdb, torch_device) + + first = pose_stack_from_backbone_coords(coords, res_types, chain_id, torch_device) + context = _build_context_for_backbone_coords(torch_device) + score_function = context._packing_score_function + dunbrack_sampler = context._dunbrack_sampler + + second = pose_stack_from_backbone_coords(coords, res_types, chain_id, torch_device) + + assert _build_context_for_backbone_coords(torch_device) is context + assert context._packing_score_function is score_function + assert context._dunbrack_sampler is dunbrack_sampler + + assert first.coords.shape == second.coords.shape + assert torch.equal(first.block_coord_offset, second.block_coord_offset) + assert torch.equal(first.real_atoms, second.real_atoms) + + co = canonical_ordering_for_atomworks() + cf_first = canonical_form_from_pose_stack(co, first) + cf_second = canonical_form_from_pose_stack(co, second) + n_res = int((res_types[0] >= 0).sum()) + for r in range(n_res): + for cf, other in ((cf_first, cf_second), (cf_second, cf_first)): + name3 = co.restype_io_equiv_classes[int(cf.res_types[0, r])] + other3 = co.restype_io_equiv_classes[int(other.res_types[0, r])] + at_map = co.restypes_atom_index_mapping[name3] + other_map = co.restypes_atom_index_mapping[other3] + for at_name in _BACKBONE_ATOM_NAMES: + if at_name not in at_map or at_name not in other_map: + continue + assert torch.equal( + cf.coords[0, r, at_map[at_name]], + other.coords[0, r, other_map[at_name]], + ) + + +@requires_cuda +def test_cpu_and_cuda_agree(ubq_pdb): + cpu = torch.device("cpu") + cuda = torch.device("cuda") + coords, res_types, chain_id = _backbone_from_pdb(ubq_pdb, cpu) + + cf_cpu = canonical_form_from_backbone_coords(coords, res_types, chain_id) + cf_cuda = canonical_form_from_backbone_coords( + coords.to(cuda), res_types.to(cuda), chain_id.to(cuda) + ) + assert torch.equal(torch.isnan(cf_cpu.coords), torch.isnan(cf_cuda.coords).cpu()) + finite = torch.isfinite(cf_cpu.coords) + torch.testing.assert_close( + cf_cpu.coords[finite], cf_cuda.coords.cpu()[finite], rtol=0.0, atol=0.0 + ) + assert torch.equal(cf_cpu.res_types, cf_cuda.res_types.cpu()) + + ps_cpu = pose_stack_from_backbone_coords( + coords, res_types, chain_id, cpu, sidechain_completion="none" + ) + ps_cuda = pose_stack_from_backbone_coords( + coords, res_types, chain_id, cuda, sidechain_completion="none" + ) + assert ps_cpu.coords.shape == ps_cuda.coords.shape + mask = ps_cpu.real_atoms & torch.isfinite(ps_cpu.coords).all(dim=-1) + torch.testing.assert_close( + ps_cpu.coords[mask], ps_cuda.coords.cpu()[mask], rtol=1e-5, atol=1e-4 + ) + + +def test_single_pose_input_is_unsqueezed(ubq_pdb, torch_device): + coords, res_types, chain_id = _backbone_from_pdb(ubq_pdb, torch_device) + + batched = pose_stack_from_backbone_coords( + coords, res_types, chain_id, torch_device, sidechain_completion="none" + ) + single = pose_stack_from_backbone_coords( + coords[0], + res_types[0], + chain_id[0], + torch_device, + sidechain_completion="none", + ) + + assert single.n_poses == batched.n_poses == 1 + assert torch.equal(single.block_type_ind, batched.block_type_ind) + finite = torch.isfinite(single.coords) & torch.isfinite(batched.coords) + assert torch.equal(single.coords[finite], batched.coords[finite]) + + +def test_output_is_on_the_requested_device(ubq_pdb, torch_device): + coords, res_types, chain_id = _backbone_from_pdb(ubq_pdb, torch.device("cpu")) + pose_stack = pose_stack_from_backbone_coords( + coords, res_types, chain_id, torch_device, sidechain_completion="none" + ) + assert pose_stack.coords.device.type == torch_device.type + + +def test_padding_positions_are_excluded(ubq_pdb, torch_device): + coords, res_types, chain_id = _backbone_from_pdb(ubq_pdb, torch_device) + n_keep = 5 + res_types = res_types.clone() + res_types[0, n_keep:] = -1 + + pose_stack = pose_stack_from_backbone_coords( + coords, res_types, chain_id, torch_device, sidechain_completion="none" + ) + assert int((pose_stack.block_type_ind[0] >= 0).sum()) == n_keep + + +def test_two_chains_are_separated(ubq_pdb, torch_device): + coords, res_types, chain_id = _backbone_from_pdb(ubq_pdb, torch_device) + n_res = int((res_types[0] >= 0).sum()) + chain_id = chain_id.clone() + chain_id[0, n_res // 2 : n_res] = 1 + + pose_stack = pose_stack_from_backbone_coords( + coords, res_types, chain_id, torch_device, sidechain_completion="none" + ) + assert int((pose_stack.block_type_ind[0] >= 0).sum()) == n_res + assert pose_stack.n_poses == 1 + assert int(pose_stack.chain_id[0, :n_res].max()) == 1 + + +def test_shape_validation(torch_device): + rt = torch.zeros((1, 4), dtype=torch.int64, device=torch_device) + ci = torch.zeros((1, 4), dtype=torch.int64, device=torch_device) + + bad = torch.zeros((1, 4, 3, 3), dtype=torch.float32, device=torch_device) + with pytest.raises(AssertionError, match="N, CA, C, O"): + canonical_form_from_backbone_coords(bad, rt, ci) + + good = torch.zeros((1, 4, 4, 3), dtype=torch.float32, device=torch_device) + with pytest.raises(AssertionError, match="res_types must be"): + canonical_form_from_backbone_coords(good, rt[:, :3], ci) + + +def test_out_of_range_res_types_are_rejected(torch_device): + coords = torch.zeros((1, 2, 4, 3), dtype=torch.float32, device=torch_device) + rt = torch.tensor([[0, 20]], dtype=torch.int64, device=torch_device) + ci = torch.zeros((1, 2), dtype=torch.int64, device=torch_device) + with pytest.raises(ValueError, match="must be in range"): + canonical_form_from_backbone_coords(coords, rt, ci)