Coverage for cvpack/reporting/meta_cv_writer.py: 100%

1""" 

2.. class:: MetaCVWriter 

3 :platform: Linux, MacOS, Windows 

4 :synopsis: A custom writer for reporting meta-collective variable data 

5 

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

7 

8""" 

9 

10import typing as t 

11 

12import openmm as mm 

13from openmm import unit as mmunit 

14 

15from ..meta_collective_variable import MetaCollectiveVariable 

16from .custom_writer import CustomWriter 

17 

18 

19class MetaCVWriter(CustomWriter): 

20 """ 

21 A custom writer for reporting meta-collective variable data. 

22 

23 Parameters 

24 ---------- 

25 metaCV 

26 The meta-collective variable whose associated values will be reported. 

27 innerValues 

28 The names of the inner variables whose values will be reported. 

29 innerMasses 

30 The names of the inner variables whose effective masses will be reported. 

31 parameters 

32 The names of the parameters whose values will be reported. 

33 derivatives 

34 The names of the parameters with respect to which the derivatives of the 

35 meta-collective variable will be reported. 

36 

37 Examples 

38 -------- 

39 >>> import cvpack 

40 >>> import openmm 

41 >>> from math import pi 

42 >>> from cvpack.reporting import StateDataReporter, MetaCVWriter 

43 >>> from openmm import app, unit 

44 >>> from sys import stdout 

45 >>> from openmmtools import testsystems 

46 >>> model = testsystems.AlanineDipeptideVacuum() 

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

48 >>> psi = cvpack.Torsion(8, 14, 16, 18, name="psi") 

49 >>> umbrella = cvpack.MetaCollectiveVariable( 

50 ... f"0.5*kappa*(min(dphi,{2*pi}-dphi)^2+min(dpsi,{2*pi}-dpsi)^2)" 

51 ... "; dphi=abs(phi-phi0); dpsi=abs(psi-psi0)", 

52 ... [phi, psi], 

53 ... unit.kilojoules_per_mole, 

54 ... name="umbrella", 

55 ... kappa=100 * unit.kilojoules_per_mole/unit.radian**2, 

56 ... phi0=5*pi/6 * unit.radian, 

57 ... psi0=-5*pi/6 * unit.radian, 

58 ... ) 

59 >>> reporter = StateDataReporter( 

60 ... stdout, 

61 ... 100, 

62 ... writers=[ 

63 ... MetaCVWriter( 

64 ... umbrella, 

65 ... values=["phi", "psi"], 

66 ... emasses=["phi", "psi"], 

67 ... parameters=["phi0", "psi0"], 

68 ... derivatives=["phi0", "psi0"], 

69 ... ), 

70 ... ], 

71 ... step=True, 

72 ... ) 

73 >>> integrator = openmm.LangevinIntegrator( 

74 ... 300 * unit.kelvin, 

75 ... 1 / unit.picosecond, 

76 ... 2 * unit.femtosecond, 

77 ... ) 

78 >>> integrator.setRandomNumberSeed(1234) 

79 >>> umbrella.addToSystem(model.system) 

80 >>> simulation = app.Simulation(model.topology, model.system, integrator) 

81 >>> simulation.context.setPositions(model.positions) 

82 >>> simulation.context.setVelocitiesToTemperature(300 * unit.kelvin, 5678) 

83 >>> simulation.reporters.append(reporter) 

84 >>> simulation.step(1000) # doctest: +SKIP 

85 #"Step","phi (rad)",...,"d[umbrella]/d[psi0] (kJ/(mol rad))" 

86 100,2.36849...,40.3718... 

87 200,2.88515...,27.9109... 

88 300,2.43112...,-12.743... 

89 400,2.96786...,3.97688... 

90 500,2.58383...,41.8782... 

91 600,2.72482...,25.2626... 

92 700,2.55836...,25.3424... 

93 800,2.71046...,11.3498... 

94 900,2.43913...,37.3804... 

95 1000,2.7584...,31.1599... 

96 """ 

97 

98 def __init__( 

99 self, 

100 metaCV: MetaCollectiveVariable, 

101 values: t.Sequence[str] = (), 

102 emasses: t.Sequence[str] = (), 

103 parameters: t.Sequence[str] = (), 

104 derivatives: t.Sequence[str] = (), 

105 ) -> None: 

106 inner_cvs = {cv.getName(): cv for cv in metaCV.getInnerVariables()} 

107 all_parameters = metaCV.getParameterDefaultValues() 

108 self._meta_cv = metaCV 

109 self._values = [inner_cvs[name] for name in values] 

110 self._masses = [inner_cvs[name] for name in emasses] 

111 self._parameters = {name: all_parameters[name] for name in parameters} 

112 self._derivatives = {name: all_parameters[name] for name in derivatives} 

113 

114 def getHeaders(self) -> t.List[str]: 

115 headers = [] 

116 

117 def add_header(name: str, unit: mmunit.Unit) -> None: 

118 headers.append(f"{name} ({unit.get_symbol()})") 

119 

120 for cv in self._values: 

121 add_header(cv.getName(), cv.getUnit()) 

122 for cv in self._masses: 

123 add_header(f"emass[{cv.getName()}]", cv.getMassUnit()) 

124 for name, quantity in self._parameters.items(): 

125 add_header(name, quantity.unit) 

126 for name, quantity in self._derivatives.items(): 

127 add_header( 

128 f"d[{self._meta_cv.getName()}]/d[{name}]", 

129 self._meta_cv.getUnit() / quantity.unit, 

130 ) 

131 return headers 

132 

133 def getValues(self, simulation: mm.app.Simulation) -> t.List[float]: 

134 context = simulation.context 

135 values = [] 

136 if self._values: 

137 inner_values = self._meta_cv.getInnerValues(context) 

138 for cv in self._values: 

139 values.append(inner_values[cv.getName()] / cv.getUnit()) 

140 if self._masses: 

141 inner_masses = self._meta_cv.getInnerEffectiveMasses(context) 

142 for cv in self._masses: 

143 values.append(inner_masses[cv.getName()] / cv.getMassUnit()) 

144 if self._parameters: 

145 parameters = self._meta_cv.getParameterValues(context) 

146 for name, quantity in self._parameters.items(): 

147 values.append(parameters[name] / quantity.unit) 

148 if self._derivatives: 

149 derivatives = self._meta_cv.getParameterDerivatives(context) 

150 for name, quantity in self._derivatives.items(): 

151 values.append( 

152 derivatives[name] / (self._meta_cv.getUnit() / quantity.unit) 

153 ) 

154 return values