#!/usr/bin/env python3 """ D2Q9 Regression Test — Poiseuille Channel + Cylinder Flow ========================================================== Uses kernel_v2.cu by default (legacy kernel.cu remains optional fallback). Produces matplotlib figures for visual validation. Usage: python tests/test_d2q9_visual.py --device 2 python tests/test_d2q9_visual.py --device 2 --cylinder python tests/test_d2q9_visual.py --device 2 --legacy """ import sys, os, argparse, time # Ensure CelerisLab package is importable sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'src')) import numpy as np import pycuda.driver as cuda import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt from matplotlib.colors import Normalize from CelerisLab.cuda import compiler from CelerisLab.common import preprocess as preproc # ━━━━━━━━━━━━━━━━━━━━━━━ Configuration ━━━━━━━━━━━━━━━━━━━━━━━ NX, NY = 1280, 512 NQ = 9 NT = 128 DIM = 2 VIS = 0.002 U0 = 0.01 RHO = 1.0 TOTAL = NX * NY # Original direction vectors (const.h ordering) E = np.array([[0,0],[1,0],[0,1],[-1,0],[0,-1],[1,1],[-1,1],[-1,-1],[1,-1]], dtype=np.int32) OPP = np.array([0, 3, 4, 1, 2, 7, 8, 5, 6], dtype=np.int32) FLUID_FLAG = 0b00000001 SOLID_FLAG = 0b00000010 INTERFACE_FLAG = 0b00001000 # ━━━━━━━━━━━━━━━━━━━━━━━ Helpers ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ def configure_macros(n_objs=0): """Write macros.h for D2Q9 regression (v2 default, legacy optional).""" lines = compiler.read_lines(compiler.kernel_path("macros.h")) defs = { 'MULT_GPU': 'False', 'NT': NT, 'X_1U': 128, 'Y_1U': 32, 'Z_1U': 1, 'LBtype': 'float', 'UX': 10, 'UY': 16, 'UZ': 1, 'NX': NX, 'NY': NY, 'NZ': 1, 'DIM': DIM, 'NQ': NQ, 'VIS': VIS, 'RHO': f'{RHO}', 'U0': U0, 'N_OBJS': n_objs, 'COLLISION_MODEL': 0, # SRT 'STREAMING_MODEL': 0, # double-buffer 'STORE_PRECISION': 0, # FP32 'USE_DDF_SHIFTING': 0, # keep unshifted for v2 defaults } for name, val in defs.items(): lines = compiler.modify_macro(lines, name, val) compiler.write_lines(compiler.kernel_path("macros.h"), lines) def extract_fields(ddf_host, use_ddf_shifting=False): """Compute rho, u, v from host DDF in original D2Q9 direction ordering.""" f = ddf_host.reshape(NQ, NY, NX) if use_ddf_shifting: rho = np.sum(f, axis=0) + RHO denom = np.full_like(rho, RHO) else: rho = np.sum(f, axis=0) denom = rho denom_safe = np.where(np.abs(denom) > 1e-12, denom, 1.0) u = (f[1] + f[5] + f[8] - f[3] - f[6] - f[7]) / denom_safe v = (f[2] + f[5] + f[6] - f[4] - f[7] - f[8]) / denom_safe u = np.where(np.abs(denom) > 1e-12, u, 0.0) v = np.where(np.abs(denom) > 1e-12, v, 0.0) return rho, u, v def analytical_poiseuille(y_arr): """Analytical parabolic profile matching InitTubeFlow.""" yy = (y_arr - 0.5 * (NY - 1)) / (NY - 2.0) return U0 * 1.5 * (1 - 4 * yy**2) def build_cylinder_data(cx, cy, radius): """Replicate driver.py add_cylinder logic for flag / delta / indx.""" flag = np.ones(TOTAL, dtype=np.uint8) # init all FLUID indx = np.zeros(TOTAL, dtype=np.int32) delta_list = [] index_offset = 0 # Build Poiseuille flag first (walls + solid borders) for y in range(NY): for x in range(NX): k = x + y * NX if y == 0 or y == NY - 1 or x == 0 or x == NX - 1: flag[k] = SOLID_FLAG # Add cylinder for x in range(int(cx - radius) - 1, int(cx + radius) + 1): for y in range(int(cy - radius) - 1, int(cy + radius) + 1): if (x - cx)**2 + (y - cy)**2 < radius**2: k = x + y * NX flag[k] = SOLID_FLAG dt = np.zeros(11, dtype=np.float32) dt[0] = np.int32(0).view(np.float32) # id_object = 0 has_interface = False for i in range(NQ): xn = x + E[i][0] yn = y + E[i][1] if (xn - cx)**2 + (yn - cy)**2 >= radius**2: has_interface = True xi, yi = preproc.find_circle_intersection( x, y, xn, yn, cx, cy, radius) d_neb = np.sqrt((xi - xn)**2 + (yi - yn)**2) e_len = np.sqrt(E[i][0]**2 + E[i][1]**2) if e_len > 0: dt[i] = d_neb / e_len if has_interface: flag[k] |= INTERFACE_FLAG dt[9] = (cy - y) / radius dt[10] = (x - cx) / radius indx[k] = index_offset delta_list.append(dt) index_offset += 11 delta = np.concatenate(delta_list) if delta_list else np.zeros(1, dtype=np.float32) return flag, indx, delta # ━━━━━━━━━━━━━━━━━━━━━━━ Simulation ━━━━━━━━━━━━━━━━━━━━━━━━━━ def run_simulation(device_id, n_steps, n_objs, flag_host, indx_host, delta_host, use_legacy=False): """Compile kernel, run LBM, return DDF on host.""" cuda.init() dev = cuda.Device(device_id) ctx = dev.make_context() print(f"[GPU {device_id}] {dev.name()}") try: configure_macros(n_objs) if use_legacy: compiler.compile_kernel() ptx_path = compiler.kernel_path("kernel.ptx") init_name = "InitTubeFlow" else: compiler.compile_kernel_v2() ptx_path = compiler.kernel_path("kernel_v2.ptx") init_name = "InitTubeFlow_v2" mod = cuda.module_from_file(ptx_path) step_fn = mod.get_function("OneStep") init_fn = mod.get_function(init_name) nbytes_ddf = TOTAL * NQ * 4 ddf_gpu = cuda.mem_alloc(nbytes_ddf) temp_gpu = cuda.mem_alloc(nbytes_ddf) flag_gpu = cuda.mem_alloc(flag_host.nbytes) indx_gpu = cuda.mem_alloc(indx_host.nbytes) delta_gpu = cuda.mem_alloc(max(delta_host.nbytes, 4)) action_host = np.zeros(max(n_objs, 1), dtype=np.float32) obs_host = np.zeros(max(n_objs * DIM, 1), dtype=np.float32) action_gpu = cuda.mem_alloc(action_host.nbytes) obs_gpu = cuda.mem_alloc(obs_host.nbytes) cuda.memcpy_htod(action_gpu, action_host) cuda.memcpy_htod(obs_gpu, obs_host) block = (NT, 1, 1) grid = (NX // NT, NY, 1) # Init Poiseuille init_fn(flag_gpu, ddf_gpu, block=block, grid=grid) ctx.synchronize() # Overwrite flag / indx / delta for cylinder case cuda.memcpy_htod(flag_gpu, flag_host) cuda.memcpy_htod(indx_gpu, indx_host) cuda.memcpy_htod(delta_gpu, delta_host) # Step loop t0 = time.time() for i in range(n_steps): step_fn(flag_gpu, ddf_gpu, temp_gpu, indx_gpu, delta_gpu, action_gpu, obs_gpu, block=block, grid=grid) ddf_gpu, temp_gpu = temp_gpu, ddf_gpu ctx.synchronize() dt = time.time() - t0 mlups = TOTAL * n_steps / dt / 1e6 print(f" {n_steps} steps in {dt:.2f}s ({mlups:.1f} MLUPS)") # Copy back ddf = np.zeros(TOTAL * NQ, dtype=np.float32) cuda.memcpy_dtoh(ddf, ddf_gpu) flag_out = np.zeros(TOTAL, dtype=np.uint8) cuda.memcpy_dtoh(flag_out, flag_gpu) return ddf, flag_out finally: ctx.pop() # ━━━━━━━━━━━━━━━━━━━━━━━ Visualization ━━━━━━━━━━━━━━━━━━━━━━━ def plot_poiseuille(ddf, flag, out_path, use_ddf_shifting=False): """3-panel figure: velocity mag, u(y) profile, pressure along centerline.""" rho, u, v = extract_fields(ddf, use_ddf_shifting=use_ddf_shifting) vel_mag = np.sqrt(u**2 + v**2) # Mask solid cells for display mask = (flag.reshape(NY, NX) & SOLID_FLAG).astype(bool) vel_mag_masked = np.ma.array(vel_mag, mask=mask) fig, axes = plt.subplots(1, 3, figsize=(18, 5)) # (a) Velocity magnitude heatmap ax = axes[0] im = ax.imshow(vel_mag_masked, origin='lower', aspect='auto', cmap='jet', extent=[0, NX, 0, NY]) plt.colorbar(im, ax=ax, label='|u|') ax.set_title('Velocity magnitude') ax.set_xlabel('x'); ax.set_ylabel('y') # (b) u(y) profile at x = NX/2 vs. analytical ax = axes[1] x_mid = NX // 2 y_arr = np.arange(NY, dtype=float) ax.plot(u[:, x_mid], y_arr, 'b-', lw=2, label='LBM') ax.plot(analytical_poiseuille(y_arr), y_arr, 'r--', lw=1.5, label='Analytical') ax.set_xlabel('u_x'); ax.set_ylabel('y') ax.set_title(f'u(y) at x={x_mid}') ax.legend() ax.grid(True, alpha=0.3) # (c) Pressure along centerline y = NY/2 ax = axes[2] y_mid = NY // 2 p = rho / 3.0 ax.plot(np.arange(NX), p[y_mid, :], 'g-', lw=1.5) ax.set_xlabel('x'); ax.set_ylabel('p = ρ/3') ax.set_title(f'Pressure along centerline (y={y_mid})') ax.grid(True, alpha=0.3) fig.suptitle(f'D2Q9 Poiseuille – NX={NX}, NY={NY}, VIS={VIS}, U0={U0}', fontsize=13) fig.tight_layout() fig.savefig(out_path, dpi=150) print(f" Saved: {out_path}") plt.close(fig) def plot_cylinder(ddf, flag, cx, cy, radius, out_path, use_ddf_shifting=False): """3-panel figure: velocity magnitude (zoom), vorticity, streamlines.""" rho, u, v = extract_fields(ddf, use_ddf_shifting=use_ddf_shifting) vel_mag = np.sqrt(u**2 + v**2) mask = (flag.reshape(NY, NX) & SOLID_FLAG).astype(bool) # Zoom window around cylinder pad = int(radius * 8) x0 = max(int(cx - pad), 0) x1 = min(int(cx + pad * 2), NX) y0 = max(int(cy - pad), 0) y1 = min(int(cy + pad), NY) fig, axes = plt.subplots(1, 3, figsize=(20, 6)) # (a) Velocity magnitude (zoomed) ax = axes[0] vm_z = np.ma.array(vel_mag[y0:y1, x0:x1], mask=mask[y0:y1, x0:x1]) im = ax.imshow(vm_z, origin='lower', aspect='equal', cmap='jet', extent=[x0, x1, y0, y1]) circ = plt.Circle((cx, cy), radius, fill=True, color='gray', alpha=0.7) ax.add_patch(circ) plt.colorbar(im, ax=ax, label='|u|') ax.set_title('Velocity magnitude') # (b) Vorticity ax = axes[1] dvdx = np.gradient(v, axis=1) dudy = np.gradient(u, axis=0) omega = dvdx - dudy om_z = np.ma.array(omega[y0:y1, x0:x1], mask=mask[y0:y1, x0:x1]) vmax = np.percentile(np.abs(omega[~mask]), 99) im = ax.imshow(om_z, origin='lower', aspect='equal', cmap='RdBu_r', extent=[x0, x1, y0, y1], vmin=-vmax, vmax=vmax) circ2 = plt.Circle((cx, cy), radius, fill=True, color='gray', alpha=0.7) ax.add_patch(circ2) plt.colorbar(im, ax=ax, label='ω') ax.set_title('Vorticity') # (c) Streamlines ax = axes[2] X, Y = np.meshgrid(np.arange(x0, x1), np.arange(y0, y1)) u_z = u[y0:y1, x0:x1].copy() v_z = v[y0:y1, x0:x1].copy() u_z[mask[y0:y1, x0:x1]] = 0 v_z[mask[y0:y1, x0:x1]] = 0 speed = np.sqrt(u_z**2 + v_z**2) ax.streamplot(X, Y, u_z, v_z, color=speed, cmap='jet', density=2.0, linewidth=0.8) circ3 = plt.Circle((cx, cy), radius, fill=True, color='gray', alpha=0.7) ax.add_patch(circ3) ax.set_xlim(x0, x1); ax.set_ylim(y0, y1) ax.set_aspect('equal') ax.set_title('Streamlines') fig.suptitle(f'D2Q9 Cylinder Flow – Re_D={U0*1.5*2*radius/VIS:.0f}, D={2*radius}', fontsize=13) fig.tight_layout() fig.savefig(out_path, dpi=150) print(f" Saved: {out_path}") plt.close(fig) # ━━━━━━━━━━━━━━━━━━━━━━━ Main ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ def main(): parser = argparse.ArgumentParser(description='D2Q9 Regression Test') parser.add_argument('--device', type=int, default=2, help='CUDA device ID (default: 2)') parser.add_argument('--legacy', action='store_true', help='Use legacy kernel.cu path (default uses kernel_v2.cu)') parser.add_argument('--cylinder', action='store_true', help='Also run cylinder flow test') parser.add_argument('--steps-pois', type=int, default=5000, help='Steps for Poiseuille (default: 5000)') parser.add_argument('--steps-cyl', type=int, default=30000, help='Steps for cylinder (default: 30000)') args = parser.parse_args() out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'output') os.makedirs(out_dir, exist_ok=True) use_ddf_shifting = bool(args.legacy) mode = 'legacy kernel.cu' if args.legacy else 'kernel_v2.cu' print(f"\n[Mode] {mode}") # ---- Test 1: Poiseuille ---- print("\n===== Test 1: Poiseuille Channel Flow =====") flag_pois = np.ones(TOTAL, dtype=np.uint8) indx_pois = np.zeros(TOTAL, dtype=np.int32) delta_pois = np.zeros(1, dtype=np.float32) ddf, flag = run_simulation(args.device, args.steps_pois, 0, flag_pois, indx_pois, delta_pois, use_legacy=args.legacy) plot_poiseuille(ddf, flag, os.path.join(out_dir, 'poiseuille_d2q9.png'), use_ddf_shifting=use_ddf_shifting) # Error metric rho, u, v = extract_fields(ddf, use_ddf_shifting=use_ddf_shifting) y_arr = np.arange(NY, dtype=float) u_ana = analytical_poiseuille(y_arr) x_mid = NX // 2 u_num = u[:, x_mid] # Interior cells only (skip walls) err = np.max(np.abs(u_num[2:-2] - u_ana[2:-2])) / np.max(np.abs(u_ana[2:-2])) print(f" L∞ relative error at x={x_mid}: {err:.2e}") # ---- Test 2: Cylinder ---- if args.cylinder: print("\n===== Test 2: Flow Around Cylinder =====") cyl_cx, cyl_cy, cyl_r = 256.0, 256.0, 32.0 flag_cyl, indx_cyl, delta_cyl = build_cylinder_data(cyl_cx, cyl_cy, cyl_r) ddf2, flag2 = run_simulation(args.device, args.steps_cyl, 1, flag_cyl, indx_cyl, delta_cyl, use_legacy=args.legacy) plot_cylinder(ddf2, flag2, cyl_cx, cyl_cy, cyl_r, os.path.join(out_dir, 'cylinder_d2q9.png'), use_ddf_shifting=use_ddf_shifting) print("\nDone.") if __name__ == '__main__': main()