Coverage for src / molecular_simulations / analysis / sasa.py: 99%

1import MDAnalysis as mda 

2from MDAnalysis.analysis.base import AnalysisBase 

3from MDAnalysis.core import groups 

4from MDAnalysis.guesser.tables import vdwradii 

5import numpy as np 

6from scipy.spatial import KDTree 

7import warnings 

8 

9warnings.filterwarnings('ignore') 

10 

11class SASA(AnalysisBase): 

12 """ 

13 Computes solvent-accessible surface area using the Shrake-Rupley algorithm. 

14 Code is adapted from the biopython and mdtraj implementations as well as an 

15 unmerged PR from MDAnalysis. Implemented as an instance of AnalysisBase to  

16 support ease of deployment and baked-in parallelism. 

17 

18 Arguments: 

19 ag (mda.AtomGroup): AtomGroup for which to perform SASA. 

20 probe_radius (float): Defaults to 1.4Å. This is pretty standard and should not 

21 be altered in most cases. 

22 n_points (int): Defaults to 256. Number of points to place on each sphere, where 

23 a higher number is more costly but accurate. 

24 """ 

25 def __init__(self, 

26 ag: mda.AtomGroup, 

27 probe_radius: float=1.4, 

28 n_points: int=256, 

29 **kwargs): 

30 if isinstance(ag, groups.UpdatingAtomGroup): 

31 raise TypeError('UpdatingAtomGroups are not valid for SASA!') 

32 

33 super(SASA, self).__init__(ag.universe.trajectory, **kwargs) 

34 

35 if not hasattr(ag, 'elements'): 

36 raise ValueError('Cannot assign atomic radii:' 

37 'Universe has no `elements` property!') 

38 

39 self.ag = ag 

40 self.probe_radius = probe_radius 

41 self.n_points = n_points 

42 

43 self.radii_dict = dict() 

44 self.radii_dict.update(vdwradii) 

45 

46 self.radii = np.vectorize(self.radii_dict.get)(self.ag.elements) 

47 self.radii += self.probe_radius 

48 self.max_radii = 2 * np.max(self.radii) 

49 

50 self.sphere = self.get_sphere() 

51 

52 def get_sphere(self) -> np.ndarray: 

53 """ 

54 Returns a fibonacci unit sphere for a given number of points. 

55 

56 Returns: 

57 (np.ndarray): Array of points in fibonacci sphere. 

58 """ 

59 dl = np.pi * (3 - np.sqrt(5)) 

60 dz = 2. / self.n_points 

61 longitude = 0 

62 z = 1 - dz / 2 

63 

64 xyz = np.zeros((self.n_points, 3), dtype=np.float32) 

65 for i in range(self.n_points): 

66 r = np.sqrt(1 - z**2) 

67 xyz[i, :] = [np.cos(longitude) * r, np.sin(longitude) * r, z] 

68 

69 z -= dz 

70 longitude += dl 

71 

72 return xyz 

73 

74 def measure_sasa(self, 

75 ag: mda.AtomGroup) -> float: 

76 """ 

77 Measures the SASA of input atomgroup. 

78 

79 Arguments: 

80 ag (mda.AtomGroup): MDAnalysis AtomGroup. 

81 

82 Returns: 

83 (float): SASA value. 

84 """ 

85 kdt = KDTree(ag.positions, 10) 

86 

87 points = np.zeros(ag.n_atoms) 

88 for i in range(ag.n_atoms): 

89 sphere = self.sphere.copy() * self.radii[i] 

90 sphere += ag.positions[i] 

91 available = self.points_available.copy() 

92 kdt_sphere = KDTree(sphere, 10) 

93 

94 for j in kdt.query_ball_point(ag.positions[i], 

95 self.max_radii, 

96 workers=-1): 

97 if j == i: 

98 continue 

99 if self.radii[j] < (self.radii[i] + self.radii[j]): 

100 available -= { 

101 n for n in kdt_sphere.query_ball_point( 

102 self.ag.positions[j], 

103 self.radii[j] 

104 ) 

105 } 

106 

107 points[i] = len(available) 

108 

109 return 4 * np.pi * self.radii**2 * points / self.n_points 

110 

111 def _prepare(self): 

112 """ 

113 Preprocessing step that is run before analyzing the trajectory. 

114 Initializes the results array. 

115 

116 Returns: 

117 None 

118 """ 

119 self.results.sasa = np.zeros(self.ag.n_residues) 

120 

121 self.points_available = set(range(self.n_points)) 

122 

123 def _single_frame(self): 

124 """ 

125 Operations called on each frame. Measures SASA for each atom and sums 

126 them per-residue, placing the results into internal output array. 

127 

128 Returns: 

129 None 

130 """ 

131 area = self.measure_sasa(self.ag) 

132 result = np.zeros(self.ag.n_residues) 

133 for i, atom in enumerate(self.ag.atoms): 

134 result[atom.resid - 1] += area[i] 

135 

136 self.results.sasa += result 

137 

138 def _conclude(self): 

139 """ 

140 Post-processing step to average the SASA per-frame. 

141 

142 Returns: 

143 None 

144 """ 

145 if self.n_frames != 0: 

146 self.results.sasa /= self.n_frames 

147 

148class RelativeSASA(SASA): 

149 """ 

150 Computese Relative SASA for a specified atomgroup. Relative SASA is defined 

151 by the measured SASA divided by the maximum accessible surface area and for 

152 proteins we define this as within a tripeptide context so as to not overestimate 

153 SASA of the amide linkage and its neighbors. 

154 

155 Arguments: 

156 ag (mda.AtomGroup): AtomGroup for which to perform SASA. 

157 probe_radius (float): Defaults to 1.4Å. This is pretty standard and should not 

158 be altered in most cases. 

159 n_points (int): Defaults to 256. Number of points to place on each sphere, where 

160 a higher number is more costly but accurate. 

161 """ 

162 def __init__(self, 

163 ag: mda.AtomGroup, 

164 probe_radius: float=1.4, 

165 n_points: int=256, 

166 **kwargs): 

167 if not hasattr(ag, 'bonds'): 

168 raise ValueError('Universe has no `bonds` property!') 

169 super(RelativeSASA, self).__init__(ag, probe_radius, n_points, **kwargs) 

170 

171 def _prepare(self): 

172 """ 

173 Preprocessing step that is run before analyzing the trajectory. 

174 Initializes the results array. 

175 

176 Returns: 

177 None 

178 """ 

179 self.results.sasa = np.zeros(self.ag.n_residues) 

180 self.results.relative_area = np.zeros(self.ag.n_residues) 

181 

182 self.points_available = set(range(self.n_points)) 

183 

184 def _single_frame(self): 

185 """ 

186 Operations called on each frame. Measures SASA for each atom and sums 

187 them per-residue. Then computes the exposed area of each tripeptide  

188 and stores the ratio of SASA / exposed area in output array. 

189 

190 Returns: 

191 None 

192 """ 

193 area = self.measure_sasa(self.ag) 

194 result = np.zeros(self.ag.n_residues) 

195 for i, atom in enumerate(self.ag.atoms): 

196 result[atom.resid - 1] += area[i] 

197 

198 self.results.sasa += result 

199 

200 for res_index in self.ag.residues.resindices: 

201 tri_peptide = self.ag.select_atoms(f'byres (bonded resindex {res_index})') 

202 

203 if len(tri_peptide) == 0: 

204 continue 

205 

206 tri_pep_area = self.measure_sasa(tri_peptide) 

207 exposed_area = sum([ 

208 a for a, _id in zip(tri_pep_area, tri_peptide.resindices) 

209 if _id == res_index 

210 ]) 

211 

212 if exposed_area != 0.: 

213 result[res_index] /= exposed_area 

214 

215 self.results.relative_area += np.array([result[_id] for _id in self.ag.residues.resindices]) 

216 

217 def _conclude(self): 

218 """ 

219 Post-processing step to average the SASA per-frame. 

220 

221 Returns: 

222 None 

223 """ 

224 if self.n_frames != 0: 

225 self.results.sasa /= self.n_frames 

226 self.results.relative_area /= self.n_frames