328 lines
12 KiB
Python
328 lines
12 KiB
Python
#!/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()
|