CelerisLab/tests/unit/test_shifted_storage.py

"""CPU host contracts; mocked CUDA executes real lifecycle, not numerical kernels."""
import ast
import types
from pathlib import Path
from contextlib import nullcontext
from dataclasses import replace
import unittest
from unittest.mock import Mock, patch
import numpy as np
from CelerisLab.lbm.population_storage import (_StoredPopulations, _WEIGHTS,
                                               _encode_physical, _decode_stored)
from CelerisLab.lbm.synchronized_remap import _remap_stored, remap_synchronized
from tests.unit.test_synchronized_remap import topology, A, B
from tests.unit import test_refinement_three_level_multipatch as backend


def shifted_fields(mesh):
    rng = np.random.default_rng(129)
    arrays = []
    for leaf in mesh.leaf_masks:
        a = rng.normal(0, 1e-7, (9, *leaf.shape)).astype(np.float32)
        a[:, ~leaf] = np.nan
        y, x = np.nonzero(leaf)
        a[0, y[:6], x[:6]] = np.array([0., -0., 2**-80, -2**-80, 2**-140, -2**-140], np.float32)
        arrays.append(a)
    return tuple(arrays)


def remap(old, new, arrays, shifted=True):
    return _remap_stored(old, new, _StoredPopulations(arrays, shifted), shifted=True,
                         base_count=7, level_steps=(7, 14, 28), phase='postcollision_g')


def ready():
    with patch.object(backend, 'CFG', replace(backend.CFG, ddf_shifting=True)):
        run, module, allocations = backend.mock_backend(children=A)
    storage = {}
    module.cuda.memcpy_htod.side_effect = lambda p, a: storage.__setitem__(id(p), np.array(a, copy=True))
    module.cuda.memcpy_dtoh.side_effect = lambda a, p: np.copyto(a, storage[id(p)])
    run._initialize_raw(_StoredPopulations(shifted_fields(topology(A)), True))
    run._step_count = 7
    return run, module, allocations, storage


class ShiftedStorage(unittest.TestCase):
    def test_codec_preserves_sub_fp32_physical_perturbation(self):
        physical = np.broadcast_to(_WEIGHTS, (9, 2, 2)).copy()
        physical[1] += 2**-35
        encoded = _encode_physical(physical, True)
        self.assertTrue((encoded[1] == np.float32(2**-35)).all())
        np.testing.assert_array_equal(_decode_stored(encoded, True), physical)
        self.assertEqual(_decode_stored(encoded, True).dtype, np.float64)
        np.testing.assert_array_equal(_encode_physical(_WEIGHTS, True), np.zeros((9,1,1), np.float32))
        self.assertEqual(_encode_physical(physical, False).dtype, np.float32)

    def test_actual_fill_loop_bitwise_copy_and_single_rounding_restriction(self):
        old, new = topology(A), topology(B)
        arrays = shifted_fields(old)
        saved = [a.tobytes() for a in arrays]
        out = remap(old, new, arrays)
        self.assertTrue(out.shifted)
        for a, b, om, nm, before in zip(arrays, out.g, old.leaf_masks, new.leaf_masks, saved):
            self.assertEqual(a[:, om & nm].tobytes(), b[:, om & nm].tobytes())
            self.assertTrue(np.isnan(b[:, ~nm]).all())
            self.assertEqual(a.tobytes(), before)
            self.assertFalse(np.shares_memory(a,b))
        y,x = np.nonzero(new.leaf_masks[2] & ~old.leaf_masks[2])
        self.assertEqual(out.g[2][:,y,x].tobytes(), arrays[1][:,y//2,x//2].tobytes())
        y,x = np.nonzero(new.leaf_masks[1] & ~old.leaf_masks[1])
        mean = sum(arrays[2][:,2*y+dy,2*x+dx].astype(np.float64)*.25
                   for dy in (0,1) for dx in (0,1))
        self.assertEqual(out.g[1][:,y,x].tobytes(), mean.astype(np.float32).tobytes())
        def integral(mesh, values):
            return sum(a[:,m].sum(axis=1,dtype=np.float64)/4**l
                       for l,(a,m) in enumerate(zip(values,mesh.leaf_masks)))
        cast_error = abs(mean.astype(np.float32).astype(np.float64)-mean).sum(axis=1)/4
        self.assertTrue(np.all(abs(integral(old, arrays)-integral(new,out.g)) <= cast_error+2e-18))

    def test_shifted_density_and_tag_validation(self):
        old,new = topology(A),topology(B)
        a = tuple(np.where(m[None], np.float32(-.001), np.float32(np.nan)).repeat(9,axis=0)
                  for m in old.leaf_masks)  # FP32 under numpy 1 and 2 promotion rules
        remap(old,new,a)
        with self.assertRaisesRegex(ValueError,'representation'): remap(old,new,a,False)
        with self.assertRaisesRegex(ValueError,'FP32'): remap(old,new,tuple(x.astype(np.float64) for x in a))
        a[0][:,old.leaf_masks[0]] = -1
        with self.assertRaisesRegex(ValueError,'positive density'): remap(old,new,a)
        rest = tuple(np.where(m[None],np.float32(0),np.float32(np.nan)).repeat(9,axis=0)
                     for m in old.leaf_masks)
        remap(old,new,rest)
        with self.assertRaisesRegex(ValueError,'positive density'):
            remap_synchronized(old,new,rest,base_count=7,level_steps=(7,14,28),phase='postcollision_g')

    def test_runtime_remap_bypasses_public_codec_and_preserves_both_buffers(self):
        run,module,allocations,storage = ready()
        clone = None
        try:
            before = run._download_raw()
            expected = remap(topology(A),topology(B),before.values)
            with patch.object(run,'download',side_effect=AssertionError('public decode used')):
                clone = run.remap_child(B)
            self.assertEqual((run.step_count,clone.step_count),(7,7))
            self.assertFalse(clone._exchange_valid)
            for level, want, got in zip(clone._levels,expected.g,clone._download_raw().values):
                self.assertEqual(want.tobytes(),got.tobytes())
                for pointer in (level.current,level.spare):
                    np.testing.assert_array_equal(storage[id(pointer)],level.layout.to_blocked(want))
            for a,b in zip(before.values,run._download_raw().values): self.assertEqual(a.tobytes(),b.tobytes())
            for a in clone.download(): self.assertEqual(a.dtype,np.float64)
        finally:
            if clone is not None: clone.close()
            run.close()

    def test_runtime_public_encode_and_raw_mismatch_fail_before_upload(self):
        run,module,_,_ = ready()
        try:
            physical = tuple(np.broadcast_to(_WEIGHTS, level.shape).copy() for level in run._levels)
            for a in physical: a[1] += 2**-35
            run.initialize(*physical)
            for a,mask in zip(run._download_raw().values,run._leaf):
                self.assertTrue((a[1,mask] == np.float32(2**-35)).all())
            self.assertEqual(run.step_count,0)
            module.cuda.memcpy_htod.reset_mock()
            with self.assertRaisesRegex(ValueError,'representation'):
                run._initialize_raw(_StoredPopulations(run._download_raw().values,False))
            module.cuda.memcpy_htod.assert_not_called()
            run._level_cfgs[0] = replace(run._level_cfgs[0],ddf_shifting=False)
            with self.assertRaisesRegex(ValueError,'mixed'): run._download_raw()
        finally: run.close()

    def test_dense_or_mixed_shifted_rejected_before_compilation(self):
        module=backend.host_module()
        cfg=replace(backend.CFG,ddf_shifting=True)
        for kwargs in ({},dict(base_block=(4,4)),dict(fine_block=(4,4))):
            with self.assertRaisesRegex(ValueError,'all-blocked|base_block requires fine_block'):
                module.FixedPatchStepper(cfg,(4,4,20,28),context=None,**kwargs)
        module.compiler.compile_kernel.assert_not_called()

    def test_uniform_block_codec_and_raw_snapshot(self):
        path = Path(backend.r.__file__).with_name('block_stepper.py')
        tree = ast.parse(path.read_text())
        tree.body = [n for n in tree.body if not (
            isinstance(n, ast.Import) and any(a.name == 'pycuda.driver' for a in n.names)
            or isinstance(n, ast.ImportFrom) and n.module == 'cuda')]
        module = types.ModuleType('CelerisLab.lbm._test_block_host')
        module.__package__ = 'CelerisLab.lbm'
        cuda, compiler = Mock(), Mock()
        module.__dict__.update(cuda=cuda,compiler=compiler)
        exec(compile(tree,str(path),'exec'),module.__dict__)
        cuda.mem_get_info.return_value=(10**9,10**9)
        cuda.mem_alloc.side_effect=lambda size: Mock()
        gpu=Mock();gpu.get_global.return_value=(Mock(),20)
        compiler.load_module.return_value=gpu
        context=Mock(activate=nullcontext,sm_arch='sm_70')
        layout=module.BlockLayout2D.build(16,16,4,4)
        cfg=replace(backend.CFG,nx=16,ny=16,ddf_shifting=True)
        run=module.PeriodicBlockStepper(cfg,layout,context=context)
        storage={}
        cuda.memcpy_htod.side_effect=lambda p,a: storage.__setitem__(id(p),np.array(a,copy=True))
        cuda.memcpy_dtoh.side_effect=lambda a,p: np.copyto(a,storage[id(p)])
        try:
            physical=np.broadcast_to(_WEIGHTS,(9,16,16)).copy();physical[1]+=2**-35
            run.initialize(physical)
            raw=run._download_raw()
            self.assertTrue(raw.shifted)
            self.assertEqual(len(raw.values),1)
            self.assertTrue((raw.values[0][1]==np.float32(2**-35)).all())
            np.testing.assert_array_equal(run.download(),physical)
        finally: run.close()

    def test_passive_shifted_body_raw_remap_preserves_support(self):
        from tests.unit.test_pinned_body_remap import mesh, body, A as rects, B as moved
        old,new=mesh(rects,(body(),)),mesh(moved,(body(),))
        result=remap(old,new,shifted_fields(old))
        self.assertTrue(result.shifted)
        self.assertTrue(np.isnan(result.g[2][:,~new.leaf_masks[2]]).all())
        changed=mesh(moved,(body(center=(12.1,14.)),))
        with self.assertRaisesRegex(ValueError,'boundary'):
            remap(old,changed,shifted_fields(old))

    def test_fp64_physical_roundtrip_cannot_replace_raw(self):
        tiny=np.full((9,1,1),np.float32(2**-80))
        self.assertFalse(np.array_equal(tiny,_encode_physical(_decode_stored(tiny,True),True)))


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