def fft_poisson(rho, dx, dy, dz):
    """
    It is assumed that the rho array is defined such that:
    axis 0 - variation along x
    axis 1 - variation along y
    axis 2 - variation along z

    This FFT solver will work only when run using periodic boundary conditions.

    Additionally, it should be noted that the physical points should be passed
    to this function. Passing the non-physical(ghost) zones will generate
    erroneous results

    NOTE: The density that is passed to this function is the charge density.
    """
    k_x = np.fft.fftfreq(rho.shape[0], dx)
    k_y = np.fft.fftfreq(rho.shape[1], dy)
    k_z = np.fft.fftfreq(rho.shape[2], dz)

    k_y, k_x, k_z = np.meshgrid(k_y, k_x, k_z)

    k_x = af.to_array(k_x)
    k_y = af.to_array(k_y)
    k_z = af.to_array(k_z)

    rho_hat = af.fft3(rho)

    potential_hat = rho_hat / (4 * np.pi**2 * (k_x**2 + k_y**2 + k_z**2))

    potential_hat[0, 0, 0] = 0

    Ex_hat = -1j * 2 * np.pi * k_x * potential_hat
    Ey_hat = -1j * 2 * np.pi * k_y * potential_hat
    Ez_hat = -1j * 2 * np.pi * k_z * potential_hat

    Ex = af.ifft3(Ex_hat)
    Ey = af.ifft3(Ey_hat)
    Ez = af.ifft3(Ez_hat)

    return (Ex, Ey, Ez)
示例#2
0
def __fftn__(a, s, axes, direction='forward'):
    if len(s) != 3 and len(s) != 2 and len(s) != 1:
        raise NotImplementedError
    if axes is not None:
        raise NotImplementedError
    if(direction == 'forward'):
        if len(s) == 3:
            fa = arrayfire.fft3(a.d_array, s[2], s[1], s[0])
        elif len(s) == 2:
            fa = arrayfire.fft2(a.d_array, s[1], s[0])
        elif len(s) == 1:
            fa = arrayfire.fft(a.d_array, s[0])
    elif direction == 'inverse':
        if len(s) == 3:
            fa = arrayfire.ifft3(a.d_array, s[2], s[1], s[0])
        elif len(s) == 2:
            fa = arrayfire.ifft2(a.d_array, s[1], s[0])
        elif len(s) == 1:
            fa = arrayfire.ifft(a.d_array, s[0])
    else:
        raise ValueError('Wrong FFT direction')
    return ndarray(pu.af_shape(fa), dtype=pu.typemap(fa.dtype()), af_array=fa)
示例#3
0
文件: fft.py 项目: daurer/afnumpy
def __fftn__(a, s, axes, direction='forward'):
    if len(s) != 3 and len(s) != 2 and len(s) != 1:
        raise NotImplementedError
    if axes is not None:
        raise NotImplementedError
    if(direction == 'forward'):
        if len(s) == 3:
            fa = arrayfire.fft3(a.d_array, s[2], s[1], s[0])
        elif len(s) == 2:
            fa = arrayfire.fft2(a.d_array, s[1], s[0])
        elif len(s) == 1:
            fa = arrayfire.fft(a.d_array, s[0])
    elif direction == 'inverse':
        if len(s) == 3:
            fa = arrayfire.ifft3(a.d_array, s[2], s[1], s[0])
        elif len(s) == 2:
            fa = arrayfire.ifft2(a.d_array, s[1], s[0])
        elif len(s) == 1:
            fa = arrayfire.ifft(a.d_array, s[0])
    else:
        raise ValueError('Wrong FFT direction')
    return ndarray(a.shape, dtype=pu.typemap(fa.dtype()), af_array=fa)
af.display(af.dft(a))
af.display(af.real(af.ifft(af.fft(a))))
af.display(af.real(af.idft(af.dft(a))))

a = af.randu(4, 4)
af.display(a)

af.display(af.fft2(a))
af.display(af.dft(a))
af.display(af.real(af.ifft2(af.fft2(a))))
af.display(af.real(af.idft(af.dft(a))))

a = af.randu(4, 4, 2)
af.display(a)

af.display(af.fft3(a))
af.display(af.dft(a))
af.display(af.real(af.ifft3(af.fft3(a))))
af.display(af.real(af.idft(af.dft(a))))

a = af.randu(10, 1)
b = af.randu(3, 1)
af.display(af.convolve1(a, b))
af.display(af.fft_convolve1(a, b))
af.display(af.convolve(a, b))
af.display(af.fft_convolve(a, b))

a = af.randu(5, 5)
b = af.randu(3, 3)
af.display(af.convolve2(a, b))
af.display(af.fft_convolve2(a, b))
示例#5
0
def simple_signal(verbose=False):
    display_func = _util.display_func(verbose)
    print_func   = _util.print_func(verbose)

    a = af.randu(10, 1)
    pos0 = af.randu(10) * 10
    display_func(af.approx1(a, pos0))

    a = af.randu(3, 3)
    pos0 = af.randu(3, 3) * 10
    pos1 = af.randu(3, 3) * 10

    display_func(af.approx2(a, pos0, pos1))

    a = af.randu(8, 1)
    display_func(a)

    display_func(af.fft(a))
    display_func(af.dft(a))
    display_func(af.real(af.ifft(af.fft(a))))
    display_func(af.real(af.idft(af.dft(a))))

    a = af.randu(4, 4)
    display_func(a)

    display_func(af.fft2(a))
    display_func(af.dft(a))
    display_func(af.real(af.ifft2(af.fft2(a))))
    display_func(af.real(af.idft(af.dft(a))))

    a = af.randu(4, 4, 2)
    display_func(a)

    display_func(af.fft3(a))
    display_func(af.dft(a))
    display_func(af.real(af.ifft3(af.fft3(a))))
    display_func(af.real(af.idft(af.dft(a))))

    a = af.randu(10, 1)
    b = af.randu(3, 1)
    display_func(af.convolve1(a, b))
    display_func(af.fft_convolve1(a, b))
    display_func(af.convolve(a, b))
    display_func(af.fft_convolve(a, b))

    a = af.randu(5, 5)
    b = af.randu(3, 3)
    display_func(af.convolve2(a, b))
    display_func(af.fft_convolve2(a, b))
    display_func(af.convolve(a, b))
    display_func(af.fft_convolve(a, b))

    a = af.randu(5, 5, 3)
    b = af.randu(3, 3, 2)
    display_func(af.convolve3(a, b))
    display_func(af.fft_convolve3(a, b))
    display_func(af.convolve(a, b))
    display_func(af.fft_convolve(a, b))


    b = af.randu(3, 1)
    x = af.randu(10, 1)
    a = af.randu(2, 1)
    display_func(af.fir(b, x))
    display_func(af.iir(b, a, x))
示例#6
0
def simple_signal(verbose=False):
    display_func = _util.display_func(verbose)
    print_func = _util.print_func(verbose)

    a = af.randu(10, 1)
    pos0 = af.randu(10) * 10
    display_func(af.approx1(a, pos0))

    a = af.randu(3, 3)
    pos0 = af.randu(3, 3) * 10
    pos1 = af.randu(3, 3) * 10

    display_func(af.approx2(a, pos0, pos1))

    a = af.randu(8, 1)
    display_func(a)

    display_func(af.fft(a))
    display_func(af.dft(a))
    display_func(af.real(af.ifft(af.fft(a))))
    display_func(af.real(af.idft(af.dft(a))))

    a = af.randu(4, 4)
    display_func(a)

    display_func(af.fft2(a))
    display_func(af.dft(a))
    display_func(af.real(af.ifft2(af.fft2(a))))
    display_func(af.real(af.idft(af.dft(a))))

    a = af.randu(4, 4, 2)
    display_func(a)

    display_func(af.fft3(a))
    display_func(af.dft(a))
    display_func(af.real(af.ifft3(af.fft3(a))))
    display_func(af.real(af.idft(af.dft(a))))

    a = af.randu(10, 1)
    b = af.randu(3, 1)
    display_func(af.convolve1(a, b))
    display_func(af.fft_convolve1(a, b))
    display_func(af.convolve(a, b))
    display_func(af.fft_convolve(a, b))

    a = af.randu(5, 5)
    b = af.randu(3, 3)
    display_func(af.convolve2(a, b))
    display_func(af.fft_convolve2(a, b))
    display_func(af.convolve(a, b))
    display_func(af.fft_convolve(a, b))

    a = af.randu(5, 5, 3)
    b = af.randu(3, 3, 2)
    display_func(af.convolve3(a, b))
    display_func(af.fft_convolve3(a, b))
    display_func(af.convolve(a, b))
    display_func(af.fft_convolve(a, b))

    b = af.randu(3, 1)
    x = af.randu(10, 1)
    a = af.randu(2, 1)
    display_func(af.fir(b, x))
    display_func(af.iir(b, a, x))
af.display(af.dft(a))
af.display(af.real(af.ifft(af.fft(a))))
af.display(af.real(af.idft(af.dft(a))))

a = af.randu(4, 4)
af.display(a)

af.display(af.fft2(a))
af.display(af.dft(a))
af.display(af.real(af.ifft2(af.fft2(a))))
af.display(af.real(af.idft(af.dft(a))))

a = af.randu(4, 4, 2)
af.display(a)

af.display(af.fft3(a))
af.display(af.dft(a))
af.display(af.real(af.ifft3(af.fft3(a))))
af.display(af.real(af.idft(af.dft(a))))

a = af.randu(10, 1)
b = af.randu(3, 1)
af.display(af.convolve1(a, b))
af.display(af.fft_convolve1(a, b))
af.display(af.convolve(a, b))
af.display(af.fft_convolve(a, b))

a = af.randu(5, 5)
b = af.randu(3, 3)
af.display(af.convolve2(a, b))
af.display(af.fft_convolve2(a, b))
示例#8
0
def simple_signal(verbose=False):
    display_func = _util.display_func(verbose)
    print_func   = _util.print_func(verbose)

    signal = af.randu(10)
    x_new  = af.randu(10)
    x_orig = af.randu(10)
    display_func(af.approx1(signal, x_new, xp = x_orig))

    signal = af.randu(3, 3)
    x_new  = af.randu(3, 3)
    x_orig = af.randu(3, 3)
    y_new  = af.randu(3, 3)
    y_orig = af.randu(3, 3)

    display_func(af.approx2(signal, x_new, y_new, xp = x_orig, yp = y_orig))

    a = af.randu(8, 1)
    display_func(a)

    display_func(af.fft(a))
    display_func(af.dft(a))
    display_func(af.real(af.ifft(af.fft(a))))
    display_func(af.real(af.idft(af.dft(a))))

    b = af.fft(a)
    af.ifft_inplace(b)
    display_func(b)
    af.fft_inplace(b)
    display_func(b)

    b = af.fft_r2c(a)
    c = af.fft_c2r(b)
    display_func(b)
    display_func(c)

    a = af.randu(4, 4)
    display_func(a)

    display_func(af.fft2(a))
    display_func(af.dft(a))
    display_func(af.real(af.ifft2(af.fft2(a))))
    display_func(af.real(af.idft(af.dft(a))))

    b = af.fft2(a)
    af.ifft2_inplace(b)
    display_func(b)
    af.fft2_inplace(b)
    display_func(b)

    b = af.fft2_r2c(a)
    c = af.fft2_c2r(b)
    display_func(b)
    display_func(c)

    a = af.randu(4, 4, 2)
    display_func(a)

    display_func(af.fft3(a))
    display_func(af.dft(a))
    display_func(af.real(af.ifft3(af.fft3(a))))
    display_func(af.real(af.idft(af.dft(a))))

    b = af.fft3(a)
    af.ifft3_inplace(b)
    display_func(b)
    af.fft3_inplace(b)
    display_func(b)

    b = af.fft3_r2c(a)
    c = af.fft3_c2r(b)
    display_func(b)
    display_func(c)

    a = af.randu(10, 1)
    b = af.randu(3, 1)
    display_func(af.convolve1(a, b))
    display_func(af.fft_convolve1(a, b))
    display_func(af.convolve(a, b))
    display_func(af.fft_convolve(a, b))

    a = af.randu(5, 5)
    b = af.randu(3, 3)
    display_func(af.convolve2(a, b))
    display_func(af.fft_convolve2(a, b))
    display_func(af.convolve(a, b))
    display_func(af.fft_convolve(a, b))

    a = af.randu(5, 5, 3)
    b = af.randu(3, 3, 2)
    display_func(af.convolve3(a, b))
    display_func(af.fft_convolve3(a, b))
    display_func(af.convolve(a, b))
    display_func(af.fft_convolve(a, b))


    b = af.randu(3, 1)
    x = af.randu(10, 1)
    a = af.randu(2, 1)
    display_func(af.fir(b, x))
    display_func(af.iir(b, a, x))

    display_func(af.medfilt1(a))
    display_func(af.medfilt2(a))
    display_func(af.medfilt(a))
示例#9
0
文件: aflib.py 项目: bfrosik/pycdi
 def fft(arr):
     return af.fft3(arr)
示例#10
0
def simple_signal(verbose=False):
    display_func = _util.display_func(verbose)

    signal = af.randu(10)
    x_new = af.randu(10)
    x_orig = af.randu(10)
    display_func(af.approx1(signal, x_new, xp=x_orig))

    signal = af.randu(3, 3)
    x_new = af.randu(3, 3)
    x_orig = af.randu(3, 3)
    y_new = af.randu(3, 3)
    y_orig = af.randu(3, 3)

    display_func(af.approx2(signal, x_new, y_new, xp=x_orig, yp=y_orig))

    a = af.randu(8, 1)
    display_func(a)

    display_func(af.fft(a))
    display_func(af.dft(a))
    display_func(af.real(af.ifft(af.fft(a))))
    display_func(af.real(af.idft(af.dft(a))))

    b = af.fft(a)
    af.ifft_inplace(b)
    display_func(b)
    af.fft_inplace(b)
    display_func(b)

    b = af.fft_r2c(a)
    c = af.fft_c2r(b)
    display_func(b)
    display_func(c)

    a = af.randu(4, 4)
    display_func(a)

    display_func(af.fft2(a))
    display_func(af.dft(a))
    display_func(af.real(af.ifft2(af.fft2(a))))
    display_func(af.real(af.idft(af.dft(a))))

    b = af.fft2(a)
    af.ifft2_inplace(b)
    display_func(b)
    af.fft2_inplace(b)
    display_func(b)

    b = af.fft2_r2c(a)
    c = af.fft2_c2r(b)
    display_func(b)
    display_func(c)

    a = af.randu(4, 4, 2)
    display_func(a)

    display_func(af.fft3(a))
    display_func(af.dft(a))
    display_func(af.real(af.ifft3(af.fft3(a))))
    display_func(af.real(af.idft(af.dft(a))))

    b = af.fft3(a)
    af.ifft3_inplace(b)
    display_func(b)
    af.fft3_inplace(b)
    display_func(b)

    b = af.fft3_r2c(a)
    c = af.fft3_c2r(b)
    display_func(b)
    display_func(c)

    a = af.randu(10, 1)
    b = af.randu(3, 1)
    display_func(af.convolve1(a, b))
    display_func(af.fft_convolve1(a, b))
    display_func(af.convolve(a, b))
    display_func(af.fft_convolve(a, b))

    a = af.randu(5, 5)
    b = af.randu(3, 3)
    display_func(af.convolve2(a, b))
    display_func(af.fft_convolve2(a, b))
    display_func(af.convolve(a, b))
    display_func(af.fft_convolve(a, b))

    c = af.convolve2NN(a, b)
    display_func(c)
    in_dims = c.dims()
    incoming_grad = af.constant(1, in_dims[0], in_dims[1])
    g = af.convolve2GradientNN(incoming_grad, a, b, c)
    display_func(g)

    a = af.randu(5, 5, 3)
    b = af.randu(3, 3, 2)
    display_func(af.convolve3(a, b))
    display_func(af.fft_convolve3(a, b))
    display_func(af.convolve(a, b))
    display_func(af.fft_convolve(a, b))

    b = af.randu(3, 1)
    x = af.randu(10, 1)
    a = af.randu(2, 1)
    display_func(af.fir(b, x))
    display_func(af.iir(b, a, x))

    display_func(af.medfilt1(a))
    display_func(af.medfilt2(a))
    display_func(af.medfilt(a))