"""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()