test: retain independent Jet CV continuation measurements

tests/validation/jet_cv_measurements.py

@@ -0,0 +1,285 @@
+"""Measurement-only continuation diagnostics. Never changes populations or commands."""
+
+import hashlib
+import json
+from pathlib import Path
+import numpy as np
+
+OFFSETS = np.array([1.5, 2.5, 3.5])
+LINEAR = np.array([2.5, -1.5, 0.0])
+QUADRATIC = np.array([4.375, -5.25, 1.875])
+
+
+def sample_steps(c):
+    """Same reference times at both scales; dense startup retained."""
+    s2 = c["scale"] ** 2
+    reference = sorted(set(range(0, 101)) | set(range(100, 3201, 4)))
+    return np.array(reference, int) * s2
+
+
+def factors(scale):
+    """Lref=one s=1 root length, Tref=one s=1 root time, rho_ref=1."""
+    return dict(
+        length=1 / scale,
+        time=1 / scale**2,
+        velocity=scale,
+        mass=1 / scale**2,
+        momentum=1 / scale,
+        impulse=1 / scale,
+        force=scale,
+        volume_rate=1.0,
+        mass_rate=1.0,
+        pressure=scale**2,
+    )
+
+
+def primitives(g, cx, cy):
+    f = np.asarray(g, np.float64)
+    rho = f.sum(axis=0)
+    mom = np.stack([np.einsum("q,qyx->yx", v, f) for v in (cx, cy)])
+    return rho, mom / rho
+
+
+def contour_terms(rho, vel, origin, rectangle, nu, h=0.5):
+    """Independent split of total outward momentum flux into advective,p,viscous."""
+    x0, y0, x1, y1 = np.rint((np.array(rectangle) - np.tile(origin, 2)) / h).astype(int)
+    duy, dux = np.gradient(vel[0], h)
+    dvy, dvx = np.gradient(vel[1], h)
+    adv = rho * vel[:, None] * vel[None, :]
+    pressure = np.zeros_like(adv)
+    pressure[0, 0] = pressure[1, 1] = (rho - 1) / 3
+    viscous = -rho * nu * np.array([[2 * dux, duy + dvx], [duy + dvx, 2 * dvy]])
+    out = []
+    for field in (adv, pressure, viscous):
+        flux = np.zeros(2)
+        for i, sign in ((x0, -1), (x1, 1)):
+            flux += (
+                sign
+                * 0.5
+                * (field[:, 0, y0:y1, i - 1] + field[:, 0, y0:y1, i]).sum(axis=1)
+                * h
+            )
+        for j, sign in ((y0, -1), (y1, 1)):
+            flux += (
+                sign
+                * 0.5
+                * (field[:, 1, j - 1, x0:x1] + field[:, 1, j, x0:x1]).sum(axis=1)
+                * h
+            )
+        out.append(flux)
+    return np.array(out)
+
+
+def ray_geometry(fluid, origin, center, radius, h=0.5):
+    z, w = np.polynomial.legendre.leggauss(32)
+    theta = np.array([np.pi / 2, 3 * np.pi / 2])[:, None] + np.pi / 12 * z
+    normal = np.stack((np.cos(theta), np.sin(theta)), axis=-1)
+    positions = np.asarray(center) + normal[:, :, None, :] * (
+        radius + h * OFFSETS[None, None, :, None]
+    )
+    local = (positions - origin) / h - 0.5
+    ix = np.floor(local[..., 0]).astype(int)
+    iy = np.floor(local[..., 1]).astype(int)
+    a = local[..., 0] - ix
+    b = local[..., 1] - iy
+    if not (
+        np.all(ix >= 0)
+        and np.all(iy >= 0)
+        and np.all(ix + 1 < fluid.shape[1])
+        and np.all(iy + 1 < fluid.shape[0])
+    ):
+        raise ValueError("Ray interpolation leaves allocated grid")
+    for dy, dx in ((0, 0), (0, 1), (1, 0), (1, 1)):
+        if not fluid[iy + dy, ix + dx].all():
+            raise ValueError("Ray stencil must be active fluid")
+    return dict(
+        ix=ix,
+        iy=iy,
+        a=a,
+        b=b,
+        normal=normal,
+        ds=w * radius * np.pi / 12,
+        positions=positions,
+        origin=np.asarray(origin),
+        h=h,
+    )
+
+
+def reconstruct_aperture(rho, vel, geo):
+    field = np.concatenate((rho[None], vel), axis=0)
+    x, y, a, b = geo["ix"], geo["iy"], geo["a"], geo["b"]
+    ray = (
+        field[:, y, x] * (1 - a) * (1 - b)
+        + field[:, y, x + 1] * a * (1 - b)
+        + field[:, y + 1, x] * (1 - a) * b
+        + field[:, y + 1, x + 1] * a * b
+    )
+    estimates = []
+    for coeff in (LINEAR, QUADRATIC):
+        wall = np.einsum("cajk,k->caj", ray, coeff)
+        un = np.einsum("caj,ajc->aj", wall[1:], geo["normal"])
+        q = np.sum(un * geo["ds"], axis=1)
+        mass = np.sum(wall[0] * un * geo["ds"], axis=1)
+        estimates.append(np.stack((q, mass), axis=-1))
+    return ray, np.array(estimates)
+
+
+class MeasurementRecorder:
+    """Retain reconstructible per-sample terms and sparse full primitive snapshots."""
+
+    def __init__(self, c, output, fluid_u, leaves, origin_a, cx, cy):
+        self.c = c
+        self.output = Path(output)
+        self.fluid_u = fluid_u
+        self.leaves = leaves
+        self.origin_a = np.asarray(origin_a)
+        self.cx = cx
+        self.cy = cy
+        self.records = []
+        self.rays = [
+            ray_geometry(fluid_u, (0.0, 0.0), np.array(c["center"]), c["diameter"] / 2),
+            ray_geometry(leaves[1], origin_a, np.array(c["center"]), c["diameter"] / 2),
+        ]
+        self.snapshots = []
+
+    def capture(self, step, ug, levels, cv_data, nearest_flux, velocity_error):
+        c = self.c
+        s = c["scale"]
+        record = dict(
+            step=step,
+            t_reference=step / s**2,
+            t_star=step * c["u"] / c["diameter"],
+            velocity_error=velocity_error,
+        )
+        arrays = [ug, *levels]
+        masks = [self.fluid_u, *self.leaves]
+        hs = [0.5, 1.0, 0.5]
+        primitive = [primitives(g, self.cx, self.cy) for g in arrays]
+        stats = []
+        for (rho, vel), mask, h in zip(primitive, masks, hs):
+            if (
+                not np.isfinite(rho[mask]).all()
+                or rho[mask].min() <= 0
+                or not np.isfinite(vel[:, mask]).all()
+            ):
+                raise RuntimeError(
+                    "Invalid density/velocity on an active composite leaf"
+                )
+            stats.append(
+                [
+                    rho[mask].min(),
+                    rho[mask].max(),
+                    np.max(abs(rho[mask] - 1)),
+                    np.linalg.norm(vel[:, mask], axis=0).max() * s,
+                    rho[mask].sum() * h * h / s**2,
+                ]
+            )
+        record["level_stats"] = stats  # uniform, AMR root leaf, AMR fine leaf
+        record["global_mass_reference"] = [stats[0][-1], stats[1][-1] + stats[2][-1]]
+        record["cv_inventory"] = []
+        record["cv_flux"] = []
+        record["momentum_flux_terms"] = []
+        record["arc_ray_values"] = []
+        record["arc_estimates"] = []
+        record["nearest_mass_flux"] = [
+            nearest_flux[k] for k in ("uniform_adapter", "amr_api")
+        ]
+        for idx, (name, p, origin) in enumerate(
+            [
+                ("uniform_adapter", primitive[0], (0.0, 0.0)),
+                ("amr_api", primitive[2], self.origin_a),
+            ]
+        ):
+            rho, vel = p
+            record["cv_inventory"].append([x[0] for x in cv_data[name]])
+            record["cv_flux"].append([x[1] for x in cv_data[name]])
+            terms = np.array(
+                [contour_terms(rho, vel, origin, r, c["nu"]) for r in c["contours"]]
+            )
+            np.testing.assert_allclose(
+                terms.sum(axis=1),
+                np.array(record["cv_flux"][-1])[:, 1:],
+                atol=2e-12,
+                rtol=2e-10,
+            )
+            record["momentum_flux_terms"].append(terms)
+            ray, estimates = reconstruct_aperture(rho, vel, self.rays[idx])
+            record["arc_ray_values"].append(ray)
+            record["arc_estimates"].append(estimates)
+        self.records.append(record)
+        if step // s**2 in (0, 100, 800, 1600, 3200) and step % s**2 == 0:
+            fields = {}
+            for name, (rho, vel), mask, h, origin in zip(
+                ("uniform", "amr_root", "amr_fine"),
+                primitive,
+                masks,
+                hs,
+                ((0.0, 0.0), (0.0, 0.0), self.origin_a),
+            ):
+                fields[name + "_rho"] = np.where(mask, rho, np.nan).astype(np.float32)
+                fields[name + "_velocity_reference"] = np.where(
+                    mask[None], vel * s, np.nan
+                ).astype(np.float32)
+                fields[name + "_origin_reference"] = np.asarray(origin) / s
+                fields[name + "_h_reference"] = np.array(h / s)
+            file = self.output / f"fields_s{s}_tref{step//s**2}.npz"
+            np.savez_compressed(file, **fields)
+            self.snapshots.append(
+                dict(
+                    file=file.name,
+                    sha256=hashlib.sha256(file.read_bytes()).hexdigest(),
+                    bytes=file.stat().st_size,
+                    t_reference=step / s**2,
+                )
+            )
+
+    def flush(self):
+        path = self.output / f"samples_s{self.c['scale']}.npz"
+        arrays = {k: np.array([r[k] for r in self.records]) for k in self.records[0]}
+        for i, g in enumerate(self.rays):
+            for k in ("normal", "ds", "positions"):
+                arrays[f"ray_{i}_{k}"] = g[k]
+        arrays["ray_offsets_fine_cells"] = OFFSETS
+        arrays["linear_coefficients"] = LINEAR
+        arrays["quadratic_coefficients"] = QUADRATIC
+        np.savez_compressed(path, **arrays)
+        metadata = dict(
+            schema=1,
+            case=self.c,
+            reference_conversion=factors(self.c["scale"]),
+            sample_file=path.name,
+            sample_sha256=hashlib.sha256(path.read_bytes()).hexdigest(),
+            sample_count=len(self.records),
+            snapshots=self.snapshots,
+            level_stats_columns=[
+                "rho_min",
+                "rho_max",
+                "max_abs_rho_minus_1",
+                "max_speed_reference",
+                "mass_reference",
+            ],
+            raw_partition_order=["wall", "aperture", "total", "blowing", "suction"],
+            raw_columns=["Jx", "Jy", "angular_impulse", "mass_into_fluid"],
+            backend_order=["uniform_adapter", "amr_api"],
+            level_order=["uniform", "amr_root_leaf", "amr_fine_leaf"],
+            momentum_term_order=["advection", "pressure", "negative_viscous_traction"],
+            arc_estimate_axes=[
+                "sample",
+                "backend",
+                "linear_or_quadratic",
+                "blowing_or_suction",
+                "Q_or_M",
+            ],
+            arc_ray_axes=[
+                "sample",
+                "backend",
+                "rho_ux_uy",
+                "aperture",
+                "gauss_point",
+                "normal_offset",
+            ],
+            warning="Extrapolation and nearest-node values are estimator-dependent, not exact wall values. Fields/ray arrays use root-lattice units except explicitly named reference arrays.",
+        )
+        (self.output / f"samples_s{self.c['scale']}.json").write_text(
+            json.dumps(metadata, indent=2) + "\n"
+        )