#!/usr/bin/env python3 """ D3Q19 SRT — Cylinder Wake Flow (Periodic Z) ============================================= Tests the v2 modular kernel (kernel_v2.cu) in 3D. Cylinder axis along z, parabolic inlet, pressure outlet, no-slip y-walls. Produces cross-section visualizations at z=NZ/2. Usage: python tests/test_d3q19_cylinder.py --device 3 python tests/test_d3q19_cylinder.py --device 3 --re 200 --steps 50000 """ import sys, os, argparse, time, struct 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 CelerisLab.cuda import compiler # ━━━━━━━━━━━━━━━━━━━━━━━ Configuration ━━━━━━━━━━━━━━━━━━━━━━━ # Reasonable 3D grid — fits in < 500 MB GPU memory NX, NY, NZ = 256, 128, 32 NQ = 19 NT = 128 DIM = 3 RHO = 1.0 CYL_CX, CYL_CY = 64.0, 64.0 # Cylinder center (x,y) CYL_R = 12.0 # Cylinder radius TOTAL = NX * NY * NZ # D3Q19 paired direction ordering (from descriptors.cuh) # 0:rest (1,2)±x (3,4)±y (5,6)±z # (7,8)±(x+y) (9,10)±(x+z) (11,12)±(y+z) # (13,14)±(x-y) (15,16)±(x-z) (17,18)±(y-z) CX = np.array([0, 1,-1, 0, 0, 0, 0, 1,-1, 1,-1, 0, 0, 1,-1, 1,-1, 0, 0], dtype=np.int32) CY = np.array([0, 0, 0, 1,-1, 0, 0, 1,-1, 0, 0, 1,-1,-1, 1, 0, 0, 1,-1], dtype=np.int32) CZ = np.array([0, 0, 0, 0, 0, 1,-1, 0, 0, 1,-1, 1,-1, 0, 0,-1, 1,-1, 1], dtype=np.int32) W = np.array([1/3] + [1/18]*6 + [1/36]*12, dtype=np.float32) FLUID_FLAG = 0x01 SOLID_FLAG = 0x02 OBSTACLE_FLAG = 0x04 # triggers half-way BB at adjacent fluid nodes # ━━━━━━━━━━━━━━━━━━━━━━━ Helpers ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ def compute_vis_omega(Re, D, U0): """Compute viscosity and omega from target Reynolds number.""" vis = U0 * D / Re omega = 1.0 / (3.0 * vis + 0.5) return vis, omega def configure_macros_3d(vis, u0, n_objs=0): """Write macros.h for D3Q19 SRT.""" lines = compiler.read_lines(compiler.kernel_path("macros.h")) defs = { 'MULT_GPU': 'False', 'NT': NT, 'X_1U': NX, 'Y_1U': NY, 'Z_1U': NZ, 'LBtype': 'float', 'UX': 1, 'UY': 1, 'UZ': 1, 'NX': NX, 'NY': NY, 'NZ': NZ, 'DIM': DIM, 'NQ': NQ, 'VIS': f'{vis:.10f}', '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, } for name, val in defs.items(): lines = compiler.modify_macro(lines, name, val) compiler.write_lines(compiler.kernel_path("macros.h"), lines) def build_cylinder_3d(cx, cy, radius): """Build flag array for 3D cylinder (axis along z, periodic z).""" flag = np.ones(TOTAL, dtype=np.uint8) * FLUID_FLAG for z in range(NZ): for y in range(NY): for x in range(NX): k = z * NY * NX + y * NX + x # Channel walls if y == 0 or y == NY - 1 or x == 0 or x == NX - 1: flag[k] = SOLID_FLAG # Cylinder body (obstacle — triggers BB at fluid neighbors) elif (x - cx)**2 + (y - cy)**2 < radius**2: flag[k] = OBSTACLE_FLAG return flag def extract_fields_3d(ddf_host, z_slice): """Extract rho, u, v, w at a given z-slice from D3Q19 DDF (v2 ordering).""" # DDF layout: f[i * TOTAL + k] where k = z*NY*NX + y*NX + x f = ddf_host.reshape(NQ, NZ, NY, NX) fz = f[:, z_slice, :, :] # shape (NQ, NY, NX) rho = np.sum(fz, axis=0) ux = np.zeros_like(rho) uy = np.zeros_like(rho) uz_field = np.zeros_like(rho) for i in range(NQ): ux += CX[i] * fz[i] uy += CY[i] * fz[i] uz_field += CZ[i] * fz[i] ux /= rho uy /= rho uz_field /= rho return rho, ux, uy, uz_field # ━━━━━━━━━━━━━━━━━━━━━━━ Simulation ━━━━━━━━━━━━━━━━━━━━━━━━━━ def run_d3q19(device_id, n_steps, vis, u0, flag_host): """Compile v2 kernel, run D3Q19 SRT, return DDF.""" omega = 1.0 / (3.0 * vis + 0.5) cuda.init() dev = cuda.Device(device_id) ctx = dev.make_context() print(f"[GPU {device_id}] {dev.name()}") try: configure_macros_3d(vis, u0) compiler.compile_kernel_v2() ptx_path = compiler.kernel_path("kernel_v2.ptx") mod = cuda.module_from_file(ptx_path) # Get kernels (extern "C" entries from kernel_v2.cu) init_fn = mod.get_function("InitTubeFlow_v2") step_fn = mod.get_function("OneStep") # Set d_params.omega via __constant__ memory params_ptr, params_size = mod.get_global("d_params") # LBMParams struct layout (see params.cuh): # Nx(4) Ny(4) Nz(4) N(8) omega(4) omega_bulk(4) fx(4) fy(4) fz(4) # rho_ref(4) u_inlet(4) n_objects(4) # Pack: unsigned int Nx, Ny, Nz; unsigned long N; float omega, omega_bulk, fx, fy, fz, rho_ref, u_inlet; unsigned int n_objects params_data = struct.pack('IIIQfffffffI', NX, NY, NZ, TOTAL, omega, 0.0, # omega, omega_bulk 0.0, 0.0, 0.0, # fx, fy, fz RHO, u0, # rho_ref, u_inlet 0) # n_objects # Pad to match struct size if len(params_data) < params_size: params_data += b'\x00' * (params_size - len(params_data)) cuda.memcpy_htod(params_ptr, params_data) # Allocate 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(TOTAL * 4) delta_gpu = cuda.mem_alloc(4) action_gpu = cuda.mem_alloc(4) obs_gpu = cuda.mem_alloc(4) # Dummy arrays cuda.memset_d32(indx_gpu, 0, TOTAL) cuda.memset_d32(delta_gpu, 0, 1) cuda.memset_d32(action_gpu, 0, 1) cuda.memset_d32(obs_gpu, 0, 1) block = (NT, 1, 1) grid = (NX // NT, NY, NZ) # Initialize parabolic flow init_fn(flag_gpu, ddf_gpu, block=block, grid=grid) ctx.synchronize() # Overwrite flags with cylinder geometry cuda.memcpy_htod(flag_gpu, flag_host) # Step loop print(f" Running {n_steps} steps (NX={NX}, NY={NY}, NZ={NZ}, omega={omega:.4f})...") 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 if (i + 1) % 5000 == 0: ctx.synchronize() elapsed = time.time() - t0 mlups = TOTAL * (i + 1) / elapsed / 1e6 print(f" step {i+1}/{n_steps} ({mlups:.1f} MLUPS)") ctx.synchronize() dt = time.time() - t0 mlups = TOTAL * n_steps / dt / 1e6 print(f" Done: {dt:.1f}s, {mlups:.1f} MLUPS") # Copy back ddf = np.zeros(TOTAL * NQ, dtype=np.float32) cuda.memcpy_dtoh(ddf, ddf_gpu) return ddf finally: ctx.pop() # ━━━━━━━━━━━━━━━━━━━━━━━ Visualization ━━━━━━━━━━━━━━━━━━━━━━━ def plot_d3q19_cylinder(ddf, flag, Re, u0, out_path): """4-panel figure at z=NZ/2: vel-mag, vorticity, streamlines, u(y) profile.""" z_mid = NZ // 2 rho, ux, uy, uz = extract_fields_3d(ddf, z_mid) vel_mag = np.sqrt(ux**2 + uy**2 + uz**2) mask2d = ((flag.reshape(NZ, NY, NX)[z_mid] & (SOLID_FLAG | OBSTACLE_FLAG)) != 0) vel_masked = np.ma.array(vel_mag, mask=mask2d) fig, axes = plt.subplots(2, 2, figsize=(16, 10)) # (a) Velocity magnitude ax = axes[0, 0] im = ax.imshow(vel_masked, origin='lower', aspect='auto', cmap='jet', extent=[0, NX, 0, NY]) circ = plt.Circle((CYL_CX, CYL_CY), CYL_R, fill=True, color='gray', alpha=0.7) ax.add_patch(circ) plt.colorbar(im, ax=ax, label='|u|') ax.set_title(f'Velocity magnitude (z={z_mid})') ax.set_xlabel('x'); ax.set_ylabel('y') # (b) Vorticity ω_z = ∂v/∂x − ∂u/∂y ax = axes[0, 1] dvdx = np.gradient(uy, axis=1) dudy = np.gradient(ux, axis=0) omega_z = dvdx - dudy om_masked = np.ma.array(omega_z, mask=mask2d) vmax = np.percentile(np.abs(omega_z[~mask2d]), 99) if np.any(~mask2d) else 1e-3 im = ax.imshow(om_masked, origin='lower', aspect='auto', cmap='RdBu_r', extent=[0, NX, 0, NY], vmin=-vmax, vmax=vmax) circ2 = plt.Circle((CYL_CX, CYL_CY), CYL_R, fill=True, color='gray', alpha=0.7) ax.add_patch(circ2) plt.colorbar(im, ax=ax, label='ω_z') ax.set_title('Vorticity ω_z') ax.set_xlabel('x'); ax.set_ylabel('y') # (c) Streamlines ax = axes[1, 0] X, Y = np.meshgrid(np.arange(NX), np.arange(NY)) ux_c = ux.copy(); ux_c[mask2d] = 0 uy_c = uy.copy(); uy_c[mask2d] = 0 speed = np.sqrt(ux_c**2 + uy_c**2) ax.streamplot(X, Y, ux_c, uy_c, color=speed, cmap='jet', density=2.5, linewidth=0.7) circ3 = plt.Circle((CYL_CX, CYL_CY), CYL_R, fill=True, color='gray', alpha=0.7) ax.add_patch(circ3) ax.set_xlim(0, NX); ax.set_ylim(0, NY) ax.set_aspect('auto') ax.set_title('Streamlines') ax.set_xlabel('x'); ax.set_ylabel('y') # (d) u_x(y) profiles at different x stations ax = axes[1, 1] y_arr = np.arange(NY) # Analytical parabolic inlet yy = (y_arr - 0.5 * (NY - 1)) / (NY - 2.0) u_ana = u0 * 1.5 * (1 - 4 * yy**2) x_stations = [NX // 8, NX // 4, NX // 2, 3 * NX // 4] for xs in x_stations: ax.plot(ux[:, xs], y_arr, label=f'x={xs}') ax.plot(u_ana, y_arr, 'k--', lw=1.5, label='Analytical inlet') ax.set_xlabel('u_x'); ax.set_ylabel('y') ax.set_title('u_x(y) profiles') ax.legend(fontsize=8) ax.grid(True, alpha=0.3) fig.suptitle(f'D3Q19 SRT Cylinder — Re={Re:.0f}, D={2*CYL_R:.0f}, ' f'Grid={NX}×{NY}×{NZ}', 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='D3Q19 SRT Cylinder Flow') parser.add_argument('--device', type=int, default=0, help='CUDA device ID (default: 0)') parser.add_argument('--re', type=float, default=100.0, help='Reynolds number based on diameter (default: 100)') parser.add_argument('--u0', type=float, default=0.04, help='Inlet characteristic velocity (default: 0.04)') parser.add_argument('--steps', type=int, default=30000, help='Number of LBM steps (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) D = 2 * CYL_R vis, omega = compute_vis_omega(args.re, D, args.u0) print(f"\n===== D3Q19 SRT Cylinder Flow =====") print(f" Re = {args.re:.0f}, D = {D:.0f}, U0 = {args.u0}") print(f" ν = {vis:.6f}, ω = {omega:.4f}") if omega > 1.95: print(f" WARNING: omega={omega:.4f} is close to 2.0, stability may be poor.") print(f" Consider reducing U0 or Re.") flag = build_cylinder_3d(CYL_CX, CYL_CY, CYL_R) n_solid = np.sum(flag == SOLID_FLAG) n_fluid = np.sum(flag == FLUID_FLAG) print(f" Grid: {NX}×{NY}×{NZ} = {TOTAL} cells (fluid: {n_fluid}, solid: {n_solid})") print(f" Memory: ~{2 * TOTAL * NQ * 4 / 1e6:.0f} MB for double-buffer DDF") ddf = run_d3q19(args.device, args.steps, vis, args.u0, flag) plot_d3q19_cylinder(ddf, flag, args.re, args.u0, os.path.join(out_dir, f'cylinder_d3q19_re{int(args.re)}.png')) print("\nDone.") if __name__ == '__main__': main()