-
Notifications
You must be signed in to change notification settings - Fork 14
/
align.py
194 lines (158 loc) · 5.85 KB
/
align.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
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
85
86
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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
import logging
import numpy as np
import parmed as pmd
logger = logging.getLogger(__name__)
def zalign(structure, mask1, mask2, save=False, filename=None):
"""Align the mask1 -- mask2 vector to the z axis.
Parameters
----------
structure : parmed.Structure
Molecular structure containing coordinates
mask1 : str
Selection of first set of atoms
mask2 : str
Selection of second set of atoms
save : bool, optional
Whether to save the coordinates (the default is False, which does nothing)
filename : str, optional
The filename for the saved coordinates (the default is None, which does nothing)
Returns
-------
parmed.Structure
A molecular structure with the coordinates aligned as specified.
"""
mask1_coordinates = structure[mask1].coordinates
mask1_masses = [atom.mass for atom in structure[mask1].atoms]
mask1_com = pmd.geometry.center_of_mass(
np.asarray(mask1_coordinates), np.asarray(mask1_masses)
)
mask2_coordinates = structure[mask2].coordinates
mask2_masses = [atom.mass for atom in structure[mask2].atoms]
mask2_com = pmd.geometry.center_of_mass(
np.asarray(mask2_coordinates), np.asarray(mask2_masses)
)
logger.info(
"Moving {} ({} atoms) to the origin...".format(mask1, len(mask1_coordinates))
)
logger.info(
"Aligning {} ({} atoms) with the z axis...".format(
mask2, len(mask2_coordinates)
)
)
axis = np.array([0.0, 0.0, 1.0])
identity = np.identity(3)
# https://math.stackexchange.com/questions/293116/rotating-one-3d-vector-to-another
# 1. Define the vector from mask1 to mask2.
mask2_com = mask2_com + -1.0 * mask1_com
# 2. Find axis and angle between the mask vector and the axis using cross
# and dot products.
try:
x = np.cross(mask2_com, axis) / np.linalg.norm(np.cross(mask2_com, axis))
except RuntimeWarning:
# The structure is already aligned and the denominator is invalid
pass
theta = np.arccos(
np.dot(mask2_com, axis) / (np.linalg.norm(mask2_com) * np.linalg.norm(axis))
)
# 3. Find the rotation matrix
A = np.array(
[[0, -1.0 * x[2], x[1]], [x[2], 0, -1.0 * x[0]], [-1.0 * x[1], x[0], 0]]
)
rotation_matrix = (
identity
+ np.dot(np.sin(theta), A)
+ np.dot((1.0 - np.cos(theta)), np.dot(A, A))
)
# This is certainly not the fastest approach, but it is explicit.
aligned_coords = np.empty_like(structure.coordinates)
for atom in range(len(structure.atoms)):
aligned_coords[atom] = structure.coordinates[atom] + -1.0 * mask1_com
aligned_coords[atom] = np.dot(rotation_matrix, aligned_coords[atom])
structure.coordinates = aligned_coords
if save:
if not filename:
logger.warning(
"Unable to save aligned coordinates (no filename provided)..."
)
else:
logger.info("Saved aligned coordinates to {}".format(filename))
# This seems to write out HETATM in place of ATOM
# We should offer the option of writing a mol2 file, directly.
structure.write_pdb(filename)
return structure
def get_theta(structure, mask1, mask2, axis):
"""Get the angle (theta) between the vector formed by two masks and an axis.
Parameters
----------
structure : parmed.Structure
Molecular structure containing coordinates
mask1 : str
Selection of first set of atoms
mask2 : str
Selection of second set of atoms
axis : str
Axis: x, y, or z
Returns
-------
float
The angle between the masks and the axis.
"""
if "x" in axis.lower():
axis = np.array([1.0, 0.0, 0.0])
elif "y" in axis.lower():
axis = np.array([0.0, 1.0, 0.0])
elif "z" in axis.lower():
axis = np.array([0.0, 0.0, 1.0])
mask1_coordinates = structure[mask1].coordinates
mask1_masses = [atom.mass for atom in structure[mask1].atoms]
mask1_com = pmd.geometry.center_of_mass(
np.asarray(mask1_coordinates), np.asarray(mask1_masses)
)
mask2_coordinates = structure[mask2].coordinates
mask2_masses = [atom.mass for atom in structure[mask2].atoms]
mask2_com = pmd.geometry.center_of_mass(
np.asarray(mask2_coordinates), np.asarray(mask2_masses)
)
vector = mask2_com + -1.0 * mask1_com
theta = np.arccos(
np.dot(vector, axis) / (np.linalg.norm(vector) * np.linalg.norm(axis))
)
return theta
def check_coordinates(structure, mask):
"""Return the coordinates of an atom selection.
Parameters
----------
structure : parmed.Structure
Molecular structure containing coordinates
mask : str
Atom selection
Returns
-------
np.array
Coordinates of the selection center of mass
"""
mask_coordinates = structure[mask].coordinates
mask_masses = [atom.mass for atom in structure[mask].atoms]
mask_com = pmd.geometry.center_of_mass(
np.asarray(mask_coordinates), np.asarray(mask_masses)
)
return mask_com
def offset_structure(structure, offset):
"""Return a structure whose coordinates have been offset.
Parameters
----------
structure : parmed.Structure
Molecular structure containing coordinates
offset : float
The offset that will be added to *every* atom in the structure
Returns
-------
:py:class:`parmed.Structure`
Coordinates of the structure offset by the given amount.
"""
offset_coords = np.empty_like(structure.coordinates)
for atom in range(len(structure.atoms)):
offset_coords[atom] = structure.coordinates[atom] + offset
structure.coordinates = offset_coords
logger.info("Added offset of {} to atomic coordinates...".format(offset))
return structure