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.

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.

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.

clone() → ConstraintSet

Return a copy with independent tensor storage.

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.

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.

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

Create an empty constraint set on a concrete device.

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.

split(index: int) → ConstraintSet

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

to(device: device) → ConstraintSet

Return a copy whose tensors and metadata use device.

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.

The PackedBlockTypes object is the bag in which scoring terms cache their tensors holding the chemical/scoring properties of the block types in use. Each term needs several tensors in order to map from block-type index to the data required to score that block type, and the construction of these tensors and moving these tensors can be slowThe idiom we follow to ensure that these tensors are preserved between score evaluations is to cache them in this object. The term will annotate the PackedBlockTypes object, pbt, using setattr(pbt, “tensor_name”, tensor) and then will later decide if the annotation has already been made using hasattr(pbt, “tensor_name”). Thus, it is more efficient to use a single PackedBlockTypes object between multiple PoseStack objects so that the expense of creating the annotations can be amortized of many score evaluations.

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.

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: Tensor[slice(None, None, None), slice(None, None, None), slice(None, None, None), slice(None, None, None), slice(None, None, None)], inter_block_bondsep64: Tensor[slice(None, None, None), slice(None, None, None), slice(None, None, None), slice(None, None, None), slice(None, None, None)], 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 – Capped bond separation between residue connections.

• inter_block_bondsep64 – 64-bit bond-separation copy.

• 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.

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

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

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

clone() → PoseStack

Deep-copy clone of this PoseStack

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].

get_constraint_set() → ConstraintSet | None

Return the optional constraint set associated with this pose stack.

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

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

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_blocks: int

Return the padded residue count per pose.

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 n_poses: int

Return the number of poses in the stack.

property n_res_per_pose: Tensor

Return the number of real residues in each pose.

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

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

split(index: int) → PoseStack

Copy one pose into a new single-pose stack.

class tmol.pose.PoseStackBuilder

Bases: object

Build heterogeneous pose stacks while preserving chemical-database identity.

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

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 built from the same chemical database.

• device – Device for the combined tensors.

Returns:

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

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.