"""CPU stencil and actual closed-interface packet budgets; no CUDA claims.""" import numpy as np import unittest from pathlib import Path from dataclasses import replace import shutil import subprocess import tempfile from CelerisLab.config import LBMConfig from CelerisLab.lbm.blocks import BlockLayout2D from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY from CelerisLab.lbm.refinement import FixedPatch2D, FixedPatchSet2D, ACTIVE, OUTER_GHOST from CelerisLab.lbm.refinement_tables import projected_explosion_tables from CelerisLab.lbm.hierarchy import build_hierarchy from tests.unit.test_hierarchy import CFG, FOUR E = np.array([D2Q9_EX, D2Q9_EY]).T W = np.array([4/9] + [1/9]*4 + [1/36]*4) def setup(union=False, bounds=(4, 4, 12, 12)): cfg = LBMConfig(nx=16, ny=16, viscosity=.1) p = (FixedPatchSet2D.from_config(cfg, [bounds]) if union else FixedPatch2D.from_config(cfg, bounds)) layout = BlockLayout2D.build(16, 16, 4, 4, periodic=(True, True)) neighbors, offsets = projected_explosion_tables(p, layout, np.where(p.covered, OUTER_GHOST, ACTIVE)) inverse = np.full(layout.nblocks * np.prod(layout.local_shape), -1, dtype=int) inverse[layout.owner.ravel()] = np.arange(256) return p, inverse[neighbors], offsets def reconstruct(c, p, neighbors, offsets, project=True): # Independent NumPy evaluation of the documented stencil; neighbors have # been decoded from actual blocked owner IDs, not recomputed geometry. values = c.reshape(9, -1)[:, neighbors] g = .5 * (values[:, :, 1::2] - values[:, :, ::2]) if project: g[1:] -= (np.einsum('qra,qa->qr', g[1:], E[1:]) / np.sum(E[1:]**2, axis=1)[:, None])[:, :, None] * E[1:, None, :] out = c.reshape(9, -1)[:, p.ghost_coarse] + np.einsum('qra,ra->qr', g, offsets) return out, g def check_legal_footprint_mean_and_characteristic(union, bounds): p, neighbors, offsets = setup(union, bounds) assert set(neighbors.ravel()) <= set(p.ghost_coarse) assert not p.covered.ravel()[neighbors].any() assert (neighbors >= 0).all() c = np.random.default_rng(72).normal(size=(9,16,16)) c[:, p.covered] = np.nan # poison every forbidden donor out, g = reconstruct(c, p, neighbors, offsets) assert np.isfinite(out).all() np.testing.assert_allclose(np.einsum('qra,qa->qr', g, E), 0, atol=5e-16) for d in np.unique(p.ghost_coarse): rows = p.ghost_coarse == d assert rows.sum() == 4 np.testing.assert_allclose(out[:, rows].mean(axis=1), c.reshape(9,-1)[:,d], atol=5e-16) for q in range(1,9): delta = offsets[rows, None, :] - offsets[None, rows, :] # Parallel displacements have zero 2D cross product. parallel = np.isclose(delta[...,0]*E[q,1]-delta[...,1]*E[q,0], 0) diff = out[q, rows][:,None]-out[q, rows][None,:] np.testing.assert_allclose(diff[parallel], 0, atol=5e-16) const = np.broadcast_to(W[:,None,None], (9,16,16)).copy() constant, _ = reconstruct(const,p,neighbors,offsets) np.testing.assert_array_equal(constant, np.broadcast_to(W[:,None], constant.shape)) def collide(f, tau): rho = f.sum(axis=0) u = np.einsum('qa,q...->a...', E, f) / rho cu = np.einsum('qa,a...->q...', E, u) eq = W[:,None] * rho * (1 + 3*cu + 4.5*cu**2 - 1.5*np.sum(u*u,axis=0)) return f + (eq-f)/tau def totals(c, f, p): q = c[:,~p.covered].sum(axis=1)+.25*f[:,p.fine_kind==ACTIVE].sum(axis=1) return np.r_[q.sum(), E.T@q] def check_full_two_hop_closed_corner_budget(union, collision, bounds): p, neighbors, offsets = setup(union,bounds) rng=np.random.default_rng(17) c=W[:,None,None]*(1+rng.normal(0,.002,(9,16,16))) f=W[:,None,None]*(1+rng.normal(0,.002,(9,*p.fine_shape))) start=totals(c,f,p) # Count exactly the eight coarse diagonal links whose source is also # uncovered, but whose midpoint lies on the closed rectangular interface. q, k=np.divmod(p.coalesce_coarse.astype(np.int64),256) y,x=np.divmod(k,16) grazing=~p.covered[y-E[q,1],x-E[q,0]] assert grazing.sum()==8 c[:,p.covered]=np.nan for _ in range(8): f.reshape(9,-1)[:,p.ghost_fine]=reconstruct(c,p,neighbors,offsets)[0] for _ in range(2): f=np.array([np.roll(f[q],(ey,ex),(0,1)) for q,(ex,ey) in enumerate(E)]) f[:,p.fine_kind==OUTER_GHOST]=np.nan if collision: f[:,p.fine_kind==ACTIVE]=collide(f[:,p.fine_kind==ACTIVE],p.tau_fine) out=np.array([np.roll(c[q],(ey,ex),(0,1)) for q,(ex,ey) in enumerate(E)]) packets=f.reshape(-1)[p.coalesce_fine].mean(axis=1) assert np.isfinite(packets[grazing]).all() out.reshape(-1)[p.coalesce_coarse]=packets if collision:out[:,~p.covered]=collide(out[:,~p.covered],p.tau_coarse) out[:,p.covered]=np.nan c=out np.testing.assert_allclose(totals(c,f,p),start,rtol=0,atol=2e-12) def check_projection_is_distinguished_by_characteristic_not_global_mass(): p,n,o=setup() c=np.random.default_rng(8).normal(size=(9,16,16)) _,g=reconstruct(c,p,n,o,project=False) assert np.max(abs(np.einsum('qra,qa->qr',g,E))) > .1 class ProjectedExplosionTests(unittest.TestCase): def test_stencil(self): for union in (False, True): for bounds in ((4,4,12,12),(2,3,9,10),(2,2,3,3)): with self.subTest(union=union,bounds=bounds): check_legal_footprint_mean_and_characteristic(union,bounds) def test_closed_interface(self): for union in (False,True): for collision in (False,True): for bounds in ((4,4,12,12),(2,3,14,12),(2,2,3,3)): with self.subTest(union=union,collision=collision,bounds=bounds): check_full_two_hop_closed_corner_budget(union,collision,bounds) def test_projection(self): check_projection_is_distinguished_by_characteristic_not_global_mass() def test_affine_and_shift_invariance(self): p,n,o=setup() y,x=np.indices((16,16)) c=np.array([.7+q*.01+x*.03+y*.05 for q in range(9)]) out,g=reconstruct(c,p,n,o) available=n[:,::2]!=n[:,1::2] gradient=np.array([.03,.05])*available expected=np.broadcast_to(gradient,(9,*gradient.shape)).copy() expected[1:]-=(np.einsum('qra,qa->qr',expected[1:],E[1:])/ np.sum(E[1:]**2,axis=1)[:,None])[:,:,None]*E[1:,None,:] np.testing.assert_allclose(g,expected,atol=2e-16) shifted,_=reconstruct(c-W[:,None,None],p,n,o) np.testing.assert_allclose(out-W[:,None],shifted,atol=5e-16) class ProjectedKernelTests(unittest.TestCase): @unittest.skipUnless(shutil.which("g++"), "g++ is required for actual-kernel CPU arithmetic") def test_actual_kernel_fp32_shifted_and_unshifted(self): p, neighbors, offsets = setup(True) parent = BlockLayout2D.build(16,16,4,4,periodic=(True,True)) child = p.fine_layout((8,4)) # deliberately unequal parent/child strides ids, offsets = projected_explosion_tables(p,parent,np.where(p.covered,OUTER_GHOST,ACTIVE)) cells = int(np.prod(parent.local_shape)) targets = child.owner.ravel()[p.ghost_fine] donors = parent.owner.ravel()[p.ghost_coarse] source = (Path(__file__).resolve().parents[2] / "src/CelerisLab/lbm/kernels/blocks.cu").read_text() kernel = source.split('extern "C" __global__ void BlockExplodeProjected(',1)[1] kernel = 'void BlockExplodeProjected(' + kernel.split('// Targets are receiver-local SOURCE slots',1)[0] def array(name, values, kind): v = np.asarray(values).ravel() text = ','.join(str(int(x)) if kind=='uint64_t' else float(x).hex()+'f' for x in v) return f"{kind} {name}[] = {{{text}}};\n" for shifted in (False,True): c=np.random.default_rng(11).normal(0,.001,(9,16,16)) if not shifted:c+=W[:,None,None] c=c.astype(np.float32) packed=np.full(parent.nblocks*9*cells,np.nan,np.float32) for q in range(9):packed[parent.slot(parent.owner.ravel(),q)]=c[q].ravel() # All halo/covered inputs are poisoned; kernel must use legal cores. for q in range(9):packed[parent.slot(parent.owner.ravel()[p.covered.ravel()],q)]=np.nan # Hex initializer needs a portable expression for poisoned NaNs. data=array('coarse',packed,'float').replace('nanf','NAN') code = '#include \n#include \n#include \nusing fpxx=float;\n' code += f'constexpr unsigned NQ=9, TOTAL_CELLS={int(np.prod(child.local_shape))};\n' code += 'struct Id {unsigned x;}; Id blockIdx{0},blockDim{1},threadIdx{0};\n' code += 'unsigned index_f(unsigned k,unsigned q){return q*TOTAL_CELLS+k;}\n' code += 'int d_cx[]={0,1,-1,0,0,1,-1,1,-1},d_cy[]={0,0,0,1,-1,1,-1,-1,1};\n'+kernel code += data+array('target',targets,'uint64_t')+array('donor',donors,'uint64_t') code += array('neighbor',ids,'uint64_t')+array('offset',offsets,'float') code += f'float fine[{child.nblocks*9*int(np.prod(child.local_shape))}]={{}};\n' code += f'int main(){{for(unsigned r=0;r<{len(targets)};++r){{threadIdx.x=r;BlockExplodeProjected(coarse,fine,target,donor,{len(targets)},{cells},neighbor,offset);}}' code += f'for(unsigned q=0;q<9;++q)for(unsigned r=0;r<{len(targets)};++r){{auto t=target[r];printf("%.9g\\n",fine[(t/TOTAL_CELLS)*9*TOTAL_CELLS+q*TOTAL_CELLS+t%TOTAL_CELLS]);}}}}' with tempfile.TemporaryDirectory() as folder: src=Path(folder)/'kernel.cpp'; binary=Path(folder)/'kernel' src.write_text(code) subprocess.run(['g++','-std=c++17','-O2',str(src),'-o',str(binary)],check=True,capture_output=True) result=np.fromstring(subprocess.check_output([str(binary)],text=True),sep=' ').reshape(9,-1) expected,_=reconstruct(c.astype(float),p,neighbors,offsets) np.testing.assert_allclose(result,expected,atol=6e-8,rtol=2e-7) UNIONS = [ [(8, 8, 24, 20)], [(8, 8, 16, 20), (16, 8, 24, 20)], # abutting [(8, 8, 18, 20), (14, 8, 24, 20)], # overlap [(8, 8, 16, 20), (17, 8, 24, 20)], # gap 1 [(8, 8, 16, 20), (18, 8, 24, 20)], # gap 2 [(8, 8, 16, 20), (19, 8, 24, 20)], # gap 3 [(8, 8, 20, 14), (8, 14, 14, 22)], # L concavity [(4, 4, 10, 10), (18, 16, 26, 26)], # separated [(4, 4, 5, 7), (4, 5, 7, 7)], # known concave exit/re-entry path ] def decode(layout, neighbors): inverse = np.full(layout.hole, -1, dtype=int) inverse[layout.owner[layout.owned]] = np.flatnonzero(layout.owned) return inverse[neighbors] def check_pair_stencil(pair, layout, kind): ids, offsets = projected_explosion_tables(pair, layout, kind) neighbors = decode(layout, ids) legal = set(map(int, pair.ghost_coarse)) assert len(np.unique(pair.ghost_fine)) == len(pair.ghost_fine) assert set(neighbors.ravel()) <= legal assert layout.owned.ravel()[neighbors].all() assert (kind.ravel()[neighbors] == ACTIVE).all() assert not pair.covered.ravel()[neighbors].any() ny, nx = pair.coarse_shape for row, donor in enumerate(pair.ghost_coarse): y, x = divmod(int(donor), nx) for axis, (dx, dy) in enumerate(((1, 0), (0, 1))): ends = [(x - dx, y - dy), (x + dx, y + dy)] expected = [] for xx, yy in ends: outside = not (0 <= xx < nx and 0 <= yy < ny) if outside and not layout.periodic[axis]: break cell = (yy % ny) * nx + xx % nx if cell not in legal: break expected.append(cell) if len(expected) != 2: expected = [int(donor)] * 2 np.testing.assert_array_equal(neighbors[row, 2*axis:2*axis+2], expected) fy, fx = divmod(int(pair.ghost_fine[row]), pair.fine_shape[1]) expected = [pair.fine_origin[0]+(fx+.5)/2-(x+.5), pair.fine_origin[1]+(fy+.5)/2-(y+.5)] np.testing.assert_array_equal(offsets[row], expected) assert set(offsets.ravel()) == {-.25, .25} assert ids.dtype == np.uint64 and ids.flags.c_contiguous assert offsets.dtype == np.float32 and offsets.flags.c_contiguous return neighbors, offsets class GenericProjectedPairs(unittest.TestCase): def test_single_window_and_singleton_union_exact_stencils(self): for bounds in ((4,4,12,12), (2,3,9,10), (2,2,3,3)): a, an, ao = setup(False, bounds) b, bn, bo = setup(True, bounds) np.testing.assert_array_equal(a.ghost_coarse, b.ghost_coarse) np.testing.assert_array_equal(an, bn) np.testing.assert_array_equal(ao, bo) # Coordinates differ between scalar and union arrays; physical rows don't. ay, ax = np.divmod(a.ghost_fine, a.fine_shape[1]) by, bx = np.divmod(b.ghost_fine, b.fine_shape[1]) np.testing.assert_array_equal(ax + 2*a.fine_origin[0], bx) np.testing.assert_array_equal(ay + 2*a.fine_origin[1], by) def test_union_decomposition_permutation_and_topologies(self): cfg = LBMConfig(nx=30, ny=30, viscosity=.1) layout = BlockLayout2D.build(30,30,6,5,periodic=(True,True)) stencils = [] for rectangles in UNIONS: pair = FixedPatchSet2D.from_config(cfg, rectangles) kind = np.where(pair.covered, OUTER_GHOST, ACTIVE) n, o = check_pair_stencil(pair, layout, kind) other = FixedPatchSet2D.from_config(cfg, rectangles[::-1]) on, oo = projected_explosion_tables(other, layout, kind) np.testing.assert_array_equal(n, decode(layout, on)) np.testing.assert_array_equal(o, oo) stencils.append((pair.ghost_fine, pair.ghost_coarse, n, o)) c = np.random.default_rng(92).normal(size=(9,30,30)) c[:, pair.covered] = np.nan values, gradients = reconstruct(c, pair, n, o) self.assertTrue(np.isfinite(values).all()) np.testing.assert_allclose(np.einsum('qra,qa->qr', gradients, E),0,atol=1e-15) # q0 keeps its full gradient; no division by zero or forced zero. expected_q0 = .5*(c[0].ravel()[n[:,1::2]]-c[0].ravel()[n[:,::2]]) np.testing.assert_array_equal(gradients[0], expected_q0) self.assertGreater(np.max(abs(gradients[0])), .1) for donor in np.unique(pair.ghost_coarse): rows = pair.ghost_coarse == donor self.assertEqual(rows.sum(),4) np.testing.assert_allclose(values[:,rows].mean(axis=1),c.reshape(9,-1)[:,donor],atol=1e-15) for q in range(1,9): delta = o[rows,None,:]-o[None,rows,:] parallel = delta[...,0]*E[q,1] == delta[...,1]*E[q,0] diff = values[q,rows][:,None]-values[q,rows][None,:] np.testing.assert_allclose(diff[parallel],0,atol=1e-15) for split in stencils[1:3]: for expected, actual in zip(stencils[0],split): np.testing.assert_array_equal(actual,expected) def test_nested_scalar_origins_sparse_holes_and_unequal_strides(self): cases = [(CFG, FOUR), (LBMConfig(nx=16,ny=16,viscosity=.02), [(4,4,12,12),(8,8,16,16),(8,8,16,16)])] for cfg, refs in cases: hierarchy = build_hierarchy(cfg,refs) for index,pair in enumerate(hierarchy.pairs): if index == 0: layout = BlockLayout2D.build(cfg.ny,cfg.nx,4,4,periodic=(True,True)) else: prior = hierarchy.pairs[index-1] layout = (prior.fine_layout((4,8)) if isinstance(prior,FixedPatchSet2D) else BlockLayout2D.build(*prior.fine_shape,4,8,periodic=(False,False))) kind = hierarchy.kinds[index] n,o = check_pair_stencil(pair,layout,kind) c = np.random.default_rng(index).normal(size=(9,*pair.coarse_shape)) c[:,~layout.owned | (kind != ACTIVE)] = np.nan self.assertTrue(np.isfinite(reconstruct(c,pair,n,o)[0]).all()) if index and isinstance(pair,FixedPatchSet2D): self.assertTrue((~layout.owned).any()) self.assertEqual(layout.periodic, (index == 0, index == 0)) def test_periodicity_is_per_axis_and_missing_pairs_repeat_center(self): cfg=LBMConfig(nx=16,ny=16,viscosity=.1) p=FixedPatch2D.from_config(cfg,(2,2,14,14)) kind=np.where(p.covered,OUTER_GHOST,ACTIVE) results=[] for periodic in ((False,False),(True,False),(False,True),(True,True)): layout=BlockLayout2D.build(16,16,4,4,periodic=periodic) n,_=check_pair_stencil(p,layout,kind) results.append(n) y,x=np.divmod(p.ghost_coarse,16) for axis,coordinate in enumerate((x,y)): edge=(coordinate==0)|(coordinate==15) if not periodic[axis]: np.testing.assert_array_equal(n[edge,2*axis:2*axis+2], np.repeat(p.ghost_coarse[edge,None],2,axis=1)) self.assertTrue((results[0]!=results[1]).any()) self.assertTrue((results[0]!=results[2]).any()) self.assertTrue((results[1]!=results[3]).any()) def test_invalid_parent_donors_reject_before_owner_lookup(self): p,_,_=setup() layout=BlockLayout2D.build(16,16,4,4) kind=np.where(p.covered,OUTER_GHOST,ACTIVE) for nonactive in (1,2,255): bad=kind.copy(); bad.ravel()[p.ghost_coarse[0]]=nonactive with self.assertRaisesRegex(ValueError,'ACTIVE'): projected_explosion_tables(p,layout,bad) with self.assertRaisesRegex(ValueError,'shape'): projected_explosion_tables(p,layout,kind[:1]) hole_layout=BlockLayout2D.build(16,16,4,4,origins=[(0,0)]) with self.assertRaisesRegex(ValueError,'owned'): projected_explosion_tables(p,hole_layout,kind) def test_generic_two_hop_mass_and_momentum_budget(self): cfg=LBMConfig(nx=30,ny=30,viscosity=.1) layout=BlockLayout2D.build(30,30,6,5) grazing_counts=set() for rectangles in UNIONS: pair=FixedPatchSet2D.from_config(cfg,rectangles) ids,o=projected_explosion_tables(pair,layout,np.where(pair.covered,OUTER_GHOST,ACTIVE)) n=decode(layout,ids) q,k=np.divmod(pair.coalesce_coarse.astype(np.int64),900) y,x=np.divmod(k,30) grazing=~pair.covered[(y-E[q,1])%30,(x-E[q,0])%30] grazing_counts.add(int(grazing.sum())) self.assertTrue(grazing.any()) for collision in (False,True): rng=np.random.default_rng(17) c=W[:,None,None]*(1+rng.normal(0,.002,(9,30,30))) f=W[:,None,None]*(1+rng.normal(0,.002,(9,*pair.fine_shape))) start=totals(c,f,pair) c[:,pair.covered]=np.nan for _ in range(4): f.reshape(9,-1)[:,pair.ghost_fine]=reconstruct(c,pair,n,o)[0] for _ in range(2): f=np.array([np.roll(f[q],(ey,ex),(0,1)) for q,(ex,ey) in enumerate(E)]) f[:,pair.fine_kind==OUTER_GHOST]=np.nan if collision: f[:,pair.fine_kind==ACTIVE]=collide(f[:,pair.fine_kind==ACTIVE],pair.tau_fine) out=np.array([np.roll(c[q],(ey,ex),(0,1)) for q,(ex,ey) in enumerate(E)]) packets=f.reshape(-1)[pair.coalesce_fine].mean(axis=1) self.assertTrue(np.isfinite(packets[grazing]).all()) out.reshape(-1)[pair.coalesce_coarse]=packets if collision: out[:,~pair.covered]=collide(out[:,~pair.covered],pair.tau_coarse) out[:,pair.covered]=np.nan c=out np.testing.assert_allclose(totals(c,f,pair),start,rtol=0,atol=4e-12) self.assertGreater(len(grazing_counts),1) # not a hard-coded eight-corner oracle