def _do_fft(self, input, direct): """ Does the fft using fftw3f It caches the plan and result of the fft so that when you do it again it will be faster This means you are limited to only using a few arrays before you exhaust the machines memory see also release_plan """ key = input.ctypes.data if key in self.plans: output, plan = self.plans[key] plan.execute() return output # we've never done an fft of this array if len(list(self.plans.keys())) > 6: print("fftw3 is making too much cache") print("needs to be programmed to reuse the same arrays") if direct == 'forward': assert input.dtype == np.float32 output = np.zeros(input.shape, np.complex64) else: assert input.dtype == np.complex64 output = np.zeros(input.shape, np.float32) fp = fftw3f.Plan(inarray=input, outarray=output, direction=direct, flags=['estimate']) fp.execute() self.plans[input.ctypes.data] = output, fp return output
def __init__(self, u,v,k, lamb = 488, n=1.51): #print k**2 #m = (u**2 + v**2) <= (n/lamb**2) #self.propFac = fftw3f.create_aligned_array(u.shape, 'complex64') #self.propFac = 1j*8*pi*sqrt(np.maximum((n/lamb)**2 - (u**2 + v**2), 0)) #self.propFac = ((2*pi*n/lamb)*sqrt(np.maximum(1 - (u**2 + v**2), 0))).astype('f') self.propFac = ((2*pi*n/lamb)*cos(.5*pi*sqrt((u**2 + v**2)))).astype('f') self.pfm =(self.propFac > 0).astype('f') self._F = fftw3f.create_aligned_array(u.shape, 'complex64') self._f = fftw3f.create_aligned_array(u.shape, 'complex64') #print('Creating plans for FFTs - this might take a while') #calculate plans for other ffts self._plan_f_F = fftw3f.Plan(self._f, self._F, 'forward', flags = FFTWFLAGS, nthreads=NTHREADS) self._plan_F_f = fftw3f.Plan(self._F, self._f, 'backward', flags = FFTWFLAGS, nthreads=NTHREADS) #self._plan_F_f = fftw3f.Plan(self._F, self._f, 'backward', flags = FFTWFLAGS, nthreads=NTHREADS) fftwWisdom.save_wisdom()
def psf_calc(self, psf, data_size): '''Precalculate the OTF etc...''' g = resizePSF(psf, data_size) #normalise g /= g[:, :, g.shape[2] / 2].sum() #keep track of our data shape self.height = data_size[0] self.width = data_size[1] self.depth = data_size[2] self.shape = data_size self.FTshape = list( self.shape[:-1] ) #[self.shape[0], self.shape[1], self.shape[2]/2 + 1] self.FTshape[-1] = self.FTshape[-1] / 2 + 1 self.g = g.astype('f4') self.g2 = 1.0 * self.g[::-1, ::-1, :] self.H = [] self.Ht = [] for i in range(self.shape[-1]): #allocate memory self.H.append( fftw3f.create_aligned_array(self.FTshape, 'complex64')) self.Ht.append( fftw3f.create_aligned_array(self.FTshape, 'complex64')) #create plans & calculate OTF and conjugate transformed OTF fftw3f.Plan(1.0 * self.g[:, :, i], self.H[i], 'forward')() fftw3f.Plan(1.0 * self.g2[:, :, i], self.Ht[i], 'forward')() self.Ht[i] /= g[:, :, i].size self.H[i] /= g[:, :, i].size
def __init__(self, ps, vox, PSSize): ps = ps.max(2) ps = ps - ps.min() ps = ps * scipy.signal.hanning( ps.shape[0])[:, None] * scipy.signal.hanning(ps.shape[1])[None, :] ps = ps / ps.sum() #PSFFileName = PSFFilename pw = (numpy.array(PSSize) - ps.shape) / 2. pw1 = numpy.floor(pw) pw2 = numpy.ceil(pw) self.cachedPSF = pad.with_constant(ps, ((pw2[0], pw1[0]), (pw2[1], pw1[1])), (0, )) self.cachedOTFH = (ifftn(self.cachedPSF) * self.cachedPSF.size).astype('complex64') self.cachedOTF2 = (self.cachedOTFH * fftn(self.cachedPSF)).astype('complex64') self.weinerFT = fftw3.create_aligned_array(self.cachedOTFH.shape, 'complex64') self.weinerR = fftw3.create_aligned_array(self.cachedOTFH.shape, 'float32') self.planForward = fftw3.Plan(self.weinerR, self.weinerFT, flags=FFTWFLAGS, nthreads=NTHREADS) self.planInverse = fftw3.Plan(self.weinerFT, self.weinerR, direction='reverse', flags=FFTWFLAGS, nthreads=NTHREADS) fftwWisdom.save_wisdom() self.otf2mean = self.cachedOTF2.mean()
def psf_calc(self, psf, data_size, subsamp=1): '''Precalculate the OTF etc...''' g = resizePSF(psf, data_size) self.subsamp = subsamp #keep track of our data shape self.height = data_size[0] self.width = data_size[1] self.depth = data_size[2] self.shape = data_size self.FTshape = [self.shape[0], self.shape[1], self.shape[2] / 2 + 1] self.g = g.astype('f4') self.g2 = 1.0 * self.g[::-1, ::-1, ::-1] self.gs = self.g[::subsamp, ::subsamp, :] self.gs2 = self.g2[::subsamp, ::subsamp, :] #allocate memory self.H = fftw3f.create_aligned_array(self.FTshape, 'complex64') self.Ht = fftw3f.create_aligned_array(self.FTshape, 'complex64') #self.f = zeros(self.shape, 'f4') #self.res = zeros(self.shape, 'f4') #self.S = zeros((size(self.f), 3), 'f4') #self._F = fftw3f.create_aligned_array(self.FTshape, 'complex64') #self._r = fftw3f.create_aligned_array(self.shape, 'f4') #S0 = self.S[:,0] #create plans & calculate OTF and conjugate transformed OTF fftw3f.Plan(self.g, self.H, 'forward')() fftw3f.Plan(self.g2, self.Ht, 'forward')() self.Ht /= g.size self.H /= g.size
def resizePSF(psf, data_size): if not psf.shape == data_size: #Expand PSF to data size by fourier domain interpolation print('Resizing PSF to match data size') g_ = fftw3f.create_aligned_array(data_size, 'complex64') H_ = fftw3f.create_aligned_array(data_size, 'complex64') sx, sy, sz = numpy.array(data_size).astype('f') / psf.shape #print sx, sy, sz OT = fftshift(fftn( fftshift(psf))) #don't bother with FFTW here as raw PSF is small if data_size[2] > 1: pr = ndimage.zoom(OT.real, [sx, sy, sz], order=1) pi = ndimage.zoom(OT.imag, [sx, sy, sz], order=1) else: #special case for 2D pr = ndimage.zoom(OT.real.squeeze(), [sx, sy], order=1).reshape(data_size) pi = ndimage.zoom(OT.imag.squeeze(), [sx, sy], order=1).reshape(data_size) H_[:] = ifftshift(pr + 1j * pi) fftw3f.Plan(H_, g_, 'backward')() #View3D(psf) #View3D(H_) #View3D(OT) #View3D(pr) g = ifftshift(g_.real) print('PSF resizing complete') else: g = psf #View3D(psf) #View3D(fftshift(numpy.abs(H_))) #View3D(fftshift(numpy.angle(H_))) #View3D(g) return g / g.sum()
def __init__(self, shape, float_type=None, options = None): """ Construct an instance of FFTTasks. Parameters ---------- shape : tuple Specify array shape for FFT. float_type : {None, 'single', 'double'} Specify floating point type. options : {None, `iocbio.utils.Options`} Specify command line options: options.float_type options.fftw_plan_flags options.fftw_threads See also -------- iocbio.ops.fft_tasks """ if VERBOSE>9: print 'Entering %s.__init__' % (self.__class__.__name__) if options is None: options = Options() flags = [options.get(fftw_plan_flags='estimate')] if float_type is None: float_type = options.get(float_type='single') self.load_wisdoms() self.shape = shape self.float_type = float_type threads = getattr(options, 'fftw_threads', 1) if float_type=='single': import fftw3f as fftw # hint: on failure to import select double float type cache = numpy.empty(shape, numpy.complex64) self.float_dtype = numpy.float32 self.complex_dtype = numpy.complex64 elif float_type=='double': import fftw3 as fftw cache = numpy.empty(shape, numpy.complex128) self.float_dtype = numpy.float64 self.complex_dtype = numpy.complex128 else: raise NotImplementedError (`float_type`) self._cache = cache self.fftw = fftw if VERBOSE: print 'Computing fftw wisdom (flags=%s, threads=%s, shape=%s, float=%s),'\ ' be patient, it may take a while..'\ % (flags, threads, shape, float_type), self._fft_plan = fftw.Plan(cache, cache, direction='forward', flags=flags, nthreads=threads) self._ifft_plan = fftw.Plan(cache, cache, direction='backward', flags=flags, nthreads=threads) if VERBOSE: print 'done' self.save_wisdoms() self.convolve_kernel_fourier = None
def get_optimal_fft_size(cls, size, return_speedup=False, max_nof_tries=None): """Compute optimal FFT size from a given size. Usually optimal FFT size is a power of two but on the other hand, achieving power of two may be memory expensive and there may exist sizes that give more efficient FFT computation. For example, if the input size is 65 then extending the FFT size to 128 is less efficient compared to extending the size, say, to 66. The method runs number of fft transforms (up-to the next power of 2) and used the actual FLOPS for finding optimal FFT size. So, the initial call may take some time but the results will be cached for subsequent calls. Note that if size>1024 then the time spent of computing fft up to sizes <=2048 can be considerable. To restrict this optimization, specify ``max_nof_tries``. Parameters ---------- size : int Specify estimate for FFT size. return_speedup : bool When True then return speed up factor. Returns ------- optimal_size : int Optimal FFT size. speedup : float Speed up factor of using optimal size. Returned only if ``return_speedup`` is True. See also -------- iocbio.ops.fft_tasks """ if VERBOSE>9: print 'Entering %s.get_optimal_fft_size' % (cls.__name__) if size==2**int(numpy.log2(size)): if return_speedup: return size, 1 return size try: import fftw3f as fftw complex_dtype = numpy.complex64 except ImportError: import fftw3 as fftw complex_dtype = numpy.complex128 max_size = 2**int(numpy.log2(size)+1) if max_nof_tries is not None: max_size = min (size+max_nof_tries, max_size) min_flops = 1e9 optimal_size = size flops_cache = cls.flops_cache for sz in range (size, max_size + 1): flops = flops_cache.get(sz) if flops is None: cache = numpy.empty((sz,), dtype=complex_dtype) plan = fftw.Plan(cache, cache, direction='forward', flags=['estimate']) iplan = fftw.Plan(cache, cache, direction='backward', flags=['estimate']) flops = sum(plan.get_flops())+ sum(iplan.get_flops()) fftw.destroy_plan (plan) fftw.destroy_plan (iplan) flops_cache[sz] = flops if flops < min_flops: min_flops = flops optimal_size = sz if return_speedup: return optimal_size, flops_cache[size] / flops_cache[optimal_size] return optimal_size
def psf_calc(self, psf, data_size): '''Precalculate the OTF etc...''' # pw = (numpy.array(data_size) - psf.shape)/2. # pw1 = numpy.floor(pw) # pw2 = numpy.ceil(pw) # # g = psf/psf.sum() # # #work out how we're going to need to pad to get the PSF the same size as our data # if pw1[0] < 0: # if pw2[0] < 0: # g = g[-pw1[0]:pw2[0]] # else: # g = g[-pw1[0]:] # # pw1[0] = 0 # pw2[0] = 0 # # if pw1[1] < 0: # if pw2[1] < 0: # g = g[-pw1[1]:pw2[1]] # else: # g = g[-pw1[1]:] # # pw1[1] = 0 # pw2[1] = 0 # # if pw1[2] < 0: # if pw2[2] < 0: # g = g[-pw1[2]:pw2[2]] # else: # g = g[-pw1[2]:] # # pw1[2] = 0 # pw2[2] = 0 # # # #do the padding # #g = pad.with_constant(g, ((pw2[0], pw1[0]), (pw2[1], pw1[1]),(pw2[2], pw1[2])), (0,)) # g_ = fftw3f.create_aligned_array(data_size, 'float32') # g_[:] = 0 # #print g.shape, g_.shape, g_[pw2[0]:-pw1[0], pw2[1]:-pw1[1], pw2[2]:-pw1[2]].shape # if pw1[2] == 0: # g_[pw2[0]:-pw1[0], pw2[1]:-pw1[1], pw2[2]:] = g # else: # g_[pw2[0]:-pw1[0], pw2[1]:-pw1[1], pw2[2]:-pw1[2]] = g # #g_[pw2[0]:-pw1[0], pw2[1]:-pw1[1], pw2[2]:-pw1[2]] = g # g = g_ print psf.sum() g = resizePSF(psf, data_size) print g.sum() #keep track of our data shape self.height = data_size[0] self.width = data_size[1] self.depth = data_size[2] self.shape = data_size print('Calculating OTF') FTshape = [self.shape[0], self.shape[1], self.shape[2] / 2 + 1] self.g = g.astype('f4') self.g2 = 1.0 * self.g[::-1, ::-1, ::-1] #allocate memory self.H = fftw3f.create_aligned_array(FTshape, 'complex64') self.Ht = fftw3f.create_aligned_array(FTshape, 'complex64') #self.f = zeros(self.shape, 'f4') #self.res = zeros(self.shape, 'f4') #self.S = zeros((size(self.f), 3), 'f4') self._F = fftw3f.create_aligned_array(FTshape, 'complex64') self._r = fftw3f.create_aligned_array(self.shape, 'f4') #S0 = self.S[:,0] #create plans & calculate OTF and conjugate transformed OTF fftw3f.Plan(self.g, self.H, 'forward')() fftw3f.Plan(self.g2, self.Ht, 'forward')() self.Ht /= g.size self.H /= g.size print('Creating plans for FFTs - this might take a while') #calculate plans for other ffts self._plan_r_F = fftw3f.Plan(self._r, self._F, 'forward', flags=FFTWFLAGS, nthreads=NTHREADS) self._plan_F_r = fftw3f.Plan(self._F, self._r, 'backward', flags=FFTWFLAGS, nthreads=NTHREADS) fftwWisdom.save_wisdom() print('Done planning')