coordinates

coordinates#

Utility functions for manipulating atomic coordinates.

prxteinmpnn.utils.coordinates

prxteinmpnn.utils.coordinates.apply_noise_to_coordinates(key, coordinates, augment_eps=0.0)[source]#

Add Gaussian noise to atomic coordinates.

Parameters:

  • coordinates (Float[Array, 'num_residues num_atoms 3']) – Atomic coordinates of the protein structure. (N, 37, 3)

  • key (Union[Key[Array, ''], UInt32[Array, '2']]) – JAX random key for stochastic operations.

  • augment_eps (float) – Standard deviation for Gaussian noise augmentation.

Return type:

tuple[Float[Array, 'num_residues num_atoms 3'], Union[Key[Array, ''], UInt32[Array, '2']]]

Returns:

Tuple of noisy coordinates and the updated JAX random key.

Example

>>> key = jax.random.PRNGKey(0)
>>> noisy_coords, new_key = apply_noise_to_coordinates(coords, key, 0.1)
prxteinmpnn.utils.coordinates.compute_backbone_coordinates(coordinates)[source]#

Compute backbone coordinates with per-residue C-beta handling using jnp.where.

Parameters:

coordinates (Float[Array, 'num_residues num_atoms 3']) – Atomic coordinates of the protein structure, shape (N, 37, 3).

Return type:

Float[Array, '4 3 3']

Returns:

Backbone coordinates with C-beta atoms computed where necessary, shape (N, 5, 3). NOTE: This deviates from the default atom37 (Nitrogen, C alpha, C, C beta, Oxygen…)

atom ordering.

Example

>>> coords = jnp.zeros((10, 37, 3))  # Example coordinates
>>> backbone_coords = compute_backbone_coordinates(coords)
>>> backbone_coords.shape
(10, 5, 3)
prxteinmpnn.utils.coordinates.compute_c_beta(alpha_to_nitrogen, carbon_to_alpha, alpha_carbon)[source]#

Compute C-beta coordinates.

Uses a linear combination of the bond vectors to estimate C-beta.

Coefficients are derived from empirical data and are used to ensure that the C-beta coordinates are consistent with the geometry of the protein backbone.

Parameters:

  • alpha_to_nitrogen (Float[Array, '3']) – Bond vector from nitrogen to alpha carbon.

  • carbon_to_alpha (Float[Array, '3']) – Bond vector from alpha carbon to carbon.

  • alpha_carbon (Float[Array, '3']) – Coordinates of the alpha carbon atom.

Return type:

Float[Array, '3']

Returns:

C-beta coordinates as an AtomicCoordinate.

Example

>>> n_to_ca = jnp.array([1.0, 0.0, 0.0])
>>> ca_to_c = jnp.array([0.0, 1.0, 0.0])
>>> ca_coords = jnp.array([0 .0, 0.0, 0.0])
>>> cb_coords = compute_c_beta(n_to_ca, ca_to_c, ca_coords)
>>> cb_coords.shape
(3,)
prxteinmpnn.utils.coordinates.compute_backbone_distance(backbone_coordinates)[source]#

Compute pairwise distances between backbone atoms.

Calculate the Euclidean distance between all pairs of backbone atom coordinates based on alpha carbon positions.

Assumes backbone_coordinates is a 3D array of shape (N, 5, 3), where N is the number of atoms, 5 is the number of backbone atoms (N, CA, C, O, N), and 3 is the spatial dimension (x, y, z).

Parameters:

  • backbone_coordinates (Float[Array, '4 3 3']) – A 3D array of shape (N, 5, 3) representing the coordinates of backbone

  • atoms.

Return type:

Float[Array, 'num_atoms num_atoms']

Returns:

A 2D array of shape (N, N) containing the pairwise distances between backbone atoms.

Example

>>> coords = jnp.zeros((10, 5, 3))  # Example coordinates
>>> distances = compute_backbone_distance(coords)
>>> distances.shape
(10, 10)
prxteinmpnn.utils.coordinates.extend_coordinate(atom_a, atom_b, atom_c, bond_length, bond_angle, dihedral_angle)[source]#

Compute the position of a fourth atom (D) given three atoms (A, B, C) and internal coordinates.

Given coordinates for atoms A, B, and C, and the desired bond length, bond angle, and dihedral angle, compute the coordinates of atom D such that:

  • |C-D| = bond_length

  • angle(B, C, D) = bond_angle

  • dihedral(A, B, C, D) = dihedral_angle

Parameters:

  • atom_a (Float[Array, '3']) – Coordinates of atom A, shape (3,).

  • atom_b (Float[Array, '3']) – Coordinates of atom B, shape (3,).

  • atom_c (Float[Array, '3']) – Coordinates of atom C, shape (3,).

  • bond_length (float) – Desired bond length between C and D.

  • bond_angle (float) – Desired bond angle (in radians) at atom C.

  • dihedral_angle (float) – Desired dihedral angle (in radians) for atoms A-B-C-D.

Return type:

Float[Array, '3']

Returns:

Coordinates of atom D, shape (3,).

Example

>>> d = extend_coordinate(a, b, c, 1.5, 2.0, 3.14)
>>> d.shape
(3,)
prxteinmpnn.utils.coordinates.compute_cb_precise(n_coord, ca_coord, c_coord)[source]#

Compute the C-beta atom position from backbone N, CA, and C coordinates.

Does so precisely using trigonometric relationships based on the backbone geometry.

Specifically, the position of the C-beta atom is determined by:

  • The bond length between the alpha carbon and the C-beta atom.

  • The bond angle between the nitrogen, alpha carbon, and C-beta atoms.

  • The dihedral angle involving the nitrogen, alpha carbon, and C-beta atoms.

Unlike the compute_c_beta function, this function does not use a linear combination of bond vectors with approximate fixed coefficients. This is more accurate and flexible for different configurations of the protein backbone, but more computationally intensive.

It is used in preparation of the atomic coordinates for the model input. It is not used in the model itself, but rather in the preprocessing of the input data to ensure that the C-beta atom is correctly placed based on the backbone structure.

Uses standard geometry for C-beta placement:

  • N-CA-CB bond length: 1.522 Å

  • N-CA-CB bond angle: 1.927 radians

  • C-N-CA-CB dihedral angle: -2.143 radians

Parameters:

  • n_coord (Float[Array, '3']) – Coordinates of the N atom, shape (3,).

  • ca_coord (Float[Array, '3']) – Coordinates of the CA atom, shape (3,).

  • c_coord (Float[Array, '3']) – Coordinates of the C atom, shape (3,).

Return type:

Float[Array, '3']

Returns:

Coordinates of the C-beta atom, shape (3,).

Example

>>> cb = compute_cb_precise(n, ca, c)
>>> cb.shape
(3,)
Contents