from rsf.proj import *

xmax=20.0
zmax=3.0

nz=500
dz=zmax/(nz-1)

# vexact
Flow('ve',None,
     '''
     spiral xmax=%g ymax=%g ny=%d |
     put label= unit=km/s
     ''' % (xmax,zmax,nz))

# vexact -> vdix (cameron)
Flow('vd x z','ve',
     '''
     ve2d nt=1001 dt=0.002 x=${TARGETS[1]} z=${TARGETS[2]} |
     put label2=Time unit2=s
     ''')

Flow('xz','x z','cmplx ${SOURCES[1]}')
Plot('xz',
     '''
     window j1=5 | transp |
     graph plotcol=7 plotfat=7 wanttitle=n scalebar=y
     yreverse=y wantaxis=n screenratio=0.34 screenht=10
     min1=4 max1=14 min2=0 max2=2
     ''')

# vdix -> vexact (cameron)
Flow('vf x00 t00','vd',
     '''
     cameron2d method=lax nz=%d dz=%g x0=${TARGETS[1]} t0=${TARGETS[2]} |
     put label2=Depth unit2=km
     ''' % (nz,dz))

# take part of domain
Flow('modl','ve',
     '''
     window min1=4. max1=14. max2=2. | transp |
     put o2=0. label1=Depth unit1=km label2=Position unit2=km
     ''')
Plot('pmodl','modl',
     '''
     grey color=j scalebar=y allpos=y barlabel=Velocity barunit=km/s
     title=Model screenratio=0.34 screenht=10 minval=2 maxval=3.4 bias=1.9
     labelsz=10 titlesz=12 titlefat=6 labelfat=6 barreverse=y clip=1.25
     ''')

Plot('modl','pmodl xz','Overlay')

Flow('vdix','vd',
     '''
     window min1=4. max1=14. max2=2. | transp |
     put o2=0. d1=0.001 label1=Time unit1=s label2=Position unit2=km
     ''')

# eikonal
Flow('eiko2 t02 x02 f02','modl','irays t0=${TARGETS[1]} x0=${TARGETS[2]} f0=${TARGETS[3]}')

# evaluate cost
Flow('cost','modl t02 x02 f02 vdix',
     '''
     t2diter s0=${SOURCES[0]} t0=${SOURCES[1]} x0=${SOURCES[2]} 
     f0=${SOURCES[3]} dix=${SOURCES[4]} what=cost
     ''')

# add Gaussian perturbation
Flow('pert','ve',
     '''
     math output="input+0.025*exp(-((x1-7.)*(x1-7.)+(x2-1.)*(x2-1.))/(0.2*0.2))" |
     window min1=4. max1=14. max2=2. | transp |
     put o2=0. label1=Depth unit1=km label2=Position unit2=km
     ''')
Plot('pert','ve',
     '''
     math output="0.025*exp(-((x1-7.)*(x1-7.)+(x2-1.)*(x2-1.))/(0.2*0.2))" |
     window min1=4. max1=14. max2=2. | transp |
     put o2=0. label1=Depth unit1=km label2=Position unit2=km |
     grey color=j scalebar=y title=Perturbation allpos=y pclip=100
     barlabel=Velocity barunit=km/s screenratio=0.34 screenht=10
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     ''')

# eikonal for perturbed model
Flow('peiko pt0 px0 pf0','pert','irays t0=${TARGETS[1]} x0=${TARGETS[2]} f0=${TARGETS[3]}')

# cost for perturbed model
Flow('pcost','pert pt0 px0 pf0 vdix',
     '''
     t2diter s0=${SOURCES[0]} t0=${SOURCES[1]} x0=${SOURCES[2]} 
     f0=${SOURCES[3]} dix=${SOURCES[4]} what=cost
     ''')

Flow('diffcost','cost pcost','add scale=1,-1 ${SOURCES[1]}')
Plot('diffcost',
     '''
     grey color=j scalebar=y barlabel=Cost screenratio=0.34 screenht=10 
     title="Exact df" barunit= minval=-0.125 maxval=0.24 clip=0.084
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     ''')

Flow('edt','pt0 t02','add scale=1,-1 ${SOURCES[1]}')
Plot('edt',
     '''
     grey color=j scalebar=y barlabel=Time screenratio=0.34 screenht=10 
     title="Exact dt" barunit=s labelsz=10 titlesz=12 titlefat=6 labelfat=6
     ''')

Flow('edx','px0 x02','add scale=1,-1 ${SOURCES[1]}')
Plot('edx',
     '''
     grey color=j scalebar=y barlabel=Position screenratio=0.34 screenht=10 
     title="Exact dx" barunit=km labelsz=10 titlesz=12 titlefat=6 labelfat=6
     ''')

# slowness-squared perturbation
Flow('ds','pert modl',
     '''
     math a=${SOURCES[0]} b=${SOURCES[1]} output="1./a/a-1./b/b"
     ''')

# linear prediction (dt)
Flow('pdt','ds modl t02 x02 f02 vdix',
     '''
     t2diter s0=${SOURCES[1]} t0=${SOURCES[2]} x0=${SOURCES[3]} 
     f0=${SOURCES[4]} dix=${SOURCES[5]} what=t
     ''')
Plot('pdt',
     '''
     grey color=j scalebar=y barlabel=Time screenratio=0.34 screenht=10 
     title="Predicted dt" barunit=s labelsz=10 titlesz=12 titlefat=6 labelfat=6
     ''')

Result('pdt','edt pdt','OverUnderIso')

# linear prediction (dx)
Flow('pdx','ds modl t02 x02 f02 vdix',
     '''
     t2diter s0=${SOURCES[1]} t0=${SOURCES[2]} x0=${SOURCES[3]} 
     f0=${SOURCES[4]} dix=${SOURCES[5]} what=x
     ''')
Plot('pdx',
     '''
     grey color=j scalebar=y barlabel=Position screenratio=0.34 screenht=10 
     title="Predicted dx" barunit=km labelsz=10 titlesz=12 titlefat=6 labelfat=6
     ''')

Result('pdx','edx pdx','OverUnderIso')

# linear prediction (cost)
Flow('predict','ds modl t02 x02 f02 vdix',
     '''
     t2diter s0=${SOURCES[1]} t0=${SOURCES[2]} x0=${SOURCES[3]} 
     f0=${SOURCES[4]} dix=${SOURCES[5]} what=linear adj=n
     ''')
Plot('predict',
     '''
     grey color=j scalebar=y barlabel=Cost screenratio=0.34 screenht=10 
     title="Predicted df" barunit= minval=-0.125 maxval=0.24 clip=0.084
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     ''')

Result('diffcost','diffcost predict','OverUnderIso')

# gradient (Ax = b)
Flow('egrad','predict modl t02 x02 f02 vdix',
     '''
     t2diter mod=${SOURCES[1]} dat=${SOURCES[1]} cgiter=10000 verb=y eps=1. shape=y rect1=30 rect2=4 
     s0=${SOURCES[1]} t0=${SOURCES[2]} x0=${SOURCES[3]} f0=${SOURCES[4]} dix=${SOURCES[5]}
     ''')
Plot('egrad',
     '''
     grey color=j scalebar=y title="Gradient (Ax = b)"
     barlabel=Slowness-squared barunit="s\^2\_/km\^2\_" screenratio=0.34 screenht=10
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     ''')

# gradient (Ax ~ b)
Flow('pgrad','diffcost modl t02 x02 f02 vdix',
     '''
     t2diter mod=${SOURCES[1]} dat=${SOURCES[1]} cgiter=10000 verb=y eps=1. shape=y rect1=30 rect2=4
     s0=${SOURCES[1]} t0=${SOURCES[2]} x0=${SOURCES[3]} f0=${SOURCES[4]} dix=${SOURCES[5]}
     ''')
Plot('pgrad',
     '''
     grey color=j scalebar=y title="Gradient (Ax ~ b)"
     barlabel=Slowness-squared barunit="s\^2\_/km\^2\_" screenratio=0.34 screenht=10
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     ''')

Result('grad','egrad pgrad','OverUnderIso')

# convert vdix to depth
Flow('vz0','vd',
     '''
     transp | time2depth velocity=$SOURCE intime=y nz=500 dz=0.00601202 | transp |
     window min1=4. max1=14. max2=2. | transp |
     put o2=0. label1=Depth unit1=km label2=Position unit2=km
     ''')

Plot('pvz0','vz0',
     '''
     grey color=j scalebar=y allpos=y title="Dix velocity converted to depth"
     barlabel=Velocity barunit=km/s barreverse=y bias=1.9 screenratio=0.34 screenht=10
     minval=2. maxval=3.4 labelsz=10 titlesz=12 titlefat=6 labelfat=6 clip=1.25
     ''')

# eikonal for prior model
Flow('beiko bt0 bx0 bf0','vz0','irays t0=${TARGETS[1]} x0=${TARGETS[2]} f0=${TARGETS[3]}')

Plot('xz0','bx0',
     '''
     contour nc=200 plotcol=7 plotfat=7 wanttitle=n scalebar=y     
     yreverse=y wantaxis=n screenratio=0.34 screenht=10
     min2=0 max2=10 min1=0 max1=2
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     ''')

Plot('vz0','pvz0 xz0','Overlay')

Result('vz0','modl vz0 pert','OverUnderIso')

# initial cost
Flow('bcost','vz0 bt0 bx0 bf0 vdix f02',
     '''
     t2diter s0=${SOURCES[0]} t0=${SOURCES[1]} x0=${SOURCES[2]} 
     f0=${SOURCES[3]} dix=${SOURCES[4]} mask=${SOURCES[5]} what=cost
     ''')

# subtract from exact cost
Flow('dcost','bf0 cost bcost',
     '''
     dd type=float | add mode=p ${SOURCES[1]} | 
     add scale=1,1 ${SOURCES[2]}
     ''')

# update
Flow('bgrad','dcost vz0 bt0 bx0 bf0 vdix f02',
     '''
     t2diter cgiter=10000 verb=y eps=1.5 shape=y rect1=5 rect2=5 s0=${SOURCES[1]}
     t0=${SOURCES[2]} x0=${SOURCES[3]} f0=${SOURCES[4]} dix=${SOURCES[5]} mask=${SOURCES[6]}
     ''')
Plot('bgrad',
     '''
     grey color=j scalebar=y title=Update
     barlabel=Slowness-squared barunit="s\^2\_/km\^2\_" screenratio=0.34 screenht=10
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     ''')

# compare
Flow('eds','modl vz0',
     '''
     math a=${SOURCES[0]} b=${SOURCES[1]} output="1./a/a-1./b/b"
     ''')
Plot('eds',
     '''
     grey color=j scalebar=y title=Exact
     barlabel=Slowness-squared barunit="s\^2\_/km\^2\_" screenratio=0.34 screenht=10
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     ''')

# recovered model
Flow('recov','vz0 bgrad',
     '''
     math a=${SOURCES[0]} b=${SOURCES[1]} output="1./a/a+b" |
     math output="1./sqrt(input)"
     ''')
Plot('recov',
     '''
     grey color=j scalebar=y allpos=y barlabel=Velocity barunit=km/s
     title=Inverted screenratio=0.34 screenht=10 minval=2 maxval=3.4 bias=1.9
     labelsz=10 titlesz=12 titlefat=6 labelfat=6 barreverse=y clip=1.25
     ''')

Result('bgrad','eds bgrad recov','OverUnderIso')

End()