-
Notifications
You must be signed in to change notification settings - Fork 10
/
Copy pathRNNwavefunction.py
200 lines (143 loc) · 8.85 KB
/
RNNwavefunction.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
import tensorflow as tf
import numpy as np
import random
class RNNwavefunction(object):
def __init__(self,systemsize_x, systemsize_y,cell=None,units=[10],scope='RNNwavefunction',seed = 111):
"""
systemsize_x: int
number of sites for x-axis
systemsize_y: int
number of sites for y-axis
cell: a tensorflow RNN cell
units: list of int
number of units per RNN layer
scope: str
the name of the name-space scope
seed: pseudo-random number generator
"""
self.graph=tf.Graph()
self.scope=scope #Label of the RNN wavefunction
self.Nx=systemsize_x #size of x direction in the 2d model
self.Ny=systemsize_y
random.seed(seed) # `python` built-in pseudo-random generator
np.random.seed(seed) # numpy pseudo-random generator
#Defining the neural network
with self.graph.as_default():
with tf.variable_scope(self.scope,reuse=tf.AUTO_REUSE):
tf.set_random_seed(seed) # tensorflow pseudo-random generator
self.rnn=cell(num_units = units[0], num_in = 2 ,name="rnn_"+str(0),dtype=tf.float64)
self.dense = tf.layers.Dense(2,activation=tf.nn.softmax,name='wf_dense', dtype = tf.float64)
def sample(self,numsamples,inputdim):
"""
generate samples from a probability distribution parametrized by a recurrent network
------------------------------------------------------------------------
Parameters:
numsamples: int
samples to be produced
inputdim: int
hilbert space dimension of one spin
------------------------------------------------------------------------
Returns:
samples: tf.Tensor of shape (numsamples,systemsize_x, systemsize_y)
the samples in integer encoding
"""
with self.graph.as_default(): #Call the default graph, used if not willing to create multiple graphs.
with tf.variable_scope(self.scope,reuse=tf.AUTO_REUSE):
#Initial input to feed to the 2drnn
self.inputdim=inputdim
self.outputdim=self.inputdim
self.numsamples=numsamples
samples=[[[] for nx in range(self.Nx)] for ny in range(self.Ny)]
rnn_states = {}
inputs = {}
for ny in range(self.Ny): #Loop over the boundary
if ny%2==0:
nx = -1
# print(nx,ny)
rnn_states[str(nx)+str(ny)]=self.rnn.zero_state(self.numsamples,dtype=tf.float64)
inputs[str(nx)+str(ny)] = tf.zeros((self.numsamples,inputdim), dtype = tf.float64)
if ny%2==1:
nx = self.Nx
# print(nx,ny)
rnn_states[str(nx)+str(ny)]=self.rnn.zero_state(self.numsamples,dtype=tf.float64)
inputs[str(nx)+str(ny)] = tf.zeros((self.numsamples,inputdim), dtype = tf.float64)
for nx in range(self.Nx): #Loop over the boundary
ny = -1
rnn_states[str(nx)+str(ny)]=self.rnn.zero_state(self.numsamples,dtype=tf.float64)
inputs[str(nx)+str(ny)] = tf.zeros((self.numsamples,inputdim), dtype = tf.float64)
#Begin sampling
for ny in range(self.Ny):
if ny%2 == 0:
for nx in range(self.Nx): #left to right
rnn_output, rnn_states[str(nx)+str(ny)] = self.rnn((inputs[str(nx-1)+str(ny)],inputs[str(nx)+str(ny-1)]), (rnn_states[str(nx-1)+str(ny)],rnn_states[str(nx)+str(ny-1)]))
output=self.dense(rnn_output)
sample_temp=tf.reshape(tf.multinomial(tf.log(output),num_samples=1),[-1,])
samples[nx][ny] = sample_temp
inputs[str(nx)+str(ny)]=tf.one_hot(sample_temp,depth=self.outputdim, dtype = tf.float64)
if ny%2 == 1:
for nx in range(self.Nx-1,-1,-1): #right to left
rnn_output, rnn_states[str(nx)+str(ny)] = self.rnn((inputs[str(nx+1)+str(ny)],inputs[str(nx)+str(ny-1)]), (rnn_states[str(nx+1)+str(ny)],rnn_states[str(nx)+str(ny-1)]))
output=self.dense(rnn_output)
sample_temp=tf.reshape(tf.multinomial(tf.log(output),num_samples=1),[-1,])
samples[nx][ny] = sample_temp
inputs[str(nx)+str(ny)]=tf.one_hot(sample_temp,depth=self.outputdim, dtype = tf.float64)
self.samples=tf.transpose(tf.stack(values=samples,axis=0), perm = [2,0,1])
return self.samples
def log_probability(self,samples,inputdim):
"""
calculate the log-probabilities of ```samples``
------------------------------------------------------------------------
Parameters:
samples: tf.Tensor
a tf.placeholder of shape (number of samples,systemsize_x,system_size_y)
containing the input samples in integer encoding
inputdim: int
dimension of the input space
------------------------------------------------------------------------
Returns:
log-probs tf.Tensor of shape (number of samples,)
the log-probability of each sample
"""
with self.graph.as_default():
self.inputdim=inputdim
self.outputdim=self.inputdim
self.numsamples=tf.shape(samples)[0]
self.outputdim=self.inputdim
samples_=tf.transpose(samples, perm = [1,2,0])
rnn_states = {}
inputs = {}
for ny in range(self.Ny): #Loop over the boundary
if ny%2==0:
nx = -1
rnn_states[str(nx)+str(ny)]=self.rnn.zero_state(self.numsamples,dtype=tf.float64)
inputs[str(nx)+str(ny)] = tf.zeros((self.numsamples,inputdim), dtype = tf.float64)
if ny%2==1:
nx = self.Nx
rnn_states[str(nx)+str(ny)]=self.rnn.zero_state(self.numsamples,dtype=tf.float64)
inputs[str(nx)+str(ny)] = tf.zeros((self.numsamples,inputdim), dtype = tf.float64)
for nx in range(self.Nx): #Loop over the boundary
ny = -1
rnn_states[str(nx)+str(ny)]=self.rnn.zero_state(self.numsamples,dtype=tf.float64)
inputs[str(nx)+str(ny)] = tf.zeros((self.numsamples,inputdim), dtype = tf.float64)
with tf.variable_scope(self.scope,reuse=tf.AUTO_REUSE):
probs = [[[] for nx in range(self.Nx)] for ny in range(self.Ny)]
#Begin estimation of log probs
for ny in range(self.Ny):
if ny%2 == 0:
for nx in range(self.Nx): #left to right
rnn_output, rnn_states[str(nx)+str(ny)] = self.rnn((inputs[str(nx-1)+str(ny)],inputs[str(nx)+str(ny-1)]), (rnn_states[str(nx-1)+str(ny)],rnn_states[str(nx)+str(ny-1)]))
output=self.dense(rnn_output)
sample_temp=tf.reshape(tf.multinomial(tf.log(output),num_samples=1),[-1,])
probs[nx][ny] = output
inputs[str(nx)+str(ny)]=tf.one_hot(samples_[nx,ny],depth=self.outputdim,dtype = tf.float64)
if ny%2 == 1:
for nx in range(self.Nx-1,-1,-1): #right to left
rnn_output, rnn_states[str(nx)+str(ny)] = self.rnn((inputs[str(nx+1)+str(ny)],inputs[str(nx)+str(ny-1)]), (rnn_states[str(nx+1)+str(ny)],rnn_states[str(nx)+str(ny-1)]))
output=self.dense(rnn_output)
sample_temp=tf.reshape(tf.multinomial(tf.log(output),num_samples=1),[-1,])
probs[nx][ny] = output
inputs[str(nx)+str(ny)]=tf.one_hot(samples_[nx,ny],depth=self.outputdim,dtype = tf.float64)
probs=tf.transpose(tf.stack(values=probs,axis=0),perm=[2,0,1,3])
one_hot_samples=tf.one_hot(samples,depth=self.inputdim, dtype = tf.float64)
self.log_probs=tf.reduce_sum(tf.reduce_sum(tf.log(tf.reduce_sum(tf.multiply(probs,one_hot_samples),axis=3)),axis=2),axis=1)
return self.log_probs