CelerisLab/tests/validation/multigrid/test_shifted_jets_gpu.py

"""Shifted velocity-Jet GPU gates. Run only on an authorized device.

PYTHONPATH=src python tests/validation/multigrid/test_shifted_jets_gpu.py -v
Uniform initial curved-body states rho=.99/1.01; amplitude 0,+.002,-.001,0;
four-level substep ledger; three-level live-Jet raw remap. No pressure actuator.
The census is an implementation-exchange ledger, not an independent physical
traction/nozzle-thrust model. Frozen inherited caps: raw rows rtol=2e-6/atol=2e-7; fluid ledger
rtol=2e-4/atol=8e-5. Observer on/off, interval accumulation and raw copies exact.
Report actual residuals too: passing old loose caps alone is not an accuracy gain.
"""
from contextlib import ExitStack
import unittest
import numpy as np
import pycuda.driver as cuda
from CelerisLab.config import LBMConfig
from CelerisLab.cuda.context import CudaContext
from CelerisLab.body.geometry.circle import CircleGeometry
from CelerisLab.body.objects import SimObject
from CelerisLab.body.jets import Jet, Tube, Parallel
from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY, SOLID
from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
from CelerisLab.lbm.population_storage import _StoredPopulations
from CelerisLab.lbm.synchronized_remap import RemapTopology, _remap_stored

CX,CY=np.asarray(D2Q9_EX),np.asarray(D2Q9_EY)
OPP=np.array([np.flatnonzero((CX==-x)&(CY==-y))[0] for x,y in zip(CX,CY)])
W=np.array([4/9,*([1/9]*4),*([1/36]*4)],np.float64)
SPEEDS=(0.,.002,-.001,0.)


def circle(center,radius):
    return SimObject(17,geometry=CircleGeometry(*center,radius),center=center,radius=radius)


def drive(center,radius):
    return Jet(Tube(center,0.,radius*1.5,radius*.32),Parallel(0.),amplitude=0.)


def initial(run,rho,u=.013,v=-.007):
    cu=3*(CX*u+CY*v)
    g=W*((rho-1)+rho*(cu+.5*cu*cu-1.5*(u*u+v*v)))
    return _StoredPopulations(tuple(np.broadcast_to(g[:,None,None],l.shape).astype(np.float32).copy()
                                    for l in run._levels),True)


def download(pointer,shape,dtype=np.float32):
    out=np.empty(shape,dtype);cuda.memcpy_dtoh(out,pointer);return out


def raw_totals(a,leaf):
    g=a[:,leaf].astype(np.float64)
    # Only differences at the same fixed leaf mask are used, so mass's unit
    # background cancels. Momentum never requires adding weights.
    return np.array([(g*CX[:,None]).sum(),(g*CY[:,None]).sum(),g.sum()])


def census(run,index,buffer):
    """Independent physical link/face census of actual staged receiver windows."""
    level=run._levels[index];layout=level.layout;b=level.boundary
    stored=download(buffer,layout.blocked_shape);dense=layout.to_dense(stored)
    leaf=run._leaf[index];ny,nx=leaf.shape
    solid=(b.host.flags.reshape(ny,nx)&SOLID)!=0
    y,x=np.nonzero(leaf);owner=layout.owner[y,x].astype(np.intp)
    cells=int(np.prod(layout.local_shape));block,cell=np.divmod(owner,cells)
    flat=stored.reshape(layout.nblocks,9,cells);external=np.zeros(3,np.float64)
    for q,(cx,cy) in enumerate(zip(CX,CY)):
        sy,sx=y-cy,x-cx;dy,dx=y+cy,x+cx
        incoming=flat[block,q,cell-cy*layout.local_shape[1]-cx].astype(np.float64)+W[q]
        outgoing=flat[block,q,cell].astype(np.float64)+W[q]
        entering=~leaf[sy,sx]&~solid[sy,sx];leaving=~leaf[dy,dx]&~solid[dy,dx]
        external+=np.array([cx,cy,1])*(incoming[entering].sum()-outgoing[leaving].sum())
    receiver,direction,_,_,rx,ry,_,_=b.host.curved_host
    owner=layout.owner.ravel()[receiver].astype(np.intp);block,cell=np.divmod(owner,cells)
    outgoing=flat[block,direction,cell].astype(np.float64)
    source=cell+CY[direction]*layout.local_shape[1]+CX[direction]
    reflected=flat[block,OPP[direction],source].astype(np.float64)
    physical_sum=outgoing+reflected+2*W[direction]
    jx,jy=CX[direction]*physical_sum,CY[direction]*physical_sum
    rows=np.column_stack((jx,jy,rx*jy-ry*jx,reflected-outgoing))
    return raw_totals(dense,leaf),external,rows


class Ledger:
    def __init__(self,test,run,index):
        self.test,self.run,self.index=test,run,index
        self.samples=[];self.max_row=0.;self.max_balance=0.
        level=run._levels[index];self.original=level.kernel;b=level.boundary
        def audited(src,dst,*args,**kwargs):
            run._stream.synchronize()
            before,external,independent=census(run,index,src)
            raw=download(b.last_raw,(len(b.tables['cell']),4))
            np.testing.assert_allclose(raw,independent,rtol=2e-6,atol=2e-7)
            self.max_row=max(self.max_row,float(abs(raw.astype(float)-independent).max()))
            self.original(src,dst,*args,**kwargs);run._stream.synchronize()
            after=level.layout.to_dense(download(dst,level.layout.blocked_shape))
            change=raw_totals(after,run._leaf[index])-before
            expected=external+np.array([-independent[:,0].sum(),-independent[:,1].sum(),independent[:,3].sum()])
            np.testing.assert_allclose(change,expected,rtol=2e-4,atol=8e-5)
            self.max_balance=max(self.max_balance,float(abs(change-expected).max()))
            self.samples.append(raw.copy())
        level.kernel=audited

    def step(self):
        run=self.run;self.samples.clear();run.step()
        b=run._level_boundaries[self.index];h=run._exchange_spacing[self.index]
        self.test.assertEqual(len(self.samples),2**self.index)
        expected=np.zeros_like(self.samples[0],dtype=np.float64)
        weights=np.array([h*h,h*h,h*h*h,h*h])
        for raw in self.samples:expected+=raw.astype(np.float64)*weights
        actual=download(b.interval_integral,expected.shape,np.float64)
        np.testing.assert_array_equal(actual,expected)
        obs=run.read_raw_exchange(0,level=self.index)
        for name in ('impulse','angular_impulse','mass_into_fluid'):
            np.testing.assert_allclose(obs['total'][name],np.asarray(obs['wall'][name])+obs['aperture'][name],rtol=0,atol=2e-13)
        self.test.assertEqual(obs['sample_count'],2**self.index)
        for a,leaf in zip(run._download_raw().values,run._leaf):
            self.test.assertTrue(np.isfinite(a[:,leaf]).all())
            self.test.assertTrue((1+a[:,leaf].sum(axis=0,dtype=np.float64)>0).all())
        return obs

    def report(self,label):
        print(f'{label}: max_abs_raw_row={self.max_row:.9g}, max_abs_fluid_ledger={self.max_balance:.9g}',flush=True)


class ShiftedJetsGPU(unittest.TestCase):
    def exact_raw(self,a,b):
        self.assertTrue(a.shifted and b.shifted)
        for x,y in zip(a.values,b.values):self.assertEqual(x.tobytes(),y.tobytes())

    def test_matching_whole_circle_parallel_velocity_preserves_uniform_state(self):
        cfg=LBMConfig(nx=32,ny=24,viscosity=.1,ddf_shifting=True)
        center=(16.,12.);radius=2.15;obj=circle(center,radius)
        # Rectangle contains every circle hit; prescribed surface velocity equals
        # the uniform fluid velocity. Geometry itself remains fixed.
        whole=Jet(Tube((center[0]-radius-1,center[1]),0.,2*radius+2,radius+1),
                  Parallel(0.),amplitude=.013)
        for rho in (.99,1.01):
            with self.subTest(rho=rho),CudaContext(0) as ctx:
                with FixedPatchStepper(cfg,context=ctx,refinements=[(8,6,24,18)],level_blocks=[(4,4)]*2,
                        level_objects=[(),(obj,)],level_jets=[(),(whole,)],observe_raw_exchange=True) as run:
                    run._initialize_raw(initial(run,rho,u=.013,v=0.))
                    expected=run._download_raw();ledger=Ledger(self,run,1)
                    self.assertTrue(run._level_boundaries[1].tables['aperture'].all())
                    for _ in range(4):ledger.step()
                    error=max(float(abs(a[:,leaf]-b[:,leaf]).max()) for a,b,leaf in
                              zip(run._download_raw().values,expected.values,run._leaf))
                    self.assertLessEqual(error,1e-6)  # inherited uniform-state cap
                    print(f'matched whole-circle rho={rho}: max_abs_raw_state={error:.9g}',flush=True)
                    ledger.report(f'matched whole-circle rho={rho}')

    def test_uniform_initial_curved_rho_both_signs_switch_and_observer(self):
        cfg=LBMConfig(nx=32,ny=24,viscosity=.1,ddf_shifting=True)
        obj=circle((16.,12.),2.15);j=drive((16.,12.),2.15)
        for rho in (.99,1.01):
            with self.subTest(rho=rho),CudaContext(0) as ctx,ExitStack() as owned:
                settings=dict(context=ctx,refinements=[(8,6,24,18)],level_blocks=[(4,4)]*2,
                              level_objects=[(),(obj,)],level_jets=[(),(j,)])
                measured=owned.enter_context(FixedPatchStepper(cfg,observe_raw_exchange=True,**settings))
                plain=owned.enter_context(FixedPatchStepper(cfg,**settings))
                measured._initialize_raw(initial(measured,rho));plain._initialize_raw(initial(plain,rho))
                ledger=Ledger(self,measured,1)
                for speed in SPEEDS:
                    before=measured._download_raw();count=measured.step_count
                    measured.set_jet_amplitude(0,speed);plain.set_jet_amplitude(0,speed)
                    self.exact_raw(before,measured._download_raw());self.assertEqual(measured.step_count,count)
                    ledger.step();plain.step()
                    self.exact_raw(measured._download_raw(),plain._download_raw())
                ledger.report(f'L1 rho={rho}')

    def test_four_level_interval_ledger(self):
        cfg=LBMConfig(nx=48,ny=48,viscosity=.02,ddf_shifting=True)
        obj=circle((16.,16.),.5);j=drive((16.,16.),.5)
        with CudaContext(0) as ctx,ExitStack() as owned:
            settings=dict(context=ctx,refinements=[[(6,6,26,26)],[(20,20,44,44)],[(56,56,72,72)]],
                          level_blocks=[(4,4)]*4,level_objects=[(),(),(),(obj,)],level_jets=[(),(),(),(j,)])
            measured=owned.enter_context(FixedPatchStepper(cfg,observe_raw_exchange=True,**settings))
            plain=owned.enter_context(FixedPatchStepper(cfg,**settings))
            measured._initialize_raw(initial(measured,1.01));plain._initialize_raw(initial(plain,1.01))
            ledger=Ledger(self,measured,3)
            for speed in SPEEDS:
                measured.set_jet_amplitude(0,speed,level=3);plain.set_jet_amplitude(0,speed,level=3)
                ledger.step();plain.step();self.exact_raw(measured._download_raw(),plain._download_raw())
            ledger.report('L3 rho=1.01')

    def test_live_jet_remap_then_both_ledgers_continue(self):
        cfg=LBMConfig(nx=48,ny=32,viscosity=.1,ddf_shifting=True)
        parents=((4,4,20,28),(24,4,44,28));a=((16,16,32,40),(56,16,72,40));b=(a[0],(60,16,76,40))
        obj=circle((12.,14.),.8);j=drive((12.,14.),.8)
        with CudaContext(0) as ctx,ExitStack() as owned:
            old=owned.enter_context(FixedPatchStepper(cfg,context=ctx,refinements=[parents,a],
                level_blocks=[(4,4)]*3,level_objects=[(),(),(obj,)],level_jets=[(),(),(j,)],
                observe_raw_exchange=True))
            old._initialize_raw(initial(old,.99));ledger=Ledger(self,old,2)
            for speed in (.002,-.001):old.set_jet_amplitude(0,speed,level=2);ledger.step()
            before=old._download_raw();count=old.step_count
            topology=lambda rects:RemapTopology.from_config(cfg,old.patch,rects,child_objects=(obj,),
                child_jets=old._level_jets[2],_shifted_blocked=True)
            expected=_remap_stored(topology(a),topology(b),before,shifted=True,base_count=count,
                level_steps=(count,2*count,4*count),phase='postcollision_g')
            new=owned.enter_context(old.remap_child(b,child_objects=(obj,)))
            self.exact_raw(before,old._download_raw())
            self.exact_raw(_StoredPopulations(expected.g,True),new._download_raw())
            self.assertEqual(new._level_jets[2][0].amplitude,-.001)
            np.testing.assert_array_equal(new._level_boundaries[2].tables['surface_velocity'],old._level_boundaries[2].tables['surface_velocity'])
            newledger=Ledger(self,new,2);ledger.step();newledger.step()
            self.assertEqual((old.step_count,new.step_count),(count+1,count+1))
            ledger.report('L2 source');newledger.report('L2 remapped')


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