Coverage for src / molecular_simulations / analysis / interaction_energy.py: 58%

1from abc import ABC, abstractmethod 

2from copy import deepcopy 

3from openmm import * 

4from openmm.app import * 

5from openmm.unit import * 

6import MDAnalysis as mda 

7from MDAnalysis.analysis.distances import contact_matrix 

8import mdtraj as md 

9import numpy as np 

10import parmed as pmd 

11from pathlib import Path 

12from pdbfixer import PDBFixer 

13import pickle 

14import gc 

15from tqdm import tqdm 

16from typing import Dict, List, Tuple, Union 

17 

18PathLike = Union[Path, str] 

19 

20class InteractionEnergy(ABC): 

21 def __init__(self): 

22 pass 

23 

24 @abstractmethod 

25 def compute(self): 

26 pass 

27 

28 @abstractmethod 

29 def energy(self): 

30 pass 

31 

32 @abstractmethod 

33 def get_selection(self): 

34 pass 

35 

36class StaticInteractionEnergy(InteractionEnergy): 

37 """ 

38 Computes the linear interaction energy between specified chain and other simulation 

39 components. Can specify a range of residues in chain to limit calculation to. Works on 

40 a static model but can be adapted to run on dynamics data. 

41 

42 Inputs: 

43 pdb (str): Path to input PDB file 

44 chain (str): Defaults to A. The chain for which to compute the energy between. 

45 Computes energy between this chain and all other components in PDB file. 

46 Set to a whitespace if there are no chains in your PDB. 

47 platform (str): Defaults to CUDA. Supports running on GPU for speed. 

48 first_residue (int | None): Defaults to None. If set, will restrict the  

49 calculation to residues beginning with resid `first_residue`. 

50 last_residue (int | None): Defaults to None. If set, will restrict the  

51 calculation to residues ending with resid `last_residue`. 

52 """ 

53 def __init__(self, 

54 pdb: str, 

55 chain: str='A', 

56 platform: str='CUDA', 

57 first_residue: Union[int, None]=None, 

58 last_residue: Union[int, None]=None): 

59 self.pdb = pdb 

60 self.chain = chain 

61 self.platform = Platform.getPlatformByName(platform) 

62 self.first = first_residue 

63 self.last = last_residue 

64 

65 def get_system(self) -> System: 

66 """ 

67 Builds implicit solvent OpenMM system. 

68 

69 Returns: 

70 (System): OpenMM system object. 

71 """ 

72 pdb = PDBFile(self.pdb) 

73 positions, topology = pdb.positions, pdb.topology 

74 forcefield = ForceField('amber14-all.xml', 'implicit/gbn2.xml') 

75 try: 

76 system = forcefield.createSystem(topology, 

77 soluteDielectric=1., 

78 solventDielectric=80.) 

79 except ValueError: 

80 positions, topology = self.fix_pdb() 

81 system = forcefield.createSystem(topology, 

82 soluteDielectric=1., 

83 solventDielectric=80.) 

84 

85 self.positions = positions 

86 self.get_selection(topology) 

87 

88 return system 

89 

90 def compute(self, 

91 positions: Union[np.ndarray, None]=None) -> None: 

92 """ 

93 Compute interaction energy of system. Can optionally provide atomic 

94 positions such that this operation can be scaled onto a trajectory of 

95 frames rather than a static model. 

96 

97 Arguments: 

98 positions (np.ndarray | None): Defaults to None. If provided, inject 

99 the positions into the OpenMM context. 

100 

101 Returns: 

102 None 

103 """ 

104 self.lj = None 

105 self.coulomb = None 

106 

107 system = self.get_system() 

108 if positions is None: 

109 positions = self.positions 

110 

111 for force in system.getForces(): 

112 if isinstance(force, NonbondedForce): 

113 force.setForceGroup(0) 

114 force.addGlobalParameter("solute_coulomb_scale", 1) 

115 force.addGlobalParameter("solute_lj_scale", 1) 

116 force.addGlobalParameter("solvent_coulomb_scale", 1) 

117 force.addGlobalParameter("solvent_lj_scale", 1) 

118 

119 for i in range(force.getNumParticles()): 

120 charge, sigma, epsilon = force.getParticleParameters(i) 

121 force.setParticleParameters(i, 0, 0, 0) 

122 if i in self.selection: 

123 force.addParticleParameterOffset("solute_coulomb_scale", i, charge, 0, 0) 

124 force.addParticleParameterOffset("solute_lj_scale", i, 0, sigma, epsilon) 

125 else: 

126 force.addParticleParameterOffset("solvent_coulomb_scale", i, charge, 0, 0) 

127 force.addParticleParameterOffset("solvent_lj_scale", i, 0, sigma, epsilon) 

128 

129 for i in range(force.getNumExceptions()): 

130 p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(i) 

131 force.setExceptionParameters(i, p1, p2, 0, 0, 0) 

132 

133 else: 

134 force.setForceGroup(2) 

135 

136 integrator = VerletIntegrator(0.001*picosecond) 

137 

138 context = Context(system, integrator, self.platform) 

139 context.setPositions(positions) 

140 

141 total_coulomb = self.energy(context, 1, 0, 1, 0) 

142 solute_coulomb = self.energy(context, 1, 0, 0, 0) 

143 solvent_coulomb = self.energy(context, 0, 0, 1, 0) 

144 total_lj = self.energy(context, 0, 1, 0, 1) 

145 solute_lj = self.energy(context, 0, 1, 0, 0) 

146 solvent_lj = self.energy(context, 0, 0, 0, 1) 

147 

148 coul_final = total_coulomb - solute_coulomb - solvent_coulomb 

149 lj_final = total_lj - solute_lj - solvent_lj 

150 

151 self.coulomb = coul_final.value_in_unit(kilocalories_per_mole) 

152 self.lj = lj_final.value_in_unit(kilocalories_per_mole) 

153 

154 def get_selection(self, 

155 topology: Topology) -> None: 

156 """ 

157 Using the poorly documented OpenMM selection language, get indices of 

158 atoms that we want to isolate for pairwise interaction energy calculation. 

159 

160 Arguments: 

161 topology (Topology): OpenMM topology object. 

162 

163 Returns: 

164 None 

165 """ 

166 if self.first is None and self.last is None: 

167 selection = [a.index 

168 for a in topology.atoms() 

169 if a.residue.chain.id == self.chain] 

170 elif self.first is not None and self.last is None: 

171 selection = [a.index 

172 for a in topology.atoms() 

173 if a.residue.chain.id == self.chain 

174 and int(self.first) <= int(a.residue.id)] 

175 elif self.first is None: 

176 selection = [a.index 

177 for a in topology.atoms() 

178 if a.residue.chain.id == self.chain 

179 and int(self.last) >= int(a.residue.id)] 

180 else: 

181 selection = [a.index 

182 for a in topology.atoms() 

183 if a.residue.chain.id == self.chain 

184 and int(self.first) <= int(a.residue.id) <= int(self.last)] 

185 

186 self.selection = selection 

187 

188 def fix_pdb(self) -> None: 

189 """ 

190 Using the OpenMM adjacent tool, PDBFixer, repair input PDB by adding 

191 hydrogens, and missing atoms such that we can actually construct an 

192 OpenMM system. 

193 

194 Returns: 

195 None 

196 """ 

197 fixer = PDBFixer(filename=self.pdb) 

198 fixer.findMissingResidues() 

199 fixer.findMissingAtoms() 

200 fixer.addMissingAtoms() 

201 fixer.addMissingHydrogens(7.0) 

202 

203 return fixer.positions, fixer.topology 

204 

205 @property 

206 def interactions(self) -> np.ndarray: 

207 """ 

208 Places the LJ and coulombic energies into an array of shape (1, 2). 

209 

210 Returns: 

211 (np.ndarray): Energy array. 

212 """ 

213 return np.vstack([self.lj, self.coulomb]) 

214 

215 @staticmethod 

216 def energy(context: Context, 

217 solute_coulomb_scale: int=0, 

218 solute_lj_scale: int=0, 

219 solvent_coulomb_scale: int=0, 

220 solvent_lj_scale: int=0) -> float: 

221 """ 

222 Computes the potential energy for provided context object. 

223 

224 Arguments: 

225 context (Context): OpenMM context object. 

226 solute_coulomb_scale (int): Defaults to 0. If 1 we will consider solute 

227 contributions to coulombic non-bonded energy. 

228 solute_lj_scale (int): Defaults to 0. If 1 we will consider solute 

229 contributions to LJ non-bonded energy. 

230 solvent_coulomb_scale (int): Defaults to 0. If 1 we will consider solvent 

231 contributions to coulombic non-bonded energy. 

232 solvent_lj_scale (int): Defaults to 0. If 1 we will consider solvent 

233 contributions to LJ non-bonded energy. 

234 

235 Returns: 

236 (float): Computed energy term. 

237 """ 

238 context.setParameter("solute_coulomb_scale", solute_coulomb_scale) 

239 context.setParameter("solute_lj_scale", solute_lj_scale) 

240 context.setParameter("solvent_coulomb_scale", solvent_coulomb_scale) 

241 context.setParameter("solvent_lj_scale", solvent_lj_scale) 

242 return context.getState(getEnergy=True, groups={0}).getPotentialEnergy() 

243 

244class InteractionEnergyFrame(StaticInteractionEnergy): 

245 """ 

246 Inherits from StaticInteractionEnergy and overloads `get_system` to allow for 

247 more easily running this analysis on a trajectory of frames. Requires the 

248 OpenMM system and topology to be built externally and passed in rather than 

249 beginning from a PDB file. 

250 

251 Arguments: 

252 system (System): OpenMM system object. 

253 top (Topology): OpenMM topology object. 

254 chain (str): Defaults to A. The chain for which to compute the energy between. 

255 Computes energy between this chain and all other components in PDB file. 

256 Set to a whitespace if there are no chains in your PDB. 

257 platform (str): Defaults to CUDA. Supports running on GPU for speed. 

258 first_residue (int | None): Defaults to None. If set, will restrict the  

259 calculation to residues beginning with resid `first_residue`. 

260 last_residue (int | None): Defaults to None. If set, will restrict the  

261 calculation to residues ending with resid `last_residue`. 

262 """ 

263 def __init__(self, 

264 system: System, 

265 top: Topology, 

266 chain: str='A', 

267 platform: str='CUDA', 

268 first_residue: Union[int, None]=None, 

269 last_residue: Union[int, None]=None): 

270 super().__init__('', chain, platform, first_residue, last_residue) 

271 self.system = system 

272 self.top = top 

273 

274 def get_system(self) -> System: 

275 """ 

276 Sets self.selection via self.get_selection and returns existing OpenMM 

277 system object. 

278 

279 Returns: 

280 (System): OpenMM system object. 

281 """ 

282 self.get_selection(self.top) 

283 return self.system 

284 

285class DynamicInteractionEnergy: 

286 """ 

287 Class for obtaining interaction energies of a trajectory. Utilizes the  

288 InteractionEnergyFrame child class to run per-frame energy calculations  

289 and orchestrates the trajectory operations. 

290 

291 Arguments: 

292 top (PathLike): Path to prmtop topology file. 

293 traj (PathLike): Path to DCD trajectory file. 

294 stride (int): Defaults to 1. The stride with which to move through the trajectory. 

295 chain (str): Defaults to A. The chain for which to compute the energy between. 

296 Computes energy between this chain and all other components in PDB file. 

297 Set to a whitespace if there are no chains in your PDB. 

298 platform (str): Defaults to CUDA. Supports running on GPU for speed. 

299 first_residue (int | None): Defaults to None. If set, will restrict the  

300 calculation to residues beginning with resid `first_residue`. 

301 last_residue (int | None): Defaults to None. If set, will restrict the  

302 calculation to residues ending with resid `last_residue`. 

303 progress_bar (bool): Defaults to False. If True a tqdm progress bar will 

304 display progress. 

305 """ 

306 def __init__(self, 

307 top: PathLike, 

308 traj: PathLike, 

309 stride: int=1, 

310 chain: str='A', 

311 platform: str='CUDA', 

312 first_residue: Union[int, None]=None, 

313 last_residue: Union[int, None]=None, 

314 progress_bar: bool=False): 

315 top = Path(top) 

316 traj = Path(traj) 

317 self.system = self.build_system(top) 

318 self.coordinates = self.load_traj(top, traj) 

319 self.stride = stride 

320 self.progress = progress_bar 

321 

322 self.IE = InteractionEnergyFrame(self.system, self.top, chain, 

323 platform, first_residue, last_residue) 

324 

325 def compute_energies(self) -> None: 

326 """ 

327 Computes the energy for each frame in trajectory, storing internally. 

328 

329 Returns: 

330 None 

331 """ 

332 n_frames = self.coordinates.shape[0] // self.stride 

333 self.energies = np.zeros((n_frames, 2)) 

334 

335 if self.progress: 

336 pbar = tqdm(total=n_frames, position=0, leave=False) 

337 

338 for i in range(n_frames): 

339 fr = i * self.stride 

340 self.IE.compute(self.coordinates[fr, :, :]) 

341 self.energies[i, 0] = self.IE.lj 

342 self.energies[i, 1] = self.IE.coulomb 

343 

344 if self.progress: 

345 pbar.update(1) 

346 

347 if self.progress: 

348 pbar.close() 

349 

350 def build_system(self, 

351 top: PathLike) -> System: 

352 """ 

353 Builds OpenMM system for both pdb and prmtop topology files. 

354 

355 Arguments: 

356 top (PathLike): Path to topology file. 

357 

358 Returns: 

359 (System): OpenMM system object. 

360 """ 

361 if top.suffix == '.pdb': 

362 top = PDBFile(str(top)).topology 

363 self.top = top 

364 forcefield = ForceField('amber14-all.xml', 'implicit/gbn2.xml') 

365 return forcefield.createSystem(top, 

366 soluteDielectric=1., 

367 solventDielectric=78.5) 

368 elif top.suffix == '.prmtop': 

369 top = AmberPrmtopFile(str(top)) 

370 self.top = top 

371 return top.createSystem(nonbondedMethod=CutoffNonPeriodic, 

372 nonbondedCutoff=2. * nanometers, 

373 constraints=HBonds) 

374 else: 

375 raise NotImplementedError(f'Error! Topology type {top} not implemented!') 

376 

377 def load_traj(self, 

378 top: PathLike, 

379 traj: PathLike) -> np.ndarray: 

380 """ 

381 Loads trajectory into mdtraj and extracts full coordinate array. 

382 

383 Arguments: 

384 top (PathLike): Path to topology file. 

385 traj (PathLike): Path to trajectory file. 

386 

387 Returns: 

388 (np.ndarray): Coordinate array of shape (n_frames, n_atoms, 3) 

389 """ 

390 return md.load(str(traj), top=str(top)).xyz 

391 

392 def setup_pbar(self) -> None: 

393 """ 

394 Builds tqdm progress bar. 

395 

396 Returns: 

397 None 

398 """ 

399 self.pbar = tqdm(total=self.coordinates.shape[0], position=0, leave=False)