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