def for_back_np_1d(axis=0):
"""Compute forward 1D FFT and then backward 1D FFT along the x-direction.
Compute this for all possible grid sizes with lengths powers of 2.
"""
length = [2**s for s in range(7, 11)]
all_shapes = list(itertools.product(length, length))
dtype = np.complex128
eps = np.finfo(dtype).eps
eps *= 10
func = np.ones
delta_r = (1, 1)
for shape in all_shapes:
grid = [func(shape, dtype=dtype)] * 2
grid_k = ttools.fft_1d(grid, delta_r, axis=axis)
grid_r = ttools.ifft_1d(grid_k, delta_r, axis=axis)
max_diff0 = np.max(abs(grid_r[0] - grid[0]))
max_diff1 = np.max(abs(grid_r[1] - grid[1]))
atoms = [ttools.calc_atoms(grid), ttools.calc_atoms(grid_r)]
assert (max_diff0 <= eps and max_diff1 <= eps), f"Max errors: \
{max_diff0}, {max_diff1}"
assert len(grid_r) == 2
assert abs(atoms[0] - atoms[1]) <= eps * math.prod(shape), \
f"\nAtom num. before/after: {atoms[0]}, {atoms[1]}; \
\nDifference: {abs(atoms[0] - atoms[1])}; \
\n2*eps*N: {eps * math.prod(shape)}."
print(f"Test `for_back_np_1d` for axis={axis} passed.")
def compare4_np_1d_2d():
"""Compute a single inverse 2D FFT and two inverse 1D FFTs.
Compute this for all possible grid sizes with lengths powers of 2.
"""
length = [2**s for s in range(7, 11)]
all_shapes = list(itertools.product(length, length))
dtype = np.complex128
eps = np.finfo(dtype).eps
eps *= 10
func = np.ones
delta_r = (1, 1)
for shape in all_shapes:
grid = [func(shape, dtype=dtype)] * 2
grid = ttools.fft_2d(grid, delta_r)
grid_r_half = ttools.ifft_1d(grid, delta_r, axis=0)
grid_r = ttools.ifft_1d(grid_r_half, delta_r, axis=1)
grid = ttools.ifft_2d(grid, delta_r)
max_diff0 = np.max(abs(grid_r[0] - grid[0]))
max_diff1 = np.max(abs(grid_r[1] - grid[1]))
atoms = [ttools.calc_atoms(grid), ttools.calc_atoms(grid_r)]
assert (max_diff0 <= eps and max_diff1 <= eps), f"\
Max errors: {max_diff0}, {max_diff1}."
assert len(grid_r) == 2
assert abs(atoms[0] - atoms[1]) <= eps * math.prod(shape), \
f"\nAtom num. before/after: {atoms[0]}, {atoms[1]}; \
\nDifference: {abs(atoms[0] - atoms[1])}; \
\n2*eps*N: {eps * math.prod(shape)}."
print("Test `compare4_np_1d_2d` passed.")
def compare_norm_torch():
"""Compare the total norm of a grid and its FFT.
Compute this for all possible grid sizes with lengths powers of 2.
"""
length = [2**s for s in range(7, 11)]
all_shapes = list(itertools.product(length, length))
dtype = torch.complex128
eps = torch.finfo(dtype).eps
eps *= 10
func = torch.ones
delta_r = (1, 1)
for shape in all_shapes:
grid = [func(shape, dtype=dtype)] * 2
grid_k = ttools.fft_2d(grid, delta_r)
vol_elem = 4 * np.pi**2 / math.prod(shape)
atoms = [ttools.calc_atoms(grid), ttools.calc_atoms(grid_k, vol_elem)]
assert abs(atoms[0] - atoms[1]) <= eps * math.prod(shape), \
f"\nAtom num. before/after: {atoms[0]}, {atoms[1]}; \
\nDifference: {abs(atoms[0] - atoms[1])}; \
\n2*eps*N: {eps * math.prod(shape)}."
print("Test `compare_norm_torch` passed.")
def for_back_torch_2d():
"""Compute forward 2D FFT and then backward 2D FFT. Assert same as initial.
Compute this for all possible grid sizes with lengths powers of 2.
"""
length = [2**s for s in range(7, 11)]
all_shapes = list(itertools.product(length, length))
dtype = torch.complex128
eps = torch.finfo(dtype).eps
eps *= 10
func = torch.ones
delta_r = (1, 1)
for shape in all_shapes:
grid = [func(shape, dtype=dtype, device='cuda')] * 2
grid_k = ttools.fft_2d(grid, delta_r)
grid_r = ttools.ifft_2d(grid_k, delta_r)
max_diff0 = torch.max(abs(grid_r[0] - grid[0]))
max_diff1 = torch.max(abs(grid_r[1] - grid[1]))
atoms = [ttools.calc_atoms(grid), ttools.calc_atoms(grid_r)]
assert (max_diff0 <= eps and max_diff1 <= eps), f"Max errors: \
{max_diff0}, {max_diff1}"
assert len(grid_r) == 2
assert abs(atoms[0] - atoms[1]) <= eps * math.prod(shape), \
f"\nAtom num. before/after: {atoms[0]}, {atoms[1]}; \
\nDifference: {abs(atoms[0] - atoms[1])}; \
\n2*eps*N: {eps * math.prod(shape)}."
print("Test `for_back_torch_2d` passed.")
