Python numpy.fft.ifftn() Examples

def impute_aligned_vols(t, v, vm, normalize=None):
    assert (normalize is not None)
    if normalize:
        v = ((v - v.mean()) / v.std())
        if (t is not None):
            t = ((t - t.mean()) / t.std())
    if (t is None):
        return v
    t_f = NF.fftshift(NF.fftn(t))
    v_f = NF.fftshift(NF.fftn(v))
    v_f[(vm == 0)] = t_f[(vm == 0)]
    v_f_if = N.real(NF.ifftn(NF.ifftshift(v_f)))
    if normalize:
        v_f_if = ((v_f_if - v_f_if.mean()) / v_f_if.std())
    if N.all(N.isfinite(v_f_if)):
        return v_f_if
    else:
        print('warning: imputation failed')
        return v 

def average(dj, mask_count_threshold):
    vol_sum = None
    mask_sum = None
    for d in dj:
        v = IF.read_mrc_vol(d['subtomogram'])
        if (not N.all(N.isfinite(v))):
            raise Exception('error loading', d['subtomogram'])
        vm = IF.read_mrc_vol(d['mask'])
        v_r = GR.rotate_pad_mean(v, angle=d['angle'], loc_r=d['loc'])
        assert N.all(N.isfinite(v_r))
        vm_r = GR.rotate_mask(vm, angle=d['angle'])
        assert N.all(N.isfinite(vm_r))
        if (vol_sum is None):
            vol_sum = N.zeros(v_r.shape, dtype=N.float64, order='F')
        vol_sum += v_r
        if (mask_sum is None):
            mask_sum = N.zeros(vm_r.shape, dtype=N.float64, order='F')
        mask_sum += vm_r
    ind = (mask_sum >= mask_count_threshold)
    vol_sum_fft = NF.fftshift(NF.fftn(vol_sum))
    avg = N.zeros(vol_sum_fft.shape, dtype=N.complex)
    avg[ind] = (vol_sum_fft[ind] / mask_sum[ind])
    avg = N.real(NF.ifftn(NF.ifftshift(avg)))
    return {'v': avg, 'm': (mask_sum / len(dj)), } 

def run_test_c2c_impl(self, shape, axes, inverse=False, fftshift=False):
        shape = list(shape)
        shape[-1] *= 2 # For complex
        known_data = np.random.normal(size=shape).astype(np.float32).view(np.complex64)
        idata = bf.ndarray(known_data, space='cuda')
        odata = bf.empty_like(idata)
        fft = Fft()
        fft.init(idata, odata, axes=axes, apply_fftshift=fftshift)
        fft.execute(idata, odata, inverse)
        if inverse:
            if fftshift:
                known_data = np.fft.ifftshift(known_data, axes=axes)
            # Note: Numpy applies normalization while CUFFT does not
            norm = reduce(lambda a, b: a * b, [known_data.shape[d]
                                               for d in axes])
            known_result = gold_ifftn(known_data, axes=axes) * norm
        else:
            known_result = gold_fftn(known_data, axes=axes)
            if fftshift:
                known_result = np.fft.fftshift(known_result, axes=axes)
        x = (np.abs(odata.copy('system') - known_result) / known_result > RTOL).astype(np.int32)
        a = odata.copy('system')
        b = known_result
        compare(odata.copy('system'), known_result) 

def lct(x1, y1, t1, v, vol, snr):
    X = len(x1)
    Y = len(y1)
    Z = len(t1)
    S = 2
    slope = np.max(x1) / (np.max(t1) * v/2)
    slope = np.max(y1) / (np.max(t1) * v/2)
    psf, fpsf = getPSF(X, Y, Z, S, slope)
    mtx, mtxi = interpMtx(Z, S, 0, np.max(t1)*v)

    def pad_array(x, S, Z, X):
        return np.pad(x, ((S*Z//2, S*Z//2), (X//2, X//2), (Y//2, Y//2)), 'constant')

    def trim_array(x, S, Z, X):
        return x[S*(Z//2)+1:-S*(Z//2)+1, X//2+1:-X//2+1, Y//2+1:-Y//2+1]

    invpsf = np.conj(fpsf) / (abs(fpsf)**2 + 1 / snr)
    tdata = np.matmul(mtx, vol)
    fdata = fftn(pad_array(tdata, S, Z, X))
    tvol = abs(trim_array(ifftn(fdata * invpsf), S, Z, X))
    vol = np.matmul(mtxi, tvol)
    return vol 

def run_lct(self, meas, X, Y, Z, S, x_max, y_max):
        self.max_dist = int(self.T * self.v/2 / self.channels[0])
        slope_x = self.dx * x_max / (self.fend/self.B * self.max_dist) * (1 + ((self.fstart/self.B)/(self.fend/self.B))**2)
        slope_y = self.dy * y_max / (self.fend/self.B * self.max_dist) * (1 + ((self.fstart/self.B)/(self.fend/self.B))**2)
        
        # params
        psf, fpsf = lct.getPSF(X, Y, Z, S, slope_x, slope_y)
        mtx, mtxi = lct.interpMtx(Z, S, self.fstart/self.B * self.max_dist, self.fend/self.B * self.max_dist)

        def pad_array(x, S, Z, X, Y):
            return np.pad(x, ((S*Z//2, S*Z//2), (Y//2, Y//2), (X//2, X//2)), 'constant')

        def trim_array(x, S, Z, X, Y):
            return x[S*int(np.floor(Z/2))+1:-S*int(np.ceil(Z/2))+1, Y//2+1:-Y//2+1, X//2+1:-X//2+1]

        invpsf = np.conj(fpsf) / (abs(fpsf)**2 + 1 / self.snr)
        tdata = np.matmul(mtx, meas.reshape((Z, -1))).reshape((-1, Y, X))

        fdata = fftn(pad_array(tdata, S, Z, X, Y))
        tvol = abs(trim_array(ifftn(fdata * invpsf), S, Z, X, Y))
        out = np.matmul(mtxi, tvol.reshape((S*Z, -1))).reshape((-1, Y, X))

        return out 

def fconv(x, otf):
    return np.real(ifftn(fftn(x) * otf)) 

def convolve(arr1, arr2, dx=None, axes=None):
    """
    Performs a centred convolution of input arrays

    Parameters
    ----------
    arr1, arr2 : `numpy.ndarray`
        Arrays to be convolved. If dimensions are not equal then 1s are appended
        to the lower dimensional array. Otherwise, arrays must be broadcastable.
    dx : float > 0, list of float, or `None` , optional
        Grid spacing of input arrays. Output is scaled by
        `dx**max(arr1.ndim, arr2.ndim)`. default=`None` applies no scaling
    axes : tuple of ints or `None`, optional
        Choice of axes to convolve. default=`None` convolves all axes

    """
    if arr2.ndim > arr1.ndim:
        arr1, arr2 = arr2, arr1
        if axes is None:
            axes = range(arr2.ndim)
    arr2 = arr2.reshape(arr2.shape + (1,) * (arr1.ndim - arr2.ndim))

    if dx is None:
        dx = 1
    elif isscalar(dx):
        dx = dx ** (len(axes) if axes is not None else arr1.ndim)
    else:
        dx = prod(dx)

    arr1 = fftn(arr1, axes=axes)
    arr2 = fftn(ifftshift(arr2), axes=axes)
    out = ifftn(arr1 * arr2, axes=axes) * dx
    return require(out, requirements="CA") 

def ifftn(a, s=None, axes=None, norm=None, **_):
        return _ifftn(a, s, axes, norm) 

def plan_ifft(A, n=None, axis=None, norm=None, **_):
        """
        Plans an ifft for repeated use. Parameters are the same as for `pyfftw`'s `ifftn`
        which are, where possible, the same as the `numpy` equivalents.
        Note that some functionality is only possible when using the `pyfftw` backend.

        Parameters
        ----------
        A : `numpy.ndarray`, of dimension `d`
            Array of same shape to be input for the ifft
        n : iterable or `None`, `len(n) == d`, optional
            The output shape of ifft (default=`None` is same as `A.shape`)
        axis : `int`, iterable length `d`, or `None`, optional
            The axis (or axes) to transform (default=`None` is all axes)
        overwrite : `bool`, optional
            Whether the input array can be overwritten during computation
            (default=False)
        planner : {0, 1, 2, 3}, optional
            Amount of effort put into optimising Fourier transform where 0 is low
            and 3 is high (default=`1`).
        threads : `int`, `None`
            Number of threads to use (default=`None` is all threads)
        auto_align_input : `bool`, optional
            If `True` then may re-align input (default=`True`)
        auto_contiguous : `bool`, optional
            If `True` then may re-order input (default=`True`)
        avoid_copy : `bool`, optional
            If `True` then may over-write initial input (default=`False`)
        norm : {None, 'ortho'}, optional
            Indicate whether ifft is normalised (default=`None`)

        Returns
        -------
        plan : function
            Returns the inverse Fourier transform of `B`, `plan() == ifftn(B)`
        B : `numpy.ndarray`, `A.shape`
            Array which should be modified inplace for ifft to be computed. If
            possible, `B is A`.
        """
        return lambda: ifftn(A, n, axis, norm), A 

def _ifftn(a, s=None, axes=None):
        return  npfft.ifftn(a, s, axes).astype(a.dtype) 

def fftconv(a, b, axes=(0, 1), origin=None):
    """Multi-dimensional convolution via the Discrete Fourier Transform.

    Compute a multi-dimensional convolution via the Discrete Fourier
    Transform. Note that the output has a phase shift relative to the
    output of :func:`scipy.ndimage.convolve` with the default `origin`
    parameter.

    Parameters
    ----------
    a : array_like
      Input array
    b : array_like
      Input array
    axes : sequence of ints, optional (default (0, 1))
      Axes on which to perform convolution
    origin : sequence of ints or None optional (default None)
      Indices of centre of `a` filter. The default of None corresponds
      to a centre at 0 on all axes of `a`

    Returns
    -------
    ab : ndarray
      Convolution of input arrays, `a` and `b`, along specified `axes`
    """

    if np.isrealobj(a) and np.isrealobj(b):
        fft = rfftn
        ifft = irfftn
    else:
        fft = fftn
        ifft = ifftn
    dims = np.maximum([a.shape[i] for i in axes], [b.shape[i] for i in axes])
    af = fft(a, dims, axes)
    bf = fft(b, dims, axes)
    ab = ifft(af * bf, dims, axes)
    if origin is not None:
        ab = np.roll(ab, -np.array(origin), axis=axes)
    return ab 

def irfftn(a, s, axes=None):
    """Multi-dimensional inverse discrete Fourier transform for real input.

    Compute the inverse of the multi-dimensional discrete Fourier
    transform for real input. This function is a wrapper for
    :func:`pyfftw.interfaces.numpy_fft.irfftn`, with an interface similar
    to that of :func:`numpy.fft.irfftn`.

    Parameters
    ----------
    a : array_like
      Input array
    s : sequence of ints
      Shape of the output along each transformed axis (input is cropped
      or zero-padded to match). This parameter is not optional because,
      unlike :func:`ifftn`, the output shape cannot be uniquely
      determined from the input shape.
    axes : sequence of ints, optional (default None)
      Axes over which to compute the inverse DFT.

    Returns
    -------
    af : ndarray
      Inverse DFT of input array
    """

    return pyfftw.interfaces.numpy_fft.irfftn(
        a, s=s, axes=axes, overwrite_input=False,
        planner_effort=pyfftw_planner_effort, threads=pyfftw_threads) 

def ifftn(a, s=None, axes=None):
    """Multi-dimensional inverse discrete Fourier transform.

    Compute the multi-dimensional inverse discrete Fourier transform.
    This function is a wrapper for :func:`pyfftw.interfaces.numpy_fft.ifftn`,
    with an interface similar to that of :func:`numpy.fft.ifftn`.

    Parameters
    ----------
    a : array_like
      Input array (can be complex)
    s : sequence of ints, optional (default None)
      Shape of the output along each transformed axis (input is cropped
      or zero-padded to match).
    axes : sequence of ints, optional (default None)
      Axes over which to compute the inverse DFT.

    Returns
    -------
    af : complex ndarray
      Inverse DFT of input array
    """

    return pyfftw.interfaces.numpy_fft.ifftn(
        a, s=s, axes=axes, overwrite_input=False,
        planner_effort=pyfftw_planner_effort, threads=pyfftw_threads) 

def potential(self, q, steps):
        hx = steps[0]
        hy = steps[1]
        hz = steps[2]
        Nx = q.shape[0]
        Ny = q.shape[1]
        Nz = q.shape[2]
        out = np.zeros((2*Nx-1, 2*Ny-1, 2*Nz-1))
        out[:Nx, :Ny, :Nz] = q
        K1 = self.sym_kernel(q.shape, steps)
        K2 = np.zeros((2*Nx-1, 2*Ny-1, 2*Nz-1))
        K2[0:Nx, 0:Ny, 0:Nz] = K1
        K2[0:Nx, 0:Ny, Nz:2*Nz-1] = K2[0:Nx, 0:Ny, Nz-1:0:-1] #z-mirror
        K2[0:Nx, Ny:2*Ny-1,:] = K2[0:Nx, Ny-1:0:-1, :]        #y-mirror
        K2[Nx:2*Nx-1, :, :] = K2[Nx-1:0:-1, :, :]             #x-mirror
        t0 = time.time()
        if pyfftw_flag:
            nthreads = int(conf.OCELOT_NUM_THREADS)
            if nthreads < 1:
                nthreads = 1
            K2_fft = pyfftw.builders.fftn(K2, axes=None, overwrite_input=False, planner_effort='FFTW_ESTIMATE',
                                       threads=nthreads, auto_align_input=False, auto_contiguous=False, avoid_copy=True)
            out_fft = pyfftw.builders.fftn(out, axes=None, overwrite_input=False, planner_effort='FFTW_ESTIMATE',
                                          threads=nthreads, auto_align_input=False, auto_contiguous=False, avoid_copy=True)
            out_ifft = pyfftw.builders.ifftn(out_fft()*K2_fft(), axes=None, overwrite_input=False, planner_effort='FFTW_ESTIMATE',
                                          threads=nthreads, auto_align_input=False, auto_contiguous=False, avoid_copy=True)
            out = np.real(out_ifft())

        else:
            out = np.real(ifftn(fftn(out)*fftn(K2)))
        t1 = time.time()
        logger.debug('fft time:' + str(t1-t0) + ' sec')
        out[:Nx, :Ny, :Nz] = out[:Nx,:Ny,:Nz]/(4*pi*epsilon_0*hx*hy*hz)
        return out[:Nx, :Ny, :Nz] 

def translation_align__given_unshifted_fft(v1f, v2f):
    cor = fftshift( N.real( ifftn( v1f * N.conj(v2f) ) ) )

    mid_co = IVU.fft_mid_co(cor.shape)
    loc = N.unravel_index( cor.argmax(), cor.shape )

    return {'loc': (loc - mid_co), 'cor': cor[loc[0], loc[1], loc[2]]}




# for each angle, do a translation search 

def dog_smooth__large_map(v, s1, s2=None):

    if s2 is None:      s2 = s1 * 1.1       # the 1.1 is according to a DoG particle picking paper
    assert      s1 < s2

    size = v.shape


    pad_width = int(N.round(s2*2))
    vp = N.pad(array=v, pad_width=pad_width, mode='reflect')

    v_fft = fftn(vp).astype(N.complex64)
    del v;      GC.collect()


    g_small = difference_of_gauss_function(size=N.array([int(N.round(s2 * 4))]*3), sigma1=s1, sigma2=s2)
    assert      N.all(N.array(g_small.shape) <= N.array(vp.shape))       # make sure we can use CV.paste_to_whole_map()

    g = N.zeros(vp.shape)
    paste_to_whole_map(whole_map=g, vol=g_small, c=None)

    g_fft_conj = N.conj(   fftn(ifftshift(g)).astype(N.complex64)   )    # use ifftshift(g) to move center of gaussian to origin
    del g;      GC.collect()

    prod_t = (v_fft * g_fft_conj).astype(N.complex64)
    del v_fft;      GC.collect()
    del g_fft_conj;      GC.collect()

    prod_t_ifft = ifftn( prod_t ).astype(N.complex64)
    del prod_t;      GC.collect()

    v_conv = N.real( prod_t_ifft )
    del prod_t_ifft;      GC.collect()
    v_conv = v_conv.astype(N.float32)

    v_conv = v_conv[(pad_width+1):(pad_width+size[0]+1), (pad_width+1):(pad_width+size[1]+1), (pad_width+1):(pad_width+size[2]+1)]
    assert      size == v_conv.shape

    return v_conv 

def apply_ctf(v, ctf):
    vc = N.real(     ifftn(  ifftshift(  ctf * fftshift(fftn(v)) ) )     )           # convolute v with ctf
    return vc 

def inv_fourier_transform(v):
    return ifftn(ifftshift(v)).real 

def vol_avg(dj, op, img_db):
    if len(dj) < op['mask_count_threshold']:        return None

    vol_sum = None
    mask_sum = None

    # temporary collection of local volume, and mask.
    for d in dj:
        v = img_db[d['subtomogram']]
        vm = img_db[d['mask']]

        v_r = GR.rotate_pad_mean(v, angle=d['angle'], loc_r=d['loc']);
        assert N.all(N.isfinite(v_r))
        vm_r = GR.rotate_mask(vm, angle=d['angle']);
        assert N.all(N.isfinite(vm_r))

        if vol_sum is None:     vol_sum = N.zeros(v_r.shape, dtype=N.float64, order='F')
        vol_sum += v_r

        if mask_sum is None:        mask_sum = N.zeros(vm_r.shape, dtype=N.float64, order='F')
        mask_sum += vm_r

    ind = mask_sum >= op['mask_count_threshold']
    if ind.sum() <= 0:  return None

    vol_sum = NF.fftshift(NF.fftn(vol_sum))
    avg = N.zeros(vol_sum.shape, dtype=N.complex)
    avg[ind] = vol_sum[ind] / mask_sum[ind]
    avg = N.real(NF.ifftn(NF.ifftshift(avg)))

    return {'v': avg, 'm': mask_sum / float(len(dj))} 

def filter_given_curve(v, curve):
    grid = GV.grid_displacement_to_center(v.shape, GV.fft_mid_co(v.shape))
    rad = GV.grid_distance_to_center(grid)
    rad = N.round(rad).astype(N.int)
    b = N.zeros(rad.shape)
    for (i, a) in enumerate(curve):
        b[(rad == i)] = a
    vf = ifftn(ifftshift((fftshift(fftn(v)) * b)))
    vf = N.real(vf)
    return vf 

def ifftnc(x, axes, ortho=True):
    tmp = fft.fftshift(x, axes=axes)
    tmp = fft.ifftn(tmp, axes=axes, norm="ortho" if ortho else None)
    return fft.ifftshift(tmp, axes=axes) 

def ifftd(I, dims=None):

    # Compute fft
    if dims is None:
        X = ifftn(I)
    elif dims == 2:
        X = ifft2(I, axes=(0, 1))
    else:
        X = ifftn(I, axes=tuple(range(dims)))

    return X