CelerisLab/tests/unit/test_shifted_storage.py

"""CPU host contracts; mocked CUDA executes real lifecycle, not numerical kernels."""
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_missing_tiles_rejected_before_compilation(self):
        module=backend.host_module()
        cfg=replace(backend.CFG,ddf_shifting=True)
        for blocks in ([None,None],[(4,4),None],[None,(4,4)]):
            with self.assertRaisesRegex(ValueError,'must be a tile'):
                module.FixedPatchStepper(cfg,context=None,refinements=[(4,4,20,28)],
                                         level_blocks=blocks)
        module.compiler.build_module.assert_not_called()

    def test_uniform_block_codec_and_raw_snapshot(self):
        module = backend.host_module()  # single blocked level: FixedPatchStepper(refinements=[])
        cuda, compiler = module.cuda, module.compiler
        cuda.mem_alloc.side_effect=lambda size: Mock()
        gpu=Mock();gpu.get_global.return_value=(Mock(),20)
        compiler.build_module.return_value=gpu
        context=Mock(activate=nullcontext,sm_arch='sm_70')
        cfg=replace(backend.CFG,nx=16,ny=16,ddf_shifting=True)
        run=module.FixedPatchStepper(cfg,context=context,refinements=[],level_blocks=[(4,4)])
        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())
            (down,)=run.download()
            np.testing.assert_array_equal(down,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()