def optimal_weights(self):
# DFT layer of x: xhat
x_to_xhat_re, x_to_xhat_im = dft_weights_2d(self.width, self.height)
# multiplication of Yhat in log space: Yhat
Xhat_to_Yhat = np.eye(self.width * self.height)
Xhat_to_Yhat_re = np.real(Xhat_to_Yhat)
Xhat_to_Yhat_im = np.imag(Xhat_to_Yhat)
# inverse DFT layer of yhat: y
yhat_to_y_re, yhat_to_y_im = idft_weights_2d(self.width, self.height)
return (x_to_xhat_re, x_to_xhat_im, Xhat_to_Yhat_re, Xhat_to_Yhat_im, yhat_to_y_re, yhat_to_y_im)
def optimal_shift_weights(self):
# DFT layer of x: xhat
x_to_xhat_re, x_to_xhat_im = dft_weights_2d(self.width, self.height)
if self.arch == 'conv':
# DFT layer of s: shat
s_to_shat_re, s_to_shat_im = dft_weights_2d(self.width, self.height)
elif self.arch == 'comp':
# shift calculation from s: shat
xgrid, ygrid = np.meshgrid(np.arange(self.width), np.arange(self.height), indexing='ij')
fx = xgrid.flatten()[:, np.newaxis]
fy = ygrid.flatten()[:, np.newaxis]
s_to_shat_x = -2j*np.pi * fx / self.width
s_to_shat_y = -2j*np.pi * fy / self.height
s_to_shat = np.hstack((s_to_shat_x, s_to_shat_y))
s_to_shat_re, s_to_shat_im = np.real(s_to_shat), np.imag(s_to_shat)
else:
assert False
# multiplication of Yhat and Shat in log space: Yhat
Xhat_to_Yhat = np.eye(self.width * self.height)
Xhat_to_Yhat_re = np.real(Xhat_to_Yhat)
Xhat_to_Yhat_im = np.imag(Xhat_to_Yhat)
Shat_to_Yhat = np.eye(self.width * self.height)
Shat_to_Yhat_re = np.real(Shat_to_Yhat)
Shat_to_Yhat_im = np.imag(Shat_to_Yhat)
# inverse DFT layer of yhat y:
yhat_to_y_re, yhat_to_y_im = idft_weights_2d(self.width, self.height)
return (x_to_xhat_re, x_to_xhat_im,
s_to_shat_re, s_to_shat_im,
Xhat_to_Yhat_re, Xhat_to_Yhat_im,
Shat_to_Yhat_re, Shat_to_Yhat_im,
yhat_to_y_re, yhat_to_y_im)
