CelerisLab/tests/validation/taylor_green/run_case.py

"""One bounded GPU case, with explicit opt-in experimental explosion."""

import argparse
import hashlib
import json
import os
import shutil
from pathlib import Path
import subprocess
import sys
import time

import numpy as np
from oracle import U, TIMES, LIMITS, initial, metrics, rectangles

NUMERICAL_REVISION = "a135252ae78b47b57d7e77d19001700b3bedac8b"
SOURCE_SHA256 = "a5dbfd804befce868067d79a23e5fd47dce7f505b3f9b1bb2145fb57cb5f2178"
BUILD_CAPTURE_SHA256 = "3dbaa6547cb0a4cb1541054f3b64ad8b6e18cc70527ac6166e56b944e9bd464a"
REFERENCE_HARNESS = "11901a7f77557a901e017a171eeae1758d584fa5"
INITIALIZATION = "physical FP64 equilibrium G at actual cell centers; rho=1+3p"
SHEAR_WAVE_SHA256 = "81d5c8ee60864372887cddf9239580c82fb6750d956038f1e3ee9de61aa4e877"
SCALED_DIAGNOSTIC = "scaled_n256"
SCALED_SHEAR = "scaled_shear_n256"  # mutually exclusive with scaled_n256 (one flow each)


def flow_functions(flow="taylor_green"):
    """(initial, metrics, limits, initialization); the default is the frozen TG oracle."""
    if flow == "taylor_green":
        return initial, metrics, LIMITS, INITIALIZATION
    if flow == "shear_wave":
        import shear_wave
        return shear_wave.initial, shear_wave.metrics, shear_wave.LIMITS, shear_wave.INITIALIZATION
    raise ValueError(f"unknown flow {flow!r}")


def finite_row(row):
    """Every scalar or list-valued diagnostic is finite (TG rows are all scalars)."""
    return all(np.isfinite(np.asarray(value, dtype=float)).all() for value in row.values())


def sample(run, masks, step, n, viscosity, initial_mass, metrics, depth, label="sample"):
    arrays = (run.download(),) if depth == 1 else run.download()
    row = metrics(arrays, masks, step, n, viscosity, initial_mass)
    if not finite_row(row):
        raise FloatingPointError(f"nonfinite {label} diagnostics")
    return row


def sample_steps(n, early_pressure=False):
    """Root synchronization steps; early ladder matches the 707afe6 diagnostic."""
    main = [round(tstar * n / U) for tstar in TIMES]
    early = list(range(0, 4 * n + 1, n // 16)) if early_pressure else []
    return main, early


def runtime_source(repo):
    """Fail before CUDA if Python or the compiler snapshot comes from another tree."""
    import CelerisLab
    from CelerisLab.cuda import compiler_v2

    paths = dict(module=Path(CelerisLab.__file__).resolve(),
                 compiler=Path(compiler_v2.__file__).resolve(),
                 kernel_root=Path(compiler_v2.kernel_path("blocks.cu")).resolve().parent)
    package = repo / "src/CelerisLab"
    expected = dict(module=package / "__init__.py", compiler=package / "cuda/compiler_v2.py",
                    kernel_root=package / "lbm/kernels")
    if paths != expected:
        raise RuntimeError("Wrong imported source or CUDA source/include root")
    return dict(checkout=str(repo), **{key: str(path) for key, path in paths.items()})


def provenance(repo):
    git = lambda *a: subprocess.check_output(
        ["git", "-C", str(repo), *a], text=True
    ).strip()
    source_hash = hashlib.sha256()
    for path in sorted((repo / "src").rglob("*")):
        if path.is_file() and "__pycache__" not in path.parts and path.suffix != ".pyc":
            source_hash.update(
                str(path.relative_to(repo)).encode() + b"\0" + path.read_bytes() + b"\0"
            )
    return dict(
        revision=git("rev-parse", "HEAD"),
        dirty=bool(git("status", "--porcelain")),
        source_tree_sha256=source_hash.hexdigest(),
        harness_sha256={
            p.name: hashlib.sha256(p.read_bytes()).hexdigest()
            for p in Path(__file__).parent.glob("*.py")
        },
    )


def verify_source(repo, revision, flow="taylor_green"):
    """Exact clean harness commit, fixed production tree, and reviewed capture helper."""
    import oracle
    from tests.validation.multigrid import check_projected_pairs as gate

    harness = repo / "tests/validation/taylor_green"
    if (Path(__file__).resolve() != harness / "run_case.py"
            or Path(oracle.__file__).resolve() != harness / "oracle.py"):
        raise RuntimeError("Wrong runner or oracle checkout")
    helper = repo / "tests/validation/multigrid/check_projected_pairs.py"
    if (Path(gate.__file__).resolve() != helper
            or hashlib.sha256(helper.read_bytes()).hexdigest() != BUILD_CAPTURE_SHA256):
        raise RuntimeError("Wrong build-capture helper")
    if gate.preflight(repo, revision, "candidate") != NUMERICAL_REVISION:
        raise RuntimeError("Wrong production revision")
    result = provenance(repo)
    if result["source_tree_sha256"] != SOURCE_SHA256:
        raise RuntimeError("Wrong numerical source fingerprint")
    result["build_capture_sha256"] = BUILD_CAPTURE_SHA256
    if flow == "shear_wave":
        import shear_wave
        module = harness / "shear_wave.py"
        if (Path(shear_wave.__file__).resolve() != module
                or hashlib.sha256(module.read_bytes()).hexdigest() != SHEAR_WAVE_SHA256):
            raise RuntimeError("Wrong shear_wave module checkout or content")
        result["shear_wave_sha256"] = SHEAR_WAVE_SHA256
    return result


def build_capture(report, save):
    from tests.validation.multigrid.check_projected_pairs import capture_builds
    return capture_builds(report, save)


def runtime_identity(ctx):
    import pycuda
    import pycuda.driver as cuda

    nvcc = shutil.which("nvcc") or "/usr/local/cuda/bin/nvcc"
    return dict(device=dict(name=ctx.device.name(), pci_bus_id=ctx.device.pci_bus_id(),
                            capability=list(ctx.device.compute_capability()),
                            driver=cuda.get_driver_version(), runtime=list(cuda.get_version())),
                nvcc_path=str(Path(nvcc).resolve()),
                nvcc_version=subprocess.check_output([nvcc, "--version"], text=True).strip(),
                python=sys.version, numpy=np.__version__, pycuda=pycuda.VERSION_TEXT)


def config(n, re):
    from CelerisLab.config import LBMConfig

    cfg = LBMConfig(
        nx=n,
        ny=n,
        viscosity=U * n / re,
        velocity=0.0,
        rho=1.0,
        lattice_model="D2Q9",
        collision="SRT",
        streaming="double_buffer",
        store_precision="FP32",
        ddf_shifting=True,
        les_enabled=False,
    )
    cfg.validate()
    return cfg


def workload(masks, base_steps):
    cells = [int(np.count_nonzero(m)) for m in masks]
    hops = [base_steps * 2**level for level in range(len(masks))]
    updates = [count * hop for count, hop in zip(cells, hops)]
    return dict(
        active_cells_by_level=cells,
        level_hops=hops,
        active_collision_updates_by_level=updates,
        active_collision_updates=sum(updates),
        root_intervals=base_steps,
        excludes="ghost streaming, halo copies, transfers, initialization and diagnostics",
    )


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--repo", type=Path, required=True)
    p.add_argument("--expected-revision", required=True)
    p.add_argument("--run-gpu", action="store_true")
    p.add_argument("--dry-run", action="store_true")
    p.add_argument("--re", type=int, choices=[100, 1000], required=True)
    p.add_argument("--n", type=int, choices=[32, 64, 128, 256, 512], required=True)
    p.add_argument("--depth", type=int, choices=[1, 2, 3], required=True)
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--explosion", choices=["constant", "projected_linear"], default="constant")
    p.add_argument("--diagnostic", choices=[SCALED_DIAGNOSTIC, SCALED_SHEAR],
                   help="explicit four-case N256/fine512 same-placement diagnostic "
                        "(scaled_n256: Taylor-Green; scaled_shear_n256: shear wave)")
    p.add_argument("--early-pressure", action="store_true",
                   help="also save 65 scalar pressure samples at t*=k*.001875, k=0..64")
    p.add_argument("--flow", choices=["taylor_green", "shear_wave"], default="taylor_green")
    args = p.parse_args()
    if args.flow == "shear_wave" and (args.re != 100 or args.depth == 3 or args.early_pressure):
        p.error("shear_wave requires --re 100, depth 1 or 2 and no --early-pressure")
    initial, metrics, limits, initialization = flow_functions(args.flow)
    if args.output.exists():
        p.error("Output already exists; choose a new path to retain prior evidence")
    if args.explosion != "constant" and args.depth != 2:
        p.error("projected_linear requires depth 2")
    if args.diagnostic:
        cases = {(256, 1, "constant"), (512, 1, "constant"),
                 (256, 2, "constant"), (256, 2, "projected_linear")}
        flow = "shear_wave" if args.diagnostic == SCALED_SHEAR else "taylor_green"
        if args.flow != flow:
            p.error(f"--diagnostic {args.diagnostic} requires --flow {flow}")
        if args.diagnostic == SCALED_DIAGNOSTIC and (
                args.re != 100 or (args.n, args.depth, args.explosion) not in cases
                or not args.early_pressure):
            p.error("scaled_n256 requires one of its four Re100 cases and --early-pressure")
        if args.diagnostic == SCALED_SHEAR and (args.n, args.depth, args.explosion) not in cases:
            p.error("scaled_shear_n256 requires one of its four Re100 cases")
    elif args.n == 512:
        p.error("N512 requires --diagnostic scaled_n256 or scaled_shear_n256")
    selected = os.environ.get("CUDA_VISIBLE_DEVICES", "")
    if not args.dry_run and (not args.run_gpu or not selected or "," in selected or selected == "-1"):
        p.error("--run-gpu and exactly one explicit CUDA_VISIBLE_DEVICES are required")
    repo = args.repo.resolve()
    sys.path[:0] = [str(repo / "src"), str(repo)]
    source = verify_source(repo, args.expected_revision, args.flow)
    runtime = runtime_source(repo)
    if args.dry_run:
        print(json.dumps(dict(source=source, runtime_source=runtime, dry_run=True,
                              **({"diagnostic": args.diagnostic} if args.diagnostic else {}))))
        return 0
    from CelerisLab.cuda.context import CudaContext
    from CelerisLab.lbm.blocks import BlockLayout2D
    from CelerisLab.lbm.block_stepper import PeriodicBlockStepper
    from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
    from CelerisLab.lbm.refinement import ACTIVE
    from CelerisLab.lbm.hierarchy import build_hierarchy, leaf_partition

    cfg = config(args.n, args.re)
    refinements = [] if args.depth == 1 else rectangles(args.n, args.depth)
    h = build_hierarchy(cfg, refinements)
    np.testing.assert_array_equal(leaf_partition(h), np.ones((args.n, args.n)))
    masks = [k == ACTIVE for k in h.kinds]
    main_steps, early_steps = sample_steps(args.n, args.early_pressure)
    report = dict(
        schema=1,
        source=source,
        numerical_revision=NUMERICAL_REVISION,
        reference_harness=REFERENCE_HARNESS,
        actual_build_sidecars=[],
        runtime_identity=None,
        runtime_source=runtime,
        historical_baseline_revision="9805383f9c50d867e8555471baf5dbda8333a652",
        explosion=args.explosion,
        n=args.n,
        re_L=args.re,
        depth=args.depth,
        refinements=refinements,
        U=U,
        mach=U * np.sqrt(3),
        viscosity=cfg.viscosity,
        tau_by_level=[0.5 + 3 * c.viscosity for c in h.configs],
        initialization=initialization,
        numerics=dict(lattice_model=cfg.lattice_model, collision=cfg.collision,
                      store_precision=cfg.store_precision, ddf_shifting=cfg.ddf_shifting,
                      streaming=cfg.streaming, les_enabled=cfg.les_enabled,
                      boundary="periodic", block_shape=[8, 8],
                      sample_phase="postcollision G at root synchronization"),
        cuda_visible_devices=selected,
        device_id=0,
        limits=limits,
        workload=workload(masks, round(TIMES[-1] * args.n / U)),
        rows=[],
        early_pressure=dict(enabled=args.early_pressure, rows=[],
                            sampled_max_pressure_l2=None),
        completed=False,
    )
    if args.flow != "taylor_green":
        import shear_wave
        report.update(flow=args.flow, mode=list(shear_wave.MODE),
                      polarization=shear_wave.E.tolist(), angle_deg=shear_wave.ANGLE_DEG,
                      phase=shear_wave.PHI, eps=shear_wave.eps(args.n))
    if args.diagnostic:
        report["diagnostic"] = args.diagnostic
    args.output.parent.mkdir(parents=True, exist_ok=True)
    def save():
        temporary = args.output.with_suffix(args.output.suffix + ".tmp")
        text = json.dumps(report, indent=2, allow_nan=False) + "\n"
        if len(text.encode()) > 1024 * 1024:
            raise RuntimeError("Scalar receipt exceeds 1 MiB cap")
        temporary.write_text(text)
        temporary.replace(args.output)

    save()
    started = time.monotonic()
    with CudaContext(0) as ctx, build_capture(report, save):
        report["runtime_identity"] = runtime_identity(ctx)
        save()
        if args.depth == 1:
            layout = BlockLayout2D.build(args.n, args.n, 8, 8, periodic=(True, True))
            run = PeriodicBlockStepper(cfg, layout, context=ctx)
        else:
            run = FixedPatchStepper(
                cfg,
                context=ctx,
                coarse_boundary="periodic",
                refinements=refinements,
                level_blocks=[(8, 8)] * args.depth,
                explosion=args.explosion,
            )
        with run:
            run.initialize(
                *[
                    initial(m.shape, level, args.n, cfg.viscosity)
                    for level, m in enumerate(masks)
                ]
            )
            initial_mass = None
            for step in sorted(set(main_steps + early_steps)):
                run.step(step - run.step_count)
                row = sample(run, masks, step, args.n, cfg.viscosity, initial_mass,
                             metrics, args.depth)
                row["step"] = step
                if initial_mass is None:
                    initial_mass = row["mass"]
                if step in main_steps:
                    report["rows"].append(row)
                if step in early_steps:
                    trace = report["early_pressure"]
                    trace["rows"].append({key: row[key] for key in ("step", "tstar", "pressure_l2")})
                    trace["sampled_max_pressure_l2"] = max(r["pressure_l2"] for r in trace["rows"])
                save()
                print(
                    json.dumps(
                        dict(
                            event="PROGRESS" if step in main_steps else "EARLY_PRESSURE",
                            n=args.n,
                            re_L=args.re,
                            depth=args.depth,
                            **(row if step in main_steps else report["early_pressure"]["rows"][-1]),
                        )
                    ),
                    flush=True,
                )
                if step == 0:
                    # Exercise real kernels/halos/transfers on every level before READY.
                    # This is the first measured trajectory step, not discarded warmup.
                    run.step(1)
                    sample(run, masks, 1, args.n, cfg.viscosity, initial_mass,
                           metrics, args.depth, "READY")
                    report["ready_after_step"] = 1
                    save()
                    print(
                        json.dumps(
                            dict(
                                event="READY",
                                step=1,
                                n=args.n,
                                re_L=args.re,
                                depth=args.depth,
                                workload=report["workload"],
                            )
                        ),
                        flush=True,
                    )
    report.update(
        completed=True,
        seconds=time.monotonic() - started,
        screening_pass=all(
            row[key] <= limit for row in report["rows"] for key, limit in limits.items()
        ),
    )
    save()
    print(
        json.dumps(
            {k: report[k] for k in ("n", "re_L", "depth", "screening_pass", "seconds")}
        ),
        flush=True,
    )
    return 0 if report["screening_pass"] else 1


if __name__ == "__main__":
    raise SystemExit(main())