from rsf.proj import *


# First, demonstrate that GPU implementation generate same result as accurate 
# as the CPU output up to floating precision (1e-7)
Flow('layer1',None,'math output=1600 n1=80 n2=160 d1=5 d2=5')
Flow('layer2',None,'math output=1800 n1=80 n2=160 d1=5 d2=5')
Flow('vel','layer1 layer2','cat axis=1 ${SOURCES[1]}')
Result('vel',
        '''
        sfput label1=Depth unit1=m label2=Distance unit2=m | 
        grey allpos=y scalebar=y color=j minval=1500 maxval=2000
        ''')

Flow('shotsa checka','vel',
        '''
        sfgenshots check=${TARGETS[1]} csdgather=n fm=15 amp=1 dt=0.0015 ns=1 ng=160 nt=800
        sxbeg=80 szbeg=80 jsx=10 jsz=0 gxbeg=0 gzbeg=3 jgx=1 jgz=0 chk=y kt=100
        ''')

Flow('shotsb checkb','vel',
        '''
        sfmodeling2d check=${TARGETS[1]} csdgather=n fm=15 amp=1 dt=0.0015 ns=1 ng=160 nt=800
        sxbeg=80 szbeg=80 jsx=10 jsz=0 gxbeg=0 gzbeg=3 jgx=1 jgz=0 chk=y kt=100
        ''')
Plot('checka','grey color=g scalebar=y')
Plot('checkb','grey color=g scalebar=y')
Plot('errab','checka checkb','add ${SOURCES[1]} scale=1,-1|grey color=g scalebar=y')
Result('allab','checka checkb errab','SideBySideIso')

# Second, demonstrate GPU is much faster especially when the model size is large!!
# Be careful: The GPU block size is 16x16. To make the comparison fair, you need 
# the size of the test model to be multiple 16x16. We only care about the time cost here. 

for i in range(1,8):
        nz=160*i
        nx=160*i
        vv="v%d"%i
        shota="shota%d"%i
        timea="timea%d"%i
        shotb="shotb%d"%i
        timeb="timeb%d"%i

        Flow(vv,None,'math output=1600 n1=%d n2=%d d1=5 d2=5'%(nz,nx))
        Flow(shota+' '+timea,vv,
                '''
                sfgenshots time=${TARGETS[1]} csdgather=n fm=15 amp=1 dt=0.0015 ns=5 ng=%d nt=1000
                sxbeg=5 szbeg=2 jsx=10 jsz=0 gxbeg=0 gzbeg=3 jgx=1 jgz=0 
                '''%(nx))
        Flow(shotb+' '+timeb,vv,
                '''
                sfmodeling2d time=${TARGETS[1]} csdgather=n fm=15 amp=1 dt=0.0015 ns=5 ng=%d nt=1000
                sxbeg=5 szbeg=2 jsx=10 jsz=0 gxbeg=0 gzbeg=3 jgx=1 jgz=0
                '''%(nx))

Flow('timealla','timea1 timea2 timea3 timea4 timea5 timea6 timea7',
        '''cat axis=1 ${SOURCES[1:7]}|put d1=160 o1=160''')
Flow('timeallb','timeb1 timeb2 timeb3 timeb4 timeb5 timeb6 timeb7',
        '''cat axis=1 ${SOURCES[1:7]}|put d1=160 o1=160''')
Plot('timeall','timealla timeallb',
        '''
        cat axis=2 ${SOURCES[1]}|
        graph dash=0,1 grid=y plotfat=3 label2=Time unit2=s 
        label1="nx(Datasize: nx=nz)" unit1= title="Time cost for different datasize"
        ''')
Plot('s0',None,'box font=2 x0=8.5 y0=6 label="CPU" xt=0.000000 yt=0.000000')
Plot('s1',None,'box font=2 x0=8.5 y0=2.5 label="GPU" xt=0.000000 yt=0.000000')
Result('timecost','timeall s0 s1','Overlay')

End()