from rsf.proj import *

# definitions
a=1.
b=15.
amp=3.5

def DSR(name,zgrid,xgrid,extra=''):
    dz=0.5/(zgrid-1)
    dx=1./(xgrid-1)

    Flow(name,None,
         '''
         math n1=%d n2=%d d1=%g d2=%g o1=0. o2=0. output="%g+%g*x1" |
         math output="input+%g*exp(-((x1-0.2)*(x1-0.2)+(x2-0.5)*(x2-0.5))/(0.1*0.1))" |
         put label1=z unit1=km label2=x unit2=km |
         dsreiko %s | window n3=1 | window n1=1 |
         spline o1=0. d1=0.001 n1=1001
         ''' % (zgrid,xgrid,dz,dx,a,b,amp,extra))

    Flow(name+'_ex',None,
         '''
         math n1=%d n2=%d d1=%g d2=%g o1=0. o2=0. output="%g+%g*x1" |
         math output="input+%g*exp(-((x1-0.2)*(x1-0.2)+(x2-0.5)*(x2-0.5))/(0.1*0.1))" |
         put label1=z unit1=km label2=x unit2=km |
         dsreiko0 %s | window n3=1 | window n1=1 |
         spline o1=0. d1=0.001 n1=1001
         ''' % (zgrid,xgrid,dz,dx,a,b,amp,extra))

def Curve(input,extra=''):
    Plot(input,
         '''
         graph plotfat=8 label1= unit1= label2= unit2= wantaxis=n
         max2=0.42 dash=1 screenratio=0.4 screenht=10 labelsz=10 titlesz=12 labelfat=6.5 titlefat=6.5
         %s
         ''' % extra)

    Plot(input+'_ex',
         '''
         graph plotfat=8 label1= unit1= label2= unit2= wantaxis=n
         max2=0.42 dash=1 screenratio=0.4 screenht=10 labelsz=10 titlesz=12 labelfat=6.5 titlefat=6.5
         %s
         ''' % extra)

Flow('modl',None,
     '''
     math n1=501 n2=1001 d1=0.001 d2=0.001 o1=0. o2=0. output="%g+%g*x1" |
     math output="input+%g*exp(-((x1-0.2)*(x1-0.2)+(x2-0.5)*(x2-0.5))/(0.1*0.1))" |
     put label1=z unit1=km label2=x unit2=km
     ''' % (a,b,amp))
Plot('modl',
     '''
     grey color=j allpos=y scalebar=y barreverse=y barunit=km/s barlabel=Velocity pclip=100
     title="Anomaly Model" screenratio=0.4 screenht=5 labelsz=6 titlesz=8 labelfat=3 titlefat=3
     ''')

Plot('rays','modl',
     '''
     rays2 yshot=0 nr=21 a0=150 amax=172.25 nt=1000 dt=0.001 |
     graph yreverse=y transp=y wantaxis=n wanttitle=n scalebar=y 
     min1=0 max1=0.5 min2=0 max2=1 plotcol=7 plotfat=3
     screenratio=0.4 screenht=5
     ''')

Flow('truth','modl','eikonal zshot=0. yshot=0. xshot=0. | window n1=1')

Plot('truth',
     '''
     graph plotfat=8 label1= unit1= label2= unit2= wantaxis=n
     max2=0.42 dash=1 screenratio=0.4 screenht=10 labelsz=10 titlesz=12 labelfat=6.5 titlefat=6.5
     wantaxis=y label1=x unit1=km label2=Traveltime unit2=s dash=0 wanttitle=n
     ''')

Result('modl','modl rays','Overlay')

# zx refine
###########

# scale 0.05 (red)
DSR('zx1',11,21)
Curve('zx1','plotcol=2 wanttitle=n')

# scale 0.01 (green)
DSR('zx2',51,101)
Curve('zx2','plotcol=4 wanttitle=n')

# scale 0.005 (white)
DSR('zx3',101,201)
Curve('zx3','plotcol=7 title="Grid Refinement in z and x"')

# compare
Plot('zx','truth zx1 zx2 zx3','Overlay')
Plot('zx_ex','truth zx1_ex zx2_ex zx3_ex','Overlay')

# z refine
##########

# scale 0.05 (red)
DSR('z1',11,101)
Curve('z1','plotcol=2 wanttitle=n')

# scale 0.01 (green)
DSR('z2',51,101)
Curve('z2','plotcol=4 wanttitle=n')

# scale 0.005 (white)
DSR('z3',101,101)
Curve('z3','plotcol=7 title="Grid Refinement in z"')

# compare
Plot('z','truth z1 z2 z3','Overlay')
Plot('z_ex','truth z1_ex z2_ex z3_ex','Overlay')

# x refine
##########

# scale 0.05 (red)
DSR('x1',51,21)
Curve('x1','plotcol=2 wanttitle=n')

# scale 0.01 (green)
DSR('x2',51,101)
Curve('x2','plotcol=4 wanttitle=n')

# scale 0.005 (white)
DSR('x3',51,201)
Curve('x3','plotcol=7 title="Grid Refinement in x"')

# compare
Plot('x','truth x1 x2 x3','Overlay')
Plot('x_ex','truth x1_ex x2_ex x3_ex','Overlay')

Result('imp','z x','OverUnderIso')
Result('exp','z_ex x_ex','OverUnderIso')

End()