CelerisLab/tests/validation/taylor_green/shear_wave.py

"""CPU-only oblique shear-wave reference and composite-leaf diagnostics.

Convention: integer mode MODE=(2,1) on the L x L periodic root domain, so the
wavevector k = (2*pi/L)*(2,1) points 26.565 deg from +x (mode (1,2) would be
63.435 deg). Velocity is polarized along E = (-1,2)/sqrt(5), perpendicular to k.
The frozen Taylor-Green oracle is reused for constants and geometry only.
"""

import math

import numpy as np
from oracle import CX, CY, W, U, TIMES, centers

MODE = (2, 1)
PHI = 0.0
E = np.array([-MODE[1], MODE[0]], dtype=float) / math.hypot(*MODE)
ANGLE_DEG = math.degrees(math.atan2(MODE[1], MODE[0]))
RE = 100  # Re = U L / nu, the same nu as Taylor-Green Re100 at equal N
FIT_FROM_TSTAR = 0.12  # late-time samples only; reduces startup sensitivity
INITIALIZATION = (
    "physical FP64 equilibrium G at actual cell centers; rho=1 (exact p const)"
)
# Per-row screens. phase_drift_abs, residual_l2 and the run-level rate screen
# use eps(n) and are preregistered diagnostic choices, not error bounds.
LIMITS = dict(velocity_l2=0.05, pressure_rel=0.02, mass_relative=1e-5)
FLOORS = dict(rate_error=0.002, residual_l2=0.002, pressure_rel=0.02)
BASELINE_FACTOR = 2.0


def wavevector(length):
    return 2 * np.pi / length * np.array(MODE, dtype=float)


def decay_rate(length, viscosity):
    """Exact amplitude decay rate per root step: nu |k|^2."""
    k = wavevector(length)
    return viscosity * float(k @ k)


def eps(n):
    """(|k| h_root)^2 with h_root = 1: dimensionless, coefficient one."""
    k = wavevector(n)
    return float(k @ k)


def exact(x, y, t, length, viscosity, amplitude=U):
    k = wavevector(length)
    a = amplitude * np.exp(-decay_rate(length, viscosity) * t)
    s = np.sin(k[0] * x + k[1] * y + PHI)
    p = np.zeros_like(s)
    return a * E[0] * s, a * E[1] * s, p  # kinematic pressure, constant


def equilibrium(rho, ux, uy):
    # Same second-order D2Q9 equilibrium as oracle.initial.
    cu = CX[:, None, None] * ux + CY[:, None, None] * uy
    return (
        W[:, None, None]
        * rho
        * (1 + 3 * cu + 4.5 * cu * cu - 1.5 * (ux * ux + uy * uy))
    )


def initial(shape, level, length, viscosity):
    x, y = centers(shape, level)
    ux, uy, p = exact(x, y, 0.0, length, viscosity)
    return equilibrium(1 + 3 * p, ux, uy)


def _leaves(arrays, masks):
    accum = []
    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)
        if not np.isfinite(values).all():
            raise FloatingPointError("nonfinite leaf populations")
        rho = values.sum(axis=0)
        if (rho <= 0).any():
            raise FloatingPointError("nonpositive leaf density")
        x, y = centers(mask.shape, level)
        area = np.full(rho.size, 0.25**level)
        lev = np.full(rho.size, level)
        accum.append((rho, CX @ values / rho, CY @ values / rho, x[mask], y[mask], area, lev))
    return [np.concatenate(v) for v in zip(*accum)]


def metrics(arrays, masks, t, length, viscosity, initial_mass=None):
    rho, ux, uy, x, y, w, lev = _leaves(arrays, masks)
    volume = w.sum()
    if not np.isclose(volume, length**2, rtol=0, atol=1e-10):
        raise ValueError("leaf partition does not cover domain once")
    k = wavevector(length)
    theta = k[0] * x + k[1] * y + PHI
    s, c = np.sin(theta), np.cos(theta)
    v = ux * E[0] + uy * E[1]
    # Weighted least squares on the actual composite quadrature (no
    # continuous-orthogonality assumption).
    gram = np.array([[w @ (s * s), w @ (s * c)], [w @ (s * c), w @ (c * c)]])
    a_hat, b_hat = np.linalg.solve(gram, [w @ (v * s), w @ (v * c)])
    ex, ey, _ = exact(x, y, t, length, viscosity)
    exact_norm = w @ (ex * ex + ey * ey)
    fit = a_hat * s + b_hat * c
    longitudinal = (ux * k[0] + uy * k[1]) / math.sqrt(k @ k)
    amplitude = U * math.exp(-decay_rate(length, viscosity) * t)
    scale = amplitude * amplitude / 4  # fixed pressure scale, see docs
    p = (rho - 1) / 3
    p = p - w @ p / volume  # one global gauge
    mass = float(w @ rho)
    return dict(
        tstar=t * U / length,
        mass=mass,
        mass_relative=(
            0.0 if initial_mass is None else abs(mass - initial_mass) / initial_mass
        ),
        a_hat=float(a_hat),
        b_hat=float(b_hat),
        amplitude_exact=amplitude,
        amplitude_ratio=float(a_hat / amplitude),
        phase_drift=float(math.atan2(b_hat, a_hat)),
        phase_drift_abs=abs(float(math.atan2(b_hat, a_hat))),
        velocity_l2=float(
            np.sqrt(w @ ((ux - ex) ** 2 + (uy - ey) ** 2) / exact_norm)
        ),
        residual_l2=float(
            np.sqrt(w @ ((ux - fit * E[0]) ** 2 + (uy - fit * E[1]) ** 2) / exact_norm)
        ),
        longitudinal_l2=float(np.sqrt(w @ (longitudinal**2) / exact_norm)),
        pressure_rel=float(np.sqrt(w @ (p * p) / volume) / scale),
        pressure_max_rel=float(np.max(np.abs(p)) / scale),
        pressure_rel_by_level=[
            float(np.sqrt(w[lev == i] @ (p[lev == i] ** 2) / w[lev == i].sum()) / scale)
            for i in range(len(masks))
        ],
        rho_min=float(rho.min()),
        rho_max=float(rho.max()),
        max_speed=float(np.max(np.hypot(ux, uy))),
        leaf_cells=int(len(rho)),
    )


def fit_rate(rows, length, viscosity):
    """Least-squares decay rate of ln(a_hat) over late samples."""
    late = [r for r in rows if r["tstar"] >= FIT_FROM_TSTAR - 1e-12]
    if len(late) < 2 or any(r["a_hat"] <= 0 for r in late):
        raise ValueError("need >= 2 late samples with positive amplitude")
    t = np.array([r["step"] for r in late], dtype=float)
    slope = np.polyfit(t, np.log([r["a_hat"] for r in late]), 1)[0]
    exact_rate = decay_rate(length, viscosity)
    return dict(rate=float(-slope), exact_rate=exact_rate,
                rate_error=abs(-slope / exact_rate - 1))


def summarize(record):
    rows = record["rows"]
    n = record["n"]
    out = fit_rate(rows, n, record["viscosity"])
    out.update(
        eps=eps(n),
        complete=record["completed"] and len(rows) == len(TIMES),
        residual_l2=max(r["residual_l2"] for r in rows),
        phase_drift_abs=max(r["phase_drift_abs"] for r in rows),
        pressure_rel=max(r["pressure_rel"] for r in rows),
        velocity_l2=max(r["velocity_l2"] for r in rows),
        mass_relative=max(r["mass_relative"] for r in rows),
    )
    out["absolute_pass"] = out["complete"] and all(
        out[key] <= out["eps"] for key in ("rate_error", "phase_drift_abs", "residual_l2")
    ) and all(out[key] <= limit for key, limit in LIMITS.items())
    return out


def assess(base_control, fine_control, refined):
    """Uniform root-spacing and finest-spacing controls versus refined runs."""
    n = base_control["n"]
    if (base_control["depth"], fine_control["depth"], fine_control["n"]) != (1, 1, 2 * n):
        raise ValueError("controls must be uniform N and uniform 2N")
    if any(r["depth"] != 2 or r["n"] != n for r in refined):
        raise ValueError("refined runs must be depth 2 at root N")
    base, fine = summarize(base_control), summarize(fine_control)
    report = dict(base_control=base, fine_control=fine, refined={}, status="PASS")
    if not (base["absolute_pass"] and fine["absolute_pass"]):
        report["status"] = "INVALID_CONTROL"
    for r in refined:
        s = summarize(r)
        s["budget"] = {
            key: max(BASELINE_FACTOR * base[key], floor) for key, floor in FLOORS.items()
        }
        s["budget_pass"] = all(s[key] <= s["budget"][key] for key in FLOORS)
        s["fine_control_ratio"] = {
            key: (s[key] / fine[key] if fine[key] > 0 else None) for key in FLOORS
        }
        s["pass"] = s["absolute_pass"] and s["budget_pass"]
        report["refined"][r["explosion"]] = s
        if report["status"] == "PASS" and not s["pass"]:
            report["status"] = "FAIL"
    return report