Coverage for ufedmm/tests/test_ufedmm.py: 100%

1""" 

2Unit and regression test for the ufedmm package. 

3""" 

4 

5# Import package, test suite, and other packages as needed 

6import io 

7import sys 

8from copy import deepcopy 

9 

10import numpy as np 

11import openmm 

12import pytest 

13from openmm import unit 

14 

15import ufedmm 

16 

17 

18def ufed_model( 

19 temp=300 * unit.kelvin, 

20 mass=50 * unit.dalton * (unit.nanometer / unit.radians) ** 2, 

21 Ks=1000 * unit.kilojoules_per_mole / unit.radians**2, 

22 Ts=1500 * unit.kelvin, 

23 limit=180 * unit.degrees, 

24 sigma=18 * unit.degrees, 

25 height=0.0, 

26 frequency=10, 

27 bias_factor=None, 

28 enforce_gridless=False, 

29 constraints=openmm.app.HBonds, 

30): 

31 model = ufedmm.AlanineDipeptideModel(constraints=constraints) 

32 s_phi = ufedmm.DynamicalVariable( 

33 "s_phi", -limit, limit, mass, Ts, model.phi, Ks, sigma=sigma 

34 ) 

35 s_psi = ufedmm.DynamicalVariable( 

36 "s_psi", -limit, limit, mass, Ts, model.psi, Ks, sigma=sigma 

37 ) 

38 return model, ufedmm.UnifiedFreeEnergyDynamics( 

39 [s_phi, s_psi], 

40 temp, 

41 height, 

42 frequency, 

43 bias_factor, 

44 enforce_gridless=enforce_gridless, 

45 ) 

46 

47 

48def test_ufedmm_imported(): 

49 """Sample test, will always pass so long as import statement worked""" 

50 assert "ufedmm" in sys.modules 

51 

52 

53def test_serialization(): 

54 model, old = ufed_model() 

55 pipe = io.StringIO() 

56 ufedmm.serialize(old, pipe) 

57 pipe.seek(0) 

58 new = ufedmm.deserialize(pipe) 

59 assert new.__repr__() == old.__repr__() 

60 for var1, var2 in zip(old.variables, new.variables): 

61 assert var1.__repr__() == var2.__repr__() 

62 

63 

64def test_variables(): 

65 model, ufed = ufed_model(constraints=None) 

66 integrator = ufedmm.MiddleMassiveNHCIntegrator( 

67 300 * unit.kelvin, 10 * unit.femtoseconds, 1 * unit.femtoseconds 

68 ) 

69 platform = openmm.Platform.getPlatformByName("Reference") 

70 simulation = ufed.simulation(model.topology, model.system, integrator, platform) 

71 simulation.context.setPositions(model.positions) 

72 simulation.context.setVelocitiesToTemperature(300 * unit.kelvin, 11234) 

73 simulation.step(20) 

74 cvs = [ 

75 value 

76 for force in simulation.context.driving_forces 

77 for value in force.getCollectiveVariableValues(simulation.context) 

78 ] 

79 xvars = simulation.context.getState(getPositions=True).getDynamicalVariables() 

80 assert cvs[0] == pytest.approx(xvars[0]) 

81 assert cvs[2] == pytest.approx(xvars[1]) 

82 state = simulation.context.getState(getPositions=True) 

83 print(type(state.getPositions(asNumpy=False))) 

84 simulation.context.setPositions(*state.getPositions(extended=True)) 

85 xvars = simulation.context.getState(getPositions=True).getDynamicalVariables() 

86 assert cvs[0] == pytest.approx(xvars[0]) 

87 assert cvs[2] == pytest.approx(xvars[1]) 

88 

89 

90def test_NHC_integrator(): 

91 model, ufed = ufed_model(constraints=None) 

92 nonbonded = next( 

93 filter(lambda f: isinstance(f, openmm.NonbondedForce), model.system.getForces()) 

94 ) 

95 nonbonded.setForceGroup(1) 

96 integrator = ufedmm.MiddleMassiveNHCIntegrator( 

97 300 * unit.kelvin, 

98 10 * unit.femtoseconds, 

99 1 * unit.femtoseconds, 

100 respa_loops=[5, 1], 

101 ) 

102 platform = openmm.Platform.getPlatformByName("Reference") 

103 simulation = ufed.simulation(model.topology, model.system, integrator, platform) 

104 simulation.context.setPositions(model.positions) 

105 simulation.context.setVelocitiesToTemperature(300 * unit.kelvin, 1234) 

106 simulation.step(100) 

107 energy = simulation.context.getState(getEnergy=True).getPotentialEnergy() 

108 assert energy / energy.unit == pytest.approx(6.18, abs=0.2) 

109 

110 

111def test_GGMT_integrator(): 

112 model, ufed = ufed_model(constraints=None) 

113 nonbonded = next( 

114 filter(lambda f: isinstance(f, openmm.NonbondedForce), model.system.getForces()) 

115 ) 

116 nonbonded.setForceGroup(1) 

117 integrator = ufedmm.MiddleMassiveGGMTIntegrator( 

118 300 * unit.kelvin, 

119 10 * unit.femtoseconds, 

120 1 * unit.femtoseconds, 

121 respa_loops=[5, 1], 

122 ) 

123 platform = openmm.Platform.getPlatformByName("Reference") 

124 simulation = ufed.simulation(model.topology, model.system, integrator, platform) 

125 simulation.context.setPositions(model.positions) 

126 simulation.context.setVelocitiesToTemperature(300 * unit.kelvin, 1234) 

127 simulation.step(100) 

128 energy = simulation.context.getState(getEnergy=True).getPotentialEnergy() 

129 assert energy / energy.unit == pytest.approx(-17.78, abs=0.2) 

130 

131 

132def test_gridded_metadynamics(): 

133 model, ufed = ufed_model(height=2.0 * unit.kilocalorie_per_mole, constraints=None) 

134 integrator = ufedmm.MiddleMassiveNHCIntegrator( 

135 300 * unit.kelvin, 10 * unit.femtoseconds, 1 * unit.femtoseconds 

136 ) 

137 platform = openmm.Platform.getPlatformByName("Reference") 

138 simulation = ufed.simulation(model.topology, model.system, integrator, platform) 

139 simulation.context.setPositions(model.positions) 

140 simulation.context.setVelocitiesToTemperature(300 * unit.kelvin, 1234) 

141 simulation.step(100) 

142 energy = simulation.context.getState(getEnergy=True).getPotentialEnergy() 

143 assert energy / energy.unit == pytest.approx(695.95, abs=0.1) 

144 

145 

146# def test_gridless_metadynamics(): 

147# model, ufed = ufed_model( 

148# height=2.0*unit.kilocalorie_per_mole, enforce_gridless=True 

149# ) 

150# integrator = ufedmm.MiddleMassiveNHCIntegrator( 

151# 300*unit.kelvin, 10*unit.femtoseconds, 1*unit.femtoseconds 

152# ) 

153# platform = openmm.Platform.getPlatformByName('Reference') 

154# simulation = ufed.simulation(model.topology, model.system, integrator, platform) 

155# simulation.context.setPositions(model.positions) 

156# simulation.context.setVelocitiesToTemperature(300*unit.kelvin, 1234) 

157# simulation.step(100) 

158# energy = simulation.context.getState(getEnergy=True).getPotentialEnergy() 

159# assert energy/energy.unit == pytest.approx(59.74040) 

160 

161 

162def test_well_tempered_metadynamics(): 

163 model, ufed = ufed_model( 

164 height=2.0 * unit.kilocalorie_per_mole, bias_factor=10, constraints=None 

165 ) 

166 integrator = ufedmm.MiddleMassiveNHCIntegrator( 

167 300 * unit.kelvin, 10 * unit.femtoseconds, 1 * unit.femtoseconds 

168 ) 

169 platform = openmm.Platform.getPlatformByName("Reference") 

170 simulation = ufed.simulation(model.topology, model.system, integrator, platform) 

171 simulation.context.setPositions(model.positions) 

172 simulation.context.setVelocitiesToTemperature(300 * unit.kelvin, 1234) 

173 simulation.step(100) 

174 energy = simulation.context.getState(getEnergy=True).getPotentialEnergy() 

175 assert energy / energy.unit == pytest.approx(153.17, abs=0.1) 

176 

177 

178def test_rmsd_forces(): 

179 model, ufed = ufed_model(constraints=None) 

180 rmsd = openmm.RMSDForce( # Level 1 

181 model.positions, np.arange(model.system.getNumParticles()) 

182 ) 

183 cvforce = openmm.CustomCVForce("rmsd + cvforce") 

184 cvforce.addCollectiveVariable("rmsd", deepcopy(rmsd)) # Level 2 

185 inner_cvforce = openmm.CustomCVForce("rmsd") 

186 inner_cvforce.addCollectiveVariable("rmsd", deepcopy(rmsd)) # Level 3 

187 cvforce.addCollectiveVariable("cvforce", inner_cvforce) 

188 model.system.addForce(rmsd) 

189 model.system.addForce(cvforce) 

190 integrator = ufedmm.MiddleMassiveNHCIntegrator( 

191 300 * unit.kelvin, 10 * unit.femtoseconds, 1 * unit.femtoseconds 

192 ) 

193 platform = openmm.Platform.getPlatformByName("Reference") 

194 simulation = ufed.simulation(model.topology, model.system, integrator, platform) 

195 simulation.context.setPositions(model.positions) 

196 simulation.context.setVelocitiesToTemperature(300 * unit.kelvin, 1234) 

197 simulation.step(1)