5 from __future__
import print_function, division
7 from numpy.random
import random, randint
8 kB = 1.3806503 * 6.0221415 / 4184.0
13 def __init__(self, nreps, inv_temps, grid, sfo_id,
14 rexlog=
'replicanums.txt', scheme=
'standard', xchg=
'random',
15 convectivelog=
'stirred.txt', tune_temps=
False,
16 tune_data={}, templog=
'temps.txt'):
19 self.replicanums = list(range(nreps))
21 self.statenums = list(range(nreps))
25 self.inv_temps = inv_temps
30 self.tune_temps = tune_temps
31 self.tune_data = tune_data
33 self.templog = templog
35 if scheme ==
"convective":
38 self.stirred[
'order'] = list(range(self.nreps))
39 self.stirred[
'order'].reverse()
41 self.stirred[
'replica'] = self.stirred[
'order'][0]
43 self.stirred[
'pos'] = 0
45 if self.stirred[
'replica'] != self.nreps - 1:
46 self.stirred[
'dir'] = 1
48 self.stirred[
'dir'] = 0
50 self.stirred[
'steps'] = 2 * (self.nreps - 1)
51 self.convectivelog = convectivelog
52 self.write_rex_stats()
54 def sort_per_state(self, inplist):
55 "sorts a replica list per state"
56 if len(inplist) != self.nreps:
57 raise ValueError(
'list has wrong size')
58 return [inplist[i]
for i
in self.replicanums]
60 def sort_per_replica(self, inplist):
61 "sorts a state list per replica number"
62 if len(inplist) != self.nreps:
63 raise ValueError(
'list has wrong size')
64 return [inplist[i]
for i
in self.statenums]
66 def get_energies(self):
67 "return replica-sorted energies"
68 return self.grid.gather(
69 self.grid.broadcast(self.sfo_id,
'm',
'evaluate',
False))
71 def gen_pairs_list_gromacs(self, direction):
72 """generate ordered list of exchange pairs based on direction.
73 direction == 0 : (0,1)(2,3)...
74 direction == 1 : (0(1,2)(3,4)...
78 if direction != 0
and direction != 1:
79 raise ValueError(direction)
82 ret = [(2 * i + direction, 2 * i + 1 + direction)
83 for i
in range(nreps // 2)]
88 def gen_pairs_list_rand(self, needed=[]):
89 "generate list of neighboring pairs of states"
93 init = [(i, i + 1)
for i
in range(nreps - 1)]
95 init.append((0, nreps - 1))
99 raise ValueError(
"wrong format for 'needed' list")
100 pairslist.append((i, i + 1))
101 init.remove((i - 1, i))
102 init.remove((i, i + 1))
103 init.remove((i + 1, i + 2))
106 i = randint(0, len(init))
111 init = [(r, q)
for (r, q)
in init
112 if (r
not in pair
and q
not in pair)]
115 if not pair == (0, nreps - 1):
116 pairslist.append(pair)
121 def gen_pairs_list_conv(self):
122 rep = self.stirred[
'replica']
123 state = self.statenums[rep]
124 pair = sorted([state, state + 2 * self.stirred[
'dir'] - 1])
125 self.stirred[
'pair'] = tuple(pair)
126 if self.xchg ==
'gromacs':
127 dir = (state + 1 + self.stirred[
'dir']) % 2
128 return self.gen_pairs_list_gromacs(dir)
129 elif self.xchg ==
'random':
130 return self.gen_pairs_list_rand(needed=[self.stirred[
'pair']])
132 raise NotImplementedError(
133 "Unknown exchange method: %s" %
136 def gen_pairs_list(self):
137 if self.scheme ==
'standard':
138 if self.xchg ==
'gromacs':
139 return self.gen_pairs_list_gromacs(self.stepno % 2)
140 elif self.xchg ==
'random':
141 return self.gen_pairs_list_rand()
143 raise NotImplementedError(
144 "unknown exchange method: %s" %
146 elif self.scheme ==
'convective':
147 return self.gen_pairs_list_conv()
149 raise NotImplementedError(
150 "unknown exchange scheme: %s" %
153 def get_cross_energies(self, pairslist):
154 "get energies assuming all exchanges have succeeded"
155 print(
"this is not needed for temperature replica-exchange")
156 raise NotImplementedError
158 def get_metropolis(self, pairslist, old_ene):
159 """compute metropolis criterion for temperature replica exchange
160 e.g. exp(Delta beta Delta E)
161 input: list of pairs, list of state-sorted energies
164 for (s1, s2)
in pairslist:
166 min(1, np.exp((old_ene[s2] - old_ene[s1]) *
167 (self.inv_temps[s2] - self.inv_temps[s1])))
170 def try_exchanges(self, plist, metrop):
173 if (metrop[couple] >= 1)
or (random() < metrop[couple]):
174 accepted.append(couple)
177 def perform_exchanges(self, accepted):
178 "exchange given state couples both in local variables and on the grid"
180 for (i, j)
in accepted:
182 ri = self.replicanums[i]
183 rj = self.replicanums[j]
184 self.statenums[ri] = j
185 self.statenums[rj] = i
187 buf = self.replicanums[i]
188 self.replicanums[i] = self.replicanums[j]
189 self.replicanums[j] = buf
191 newtemps = self.sort_per_replica(self.inv_temps)
192 states = self.grid.gather(
193 self.grid.broadcast(self.sfo_id,
'get_state'))
194 for (i, j)
in accepted:
195 ri = self.replicanums[i]
196 rj = self.replicanums[j]
198 states[ri] = states[rj]
200 for temp, state
in zip(newtemps, states):
201 state[
'inv_temp'] = temp
203 self.grid.scatter(self.sfo_id,
'set_state', states))
205 def write_rex_stats(self):
206 "write replica numbers as a function of state"
207 fl = open(self.logfile,
'a')
208 fl.write(
'%8d ' % self.stepno)
209 fl.write(
' '.join([
'%2d' % (i + 1)
for i
in self.replicanums]))
212 if self.scheme ==
'convective':
213 fl = open(self.convectivelog,
'a')
214 fl.write(
'%5d ' % self.stepno)
215 fl.write(
'%2d ' % (self.stirred[
'replica'] + 1))
216 fl.write(
'%2d ' % self.stirred[
'dir'])
217 fl.write(
'%2d\n' % self.stirred[
'steps'])
219 if self.tune_temps
and len(self.rn_history) == 1:
220 fl = open(self.templog,
'a')
221 fl.write(
'%5d ' % self.stepno)
222 fl.write(
' '.join([
'%.3f' % i
for i
in self.inv_temps]))
227 """use TuneRex to optimize temp set. Temps are optimized every
228 'rate' steps and 'method' is used. Data is accumulated as long as
229 the temperatures weren't optimized.
230 td keys that should be passed to init:
231 rate : the rate at which to try tuning temps
233 targetAR : for ar only, target acceptance rate
234 alpha : Type I error to use.
238 self.rn_history.append([i
for i
in self.replicanums])
240 if len(self.rn_history) % td[
'rate'] == 0\
241 and len(self.rn_history) > 0:
242 temps = [1 / (kB * la)
for la
in self.inv_temps]
244 if td[
'method'] ==
'ar':
246 kwargs[
'targetAR'] = td[
'targetAR']
248 kwargs[
'alpha'] = td[
'alpha']
249 if 'dumb_scale' in td:
250 kwargs[
'dumb_scale'] = td[
'dumb_scale']
251 indicators = TuneRex.compute_indicators(
252 np.transpose(self.rn_history))
253 changed, newparams = TuneRex.tune_params_ar(
254 indicators, temps, **kwargs)
255 elif td[
'method'] ==
'flux':
257 kwargs[
'alpha'] = td[
'alpha']
258 changed, newparams = TuneRex.tune_params_flux(
259 np.transpose(self.rn_history),
264 self.inv_temps = [1 / (kB * t)
for t
in newparams]
266 def do_bookkeeping_before(self):
268 if self.scheme ==
'convective':
272 st[
'pos'] = (st[
'pos'] + 1) % self.nreps
273 st[
'replica'] = st[
'order'][st[
'pos']]
275 st[
'steps'] = 2 * (self.nreps - 1)
277 state = self.statenums[rep]
279 if state == self.nreps - 1:
285 def do_bookkeeping_after(self, accepted):
286 if self.scheme ==
'convective':
287 rep = self.stirred[
'replica']
288 state = self.statenums[rep]
289 dir = 2 * self.stirred[
'dir'] - 1
290 expected = (min(state, state + dir), max(state, state + dir))
291 if self.stirred[
'pair']
in accepted:
292 self.stirred[
'steps'] -= 1
294 def replica_exchange(self):
295 "main entry point for replica-exchange"
297 self.do_bookkeeping_before()
302 energies = self.sort_per_state(self.get_energies())
304 plist = self.gen_pairs_list()
306 metrop = self.get_metropolis(plist, energies)
308 accepted = self.try_exchanges(plist, metrop)
310 self.perform_exchanges(accepted)
312 self.do_bookkeeping_after(accepted)
def __init__
input a list of particles, the slope and theta of the sigmoid potential theta is the cutoff distance ...