from rsf.proj import *

# constant velocity-gradient model (horizontal gradient)
v0 = 2.
a = 0.5

Flow('model',None,
     '''
     math
     n1=401  d1=0.01 o1=0.  label1=Depth unit1=km
     n2=401  d2=0.01 o2=-2. label2=Position unit2=km
     output="%g+%g*x2"
     ''' % (v0,a))
Plot('model',
     '''
     grey color=j scalebar=y allpos=y title=Model screenratio=0.87
     labelsz=10 titlesz=12 barlabel=Velocity barunit=km/s barreverse=y
     titlefat=6 labelfat=6
     ''')

# analytical traveltime
Flow('analt','model',
     '''
     math output="1./%g*acosh((%g*%g*(x1*x1+x2*x2))/(2.*%g*(%g*x2+%g))+1.)"
     ''' % (a,a,a,v0,a,v0))

Plot('analt',
     '''
     grey color=j scalebar=y allpos=y title="Analytical Traveltime" screenratio=0.87
     labelsz=10 titlesz=12 barlabel=Traveltime barunit=s barreverse=y
     titlefat=6 labelfat=6
     ''')

# analytical derivative (relative coordinate)
Flow('atdl','model',
     '''
     math output="-(%g*x2+2.*%g)*sqrt((%g*%g*(x1*x1+x2*x2))/(%g*%g*x1*x1+(%g*x2+2.*%g)*(%g*x2+2.*%g)))/(%g*(%g*x2+%g))"
     ''' % (a,v0,a,a,a,a,a,v0,a,v0,v0,a,v0))
Plot('atdl',
     '''
     grey color=j scalebar=y title="Analytical dT/dx" screenratio=0.87
     labelsz=10 titlesz=12 barlabel=Derivative barunit=s/km
     titlefat=6 labelfat=6
     ''')

# analytical derivative (source)
# atdx has a singularity at source
Flow('atdx','model',
     '''
     math output="%g*(%g*x2*x2-%g*x1*x1+2.*%g*x2)/(%g*x2+%g)/sqrt(%g*%g*%g*%g*(x1*x1+x2*x2)*(x1*x1+x2*x2)+4.*%g*%g*%g*(%g*x2+%g)*(x1*x1+x2*x2))" |
     cut n1=1 n2=1 f2=200
     ''' % (a,a,a,v0,a,v0,a,a,a,a,v0,a,a,a,v0))

Flow('atds','atdl atdx','add scale=1,-1 ${SOURCES[1]}')
Plot('atds',
     '''
     grey color=j scalebar=y title="Analytical dT/ds" screenratio=0.87
     labelsz=10 titlesz=12 barlabel=Derivative barunit=s/km barreverse=y
     titlefat=6 labelfat=6
     ''')

Result('model','model atds','SideBySideIso')

# computed (source sampling will affect FD result)
nshot=5

Flow('yshot',None,'math n1=%d o1=-2. d1=%g output=x1' % (nshot,4./(nshot-1)))
Flow('szero','yshot','math output=0.')
Flow('shots','szero yshot','cat axis=2 ${SOURCES[1]} ${SOURCES[0]} | transp')

Flow('time ctdl ctds','model shots',
     '''
     put n3=1 d3=0.01 o3=0. label3= unit3= |
     eikods shotfile=${SOURCES[1]} tdl1=${TARGETS[1]} tds1=${TARGETS[2]} |
     put o4=-2. d4=%g
     ''' % (4./(nshot-1)))

Flow('time1','time','window n4=1 f4=%d' % ((nshot-1)/2-1))
Flow('time2','time','window n4=1 f4=%d' % ((nshot-1)/2+1))
Flow('fdtds','time1 time2',
     'add scale=-1,1 ${SOURCES[1]} | scale dscale=%g' % ((nshot-1)/8.))

# compare tds error
Flow('cdiff','ctds atds',
     '''
     window n4=1 f4=%d | add scale=1,-1 ${SOURCES[1]}
     ''' % ((nshot-1)/2))
Plot('cdiff',
     '''
     grey color=j scalebar=y title="Error Eikonal-based" screenratio=0.87
     labelsz=10 titlesz=12 barlabel=Derivative barunit=s/km barreverse=y
     maxval=0.5 minval=-0.5 clip=0.2
     titlefat=6 labelfat=6
     ''')

Flow('fddiff','fdtds atds',
     '''
     window n3=1 f3=%d | add scale=1,-1 ${SOURCES[1]}
     ''' % ((nshot-1)/2))
Plot('fddiff',
     '''
     grey color=j scalebar=y title="Error finite-difference" screenratio=0.87
     labelsz=10 titlesz=12 barlabel=Derivative barunit=s/km barreverse=y
     maxval=0.5 minval=-0.5 clip=0.2
     titlefat=6 labelfat=6
     ''')

Result('diff','cdiff fddiff','SideBySideIso')

# compare tdl error
Plot('ldiff','ctdl atdl',
     '''
     window n4=1 f4=%d | add scale=1,-1 ${SOURCES[1]} | 
     grey color=j scalebar=y title="Error Eikonal-based" screenratio=0.87
     labelsz=10 titlesz=12 barlabel=Derivative barunit=s/km barreverse=y
     titlefat=6 labelfat=6
     ''' % ((nshot-1)/2))

# compare tdy error
Flow('atdy','atds atdl','add scale=-1,1 ${SOURCES[1]}')
Flow('ctdy','ctds ctdl','add scale=-1,1 ${SOURCES[1]}')

Plot('ydiff','ctdy atdy',
     '''
     window n4=1 f4=%d | add scale=1,-1 ${SOURCES[1]} | 
     grey color=j scalebar=y title="Error Eikonal-based" screenratio=0.87
     labelsz=10 titlesz=12 barlabel=Derivative barunit=s/km barreverse=y
     titlefat=6 labelfat=6
     ''' % ((nshot-1)/2))

# Hermite interpolation
Flow('hermite','time ctds','tinterp deriv=${SOURCES[1]} ns=1 ds=0.25 os=0.25')

# linear interpolation
Flow('linear','time','tinterp type=linear ns=1 ds=0.25 os=0.25')

# partial interpolation
Flow('partial','time','tinterp type=partial ns=1 ds=0.25 os=0.25')

# compare error
Flow('ytemp',None,'math n1=1 o1=0.25 d1=0.25 output=x1')
Flow('stemp','ytemp','math output=0.')
Flow('tshot','stemp ytemp','cat axis=2 ${SOURCES[1]} ${SOURCES[0]} | transp')

Flow('eiko','model tshot',
     '''
     put n3=1 d3=0.01 o3=0. label3= unit3= |
     eikonal shotfile=${SOURCES[1]} |
     put d4=0.25 o4=0.25
     ''')
Plot('eiko',
     '''
     grey color=j scalebar=y title="Exact Traveltime" screenratio=0.87
     labelsz=10 titlesz=12 barlabel=Traveltime barunit=s
     allpos=y barreverse=y titlefat=6 labelfat=6
     ''')

Flow('herror','hermite eiko','add scale=1,-1 ${SOURCES[1]}')
Plot('herror',
     '''
     grey color=j scalebar=y title="Error Hermite" screenratio=0.87
     labelsz=10 titlesz=12 barlabel=Traveltime barunit=s
     maxval=0.1 minval=-0.01 clip=0.1
     titlefat=6 labelfat=6
     ''')

Flow('lerror','linear eiko','add scale=1,-1 ${SOURCES[1]}')
Plot('lerror',
     '''
     grey color=j scalebar=y title="Error linear" screenratio=0.87
     labelsz=10 titlesz=12 barlabel=Traveltime barunit=s
     maxval=0.1 minval=-0.01 clip=0.1
     titlefat=6 labelfat=6
     ''')

Flow('perror','partial eiko','add scale=1,-1 ${SOURCES[1]}')
Plot('perror',
     '''
     grey color=j scalebar=y title="Error shift" screenratio=0.87
     labelsz=10 titlesz=12 barlabel=Traveltime barunit=s
     maxval=0.1 minval=-0.01 clip=0.1
     titlefat=6 labelfat=6
     ''')

Result('ierror','eiko herror lerror perror','TwoRows')

# source sampling refinement
zpos=150
xpos=150

Flow('atds1','atds','window n1=1 f1=%d n2=1 f2=%d' % (zpos,xpos))

sfddiff0 = []
sfddiff1 = []

for n in (5,9,17,33):
    sfddiff0.append('sfddiff0_'+str(n))
    sfddiff1.append('sfddiff1_'+str(n))

    Flow('yshot_'+str(n),None,'math n1=%d o1=-2. d1=%g output=x1' % (n,4./(n-1)))
    Flow('szero_'+str(n),'yshot_'+str(n),'math output=0.')
    Flow('shots_'+str(n),['szero_'+str(n),'yshot_'+str(n)],'cat axis=2 ${SOURCES[1]} ${SOURCES[0]} | transp')

    Flow(['stime_'+str(n),'sctdl_'+str(n),'sctds_'+str(n)],['model','shots_'+str(n)],
         '''
         put n3=1 d3=0.01 o3=0. label3= unit3= |
         eikods shotfile=${SOURCES[1]} tdl1=${TARGETS[1]} tds1=${TARGETS[2]} |
         put o4=-2. d4=%g
         ''' % (4./(n-1)))

    Flow('sfdtds1_'+str(n),'stime_'+str(n),'window | window n3=1 f3=%d' % ((n-1)/2-1))
    Flow('sfdtds2_'+str(n),'stime_'+str(n),'window | window n3=1 f3=%d' % ((n-1)/2))
    Flow('sfdtds3_'+str(n),'stime_'+str(n),'window | window n3=1 f3=%d' % ((n-1)/2+1))

    Flow('sfdtds_0_'+str(n),['sfdtds1_'+str(n),'sfdtds2_'+str(n)],
         '''
         add scale=-1,1 ${SOURCES[1]} | scale dscale=%g
         ''' % ((n-1.)/4.))
    Flow('sfdtds_1_'+str(n),['sfdtds1_'+str(n),'sfdtds3_'+str(n)],
         '''
         add scale=-1,1 ${SOURCES[1]} | scale dscale=%g
         ''' % ((n-1.)/8.))

    Flow('sfddiff0_'+str(n),['sfdtds_0_'+str(n),'atds1'],
         '''
         window n1=1 f1=%d n2=1 f2=%d |
         add scale=1,-1 ${SOURCES[1]}
         ''' % (zpos,xpos))
    Flow('sfddiff1_'+str(n),['sfdtds_1_'+str(n),'atds1'],
         '''
         window n1=1 f1=%d n2=1 f2=%d |
         add scale=1,-1 ${SOURCES[1]}
         ''' % (zpos,xpos))

Flow('sfddiff0',sfddiff0,
     '''
     cat axis=1 ${SOURCES[1:%d]} |
     math output="log(abs(input))"
     ''' % len(sfddiff0))
Flow('sfddiff1',sfddiff1,
     '''
     cat axis=1 ${SOURCES[1:%d]} |
     math output="log(abs(input))"
     ''' % len(sfddiff1))

Flow('saxis',None,'math n1=4 d1=1. o1=0. output="4.*2.^x1+1."')

Plot('sfddiff0','saxis sfddiff0',
     '''
     cmplx ${SOURCES[1]} | graph plotcol=6 plotfat=5
     label1="Number of Sources" unit1= label2="Log(absolute error)" unit2=
     title="Source Sampling Refinement"
     screenratio=0.5 screenht=7 labelsz=4 titlesz=5.5
     min1=4 max1=34 min2=-7 max2=-1. labelsz=6 titlesz=8
     titlefat=4 labelfat=4
     ''')
Plot('sfddiff1','saxis sfddiff1',
     '''
     cmplx ${SOURCES[1]} | graph plotcol=4 plotfat=5
     wantaxis=n wanttitle=n dash=2
     screenratio=0.5 screenht=7
     min1=4 max1=34 min2=-7 max2=-1.
     ''')

Flow('scdiff1','cdiff',
     '''
     window n1=1 f1=%d n2=1 f2=%d | math output="log(abs(input))" |
     spray axis=1 n=4 d=0.01 o=2.
     ''' % (zpos,xpos))

Plot('scdiff1','saxis scdiff1',
     '''
     cmplx ${SOURCES[1]} | graph plotcol=2 plotfat=5
     wantaxis=n wanttitle=n dash=3
     screenratio=0.5 screenht=7
     min1=4 max1=34 min2=-7 max2=-1.
     ''')

Result('sfddiff','sfddiff0 sfddiff1 scdiff1','Overlay')

# grid sampling refinement
gfddiff0 = []
gfddiff1 = []
gcdiff1 = []

for g in (201,401,801,1601):
    gfddiff0.append('gfddiff0_'+str(g))
    gfddiff1.append('gfddiff1_'+str(g))
    gcdiff1.append('gcdiff1_'+str(g))

    Flow('model'+str(g),None,
         '''
         math
         n1=%d  d1=%g o1=0.  label1=x1 unit1=km
         n2=%d  d2=%g o2=-2. label2=x2 unit2=km
         output="%g+%g*x2"
         ''' % (g,4./(g-1),g,4./(g-1),v0,a))

    Flow(['gtime_'+str(g),'gctdl_'+str(g),'gctds_'+str(g)],['model'+str(g),'shots'],
         '''
         put n3=1 d3=%g o3=0. label3= unit3= |
         eikods shotfile=${SOURCES[1]} tdl1=${TARGETS[1]} tds1=${TARGETS[2]} |
         put o4=-2. d4=1.
         ''' % (4./(g-1)))

    Flow('gfdtds1_'+str(g),'gtime_'+str(g),'window | window n3=1 f3=1')
    Flow('gfdtds2_'+str(g),'gtime_'+str(g),'window | window n3=1 f3=2')
    Flow('gfdtds3_'+str(g),'gtime_'+str(g),'window | window n3=1 f3=3')

    Flow('gfdtds_0_'+str(g),['gfdtds1_'+str(g),'gfdtds2_'+str(g)],
         '''
         add scale=-1,1 ${SOURCES[1]} | scale dscale=1.
         ''')
    Flow('gfdtds_1_'+str(g),['gfdtds1_'+str(g),'gfdtds3_'+str(g)],
         '''
         add scale=-1,1 ${SOURCES[1]} | scale dscale=0.5
         ''')

    Flow('gfddiff0_'+str(g),['gfdtds_0_'+str(g),'atds1'],
         '''
         window n1=1 f1=%d n2=1 f2=%d | put d1=0.01 d2=0.01 |
         add scale=1,-1 ${SOURCES[1]}
         ''' % (zpos*(g-1)/400,xpos*(g-1)/400))
    Flow('gfddiff1_'+str(g),['gfdtds_1_'+str(g),'atds1'],
         '''
         window n1=1 f1=%d n2=1 f2=%d | put d1=0.01 d2=0.01 |
         add scale=1,-1 ${SOURCES[1]}
         ''' % (zpos*(g-1)/400,xpos*(g-1)/400))
    Flow('gcdiff1_'+str(g),['gctds_'+str(g),'atds1'],
         '''
         window n1=1 f1=%d n2=1 f2=%d n4=1 f4=2 | put d1=0.01 d2=0.01 |
         add scale=1,-1 ${SOURCES[1]}
         ''' % (zpos*(g-1)/400,xpos*(g-1)/400))

Flow('gfddiff0',gfddiff0,
     '''
     cat axis=1 ${SOURCES[1:%d]} |
     math output="log(abs(input))"
     ''' % len(gfddiff0))
Flow('gfddiff1',gfddiff1,
     '''
     cat axis=1 ${SOURCES[1:%d]} |
     math output="log(abs(input))"
     ''' % len(gfddiff1))
Flow('gcdiff1',gcdiff1,
     '''
     cat axis=1 ${SOURCES[1:%d]} |
     math output="log(abs(input))"
     ''' % len(gcdiff1))

Flow('gaxis',None,'math n1=4 d1=1. o1=1. output="100.*2.^x1+1."')

Plot('gfddiff0','gaxis gfddiff0',
     '''
     cmplx ${SOURCES[1]} | graph plotcol=6 plotfat=5
     label1="Number of Gird-points" unit1= label2="Log(absolute error)" unit2=
     title="Grid Spacing Refinement"
     screenratio=0.5 screenht=7 labelsz=4 titlesz=5.5
     min1=151 max1=1651 min2=-7 max2=-1. labelsz=6 titlesz=8
     titlefat=4 labelfat=4
     ''')
Plot('gfddiff1','gaxis gfddiff1',
     '''
     cmplx ${SOURCES[1]} | graph plotcol=4 plotfat=5
     wantaxis=n wanttitle=n dash=2
     screenratio=0.5 screenht=7
     min1=151 max1=1651 min2=-7 max2=-1.
     ''')
Plot('gcdiff1','gaxis gcdiff1',
     '''
     cmplx ${SOURCES[1]} | graph plotcol=2 plotfat=5
     wantaxis=n wanttitle=n dash=3
     screenratio=0.5 screenht=7
     min1=151 max1=1651 min2=-7 max2=-1.
     ''')

Result('gfddiff','gfddiff0 gfddiff1 gcdiff1','Overlay')

End()