CelerisLab/tests/validation/multigrid/test_exchange_window.py

"""Optional observation-window exchange ledger: host contract (mocked driver) and GPU gate.

Run CPU and GPU classes in separate processes; GPU callers select the device.
GPU tolerance (conservative engineering tolerance, not a universal proof): the window ledger
sums h*dt*raw per link over fine substeps, while the reference sums per-base interval ledgers
over links and then over base steps. Both are FP64 reorderings of the same FP32-exact substep
contributions, so the scale is the sum of the ACTUAL absolute substep contributions (read from
last_raw after every fine hop), not of per-base totals that could hide substep cancellation:
|window - reference| <= 4 * m * eps * sum|substep terms| + 1e-300, m = links * substeps.
"""
from contextlib import ExitStack
import unittest
from unittest.mock import patch

import numpy as np
import pycuda.driver as cuda

from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
from test_refinement_blocked_base import mocked_backend
from test_refinement_blocked_fine import BOUNDS, circle, config

EPS = np.finfo(np.float64).eps


def build(m, window=True, raw=True):
    return FixedPatchStepper(config("periodic"), BOUNDS, context=m["context"], base_block=(3, 4),
                             fine_block=(16, 20), fine_objects=[circle()],
                             observe_raw_exchange=raw, observe_window_exchange=window)


def same_bits(test, a, b):
    """Identical NaN masks, no infinities, identical finite FP32 bits (incl. signed zero)."""
    a, b = np.asarray(a), np.asarray(b)
    test.assertEqual(a.dtype, b.dtype)
    test.assertFalse(np.isinf(a).any() or np.isinf(b).any())
    np.testing.assert_array_equal(np.isnan(a), np.isnan(b))
    finite = np.isfinite(a)
    np.testing.assert_array_equal(a[finite].view(np.uint32), b[finite].view(np.uint32))


def accumulations(m):
    return [c for c in m["calls"] if c[0] == "BlockAccumulateExchange"]


class ExchangeWindowCPU(unittest.TestCase):
    def test_constructor_validation(self):
        with mocked_backend() as m:
            with self.assertRaisesRegex(ValueError, "requires observe_raw_exchange"):
                build(m, window=True, raw=False)
            with self.assertRaisesRegex(ValueError, "must be bool"):
                build(m, window=1)

    def test_state_machine_and_launches(self):
        with mocked_backend() as m, ExitStack() as stack:
            stack.enter_context(patch.object(cuda, "memset_d8"))
            stack.enter_context(patch.object(cuda, "memset_d8_async"))
            stack.enter_context(patch.object(cuda, "memcpy_dtoh",
                                             side_effect=lambda out, src: out.fill(0)))
            backend = build(m)
            b = backend._level_boundaries[1]
            self.assertIsNotNone(b.window_integral)
            with self.assertRaises(RuntimeError):
                backend.start_exchange_window()  # not initialized
            arrays = [np.full(level.shape, 1 / 9, np.float32) for level in backend._levels]
            backend.initialize(*arrays)
            with self.assertRaisesRegex(RuntimeError, "no valid exchange window"):
                backend.read_exchange_window(0)
            backend.step(1)
            self.assertEqual(len(accumulations(m)), 2)  # closed window: interval ledger only
            backend.start_exchange_window()
            with self.assertRaisesRegex(RuntimeError, "at least one completed"):
                backend.read_exchange_window(0)
            m["calls"].clear()
            backend.step(2)
            calls = accumulations(m)
            self.assertEqual(len(calls), 8)  # two fine hops x two ledgers x two base steps
            self.assertEqual({id(c[1][1]) for c in calls},
                             {id(b.interval_integral), id(b.window_integral)})
            out = backend.read_exchange_window(0)
            self.assertEqual((out["time_start"], out["time_end"], out["sample_count"]), (1., 3, 4))
            self.assertEqual(backend.read_exchange_window(0)["time_end"], 3)  # non-destructive
            backend.initialize(*arrays)  # invalidates
            with self.assertRaisesRegex(RuntimeError, "no valid exchange window"):
                backend.read_exchange_window(0)
            backend.start_exchange_window()
            with patch.object(backend, "_hop", side_effect=RuntimeError("injected")):
                with self.assertRaises(RuntimeError):
                    backend.step(1)
            self.assertFalse(backend._window_valid)
            with self.assertRaises(RuntimeError):
                backend.read_exchange_window(0)  # failed status

    def test_disabled_default_has_no_window_cost(self):
        with mocked_backend() as m, ExitStack() as stack:
            stack.enter_context(patch.object(cuda, "memset_d8_async"))
            backend = build(m, window=False)
            self.assertIsNone(backend._level_boundaries[1].window_integral)
            backend.initialize(*[np.full(l.shape, 1 / 9, np.float32) for l in backend._levels])
            backend.step(1)
            self.assertEqual(len(accumulations(m)), 2)
            with self.assertRaisesRegex(RuntimeError, "disabled"):
                backend.start_exchange_window()


class ExchangeWindowGPU(unittest.TestCase):
    def test_window_matches_summed_base_reads_and_populations_unchanged(self):
        from test_raw_exchange_gpu import BOUNDS as RB, CFG, circle as rc, jet, state
        from CelerisLab.cuda.context import CudaContext
        with CudaContext(0) as ctx:
            kw = dict(context=ctx, fine_block=(16, 20), fine_objects=[rc()], jets=[jet()],
                      observe_raw_exchange=True)
            with FixedPatchStepper(CFG, RB, observe_window_exchange=True, **kw) as windowed, \
                 FixedPatchStepper(CFG, RB, **kw) as plain:
                arrays = [state(level.shape) for level in windowed._levels]
                windowed.initialize(*arrays)
                plain.initialize(*arrays)
                windowed.step(2)
                plain.step(2)
                windowed.start_exchange_window()
                b = windowed._level_boundaries[1]
                count = len(b.tables["cell"])
                factors = np.array([.25, .25, .125, .25])  # h*dt, h*h*dt, h*h for level-1 h=.5
                scale = np.zeros(4)
                substeps = 0
                kernel = windowed._levels[1].kernel

                def audited(src, dst, *args, **kw):
                    nonlocal scale, substeps
                    windowed._stream.synchronize()  # BFL and both accumulations are complete
                    raw = np.empty((count, 4), np.float32)
                    cuda.memcpy_dtoh(raw, b.last_raw)
                    scale = scale + (np.abs(raw.astype(np.float64)) * factors).sum(axis=0)
                    substeps += 1
                    kernel(src, dst, *args, **kw)
                windowed._levels[1].kernel = audited
                keys = ("impulse", "angular_impulse", "mass_into_fluid")
                summed = {part: [0., 0., 0., 0.] for part in ("total", "wall", "aperture")}
                steps = 0
                for speed in (0., .002, .002, -.001, 0., 0.):  # amplitude updates stay in-window
                    windowed.set_jet_amplitude(0, speed)
                    plain.set_jet_amplitude(0, speed)
                    windowed.step(1)
                    plain.step(1)
                    steps += 1
                    obs = windowed.read_raw_exchange(0)
                    for part in summed:
                        o = obs[part]
                        summed[part] = [summed[part][0] + o["impulse"][0], summed[part][1] + o["impulse"][1],
                                        summed[part][2] + o["angular_impulse"],
                                        summed[part][3] + o["mass_into_fluid"]]
                    for a, e in zip(windowed._download_raw().values, plain._download_raw().values):
                        same_bits(self, a, e)
                out = windowed.read_exchange_window(0)
                self.assertEqual((out["time_end"] - out["time_start"], out["sample_count"]),
                                 (steps, 2 * steps))
                self.assertEqual(substeps, 2 * steps)
                bound = 4 * (count * substeps) * EPS * scale + 1e-300
                for part in summed:
                    o = out[part]
                    got = np.array([o["impulse"][0], o["impulse"][1], o["angular_impulse"],
                                    o["mass_into_fluid"]])
                    diff = np.abs(got - np.array(summed[part]))
                    self.assertTrue((diff <= bound).all(), f"{part}: {diff} > {bound}")
                    np.testing.assert_array_equal(o["force"], o["impulse"] / steps)
                again = windowed.read_exchange_window(0)
                np.testing.assert_array_equal(again["total"]["impulse"], out["total"]["impulse"])
                windowed.initialize(*arrays)
                with self.assertRaises(RuntimeError):
                    windowed.read_exchange_window(0)


if __name__ == "__main__": unittest.main()