Coverage for openxps/metadynamics.py: 88%

1""" 

2.. module:: openxps.metadynamics 

3 :platform: Linux, Windows, macOS 

4 :synopsis: Bias potentials applied to dynamical variables. 

5 

6.. moduleauthor:: Charlles Abreu <craabreu@gmail.com> 

7 

8""" 

9 

10import typing as t 

11from dataclasses import dataclass, field 

12 

13from cvpack.serialization import Serializable 

14from openmm import app as mmapp 

15from openmm import unit as mmunit 

16 

17from .bounds import NoBounds 

18from .dynamical_variable import DynamicalVariable 

19from .system import ExtendedSpaceSystem 

20from .utils import preprocess_args 

21 

22 

23class _SimulationWrapper: 

24 def __init__(self, simulation: mmapp.Simulation): 

25 self._simulation = simulation 

26 self.context = simulation.context.getExtensionContext() 

27 self.step = simulation.step 

28 

29 @property 

30 def currentStep(self): 

31 return self._simulation.currentStep 

32 

33 

34@dataclass(frozen=True) 

35class ExtendedSpaceBiasVariable(Serializable): 

36 """ 

37 A dynamical variable augmented with a Gaussian bias kernel. 

38 

39 Parameters 

40 ---------- 

41 dynamical_variable 

42 The dynamical variable to be augmented. 

43 sigma 

44 The bandwidth (standard deviation) of the Gaussian bias kernel. 

45 grid_width 

46 The width of the grid for the bias kernel. If ``None``, the grid width is 

47 automatically determined based on the dynamical variable's bounds and the 

48 bandwidth. 

49 

50 Example 

51 ------- 

52 >>> import openxps as xps 

53 >>> from openmm import unit 

54 >>> dv = xps.DynamicalVariable( 

55 ... "phi", 

56 ... unit.radian, 

57 ... 3 * unit.dalton*(unit.nanometer/unit.radian)**2, 

58 ... xps.PeriodicBounds(-180, 180, unit.degree) 

59 ... ) 

60 >>> bias_variable = xps.ExtendedSpaceBiasVariable(dv, 18 * unit.degree) 

61 >>> bias_variable.dynamical_variable 

62 DynamicalVariable(name='phi', unit=rad, mass=3 nm**2 Da/(rad**2), bounds=...) 

63 """ 

64 

65 dynamical_variable: DynamicalVariable 

66 sigma: mmunit.Quantity 

67 grid_width: t.Optional[int] = field(default_factory=lambda: None) 

68 

69 def __post_init__(self) -> None: 

70 if not mmunit.is_quantity(self.sigma): 

71 raise TypeError("sigma must be a quantity") 

72 if not self.sigma.unit.is_compatible(self.dynamical_variable.unit): 

73 raise ValueError( 

74 "sigma must be compatible with the dynamical variable's unit" 

75 ) 

76 if not (self.grid_width is None or isinstance(self.grid_width, int)): 

77 raise TypeError("grid_width must be an integer or None") 

78 

79 

80ExtendedSpaceBiasVariable.__init__ = preprocess_args(ExtendedSpaceBiasVariable.__init__) 

81 

82ExtendedSpaceBiasVariable.registerTag("!openxps.metadynamics.ExtendedSpaceBiasVariable") 

83 

84 

85class ExtendedSpaceMetadynamics(mmapp.Metadynamics): 

86 """ 

87 Performs Extended Phase-Space (XPS) Metadynamics simulations, in which a bias 

88 potential is applied to dynamical variables (DVs) to enhance sampling of the 

89 physical coordinates. 

90 

91 Parameters 

92 ---------- 

93 system 

94 The :class:`ExtendedSpaceSystem` to be used in the XPS simulation. 

95 variables 

96 A sequence of :class:`ExtendedSpaceBiasVariable` objects to specify the biased 

97 dynamical variables and their corresponding Gaussian kernels. 

98 temperature 

99 The temperature at which the simulation is being run. This is used in computing 

100 the free energy. 

101 biasFactor 

102 Used in scaling the height of the Gaussians added to the bias. The dynamical 

103 variables are sampled as if the effective temperature of the simulation were 

104 temperature*biasFactor. 

105 height 

106 The initial height of the Gaussians to add (in units of energy). 

107 frequency 

108 The interval in time steps at which Gaussians should be added to the bias 

109 potential. 

110 saveFrequency 

111 The interval in time steps at which to write out the current biases to disk. At 

112 the same time it writes biases, it also checks for updated biases written by 

113 other processes and loads them in. This must be a multiple of frequency. 

114 biasDir 

115 The directory to which biases should be written, and from which biases written 

116 by other processes should be loaded. 

117 

118 Example 

119 ------- 

120 >>> import openxps as xps 

121 >>> from openmm import unit 

122 >>> import cvpack 

123 >>> from math import pi 

124 >>> import openmm 

125 >>> from openmmtools import testsystems 

126 >>> model = testsystems.AlanineDipeptideVacuum() 

127 >>> phi = cvpack.Torsion(6, 8, 14, 16, name="phi") 

128 >>> mass = 3 * unit.dalton*(unit.nanometer/unit.radian)**2 

129 >>> phi0 = xps.DynamicalVariable("phi0", unit.radian, mass, xps.CircularBounds()) 

130 >>> phi = cvpack.Torsion(6, 8, 14, 16, name="phi") 

131 >>> kappa = 1000 * unit.kilojoules_per_mole / unit.radian**2 

132 >>> harmonic_force = xps.HarmonicCoupling(phi, phi0, kappa) 

133 >>> system = xps.ExtendedSpaceSystem(model.system, harmonic_force) 

134 >>> bias_variable = xps.ExtendedSpaceBiasVariable(phi0, 18 * unit.degrees) 

135 >>> temperature = 300 * unit.kelvin 

136 >>> metadynamics = xps.ExtendedSpaceMetadynamics( 

137 ... system=system, 

138 ... variables=[bias_variable], 

139 ... temperature=temperature, 

140 ... biasFactor=5, 

141 ... height=2 * unit.kilojoule_per_mole, 

142 ... frequency=100, 

143 ... ) 

144 >>> simulation = xps.ExtendedSpaceSimulation( 

145 ... model.topology, 

146 ... system, 

147 ... xps.LockstepIntegrator( 

148 ... openmm.LangevinMiddleIntegrator( 

149 ... temperature, 1 / unit.picosecond, 4 * unit.femtosecond 

150 ... ) 

151 ... ), 

152 ... openmm.Platform.getPlatformByName("Reference"), 

153 ... ) 

154 >>> context = simulation.context 

155 >>> context.setPositions(model.positions) 

156 >>> context.setVelocitiesToTemperature(temperature, 1234) 

157 >>> context.setDynamicalVariableValues([180 * unit.degree]) 

158 >>> context.setDynamicalVariableVelocitiesToTemperature(temperature, 1234) 

159 >>> metadynamics.step(simulation, 100) 

160 >>> metadynamics.getFreeEnergy() 

161 [... ... ...] kJ/mol 

162 """ 

163 

164 def __init__( # noqa: PLR0913 

165 self, 

166 system: ExtendedSpaceSystem, 

167 variables: t.Sequence[ExtendedSpaceBiasVariable], 

168 temperature: mmunit.Quantity, 

169 biasFactor: float, 

170 height: mmunit.Quantity, 

171 frequency: int, 

172 saveFrequency: t.Optional[int] = None, 

173 biasDir: t.Optional[str] = None, 

174 ) -> None: 

175 system_dvs = system.getDynamicalVariables() 

176 bias_variables = [] 

177 for variable in variables: 

178 dv = variable.dynamical_variable.in_md_units() 

179 if dv not in system_dvs: 

180 raise ValueError(f"Dynamical variable {dv.name} not found in system") 

181 if isinstance(dv.bounds, NoBounds): 

182 raise ValueError(f"Dynamical variable {dv.name} has no bounds") 

183 index = system_dvs.index(dv) 

184 bias_variables.append( 

185 mmapp.BiasVariable( 

186 dv.createCollectiveVariable(index), 

187 dv.bounds.lower * dv.bounds.unit, 

188 dv.bounds.upper * dv.bounds.unit, 

189 variable.sigma, 

190 dv.isPeriodic(), 

191 variable.grid_width, 

192 ) 

193 ) 

194 super().__init__( 

195 system.getExtensionSystem(), 

196 bias_variables, 

197 temperature, 

198 biasFactor, 

199 height, 

200 frequency, 

201 saveFrequency, 

202 biasDir, 

203 ) 

204 

205 def step(self, simulation, steps): 

206 super().step(_SimulationWrapper(simulation), steps) 

207 

208 def getCollectiveVariables(self, simulation): 

209 return self._force.getCollectiveVariableValues( 

210 simulation.context.getExtensionContext() 

211 )