-
Notifications
You must be signed in to change notification settings - Fork 6
/
tf_observables.py
259 lines (226 loc) · 8.7 KB
/
tf_observables.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
'''
This file was directly copied from the
codebase at https://github.com/sklus/d3s
'''
import math
# import numpy as np
import tensorflow as tf
from scipy.spatial import distance
def identity(x):
'''
Identity function.
'''
return x
class monomials(object):
'''
Computation of monomials in d dimensions.
'''
def __init__(self, p):
'''
The parameter p defines the maximum order of the monomials.
'''
self.p = p
def __call__(self, x):
'''
Evaluate all monomials of order up to p for all data points in x.
'''
[d, m] = x.shape # d = dimension of state space, m = number of test points
c = allMonomialPowers(d, self.p) # matrix containing all powers for the monomials
n = c.shape[1] # number of monomials
y = [[1.0 for _ in range(m)] for _ in range(n)]
for i in range(n):
for j in range(d):
y[i] = y[i] * tf.math.pow(tf.cast(x[j], tf.float32), c[j, i])
return tf.stack(y)
def diff(self, x):
'''
Compute partial derivatives for all data points in x.
'''
[d, m] = x.shape # d = dimension of state space, m = number of test points
c = allMonomialPowers(d, self.p) # matrix containing all powers for the monomials
n = c.shape[1] # number of monomials
y = tf.zeros([n, d, m])
for i in range(n): # for all monomials
for j in range(d): # for all dimensions
e = c[:, i].copy() # exponents of ith monomial
a = e[j]
e[j] = e[j] - 1 # derivative w.r.t. j
if tf.reduce_any(e < 0):
continue # nothing to do, already zero
y[i, j, :] = a*tf.ones([1, m])
for k in range(d):
y[i, j, :] = y[i, j, :] * tf.math.pow(x[k, :], e[k])
return y
def ddiff(self, x):
'''
Compute second order derivatives for all data points in x.
'''
[d, m] = x.shape # d = dimension of state space, m = number of test points
c = allMonomialPowers(d, self.p) # matrix containing all powers for the monomials
n = c.shape[1] # number of monomials
y = tf.zeros([n, d, d, m])
for i in range(n): # for all monomials
for j1 in range(d): # for all dimensions
for j2 in range(d): # for all dimensions
e = c[:, i].copy() # exponents of ith monomial
a = e[j1]
e[j1] = e[j1] - 1 # derivative w.r.t. j1
a *= e[j2]
e[j2] = e[j2] - 1 # derivative w.r.t. j2
if tf.reduce_any(e < 0):
continue # nothing to do, already zero
y[i, j1, j2, :] = a*tf.ones([1, m])
for k in range(d):
y[i, j1, j2, :] = y[i, j1, j2, :] * tf.math.pow(x[k, :], e[k])
return y
def __repr__(self):
return 'Monomials of order up to %d.' % self.p
def display(self, alpha, d, name = None, eps = 1e-6):
'''
Display the polynomial with coefficients alpha.
'''
c = allMonomialPowers(d, self.p) # matrix containing all powers for the monomials
if name != None: print(name + ' = ', end = '')
ind, = tf.where(abs(alpha) > eps)
k = ind.shape[0]
if k == 0: # no nonzero coefficients
print('0')
return
for i in range(k):
if i == 0:
print('%.5f' % alpha[ind[i]], end = '')
else:
if alpha[ind[i]] > 0:
print(' + %.5f' % alpha[ind[i]], end = '')
else:
print(' - %.5f' % -alpha[ind[i]], end = '')
self._displayMonomial(c[:, ind[i]])
print('')
def _displayMonomial(self, p):
d = p.shape[0]
if tf.reduce_all(p == 0):
print('1', end = '')
else:
for j in range(d):
if p[j] == 0:
continue;
if p[j] == 1:
print(' x_%d' % (j+1), end = '')
else:
print(' x_%d^%d' % (j+1, p[j]), end = '')
class indicators(object):
'''
Indicator functions for box discretization Omega.
'''
def __init__(self, Omega):
self.Omega = Omega
def __call__(self, x):
[d, m] = x.shape # d = dimension of state space, m = number of test points
n = self.Omega.numBoxes()
y = tf.zeros([n, m])
for i in range(m):
ind = self.Omega.index(x[:, i])
pass
if ind == -1:
continue
y[ind, i] = 1
return y
def __repr__(self):
return 'Indicator functions for box discretization.'
class gaussians(object):
'''
Gaussians whose centers are the centers of the box discretization Omega.
sigma: width of Gaussians
'''
def __init__(self, Omega, sigma=1):
self.Omega = Omega
self.sigma = sigma
def __call__(self, x):
'''
Evaluate Gaussians for all data points in x.
'''
c = self.Omega.midpointGrid()
D = distance.cdist(c.T, x.T, 'sqeuclidean')
y = tf.math.exp(-1/(2*self.sigma**2)*D)
return y
def diff(self, x):
'''
Compute partial derivatives for all data points in x.
'''
[d, m] = x.shape # d = dimension of state space, m = number of test points
n = self.Omega.numBoxes() # number of basis functions
c = self.Omega.midpointGrid()
D = distance.cdist(c.T, x.T, 'sqeuclidean')
y = tf.zeros([n, d, m])
for i in range(n): # for all Gaussians
for j in range(d): # for all dimensions
y[i, j, :] = -2/(2*self.sigma**2) * (x[j, :] - c[j, i]) * tf.math.exp(-1/(2*self.sigma**2)*D[i, :])
return y
def ddiff(self, x):
'''
Compute second order derivatives for all data points in x.
'''
[d, m] = x.shape # d = dimension of state space, m = number of test points
n = self.Omega.numBoxes() # number of basis functions
c = self.Omega.midpointGrid()
D = distance.cdist(c.T, x.T, 'sqeuclidean')
y = tf.zeros([n, d, d, m])
for i in range(n): # for all Gaussians
for j1 in range(d): # for all dimensions
for j2 in range(d): # for all dimensions
if j1 == j2:
y[i, j1, j2, :] = ( -2/(2*self.sigma**2) + 4/(4*self.sigma**4) * (x[j1, :] - c[j1, i])**2 ) * tf.math.exp(-1/(2*self.sigma**2)*D[i, :])
else:
y[i, j1, j2, :] = ( 4/(4*self.sigma**4) * (x[j1, :] - c[j1, i]) * (x[j2, :] - c[j2, i]) ) * tf.math.exp(-1/(2*self.sigma**2)*D[i, :])
return y
def __repr__(self):
return 'Gaussian functions for box discretization with bandwidth %f.' % self.sigma
# auxiliary functions
def nchoosek(n, k):
'''
Computes binomial coefficients.
'''
return tf.math.exp(tf.math.lgamma(float(n+1))) / tf.math.exp(tf.math.lgamma(float(k+1))) / tf.math.exp(tf.math.lgamma(float(n-k+1)))
def nextMonomialPowers(x):
'''
Returns powers for the next monomial. Implementation based on John Burkardt's MONOMIAL toolbox, see
http://people.sc.fsu.edu/~jburkardt/m_src/monomial/monomial.html.
'''
m = len(x)
j = 0
for i in range(1, m): # find the first index j > 1 s.t. x[j] > 0
if x[i] > 0:
j = i
break
if j == 0:
t = x[0]
x[0] = 0
x[m - 1] = t + 1
elif j < m - 1:
x[j] = x[j] - 1
t = x[0] + 1
x[0] = 0
x[j-1] = x[j-1] + t
elif j == m - 1:
t = x[0]
x[0] = 0
x[j - 1] = t + 1
x[j] = x[j] - 1
return x
def allMonomialPowers(d, p):
'''
All monomials in d dimensions of order up to p.
'''
# Example: For d = 3 and p = 2, we obtain
# [[ 0 1 0 0 2 1 1 0 0 0]
# [ 0 0 1 0 0 1 0 2 1 0]
# [ 0 0 0 1 0 0 1 0 1 2]]
n = tf.cast(tf.round(nchoosek(p + d, p)), tf.int32) # number of monomials
x = [0.0 for _ in range(d)] # tf.zeros(d) # vector containing powers for the monomials, initially zero
c = [[0 for _ in range(d)]] # tf.zeros([d, n]) # matrix containing all powers for the monomials
def computeC(i):
return tf.cast(nextMonomialPowers(x), tf.float32)
c = tf.map_fn(fn=computeC, elems=tf.range(1, n), dtype=tf.float32)
c = tf.transpose(c)
c = c[::-1] # flip array in the up/down direction
return c