Kinematics#

Fold forests, move maps, and kinematic coordinate operations.

Kinematic trees, fold forests, move maps, and coordinate transforms.

class tmol.kinematics.BondDOF(raw: Tensor[..., 4])[source]#

Bases: TensorGroup, ConvertAttrs

A bond dof view of KinDOF.

class tmol.kinematics.BondDOFTypes(*values)[source]#

Bases: IntEnum

Indices of bond dof types within KinDOF.raw.

class tmol.kinematics.CartesianMoveMap(coord_mask: Tensor | None = None)[source]#

Bases: object

Move map for Cartesian-space minimization.

A lightweight wrapper around a coordinate mask that can be passed as the move_map argument to fast_relax() when using a Cartesian min_fn (e.g. one built around run_cart_min()).

Unlike the full MoveMap, which describes freedom in internal (torsion/jump) DOF space, CartesianMoveMap works directly in atomic Cartesian space. The min_fn is responsible for extracting coord_mask and passing it to CartesianSfxnNetwork.

Parameters:: coord_mask – Boolean tensor of shape [n_poses, max_n_atoms] indicating which atoms are free to move. None means all atoms are free to move (the default behaviour of CartesianSfxnNetwork).

tmol.kinematics.CoordArray#: alias of Tensor

class tmol.kinematics.DOFMetadata(node_idx: Tensor, dof_idx: Tensor, dof_type: Tensor, parent_id: Tensor, child_id: Tensor)[source]#

Bases: TensorGroup, ConvertAttrs

The location, type, and descriptive ids of valid dofs within a KinForest.

Descriptive entries for dofs within a KinForest, this provides a 1-d structure to select and report a subset of entries within a KinDOF buffer. DOFMetadata sets are used to indicate mobile vs fixed dofs for KinematicOp dof to coordinate functions.

DOFMetadata supports isomorphic conversion between a DataFrame and TensorGroup representation to support symbolic selection. This converts the IntEnum encoded “dof_type” entry into a string categorical column.

The DOFMetadata data members, just like the KinForest, suffer from the same confusion about what an index represents because there are two ways to index the data:

The “Target Order” (TO) that refers to the index of an atom in the PoseStack it came from where the coordinate tensor is squashed to (N,3)
The “KinForest Order” (KFO) that refers to the order that an atom’s node appears in the KinForest; this second ordering puts the index of any child atom after the index for any parent atom

The DOFMetadata class indexes all available DOFs in the system. There are 9 possible DOFs per atom (either 3 for BondedAtoms or 9 for JumpAtoms), but in actuality, there are many fewer valid DOFs. The DOFMetadata class indexes valid DOFs.

For each valid DOF i, there’s: - node_idx[i]: the KFO index of the atom that DOF i belongs to - dof_idx[i]: the index between 0-8 for DOF i on its atom - dof_type[i]: the DOF type (either a BondDOFType or a JumpDOFType) for DOF i - parent_id[i]: the TO index for the parent to node_idx[i] for DOF i - child_id[i]: the TO index for node_idx[i] for DOF i

The DOFMetadata class is primarily used to index into torch tensors in python, and therefore all of its dtypes are 64-bit integers.

classmethod for_kinforest(kinforest: KinForest)[source]#: Return all valid dofs within a KinForest.

classmethod from_frame(frame)[source]#: Convert from DataFrame to metadata, discarding any unneeded columns.

class tmol.kinematics.DOFTypes(*values)[source]#

Bases: IntEnum

High-level class of kinematic DOF types.

class tmol.kinematics.EdgeType(*values)[source]#

Bases: IntEnum

Kinds of directed connections represented in a fold forest.

class tmol.kinematics.FoldForest(max_n_edges: int, n_edges: NDArray, edges: NDArray[slice(None, None, None), slice(None, None, None), 4])[source]#

Bases: object

The fold forest will define the fold trees for the poses in a PoseStack. Each tensor in the class has its first dimension over the number of poses.

The primary definition of a FoldTree is the Edge. The Edge defines a connection between two parts of a Pose. The three types of edges are 1. polymer edges (analgous to the previously named “peptide edges” from Rosetta++ and Rosetta3), 2. jump edges which connect any pair of residues in the Pose, and 3. root-jump edges, which originate at the explicit virtual root and connect to a particular residue. A polymer edge spans a contiguous range of polymeric block types where the “up” connection of residue i is connected to the “down” connection of residue i+1 for all i in the range between the “start” and “end” blocks.

Each edge is described by a 4-tuple of integers (type, start, end, jump-index); where type is one of the EdgeType enum values, start is the index of the upstream residue of the edge, end is the index of the downstream residue of the edge, and jump-index is used to assign an id to any particular jump edge; jump-edge indices must be unique and ascending from 0 to n_jumps-1. “Root jump” edges take their “identity” from the downstream residue of the edge, so they do not need an index.

The FoldForest in tmol differs from the FoldTree in Rosetta3 in that there is always a virtual root at the origin and any residue (block) may be connected to this root by a “root jump”. Such root-jump residues are defined by listing the residue that the root is connected to as the “end” residue; the “start” residue field should be left as -1. An example FoldForest for a ten-residue protein might be:

(polymer, 0, 4)
(jump   , 0, 7)
(polymer, 7, 9)
(polymer, 7, 6)
(root-jump, -1, 0)
(root-jump, -1, 5)

where both residues are 0 and 5 are connected to the root.

Note that in the MoveMap, the root-jumps are distinct from the non-root-jumps.

classmethod from_edges(edges: NDArray[slice(None, None, None), slice(None, None, None), 4]) → FoldForest[source]#

Construct a fold forest from a padded edge array.

Parameters:: edges – Edge records shaped [pose, edge, 4]. Unused entries have -1 in their edge-type column.
Returns:: A fold forest with per-pose edge counts inferred from edges.

classmethod reasonable_fold_forest(pose_stack: PoseStack)[source]#

Create a fold forest for each pose using only backbone (up/down) polymer connectivity.

Each biological chain (same chain_id) is rooted with a single root-jump to its first residue. Polymer gaps within that chain (chain breaks) become ordinary jump edges connecting the last residue before the gap to the first residue after it. Gaps between different biological chains produce separate root-jumps. Cyclic polymers (C→N cyclisation) are broken at the bond entering the lowest-index residue; that bond is dropped to keep the forest a valid tree. Non-polymer connections (disulfides, etc.) are ignored.

class tmol.kinematics.JumpDOF(raw: Tensor[..., 9])[source]#

Bases: TensorGroup, ConvertAttrs

A jump dof view of KinDOF.

class tmol.kinematics.JumpDOFTypes(*values)[source]#

Bases: IntEnum

Indices of jump dof types within KinDOF.raw.

class tmol.kinematics.KinDOF(raw: Tensor[..., 9])[source]#

Bases: TensorGroup, ConvertAttrs

Internal coordinate data.

The KinDOF data structure holds two logical views: the “raw” view a sparsely populated [n,9] tensor of DOF values and a set of named property accessors providing access to specific entries within this array. This is logically equivalent a C union datatype, the interpretation of an entry in the DOF buffer depends on the type of the corresponding KinForest entry.

class tmol.kinematics.KinForest(id: Tensor, doftype: Tensor, parent: Tensor, frame_x: Tensor, frame_y: Tensor, frame_z: Tensor)[source]#

Bases: TensorGroup, ConvertAttrs

A collection of atom-level kinematic trees, each of which can be processed in parallel.

A kinematic description of a collection of atom locations, each atom location corresponding to a node within a tree. The root of each tree in this forest is built from a jump from the global reference frame at the origin. (The global reference frame will later be treated as a node in the forest, effectively linking all the trees, but this is a minor technical detail; best to think of this as several independent trees than a single tree). Every other node corresponds to a derived orientation, with an atomic coordinate at the center of the frame.

Each node in the tree is connected by one of two “node types”:

1) Jump nodes, representing an arbitrary rigid body transform between two reference frames via six degrees of freedom, 3 translational and 3 rotational.

2) Bond nodes, representing the relationships between two atom reference frames via three bond degrees of freedom: the translation from the parent to the child along the bond axis (bond length, d), the rotation from the grand-parent-to-parent bond axis to the bond axis (an improper bond angle, theta), and the rotation about the grand-parent-to-parent bond axis (bond torsion, phi). Bond nodes include an additional, redundent, degree of freedom representing concerted rotation of all downstream atoms about the parent-to-self bond. These DOFs are used to represent the torsions that alter the location of several children. For example, chi1 is represented as the 4th DOF for the CB atom of LEU. A rotation about the CA–>CB bond axis will spin CG, HB1 and HB2. In this scheme, the phi DOF would be 0 for CG, 120 for HB1 and 240 for HB2. This differs from the Rosetta3 implementation of downstream-dihedral propagation where Chi1 would live as the phi DOF of CG, and CG’s rotation would carry forward to HB1 and HB2 (requiring that CG be the first child of CB).

The atoms in the KinForest have their own order that is distinct from the ordering in the target (e.g. a PoseStack) where there might be gaps between sets of atoms (e.g. because each Pose in the stack has a different number of atoms, so a contiguous block of atom indices from 0-100 might have a gap before the next contiguous block begins at 150). When working with a KinForest, remembering what order and array’s indices is in (the kin-forest order (KFO) or the target order (TO)) and what a value/index read out of an array represents (is the index an index in KFO or TO?) is very challenging. The documentation for these arrays includes whether the arrays are indexed in KFO or TO and whether the values they hold are KFO or TO indices.

The KinForest data structure itself is frozen and can not be modified post construction. KinForests are not typically built directly by users, but rather are constructed as part of building the KinModuleData class by construct_kin_module_data that lives in scan_ordering.py. The _KinematicBuilder factory class is now deprecated but use to be responsible for construction of KinForests

Indices::: id = the TO index in KFO; i.e. kin_forest_order_2_target_order parent = KFO index of the parent, in KFO frame_x = KFO index of self, in KFO frame_y = KFO index of parent, in KFO frame_z = KFO index of grandparent, in KFO

classmethod node(id: int, doftype: NodeType, parent: int, frame_x: int, frame_y: int, frame_z: int)[source]#: Construct a single node from element values.

classmethod root_node()[source]#: The global/root kinematic node at KinForest[0].

class tmol.kinematics.KinForestScanData(nodes: Tensor, scans: Tensor, gens: Tensor)[source]#

Bases: TensorGroup, ConvertAttrs

Node, scan-start, and generation-start tensors for a kinematic scan.

class tmol.kinematics.KinematicModuleData(forest: KinForest, scan_data_fw: KinForestScanData, scan_data_bw: KinForestScanData, block_in_and_first_out: Tensor[slice(None, None, None), slice(None, None, None)], keep_atom_fixed: Tensor[slice(None, None, None), slice(None, None, None)], pose_stack_atom_for_jump: Tensor[slice(None, None, None), slice(None, None, None), 2], pose_stack_atom_for_root_jump: Tensor[slice(None, None, None), 2])[source]#

Bases: object

Topology and traversal tensors required by a pose kinematics module.

class tmol.kinematics.MinimizerMap(pose_stack: PoseStack, kmd: KinematicModuleData, mm: MoveMap)[source]#

Bases: object

Resolved per-DOF movement mask for a pose’s kinematic forest.

class tmol.kinematics.MoveMap(n_poses: int, max_n_blocks: int, max_n_named_torsions: int, max_n_atoms_per_block: int, device: device | None = None)[source]#

Bases: object

Hierarchical internal-coordinate degrees of freedom for a pose stack.

classmethod from_pose_stack(ps: PoseStack) → MoveMap[source]#: Construct a move map sized for ps with all movement disabled.

set_move_all_jump_dofs_for_jump(pose_selection: Tensor | int, jump_selection: Tensor | int | None = None, value: bool = True) → None[source]#

Enable or disable all jump dofs for a particular set of jumps on particular poses.

If jump_selection is None, then the two dimensional settings tensor will be indexed by the pose_selection tensor only; if both are not None, then the settings tensor will be indexed by both the pose_selection tensor and the jump_selection tensor.

set_move_all_jump_dofs_for_root_jump(pose_selection: Tensor | int, root_jump_selection: Tensor | int | None = None, value: bool = True) → None[source]#

Enable or disable all jump dofs for a particular set of root-jumps on particular poses.

If root_jump_selection is None, then the two dimensional settings tensor will be indexed by the pose_selection tensor only; if both are not None, then the settings tensor will be indexed by both the pose_selection tensor and the jump_selection tensor.

set_move_all_mc_tors_for_blocks(pose_selection: Tensor | int, block_selection: Tensor | int | None = None, value: bool = True) → None[source]#

Enable or disable all DOFs for a partiular set of blocks on particular poses.

If block_selection is None, then the two dimensional settings tensor will be indexed by the pose_selection tensor only; if both are not None, then the tensor will be indexed by the pose_selection tensor and the block_selection tensor.

Valid combinations of pose_selection and block_selection are, e.g.:

pose_selection: int, block_selection: int == a particular block on a particular pose
pose_selection: int, block_selection: None == all blocks on a particular pose
pose_selection: Tensor[bool][n_poses, max_n_blocks], block_selection: None == pose/block pairs encoded in “pose_selection” tensor
pose_selection: Tensor[int][N], block_selection: Tensor[int][N] == different blocks on different poses, selected by index

Enable or disable partiular atom dofs for a particular set of blocks on particular poses.

Either only “pose_selection” should be not None or all four “selection” variables should be not None; in the former case, the settings tensor, self.move_atom_dof, will be indexed solely by the pose_selection tensor; in the latter case, the settings tensor will be indexed by all four selection tensors.

This function offers the finest grain control over which dofs should be minimized and settings made using this function will override any settings made using the other settings tensors.

Enable or disable all jump dofs for a particular set of jumps on particular poses.

Enable or disable all jump dofs for a particular set of root-jumps on particular poses.

Enable or disable partiular main-chain torsions for a particular set of blocks on particular poses.

Valid combinations of block_selection and tor_selection are:

pose_selection: int, block_selection: int, tor_selection: int == a single DOF
pose_selection: Tensor[bool][n_poses, max_n_blocks, max_n_dofs], block_selection: None, tor_selection: None == pose/block/tor triples encoded in “pose_selection” tensor
pose_selection: Tensor[int][N], block_selection: Tensor[int][N], dof_selection: Tensor[int][N] == different torsions on different blocks on different poses, selected by index

set_move_named_torsion_for_blocks(pose_selection: Tensor | int, block_selection: Tensor | int | None = None, tor_selection: Tensor | int | None = None, value: bool = True) → None[source]#: Enable or disable partiular named-torsions for a particular set of blocks on particular poses.

set_move_sc_tor_for_blocks(pose_selection: Tensor | int, block_selection: Tensor | int | None = None, dof_selection: Tensor | int | None = None, value: bool = True) → None[source]#: Enable or disable partiular side-chain torsions for a particular set of blocks on particular poses.

class tmol.kinematics.NodeType(*values)[source]#

Bases: IntEnum

KinForest node types.

class tmol.kinematics.PoseStackKinematicsModule(pose_stack: PoseStack, fold_forest: FoldForest)[source]#

Bases: Module

torch.autograd compatible forward kinematic operator for PoseStack.

Perform forward (dof to coordinate) kinematics within torch.autograd compute graph. Provides support for forward kinematics over of a subset of source dofs, as specified by the provided DOFMetadata entries.

The kinematic system maps between the natm x 9 internal coordinate frame and the natm x 3 coordinate frame. Some of this natm x 9 array is unused or is redundant but this is not known by the kinematic module.

See KinDOF for a description of the internal coordinate representation.

forward(dofs)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

tmol.kinematics.construct_kin_module_data_for_pose(pose_stack: PoseStack, fold_forest_edges: Tensor[slice(None, None, None), slice(None, None, None), 4])[source]#

Construct the indexed forest and scan order for batched poses.

Parameters:

pose_stack – Poses whose atoms will form the kinematic forest.
fold_forest_edges – Fold-forest edges shaped [pose, edge, 4].

Returns:

Immutable topology and traversal data for forward and backward scans.

tmol.kinematics.inverseKin(kinforest: KinForest, coords: Tensor[slice(None, None, None), 3], requires_grad=False) → KinDOF[source]#: xyzs -> HTs, dofs - “backward” kinematics

tmol.kinematics.set_named_torsions(pose_stack: PoseStack, poses, blocks, names, values, degrees: bool = True, fold_forest=None) → PoseStack[source]#

Set named torsions, returning a new PoseStack.

poses, blocks, names and values are either all scalars or equal-length sequences; a batch is applied in one kinematic pass. Movement follows the fold forest, so which side of each bond stays fixed depends on it; reasonable_fold_forest roots each chain at its first residue.

tmol.kinematics.validate_fold_forest(n_blocks: NDArray, edges: NDArray[slice(None, None, None), slice(None, None, None), 4])[source]#

Validate batched fold-forest edges and raise with pose-specific details.

Parameters:

n_blocks – Number of real blocks in each pose, shaped [pose].
edges – Fold-forest edges shaped [pose, edge, 4].

Raises:

ValueError – If an edge is invalid, connectivity is incomplete or cyclic, or jump numbering is inconsistent.