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_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 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 dimension_reduction_randomized_pca_stack_diff(avg_vol_key, data):
  """
  :TODO: documentation
  :TODO: rename
  :TODO: move somewhere else
  """
  vol_avg = get_mrc(avg_vol_key)
  vol_avg_fft = fftshift(fftn(vol_avg))

  mat = np.zeros((len(data), vol_avg.size))

  for i, (vk, mk, ang, loc) in enumerate(data):
    # load vol/mask.
    vol  = get_mrc(vk)
    mask = get_mrc(mk)

    # rotate vol and mask according to angle/loc.
    vol_r  = core.rotate_vol_pad_mean(vol, ang, loc)
    mask_r = core.rotate_mask(mask, ang)

    v_dif = np.real(ifftn(ifftshift(fftshift(fftn(vol_r) - vol_avg_fft) * mask_r)))

    mat[i, :] = v_dif.flatten()

  return mat 

def test_UnscaledFFT_3d(backend, M, N, K, B ):
    b = backend()

    # forward
    x = b.rand_array( (M*N*K, B) )
    y = b.rand_array( (M*N*K, B) )
    x_h = x.to_host().reshape( (M,N,K,B), order='F' )

    A = b.UnscaledFFT( (M,N,K), dtype=x.dtype )

    A.eval(y, x)

    y_exp = np.fft.fftn( x_h, axes=(0,1,2) )
    y_act = y.to_host().reshape( (M,N,K,B), order='F' )
    npt.assert_allclose(y_act, y_exp, rtol=1e-2)

    # adjoint
    x = b.rand_array( (M*N*K, B) )
    y = b.rand_array( (M*N*K, B) )
    x_h = x.to_host().reshape( (M,N,K,B), order='F' )

    A.H.eval(y, x)

    y_exp = np.fft.ifftn( x_h, axes=(0,1,2) ) * (M*N*K)
    y_act = y.to_host().reshape( (M,N,K,B), order='F' )
    npt.assert_allclose(y_act, y_exp, rtol=1e-2) 

def test_UnscaledFFT_2d(backend, M, N, B ):
    b = backend()

    # forward
    x = b.rand_array( (M*N, B) )
    y = b.rand_array( (M*N, B) )
    x_h = x.to_host().reshape( (M,N,B), order='F' )

    A = b.UnscaledFFT( (M,N), dtype=x.dtype )

    A.eval(y, x)

    y_exp = np.fft.fftn( x_h, axes=(0,1) )
    y_act = y.to_host().reshape( (M,N,B), order='F' )
    npt.assert_allclose(y_act, y_exp, rtol=1e-2)

    # adjoint
    x = b.rand_array( (M*N, B) )
    y = b.rand_array( (M*N, B) )
    x_h = x.to_host().reshape( (M,N,B), order='F' )

    A.H.eval(y, x)

    y_exp = np.fft.ifftn( x_h, axes=(0,1) ) * (M*N)
    y_act = y.to_host().reshape( (M,N,B), order='F' )
    npt.assert_allclose(y_act, y_exp, rtol=1e-2) 

def test_UnscaledFFT_1d(backend, M, B ):
    b = backend()

    # forward
    x = b.rand_array( (M, B) )
    y = b.rand_array( (M, B) )
    x_h = x.to_host().reshape( (M,B), order='F' )

    A = b.UnscaledFFT( (M,), dtype=x.dtype )

    A.eval(y, x)

    y_exp = np.fft.fftn( x_h, axes=(0,) )
    y_act = y.to_host().reshape( (M,B), order='F' )
    npt.assert_allclose(y_act, y_exp, rtol=1e-2)

    # adjoint
    x = b.rand_array( (M, B) )
    y = b.rand_array( (M, B) )
    x_h = x.to_host().reshape( (M,B), order='F' )

    A.H.eval(y, x)

    y_exp = np.fft.ifftn( x_h, axes=(0,) ) * M
    y_act = y.to_host().reshape( (M,B), order='F' )
    npt.assert_allclose(y_act, y_exp, rtol=1e-2) 

def test_CenteredFFT(backend, M, N, K, B ):
    from numpy.fft import fftshift, ifftshift, fftn, ifftn

    b = backend()
    A = b.FFTc( (M,N,K), dtype=np.dtype('complex64') )

    # forward
    ax = (0,1,2)
    x = b.rand_array( (M*N*K,B) )
    y = b.rand_array( (M*N*K,B) )
    x_h = x.to_host().reshape( (M,N,K,B), order='F' )

    A.eval(y, x)

    y_act = y.to_host().reshape( (M,N,K,B), order='F' )
    y_exp = fftshift( fftn( ifftshift(x_h, axes=ax), axes=ax, norm='ortho'), axes=ax)
    npt.assert_allclose(y_act, y_exp, rtol=1e-2)

    # adjoint
    x = b.rand_array( (M*N*K,B) )
    y = b.rand_array( (M*N*K,B) )
    x_h = x.to_host().reshape( (M,N,K,B), order='F' )

    A.H.eval(y, x)

    y_act = y.to_host().reshape( (M,N,K,B), order='F' )
    y_exp = fftshift( ifftn( ifftshift(x_h, axes=ax), axes=ax, norm='ortho'), axes=ax)
    npt.assert_allclose(y_act, y_exp, rtol=1e-2) 

def test_plan_call():
    for shape in tested_shapes:
        plan = Plan(
            input_array=ranf_unit_complex(shape),
            output_array=numpy.empty(shape, dtype=numpy.complex),
            direction=Direction.forward,
        )
        testing.assert_allclose(
            plan(),
            fft.fftn(plan.input_array)
        )
        testing.assert_allclose(
            plan(normalize=True),
            fft.fftn(plan.input_array) / plan.input_array.size
        )
        plan = Plan(
            input_array=ranf_unit_complex(shape),
            output_array=numpy.empty(shape, dtype=numpy.complex),
            direction=Direction.backward
        )
        testing.assert_allclose(
            plan(),
            fft.ifftn(plan.input_array) * plan.input_array.size
        )
        testing.assert_allclose(
            plan(normalize=True),
            fft.ifftn(plan.input_array)
        ) 

def ifftnc(x, axes):
    tmp = fft.fftshift(x, axes=axes)
    tmp = fft.ifftn(tmp, axes=axes)
    return fft.ifftshift(tmp, axes=axes) 

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:
            nthread = multiprocessing.cpu_count()
            K2_fft = pyfftw.builders.fftn(K2, axes=None, overwrite_input=False, planner_effort='FFTW_ESTIMATE',
                                       threads=nthread, 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=nthread, 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=nthread, 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 spectralcutoff(u, kappa, cx, cy, cz):
  nx = len(u[:,0,0])
  ny = len(u[0,:,0])
  nz = len(u[0,0,:])  
  nt= nx*ny*nz  
  uh = fftn(u)/nt  
  for i in range(0,nx):
    for j in range (0,ny):
      for k in range (0,nz):      
        rk = sqrt(cx*i*i + cy*j*j + cz*k*k)
        #rk = int(np.round(rk))
        if(rk >= kappa):
          uh[i,j,k] = 0.0
  ureal = ifftn(uh)*nt
  return ureal.real

#def spectralcutoff(u, kappa):
#  nx = len(u[:,0,0])
#  ny = len(u[0,:,0])
#  nz = len(u[0,0,:])  
#  nt= nx*ny*nz  
#  uh = fftn(u)/nt  
#  for i in range(0,nx):
#    for j in range (0,ny):
#      for k in range (0,nz):      
#        rk = sqrt(i*i + j*j + k*k)
#        #rk = int(np.round(rk))
#        if(rk >= kappa):
#          uh[i,j,k] = 0.0
#  ureal = ifftn(uh)*nt
#  return ureal.real 

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 

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 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 inv_fourier_transform(v):
    return ifftn(ifftshift(v)).real 

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

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 fconv(x, otf):
    return np.real(ifftn(fftn(x) * otf)) 

def masked_difference_given_vol_avg_fft(v, m, ang, loc, vol_avg_fft, vol_mask_avg, smoothing_gauss_sigma=0):
  """
  :TODO: documentation

  # mxu: IMPORTANT: correction of dimension reduction procedure according to the paper  Heumann12

  :param v:
  :param m:
  :param ang:
  :param loc:
  :param vol_avg_fft:
  :param vol_mask_avg:
  :param smoothing_gauss_sigma:
  """
  v_r = core.rotate_vol_pad_mean(v,ang,loc)
  m_r = core.rotate_mask(m, ang)

  v_r_fft = fftshift(fftn(v_r))
  v_r_msk_dif = np.real(ifftn(ifftshift( (v_r_fft - vol_avg_fft) * m_r * vol_mask_avg )))

  if smoothing_gauss_sigma > 0:
    # mxu:
    # when the subtomograms are very noisy, PCA alone may not work well.
    # Because PCA only process vectors but does consider spetial
    # structures. In such case we can make use spatial information as an
    # additional layer of constrain (or approsimation model) by apply
    # certain amount of gaussian smoothing
    v_r_msk_dif = filtering.gaussian_smoothing(v=v_r_msk_dif, sigma=smoothing_gauss_sigma)

  return (v_r_msk_dif, m_r) 

def gaussian_smoothing(v, sigma):
  """Smooth volume v, with spherical gaussian filter.

  Gaussian filter has single parameter sigma, which is symmetric standard
  deviation.

  Returns the volume v, after Gaussian smoothing is applied.
  """

  g = gauss_function(size=v.shape, sigma=sigma)
  # use ifftshift(g) to move center of gaussian to origin
  g_fft = fftn(ifftshift(g))
  v_conv = np.real( ifftn( fftn(v) * np.conj( g_fft ) ) )
  return v_conv 

def lowpass_smoothing(v, high, sigma=0):
  """Apply a low pass filter to volume v.

  Apply the filter in Fourier space.  Trim all frequencies higher then :high:.
  Frequencies are measured in pixel units.  If sigma > 0, Gaussian
  smoothing is applied frequencies on the edge.  So instead of a sharp cutoff
  at frequency :high:, there is decay from :high: to :high:+2*sigma.
  """
  dist_sq = grid_distance_sq_to_center(v)
  mask = (dist_sq <= high*high)

  v_fft = fftshift(fftn(v))
  v_fft_filt = v_fft * mask
  v_filt = np.real(ifftn(ifftshift(v_fft_filt)))

  if sigma > 0:
    return np.real(ifftn(ifftshift(spheremask(fftshift(fftn(v_filt)), high, sigma))))
  return v_filt


#def low_bandpass_smoothing(v, cutoff, order=5):
#  b,a = butter(order, Wn=cutoff)
#  # Default filter is lowpass.
#  # Wn is point of 1/sqrt(2) signal loss in passband.  Normalized on 0,1 where
#  # 1 is nyquist freq.
#
#  # Next apply the smoothing.
#  return lfilter(b,a,v) 

def vol_avg_fft_reduce(vm_in, vol_shape, n_vol, pass_dir):
  """
  Collect the output from the local vol_avg_fft_map tep.  Complete the sum, and
  write a final volume and mask average.

  :param vm_in:    The list of volume and maks pairs that are the subtotals
            computed so far.
  :param vol_shape:  3 element vector describing the shape of the volumes and
            masks.
  :param n_vol:    The total number of volumes that contributed to the
            subtotals we have, (necessary for the average calculation)
  :pass_dir:       Where to write the output.

  :returns: A tuple containing the volume and mask file names.  The data is
  computed and then written to disk in the pass_dir.

  """

  # TODO: what is the correct value?
  mask_threshold = 1.0

  fft_sum  = np.zeros(vol_shape, dtype=np.complex128, order='F')
  mask_sum = np.zeros(vol_shape, dtype=np.float64,  order='F')

  for vk,mk in vm_in:
    vol = np.load(vk)
    fft_sum += vol

    mask = np.load(mk)
    mask_sum += mask

  # This avoids RuntimeWarning: invalid value encountered in divide.
  flag_t = (mask_sum >= mask_threshold)
  tem_fft = np.zeros(vol_shape, dtype=np.complex128, order='F')
  tem_fft[flag_t] = fft_sum[flag_t] / mask_sum[flag_t]
  #tem_fft[mask_sum < mask_threshold] = 0

  vol_avg = np.asfortranarray(np.real(ifftn(ifftshift(tem_fft))))
  mask_avg = mask_sum / n_vol

  (v_fh, v_name) = tempfile.mkstemp(prefix='tm_tmp_vafrv_', suffix='.mrc', dir=pass_dir)
  (m_fh, m_name) = tempfile.mkstemp(prefix='tm_tmp_vafrm_', suffix='.mrc', dir=pass_dir)
  os.close(v_fh)
  os.close(m_fh)
  core.write_mrc(vol_avg, v_name)
  core.write_mrc(mask_avg, m_name)

  return v_name, m_name