/
electron_sampler.py
313 lines (265 loc) · 14.5 KB
/
electron_sampler.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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
import numpy as np
from typing import Optional, Callable
class ElectronSampler:
"""
This class enables to initialize electron's position using gauss distribution around a nucleus and update using Markov Chain Monte-Carlo(MCMC) moves.
Using the probability obtained from the square of magnitude of wavefunction of a molecule/atom, MCMC steps can be performed to get the electron's positions and further update the wavefunction.
This method is primarily used in methods like Variational Monte Carlo to sample electrons around the nucleons.
Sampling can be done in 2 ways:
-Simultaneous: All the electrons' positions are updated all at once.
-Single-electron: MCMC steps are performed only a particular electron, given their index value.
Further these moves can be done in 2 methods:
-Symmetric: In this configuration, the standard deviation for all the steps are uniform.
-Asymmetric: In this configuration, the standard deviation are not uniform and typically the standard deviation is obtained a function like harmonic distances, etc.
Irrespective of these methods, the initialization is done uniformly around the respective nucleus and the number of electrons specified.
Example
-------
>>> from deepchem.utils.electron_sampler import ElectronSampler
>>> test_f = lambda x: 2*np.log(np.random.uniform(low=0,high=1.0,size=np.shape(x)[0]))
>>> distribution=ElectronSampler(central_value=np.array([[1,1,3],[3,2,3]]),f=test_f,seed=0,batch_no=2,steps=1000,)
>>> distribution.gauss_initialize_position(np.array([[1],[2]]))
>> print(distribution.x)
[[[[1.03528105 1.00800314 3.01957476]]
[[3.01900177 1.99697286 2.99793562]]
[[3.00821197 2.00288087 3.02908547]]]
[[[1.04481786 1.03735116 2.98045444]]
[[3.01522075 2.0024335 3.00887726]]
[[3.00667349 2.02988158 2.99589683]]]]
>>> distribution.move()
0.5115
>> print(distribution.x)
[[[[-0.32441754 1.23330263 2.67927645]]
[[ 3.42250997 2.23617126 3.55806632]]
[[ 3.37491385 1.54374006 3.13575241]]]
[[[ 0.49067726 1.03987841 3.70277884]]
[[ 3.5631939 1.68703947 2.5685874 ]]
[[ 2.84560249 1.73998364 3.41274181]]]]
"""
def __init__(self,
central_value: np.ndarray,
f: Callable[[np.ndarray], np.ndarray],
batch_no: int = 10,
x: np.ndarray = np.array([]),
steps: int = 200,
steps_per_update: int = 10,
seed: Optional[int] = None,
symmetric: bool = True,
simultaneous: bool = True):
"""
Parameters
----------
central_value: np.ndarray
Contains each nucleus' coordinates in a 2D array. The shape of the array should be(number_of_nucleus,3).Ex: [[1,2,3],[3,4,5],..]
f:Callable[[np.ndarray],np.ndarray]
A function that should give the twice the log probability of wavefunction of the molecular system when called. Should taken in a 4D array of electron's positions(x) as argument and return a numpy array containing the log probabilities of each batch.
batch_no: int, optional (default 10)
Number of batches of the electron's positions to be initialized.
x: np.ndarray, optional (default np.ndarray([]))
Contains the electron's coordinates in a 4D array. The shape of the array should be(batch_no,no_of_electrons,1,3). Can be a 1D empty array, when electron's positions are yet to be initialized.
steps: int, optional (default 10)
The number of MCMC steps to be performed when the moves are called.
steps_per_update: int (default 10)
The number of steps after which the parameters of the MCMC gets updated.
seed: int, optional (default None)
Random seed to use.
symmetric: bool, optional(default True)
If true, symmetric moves will be used, else asymmetric moves will be followed.
simultaneous: bool, optional(default True)
If true, MCMC steps will be performed on all the electrons, else only a single electron gets updated.
Attributes
----------
sampled_electrons: np.ndarray
Keeps track of the sampled electrons at every step, must be empty at start.
"""
self.x = x
self.f = f
self.num_accept = 0
self.symmetric = symmetric
self.simultaneous = simultaneous
self.steps = steps
self.steps_per_update = steps_per_update
self.central_value = central_value
self.batch_no = batch_no
self.sampled_electrons: np.ndarray = np.array([])
if seed is not None:
seed = int(seed)
np.random.seed(seed)
def harmonic_mean(self, y: np.ndarray) -> np.ndarray:
"""Calculates the harmonic mean of the value 'y' from the self.central value. The numpy array returned is typically scaled up to get the standard deviation matrix.
Parameters
----------
y: np.ndarray
Containing the data distribution. Shape of y should be (batch,no_of_electron,1,3)
Returns
-------
np.ndarray
Contains the harmonic mean of the data distribution of each batch. Shape of the array obtained (batch_no, no_of_electrons,1,1)
"""
diff = y - self.central_value
distance = np.linalg.norm(diff, axis=-1, keepdims=True)
return 1.0 / np.mean(1.0 / distance, axis=-2, keepdims=True)
def log_prob_gaussian(self, y: np.ndarray, mu: np.ndarray,
sigma: np.ndarray) -> np.ndarray:
"""Calculates the log probability of a gaussian distribution, given the mean and standard deviation
Parameters
----------
y: np.ndarray
data for which the log normal distribution is to be found
mu: np.ndarray
Means wrt which the log normal is calculated. Same shape as x or should be brodcastable to x
sigma: np.ndarray,
The standard deviation of the log normal distribution. Same shape as x or should be brodcastable to x
Returns
-------
np.ndarray
Log probability of gaussian distribution, with the shape - (batch_no,).
"""
numer = np.sum((-0.5 * ((y - mu)**2) / (sigma**2)), axis=(1, 2, 3))
denom = y.shape[-1] * np.sum(np.log(sigma), axis=(1, 2, 3))
return numer - denom
def gauss_initialize_position(self,
no_sample: np.ndarray,
stddev: float = 0.02):
"""Initializes the position around a central value as mean sampled from a gauss distribution and updates self.x.
Parameters
----------
no_sample: np.ndarray,
Contains the number of samples to initialize under each mean. should be in the form [[3],[2]..], where here it means 3 samples and 2 samples around the first entry and second entry,respectively in self.central_value is taken.
stddev: float, optional (default 0.02)
contains the stddev with which the electrons' coordinates are initialized
"""
mean = self.central_value[0]
specific_sample = no_sample[0][0]
ndim = np.shape(self.central_value)[1]
self.x = np.random.normal(mean, stddev,
(self.batch_no, specific_sample, 1, ndim))
end = np.shape(self.central_value)[0]
for i in range(1, end):
mean = self.central_value[i]
specific_sample = no_sample[i][0]
self.x = np.append(
self.x,
np.random.normal(mean, stddev,
(self.batch_no, specific_sample, 1, ndim)),
axis=1)
def electron_update(self, lp1, lp2, move_prob, ratio, x2) -> np.ndarray:
"""
Performs sampling & parameter updates of electrons and appends the sampled electrons to self.sampled_electrons.
Parameters
----------
lp1: np.ndarray
Log probability of initial parameter state.
lp2: np.ndarray
Log probability of the new sampled state.
move_prob: np.ndarray
Sampled log probabilty of the electron moving from the initial to final state, sampled assymetrically or symetrically.
ratio: np.ndarray
Ratio of lp1 and lp2 state.
x2: np.ndarray
Numpy array of the new sampled electrons.
Returns
-------
lp1: np.ndarray
The update log probability of initial parameter state.
"""
cond = move_prob < ratio
tmp_sampled = np.where(cond[:, None, None, None], x2, self.x)
if (self.steps % self.steps_per_update) == 0:
self.x = tmp_sampled
lp1 = np.where(cond, lp2, lp1)
if (np.shape(self.sampled_electrons)[0] == 0):
self.sampled_electrons = tmp_sampled
else:
self.sampled_electrons = np.concatenate(
(self.sampled_electrons, tmp_sampled))
self.num_accept += np.sum(cond)
return lp1
def move(self,
stddev: float = 0.02,
asymmetric_func: Optional[Callable[[np.ndarray],
np.ndarray]] = None,
index: Optional[int] = None) -> float:
"""Performs Metropolis-Hasting move for self.x(electrons). The type of moves to be followed -(simultaneous or single-electron, symmetric or asymmetric) have been specified when calling the class.
The self.x array is replaced with a new array at the end of each step containing the new electron's positions.
Parameters
----------
asymmetric_func: Callable[[np.ndarray],np.ndarray], optional(default None)
Should be specified for an asymmetric move.The function should take in only 1 argument- y: a numpy array wrt to which mean should be calculated.
This function should return the mean for the asymmetric proposal. For ferminet, this function is the harmonic mean of the distance between the electron and the nucleus.
stddev: float, optional (default 0.02)
Specifies the standard deviation in the case of symmetric moves and the scaling factor of the standard deviation matrix in the case of asymmetric moves.
index: int, optional (default None)
Specifies the index of the electron to be updated in the case of a single electron move.
Returns
-------
float
accepted move ratio of the MCMC steps.
"""
self.sampled_electrons = np.array([])
lp1 = self.f(self.x) # log probability of self.x state
if self.simultaneous:
if self.symmetric:
for i in range(self.steps):
x2 = np.random.normal(self.x, stddev, self.x.shape)
lp2 = self.f(x2) # log probability of x2 state
move_prob = np.log(
np.random.uniform(low=0,
high=1.0,
size=np.shape(self.x)[0]))
ratio = lp2 - lp1
lp1 = self.electron_update(lp1, lp2, move_prob, ratio, x2)
elif asymmetric_func is not None:
for i in range(self.steps):
std = stddev * asymmetric_func(self.x)
x2 = np.random.normal(self.x, std, self.x.shape)
lp2 = self.f(x2) # log probability of x2 state
lq1 = self.log_prob_gaussian(self.x, x2,
std) # forward probability
lq2 = self.log_prob_gaussian(
x2, self.x,
stddev * asymmetric_func(x2)) # backward probability
ratio = lp2 + lq2 - lq1 - lp1
move_prob = np.log(
np.random.uniform(low=0,
high=1.0,
size=np.shape(self.x)[0]))
lp1 = self.electron_update(lp1, lp2, move_prob, ratio, x2)
elif index is not None:
index = int(index)
x2 = np.copy(self.x)
altered_shape = (self.batch_no, 1, np.shape(self.x)[3])
if self.symmetric:
for i in range(self.steps):
x2[:, index, :, :] = np.random.normal(x2[:, index, :, :],
stddev,
size=altered_shape)
lp2 = self.f(x2) # log probability of x2 state
ratio = lp2 - lp1
move_prob = np.log(
np.random.uniform(low=0,
high=1.0,
size=np.shape(self.x)[0]))
lp1 = self.electron_update(lp1, lp2, move_prob, ratio, x2)
elif asymmetric_func is not None:
init_dev = stddev * asymmetric_func(
self.x) # initial standard deviation matrix
for i in range(self.steps):
std = stddev * asymmetric_func(self.x[:, index, :, :])
x2[:, index, :, :] = np.random.normal(x2[:, index, :, :],
std,
size=altered_shape)
lp2 = self.f(x2) # log probability of x2 state
init_dev[:, index, :, :] = std
lq1 = self.log_prob_gaussian(
self.x, x2, init_dev) # forward probability
lq2 = self.log_prob_gaussian(
x2, self.x,
stddev * asymmetric_func(x2)) # backward probability
ratio = lp2 + lq2 - lq1 - lp1
move_prob = np.log(
np.random.uniform(low=0,
high=1.0,
size=np.shape(self.x)[0]))
lp1 = self.electron_update(lp1, lp2, move_prob, ratio, x2)
return self.num_accept / (
(i + 1) * np.shape(self.x)[0]) # accepted move ratio