from rsf.proj import *
import math

s02 = 1.
qx  = 0.052

min2 = 0.5
max2 = 6.5

# velocity model
Flow('slow2',None,
     '''
     math n1=101 n2=361 d1=0.02 d2=0.02 output="%g-2*%g*x2"
     label1=Depth unit1=km label2=Position unit2=km 
     label=Slowness-squared unit="s\^2\_/km\^2\_"
     ''' % (s02,qx))

Flow('vel','slow2',
     'math output="1./sqrt(input)" | put label=Velocity unit=km/s')
Plot('slow2',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y allpos=y title=Model pclip=100
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     minval=0.12 maxval=0.95 clip=0.95 screenratio=0.65 screenht=9
     ''' % (min2,max2))

# analytical sigma in (z,x)
Flow('depsigma','slow2',
     '''
     math output="sqrt(((%g-2*%g*x2)-sqrt((%g-2*%g*x2)^2-4*%g*%g*x1*x1))/(2*%g*%g))" |
     put label=Sigma unit=
     ''' % (s02,qx,s02,qx,qx,qx,qx,qx))

# analytical x0
Flow('ax0','depsigma',
     '''
     math output="x2+%g*input*input/2" |
     put label=Position unit=km
     ''' % qx)

Plot('ax0',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y allpos=y barreverse=y
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     title="Analytical x\_0\^" screenratio=0.65 screenht=9
     ''' % (min2,max2))
Plot('cax0','ax0',
     '''
     window min2=%g max2=%g |
     contour nc=100 scalebar=y plotcol=7 plotfat=7
     wantaxis=n wanttitle=n screenratio=0.65 screenht=9
     ''' % (min2,max2))
Plot('x0','ax0 cax0','Overlay')

# analytical t0
Flow('at0','ax0 depsigma',
     '''
     math sigma=${SOURCES[1]}
     output="(%g-2*%g*input)*sigma+%g*%g*sigma*sigma*sigma/3" |
     put label=Time unit=s
     ''' % (s02,qx,qx,qx))

Plot('at0',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y allpos=y barreverse=y
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     title="Analytical t\_0\^" screenratio=0.65 screenht=9
     ''' % (min2,max2))
Plot('cat0','at0',
     '''
     window min2=%g max2=%g |
     contour nc=40 scalebar=y plotcol=7 plotfat=7
     wantaxis=n wanttitle=n screenratio=0.65 screenht=9
     ''' % (min2,max2))
Plot('t0','at0 cat0','Overlay')

Result('hs2analy','x0 t0','OverUnderAniso')

# analytical geometrical spreading
Flow('ageo','slow2',
     '''
     math output="2*((%g-2*%g*x2)^2-4*%g*%g*x1*x1)/(%g-2*%g*x2)/(%g-2*%g*x2+sqrt((%g-2*%g*x2)^2-4*%g*%g*x1*x1))"
     ''' % (s02,qx,qx,qx,s02,qx,s02,qx,s02,qx,qx,qx))

Plot('ageo',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y allpos=y bias=0.66
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     barlabel="Q\^2\_" barunit=
     title="Geometrical Spreading" screenratio=0.65 screenht=9
     ''' % (min2,max2))

# analytical sigma in (t0,x0)
Flow('timsigma',None,
     '''
     math n1=626 n2=361 d1=0.004 d2=0.02
     output="((3*x1+sqrt(9*x1*x1+4*(%g-2*%g*x2)^3/(%g*%g)))/(2*%g*%g))^(1./3)-(%g-2*%g*x2)/%g*(2/(3*%g*x1+sqrt(9*%g*%g*x1*x1+4*(%g-2*%g*x2)^3)))^(1./3)"
     label1=Time unit1=s label2=Position unit2=km 
     label=Sigma unit= |
     cut n1=1
     ''' % (s02,qx,qx,qx,qx,qx,s02,qx,qx,qx,qx,qx,s02,qx))

# analytical Dix
Flow('dix','timsigma',
     '''
     math output="sqrt(%g-2*%g*x2)/(%g-2*%g*x2-%g*%g*input*input)" |
     put label=Velocity unit=km/s
     ''' % (s02,qx,s02,qx,qx,qx))

Plot('dix',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y allpos=y bias=1
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     title="Dix Velocity" pclip=100
     screenratio=0.65 screenht=9
     ''' % (min2,max2))

Result('hs2grad','ageo dix','OverUnderAniso')

# convert Dix to depth
Flow('init','dix',
     '''
     time2depth velocity=$SOURCE intime=y twoway=n nz=101 dz=0.02 |
     put label1=Depth unit1=km
     ''')

# mask
Flow('mask','ax0',
     '''
     math output=-1 | cut n2=25 | cut n2=35 f2=326 | dd type=int
     ''')

# inversion
Flow('t2d tt2d xt2d ft2d gt2d ct2d','init dix mask',
     '''
     tdconvert niter=3 verb=n cgiter=2000 eps=4 shape=y rect1=4 rect2=10 dix=${SOURCES[1]} mask=${SOURCES[2]}
     t0=${TARGETS[1]} x0=${TARGETS[2]} f0=${TARGETS[3]} grad=${TARGETS[4]} cost=${TARGETS[5]}
     ''')

# cost
Plot('cost0','ct2d',
     '''
     window min2=%g max2=%g | window n3=1 f3=0 | 
     grey color=j scalebar=y barreverse=y barlabel=Cost barunit=
     title="Initial f" minval=-0.45 maxval=0.63 clip=0.63
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     screenratio=0.65 screenht=9
     ''' % (min2,max2))
Plot('cost3','ct2d',
     '''
     window min2=%g max2=%g | window n3=1 f3=3 | 
     grey color=j scalebar=y barreverse=y barlabel=Cost barunit=
     title="Final f" minval=-0.45 maxval=0.63 clip=0.63
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     screenratio=0.65 screenht=9
     ''' % (min2,max2))

Result('hs2cost','cost0 cost3','OverUnderAniso')

# error
Plot('error0','vel init',
     '''
     add scale=-1,1 ${SOURCES[1]} | window min2=%g max2=%g |
     grey color=j scalebar=y barreverse=y
     title="Initial Model Error" minval=-0.004 maxval=0.176 clip=0.176
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     screenratio=0.65 screenht=9
     ''' % (min2,max2))
Plot('error3','vel t2d',
     '''
     add scale=-1,1 ${SOURCES[1]} | window min2=%g max2=%g |
     grey color=j scalebar=y barreverse=y
     title="Final Model Error" minval=-0.004 maxval=0.176 clip=0.176
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     screenratio=0.65 screenht=9
     ''' % (min2,max2))

Result('hs2error','error0 error3','OverUnderAniso')

End()