最近在学taichi，想用taichi 模拟相场模型(PFM)，但PFM模型中普遍使用半隐式傅里叶谱方法（ semi-implicit Fourier spectral algorithm ）求解AC和CH方程， 以提高模拟效率和准确性。目前貌似 taichi库中没有 fft和 ifft， 因此，写了一个简单的文件，用以实现fft和ifft，具体程序如下。程序中， (i)fft_ti和fft_np的结果做了一点对比，不是特别的准确，如果有错误和优化，欢迎提出建议。 仅供参考。
程序如下：
import time
import math
import numpy as np
import taichi as ti
import taichi.math as tm
ti.init(arch=ti.gpu, device_memory_fraction=0.95)
fft -numpy
np.random.seed = int(time.time())
nx, ny, nz = 64, 64, 64
a0 = np.random.randint(1000, size=(nx, ny, nz))
ak_np0 = np.fft.fftn(a0)
inver fft -numpy
a0_complex = np.random.randint(1000, size=(nx, ny, nz)) + 1j * np.random.randint(1000, size=(nx, ny, nz))
real0, imag0 = a0_complex.real, a0_complex.imag
a_np = np.fft.ifftn(a0_complex).real
fft -taichi
a = ti.field(dtype=ti.int32, shape=(nx, ny, nz))
a.from_numpy(a0)
ifft -taichi
real_ti = ti.field(dtype=ti.int32, shape=(nx, ny, nz))
imag_ti = ti.field(dtype=ti.int32, shape=(nx, ny, nz))
real_ti.from_numpy(real0), imag_ti.from_numpy(imag0)
fft -taichi
ak_fft = ti.Struct.field({
‘a0’: tm.vec2,
‘fftx’: tm.vec2,
‘ffty’: tm.vec2,
‘fftz’: tm.vec2}, shape=(nx, ny, nz))
ifft -taichi
ak_ifft = ti.Struct.field({
‘a0’: tm.vec2,
‘ifftx’: tm.vec2,
‘iffty’: tm.vec2,
‘ifftz’: tm.vec2}, shape=(nx, ny, nz))
@ti.func
def fft(var, var_fft, dim):
if dim == 0:
for ix, iy, iz, iix in ti.ndrange(nx, ny, nz, nx):
var_fft[ix, iy, iz][0] += (var[iix, iy, iz][0] * tm.cos(2 * math.pi / nx * ix * iix) +
var[iix, iy, iz][1] * tm.sin(2 * math.pi / nx * ix * iix))
var_fft[ix, iy, iz][1] += (var[iix, iy, iz][1] * tm.cos(2 * math.pi / nx * ix * iix) -
var[iix, iy, iz][0] * tm.sin(2 * math.pi / nx * ix * iix))
elif dim == 1:
for ix, iy, iz, iiy in ti.ndrange(nx, ny, nz, ny):
var_fft[ix, iy, iz][0] += (var[ix, iiy, iz][0] * tm.cos(2 * math.pi / ny * iy * iiy) +
var[ix, iiy, iz][1] * tm.sin(2 * math.pi / ny * iy * iiy))
var_fft[ix, iy, iz][1] += (var[ix, iiy, iz][1] * tm.cos(2 * math.pi / ny * iy * iiy) -
var[ix, iiy, iz][0] * tm.sin(2 * math.pi / ny * iy * iiy))
else:
for ix, iy, iz, iiz in ti.ndrange(nx, ny, nz, nz):
var_fft[ix, iy, iz][0] += (var[ix, iy, iiz][0] * tm.cos(2 * math.pi / nz * iz * iiz) +
var[ix, iy, iiz][1] * tm.sin(2 * math.pi / nz * iz * iiz))
var_fft[ix, iy, iz][1] += (var[ix, iy, iiz][1] * tm.cos(2 * math.pi / nz * iz * iiz) -
var[ix, iy, iiz][0] * tm.sin(2 * math.pi / nz * iz * iiz))
@ti.func
def ifft(var, var_ifft, dim):
if dim == 0:
for ix, iy, iz, iix in ti.ndrange(nx, ny, nz, nx):
var_ifft[ix, iy, iz][0] += (var[iix, iy, iz][0] * tm.cos(2 * math.pi / nx * ix * iix) -
var[iix, iy, iz][1] * tm.sin(2 * math.pi / nx * ix * iix)) / nx
var_ifft[ix, iy, iz][1] += (var[iix, iy, iz][1] * tm.cos(2 * math.pi / nx * ix * iix) +
var[iix, iy, iz][0] * tm.sin(2 * math.pi / nx * ix * iix)) / nx
elif dim == 1:
for ix, iy, iz, iiy in ti.ndrange(nx, ny, nz, ny):
var_ifft[ix, iy, iz][0] += (var[ix, iiy, iz][0] * tm.cos(2 * math.pi / ny * iy * iiy) -
var[ix, iiy, iz][1] * tm.sin(2 * math.pi / ny * iy * iiy)) / ny
var_ifft[ix, iy, iz][1] += (var[ix, iiy, iz][1] * tm.cos(2 * math.pi / ny * iy * iiy) +
var[ix, iiy, iz][0] * tm.sin(2 * math.pi / ny * iy * iiy)) / ny
else:
for ix, iy, iz, iiz in ti.ndrange(nx, ny, nz, nz):
var_ifft[ix, iy, iz][0] += (var[ix, iy, iiz][0] * tm.cos(2 * math.pi / nz * iz * iiz) -
var[ix, iy, iiz][1] * tm.sin(2 * math.pi / nz * iz * iiz)) / nz
var_ifft[ix, iy, iz][1] += (var[ix, iy, iiz][1] * tm.cos(2 * math.pi / nz * iz * iiz) +
var[ix, iy, iiz][0] * tm.sin(2 * math.pi / nz * iz * iiz)) / nz
@ti.kernel
def test():
for i in ti.grouped(ak_ifft):
# ak_fft[i].a0[0] = a[i]
ak_ifft[i].a0[0] = real_ti[i]
ak_ifft[i].a0[1] = imag_ti[i]
# fft(ak_fft.a0, ak_fft.fftx, 0)
# fft(ak_fft.fftx, ak_fft.ffty, 1)
# fft(ak_fft.ffty, ak_fft.fftz, 2)
ifft(ak_ifft.a0, ak_ifft.ifftx, 0)
ifft(ak_ifft.ifftx, ak_ifft.iffty, 1)
ifft(ak_ifft.iffty, ak_ifft.ifftz, 2)
test()
#--------------- for check----------------------------
fft
ak_ti = ak_fft.fftz.to_numpy()
ak_np = np.zeros([nx, ny, nz, 2])
ak_np[:, :, :, 0] = ak_np0.real
ak_np[:, :, :, 1] = ak_np0.imag
ak_ti0, ak_ti1 = ak_ti[:, :, :, 0], ak_ti[:, :, :, 1]
ak_np0, ak_np1 = ak_np[:, :, :, 0], ak_np[:, :, :, 1]
ifft
a0_ti = ak_ifft.ifftz.to_numpy()
a_ti = a0_ti[:, :, :, 0]
diff_a = a_np - a_ti
ratio = a_np / a_ti
print(np.max(diff_a), np.min(diff_a))
print(np.max(ratio), np.min(ratio))