363 lines
13 KiB
Python
363 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
D2Q9 Regression Test — Poiseuille Channel + Cylinder Flow
|
||
==========================================================
|
||
Uses the ORIGINAL kernel.cu with same grid / BCs.
|
||
Produces matplotlib figures for visual validation.
|
||
|
||
Usage:
|
||
python tests/test_d2q9_visual.py --device 2
|
||
python tests/test_d2q9_visual.py --device 2 --cylinder
|
||
"""
|
||
|
||
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 to match original kernel settings."""
|
||
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,
|
||
}
|
||
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):
|
||
"""Compute rho, u, v from host DDF (original const.h ordering).
|
||
|
||
The original kernel uses DDF-shifting: stores f_shifted = f - w_i*RHO.
|
||
So sum(f_shifted) = rho - RHO (~0 for incompressible flow),
|
||
and momentum = sum(e_x * f_shifted) works because sum(w_i * e_x) = 0.
|
||
"""
|
||
f = ddf_host.reshape(NQ, NY, NX)
|
||
rho = np.sum(f, axis=0) + RHO # un-shift: rho = sum(f_shifted) + RHO
|
||
u = (f[1] + f[5] + f[8] - f[3] - f[6] - f[7]) / RHO
|
||
v = (f[2] + f[5] + f[6] - f[4] - f[7] - f[8]) / RHO
|
||
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):
|
||
"""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)
|
||
compiler.compile_kernel()
|
||
ptx_path = compiler.kernel_path("kernel.ptx")
|
||
mod = cuda.module_from_file(ptx_path)
|
||
step_fn = mod.get_function("OneStep")
|
||
init_fn = mod.get_function("InitTubeFlow")
|
||
|
||
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):
|
||
"""3-panel figure: velocity mag, u(y) profile, pressure along centerline."""
|
||
rho, u, v = extract_fields(ddf)
|
||
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):
|
||
"""3-panel figure: velocity magnitude (zoom), vorticity, streamlines."""
|
||
rho, u, v = extract_fields(ddf)
|
||
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('--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)
|
||
|
||
# ---- 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)
|
||
plot_poiseuille(ddf, flag, os.path.join(out_dir, 'poiseuille_d2q9.png'))
|
||
|
||
# Error metric
|
||
rho, u, v = extract_fields(ddf)
|
||
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)
|
||
plot_cylinder(ddf2, flag2, cyl_cx, cyl_cy, cyl_r,
|
||
os.path.join(out_dir, 'cylinder_d2q9.png'))
|
||
|
||
print("\nDone.")
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|