from rsf.proj import *

par = dict(
    nx=200,  ox=-10, dx=0.1,  lx='x', ux='km',
    nz=200,  oz=-10, dz=0.1,  lz='z', uz='km',
    nt=1000, ot=0,   dt=0.01, # traveltime
    ng=1801, og=-90, dg=0.1,  # angle
    sig=1.5,                  # stdev
    xsou=0, zsou=-10,         # coords
    labelattr="titlefat=3 labelsz=5 labelfat=3"
    )
par['xmin']=par['ox']
par['xmax']=par['ox'] + (par['nx']-1)*par['dx']
par['zmin']=par['oz']
par['zmax']=par['oz'] + (par['nz']-1)*par['dz']
par['ratio']=(par['nz']-1)*par['dz']/(par['nx']-1)/par['dx']

# ------------------------------------------------------------

# plot rays and wavefronts
def hwtplot(custom,par):
    return '''
    graph 
    title="" screenratio=%(ratio)g
    plotcol=0 plotfat=3 wantaxis=n yreverse=y
    min1=%(zmin)g max1=%(zmax)g
    min2=%(xmin)g max2=%(xmax)g
    %(labelattr)s
    '''%par + custom

# plot traveltime contour lines
def fmeplot(custom,par):
    return '''
    contour 
    title="" screenratio=%(ratio)g
    nc=100 plotcol=6 plotfat=3
    labelrot=n wantaxis=n wanttitle=n
    %(labelattr)s
    '''%par + custom

# ------------------------------------------------------------

# make Gaussian function
Flow('gg',None,
     '''
     math output="exp(-(x1^2+x2^2)/(2*%(sig)g^2))"
     n1=%(nz)d d1=%(dz)g o1=%(oz)g
     n2=%(nx)d d2=%(dx)g o2=%(ox)g |
     put label1=%(lz)s unit1=%(uz)s 
         label2=%(lx)s unit2=%(ux)s
     '''%par)

# plot 2D Gaussian function
Result('gg',
       '''
       grey title="2D Gaussian" 
       pclip=100 screenratio=%(ratio)g
       %(labelattr)s
       '''%par)

# plot 1D Gaussian function
Result('gg0','gg',
       '''
       window n1=1 min1=0 |
       graph title="1D Gaussian"
       plotfat=10 screenratio=0.5 screenht=7
       %(labelattr)s
       '''%par)

# ------------------------------------------------------------

# make velocity
Flow('vel','gg',
     'math output="3.0-input"')

# plot velocity
Plot('vel',
     '''
     grey title="" pclip=100 color=g
     screenratio=%(ratio)g mean=y %(labelattr)s
     '''%par)
Result('vel','grey title="" pclip=100 screenratio=%(ratio)g mean=y %(labelattr)s'%par)
# ------------------------------------------------------------

# make eikonal traveltimes
Flow(  'fme','vel',
       'eikonal zshot=%(zsou)g yshot=%(xsou)g'%par)

# plot eikonal traveltimes
Plot(  'fme',fmeplot('plotcol=1',par))
Result('fme',['vel','fme'],'Overlay')

# ------------------------------------------------------------

# make rays and wavefronts
Flow('hwt','vel',
     '''
     hwt2d xsou=%(xsou)g zsou=%(zsou)g
     nt=1000 ot=0   dt=0.01
     ng=1801 og=-90 dg=0.1
     '''%par)

# plot rays and wavefronts
Plot('ray','hwt',
     'transp | window f1=25 j1=25 j2=20 |'
     + hwtplot('plotcol=5',par))
Plot('wft','hwt',
     'window f2=50 j2=25 j1=20 |'
     + hwtplot('plotcol=6',par))
Result('hwt',['vel','wft','ray'],'Overlay')

# ------------------------------------------------------------

End()