CelerisLab/tests/test_stability_matrix.py
2026-04-17 21:50:38 +08:00

756 lines
26 KiB
Python

#!/usr/bin/env python3
"""
Stability Matrix Test
=====================
Tests three collision models (SRT/TRT/MRT) at low and high Re (with/without LES),
plus Esoteric-Pull streaming at low Re with SRT.
Outputs:
- Flow-field images (velocity, vorticity, streamlines) for each case
- Diagnostic JSON with stability metrics
- EsoPull vs double-buffer comparison plots
Usage:
python3 tests/test_stability_matrix.py [--device 0] [--steps 2000]
"""
import argparse
import json
import math
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
# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
FLUID = 0x01
SOLID = 0x02
OBSTACLE = 0x20 # fixed: was 0x04
COLLISION_NAMES = {0: "SRT", 1: "TRT", 2: "MRT"}
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def compute_vis_omega(re, diameter, u0):
vis = u0 * diameter / re
omega = 1.0 / (3.0 * vis + 0.5)
return vis, omega
def lattice_weights(nq):
if nq == 9:
return np.array([4/9] + [1/9]*4 + [1/36]*4, dtype=np.float32)
if nq == 19:
return np.array([1/3] + [1/18]*6 + [1/36]*12, dtype=np.float32)
raise ValueError(f"nq={nq}")
def build_flags_2d(nx, ny, cx, cy, radius):
flag = np.ones(nx * ny, dtype=np.uint8) * FLUID
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
elif (x - cx)**2 + (y - cy)**2 < radius**2:
flag[k] = OBSTACLE
return flag
def set_macros(nx, ny, dim, nq, vis, u0, collision_model, use_les, streaming_model,
omega_collision_max=1.999, inlet_profile=1, trt_magic_param=0.1875,
les_cs=0.16):
"""Write config/*.h files used by kernel_v2.cu."""
cfg_dir = os.path.join(os.path.dirname(compiler.kernel_path("config.h")), "config")
# config_grid.h
with open(os.path.join(cfg_dir, "config_grid.h"), "w") as f:
f.write(f"""\
// AUTO-GENERATED by test_stability_matrix.py
#ifndef CELERIS_CONFIG_GRID_H
#define CELERIS_CONFIG_GRID_H
#define NT 128
#define MULT_GPU 0
#define NX {nx}
#define NY {ny}
#define NZ 1
#define DIM {dim}
#define NQ {nq}
#endif
""")
# config_physics.h
with open(os.path.join(cfg_dir, "config_physics.h"), "w") as f:
f.write(f"""\
// AUTO-GENERATED by test_stability_matrix.py
#ifndef CELERIS_CONFIG_PHYSICS_H
#define CELERIS_CONFIG_PHYSICS_H
#define LBtype float
#define VIS {vis:.10f}
#define RHO 1.0
#define U0 {u0}
#define PI 3.141592653589793238
#define FLUID 0x01
#define SOLID 0x02
#define GAS 0x04
#define INTERFACE 0x08
#define SENSOR 0x10
#define OBSTACLE 0x20
#define V_TAYLOR 1
#endif
""")
# config_method.h
with open(os.path.join(cfg_dir, "config_method.h"), "w") as f:
f.write(f"""\
// AUTO-GENERATED by test_stability_matrix.py
#ifndef CELERIS_CONFIG_METHOD_H
#define CELERIS_CONFIG_METHOD_H
#define COLLISION_MODEL {collision_model}
#define STREAMING_MODEL {streaming_model}
#define STORE_PRECISION 0
#define USE_DDF_SHIFTING 0
#define USE_LES {int(use_les)}
#define LES_CS {les_cs:.6f}f
#define INLET_PROFILE {int(inlet_profile)}
#define OUTLET_MODE 0
#define OUTLET_BLEND_ALPHA 0.700f
#define OUTLET_BACKFLOW_CLAMP 1
#define OMEGA_COLLISION_MIN 0.01f
#define OMEGA_COLLISION_MAX {float(omega_collision_max):.3f}f
#define TRT_MAGIC_PARAM {float(trt_magic_param):.6f}f
#endif
""")
# config_objects.h
with open(os.path.join(cfg_dir, "config_objects.h"), "w") as f:
f.write("""\
// AUTO-GENERATED by test_stability_matrix.py
#ifndef CELERIS_CONFIG_OBJECTS_H
#define CELERIS_CONFIG_OBJECTS_H
#define N_OBJS 0
#endif
""")
def pack_d_params(nx, ny, omega, u0):
"""Pack LBMParams struct for __constant__ memory upload."""
return struct.pack(
"IIIQfffffffI",
nx, ny, 1, # Nx, Ny, Nz
nx * ny, # N
omega, # omega
1.1, # omega_bulk
0.0, 0.0, 0.0, # fx, fy, fz
1.0, # rho_ref
u0, # u_inlet
0, # n_objects
)
def impose_rest_on_nonfluid(flag, host_ddf, nq, nx, ny):
w = lattice_weights(nq)
f = host_ddf.reshape(nq, ny, nx)
nonfluid = flag.reshape(ny, nx) != FLUID
for i in range(nq):
f[i, nonfluid] = w[i]
return host_ddf
def compute_macros_2d(host_ddf, nq, nx, ny, flag):
"""Compute rho, ux, uy from DDF."""
cx9 = [0, 1, -1, 0, 0, 1, -1, 1, -1]
cy9 = [0, 0, 0, 1, -1, 1, -1, -1, 1]
f = host_ddf.reshape(nq, ny, nx)
rho = np.sum(f, axis=0)
ux = np.zeros_like(rho)
uy = np.zeros_like(rho)
for i in range(nq):
ux += cx9[i] * f[i]
uy += cy9[i] * f[i]
rho_safe = np.where(np.abs(rho) > 1e-12, rho, 1.0)
ux /= rho_safe
uy /= rho_safe
return rho, ux, uy
def diagnose(rho, ux, uy, flag, nx, ny):
"""Compute stability diagnostics."""
fluid = flag.reshape(ny, nx) == FLUID
nan_count = int(np.isnan(rho).sum())
rho_min = float(np.nanmin(rho))
rho_max = float(np.nanmax(rho))
mass = float(np.nansum(rho[fluid]))
vel = np.sqrt(ux**2 + uy**2)
# Ma check
ma_max = float(np.nanmax(vel[fluid])) * math.sqrt(3.0) if np.any(fluid) else 0.0
# Vorticity RMS in wake region
vort = np.gradient(uy, axis=1) - np.gradient(ux, axis=0)
wake_mask = fluid & (np.arange(nx)[None, :] > nx // 3)
vort_rms = float(np.sqrt(np.nanmean(vort[wake_mask]**2))) if np.any(wake_mask) else 0.0
stable = nan_count == 0 and rho_min > 0.0 and rho_max < 2.0
return {
"nan_count": nan_count,
"rho_min": rho_min,
"rho_max": rho_max,
"mass": mass,
"ma_max": ma_max,
"vort_rms": vort_rms,
"stable": stable,
}
def plot_flow(rho, ux, uy, flag, nx, ny, title, out_path):
"""Plot velocity magnitude, vorticity, and streamlines."""
fluid_mask = flag.reshape(ny, nx) != FLUID
vel = np.sqrt(ux**2 + uy**2)
vel_m = np.ma.array(vel, mask=fluid_mask)
vort = np.gradient(uy, axis=1) - np.gradient(ux, axis=0)
vort_m = np.ma.array(vort, mask=fluid_mask)
fig, axes = plt.subplots(1, 3, figsize=(18, 5))
# Velocity magnitude
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")
# Vorticity
vals = vort[~fluid_mask]
if vals.size > 0:
vmax = max(float(np.percentile(np.abs(vals), 99)), 1e-8)
else:
vmax = 1e-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")
# Streamlines
X, Y = np.meshgrid(np.arange(nx), np.arange(ny))
ux_s = np.ma.array(ux, mask=fluid_mask)
uy_s = np.ma.array(uy, mask=fluid_mask)
speed = np.ma.sqrt(ux_s**2 + uy_s**2)
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(title, fontsize=13)
fig.tight_layout()
fig.savefig(out_path, dpi=150)
plt.close(fig)
return out_path
# ---------------------------------------------------------------------------
# Case runner: double-buffer
# ---------------------------------------------------------------------------
def run_double_buffer(device_id, cfg, out_dir):
"""Run a case with standard double-buffer streaming."""
nx, ny = cfg["nx"], cfg["ny"]
nq = cfg["nq"]
n = nx * ny
set_macros(nx, ny, cfg["dim"], nq, cfg["vis"], cfg["u0"],
cfg["collision_model"], cfg["use_les"], streaming_model=0,
omega_collision_max=cfg.get("omega_max", 1.999),
trt_magic_param=cfg.get("trt_magic", 0.1875))
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")
# Upload d_params
params_ptr, params_size = mod.get_global("d_params")
params_data = pack_d_params(nx, ny, cfg["omega"], cfg["u0"])
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, 1)
init_fn(d_flag, d_fi, block=block, grid=grid)
cuda.memcpy_htod(d_flag, cfg["flag"])
host0 = np.empty(n * nq, dtype=np.float32)
cuda.memcpy_dtoh(host0, d_fi)
host0 = impose_rest_on_nonfluid(cfg["flag"], host0, nq, nx, ny)
cuda.memcpy_htod(d_fi, host0)
cuda.memcpy_htod(d_fi2, host0)
steps = cfg["steps"]
report = max(steps // 5, 1)
t0 = time.time()
diverged_step = None
for s in range(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 (s + 1) % report == 0:
cuda.Context.synchronize()
h = np.empty(n * nq, dtype=np.float32)
cuda.memcpy_dtoh(h, d_fi)
rho_c = h.reshape(nq, ny, nx).sum(axis=0)
nc = int(np.isnan(rho_c).sum())
center = float(rho_c[ny // 2, nx // 2])
print(f" step {s+1:6d}: rho_center={center:.6f} nan={nc}")
if nc > 0:
diverged_step = s + 1
break
cuda.Context.synchronize()
elapsed = time.time() - t0
host = np.empty(n * nq, dtype=np.float32)
cuda.memcpy_dtoh(host, d_fi)
rho, ux, uy = compute_macros_2d(host, nq, nx, ny, cfg["flag"])
diag = diagnose(rho, ux, uy, cfg["flag"], nx, ny)
diag["elapsed"] = elapsed
diag["mlups"] = n * steps / elapsed / 1e6 if elapsed > 0 else 0
diag["diverged_step"] = diverged_step
tag = cfg["tag"]
plot_path = plot_flow(rho, ux, uy, cfg["flag"], nx, ny, tag,
os.path.join(out_dir, f"{tag}.png"))
diag["plot"] = plot_path
return diag
finally:
ctx.pop()
# ---------------------------------------------------------------------------
# Case runner: Esoteric-Pull (single buffer)
# ---------------------------------------------------------------------------
def run_esopull(device_id, cfg, out_dir):
"""Run a case with Esoteric-Pull single-buffer streaming."""
nx, ny = cfg["nx"], cfg["ny"]
nq = cfg["nq"]
n = nx * ny
set_macros(nx, ny, cfg["dim"], nq, cfg["vis"], cfg["u0"],
cfg["collision_model"], cfg["use_les"], streaming_model=1,
omega_collision_max=cfg.get("omega_max", 1.999),
trt_magic_param=cfg.get("trt_magic", 0.1875))
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("InitEsoPull")
step_fn = mod.get_function("EsoPullStep")
# Upload d_params
params_ptr, params_size = mod.get_global("d_params")
params_data = pack_d_params(nx, ny, cfg["omega"], cfg["u0"])
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_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, 1)
init_fn(d_flag, d_fi, block=block, grid=grid)
cuda.memcpy_htod(d_flag, cfg["flag"])
# Note: for EsoPull, we don't impose_rest_on_nonfluid on the raw
# DDF because the data is stored in esoteric layout. InitEsoPull
# already stores rest equilibrium for solid nodes.
steps = cfg["steps"]
report = max(steps // 5, 1)
t0 = time.time()
diverged_step = None
for s in range(steps):
t_val = np.uint64(s) # timestep counter for load/store parity
step_fn(d_fi, d_flag, d_indx, d_delta, d_action, d_obs,
t_val, block=block, grid=grid)
if (s + 1) % report == 0:
cuda.Context.synchronize()
# For diagnostics, download raw DDF and decode from esopull layout
h = np.empty(n * nq, dtype=np.float32)
cuda.memcpy_dtoh(h, d_fi)
# Esoteric layout: at this point the DDF is in post-store layout
# for timestep s. To compute macros we need to "undo" the esoteric
# read pattern. A simpler approach: compute rho = sum(fi) per node.
# Because sum is invariant under slot permutation, rho is correct.
# But ux/uy need correct direction assignment.
# For diagnostic, use a simple sum-based stability check.
f_arr = h.reshape(nq, ny, nx)
rho_c = f_arr.sum(axis=0)
nc = int(np.isnan(rho_c).sum())
center = float(rho_c[ny // 2, nx // 2])
print(f" step {s+1:6d}: rho_center={center:.6f} nan={nc}")
if nc > 0:
diverged_step = s + 1
break
cuda.Context.synchronize()
elapsed = time.time() - t0
# For final macros, do one more step that also writes to rho/u arrays.
# But we don't have UpdateMacro for EsoPull yet. Instead, use the
# approach: run a "read-only" macro computation from the esoteric layout.
# For correctness, we load from the proper esoteric positions on host.
h = np.empty(n * nq, dtype=np.float32)
cuda.memcpy_dtoh(h, d_fi)
rho, ux, uy = _decode_esopull_macros(h, nq, nx, ny, cfg["flag"], steps)
diag = diagnose(rho, ux, uy, cfg["flag"], nx, ny)
diag["elapsed"] = elapsed
diag["mlups"] = n * steps / elapsed / 1e6 if elapsed > 0 else 0
diag["diverged_step"] = diverged_step
tag = cfg["tag"]
plot_path = plot_flow(rho, ux, uy, cfg["flag"], nx, ny, tag,
os.path.join(out_dir, f"{tag}.png"))
diag["plot"] = plot_path
return diag
finally:
ctx.pop()
def _decode_esopull_macros(host_ddf, nq, nx, ny, flag, last_t):
"""Decode macroscopic quantities from esoteric-pull layout on host.
After step t (0-based), the store was done at parity t.
The next load would use parity t+1. To read correct DDFs we mimic
load_f_esopull at t_read = last_t (the parity of the *next* step to execute).
"""
fi = host_ddf.reshape(nq, ny * nx) # fi[direction, node]
t_read = last_t # parity for the load that would happen next
cx9 = np.array([0, 1, -1, 0, 0, 1, -1, 1, -1], dtype=np.float32)
cy9 = np.array([0, 0, 0, 1, -1, 1, -1, -1, 1], dtype=np.float32)
# Compute neighbor table once
j_table = np.zeros((nq, ny * nx), dtype=np.int64)
for y in range(ny):
for x in range(nx):
k = y * nx + x
xp = (x + 1) % nx
xm = (x - 1) % nx
yp = (y + 1) % ny
ym = (y - 1) % ny
j_table[0, k] = k
j_table[1, k] = yp * nx + xp if nq > 1 else k # placeholder
j_table[2, k] = ym * nx + xm if nq > 2 else k
# D2Q9 neighbors: j[i] = neighbor in direction c_i
if nq == 9:
j_table[1, k] = y * nx + xp # +x
j_table[2, k] = y * nx + xm # -x
j_table[3, k] = yp * nx + x # +y
j_table[4, k] = ym * nx + x # -y
j_table[5, k] = yp * nx + xp # +x+y
j_table[6, k] = ym * nx + xm # -x-y
j_table[7, k] = ym * nx + xp # +x-y
j_table[8, k] = yp * nx + xm # -x+y
n = nx * ny
f_decoded = np.zeros((nq, n), dtype=np.float32)
f_decoded[0] = fi[0]
for i in range(1, nq, 2):
if t_read & 1:
# Odd: f[i] from fi[n, i], f[i+1] from fi[j[i], i+1]
f_decoded[i] = fi[i]
f_decoded[i + 1] = fi[i + 1, j_table[i]]
else:
# Even: f[i] from fi[n, i+1], f[i+1] from fi[j[i], i]
f_decoded[i] = fi[i + 1]
f_decoded[i + 1] = fi[i, j_table[i]]
f_decoded = f_decoded.reshape(nq, ny, nx)
rho = f_decoded.sum(axis=0)
rho_safe = np.where(np.abs(rho) > 1e-12, rho, 1.0)
ux = np.zeros_like(rho)
uy = np.zeros_like(rho)
for i in range(nq):
ux += cx9[i] * f_decoded[i]
uy += cy9[i] * f_decoded[i]
ux /= rho_safe
uy /= rho_safe
return rho, ux, uy
# ---------------------------------------------------------------------------
# Case builders
# ---------------------------------------------------------------------------
def build_cases(steps_low, steps_high):
"""Build the full test matrix."""
# Grid params (moderate size for fast testing)
nx, ny = 384, 192
cx_ob, cy_ob, radius = 96.0, 96.0, 18.0
u0 = 0.04
cases = []
for re_val, re_label, n_steps, use_les in [
(100.0, "Re100", steps_low, False),
(100.0, "Re100", steps_low, True),
(3000.0, "Re3000", steps_high, False),
(3000.0, "Re3000", steps_high, True),
]:
for cm in (0, 1, 2):
vis, omega = compute_vis_omega(re_val, 2.0 * radius, u0)
les_tag = "LES" if use_les else "noLES"
cm_name = COLLISION_NAMES[cm]
tag = f"DB_{re_label}_{cm_name}_{les_tag}"
cases.append({
"tag": tag,
"nx": nx, "ny": ny,
"dim": 2, "nq": 9,
"cx": cx_ob, "cy": cy_ob, "radius": radius,
"flag": build_flags_2d(nx, ny, cx_ob, cy_ob, radius),
"u0": u0,
"vis": vis,
"omega": omega,
"collision_model": cm,
"use_les": use_les,
"steps": n_steps,
"streaming": "double_buffer",
"omega_max": 1.999,
"trt_magic": 0.1875,
})
# EsoPull case: low Re, SRT only
re_eso = 100.0
vis_eso, omega_eso = compute_vis_omega(re_eso, 2.0 * radius, u0)
cases.append({
"tag": "EsoPull_Re100_SRT_noLES",
"nx": nx, "ny": ny,
"dim": 2, "nq": 9,
"cx": cx_ob, "cy": cy_ob, "radius": radius,
"flag": build_flags_2d(nx, ny, cx_ob, cy_ob, radius),
"u0": u0,
"vis": vis_eso,
"omega": omega_eso,
"collision_model": 0,
"use_les": False,
"steps": steps_low,
"streaming": "esopull",
"omega_max": 1.999,
"trt_magic": 0.1875,
})
return cases
# ---------------------------------------------------------------------------
# Comparison plot: EsoPull vs DoubleBuffer
# ---------------------------------------------------------------------------
def plot_comparison(results, out_dir):
"""Compare EsoPull and DoubleBuffer at matching Re/collision settings."""
eso_key = "EsoPull_Re100_SRT_noLES"
db_key = "DB_Re100_SRT_noLES"
eso = results.get(eso_key)
db = results.get(db_key)
if eso is None or db is None:
return None
fig, axes = plt.subplots(2, 3, figsize=(18, 10))
fig.suptitle("EsoPull vs DoubleBuffer — Re100 SRT noLES", fontsize=14)
labels = ["DoubleBuffer", "EsoPull"]
for row, (r, label) in enumerate([(db, labels[0]), (eso, labels[1])]):
vel_img = plt.imread(r["plot"]) if os.path.exists(r["plot"]) else None
if vel_img is not None:
axes[row, 0].imshow(vel_img)
axes[row, 0].set_title(f"{label}: flow field")
axes[row, 0].axis("off")
else:
axes[row, 0].text(0.5, 0.5, f"No image for {label}",
ha="center", va="center", transform=axes[row, 0].transAxes)
axes[row, 0].set_title(label)
# Metrics bar chart
metrics = {
"rho_min": r.get("rho_min", 0),
"rho_max": r.get("rho_max", 0),
"ma_max": r.get("ma_max", 0),
"vort_rms": r.get("vort_rms", 0),
}
bars = list(metrics.keys())
vals = [float(metrics[b]) for b in bars]
axes[row, 1].barh(bars, vals, color=["steelblue", "salmon", "green", "purple"])
axes[row, 1].set_title(f"{label}: diagnostics")
# Stability text
text_lines = [
f"stable: {r.get('stable', '?')}",
f"nan_count: {r.get('nan_count', '?')}",
f"mass: {r.get('mass', 0):.2f}",
f"MLUPS: {r.get('mlups', 0):.1f}",
f"diverged_step: {r.get('diverged_step', 'None')}",
]
axes[row, 2].text(0.1, 0.5, "\n".join(text_lines), fontsize=12,
family="monospace", va="center",
transform=axes[row, 2].transAxes)
axes[row, 2].set_title(f"{label}: summary")
axes[row, 2].axis("off")
fig.tight_layout()
cmp_path = os.path.join(out_dir, "esopull_vs_doublebuffer.png")
fig.savefig(cmp_path, dpi=150)
plt.close(fig)
return cmp_path
# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
parser = argparse.ArgumentParser(description="Stability matrix test")
parser.add_argument("--device", type=int, default=0)
parser.add_argument("--steps-low", type=int, default=3000,
help="Steps for low-Re cases")
parser.add_argument("--steps-high", type=int, default=6000,
help="Steps for high-Re cases")
parser.add_argument("--only-esopull", action="store_true",
help="Only run the EsoPull test")
args = parser.parse_args()
# Backup config/*.h files (kernel_v2.cu uses config.h, NOT macros.h)
cfg_dir = os.path.join(os.path.dirname(compiler.kernel_path("config.h")), "config")
config_files = ["config_grid.h", "config_physics.h", "config_method.h", "config_objects.h"]
config_backups = {}
for cf in config_files:
path = os.path.join(cfg_dir, cf)
with open(path, "r") as f:
config_backups[path] = f.read()
out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)),
"..", "output", "stability_matrix")
os.makedirs(out_dir, exist_ok=True)
cases = build_cases(args.steps_low, args.steps_high)
if args.only_esopull:
cases = [c for c in cases if c["streaming"] == "esopull"]
results = {}
try:
for i, cfg in enumerate(cases):
tag = cfg["tag"]
streaming = cfg["streaming"]
print(f"\n[{i+1}/{len(cases)}] {tag}")
print(f" Re={cfg['u0']*2*cfg['radius']/cfg['vis']:.0f}, "
f"omega={cfg['omega']:.4f}, "
f"collision={COLLISION_NAMES[cfg['collision_model']]}, "
f"LES={cfg['use_les']}, streaming={streaming}")
if streaming == "esopull":
diag = run_esopull(args.device, cfg, out_dir)
else:
diag = run_double_buffer(args.device, cfg, out_dir)
diag["tag"] = tag
diag["streaming"] = streaming
diag["collision"] = COLLISION_NAMES[cfg["collision_model"]]
diag["use_les"] = cfg["use_les"]
diag["re"] = cfg["u0"] * 2 * cfg["radius"] / cfg["vis"]
results[tag] = diag
status = "PASS" if diag["stable"] else "FAIL"
print(f" => {status}: rho=[{diag['rho_min']:.4f}, {diag['rho_max']:.4f}], "
f"nan={diag['nan_count']}, ma_max={diag['ma_max']:.4f}, "
f"MLUPS={diag['mlups']:.1f}")
# Comparison plot
cmp_path = plot_comparison(results, out_dir)
if cmp_path:
print(f"\nComparison plot: {cmp_path}")
# Summary table
print("\n" + "=" * 100)
print(f"{'Tag':<35s} {'Stream':<8s} {'Col':<5s} {'LES':<5s} "
f"{'Re':>6s} {'Stable':>7s} {'rho_min':>9s} {'rho_max':>9s} "
f"{'Ma_max':>8s} {'MLUPS':>7s}")
print("-" * 100)
for tag, r in results.items():
print(f"{tag:<35s} {r['streaming']:<8s} {r['collision']:<5s} "
f"{'Y' if r['use_les'] else 'N':<5s} "
f"{r['re']:6.0f} {'PASS' if r['stable'] else 'FAIL':>7s} "
f"{r['rho_min']:9.5f} {r['rho_max']:9.5f} "
f"{r['ma_max']:8.5f} {r['mlups']:7.1f}")
print("=" * 100)
# Save JSON
json_path = os.path.join(out_dir, "stability_matrix_results.json")
json_results = {}
for k, v in results.items():
jr = {}
for rk, rv in v.items():
if isinstance(rv, (np.integer, np.floating)):
jr[rk] = float(rv)
elif isinstance(rv, np.bool_):
jr[rk] = bool(rv)
else:
jr[rk] = rv
json_results[k] = jr
with open(json_path, "w") as f:
json.dump(json_results, f, indent=2)
print(f"\nResults saved: {json_path}")
finally:
for path, content in config_backups.items():
with open(path, "w") as f:
f.write(content)
if __name__ == "__main__":
main()