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)))