CelerisLab/tests/validation/taylor_green/shear_nested.py

"""Fixed three-level/two-region shear preset (opt-in nested_shear_n128).

Four fresh cases: U128, U512 (finest control = 4N), constant R128 depth 3,
projected R128 depth 3. Same wave, phase, Re100/U=0.03, initialization, five
sample times, eps(root N), limits and floors as shear_wave. This changes
placement and refined volume as well as depth: no isolated-level cost or
resolution-trend label. shear_wave.py is reused unchanged (historical receipt
identities stay valid); only the depth/finest-control contract is generalized
here.
"""

import hashlib
import json
from pathlib import Path

import numpy as np

import shear_wave as sw
from oracle import CX, CY, U, centers

ROOT_N = 128
DEPTH = 3
FINEST_FACTOR = 2 ** (DEPTH - 1)  # finest-spacing uniform control is 4N
# Parent-frame (x0, y0, x1, y1) rectangles; level-wide, zero-origin unions.
REFINEMENTS = [
    [(16, 16, 60, 112), (68, 24, 112, 120)],     # L0 -> L1, root coordinates
    [(44, 44, 108, 212), (148, 60, 212, 228)],   # L1 -> L2, L1 coordinates
]
CASES = ((100, ROOT_N, 1, "constant"), (100, FINEST_FACTOR * ROOT_N, 1, "constant"),
         (100, ROOT_N, DEPTH, "constant"), (100, ROOT_N, DEPTH, "projected_linear"))
ROOT_STEPS = 16 * ROOT_N
LEAVES = [7936, 12288, 86016]
HOPS = [ROOT_STEPS * 2**level for level in range(DEPTH)]  # [2048, 4096, 8192]
WORK = sum(c * h for c, h in zip(LEAVES, HOPS))           # 771227648 per refined case
NU = [U * ROOT_N / sw.RE * 2**level for level in range(DEPTH)]  # [.0384, .0768, .1536]
TAU = [0.5 + 3 * nu for nu in NU]                                # [.6152, .7304, .9608]
SNAPSHOT_MAX_BYTES = 8 * 1024 * 1024
SNAPSHOT_KEYS = ("level", "i", "j", "rho", "ux", "uy", "meta")


def leaf_masks():
    """Independent numpy leaf masks of the fixed preset (root, L1, L2 level-wide grids)."""
    masks, covered = [], None
    for level in range(DEPTH):
        size = ROOT_N * 2**level
        inside = np.ones((size, size), bool) if level == 0 else covered
        if level < DEPTH - 1:
            child = np.zeros((size, size), bool)
            for x0, y0, x1, y1 in REFINEMENTS[level]:
                child[y0:y1, x0:x1] = True
            masks.append(inside & ~child)
            covered = np.kron(child, np.ones((2, 2), bool))
        else:
            masks.append(inside)
    return masks


def assess(base_control, finest_control, refined):
    """shear_wave.assess generalized to depth 3 with a 4N finest control.

    Same summaries, screens, budgets and outcome semantics; only the
    (depth, finest-control) shape contract differs.
    """
    if base_control is None or finest_control is None or not refined:
        return dict(status="INVALID_CONTROL", reason="missing control or refined run")
    n = base_control["n"]
    if (base_control["depth"], finest_control["depth"], finest_control["n"]) != (
            1, 1, FINEST_FACTOR * n):
        raise ValueError("controls must be uniform N and uniform 4N")
    if any(r["depth"] != DEPTH or r["n"] != n for r in refined):
        raise ValueError("refined runs must be depth 3 at root N")
    if len({r.get("explosion") for r in refined}) != len(refined):
        raise ValueError("duplicate refined explosion scheme")
    sources = {repr(r.get("source")) for r in [base_control, finest_control, *refined]}
    base, fine = sw._usable(base_control), sw._usable(finest_control)
    if base is None or fine is None or len(sources) != 1:
        return dict(status="INVALID_CONTROL", reason="incomplete/invalid control or mixed source")
    controls_pass = base["own_screens_pass"] and fine["own_screens_pass"]
    report = dict(base_control=base, finest_control=fine, refined={},
                  controls_pass=controls_pass)
    statuses = []
    for r in refined:
        try:
            s = sw.summarize(r)
        except (KeyError, TypeError, ValueError) as error:
            report["refined"][r["explosion"]] = dict(valid=False, reason=str(error))
            statuses.append("INVALID_RUN")
            continue
        s["valid"] = True
        s["budget"] = {
            key: (None if base[key] is None else max(sw.BASELINE_FACTOR * base[key], floor))
            for key, floor in sw.FLOORS.items()
        }
        s["budget_pass"] = all(
            s[key] is not None and s["budget"][key] is not None and s[key] <= s["budget"][key]
            for key in sw.FLOORS)
        s["control_ratios"] = {
            name: {key: (s[key] / c[key] if s[key] is not None and c[key] else None)
                   for key in (*sw.FLOORS, "spurious_pressure_rms")}
            for name, c in (("root", base), ("finest", fine))
        }
        s["controlled_pass"] = controls_pass and s["own_screens_pass"] and s["budget_pass"]
        report["refined"][r["explosion"]] = s
        statuses.append("PASS" if s["controlled_pass"] else "FAIL")
    if "INVALID_RUN" in statuses:
        report["status"] = "INVALID_RUN"
    elif not controls_pass:
        report["status"] = "CONTROL_SCREEN_FAIL"
    else:
        report["status"] = "PASS" if all(v == "PASS" for v in statuses) else "FAIL"
    return report


def write_snapshot(path, arrays, masks, meta):
    """Final macro snapshot: packed ACTIVE leaves (level, i, j), rho, ux, uy only.

    No populations, no restart state, no time series. Returns the receipt block;
    a file over SNAPSHOT_MAX_BYTES is removed and reported, never kept.
    """
    path = Path(path)
    if path.exists():
        raise FileExistsError(path)
    cols = {k: [] for k in ("level", "i", "j", "rho", "ux", "uy")}
    counts = []
    for level, (g, mask) in enumerate(zip(arrays, masks)):
        if g.shape != (9, *mask.shape):
            raise ValueError("population/mask shape mismatch")
        values = g[:, mask].astype(float)
        rho = values.sum(axis=0)
        j, i = np.nonzero(mask)  # row = y index, column = x index
        cols["level"].append(np.full(rho.size, level, np.uint8))
        cols["i"].append(i.astype(np.uint16))
        cols["j"].append(j.astype(np.uint16))
        cols["rho"].append(rho)
        cols["ux"].append(CX @ values / rho)
        cols["uy"].append(CY @ values / rho)
        counts.append(int(rho.size))
    packed = {k: np.concatenate(v) for k, v in cols.items()}
    if not all(np.isfinite(packed[k]).all() for k in ("rho", "ux", "uy")):
        raise FloatingPointError("nonfinite snapshot macros")
    meta = dict(meta, leaf_counts=counts, cell_center="h*(index+0.5), h=2**-level",
                fields="rho, ux, uy from FP64 physical populations at ACTIVE leaves")
    with path.open("xb") as handle:
        np.savez_compressed(handle, meta=np.array(json.dumps(meta, sort_keys=True)), **packed)
    size = path.stat().st_size
    block = dict(file=path.name, bytes=size, leaf_counts=counts, max_bytes=SNAPSHOT_MAX_BYTES,
                 step=meta.get("step"), tstar=meta.get("tstar"))
    if size > SNAPSHOT_MAX_BYTES:
        path.unlink()
        return dict(block, kept=False, sha256=None, reason="exceeds snapshot byte cap")
    block.update(kept=True, sha256=hashlib.sha256(path.read_bytes()).hexdigest())
    return block


def read_snapshot(path):
    """Validate a snapshot: exact keys, no populations, consistent packed arrays."""
    with np.load(path, allow_pickle=False) as data:
        if tuple(sorted(data.files)) != tuple(sorted(SNAPSHOT_KEYS)):
            raise ValueError("unexpected snapshot contents")
        out = {k: data[k] for k in data.files}
    meta = json.loads(str(out.pop("meta")))
    n = len(out["rho"])
    if any(len(v) != n for v in out.values()) or sum(meta["leaf_counts"]) != n:
        raise ValueError("inconsistent snapshot arrays")
    return meta, out


def snapshot_meta(step, viscosity, explosion, source_revision):
    """Exact time/phase/config identity stored inside the snapshot."""
    return dict(step=step, tstar=step * U / ROOT_N, phase=sw.PHI, mode=list(sw.MODE),
                polarization=sw.E.tolist(), amplitude=U, re_L=sw.RE, n=ROOT_N, depth=DEPTH,
                explosion=explosion, viscosity=viscosity, tau_by_level=TAU,
                refinements=REFINEMENTS, source_revision=source_revision,
                origin="level-wide zero-origin grids", time="root synchronization, postcollision")