diff --git a/verify_fft_equivalence.py b/verify_fft_equivalence.py new file mode 100644 index 0000000..a5a17ab --- /dev/null +++ b/verify_fft_equivalence.py @@ -0,0 +1,102 @@ +""" +Verify that the OpenCL FFT convolution in solve_opencl / solve_opencl_batch2 +produces the same result as scipy.signal.fftconvolve(f, G, mode='same'). + +This script simulates the OpenCL logic with numpy FFTs so we can compare +directly without needing a GPU. +""" +import numpy as np +from scipy.signal import fftconvolve + +def make_green_2d(H, W, dx=1.0): + """Reproduce PoissonSolver._create_radialGrid for 2D.""" + rdims = (2*H - 1, 2*W - 1) + R = np.indices(rdims, dtype=np.float64) + R[0] += 1 - H + R[1] += 1 - W + R[0] *= dx + R[1] *= dx + np.hypot(R[1], R[0], R[0]) + R = R[0] + zval = 0.5 * (dx + dx) / 2.0 + R[H-1, W-1] = zval + # Apply 2D Green's function: log(r)/(2*pi) + G = np.log(R) * (0.5 / np.pi) + return G + +def test_equivalence(H, W, dx=1.0, seed=42): + np.random.seed(seed) + f = np.random.randn(H, W).astype(np.float64) + G = make_green_2d(H, W, dx) + gdx_prod = dx ** 2 + + # --- CPU reference: scipy fftconvolve --- + p_cpu = fftconvolve(f, G, mode='same') * gdx_prod + + # --- Simulate what OpenCL solve_opencl does --- + f32 = f.astype(np.float32) + G32 = G.astype(np.float32) + full_shape = (2*H - 1, 2*W - 1) + + # G already has shape full_shape, no padding needed + # f is zero-padded to full_shape at [0:H, 0:W] + f_pad = np.zeros(full_shape, dtype=np.float32) + f_pad[:H, :W] = f32 + + # FFT at full_shape, multiply, iFFT + F_f = np.fft.rfftn(f_pad) + F_g = np.fft.rfftn(G32) + conv_circ = np.fft.irfftn(F_f * F_g, s=full_shape) + + # OLD OpenCL extraction (BUGGY): same_starts = [(full_shape[i] - fshape[i]) // 2] + same_starts_old = [(full_shape[0] - H) // 2, (full_shape[1] - W) // 2] + p_old = conv_circ[ + same_starts_old[0]:same_starts_old[0]+H, + same_starts_old[1]:same_starts_old[1]+W + ].astype(np.float32) * np.float32(gdx_prod) + + # FIXED OpenCL extraction: same_starts = [s - 1 for s in fshape] + same_starts_fixed = [H - 1, W - 1] + p_fixed = conv_circ[ + same_starts_fixed[0]:same_starts_fixed[0]+H, + same_starts_fixed[1]:same_starts_fixed[1]+W + ].astype(np.float32) * np.float32(gdx_prod) + + # --- Also test with correct FFT size (3H-2) --- + correct_fft_shape = (H + (2*H-1) - 1, W + (2*W-1) - 1) # (3H-2, 3W-2) + F_f_full = np.fft.rfftn(f32, s=correct_fft_shape) + F_g_full = np.fft.rfftn(G32, s=correct_fft_shape) + conv_full = np.fft.irfftn(F_f_full * F_g_full, s=correct_fft_shape) + # scipy-style 'same' extraction from full linear conv + start_full = [(correct_fft_shape[0] - H) // 2, (correct_fft_shape[1] - W) // 2] + p_full_linear = conv_full[ + start_full[0]:start_full[0]+H, + start_full[1]:start_full[1]+W + ].astype(np.float32) * np.float32(gdx_prod) + + # --- Compare --- + p_cpu32 = p_cpu.astype(np.float32) + scale = max(np.abs(p_cpu32).max(), 1e-10) + + diff_old = np.abs(p_cpu32 - p_old) + diff_fixed = np.abs(p_cpu32 - p_fixed) + diff_full = np.abs(p_cpu32 - p_full_linear) + + print(f"=== Test: f({H},{W}), G({2*H-1},{2*W-1}), full_shape={full_shape} ===") + print(f" scipy fftconvolve full linear conv shape: {correct_fft_shape}") + print(f" p_cpu range: [{p_cpu32.min():.4f}, {p_cpu32.max():.4f}]") + print() + print(f" [OLD BUG] same_starts={same_starts_old}:") + print(f" max_abs_diff = {diff_old.max():.6e} rel_max = {diff_old.max()/scale:.6e}") + print(f" [FIXED] same_starts={same_starts_fixed}:") + print(f" max_abs_diff = {diff_fixed.max():.6e} rel_max = {diff_fixed.max()/scale:.6e}") + print(f" [REF] Full linear conv (FFT size {correct_fft_shape}):") + print(f" max_abs_diff = {diff_full.max():.6e} rel_max = {diff_full.max()/scale:.6e}") + print() + + return diff_old.max(), diff_fixed.max(), diff_full.max() + +if __name__ == '__main__': + for size in [16, 32, 64, 128, 256, 384]: + d_ocl, d_cor, d_full = test_equivalence(size, size) + print()