-
Notifications
You must be signed in to change notification settings - Fork 0
/
odld_system.py
150 lines (116 loc) · 5.63 KB
/
odld_system.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
from __future__ import print_function, division
import numpy, time
from west.propagators import WESTPropagator
from west.systems import WESTSystem
from westpa.binning import RectilinearBinMapper, VoronoiBinMapper
from scipy.spatial.distance import cdist
from numpy import sin, cos, exp
from numpy.random import normal as random_normal
import numpy as np
PI = numpy.pi
pcoord_len = 21
pcoord_dtype = np.float32
def dfunc(p, centers):
ds = cdist(np.array([p]), centers)
return np.array(ds[0], dtype=p.dtype)
class ODLDPropagator(WESTPropagator):
def __init__(self, rc=None):
super(ODLDPropagator, self).__init__(rc)
self.coord_len = pcoord_len
self.coord_dtype = pcoord_dtype
self.coord_ndim = 2
self.initial_pcoord = np.array([2.0, 2.0], dtype=self.coord_dtype)
self.sigma = 0.1
self.well_1 = [2,2,2,0.6,0.6]
self.well_2 = [3,8,8,0.6,0.6]
self.well_3 = [1,7.1,2,0.5,20]
self.x0 = np.ones(self.coord_ndim)
# Implement a reflecting boundary at this x value
# (or None, for no reflection)
self.reflect_at_1 = 10.0
self.reflect_at_2 = 0.0
def get_pcoord(self, state):
'''Get the progress coordinate of the given basis or initial state.'''
state.pcoord = self.initial_pcoord.copy()
def gen_istate(self, basis_state, initial_state):
initial_state.pcoord = self.initial_pcoord.copy()
initial_state.istate_status = initial_state.ISTATE_STATUS_PREPARED
return initial_state
def grad_x(self, A, sig_x, sig_y, x0, y0, x, y):
var_x = 2*sig_x*sig_x
var_y = 2*sig_y*sig_y
pre = (2*A * (x-x0))/var_x
ex = np.exp(((-1/var_x)*(x-x0)**2)+((-1/var_y)*(y-y0)**2))
return pre*ex
def grad_y(self, A, sig_x, sig_y, x0, y0, x, y):
var_x = 2*sig_x*sig_x
var_y = 2*sig_y*sig_y
pre = (2*A * (y-y0))/var_y
ex = np.exp(((-1/var_x)*(x-x0)**2)+((-1/var_y)*(y-y0)**2))
return pre*ex
def propagate(self, segments):
start_time = time.time()
well_1, well_2, well_3 = self.well_1, self.well_2, self.well_3
sig = self.sigma
gradfactor = 2*sig*sig
cos45 = np.cos(np.pi/4)
sin45 = np.sin(np.pi/4)
n_segs = len(segments)
coords = np.empty((n_segs, self.coord_len, self.coord_ndim),
dtype=self.coord_dtype)
for iseg, segment in enumerate(segments):
coords[iseg,:] = segment.pcoord[:]
coord_len = self.coord_len
reflect_at_1 = self.reflect_at_1
reflect_at_2 = self.reflect_at_2
all_displacements = np.zeros((n_segs, self.coord_len, self.coord_ndim),
dtype=self.coord_dtype)
for istep in xrange(1,coord_len):
x = coords[:,istep-1,:]
all_displacements[:,istep,:] = displacements = \
random_normal(scale=sig, size=(n_segs,self.coord_ndim))
# We have 3 wells
mod_x = (cos45*x[:,0]+sin45*x[:,1])
mod_y = (cos45*x[:,0]-sin45*x[:,1])
grad_x = self.grad_x(well_1[0], well_1[3], well_1[4], well_1[1], well_1[2], x[:,0], x[:,1]) + \
self.grad_x(well_2[0], well_2[3], well_2[4], well_2[1], well_2[2], x[:,0], x[:,1]) + \
self.grad_x(well_3[0], well_3[3], well_3[4], well_3[1], well_3[2], mod_x, mod_y)
grad_y = self.grad_y(well_1[0], well_1[3], well_1[4], well_1[1], well_1[2], x[:,0], x[:,1]) + \
self.grad_y(well_2[0], well_2[3], well_2[4], well_2[1], well_2[2], x[:,0], x[:,1]) + \
self.grad_y(well_3[0], well_3[3], well_3[4], well_3[1], well_3[2], mod_x, mod_y)
grad = np.array([grad_x,grad_y]).T
newx = x - gradfactor*grad + displacements
if reflect_at_1 is not None:
# Anything that has moved beyond reflect_at must move back that much
# boolean array of what to reflect
to_reflect_1 = newx > reflect_at_1
to_reflect_2 = newx < reflect_at_2
# how far the things to reflect are beyond our boundary
reflect_by_1 = newx[to_reflect_1] - reflect_at_1
reflect_by_2 = newx[to_reflect_2] - reflect_at_2
# subtract twice how far they exceed the boundary by
# puts them the same distance from the boundary, on the other side
newx[to_reflect_1] -= 2*reflect_by_1
newx[to_reflect_2] -= 2*reflect_by_2
coords[:,istep,:] = newx
cputime_end = time.time()
for iseg, segment in enumerate(segments):
segment.pcoord[...] = coords[iseg,:]
segment.data['displacement'] = all_displacements[iseg]
segment.cputime = cputime_end - start_time
end_time = time.time()
segment.walltime = end_time - start_time
segment.status = segment.SEG_STATUS_COMPLETE
return segments
class ODLDSystem(WESTSystem):
def initialize(self):
self.pcoord_ndim = 2
self.pcoord_dtype = pcoord_dtype
self.pcoord_len = pcoord_len
nbins = 1
self.nbins = nbins
centers = np.zeros((self.nbins,self.pcoord_ndim),dtype=self.pcoord_dtype)
centers[:,:] = 1
self.bin_mapper = VoronoiBinMapper(dfunc, centers)
self.bin_target_counts = np.empty((self.bin_mapper.nbins,), np.int)
self.bin_target_counts[...] = 10