def get_highPass(img, mode='gradient', paras={}):
if mode == 'gradient':
imgY, imgX = tf.image.image_gradients(img)
return tf.sqrt(tf.square(imgX) + tf.square(imgY), name='xHighPass')
elif mode == 'fft':
# n, h, w, c = tf.shape(img)
n = 5
h = 320
w = 320
c = 1
# freq_thershold = paras.get('freq_thershold', int(h/3))
# # freq_thershold = int(h/3)
freq_thershold = 8
# img = tf.reshape(img[:,:,:,0], [5, 320,320,1])
img_fft = tf.fft3d(tf.cast(img, tf.complex64))
# left = freq_thershold; right = w - freq_thershold
# up = freq_thershold; down = h - freq_thershold
mask = np.zeros((h, w, c))
mask[freq_thershold:h - freq_thershold, :, :] = 1
mask[:, freq_thershold:w - freq_thershold, :] = 1
# maskHP = np.zeros((h, w))
# maskHP[freqThreshold_HP:h-freqThreshold_HP, :] = 1
# maskHP[:, freqThreshold_HP:w-freqThreshold_HP] = 1
# mask = np.ones([h, w, c])
# mask[up:down, left:right, :] = 0
# plt.imsave('./highPass_mask.png', mask, cmap='gray')
# plt.imsave('./highPass_masked.png', mask, cmap='gray')
return tf.real(tf.ifft3d(tf.multiply(img_fft, mask)))
# return tf.ifft(tf.multiply(img_fft, mask))
# return tf.real(tf.ifft(tf.fft(tf.cast(img, tf.complex64))))
elif mode == 'fft_Gaussian':
mask = 1 - get_Gaussian_mask(0.007, 0.001, 320, 320)
img_fft = tf.fft3d(tf.cast(img, tf.complex64))
return tf.real(tf.ifft3d(tf.multiply(img_fft, mask)))
def Shrinkwrap(self, sigma, thresh):
dist = tf.contrib.distributions.MultivariateNormalDiag(
[0., 0., 0.], [sigma, sigma, sigma])
kernelFFT = tf.fft3d(
tf.cast(tf.reshape(dist.prob(self._domain), self._probSize),
tf.complex64))
blurred = tf.abs(
tf.ifft3d(
tf.multiply(
tf.fft3d(tf.cast(tf.abs(self._cImage), tf.complex64)),
kernelFFT)))
blurred = tf.concat((blurred[(self._probSize[0] // 2):, :, :],
blurred[:(self._probSize[0] // 2), :, :]),
axis=0)
blurred = tf.concat((blurred[:, (self._probSize[1] // 2):, :],
blurred[:, :(self._probSize[1] // 2), :]),
axis=1)
blurred = tf.concat((blurred[:, :, (self._probSize[2] // 2):],
blurred[:, :, :(self._probSize[2] // 2)]),
axis=2)
self._support = tf.cast(blurred > thresh * tf.reduce_max(blurred),
tf.complex64)
self._support_comp = tf.cast(
blurred <= thresh * tf.reduce_max(blurred), tf.complex64)
return
def defineShrinkwrap(self, varDict, gpu):
# Array coordinates
x, y, z = np.meshgrid(list(range(varDict['cImage'].shape[0])),
list(range(varDict['cImage'].shape[1])),
list(range(varDict['cImage'].shape[2])))
y = y.max() - y
x = x - x.mean()
y = y - y.mean()
z = z - z.mean()
with tf.device('/gpu:%d' % gpu):
# these are Tensorflow variables
self._x = tf.constant(fftshift(x), dtype=tf.float32, name='x')
self._y = tf.constant(fftshift(y), dtype=tf.float32, name='y')
self._z = tf.constant(fftshift(z), dtype=tf.float32, name='z')
self._blurred = tf.Variable(np.zeros(varDict['support'].shape),
dtype=tf.complex64,
name='blurred')
self._dist = tf.Variable(tf.zeros(self._x.shape, dtype=tf.float32),
name='dist')
# These are shrinkwrap-specific symbolic ops
with tf.name_scope('Shrinkwrap'):
self._getNewDist = tf.assign(
self._dist,
tf.exp(self._neg * (self._x * self._x + self._y * self._y +
self._z * self._z) /
(self._sigma * self._sigma)),
name='getNewDist')
# self._copyDistToRollBuffer = tf.assign( self._rollBuffer, tf.cast( self._dist, dtype=tf.complex64 ), name='CopyDistToRollBuffer' )
# self._retrieveDistFromRollBuffer = tf.assign( self._blurred, self._rollBuffer, name='retrieveDistFromRollBuffer' )
# self._copyImageToRollBuffer = tf.assign( self._rollBuffer, tf.cast( tf.abs( self._cImage ), dtype=tf.complex64 ), name='copyImgToRollBuffer' )
# self._convolveRollBufferWithBlur = tf.assign( self._rollBuffer, self._rollBuffer*self._blurred, name='convolve' )
# self._retrieveBlurred = tf.assign( self._blurred, self._rollBuffer, name='retrieveBlurred' )
self._blurShape = tf.assign(
self._blurred,
tf.ifft3d(
tf.fft3d( tf.cast( self._dist, dtype=tf.complex64 ) ) *\
tf.fft3d( tf.cast( tf.abs( self._cImage ), dtype=tf.complex64 ) )
),
name='blurShape'
)
self._updateSupport = tf.assign(
self._support,
tf.cast(
tf.abs(self._blurred) >
self._thresh * tf.reduce_max(tf.abs(self._blurred)),
tf.complex64),
name='updateSup')
self._updateSupComp = tf.assign(
self._support_comp,
tf.cast(
tf.abs(self._blurred) <=
self._thresh * tf.reduce_max(tf.abs(self._blurred)),
tf.complex64),
name='updateSupComp')
def my_ift3d(tensor, scaling=1.):
"""
fftshift(ifft(ifftshift(a)))
Applies shifts to work with arrays whose "zero" is in the middle
instead of the first element.
Uses standard normalization of fft unlike dip_image.
"""
return fftshift(tf.ifft3d(ifftshift(tensor))) * scaling
def __SFKernel__(self, mycount, myimage, myimage_fft_mod):
myimage = 2. * (self._support * myimage) - myimage
myimage = tf.ifft3d(
tf.multiply(
self._modulus,
tf.exp(
tf.complex(tf.zeros(myimage.shape),
tf.angle(tf.fft3d(myimage))))))
myimage = 2. * (self._support * myimage) - myimage
mycount -= 1
return mycount, myimage, myimage_fft_mod
def ifft(self, k):
rank = len(k.shape) - 2
assert rank >= 1
if rank == 1:
return tf.stack([tf.ifft(c) for c in tf.unstack(k, axis=-1)], axis=-1)
elif rank == 2:
return tf.stack([tf.ifft2d(c) for c in tf.unstack(k, axis=-1)], axis=-1)
elif rank == 3:
return tf.stack([tf.ifft3d(c) for c in tf.unstack(k, axis=-1)], axis=-1)
else:
raise NotImplementedError('n-dimensional inverse FFT not implemented.')
def conv3d_fft_tf(vol, otf):
input = tf.complex(vol, tf.zeros(vol.shape, dtype=tf.float32))
input = tf.cast(input, dtype=tf.complex64)
otf = tf.cast(otf, dtype=tf.complex64)
vol_fft = tf.fft3d(input)
vol_fftshift = hp.fftshift3d(vol_fft)
vol_fftshift = tf.multiply(vol_fftshift, otf)
vol_fftshift = hp.ifftshift3d(vol_fftshift)
vol_fft = tf.ifft3d(vol_fftshift)
return abssqr_tf(vol_fft)
def __HIOKernel__(self, mycount, myimage, myimage_fft_mod):
origImage = tf.identity(myimage)
myimage = tf.ifft3d(
tf.multiply(
self._modulus,
tf.exp(
tf.complex(tf.zeros(myimage.shape),
tf.angle(tf.fft3d(myimage))))))
myimage = tf.multiply(self._support, myimage) + tf.multiply(
self._support_comp, origImage - self._beta * myimage)
mycount -= 1
return mycount, myimage, myimage_fft_mod
def __ERKernel__(self, mycount, myimage, myimage_fft_mod):
myimage = tf.ifft3d(
tf.multiply(
self._modulus,
tf.exp(
tf.complex(tf.zeros(myimage.shape),
tf.angle(tf.fft3d(myimage))))))
myimage = tf.multiply(myimage, self._support)
myimage_fft_mod = tf.cast(tf.abs(tf.fft3d(myimage)),
dtype=tf.complex64)
mycount -= 1
return mycount, myimage, myimage_fft_mod
def ifft3d_b01c(x_b01c):
#x_shape=tf.shape(x_b01c)
x = int(x_b01c.get_shape()[1])
y = int(x_b01c.get_shape()[2])
z = int(x_b01c.get_shape()[3])
# fft2d for b01c type images. fft2d performs only for inner most 2 dims, so
# we need transpose to put w, h dim to the final 2 dims.
x_bc01 = tf.transpose(fftshift_b01c(x_b01c), (0, 4, 1, 2, 3))
ifft_x_bc01 = tf.ifft3d(x_bc01)
ifft_x_b01c = tf.transpose(ifft_x_bc01, (0, 2, 3, 4, 1))
return ifft_x_b01c * x * y * z
def conv3d_fft_tf(vol, otf):
''' convolve given volume with OTF
Requirement/Assumption:
volumne AND OTF are not shifted
'''
# input = tf.complex(vol, tf.zeros(vol.shape, dtype=tf.float32))
input = tf.cast(vol, dtype=tf.complex64)
otf = tf.cast(otf, dtype=tf.complex64)
vol_fft = tf.fft3d(input)
# vol_fftshift = hp.fftshift3d(vol_fft)
vol_fftshift = tf.multiply(vol_fft, otf)
# vol_fftshift = hp.ifftshift3d(vol_fftshift)
vol_fft = tf.ifft3d(vol_fftshift)
return tf.real(vol_fft)
def get_lowPass(img, mode='average', paras={}):
if mode == 'average':
size = paras.get('size', 7)
lowPassFilter_C3 = tf.constant(1 / size**2,
shape=[size, size, 3, 3],
name='lowPass_filter_C1')
return tf.nn.conv2d(img,
lowPassFilter_C3,
strides=[1, 1, 1, 1],
padding='SAME',
name='xlowPass')
elif mode == 'fft':
# n, h, w, c = img.shape
n = 5
h = 320
w = 320
c = 1
# img = tf.reshape(img[:,:,:,0], [5, 320,320,1])
# freq_thershold = paras.get('freq_thershold', int(h/3))
freq_thershold = 8
img_fft = tf.fft3d(tf.cast(img, tf.complex64))
# left = freq_thershold; right = w - freq_thershold
# up = freq_thershold; down = h - freq_thershold
# mask = np.ones([h, w, c])
# mask = np.zeros([h, w, c])
# mask[up:down, left:right, :] = 1
mask = np.ones((h, w, c))
mask[freq_thershold:h - freq_thershold, :, :] = 0
mask[:, freq_thershold:w - freq_thershold, :] = 0
return tf.real(tf.ifft3d(tf.multiply(img_fft, mask)))
# return tf.ifft(tf.multiply(img_fft, mask))
# return tf.real(tf.ifft(tf.fft(tf.cast(img, tf.complex64))))
elif mode == 'fft_Gaussian':
mask = get_Gaussian_mask(0.007, 0.001, 320, 320)
img_fft = tf.fft3d(tf.cast(img, tf.complex64))
return tf.real(tf.ifft3d(tf.multiply(img_fft, mask)))
def gpu_gaussian_random_field(size=32, scale=1, length=48 * 28):
shape = (size, size, length)
amplitude = tf.constant(form_spectral_matrix(shape), dtype=tf.float32)
complex_amplitude = tf.complex(amplitude, tf.zeros(shape,
dtype=tf.float32))
random_noise = tf.random_normal(shape=shape,
stddev=scale,
dtype=tf.float32)
zeros = tf.zeros(shape, dtype=tf.float32)
complex_noise = tf.complex(random_noise, zeros)
noise_spectrum = tf.fft3d(complex_noise)
convolved = tf.multiply(complex_amplitude, noise_spectrum)
simulation = tf.ifft3d(convolved)
with tf.Session() as sess:
result = sess.run(simulation)
return result
def initialize_k_space_domain(self):
if self.use_spatial_patching:
recon_shape = self.recon_shape_full
else:
recon_shape = self.recon_shape
self.DT_recon_r = tf.get_variable(
shape=recon_shape,
dtype=tf.float32,
initializer=tf.random_uniform_initializer(0, 1e-7),
name='recon_real_k')
self.DT_recon_i = tf.get_variable(
shape=recon_shape,
dtype=tf.float32,
initializer=tf.random_uniform_initializer(0, 1e-7),
name='recon_imag_k')
self.k_space = tf.complex(self.DT_recon_r, self.DT_recon_i)
self.DT_recon = self.tf_ifftshift3(
tf.ifft3d(self.tf_fftshift3(self.k_space)))
def upsample_FT(inputs, upsam_size, scope, data_format='channels_last'):
if data_format == 'channels_first' or data_format == 'NCDHW':
raise RuntimeError("This has not been tested for channels_first")
sh = np.array(inputs.shape.as_list())
f_ny_old = sh // 2 #nyqvist frequency of original tensor
with tf.name_scope(scope):
t_cmplx = tf.complex(inputs, tf.zeros(inputs.shape))
t_cmplx_ft = tf.fft3d(t_cmplx)
t_cmplx_ft_pad = tf.manip.roll(t_cmplx_ft, f_ny_old, axis=(0, 1, 2))
t_cmplx_ft_pad = tf.pad(
t_cmplx_ft_pad, ((0, (upsam_size[0] - 1) * t_cmplx_ft.shape[0]),
(0, (upsam_size[1] - 1) * t_cmplx_ft.shape[1]),
(0, (upsam_size[2] - 1) * t_cmplx_ft.shape[2])),
'constant')
t_cmplx_ft_pad = tf.manip.roll(t_cmplx_ft_pad,
-f_ny_old,
axis=(0, 1, 2))
t_upsam = tf.real(tf.ifft3d(t_cmplx_ft_pad))
# the test found a significant imag part though --> bc of hard edge?
return t_upsam
def fftconvolve3d(x, y, padding):
# FIXME SAME will not work correctly
# FIXME specifically designed for normxcorr (need to work more to make it general)
# Read shapes
x_shape = np.array(tuple(x.get_shape().as_list()), dtype=np.int32)
y_shape = np.array(tuple(y.get_shape().as_list()), dtype=np.int32)
# Construct paddings and pad
x_shape[1:4] = x_shape[1:4] - 1
y_pad = [[0, 0], [0, x_shape[1]], [0, x_shape[2]], [0, x_shape[3]]]
y_shape[1:4] = y_shape[1:4] - 1
x_pad = [[0, 0], [0, y_shape[1]], [0, y_shape[2]], [0, y_shape[3]]]
x = tf.pad(x, x_pad)
y = tf.pad(y, y_pad)
y = tf.cast(y, tf.complex64, name='complex_Y')
x = tf.cast(x, tf.complex64, name='complex_X')
convftt = tf.real(
tf.ifft3d(tf.multiply(tf.fft3d(x), tf.fft3d(y), name='fft_mult')))
print(convftt.get_shape())
#Slice correctly based on requirements
if padding == 'VALID':
begin = [0, y_shape[1], y_shape[2], y_shape[3]]
size = [
x_shape[0], x_shape[1] - y_shape[1], x_shape[2] - y_shape[1], 1
]
if padding == 'SAME':
begin = [0, y_shape[1] / 2 - 1, y_shape[2] / 2 - 1, y_shape[3] - 1]
size = x_shape #[-1, x_shape[0], x_shape[1]]
z = tf.slice(convftt, begin, size)
z = tf.squeeze(z)
return z
def phase_corr(ashape, bshape, filter_shape=None):
"""Construct a TensorFlow op to compute phase correlation.
Parameters
----------
ashape : tuple of ints
Shape of input array.
bshape : tuple of ints
Shape of input array.
filter_shape : tuple
Shape of filter array. Optional. If not given, the window filter is
not applied.
Returns
-------
phase_corr : tf.Operation
The op to be run to compute phase correlation. When running the op,
values for the following placeholders must be fed:
`input/a_ph:0`, `input/b_ph:0`, `input/filter_ph:0`.
"""
my_filter_t = None
with tf.name_scope('input'):
aph = tf.placeholder(dtype=tf.uint16, shape=ashape, name='a_ph')
bph = tf.placeholder(dtype=tf.uint16, shape=bshape, name='b_ph')
if filter_shape is not None:
my_filter_t = tf.placeholder(dtype=tf.float32,
shape=filter_shape,
name='filter_ph')
at = tf.to_float(aph)
bt = tf.to_float(bph)
with tf.name_scope('subtract_mean'):
at -= tf.reduce_mean(at)
bt -= tf.reduce_mean(bt)
if filter_shape is not None:
with tf.name_scope('window_filter'):
at = at * my_filter_t
bt = bt * my_filter_t
with tf.name_scope('FFT'):
ac = tf.cast(at, tf.complex64, name='to_complex')
bc = tf.cast(bt, tf.complex64, name='to_complex')
aft = tf.fft3d(ac)
bft = tf.fft3d(bc)
with tf.name_scope('cross_power_spectrum'):
prod = aft * tf.conj(bft)
prodnorm = tf.abs(prod, name='norm')
ratio = prod / tf.cast(prodnorm, tf.complex64, name='to_complex')
with tf.name_scope('phase_correlation'):
ifft = tf.ifft3d(ratio)
phase_corr = tf.square(tf.real(ifft) + tf.square(tf.imag(ifft)))
phase_corr = tf.sqrt(phase_corr)
return phase_corr
def initializeGaussianPCF(self,
vardict,
array_shape,
gpu,
learning_rate=1.e-1):
x, y, z = np.meshgrid(list(range(array_shape[0])),
list(range(array_shape[1])),
list(range(array_shape[2])))
y = y.max() - y
x = (fftshift(x - x.mean())).reshape(1, -1)
y = (fftshift(y - y.mean())).reshape(1, -1)
z = (fftshift(z - z.mean())).reshape(1, -1)
pts = np.concatenate((x, y, z), axis=0)
if 'initial_guess' not in vardict.keys():
l1p, l2p, l3p, psip, thetap, phip = 2., 2., 2., 0., 0., 0.
else:
l1p, l2p, l3p, psip, thetap, phip = tuple(vardict['initial_guess'])
with tf.device('/gpu:%d' % gpu):
with tf.name_scope('GaussianPCF'):
self._roll = [n // 2 for n in self._probSize
] # defined in GPUModule_Base
with tf.name_scope('Constants'):
self._coherentEstimate = tf.Variable(
tf.zeros(self._cImage.shape, dtype=tf.float32),
name='coherentEstimate')
self._intensity = tf.constant(vardict['modulus']**2,
dtype=tf.float32,
name='Measurement')
self._q = tf.constant(pts,
dtype=tf.float32,
name='domainPoints')
self._v0 = tf.constant(np.array([1., 0.,
0.]).reshape(-1, 1),
dtype=tf.float32)
self._v1 = tf.constant(np.array([0., 1.,
0.]).reshape(-1, 1),
dtype=tf.float32)
self._v2 = tf.constant(np.array([0., 0.,
1.]).reshape(-1, 1),
dtype=tf.float32)
self._nskew0 = tf.constant(np.array([[0., 0., 0.],
[0., 0., -1.],
[0., 1., 0.]]),
dtype=tf.float32)
self._nskew1 = tf.constant(np.array([[0., 0., 1.],
[0., 0., 0.],
[-1., 0., 0.]]),
dtype=tf.float32)
self._nskew2 = tf.constant(np.array([[0., -1., 0.],
[1., 0., 0.],
[0., 0., 0.]]),
dtype=tf.float32)
self._one = tf.constant(1., dtype=tf.float32)
self._neg = tf.constant(-0.5, dtype=tf.float32)
self._twopi = tf.constant((2 * np.pi)**(3. / 2.),
dtype=tf.float32)
self._I = tf.eye(3)
with tf.name_scope('Parameters'):
self._l1 = tf.Variable(l1p,
dtype=tf.float32,
name='Lambda1') #
self._l2 = tf.Variable(
l2p, dtype=tf.float32, name='Lambda2'
) # Sqrt of eigenvalues of covariance matrix
self._l3 = tf.Variable(l3p,
dtype=tf.float32,
name='Lambda3') #
self._psi = tf.Variable(
psip, dtype=tf.float32,
name='Psi') # Rotation angle of eigenbasis
self._theta = tf.Variable(
thetap, dtype=tf.float32,
name='Theta') # Polar angle of rotation axis
self._phi = tf.Variable(
phip, dtype=tf.float32,
name='Phi') # Azimuth angle of rotation axis
with tf.name_scope('Auxiliary'):
self._FreqSupportMask = tf.placeholder(
dtype=tf.float32, shape=self._intensity.shape)
self._mD = tf.diag([self._l1, self._l2, self._l3])
self._n0 = tf.sin(self._theta) * tf.cos(self._phi)
self._n1 = tf.sin(self._theta) * tf.sin(self._phi)
self._n2 = tf.cos(self._theta)
self._n = self._n0 * self._v0 + self._n1 * self._v1 + self._n2 * self._v2
self._nskew = self._n0 * self._nskew0 + self._n1 * self._nskew1 + self._n2 * self._nskew2
self._R = tf.cos( self._psi )*self._I +\
tf.sin( self._psi )*self._nskew +\
( self._one - tf.cos( self._psi ) )*tf.matmul( self._n, tf.transpose( self._n ) )
self._C = tf.matmul(
self._R,
tf.matmul(tf.matmul(self._mD, self._mD),
tf.transpose(self._R)))
with tf.name_scope('Blurring'):
self._pkfft = tf.Variable(np.zeros(self._probSize),
dtype=tf.complex64,
name='pkfft')
self._blurKernel = tf.reshape(
tf.exp(self._neg * tf.reduce_sum(
self._q * tf.matmul(self._C, self._q), axis=0)),
shape=self._coherentEstimate.shape) * (
self._l1 * self._l2 * self._l3) / self._twopi
self._tf_intens_f = tf.fft3d(
tf.cast(self._coherentEstimate, dtype=tf.complex64))
self._tf_blur_f = tf.fft3d(
tf.cast(self._blurKernel, dtype=tf.complex64))
self._tf_prod_f = self._tf_intens_f * self._tf_blur_f
self._imgBlurred = tf.abs(tf.ifft3d(self._tf_prod_f))
with tf.name_scope('Optimizer'):
self._var_list = [
self._l1, self._l2, self._l3, self._psi, self._theta,
self._phi
]
self._poissonNLL = tf.reduce_mean(
self._FreqSupportMask *
(self._imgBlurred -
self._intensity * tf.log(self._imgBlurred)))
self._poissonOptimizer = tf.train.AdagradOptimizer(
learning_rate=vardict['pcc_learning_rate'],
name='poissonOptimize')
self._trainPoisson = self._poissonOptimizer.minimize(
self._poissonNLL, var_list=self._var_list)
self._currentGradients = [
n[0] for n in self._poissonOptimizer.compute_gradients(
self._poissonNLL, var_list=self._var_list)
]
# self._gaussNLL = tf.reduce_mean(
# self._FreqSupportMask * ( tf.sqrt( self._imgBlurred ) - tf.sqrt( self._intensity ) )**2
# )
# self._gaussOptimizer = tf.train.AdagradOptimizer( learning_rate=vardict[ 'pcc_learning_rate' ], name='gaussOptimize' )
# self._trainGauss = self._gaussOptimizer.minimize( self._gaussNLL, var_list=self._var_list )
# self._currentGradients = [
# n[0] for n in self._gaussOptimizer.compute_gradients( self._gaussNLL, var_list=self._var_list )
# ]
with tf.name_scope('Preparation'):
self._getCoherentEstimate = tf.assign(
self._coherentEstimate,
tf.cast(self._intermedInt, dtype=tf.float32),
name='getCoherentEstimate')
self._getPCFKernelFFT = tf.assign(
self._pkfft,
tf.fft3d(tf.cast(self._blurKernel,
dtype=tf.complex64)),
name='getPCFKernelFFT')
with tf.name_scope('Convolution'):
self._convolveWithCoherentEstimate = tf.assign(
self._intermedInt,
tf.cast(tf.abs(
tf.ifft3d(self._pkfft *
tf.fft3d(self._intermedInt))),
dtype=tf.complex64),
name='convolveWithCoherentEstimate')
self._progress = []
def defineBaseVariables(self, varDict, gpu):
# Tensorflow variables specified here.
with tf.device('/gpu:%d' % gpu):
self._modulus = tf.constant(varDict['modulus'],
dtype=tf.complex64,
name='mod_measured')
self._support = tf.Variable(varDict['support'],
dtype=tf.complex64,
name='sup')
self._support_comp = tf.Variable(1. - varDict['support'],
dtype=tf.complex64,
name='Support_comp')
self._cImage = tf.Variable(varDict['cImage'],
dtype=tf.complex64,
name='Image')
self._buffImage = tf.Variable(varDict['cImage'],
dtype=tf.complex64,
name='buffImage')
self._beta = tf.constant(varDict['beta'],
dtype=tf.complex64,
name='beta')
self._probSize = self._cImage.shape
self._thresh = tf.placeholder(dtype=tf.float32, name='thresh')
self._sigma = tf.placeholder(dtype=tf.float32, name='sigma')
self._neg = tf.constant(-0.5, dtype=tf.float32)
self._intermedFFT = tf.Variable(tf.zeros(self._cImage.shape,
dtype=tf.complex64),
name='intermedFFT')
self._intermedInt = tf.Variable(tf.zeros(self._cImage.shape,
dtype=tf.complex64),
name='intermedInt')
with tf.name_scope('Support'):
self._supproject = tf.assign(self._cImage,
self._cImage * self._support,
name='supProject')
# These are defined only if high-energy phasing is required.
if 'bin_left' in varDict.keys():
bL = varDict['bin_left']
sh = bL.shape
self._binL = tf.constant(bL.reshape(sh[0], sh[1], 1).repeat(
varDict['modulus'].shape[-1], axis=2),
dtype=tf.complex64,
name='binL')
self._binR = tf.constant(bL.T.reshape(sh[1], sh[0], 1).repeat(
varDict['modulus'].shape[-1], axis=2),
dtype=tf.complex64,
name='binR')
self._scale = tf.constant(varDict['scale'],
dtype=tf.complex64,
name='scale')
self._binned = tf.Variable(tf.zeros(self._modulus.shape,
dtype=tf.complex64),
name='binned')
self._expanded = tf.Variable(tf.zeros(self._support.shape,
dtype=tf.complex64),
name='expanded')
self._scaled = tf.Variable(tf.zeros(self._modulus.shape,
dtype=tf.complex64),
name='scaled')
with tf.name_scope('highEnergy'):
self._binThis = tf.assign(
self._binned,
tf.transpose(
tf.matmul(
tf.matmul(
tf.transpose(self._binL, [2, 0, 1]),
tf.transpose(
tf.cast(tf.square(
tf.abs(self._intermedFFT)),
dtype=tf.complex64),
[2, 0, 1])),
tf.transpose(self._binR, [2, 0, 1])),
[1, 2, 0]),
name='Binning')
self._scaleThis = tf.assign(self._scaled,
tf.divide(
self._modulus,
tf.sqrt(self._binned)),
name='Scaling')
self._expandThis = tf.assign(
self._expanded,
tf.transpose(
tf.matmul(
tf.matmul(
tf.transpose(self._binR, [2, 0, 1]),
tf.transpose(self._scaled, [2, 0, 1])),
tf.transpose(self._binL, [2, 0, 1])),
[1, 2, 0]),
name='Expansion')
self._HEImgUpdate = tf.assign(
self._cImage,
tf.multiply(
self._support,
tf.ifft3d(self._scale * tf.multiply(
self._expanded, self._intermedFFT))),
name='HEImgUpdate')
self._HEImgCorrect = tf.assign(
self._cImage,
self._cImage + tf.multiply(
self._support_comp,
self._buffImage - self._beta * self._cImage),
name='HEImgCorrect')
else: # regular phasing
with tf.name_scope('ER'):
self._modproject = tf.assign(
self._cImage,
tf.ifft3d(
tf.divide(self._modulus, tf.sqrt(
self._intermedInt)) * tf.fft3d(self._cImage)),
name='modProject')
return
##im_ft_pad[:f_ny_old[0], :f_ny_old[1], :f_ny_old[2] ] = im_ft[:f_ny_old[0], :f_ny_old[1], :f_ny_old[2]]
##im_ft_pad[:f_ny_old[0], :f_ny_old[1], -f_ny_old[2]:] = im_ft[:f_ny_old[0], :f_ny_old[1], -f_ny_old[2]:]
##im_ft_pad[:f_ny_old[0], -f_ny_old[1]:, :f_ny_old[2] ] = im_ft[:f_ny_old[0], -f_ny_old[1]:, :f_ny_old[2]]
##im_ft_pad[:f_ny_old[0], -f_ny_old[1]:, -f_ny_old[2]:] = im_ft[:f_ny_old[0], -f_ny_old[1]:, -f_ny_old[2]:]
##im_ft_pad[-f_ny_old[0]:, :f_ny_old[1], :f_ny_old[2] ] = im_ft[-f_ny_old[0]:, :f_ny_old[1], :f_ny_old[2]]
##im_ft_pad[-f_ny_old[0]:, :f_ny_old[1], -f_ny_old[2]:] = im_ft[-f_ny_old[0]:, :f_ny_old[1], -f_ny_old[2]:]
##im_ft_pad[-f_ny_old[0]:, -f_ny_old[1]:, :f_ny_old[2] ] = im_ft[-f_ny_old[0]:, -f_ny_old[1]:, :f_ny_old[2]]
##im_ft_pad[-f_ny_old[0]:, -f_ny_old[1]:, -f_ny_old[2]:] = im_ft[-f_ny_old[0]:, -f_ny_old[1]:, -f_ny_old[2]:]
im_ft_pad = tf.manip.roll(im_ft, f_ny_old, axis=(0, 1, 2))
im_ft_pad = tf.pad(im_ft_pad, ((0, (upsam_factor - 1) * im_ft.shape[0]),
(0, (upsam_factor - 1) * im_ft.shape[1]),
(0, (upsam_factor - 1) * im_ft.shape[2])),
'constant')
im_ft_pad = tf.manip.roll(im_ft_pad, -f_ny_old, axis=(0, 1, 2))
im_upsam = tf.ifft3d(im_ft_pad)
with tf.Session() as sess:
im_upsamc, im_ft_padc, im_ftc, imc = sess.run(
[im_upsam, im_ft_pad, im_ft, im])
print(np.abs(im_upsamc.imag).max())
print(np.sum(np.abs(im_upsamc.imag)))
print(np.abs(im_upsamc.real).max())
print(np.sum(np.abs(im_upsamc.real)))
fig, axes = plt.subplots(2, 2)
axes[0, 0].imshow(imc[..., 7])
axes[0, 1].imshow(np.log(np.abs(im_ftc[..., 0]) + 1))
axes[1, 0].imshow(np.log(np.abs(im_ft_padc[..., 0]) + 1))
axes[1, 1].imshow(im_upsamc.real[..., upsam_factor * 7])
Vk = tf.placeholder("complex64", shape=(N, N, N))
nkplace = tf.placeholder("complex64", shape=(N, N, N))
nrplace = tf.placeholder("complex64", shape=(N, N, N))
nk = tf.Variable(nkplace)
nr = tf.Variable(nrplace)
''' Control time step'''
dt = 0.0001
dtr = tf.constant(dt, dtype=tf.complex64)
Tr = tf.constant(Tr1, dtype=tf.complex64)
'''Update rule '''
SS = tf.log(nr / (1 - nr))
Sk = tf.fft3d(SS)
dnkdt = Lk * (Vk * nk / T0 + Tr * (Sk))
nk1 = nk + dnkdt * dtr
nr_ = tf.ifft3d(nk1)
step1 = tf.group(
nk.assign(nk1),
nr.assign(nr_),
)
'''Calculate order parameter '''
def make_kernel(a):
"""Transform a 3D array into a convolution kernel"""
a = np.asarray(a)
a = a.reshape(list(a.shape) + [1, 1])
return tf.constant(a, dtype="complex64")
def getGaussianPCF(noisy_data,
coherent_estimate,
domain,
learning_rate=1.e-1,
recip_basis=np.eye(3),
min_iterations=200,
max_iterations=10000,
iterations_per_checkpoint=50,
tol=1.e-4):
l1p, l2p, l3p = np.random.rand(), np.random.rand(), np.random.rand()
psip, thetap, phip = 2. * np.pi * np.random.rand(), np.pi * np.random.rand(
), 2. * np.pi * np.random.rand()
with tf.device('/gpu:0'):
roll_ = [n // 2 for n in coherent_estimate.shape]
# defining constants
with tf.name_scope('Constants'):
noisyData = tf.constant(noisy_data,
dtype=tf.float32,
name='noisyData')
coherentEstimate = tf.constant(coherent_estimate,
dtype=tf.float32,
name='coherentEstimate')
q = tf.constant(domain, dtype=tf.float32, name='domainPoints')
v0 = tf.constant(np.array([1., 0., 0.]).reshape(-1, 1),
dtype=tf.float32)
v1 = tf.constant(np.array([0., 1., 0.]).reshape(-1, 1),
dtype=tf.float32)
v2 = tf.constant(np.array([0., 0., 1.]).reshape(-1, 1),
dtype=tf.float32)
nskew0 = tf.constant(np.array([[0., 0., 0.], [0., 0., -1.],
[0., 1., 0.]]),
dtype=tf.float32)
nskew1 = tf.constant(np.array([[0., 0., 1.], [0., 0., 0.],
[-1., 0., 0.]]),
dtype=tf.float32)
nskew2 = tf.constant(np.array([[0., -1., 0.], [1., 0., 0.],
[0., 0., 0.]]),
dtype=tf.float32)
one = tf.constant(1., dtype=tf.float32)
neg = tf.constant(-0.5, dtype=tf.float32)
twopi = tf.constant((2 * np.pi)**(3. / 2.), dtype=tf.float32)
I = tf.eye(3)
# defining the 6 optimization parameters
with tf.name_scope('Parameters'):
l1 = tf.Variable(l1p, dtype=tf.float32, name='Lambda1') #
l2 = tf.Variable(
l2p, dtype=tf.float32,
name='Lambda2') # Sqrt of eigenvalues of covariance matrix
l3 = tf.Variable(l3p, dtype=tf.float32, name='Lambda3') #
psi = tf.Variable(psip, dtype=tf.float32,
name='Psi') # Rotation angle of eigenbasis
theta = tf.Variable(thetap, dtype=tf.float32,
name='Theta') # Polar angle of rotation axis
phi = tf.Variable(phip, dtype=tf.float32,
name='Phi') # Azimuth angle of rotation axis
# everything else
with tf.name_scope('Auxiliary'):
mD = tf.diag([l1, l2, l3])
n0 = tf.sin(theta) * tf.cos(phi)
n1 = tf.sin(theta) * tf.sin(phi)
n2 = tf.cos(theta)
n = n0 * v0 + n1 * v1 + n2 * v2
nskew = n0 * nskew0 + n1 * nskew1 + n2 * nskew2
R = tf.cos(psi) * I + tf.sin(psi) * nskew + (
one - tf.cos(psi)) * tf.matmul(n, tf.transpose(n))
C = tf.matmul(R, tf.matmul(tf.matmul(mD, mD), tf.transpose(R)))
with tf.name_scope('Blurring'):
blurKernel = tf.reshape(
tf.exp(neg * tf.reduce_sum(q * tf.matmul(C, q), axis=0)),
shape=coherentEstimate.shape) * (l1 * l2 * l3) / twopi
tf_intens_f = tf.fft3d(
tf.cast(coherentEstimate, dtype=tf.complex64))
tf_blur_f = tf.fft3d(tf.cast(blurKernel, dtype=tf.complex64))
tf_prod_f = tf_intens_f * tf_blur_f
tf_prod_rolledIJK = tf.ifft3d(tf_prod_f)
tf_prod_rolledJK = tf.concat((tf_prod_rolledIJK[roll_[0]:, :, :],
tf_prod_rolledIJK[:roll_[0], :, :]),
axis=0)
tf_prod_rolledK = tf.concat((tf_prod_rolledJK[:, roll_[1]:, :],
tf_prod_rolledJK[:, :roll_[1], :]),
axis=1)
imgBlurred = tf.abs(
tf.concat((tf_prod_rolledK[:, :, roll_[2]:],
tf_prod_rolledK[:, :, :roll_[2]]),
axis=2))
# imgBlurred = tf.reshape(
# tf.nn.convolution(
# tf.reshape( coherentEstimate, [ coherentEstimate.shape[0], coherentEstimate.shape[1], coherentEstimate.shape[2], 1, 1 ] ),
# tf.reshape( blurKernel, [ blurKernel.shape[0], blurKernel.shape[1], blurKernel.shape[2], 1, 1 ] ),
# padding='SAME'
# ),
# shape=coherentEstimate.shape
# )
with tf.name_scope('Optimizer'):
poissonNLL = tf.reduce_mean(imgBlurred -
noisyData * tf.log(imgBlurred))
var_list = [l1, l2, l3, psi, theta, phi]
# poissonOptimizer = tf.train.MomentumOptimizer( learning_rate=1.e-5, momentum=0.99, use_nesterov=True, name='PoissonOptimize' )
poissonOptimizer = tf.train.AdagradOptimizer(
learning_rate=learning_rate, name='PoissonOptimize')
trainPoisson = poissonOptimizer.minimize(poissonNLL,
var_list=var_list)
currentGradients = [
n[0]
for n in poissonOptimizer.compute_gradients(poissonNLL,
var_list=var_list)
]
session = tf.Session()
session.run(tf.global_variables_initializer())
progress = []
this_grad = np.linalg.norm(np.array(session.run(currentGradients)))
normalizr = session.run(tf.reduce_sum(blurKernel))
checkpoint = session.run(var_list)
obj = session.run(poissonNLL)
checkpoint.extend([obj, this_grad, normalizr])
progress.append(checkpoint)
n_iter = 1
start = time.time()
# for n_iter in tqdm( list( range( max_iter ) ) ):
while n_iter < min_iterations or (this_grad > tol
and n_iter < max_iterations):
session.run(trainPoisson)
this_grad = np.linalg.norm(np.array(session.run(currentGradients)))
if n_iter % iterations_per_checkpoint == 0: # store progress every 'iterations_per_checkpoint' iterations.
normalizr = session.run(tf.reduce_sum(blurKernel))
checkpoint = session.run(var_list)
obj = session.run(poissonNLL)
checkpoint.extend([obj, this_grad, normalizr])
progress.append(checkpoint)
n_iter += 1
normalizr = session.run(tf.reduce_sum(blurKernel))
checkpoint = session.run(var_list)
obj = session.run(poissonNLL)
checkpoint.extend([obj, this_grad, normalizr])
progress.append(checkpoint)
stop = time.time()
if n_iter >= max_iterations - 1:
print('Warning: Max. number of iterations reached (%f).' %
max_iterations)
else:
print('Converged in %d iterations.' % n_iter)
print('Time taken = %f sec' % (stop - start))
blur_final, imgB_final, C_final = session.run([blurKernel, imgBlurred, C])
return progress, blur_final, imgB_final, recip_basis.T @ C_final @ recip_basis
def test_IFFT3D(self):
# only defined for gpu
if DEVICE == GPU:
t = tf.ifft3d(self.random(3, 4, 5, complex=True))
self.check(t)
def _ftconvolve(tensor1, tensor2):
return tf.real(tf.ifft3d(tf.fft3d(tensor1) * tf.fft3d(tensor2)))
