deepfold.utils.geometry package

Submodules

deepfold.utils.geometry.quat_rigid module

class deepfold.utils.geometry.quat_rigid.QuatRigid(c_hidden, full_quat)

Bases: Module

forward(activations: Tensor) → Rigid3Array

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.

deepfold.utils.geometry.rigid_matrix_vector module

Rigid3Array Transformations represented by a Matrix and a Vector.

class deepfold.utils.geometry.rigid_matrix_vector.Rigid3Array(rotation: Rot3Array, translation: Vec3Array)

Bases: object

Rigid Transformation, i.e. element of special euclidean group.

apply(point: Tensor) → Tensor

apply_inverse_to_point(point: Vec3Array) → Vec3Array

Apply inverse Rigid3Array transform to point.

apply_to_point(point: Vec3Array) → Vec3Array

Apply Rigid3Array transform to point.

classmethod cat(rigids: List[Rigid3Array], dim: int) → Rigid3Array

compose(other_rigid)

compose_rotation(other_rotation)

cuda() → Rigid3Array

property device: device

property dtype: dtype

classmethod from_array(array)

classmethod from_array4x4(array: Tensor) → Rigid3Array

Construct Rigid3Array from homogeneous 4x4 array.

classmethod from_tensor_4x4(array)

classmethod identity(shape, device) → Rigid3Array

Return identity Rigid3Array of given shape.

inverse() → Rigid3Array

Return Rigid3Array corresponding to inverse transform.

invert_apply(point: Tensor) → Tensor

map_tensor_fn(fn) → Rigid3Array

reshape(new_shape) → Rigid3Array

rotation: Rot3Array

scale_translation(factor: float | Tensor) → Rigid3Array

Scale translation in Rigid3Array by ‘factor’.

property shape: Size

stop_rot_gradient() → Rigid3Array

to_tensor() → Tensor

to_tensor_4x4() → Tensor

translation: Vec3Array

unsqueeze(dim: int)

deepfold.utils.geometry.rotation_matrix module

Rot3Array Matrix Class.

class deepfold.utils.geometry.rotation_matrix.Rot3Array(xx: Tensor, xy: Tensor, xz: Tensor, yx: Tensor, yy: Tensor, yz: Tensor, zx: Tensor, zy: Tensor, zz: Tensor)

Bases: object

Rot3Array Matrix in 3 dimensional Space implemented as struct of arrays.

apply_inverse_to_point(point: Vec3Array) → Vec3Array

Applies inverse Rot3Array to point.

apply_to_point(point: Vec3Array) → Vec3Array

Applies Rot3Array to point.

classmethod cat(rots: List[Rot3Array], dim: int) → Rot3Array

classmethod from_array(array: Tensor) → Rot3Array

Construct Rot3Array Matrix from array of shape. […, 3, 3].

classmethod from_quaternion(w: Tensor, x: Tensor, y: Tensor, z: Tensor, normalize: bool = True, eps: float = 1e-06) → Rot3Array

Construct Rot3Array from components of quaternion.

classmethod from_two_vectors(e0: Vec3Array, e1: Vec3Array) → Rot3Array

Construct Rot3Array from two Vectors.

Rot3Array is constructed such that in the corresponding frame ‘e0’ lies on the positive x-Axis and ‘e1’ lies in the xy plane with positive sign of y.

Parameters:

• e0 – Vector

• e1 – Vector

Returns:

Rot3Array

classmethod identity(shape, device) → Rot3Array

Returns identity of given shape.

inverse() → Rot3Array

Returns inverse of Rot3Array.

map_tensor_fn(fn) → Rot3Array

reshape(new_shape)

stop_gradient() → Rot3Array

to_tensor() → Tensor

Convert Rot3Array to array of shape […, 3, 3].

unsqueeze(dim: int)

xx: Tensor

xy: Tensor

xz: Tensor

yx: Tensor

yy: Tensor

yz: Tensor

zx: Tensor

zy: Tensor

zz: Tensor

deepfold.utils.geometry.utils module

Utils for geometry library.

deepfold.utils.geometry.utils.get_field_names(cls)

deepfold.utils.geometry.vector module

Vec3Array Class.

class deepfold.utils.geometry.vector.Vec3Array(x: 'torch.Tensor', y: 'torch.Tensor', z: 'torch.Tensor')

Bases: object

classmethod cat(vecs: List[Vec3Array], dim: int) → Vec3Array

clone() → Vec3Array

cross(other: Vec3Array) → Vec3Array

Compute cross product between ‘self’ and ‘other’.

dot(other: Vec3Array) → float | Tensor

Compute dot product between ‘self’ and ‘other’.

classmethod from_array(tensor)

map_tensor_fn(fn) → Vec3Array

norm(epsilon: float = 1e-06) → float | Tensor

Compute Norm of Vec3Array, clipped to epsilon.

norm2()

normalized(epsilon: float = 1e-06) → Vec3Array

Return unit vector with optional clipping.

reshape(new_shape) → Vec3Array

property shape

sum(dim: int) → Vec3Array

to_tensor() → Tensor

unsqueeze(dim: int)

x: Tensor

y: Tensor

z: Tensor

classmethod zeros(shape, device='cpu')

Return Vec3Array corresponding to zeros of given shape.

deepfold.utils.geometry.vector.cross(vector1: Vec3Array, vector2: Vec3Array) → Vec3Array

deepfold.utils.geometry.vector.dihedral_angle(a: Vec3Array, b: Vec3Array, c: Vec3Array, d: Vec3Array) → float | Tensor

Computes torsion angle for a quadruple of points.

For points (a, b, c, d), this is the angle between the planes defined by points (a, b, c) and (b, c, d). It is also known as the dihedral angle.

Parameters:

• a – A Vec3Array of coordinates.

• b – A Vec3Array of coordinates.

• c – A Vec3Array of coordinates.

• d – A Vec3Array of coordinates.

Returns:

[-pi, pi].

Return type:

A tensor of angles in radians

deepfold.utils.geometry.vector.dot(vector1: Vec3Array, vector2: Vec3Array) → float | Tensor

deepfold.utils.geometry.vector.euclidean_distance(vec1: Vec3Array, vec2: Vec3Array, epsilon: float = 1e-06) → float | Tensor

Computes euclidean distance between ‘vec1’ and ‘vec2’.

Parameters:

• vec1 – Vec3Array to compute euclidean distance to

• vec2 – Vec3Array to compute euclidean distance from, should be broadcast compatible with ‘vec1’

• epsilon – distance is clipped from below to be at least epsilon

Returns:

Array of euclidean distances; shape will be result of broadcasting ‘vec1’ and ‘vec2’

deepfold.utils.geometry.vector.norm(vector: Vec3Array, epsilon: float = 1e-06) → float | Tensor

deepfold.utils.geometry.vector.normalized(vector: Vec3Array, epsilon: float = 1e-06) → Vec3Array

deepfold.utils.geometry.vector.square_euclidean_distance(vec1: Vec3Array, vec2: Vec3Array, epsilon: float = 1e-06) → float | Tensor

Computes square of euclidean distance between ‘vec1’ and ‘vec2’.

Parameters:

• vec1 – Vec3Array to compute distance to

• vec2 – Vec3Array to compute distance from, should be broadcast compatible with ‘vec1’

• epsilon – distance is clipped from below to be at least epsilon

Returns:

Array of square euclidean distances; shape will be result of broadcasting ‘vec1’ and ‘vec2’

Module contents

class deepfold.utils.geometry.Rigid(rots: Rotation | None, trans: Tensor | None)

Bases: object

A class representing a rigid transformation. Little more than a wrapper around two objects: a Rotation object and a [*, 3] translation Designed to behave approximately like a single torch tensor with the shape of the shared batch dimensions of its component parts.

apply(pts: Tensor) → Tensor

Applies the transformation to a coordinate tensor.

Parameters:

pts – A [*, 3] coordinate tensor.

Returns:

The transformed points.

apply_rot_fn(fn: Callable[[Rotation], Rotation]) → Rigid

Applies a Rotation -> Rotation function to the stored rotation object.

Parameters:

fn – A function of type Rotation -> Rotation

Returns:

A transformation object with a transformed rotation.

apply_trans_fn(fn: Callable[[Tensor], Tensor]) → Rigid

Applies a Tensor -> Tensor function to the stored translation.

Parameters:

fn – A function of type Tensor -> Tensor to be applied to the translation

Returns:

A transformation object with a transformed translation.

static cat(ts: Sequence[Rigid], dim: int) → Rigid

Concatenates transformations along a new dimension.

Parameters:

• ts – A list of T objects

• dim – The dimension along which the transformations should be concatenated

Returns:

A concatenated transformation object

compose(r: Rigid) → Rigid

Composes the current rigid object with another.

Parameters:

r – Another Rigid object

Returns:

The composition of the two transformations

compose_q_update_vec(q_update_vec: Tensor) → Rigid

Composes the transformation with a quaternion update vector of shape [*, 6], where the final 6 columns represent the x, y, and z values of a quaternion of form (1, x, y, z) followed by a 3D translation.

Parameters:

q_vec – The quaternion update vector.

Returns:

The composed transformation.

cuda() → Rigid

Moves the transformation object to GPU memory

Returns:

A version of the transformation on GPU

property device: device

Returns the device on which the Rigid’s tensors are located.

Returns:

The device on which the Rigid’s tensors are located

property dtype: dtype

Returns the dtype of the Rigid tensors.

Returns:

The dtype of the Rigid tensors

static from_3_points(p_neg_x_axis: Tensor, origin: Tensor, p_xy_plane: Tensor, eps: float = 1e-08) → Rigid

Implements algorithm 21. Constructs transformations from sets of 3 points using the Gram-Schmidt algorithm.

Parameters:

• p_neg_x_axis – [*, 3] coordinates

• origin – [*, 3] coordinates used as frame origins

• p_xy_plane – [*, 3] coordinates

• eps – Small epsilon value

Returns:

A transformation object of shape [*]

static from_tensor_4x4(t: Tensor) → Rigid

Constructs a transformation from a homogenous transformation tensor.

Parameters:

t – [*, 4, 4] homogenous transformation tensor

Returns:

T object with shape [*]

static from_tensor_7(t: Tensor, normalize_quats: bool = False) → Rigid

get_rots() → Rotation

Getter for the rotation.

Returns:

The rotation object

get_trans() → Tensor

Getter for the translation.

Returns:

The stored translation

static identity(shape: Tuple[int], dtype: dtype | None = None, device: device | None = None, requires_grad: bool = True, fmt: str = 'quat') → Rigid

Constructs an identity transformation.

Parameters:

• shape – The desired shape

• dtype – The dtype of both internal tensors

• device – The device of both internal tensors

• requires_grad – Whether grad should be enabled for the internal tensors

Returns:

The identity transformation

invert() → Rigid

Inverts the transformation.

Returns:

The inverse transformation.

invert_apply(pts: Tensor) → Tensor

Applies the inverse of the transformation to a coordinate tensor.

Parameters:

pts – A [*, 3] coordinate tensor

Returns:

The transformed points.

static make_transform_from_reference(n_xyz, ca_xyz, c_xyz, eps=1e-20)

Returns a transformation object from reference coordinates.

Note that this method does not take care of symmetries. If you provide the atom positions in the non-standard way, the N atom will end up not at [-0.527250, 1.359329, 0.0] but instead at [-0.527250, -1.359329, 0.0]. You need to take care of such cases in your code.

Parameters:

• n_xyz – A [*, 3] tensor of nitrogen xyz coordinates.

• ca_xyz – A [*, 3] tensor of carbon alpha xyz coordinates.

• c_xyz – A [*, 3] tensor of carbon xyz coordinates.

Returns:

A transformation object. After applying the translation and rotation to the reference backbone, the coordinates will approximately equal to the input coordinates.

map_tensor_fn(fn: Callable[[Tensor], Tensor]) → Rigid

Apply a Tensor -> Tensor function to underlying translation and rotation tensors, mapping over the translation/rotation dimensions respectively.

Parameters:

fn – A Tensor -> Tensor function to be mapped over the Rigid

Returns:

The transformed Rigid object

scale_translation(trans_scale_factor: float) → Rigid

Scales the translation by a constant factor.

Parameters:

trans_scale_factor – The constant factor

Returns:

A transformation object with a scaled translation.

property shape: Size

Returns the shape of the shared dimensions of the rotation and the translation.

Returns:

The shape of the transformation

stop_rot_gradient() → Rigid

Detaches the underlying rotation object

Returns:

A transformation object with detached rotations

to_tensor_4x4() → Tensor

Converts a transformation to a homogenous transformation tensor.

Returns:

A [*, 4, 4] homogenous transformation tensor

to_tensor_7() → Tensor

Converts a transformation to a tensor with 7 final columns, four for the quaternion followed by three for the translation.

Returns:

A [*, 7] tensor representation of the transformation

unsqueeze(dim: int) → Rigid

Analogous to torch.unsqueeze. The dimension is relative to the shared dimensions of the rotation/translation.

Parameters:

dim – A positive or negative dimension index.

Returns:

The unsqueezed transformation.

class deepfold.utils.geometry.Rigid3Array(rotation: Rot3Array, translation: Vec3Array)

Bases: object

Rigid Transformation, i.e. element of special euclidean group.

apply(point: Tensor) → Tensor

apply_inverse_to_point(point: Vec3Array) → Vec3Array

Apply inverse Rigid3Array transform to point.

apply_to_point(point: Vec3Array) → Vec3Array

Apply Rigid3Array transform to point.

classmethod cat(rigids: List[Rigid3Array], dim: int) → Rigid3Array

compose(other_rigid)

compose_rotation(other_rotation)

cuda() → Rigid3Array

property device: device

property dtype: dtype

classmethod from_array(array)

classmethod from_array4x4(array: Tensor) → Rigid3Array

Construct Rigid3Array from homogeneous 4x4 array.

classmethod from_tensor_4x4(array)

classmethod identity(shape, device) → Rigid3Array

Return identity Rigid3Array of given shape.

inverse() → Rigid3Array

Return Rigid3Array corresponding to inverse transform.

invert_apply(point: Tensor) → Tensor

map_tensor_fn(fn) → Rigid3Array

reshape(new_shape) → Rigid3Array

rotation: Rot3Array

scale_translation(factor: float | Tensor) → Rigid3Array

Scale translation in Rigid3Array by ‘factor’.

property shape: Size

stop_rot_gradient() → Rigid3Array

to_tensor() → Tensor

to_tensor_4x4() → Tensor

translation: Vec3Array

unsqueeze(dim: int)

class deepfold.utils.geometry.Rot3Array(xx: Tensor, xy: Tensor, xz: Tensor, yx: Tensor, yy: Tensor, yz: Tensor, zx: Tensor, zy: Tensor, zz: Tensor)

Bases: object

Rot3Array Matrix in 3 dimensional Space implemented as struct of arrays.

apply_inverse_to_point(point: Vec3Array) → Vec3Array

Applies inverse Rot3Array to point.

apply_to_point(point: Vec3Array) → Vec3Array

Applies Rot3Array to point.

classmethod cat(rots: List[Rot3Array], dim: int) → Rot3Array

classmethod from_array(array: Tensor) → Rot3Array

Construct Rot3Array Matrix from array of shape. […, 3, 3].

classmethod from_quaternion(w: Tensor, x: Tensor, y: Tensor, z: Tensor, normalize: bool = True, eps: float = 1e-06) → Rot3Array

Construct Rot3Array from components of quaternion.

classmethod from_two_vectors(e0: Vec3Array, e1: Vec3Array) → Rot3Array

Construct Rot3Array from two Vectors.

Rot3Array is constructed such that in the corresponding frame ‘e0’ lies on the positive x-Axis and ‘e1’ lies in the xy plane with positive sign of y.

Parameters:

• e0 – Vector

• e1 – Vector

Returns:

Rot3Array

classmethod identity(shape, device) → Rot3Array

Returns identity of given shape.

inverse() → Rot3Array

Returns inverse of Rot3Array.

map_tensor_fn(fn) → Rot3Array

reshape(new_shape)

stop_gradient() → Rot3Array

to_tensor() → Tensor

Convert Rot3Array to array of shape […, 3, 3].

unsqueeze(dim: int)

xx: Tensor

xy: Tensor

xz: Tensor

yx: Tensor

yy: Tensor

yz: Tensor

zx: Tensor

zy: Tensor

zz: Tensor

class deepfold.utils.geometry.Rotation(rot_mats: Tensor | None = None, quats: Tensor | None = None, normalize_quats: bool = True)

Bases: object

A 3D rotation. Depending on how the object is initialized, the rotation is represented by either a rotation matrix or a quaternion, though both formats are made available by helper functions. To simplify gradient computation, the underlying format of the rotation cannot be changed in-place. Like Rigid, the class is designed to mimic the behavior of a torch Tensor, almost as if each Rotation object were a tensor of rotations, in one format or another.

apply(pts: Tensor) → Tensor

Apply the current Rotation as a rotation matrix to a set of 3D coordinates.

Parameters:

pts – A [*, 3] set of points

Returns:

[*, 3] rotated points

static cat(rs: Sequence[Rotation], dim: int) → Rotation

Concatenates rotations along one of the batch dimensions. Analogous to torch.cat().

Note that the output of this operation is always a rotation matrix, regardless of the format of input rotations.

Parameters:

• rs – A list of rotation objects

• dim – The dimension along which the rotations should be concatenated

Returns:

A concatenated Rotation object in rotation matrix format

compose_q(r: Rotation, normalize_quats: bool = True) → Rotation

Compose the quaternions of the current Rotation object with those of another.

Depending on whether either Rotation was initialized with quaternions, this function may call torch.linalg.eigh.

Parameters:

r – An update rotation object

Returns:

An updated rotation object

compose_q_update_vec(q_update_vec: Tensor, normalize_quats: bool = True) → Rotation

Returns a new quaternion Rotation after updating the current object’s underlying rotation with a quaternion update, formatted as a [*, 3] tensor whose final three columns represent x, y, z such that (1, x, y, z) is the desired (not necessarily unit) quaternion update.

Parameters:

• q_update_vec – A [*, 3] quaternion update tensor

• normalize_quats – Whether to normalize the output quaternion

Returns:

An updated Rotation

compose_r(r: Rotation) → Rotation

Compose the rotation matrices of the current Rotation object with those of another.

Parameters:

r – An update rotation object

Returns:

An updated rotation object

cuda() → Rotation

Analogous to the cuda() method of torch Tensors

Returns:

A copy of the Rotation in CUDA memory

detach() → Rotation

Returns a copy of the Rotation whose underlying Tensor has been detached from its torch graph.

Returns:

A copy of the Rotation whose underlying Tensor has been detached from its torch graph

property device: device

The device of the underlying rotation

Returns:

The device of the underlying rotation

property dtype: dtype

Returns the dtype of the underlying rotation.

Returns:

The dtype of the underlying rotation

get_cur_rot() → Tensor

Return the underlying rotation in its current form

Returns:

The stored rotation

get_quats() → Tensor

Returns the underlying rotation as a quaternion tensor.

Depending on whether the Rotation was initialized with a quaternion, this function may call torch.linalg.eigh.

Returns:

The rotation as a quaternion tensor.

get_rot_mats() → Tensor

Returns the underlying rotation as a rotation matrix tensor.

Returns:

The rotation as a rotation matrix tensor

static identity(shape, dtype: dtype | None = None, device: device | None = None, requires_grad: bool = True, fmt: str = 'quat') → Rotation

Returns an identity Rotation.

Parameters:

• shape – The “shape” of the resulting Rotation object. See documentation for the shape property

• dtype – The torch dtype for the rotation

• device – The torch device for the new rotation

• requires_grad – Whether the underlying tensors in the new rotation object should require gradient computation

• fmt – One of “quat” or “rot_mat”. Determines the underlying format of the new object’s rotation

Returns:

A new identity rotation

invert() → Rotation

Returns the inverse of the current Rotation.

Returns:

The inverse of the current Rotation

invert_apply(pts: Tensor) → Tensor

The inverse of the apply() method.

Parameters:

pts – A [*, 3] set of points

Returns:

[*, 3] inverse-rotated points

map_tensor_fn(fn: Callable[[Tensor], Tensor]) → Rotation

Apply a Tensor -> Tensor function to underlying rotation tensors, mapping over the rotation dimension(s). Can be used e.g. to sum out a one-hot batch dimension.

Parameters:

fn – A Tensor -> Tensor function to be mapped over the Rotation

Returns:

The transformed Rotation object

property requires_grad: bool

Returns the requires_grad property of the underlying rotation

Returns:

The requires_grad property of the underlying tensor

property shape: Size

Returns the virtual shape of the rotation object. This shape is defined as the batch dimensions of the underlying rotation matrix or quaternion. If the Rotation was initialized with a [10, 3, 3] rotation matrix tensor, for example, the resulting shape would be [10].

Returns:

The virtual shape of the rotation object

to(device: device | None, dtype: dtype | None) → Rotation

Analogous to the to() method of torch Tensors

Parameters:

• device – A torch device

• dtype – A torch dtype

Returns:

A copy of the Rotation using the new device and dtype

unsqueeze(dim: int) → Rotation

Analogous to torch.unsqueeze. The dimension is relative to the shape of the Rotation object.

Parameters:

dim – A positive or negative dimension index.

Returns:

The unsqueezed Rotation.

class deepfold.utils.geometry.Vec3Array(x: 'torch.Tensor', y: 'torch.Tensor', z: 'torch.Tensor')

Bases: object

classmethod cat(vecs: List[Vec3Array], dim: int) → Vec3Array

clone() → Vec3Array

cross(other: Vec3Array) → Vec3Array

Compute cross product between ‘self’ and ‘other’.

dot(other: Vec3Array) → float | Tensor

Compute dot product between ‘self’ and ‘other’.

classmethod from_array(tensor)

map_tensor_fn(fn) → Vec3Array

norm(epsilon: float = 1e-06) → float | Tensor

Compute Norm of Vec3Array, clipped to epsilon.

norm2()

normalized(epsilon: float = 1e-06) → Vec3Array

Return unit vector with optional clipping.

reshape(new_shape) → Vec3Array

property shape

sum(dim: int) → Vec3Array

to_tensor() → Tensor

unsqueeze(dim: int)

x: Tensor

y: Tensor

z: Tensor

classmethod zeros(shape, device='cpu')

Return Vec3Array corresponding to zeros of given shape.

deepfold.utils.geometry.dihedral_angle(a: Vec3Array, b: Vec3Array, c: Vec3Array, d: Vec3Array) → float | Tensor

Computes torsion angle for a quadruple of points.

For points (a, b, c, d), this is the angle between the planes defined by points (a, b, c) and (b, c, d). It is also known as the dihedral angle.

Parameters:

• a – A Vec3Array of coordinates.

• b – A Vec3Array of coordinates.

• c – A Vec3Array of coordinates.

• d – A Vec3Array of coordinates.

Returns:

[-pi, pi].

Return type:

A tensor of angles in radians

deepfold.utils.geometry.euclidean_distance(vec1: Vec3Array, vec2: Vec3Array, epsilon: float = 1e-06) → float | Tensor

Computes euclidean distance between ‘vec1’ and ‘vec2’.

Parameters:

• vec1 – Vec3Array to compute euclidean distance to

• vec2 – Vec3Array to compute euclidean distance from, should be broadcast compatible with ‘vec1’

• epsilon – distance is clipped from below to be at least epsilon

Returns:

Array of euclidean distances; shape will be result of broadcasting ‘vec1’ and ‘vec2’

deepfold.utils.geometry.square_euclidean_distance(vec1: Vec3Array, vec2: Vec3Array, epsilon: float = 1e-06) → float | Tensor

Computes square of euclidean distance between ‘vec1’ and ‘vec2’.

Parameters:

• vec1 – Vec3Array to compute distance to

• vec2 – Vec3Array to compute distance from, should be broadcast compatible with ‘vec1’

• epsilon – distance is clipped from below to be at least epsilon

Returns:

Array of square euclidean distances; shape will be result of broadcasting ‘vec1’ and ‘vec2’