"""CPU stencil and actual closed-interface packet budgets; no CUDA claims.""" import numpy as np import unittest from pathlib import Path 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 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) 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 def check_multiple_rectangles_rejected(): cfg=LBMConfig(nx=24,ny=16,viscosity=.1) p=FixedPatchSet2D.from_config(cfg, [(4,4,8,12),(12,4,16,12)]) layout=BlockLayout2D.build(16,24,4,4,periodic=(True,True)) with unittest.TestCase().assertRaisesRegex(ValueError, 'one rectangle'): projected_explosion_tables(p,layout) 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) def test_reject(self): check_multiple_rectangles_rejected() 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) 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)