Python pyfftw.n_byte_align函数代码示例

    def ShiftConvNumbaFFT(h, N, M, xdtype=np.complex_, powerof2=True):
        # implements Doppler filter:
        # y[n, p] = SUM_k (exp(2*pi*j*n*(k - (L-1))/N) * h[k]) * x[p - k]
        #         = SUM_k (exp(-2*pi*j*n*k/N) * s*[k]) * x[p - (L-1) + k]
        L = len(h)
        outlen = M + L - 1
        nfft = outlen
        if powerof2:
            nfft = pow2(nfft)

        dopplermat = np.exp(2*np.pi*1j*np.arange(N)[:, np.newaxis]*(np.arange(L) - (L - 1))/N)
        dopplermat.astype(np.result_type(h.dtype, np.complex64)) # cast to complex type with precision of h
        hbank = h*dopplermat
        # speed not critical here, just use numpy fft
        hbankpad = zero_pad(hbank, nfft)
        H = np.fft.fft(hbankpad) / nfft # divide by nfft b/c FFTW's ifft does not do this

        xcdtype = np.result_type(xdtype, np.complex64) # cast to complex type with precision of x
        xpad = pyfftw.n_byte_align(np.zeros(nfft, xcdtype), 16)
        X = pyfftw.n_byte_align(np.zeros(nfft, xcdtype), 16)
        xfft = pyfftw.FFTW(xpad, X, threads=_THREADS)

        ydtype = np.result_type(H.dtype, xcdtype)
        Y = pyfftw.n_byte_align_empty(H.shape, 16, ydtype)
        y = pyfftw.n_byte_align_empty(H.shape, 16, ydtype)
        ifft = pyfftw.FFTW(Y, y, direction='FFTW_BACKWARD', threads=_THREADS)

        xtype = numba.__getattribute__(str(np.dtype(xdtype)))

        #htype = numba.__getattribute__(str(H.dtype))
        #xctype = numba.__getattribute__(str(X.dtype))
        #ytype = numba.__getattribute__(str(Y.dtype))
        #@jit(argtypes=[htype[:, ::1], xctype[::1], ytype[:, ::1], xtype[::1]])
        #def fun(H, X, Y, x):
            #xpad[:M] = x
            #xfft.execute() # input is xpad, output is X
            #Y[:, :] = H*X # need expression optimized by numba but that writes into Y
            #ifft.execute() # input is Y, output is y

            #yc = np.array(y)[:, :outlen] # need a copy, which np.array provides
            #return yc

        #@dopplerbank_dec(h, N, M, nfft=nfft, H=H)
        #def shiftconv_numba_fft(x):
            #return fun(H, X, Y, x)

        #@jit(argtypes=[xtype[::1]])
        @jit
        def shiftconv_numba_fft(x):
            xpad[:M] = x
            xfft.execute() # input is xpad, output is X
            Y[:, :] = X*H # need expression optimized by numba but that writes into Y
            ifft.execute() # input is Y, output is y

            yc = np.array(y[:, :outlen]) # need a copy, which np.array provides
            return yc

        shiftconv_numba_fft = dopplerbank_dec(h, N, M, nfft=nfft, H=H)(shiftconv_numba_fft)

        return shiftconv_numba_fft

def SweepSpectraStridedInput(h, N, M, xdtype=np.complex_):
    # implements Doppler filter:
    # y[n, p] = SUM_k exp(2*pi*j*n*(k - (L-1))/N) * (h[k] * x[p - k])
    #         = SUM_k exp(-2*pi*j*n*k/N) * (s*[k] * x[p - (L-1) + k])
    L = len(h)
    outlen = M + L - 1
    # when N < L, still need to take FFT with nfft >= L so we don't lose data
    # then subsample to get our N points that we desire
    step = L // N + 1
    nfft = N*step

    hrev = h[::-1]
    xpad = np.zeros(M + 2*(L - 1), xdtype) # x[0] at xpad[L - 1]
    xshifted = np.lib.stride_tricks.as_strided(xpad,
                                               (outlen, L),
                                               (xpad.itemsize, xpad.itemsize))

    demodpad = np.zeros((outlen, nfft), np.result_type(xdtype, h.dtype, np.complex64))
    demodpad = pyfftw.n_byte_align(demodpad, 16)
    y = pyfftw.n_byte_align(np.zeros_like(demodpad), 16)
    fft = pyfftw.FFTW(demodpad, y, threads=_THREADS)

    @dopplerbank_dec(h, N, M)
    def sweepspectra_strided_input(x):
        xpad[(L - 1):outlen] = x
        np.multiply(xshifted, hrev, demodpad[:, :L])
        fft.execute() # input is demodpad, output is y
        yc = np.array(y[:, ::step].T) # we need a copy, which np.array provides
        return yc

    return sweepspectra_strided_input

def SweepSpectraStridedTapsMod(h, N, M, xdtype=np.complex_):
    """Doppler bank where the signal is downshift modulated before filtering."""
    # implements Doppler filter:
    # y[n, p] = SUM_k (h[p - k] * x[k]) * exp(-2*pi*j*n*k/N)
    #         = SUM_k (s*[k + (L-1) - p] * x[k]) * exp(-2*pi*j*n*k/N)
    L = len(h)
    # when N < M, still need to take FFT with nfft >= M so we don't lose data
    # then subsample to get our N points that we desire
    step = M // N + 1
    nfft = N*step

    hpad = np.hstack((np.zeros(M - 1, h.dtype), h, np.zeros(M - 1, h.dtype)))
    hmat = np.lib.stride_tricks.as_strided(hpad[(M - 1):], (M + L - 1, M),
                                           (hpad.itemsize, -hpad.itemsize))

    demodpad = np.zeros((M + L - 1, nfft), np.result_type(xdtype, h.dtype, np.complex64))
    demodpad = pyfftw.n_byte_align(demodpad, 16)
    y = pyfftw.n_byte_align(np.zeros_like(demodpad), 16)
    fft = pyfftw.FFTW(demodpad, y, threads=_THREADS)

    @dopplerbank_dec(h, N, M, hmat=hmat)
    def sweepspectra_strided_taps_mod(x):
        np.multiply(hmat, x, demodpad[:, :M])
        fft.execute() # input is demodpad, output is y
        yc = np.array(y[:, ::step].T) # need a copy, which np.array provides
        return yc

    return sweepspectra_strided_taps_mod

def SweepSpectraCython(h, N, M, xdtype=np.complex_):
    # implements Doppler filter:
    # y[n, p] = SUM_k exp(2*pi*j*n*(k - (L-1))/N) * (h[k] * x[p - k])
    #         = SUM_k exp(-2*pi*j*n*k/N) * (s*[k] * x[p - (L-1) + k])
    L = len(h)
    outlen = M + L - 1
    # when N < L, still need to take FFT with nfft >= L so we don't lose data
    # then subsample to get our N points that we desire
    step = L // N + 1
    nfft = N*step

    # ensure that h is C-contiguous as required by the Cython function
    h = np.asarray(h, order='C')

    # make sure xdtype is a dtype object
    xdtype = np.dtype(xdtype)

    demodpad = np.zeros((outlen, nfft), np.result_type(xdtype, h.dtype, np.complex64))
    demodpad = pyfftw.n_byte_align(demodpad, 16)
    y = pyfftw.n_byte_align(np.zeros_like(demodpad), 16)
    fft = pyfftw.FFTW(demodpad, y, threads=_THREADS)

    sweepspectra_cython = libdopplerbanks.SweepSpectraCython(h, demodpad, y, fft,
                                                             step, N, M, xdtype)
    sweepspectra_cython = dopplerbank_dec(h, N, M)(sweepspectra_cython)

    return sweepspectra_cython

    def SweepSpectraNumba(h, N, M, xdtype=np.complex_):
        # implements Doppler filter:
        # y[n, p] = SUM_k exp(2*pi*j*n*(k - (L-1))/N) * (h[k] * x[p - k])
        #         = SUM_k exp(-2*pi*j*n*k/N) * (s*[k] * x[p - (L-1) + k])
        L = len(h)
        outlen = M + L - 1
        # when N < L, still need to take FFT with nfft >= L so we don't lose data
        # then subsample to get our N points that we desire
        step = L // N + 1
        nfft = N*step

        hrev = h[::-1]
        xpad = np.zeros(M + 2*(L - 1), xdtype) # x[0] at xpad[L - 1]

        demodpad = np.zeros((outlen, nfft), np.result_type(xdtype, h.dtype, np.complex64))
        demodpad = pyfftw.n_byte_align(demodpad, 16)
        y = pyfftw.n_byte_align(np.zeros_like(demodpad), 16)
        fft = pyfftw.FFTW(demodpad, y, threads=_THREADS)

        xtype = numba.__getattribute__(str(np.dtype(xdtype)))

        #@jit(argtypes=[xtype[::1]])
        @jit
        def sweepspectra_numba(x):
            xpad[(L - 1):outlen] = x
            for p in range(outlen):
                demodpad[p, :L] = hrev*xpad[p:(p + L)]
            fft.execute() # input is demodpad, output is y
            yc = np.array(y[:, ::step].T) # we need a copy, which np.array provides
            return yc

        sweepspectra_numba = dopplerbank_dec(h, N, M)(sweepspectra_numba)

        return sweepspectra_numba

    def acquire(self, n):
        # Leave all array allocation beyond storage to subclasses

        # Set up counters
        self.n = n
        self.i = -1     # First thing it does is get incremented
        self.start_countup = 0

        # Set up storage
        self.storage = np.zeros((self.n, self.x), dtype=np.int32)

        # Set up fft arrays
        self.ft_in = fftw.n_byte_align(np.zeros((n, self.x), dtype=np.complex128), 16)
        self.ft_out = fftw.n_byte_align(np.zeros((n, self.x), dtype=np.complex128), 16)
        self.ift_in = fftw.n_byte_align(np.zeros((n, self.x), dtype=np.complex128), 16)
        self.ift_out = fftw.n_byte_align(np.zeros((n, self.x), dtype=np.complex128), 16)
        self.high_pass = np.zeros((n, self.x), dtype=np.complex128)
        self.high_pass[0,:] = 1
        self.fft = fftw.FFTW(self.ft_in, self.ft_out, axes=(0,), direction='FFTW_FORWARD',
                            flags=('FFTW_MEASURE',), threads=1, planning_timelimit=None)
        self.ifft = fftw.FFTW(self.ift_in, self.ift_out, axes=(0,), direction='FFTW_BACKWARD',
                            flags=('FFTW_MEASURE',), threads=1, planning_timelimit=None)
        # self.fft = fftw.FFTW(self.ft_in, self.ft_out, axes=(0), direction='FFTW_FORWARD')
        # self.ifft = fftw.FFTW(self.ift_in, self.ift_out, axes=(0), direction='FFTW_BACKWARD')
        

        # Get ready to calculate fps
        self.plot_times = []

        # Go
        self.new_image.connect(self.send_plot)
        self.running = True

def phrt(im, plan, method=4, nr_threads=2):
	"""Process a tomographic projection image with the selected phase retrieval algorithm.

	Parameters
	----------
	im : array_like
		Flat corrected image data as numpy array.

	plan : structure
		Structure with pre-computed data (see prepare_plan function)

	method : int 
		Phase retrieval filter {1 = TIE (default), 2 = CTF, 3 = CTF first-half sine, 4 = Quasiparticle, 
		5 = Quasiparticle first half sine}.

	nr_threads : int 
		Number of threads to be used in the computation of FFT by PyFFTW

	Credits
	-------
	Julian Moosmann, KIT (Germany) is acknowledged for this code
		
	"""
	# Extract plan values:
	dim0_o = plan['dim0']
	dim1_o = plan['dim1']
	n_pad0 = plan['npad0']
	n_pad1 = plan['npad1']
	marg0  = (n_pad0 - dim0_o) / 2
	marg1  = (n_pad1 - dim1_o) / 2
	filter = plan['filter']	
	
	# Pad image (if required):	
	im  = padImage(im, n_pad0, n_pad1) 

	# (Un)comment the following two lines to use PyFFTW:
	n_byte_align(im, simd_alignment) 
	im = fft2(im - 1, threads=nr_threads)			

	# (Un)comment the following line to use NumPy:	
	#im = fft2(im - 1)			

	# Apply phase retrieval filter:
	im = filter * im   

	# (Un)comment the following two lines to use PyFFTW:
	n_byte_align(im, simd_alignment)
	im = real(ifft2(im, threads=nr_threads))	

	# (Un)comment the following line to use NumPy:	
	#im = real(ifft2(im))	

	# Return the negative:
	im = - im

	# Return cropped output:
	return im[marg0:dim0_o + marg0, marg1:dim1_o + marg1]   

    def NumbaFFTW(h, M, xdtype=np.complex_, powerof2=True):
        L = len(h)
        outlen = M + L - 1
        nfft = outlen
        if powerof2:
            nfft = pow2(nfft)

        outdtype = np.result_type(h.dtype, xdtype)
        fftdtype = np.result_type(outdtype, np.complex64) # output is always complex, promote using smallest

        # speed not critical here, just use numpy fft
        # cast to outdtype so we use same type of fft as when transforming x
        hpad = zero_pad(h, nfft).astype(outdtype)
        if np.iscomplexobj(hpad):
            H = np.fft.fft(hpad)
        else:
            H = np.fft.rfft(hpad)
        H = (H / nfft).astype(fftdtype) # divide by nfft b/c FFTW's ifft does not do this

        xpad = pyfftw.n_byte_align(np.zeros(nfft, outdtype), 16) # outdtype so same type fft as h->H
        X = pyfftw.n_byte_align(np.zeros(len(H), fftdtype), 16) # len(H) b/c rfft may be used
        xfft = pyfftw.FFTW(xpad, X, threads=_THREADS)

        y = pyfftw.n_byte_align_empty(nfft, 16, outdtype)
        ifft = pyfftw.FFTW(X, y, direction='FFTW_BACKWARD', threads=_THREADS)

        xtype = numba.__getattribute__(str(np.dtype(xdtype)))
        outtype = numba.__getattribute__(str(outdtype))
        ffttype = numba.__getattribute__(str(fftdtype))

        #@jit(restype=outtype[::1],
            #argtypes=[outtype[::1], ffttype[::1], ffttype[::1], outtype[::1], xtype[::1]])
        #def filt(xpad, X, H, y, x):
            #xpad[:M] = x
            #xfft.execute() # input in xpad, result in X
            #X[:] = H*X
            #ifft.execute() # input in X, result in y
            #yc = y[:outlen].copy()
            #return yc

        #@filter_dec(h, M, nfft=nfft, H=H)
        #def numba_fftw(x):
            #return filt(xpad, X, H, y, x)

        #@jit(argtypes=[xtype[::1]])
        @jit
        def numba_fftw(x):
            xpad[:M] = x
            xfft.execute() # input in xpad, result in X
            X[:] = H*X # want expression that is optimized by numba but writes into X
            ifft.execute() # input in X, result in y
            yc = y[:outlen].copy()
            return yc

        numba_fftw = filter_dec(h, M, nfft=nfft, H=H)(numba_fftw)

        return numba_fftw

    def test_call_with_unaligned(self):
        '''Make sure the right thing happens with unaligned data.
        '''
        input_array = (numpy.random.randn(*self.input_array.shape) 
                + 1j*numpy.random.randn(*self.input_array.shape))

        output_array = self.fft(
                input_array=n_byte_align(input_array.copy(), 16)).copy()

        input_array = n_byte_align(input_array, 16)
        output_array = n_byte_align(output_array, 16)

        # Offset by one from 16 byte aligned to guarantee it's not
        # 16 byte aligned
        a = n_byte_align(input_array.copy(), 16)
        a__ = n_byte_align_empty(
                numpy.prod(a.shape)*a.itemsize+1, 16, dtype='int8')
        
        a_ = a__[1:].view(dtype=a.dtype).reshape(*a.shape)
        a_[:] = a

        # Create a different second array the same way
        b = n_byte_align(output_array.copy(), 16)
        b__ = n_byte_align_empty(
                numpy.prod(b.shape)*a.itemsize+1, 16, dtype='int8')
        
        b_ = b__[1:].view(dtype=b.dtype).reshape(*b.shape)
        b_[:] = a

        # Set up for the first array
        fft = FFTW(input_array, output_array)
        a_[:] = a
        output_array = fft().copy()

        # Check a_ is not aligned...
        self.assertRaisesRegex(ValueError, 'Invalid input alignment', 
                self.fft.update_arrays, *(a_, output_array))

        # and b_ too
        self.assertRaisesRegex(ValueError, 'Invalid output alignment', 
                self.fft.update_arrays, *(input_array, b_))
        
        # But it should still work with the a_
        fft(a_)

        # However, trying to update the output will raise an error
        self.assertRaisesRegex(ValueError, 'Invalid output alignment', 
                self.fft.update_arrays, *(input_array, b_))

        # Same with SIMD off
        fft = FFTW(input_array, output_array, flags=('FFTW_UNALIGNED',))
        fft(a_)
        self.assertRaisesRegex(ValueError, 'Invalid output alignment', 
                self.fft.update_arrays, *(input_array, b_))

    def _create_ifft_plan(self):
        x = pyfftw.n_byte_align(
            np.zeros((self.P, self.N), self.xydtype),
            pyfftw.simd_alignment
        )
        ifft_out = pyfftw.n_byte_align(
            np.zeros_like(x),
            pyfftw.simd_alignment
        )
        ifft = pyfftw.FFTW(x, ifft_out, direction='FFTW_BACKWARD')

        return ifft

    def _create_fft_plan(self):
        delaymult_out = pyfftw.n_byte_align(
            np.zeros((self.P, self.N), self.xydtype),
            pyfftw.simd_alignment
        )
        fft_out = pyfftw.n_byte_align(
            np.zeros_like(delaymult_out),
            pyfftw.simd_alignment
        )
        fft = pyfftw.FFTW(delaymult_out, fft_out)

        return fft

    def test_update_data_with_alignment_error(self):
        in_shape = self.input_shapes['2d']
        out_shape = self.output_shapes['2d']

        byte_error = 1
        
        axes=(-1,)
        a, b = self.create_test_arrays(in_shape, out_shape)

        a = n_byte_align(a, 16)
        b = n_byte_align(b, 16)

        fft, ifft = self.run_validate_fft(a, b, axes)
        
        a, b = self.create_test_arrays(in_shape, out_shape)

        # Offset from 16 byte aligned to guarantee it's not
        # 16 byte aligned
        a__ = n_byte_align_empty(
                numpy.prod(in_shape)*a.itemsize+byte_error, 
                16, dtype='int8')
        
        a_ = (a__[byte_error:]
                .view(dtype=self.input_dtype).reshape(*in_shape))
        a_[:] = a 
        
        b__ = n_byte_align_empty(
                numpy.prod(out_shape)*b.itemsize+byte_error, 
                16, dtype='int8')
        
        b_ = (b__[byte_error:]
                .view(dtype=self.output_dtype).reshape(*out_shape))
        b_[:] = b
     
        with self.assertRaisesRegex(ValueError, 'Invalid output alignment'):
            self.run_validate_fft(a, b_, axes, fft=fft, ifft=ifft, 
                    create_array_copies=False)

        with self.assertRaisesRegex(ValueError, 'Invalid input alignment'):
            self.run_validate_fft(a_, b, axes, fft=fft, ifft=ifft, 
                    create_array_copies=False)

        # Should also be true for the unaligned case
        fft, ifft = self.run_validate_fft(a, b, axes, 
                force_unaligned_data=True)

        with self.assertRaisesRegex(ValueError, 'Invalid output alignment'):
            self.run_validate_fft(a, b_, axes, fft=fft, ifft=ifft, 
                    create_array_copies=False)

        with self.assertRaisesRegex(ValueError, 'Invalid input alignment'):
            self.run_validate_fft(a_, b, axes, fft=fft, ifft=ifft, 
                    create_array_copies=False)

    def test_call_with_keyword_output_update(self):
        """Test the class call with a keyword output update.
        """
        output_array = n_byte_align(
            (numpy.random.randn(*self.output_array.shape) + 1j * numpy.random.randn(*self.output_array.shape)),
            simd_alignment,
        )

        returned_output_array = self.fft(output_array=n_byte_align(output_array.copy(), simd_alignment)).copy()

        self.update_arrays(self.input_array, output_array)
        self.fft.execute()

        self.assertTrue(numpy.alltrue(returned_output_array == output_array))

def tiehom(im, plan, nr_threads=2):
	"""Process a tomographic projection image with the TIE-HOM (Paganin's) phase retrieval algorithm.

	Parameters
	----------
	im : array_like
		Flat corrected image data as numpy array.

	plan : structure
		Structure with pre-computed data (see tiehom_plan function).

	nr_threads : int 
		Number of threads to be used in the computation of FFT by PyFFTW (default = 2).
		
	"""
	# Extract plan values:
	dim0_o = plan['dim0']
	dim1_o = plan['dim1']
	n_pad0 = plan['npad0']
	n_pad1 = plan['npad1']
	marg0 = (n_pad0 - dim0_o) / 2
	marg1 = (n_pad1 - dim1_o) / 2
	den = plan['den']
	mu = plan['mu']

	# Pad image (if required):	
	im = padImage(im, n_pad0, n_pad1) 

	# (Un)comment the following two lines to use PyFFTW:
	n_byte_align(im, simd_alignment) 
	im = rfft2(im, threads=nr_threads)			

	# (Un)comment the following line to use NumPy:
	#im = rfft2(im)

	# Apply Paganin's (pre-computed) formula:
	im = im / den
		
	# (Un)comment the following two lines to use PyFFTW:
	n_byte_align(im, simd_alignment)
	im = irfft2(im, threads=nr_threads)
		
	# (Un)comment the following line to use NumPy:
	#im = irfft2(im)
				
	im = im.astype(float32)		
	im = -1 / mu * nplog(im)    		

	# Return cropped output:
	return im[marg0:dim0_o + marg0, marg1:dim1_o + marg1] 

    def test_call_with_keyword_input_update(self):
        '''Test the class call with a keyword input update.
        '''
        input_array = n_byte_align(
                numpy.random.randn(*self.input_array.shape) 
                    + 1j*numpy.random.randn(*self.input_array.shape), 16)

        output_array = self.fft(
            input_array=n_byte_align(input_array.copy(), 16)).copy()

        self.update_arrays(input_array, self.output_array)
        self.fft.execute()

        self.assertTrue(numpy.alltrue(output_array == self.output_array))