Pose

PoseStack stores a padded batch of molecular coordinates and block topology. PoseStackBuilder combines compatible poses, while PackedBlockTypes holds the shared block-type vocabulary. Constraint data is stored in ConstraintSet.

Pose stacks, packed residue types, and structural metadata.

class tmol.pose.ConstraintSet(device: device, n_poses: int, constraint_function_inds: Tensor, constraint_atoms: Tensor[slice(None, None, None), 4, 3], constraint_params: Tensor[slice(None, None, None), slice(None, None, None)], constraint_num_unique_blocks: Tensor, constraint_unique_blocks: Tensor[slice(None, None, None), 3], constraint_functions: tuple[Callable[[Tensor[slice(None, None, None), slice(None, None, None), 3], Tensor[slice(None, None, None), slice(None, None, None)]], Tensor], ...])

Bases: object

Immutable batched coordinate constraints and their tensor storage.

classmethod create_empty(device: device, n_poses: int) → ConstraintSet

Create an empty constraint set on a concrete device.

classmethod concatenate(constraint_sets: Sequence[ConstraintSet | None], from_multiple_pose_stacks: bool = True, n_poses: int | None = None, ps_offset: Tensor | None = None) → ConstraintSet | None

Combine constraint sets while deduplicating their scoring functions.

Parameters:

• constraint_sets – Constraint sets to combine; None entries are retained when calculating pose offsets.

• from_multiple_pose_stacks – Shift pose indices between inputs when true. When false, all inputs describe the same pose batch.

• n_poses – Pose count for the result, inferred when omitted.

• ps_offset – Per-input pose offsets, inferred when omitted.

Returns:

The combined constraint set, or None when every input is None.

clone() → ConstraintSet

Return a copy with independent tensor storage.

to(device: device) → ConstraintSet

Return a copy whose tensors and metadata use device.

split(index: int) → ConstraintSet

Split out a single pose’s worth of constraints from a batch.

static count_unique_blocks(atom_indices: Tensor[slice(None, None, None), slice(None, None, None), 3]) → Tensor

Count distinct consecutive blocks referenced by each constraint.

add_constraints_to_all_poses(fn: Callable[[Tensor[slice(None, None, None), slice(None, None, None), 3], Tensor[slice(None, None, None), slice(None, None, None)]], Tensor], atom_indices: Tensor[slice(None, None, None), slice(None, None, None), slice(None, None, None)], params: Tensor[slice(None, None, None), slice(None, None, None)] | None = None) → ConstraintSet

Add the same constraints to every pose in the batch.

Parameters:

• fn – Function that scores coordinates and per-constraint parameters.

• atom_indices – Integer [constraint, atom, block/atom] indices, or full [constraint, atom, pose/block/atom] indices.

• params – Optional float [constraint, parameter] values.

Returns:

A new constraint set containing the replicated constraints.

add_constraints(fn: Callable[[Tensor[slice(None, None, None), slice(None, None, None), 3], Tensor[slice(None, None, None), slice(None, None, None)]], Tensor], atom_indices: Tensor[slice(None, None, None), slice(None, None, None), slice(None, None, None)], params: Tensor[slice(None, None, None), slice(None, None, None)] | None = None) → ConstraintSet

Return a constraint set containing the existing and new constraints.

Parameters:

• fn – Function that scores coordinates and per-constraint parameters.

• atom_indices – Integer [constraint, atom, pose/block/atom] indices. Omitting the pose column applies each constraint to every pose.

• params – Optional float [constraint, parameter] values.

Returns:

A new constraint set containing the added constraints.

static replicate_constraints(n_poses: int, c_atms: Tensor[slice(None, None, None), slice(None, None, None), 3], c_params: Tensor[slice(None, None, None), slice(None, None, None)] | None) → tuple[Tensor[slice(None, None, None), slice(None, None, None), 3], Tensor[slice(None, None, None), slice(None, None, None)] | None]

Replicate pose-independent constraints and optional parameters.

class tmol.pose.InterBlockBondsep(near_blocks: Tensor[slice(None, None, None), slice(None, None, None), slice(None, None, None), 2], bondsep: Tensor[slice(None, None, None), slice(None, None, None), slice(None, None, None), slice(None, None, None), slice(None, None, None)])

Bases: object

Connection bond separations below the cap of nearby block pairs; each row lists its block2 slots ascending, then empty slots (block -1, the cap).

property shape: tuple[int, int, int, int, int]

Shape of the equivalent dense table.

classmethod empty(n_poses: int, max_n_blocks: int, max_n_conn: int, device: device, n_slots: int) → InterBlockBondsep

Separations for blocks that are all at least the cap apart.

classmethod from_bonded_graph(counts: Tensor[slice(None, None, None), slice(None, None, None)], intra_separation: Tensor[slice(None, None, None), slice(None, None, None), slice(None, None, None), slice(None, None, None)], connections: Tensor[slice(None, None, None), slice(None, None, None), slice(None, None, None), 2]) → InterBlockBondsep

Shortest paths below the cap between [pose, block, conn] connections, one bond per connections partner (-1: none) plus intra_separation in a block.

classmethod concatenate(parts: Sequence[InterBlockBondsep], max_n_blocks: int, max_n_conn: int) → InterBlockBondsep

Stack the poses of several tables, padding blocks and connections.

to_dense() → Tensor[slice(None, None, None), slice(None, None, None), slice(None, None, None), slice(None, None, None), slice(None, None, None)]

Materialize the [pose, block1, block2, conn1, conn2] table.

select_poses(start: int, stop: int) → InterBlockBondsep

View of poses start to stop.

class tmol.pose.PackedBlockTypes(chem_db: PatchedChemicalDatabase, restype_set: ResidueTypeSet, active_block_types: Sequence[RefinedResidueType], restype_index: Index, max_n_atoms: int, n_atoms: Tensor, atom_is_real: Tensor[slice(None, None, None), slice(None, None, None)], atom_is_hydrogen: Tensor[slice(None, None, None), slice(None, None, None)], atom_downstream_of_conn: Tensor[slice(None, None, None), slice(None, None, None), slice(None, None, None)], atom_paths_from_conn: Tensor[slice(None, None, None), slice(None, None, None), 13, 3], max_n_torsions: int, n_torsions: Tensor, torsion_is_real: Tensor[slice(None, None, None), slice(None, None, None)], torsion_uaids: Tensor[slice(None, None, None), slice(None, None, None), 3], is_torsion_mc: Tensor[slice(None, None, None), slice(None, None, None)], n_mc_torsions: Tensor, mc_torsion_is_real: Tensor[slice(None, None, None), slice(None, None, None)], mc_torsions: Tensor[slice(None, None, None), slice(None, None, None)], n_sc_torsions: Tensor, sc_torsion_is_real: Tensor[slice(None, None, None), slice(None, None, None)], sc_torsions: Tensor[slice(None, None, None), slice(None, None, None)], which_mcsc_torsions: Tensor[slice(None, None, None), slice(None, None, None)], max_n_bonds: int, n_bonds: Tensor, bond_is_real: Tensor[slice(None, None, None), slice(None, None, None)], bond_indices: Tensor[slice(None, None, None), slice(None, None, None), 2], max_n_conn: int, n_conn: Tensor, conn_is_real: Tensor[slice(None, None, None), slice(None, None, None)], conn_atom: Tensor[slice(None, None, None), slice(None, None, None)], down_conn_inds: Tensor, up_conn_inds: Tensor, polymeric_conn_inds: Tensor[slice(None, None, None), 2], default_jump_connection_atom_inds: Tensor, device: device)

Bases: object

A class to aggregate the properties for a collection of residue types.

The PackedBlockTypes object holds an ordered set of residue types (specifically, RefinedResidueTypes); once constructed, this order will not change, so residue types may be referred to by index within this object.

Score terms cache derived tensors on this object to share annotation work across poses with the same ordered residue types. Topology annotations can be reused for the object’s lifetime. Parameter-dependent annotations must also identify their chemical/scoring sources, settings and tensor device; the presence of an attribute alone does not establish a valid cache hit. Keep historical configurations bounded, and let rendered scoring modules capture their requested annotation tensors instead of reading whichever configuration happens to be stored here during a later forward call.

Annotation process: There are three steps to the annotation process. 1) Terms annotate individual block types, 2) terms aggregate (concattenate) annotations of the individual block types for the packed_block_type, and 3) terms retrieve their annotations. 1) Typically, score terms will create one annotation for each of the RefinedResidueType objects that the PackedBlockTypes object holds in their method named “setup_block_type,” and cache these annotations on each RefinedResiduetype. The idiom we use is for these residue-type annotations to be held in numpy arrays (on the CPU). 2) The terms will then aggregate the annotations for each of the individual residue types into a single torch Tensor, one tensor for each property (or property group) the term needs, in their method named “setup_packed_block_types”; the idiom is for these annotations to be moved to the PackedBlockType’s device in this step. 3) Finally, each term will retrieve the cached annotations from the PackedBlockType in their “render_whole_pose_scoring_module” method.

classmethod from_restype_list(chem_db: PatchedChemicalDatabase, restype_set: ResidueTypeSet, active_block_types: Sequence[RefinedResidueType], device: device) → PackedBlockTypes

Pack residue-type metadata onto a concrete torch device.

Parameters:

• chem_db – Chemical database shared by the residue types.

• restype_set – Source residue-type collection.

• active_block_types – Ordered residue types to pack.

• device – Target device. An unindexed CUDA device resolves to the current CUDA device.

Returns:

Packed residue-type tensors and their concrete device metadata.

Raises:

ValueError – If a residue type has more than MAX_N_CONN connections.

tmol.pose.annotate_packed_block_types_w_dslf_conn_inds(pbt: PackedBlockTypes)

Annotate each block type with its disulfide connection index, or -1.

tmol.pose.annotate_packed_block_types_w_kinematic_conns(pbt: PackedBlockTypes)

Mark, per block type, which connections carry motion between blocks.

class tmol.pose.PoseStack(packed_block_types: PackedBlockTypes, coords: Tensor[slice(None, None, None), slice(None, None, None), 3], block_coord_offset: Tensor[slice(None, None, None), slice(None, None, None)], block_coord_offset64: Tensor[slice(None, None, None), slice(None, None, None)], inter_residue_connections: Tensor[slice(None, None, None), slice(None, None, None), slice(None, None, None), 2], inter_residue_connections64: Tensor[slice(None, None, None), slice(None, None, None), slice(None, None, None), 2], inter_block_bondsep: InterBlockBondsep, block_type_ind: Tensor[slice(None, None, None), slice(None, None, None)], block_type_ind64: Tensor[slice(None, None, None), slice(None, None, None)], chain_id: Tensor[slice(None, None, None), slice(None, None, None)], chain_id64: Tensor[slice(None, None, None), slice(None, None, None)], pdb_info: PDBInfo, constraint_set: ConstraintSet | None, device: device, split_block_mapping: SplitBlockMapping | None = None)

Bases: object

Batch of molecular systems with shared residue-type definitions.

Chemistry and connectivity are fixed after construction, while coords may be updated in place during minimization. Use attr.evolve() for structural metadata changes and replace or clone coords to avoid sharing mutable coordinate storage between pose stacks.

Parameters:

• packed_block_types – Residue types and their score-term annotations.

• coords – Cartesian coordinates shaped [pose, atom, xyz].

• block_coord_offset – Per-residue atom offsets shaped [pose, residue].

• block_coord_offset64 – 64-bit copy of block_coord_offset for PyTorch.

• inter_residue_connections – Connected residue and connection indices.

• inter_residue_connections64 – 64-bit connection-index copy.

• inter_block_bondsep – Bond separation between residue connections, capped at MAX_SIG_BOND_SEPARATION and stored only for nearby residue pairs.

• block_type_ind – Packed block-type index for each residue; -1 is padding.

• block_type_ind64 – 64-bit block-type-index copy.

• chain_id – Chain index for each residue.

• chain_id64 – 64-bit chain-index copy.

• pdb_info – Source labels, occupancy, and B-factor metadata.

• constraint_set – Optional geometric constraints.

• device – Device holding all pose tensors.

• split_block_mapping – Optional mapping back to pre-split residue blocks.

property n_poses: int

Return the number of poses in the stack.

property max_n_blocks: int

Return the padded residue count per pose.

property max_n_atoms: int

Return the maximum atom count among packed residue types.

property max_n_block_atoms: int

Return the maximum atoms in any packed residue type.

property max_n_pose_atoms: int

Return the padded atom dimension of each pose.

property n_ats_per_block: Tensor[slice(None, None, None), slice(None, None, None)]

Return the number of atoms in each block

property real_atoms: Tensor[slice(None, None, None), slice(None, None, None)]

Return the mask of real, non-padding atoms in coords.

clone() → PoseStack

Deep-copy clone of this PoseStack

clone_sharing_topology() → PoseStack

Copy per-pose state but share this stack’s immutable bond topology.

clone() deep-copies every tensor, including the connection tensors, even though nothing ever mutates them in place: they are built once by PoseStackBuilder and thereafter only read.

Callers that vary coordinates while holding topology fixed – the FastRelax accept/restore loop, for example – should use this instead. Sharing is also what lets a rendered scoring module recognise the result as the same topology and reuse its cached parameters.

Returns:

A PoseStack whose coordinate and block state are independently owned, and whose connection topology aliases this stack’s.

split(index: int) → PoseStack

Copy one pose into a new single-pose stack.

expand_coords() → tuple[Tensor[slice(None, None, None), slice(None, None, None), slice(None, None, None), 3], Tensor[slice(None, None, None), slice(None, None, None), slice(None, None, None)]]

Expand packed coordinates into residue-major layout.

Returns:

Coordinates shaped [pose, residue, residue_atom, xyz] and the corresponding real-atom mask shaped [pose, residue, residue_atom].

property n_res_per_pose: Tensor

Return the number of real residues in each pose.

is_real_block(pose_ind: int, block_ind: int) → Tensor

Return a scalar boolean tensor indicating whether a block is real.

block_type(pose_ind: int, block_ind: int) → RefinedResidueType

Look up the block type for a particular pose and block and retrieve it from the PackedBlockTypes object. is_real_block must return True

get_constraint_set() → ConstraintSet | None

Return the optional constraint set associated with this pose stack.

block_identity_map() → Tensor[slice(None, None, None), slice(None, None, None)]

Return each residue’s padded block index for every pose.

class tmol.pose.PoseStackBuilder

Bases: object

Build heterogeneous pose stacks while preserving chemical-database identity.

classmethod from_poses(pose_stacks: List[PoseStack], device: device) → PoseStack

Combine one or more pose stacks on a common device.

Parameters:

• pose_stacks – Pose stacks whose chemical databases are one database or extensions of one another.

• device – Device for the combined tensors.

Returns:

A padded pose stack containing every input pose in order. Split- block mappings are concatenated when present.

classmethod from_block_type_names(packed_block_types: PackedBlockTypes, sequences, chain_lengths)

Construct a zero-coordinate PoseStack from per-pose block type names.

A name may carry a non-polymeric connection as “NAME–conn-label”, where “conn” is the name of an inter-residue connection on that block type and “label” pairs up the two partners. E.g. a pose with two disulfides:

AAAA[CYD–dslf-first]AAA[CYD–dslf-second]AAA … AA[CYD–dslf-second]AAAA[CYD–dslf-first]AAA

class tmol.pose.SeqToken(letter: str, name: str | None = None, smiles: str | None = None)

Bases: object

One residue position in a sequence string.

class tmol.pose.SplitBlockEntry(pose_ind: int, block_ind: int, group_ind: int, orig_block_type_ind: int, split_to_orig_atom_inds: ndarray, orig_residue_label: int, orig_chain_label: str, orig_ins_code: str)

Bases: object

Mapping for one split (fragment) block back to its original block type.

All entries with the same (pose_ind, group_ind) were split from the same original block and should be merged back together when unsplitting.

split_to_orig_atom_inds[i] gives the index of the i-th split-block atom within the original block type’s atom list.

orig_residue_label, orig_chain_label, and orig_ins_code record the PDB residue identity of the original (pre-split) block so that the original numbering can be restored without reference to any separate ligand-specific mapping object.

class tmol.pose.SplitBlockMapping(entries: tuple[SplitBlockEntry, ...])

Bases: object

All split-block records across every pose in a PoseStack.

One SplitBlockEntry per (pose, fragment-block). Entries are not required to be sorted; callers should group by (pose_ind, group_ind) to find all fragments of a given original block.

split(pose_index: int) → SplitBlockMapping

Return the mapping for a single pose, with pose_ind reset to 0.

tmol.pose.get_named_torsions(pose_stack: PoseStack, poses=None, blocks=None, names=None, degrees: bool = True)

Measure named torsions.

Returns a single float when poses, blocks and names are all scalars, a {name: value} dict when only names is left open, and otherwise a list-of-lists of such dicts indexed [pose][block]. Torsions that reach an absent neighboring residue measure as nan.

tmol.pose.get_torsion_names(pose_stack: PoseStack, pose: int, block: int) → List[str]

Names of the torsions defined on a block’s type, in database order.

tmol.pose.resolve_block_type_names(tokens: List[List[SeqToken]], chain_lengths: List[List[int]], restype_set, ligand_names: Dict[str, str] | None = None, termini: bool = True) → Tuple[List[List[str]], List[List[int]]]

Turn tokens into block type names and chain lengths.

Chains are split wherever a polymer bond cannot form, so ligands and waters end up in chains of their own.

tmol.pose.smiles_in_tokens(tokens: List[List[SeqToken]]) → List[str]

Unique SMILES strings appearing in the tokenized sequences.

tmol.pose.tokenize_sequences(seqs: str | Sequence[str]) → Tuple[List[List[SeqToken]], List[List[int]]]

Split sequence strings into per-pose tokens and chain lengths.

Public defaults

DEFAULT_ATOM_B_FACTOR	Convert a string or number to a floating-point number, if possible.
DEFAULT_ATOM_OCCUPANCY	Convert a string or number to a floating-point number, if possible.