CelerisLab/tests/test_high_re_validation.py
2026-03-29 22:16:20 +08:00

548 lines
19 KiB
Python

#!/usr/bin/env python3
"""
High-Re Validation (kernel_v2)
==============================
Unified validation script for high-Re runs with optional LES.
Default targets:
- 2D D2Q9: Re=5000
- 3D D3Q19: Re=3000
The script configures macros.h temporarily, compiles kernel_v2, runs the case,
and restores macros.h automatically.
"""
import argparse
import json
import os
import struct
import sys
import time
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "src"))
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pycuda.driver as cuda
from CelerisLab.cuda import compiler
FLUID_FLAG = 0x01
SOLID_FLAG = 0x02
OBSTACLE_FLAG = 0x04
def collision_name(model):
return {0: "SRT", 1: "TRT", 2: "MRT"}.get(model, f"M{model}")
def make_case_tag(cfg):
les_tag = "LES" if cfg["use_les"] else "NoLES"
return (
f"{cfg['name']}_Re{int(cfg['target_re'])}_"
f"{collision_name(cfg['collision_model'])}_{les_tag}_"
f"OM{int(cfg['outlet_mode'])}_WMAX{cfg['omega_collision_max']:.3f}"
)
def validate_case(rho):
nan_count = int(np.isnan(rho).sum())
if nan_count > 0:
return False, "NaN detected"
rho_min = float(np.min(rho))
rho_max = float(np.max(rho))
if rho_min <= 0.0:
return False, "Non-positive density"
if rho_max >= 2.0:
return False, "Density blow-up"
return True, "OK"
def plot_case(cfg, host_ddf, out_dir):
nq = cfg["nq"]
nx, ny, nz = cfg["nx"], cfg["ny"], cfg["nz"]
flag = cfg["flag"]
tag = make_case_tag(cfg)
out_path = os.path.join(out_dir, f"{tag}.png")
if nq == 9:
f = host_ddf.reshape(nq, ny, nx)
rho = np.sum(f, axis=0)
ux = np.zeros_like(rho)
uy = np.zeros_like(rho)
cx = [0, 1, -1, 0, 0, 1, -1, 1, -1]
cy = [0, 0, 0, 1, -1, 1, -1, -1, 1]
for i in range(nq):
ux += cx[i] * f[i]
uy += cy[i] * f[i]
rho_safe = np.where(np.abs(rho) > 1.0e-12, rho, 1.0)
ux /= rho_safe
uy /= rho_safe
vel = np.sqrt(ux * ux + uy * uy)
mask = flag.reshape(ny, nx) != FLUID_FLAG
vel_m = np.ma.array(vel, mask=mask)
vort = np.gradient(uy, axis=1) - np.gradient(ux, axis=0)
vort_m = np.ma.array(vort, mask=mask)
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
im0 = axes[0].imshow(vel_m, origin="lower", aspect="auto", cmap="turbo")
plt.colorbar(im0, ax=axes[0], label="|u|")
axes[0].set_title("Velocity Magnitude")
vmax = np.percentile(np.abs(vort[~mask]), 99) if np.any(~mask) else 1e-6
vmax = max(vmax, 1.0e-6)
im1 = axes[1].imshow(vort_m, origin="lower", aspect="auto", cmap="RdBu_r", vmin=-vmax, vmax=vmax)
plt.colorbar(im1, ax=axes[1], label="vorticity")
axes[1].set_title("Vorticity")
X, Y = np.meshgrid(np.arange(nx), np.arange(ny))
ux_s = np.ma.array(ux, mask=mask)
uy_s = np.ma.array(uy, mask=mask)
speed = np.ma.sqrt(ux_s * ux_s + uy_s * uy_s)
axes[2].streamplot(X, Y, ux_s, uy_s, color=speed, cmap="viridis", density=2.0, linewidth=0.7)
axes[2].set_xlim(0, nx)
axes[2].set_ylim(0, ny)
axes[2].set_title("Streamlines")
fig.suptitle(tag)
fig.tight_layout()
fig.savefig(out_path, dpi=150)
plt.close(fig)
return out_path
# D3Q19: visualize mid-z slice
f = host_ddf.reshape(nq, nz, ny, nx)
z0 = nz // 2
fs = f[:, z0, :, :]
rho = np.sum(fs, axis=0)
ux = np.zeros_like(rho)
uy = np.zeros_like(rho)
uz = np.zeros_like(rho)
cx = np.array([0, 1,-1, 0, 0, 0, 0, 1,-1, 1,-1, 0, 0, 1,-1, 1,-1, 0, 0])
cy = np.array([0, 0, 0, 1,-1, 0, 0, 1,-1, 0, 0, 1,-1,-1, 1, 0, 0, 1,-1])
cz = np.array([0, 0, 0, 0, 0, 1,-1, 0, 0, 1,-1, 1,-1, 0, 0,-1, 1,-1, 1])
for i in range(nq):
ux += cx[i] * fs[i]
uy += cy[i] * fs[i]
uz += cz[i] * fs[i]
rho_safe = np.where(np.abs(rho) > 1.0e-12, rho, 1.0)
ux /= rho_safe
uy /= rho_safe
uz /= rho_safe
vel = np.sqrt(ux * ux + uy * uy + uz * uz)
mask3 = flag.reshape(nz, ny, nx)[z0] != FLUID_FLAG
vel_m = np.ma.array(vel, mask=mask3)
vort = np.gradient(uy, axis=1) - np.gradient(ux, axis=0)
vort_m = np.ma.array(vort, mask=mask3)
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
im0 = axes[0].imshow(vel_m, origin="lower", aspect="auto", cmap="turbo")
plt.colorbar(im0, ax=axes[0], label="|u|")
axes[0].set_title("Velocity Magnitude (z-mid)")
vmax = np.percentile(np.abs(vort[~mask3]), 99) if np.any(~mask3) else 1.0e-6
vmax = max(vmax, 1.0e-6)
im1 = axes[1].imshow(vort_m, origin="lower", aspect="auto", cmap="RdBu_r", vmin=-vmax, vmax=vmax)
plt.colorbar(im1, ax=axes[1], label="vorticity")
axes[1].set_title("Vorticity (z-mid)")
X, Y = np.meshgrid(np.arange(nx), np.arange(ny))
ux_s = np.ma.array(ux, mask=mask3)
uy_s = np.ma.array(uy, mask=mask3)
speed = np.ma.sqrt(ux_s * ux_s + uy_s * uy_s)
axes[2].streamplot(X, Y, ux_s, uy_s, color=speed, cmap="viridis", density=2.0, linewidth=0.7)
axes[2].set_xlim(0, nx)
axes[2].set_ylim(0, ny)
axes[2].set_title("Streamlines (z-mid)")
fig.suptitle(tag)
fig.tight_layout()
fig.savefig(out_path, dpi=150)
plt.close(fig)
return out_path
def compute_vis_omega(reynolds, diameter, u0):
vis = u0 * diameter / reynolds
omega = 1.0 / (3.0 * vis + 0.5)
return vis, omega
def set_macros(nx, ny, nz, dim, nq, vis, u0, collision_model, use_les, les_cs,
outlet_mode, outlet_backflow_clamp, outlet_blend_alpha,
omega_collision_max):
lines = compiler.read_lines(compiler.kernel_path("macros.h"))
defs = {
"MULT_GPU": "False",
"NT": 128,
"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": "1.0",
"U0": u0,
"N_OBJS": 0,
"COLLISION_MODEL": collision_model,
"STREAMING_MODEL": 0,
"STORE_PRECISION": 0,
"USE_DDF_SHIFTING": 0,
"USE_LES": int(use_les),
"LES_CS": f"{les_cs:.6f}f",
"INLET_PROFILE": 0,
"OUTLET_MODE": int(outlet_mode),
"OUTLET_BACKFLOW_CLAMP": int(outlet_backflow_clamp),
"OUTLET_BLEND_ALPHA": f"{float(outlet_blend_alpha):.3f}f",
"OMEGA_COLLISION_MAX": f"{float(omega_collision_max):.3f}f",
}
for name, value in defs.items():
lines = compiler.modify_macro(lines, name, value)
compiler.write_lines(compiler.kernel_path("macros.h"), lines)
def build_flags_2d(nx, ny, cx, cy, radius):
n = nx * ny
flag = np.ones(n, dtype=np.uint8) * FLUID_FLAG
for y in range(ny):
for x in range(nx):
k = y * nx + x
if y == 0 or y == ny - 1 or x == 0 or x == nx - 1:
flag[k] = SOLID_FLAG
elif (x - cx) ** 2 + (y - cy) ** 2 < radius ** 2:
flag[k] = OBSTACLE_FLAG
return flag
def build_flags_3d(nx, ny, nz, cx, cy, radius):
n = nx * ny * nz
flag = np.ones(n, 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
if y == 0 or y == ny - 1 or x == 0 or x == nx - 1:
flag[k] = SOLID_FLAG
elif (x - cx) ** 2 + (y - cy) ** 2 < radius ** 2:
flag[k] = OBSTACLE_FLAG
return flag
def run_case(device_id, cfg):
nx, ny, nz = cfg["nx"], cfg["ny"], cfg["nz"]
dim, nq = cfg["dim"], cfg["nq"]
n = nx * ny * nz
set_macros(
nx=nx,
ny=ny,
nz=nz,
dim=dim,
nq=nq,
vis=cfg["vis"],
u0=cfg["u0"],
collision_model=cfg["collision_model"],
use_les=cfg["use_les"],
les_cs=cfg["les_cs"],
outlet_mode=cfg["outlet_mode"],
outlet_backflow_clamp=cfg["outlet_backflow_clamp"],
outlet_blend_alpha=cfg["outlet_blend_alpha"],
omega_collision_max=cfg["omega_collision_max"],
)
compiler.compile_kernel_v2()
cuda.init()
dev = cuda.Device(device_id)
ctx = dev.make_context()
try:
mod = cuda.module_from_file(compiler.kernel_path("kernel_v2.ptx"))
init_fn = mod.get_function("InitTubeFlow_v2")
step_fn = mod.get_function("OneStep")
params_ptr, params_size = mod.get_global("d_params")
params_data = struct.pack(
"IIIQfffffffI",
nx,
ny,
nz,
n,
cfg["omega"],
1.1,
0.0,
0.0,
0.0,
1.0,
cfg["u0"],
0,
)
if len(params_data) < params_size:
params_data += b"\x00" * (params_size - len(params_data))
cuda.memcpy_htod(params_ptr, params_data)
fsize = n * nq * 4
d_fi = cuda.mem_alloc(fsize)
d_fi2 = cuda.mem_alloc(fsize)
d_flag = cuda.mem_alloc(n)
d_indx = cuda.mem_alloc(n * 4)
d_delta = cuda.mem_alloc(4)
d_action = cuda.mem_alloc(4)
d_obs = cuda.mem_alloc(4)
cuda.memset_d32(d_indx, 0, n)
cuda.memset_d32(d_delta, 0, 1)
cuda.memset_d32(d_action, 0, 1)
cuda.memset_d32(d_obs, 0, 1)
block = (128, 1, 1)
grid = ((nx + 127) // 128, ny, nz)
init_fn(d_flag, d_fi, block=block, grid=grid)
cuda.memcpy_dtod(d_fi2, d_fi, fsize)
cuda.memcpy_htod(d_flag, cfg["flag"])
t0 = time.time()
for step in range(cfg["steps"]):
step_fn(d_flag, d_fi, d_fi2, d_indx, d_delta, d_action, d_obs, block=block, grid=grid)
d_fi, d_fi2 = d_fi2, d_fi
if (step + 1) % cfg["report_every"] == 0:
cuda.Context.synchronize()
host = np.empty(n * nq, dtype=np.float32)
cuda.memcpy_dtoh(host, d_fi)
if nq == 9:
rho = host.reshape(nq, ny, nx).sum(axis=0)
c = float(rho[ny // 2, nx // 2])
else:
rho = host.reshape(nq, nz, ny, nx).sum(axis=0)
c = float(rho[nz // 2, ny // 2, nx // 2])
nan_count = int(np.isnan(rho).sum())
print(f" step {step+1:7d}: rho_center={c:.6f}, nan={nan_count}")
if nan_count > 0:
break
cuda.Context.synchronize()
elapsed = time.time() - t0
host = np.empty(n * nq, dtype=np.float32)
cuda.memcpy_dtoh(host, d_fi)
if nq == 9:
rho = host.reshape(nq, ny, nx).sum(axis=0)
center = float(rho[ny // 2, nx // 2])
else:
rho = host.reshape(nq, nz, ny, nx).sum(axis=0)
center = float(rho[nz // 2, ny // 2, nx // 2])
ok, reason = validate_case(rho)
plot_path = None
if cfg.get("save_plot", True):
plot_path = plot_case(cfg, host, cfg["out_dir"])
return {
"case_tag": make_case_tag(cfg),
"name": cfg["name"],
"target_re": cfg["target_re"],
"steps": cfg["steps"],
"mlups": float(n * cfg["steps"] / elapsed / 1e6),
"nan_count": int(np.isnan(rho).sum()),
"rho_center": center,
"rho_min": float(np.nanmin(rho)),
"rho_max": float(np.nanmax(rho)),
"omega": cfg["omega"],
"vis": cfg["vis"],
"collision_model": cfg["collision_model"],
"use_les": bool(cfg["use_les"]),
"les_cs": float(cfg["les_cs"]),
"outlet_mode": int(cfg["outlet_mode"]),
"outlet_backflow_clamp": int(cfg["outlet_backflow_clamp"]),
"outlet_blend_alpha": float(cfg["outlet_blend_alpha"]),
"omega_collision_max": float(cfg["omega_collision_max"]),
"pass": bool(ok),
"reason": reason,
"plot_path": plot_path,
}
finally:
ctx.pop()
def build_case_2d(re2d, steps2d, collision_model, use_les, les_cs, out_dir,
outlet_mode, outlet_backflow_clamp, outlet_blend_alpha,
omega_collision_max):
nx, ny, nz = 512, 256, 1
cx, cy, radius = 128.0, 128.0, 24.0
u0 = 0.03
vis, omega = compute_vis_omega(re2d, 2.0 * radius, u0)
return {
"name": "2D_D2Q9_highRe",
"dim": 2,
"nq": 9,
"nx": nx,
"ny": ny,
"nz": nz,
"flag": build_flags_2d(nx, ny, cx, cy, radius),
"u0": u0,
"vis": vis,
"omega": omega,
"steps": steps2d,
"report_every": max(steps2d // 10, 1),
"collision_model": collision_model,
"use_les": use_les,
"les_cs": les_cs,
"outlet_mode": int(outlet_mode),
"outlet_backflow_clamp": int(outlet_backflow_clamp),
"outlet_blend_alpha": float(outlet_blend_alpha),
"omega_collision_max": float(omega_collision_max),
"target_re": re2d,
"save_plot": True,
"out_dir": out_dir,
}
def build_case_3d(re3d, steps3d, collision_model, use_les, les_cs, out_dir,
outlet_mode, outlet_backflow_clamp, outlet_blend_alpha,
omega_collision_max):
nx, ny, nz = 256, 128, 32
cx, cy, radius = 64.0, 64.0, 12.0
u0 = 0.04
vis, omega = compute_vis_omega(re3d, 2.0 * radius, u0)
return {
"name": "3D_D3Q19_highRe",
"dim": 3,
"nq": 19,
"nx": nx,
"ny": ny,
"nz": nz,
"flag": build_flags_3d(nx, ny, nz, cx, cy, radius),
"u0": u0,
"vis": vis,
"omega": omega,
"steps": steps3d,
"report_every": max(steps3d // 10, 1),
"collision_model": collision_model,
"use_les": use_les,
"les_cs": les_cs,
"outlet_mode": int(outlet_mode),
"outlet_backflow_clamp": int(outlet_backflow_clamp),
"outlet_blend_alpha": float(outlet_blend_alpha),
"omega_collision_max": float(omega_collision_max),
"target_re": re3d,
"save_plot": True,
"out_dir": out_dir,
}
def build_comprehensive_cases(args, out_dir):
cases = []
# Coverage matrix at moderate Re to verify all changed pathways.
for cm in (0, 1, 2):
for les in (0, 1):
cases.append(build_case_2d(re2d=200.0, steps2d=args.matrix_steps2d,
collision_model=cm, use_les=les,
les_cs=args.les_cs, out_dir=out_dir,
outlet_mode=args.outlet_mode,
outlet_backflow_clamp=1,
outlet_blend_alpha=args.outlet_blend_alpha,
omega_collision_max=args.omega_collision_max))
cases.append(build_case_3d(re3d=200.0, steps3d=args.matrix_steps3d,
collision_model=cm, use_les=les,
les_cs=args.les_cs, out_dir=out_dir,
outlet_mode=args.outlet_mode,
outlet_backflow_clamp=1,
outlet_blend_alpha=args.outlet_blend_alpha,
omega_collision_max=args.omega_collision_max))
return cases
def main():
parser = argparse.ArgumentParser(description="High-Re validation for kernel_v2")
parser.add_argument("--device", type=int, default=0)
parser.add_argument("--re2d", type=float, default=5000.0)
parser.add_argument("--re3d", type=float, default=3000.0)
parser.add_argument("--steps2d", type=int, default=10000)
parser.add_argument("--steps3d", type=int, default=20000)
parser.add_argument("--collision", type=int, default=1, choices=[0, 1, 2],
help="0=SRT, 1=TRT, 2=MRT")
parser.add_argument("--use-les", action="store_true", default=True,
help="Enable Smagorinsky LES")
parser.add_argument("--no-les", action="store_false", dest="use_les")
parser.add_argument("--les-cs", type=float, default=0.16)
parser.add_argument("--outlet-mode", type=int, default=0, choices=[0, 1, 2],
help="0=non-equilibrium extrapolation, 1=zero-gradient copy, 2=damped blend")
parser.add_argument("--outlet-blend-alpha", type=float, default=0.70,
help="Blend alpha for outlet-mode 2")
parser.add_argument("--omega-collision-max", type=float, default=1.999,
help="Upper clamp for collision omega")
parser.add_argument("--only", choices=["2d", "3d", "both"], default="both")
parser.add_argument("--comprehensive", action="store_true",
help="Run coverage matrix: SRT/TRT/MRT x LES on/off for 2D and 3D")
parser.add_argument("--matrix-steps2d", type=int, default=1000)
parser.add_argument("--matrix-steps3d", type=int, default=600)
args = parser.parse_args()
macro_path = compiler.kernel_path("macros.h")
macro_backup = compiler.read_lines(macro_path)
out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "output")
os.makedirs(out_dir, exist_ok=True)
out_json = os.path.join(out_dir, "high_re_validation_summary.json")
try:
results = []
if args.only in ("2d", "both"):
c2 = build_case_2d(args.re2d, args.steps2d, args.collision, args.use_les,
args.les_cs, out_dir, args.outlet_mode, 1,
args.outlet_blend_alpha, args.omega_collision_max)
print("\n=== Running 2D high-Re case ===")
print(f" target Re={args.re2d:.1f}, vis={c2['vis']:.6e}, omega={c2['omega']:.6f}")
results.append(run_case(args.device, c2))
if args.only in ("3d", "both"):
c3 = build_case_3d(args.re3d, args.steps3d, args.collision, args.use_les,
args.les_cs, out_dir, args.outlet_mode, 1,
args.outlet_blend_alpha, args.omega_collision_max)
print("\n=== Running 3D high-Re case ===")
print(f" target Re={args.re3d:.1f}, vis={c3['vis']:.6e}, omega={c3['omega']:.6f}")
results.append(run_case(args.device, c3))
if args.comprehensive:
print("\n=== Running comprehensive coverage matrix ===")
for cfg in build_comprehensive_cases(args, out_dir):
print(f" {cfg['name']} Re={cfg['target_re']:.1f} "
f"{collision_name(cfg['collision_model'])} LES={int(cfg['use_les'])}")
results.append(run_case(args.device, cfg))
with open(out_json, "w", encoding="utf-8") as f:
json.dump(results, f, indent=2)
print("\n=== Summary ===")
n_pass = 0
for r in results:
if r["pass"]:
n_pass += 1
print(f"{r['name']}: nan={r['nan_count']}, rho_center={r['rho_center']:.6f}, "
f"rho[min,max]=[{r['rho_min']:.6f}, {r['rho_max']:.6f}], "
f"MLUPS={r['mlups']:.1f}, pass={r['pass']} ({r['reason']})")
if r.get("plot_path"):
print(f" plot: {r['plot_path']}")
print(f"Pass rate: {n_pass}/{len(results)}")
print(f"Saved: {out_json}")
finally:
compiler.write_lines(macro_path, macro_backup)
if __name__ == "__main__":
main()