Skip to content

geometry #

Common functions for computing the internal coordinates (e.g. bond lengths).

Functions:

  • compute_bond_vectors –

    Computes the vectors between each atom pair specified by the idxs as

  • compute_distances –

    Computes the distances between each pair of coordinates.

  • compute_angles –

    Computes the angles [rad] between each specified triplet of indices.

  • compute_dihedrals –

    Computes the angles [rad] between each specified quartet of indices.

compute_bond_vectors #

compute_bond_vectors(
    coords: ndarray, idxs: ndarray
) -> tuple[ndarray, ndarray]

Computes the vectors between each atom pair specified by the idxs as well as their norms.

Parameters:

  • coords (ndarray) –

    The coordinates with shape=(n_coords, 3).

  • idxs (ndarray) –

    The indices of the coordinates to compute the distances between with shape=(n_pairs, 2).

Returns:

  • tuple[ndarray, ndarray] –

    A tuple of the vectors with shape=shape=(n_pairs, 3) and norms with shape=(n_pairs,).

Source code in femto/md/utils/geometry.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
def compute_bond_vectors(
    coords: numpy.ndarray, idxs: numpy.ndarray
) -> tuple[numpy.ndarray, numpy.ndarray]:
    """Computes the vectors between each atom pair specified by the ``idxs`` as
    well as their norms.

    Args:
        coords: The coordinates with ``shape=(n_coords, 3)``.
        idxs: The indices of the coordinates to compute the distances between with
            ``shape=(n_pairs, 2)``.

    Returns:
        A tuple of the vectors with shape=``shape=(n_pairs, 3)`` and norms with
        ``shape=(n_pairs,)``.
    """

    if len(idxs) == 0:
        return numpy.ndarray([]), numpy.ndarray([])

    directions = coords[idxs[:, 1]] - coords[idxs[:, 0]]
    distances = numpy.linalg.norm(directions, axis=1)

    return directions, distances

compute_distances #

compute_distances(
    coords: ndarray, idxs: ndarray
) -> ndarray

Computes the distances between each pair of coordinates.

Parameters:

  • coords (ndarray) –

    The coordinates with shape=(n_coords, 3).

  • idxs (ndarray) –

    The indices of the coordinates to compute the distances between with shape=(n_pairs, 2).

Returns:

  • ndarray –

    The distances with shape=(n_pairs,).

Source code in femto/md/utils/geometry.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
def compute_distances(coords: numpy.ndarray, idxs: numpy.ndarray) -> numpy.ndarray:
    """Computes the distances between each pair of coordinates.

    Args:
        coords: The coordinates with ``shape=(n_coords, 3)``.
        idxs: The indices of the coordinates to compute the distances between with
            ``shape=(n_pairs, 2)``.

    Returns:
        The distances with ``shape=(n_pairs,)``.
    """

    return compute_bond_vectors(coords, idxs)[1]

compute_angles #

compute_angles(coords: ndarray, idxs: ndarray) -> ndarray

Computes the angles [rad] between each specified triplet of indices.

Parameters:

  • coords (ndarray) –

    The coordinates with shape=(n_coords, 3) or shape=(n_frames, n_coords, 3).

  • idxs (ndarray) –

    The indices of the coordinates to compute the angles between with shape=(n_pairs, 3).

Returns:

  • ndarray –

    The angles with shape=(n_pairs,) or shape=(n_frames, n_pairs).

Source code in femto/md/utils/geometry.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def compute_angles(coords: numpy.ndarray, idxs: numpy.ndarray) -> numpy.ndarray:
    """Computes the angles [rad] between each specified triplet of indices.

    Args:
        coords: The coordinates with ``shape=(n_coords, 3)`` or
            ``shape=(n_frames, n_coords, 3)``.
        idxs: The indices of the coordinates to compute the angles between with
            ``shape=(n_pairs, 3)``.

    Returns:
        The angles with ``shape=(n_pairs,)`` or ``shape=(n_frames, n_pairs)``.
    """

    if len(idxs) == 0:
        return numpy.ndarray([])

    is_batched = coords.ndim == 3

    if not is_batched:
        coords = coords[None, :, :]

    vector_ab = coords[:, idxs[:, 1]] - coords[:, idxs[:, 0]]
    vector_ac = coords[:, idxs[:, 1]] - coords[:, idxs[:, 2]]

    # tan theta = sin theta / cos theta
    #
    # ||a x b|| = ||a|| ||b|| sin theta
    #   a . b   = ||a|| ||b|| cos theta
    #
    # => tan theta = (a x b) / (a . b)
    angles = numpy.arctan2(
        numpy.linalg.norm(numpy.cross(vector_ab, vector_ac, axis=-1), axis=-1),
        (vector_ab * vector_ac).sum(axis=-1),
    )

    if not is_batched:
        angles = angles[0]

    return angles

compute_dihedrals #

compute_dihedrals(
    coords: ndarray, idxs: ndarray
) -> ndarray

Computes the angles [rad] between each specified quartet of indices.

Parameters:

  • coords (ndarray) –

    The coordinates with shape=(n_coords, 3) or shape=(n_frames, n_coords, 3).

  • idxs (ndarray) –

    The indices of the coordinates to compute the dihedrals between with shape=(n_pairs, 4).

Returns:

  • ndarray –

    The dihedrals with shape=(n_pairs,) or shape=(n_frames, n_pairs).

Source code in femto/md/utils/geometry.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def compute_dihedrals(coords: numpy.ndarray, idxs: numpy.ndarray) -> numpy.ndarray:
    """Computes the angles [rad] between each specified quartet of indices.

    Args:
        coords: The coordinates with ``shape=(n_coords, 3)`` or
            ``shape=(n_frames, n_coords, 3)``.
        idxs: The indices of the coordinates to compute the dihedrals between with
            ``shape=(n_pairs, 4)``.

    Returns:
        The dihedrals with ``shape=(n_pairs,)`` or ``shape=(n_frames, n_pairs)``.
    """

    if len(idxs) == 0:
        return numpy.ndarray([])

    is_batched = coords.ndim == 3

    if not is_batched:
        coords = coords[None, :, :]

    vector_ab = coords[:, idxs[:, 0]] - coords[:, idxs[:, 1]]
    vector_cb = coords[:, idxs[:, 2]] - coords[:, idxs[:, 1]]
    vector_cd = coords[:, idxs[:, 2]] - coords[:, idxs[:, 3]]

    vector_ab_cross_cb = numpy.cross(vector_ab, vector_cb, axis=-1)
    vector_cb_cross_cd = numpy.cross(vector_cb, vector_cd, axis=-1)

    vector_cb_norm = numpy.linalg.norm(vector_cb, axis=-1)[:, :, None]

    y = (
        numpy.cross(vector_ab_cross_cb, vector_cb_cross_cd, axis=-1)
        * vector_cb
        / vector_cb_norm
    ).sum(axis=-1)

    x = (vector_ab_cross_cb * vector_cb_cross_cd).sum(axis=-1)

    phi = numpy.arctan2(y, x)

    if not is_batched:
        phi = phi[0]

    return phi