## 
 # elastic modeling; wavefield separation
 ## for a heterogeneous VTI model

from rsf.proj import *
import sys
#sys.path.append('Python')
import fdmod,pot,fdd,spk,stiff
# ------------------------------------------------------------
par = {
    'nx':600, 'ox':0, 'dx':0.002,  'lx':'x', 'ux':'km',
    'nz':600, 'oz':0, 'dz':0.002,  'lz':'z', 'uz':'km',
    'nt':2401,'ot':0, 'dt':0.0002, 'lt':'t', 'ut':'sec',
    'kt':150,
    'jsnap':200,
    'height':10,
    'nb':0,
    'frq':75,
    'ratio':1
    }
fdmod.param(par)
par['labelattr']=par['labelattr']+'''
titlesz=15 labelsz=9
n2tic=7 o2num=0 d2num=0.2
n1tic=7 o1num=0 d1num=0.2
'''
nframe=12

# ------------------------------------------------------------
# source/receiver coordinates
# ------------------------------------------------------------
fdmod.point('ss',
            par['ox']+(par['nx']/2*par['dx']),
            par['oz']+(par['nz']/2*par['dz']),par)
fdmod.horizontal('rr',0,par)

Plot('rr',fdmod.rrplot('xll=2.2',par))
Plot('ss',fdmod.ssplot('xll=2.2',par))

# ------------------------------------------------------------
# model parameters
# ------------------------------------------------------------
Flow('zero',None,
     '''
     spike nsp=1 mag=0.0
     n1=%(nz)d o1=%(oz)g d1=%(dz)g 
     n2=%(nx)d o2=%(ox)g d2=%(dx)g |
     put label1=%(lz)s label2=%(lx)s unit1=%(uz)s unit2=%(ux)s
     ''' % par)

sig=0.05
Flow('vp',     'zero',
     '''
     math output="exp(-( (x1-0.65)*(x1-0.65)+(x2-0.65)*(x2-0.65))/(2*%g*%g))" |
     scale rscale=-1.0 |
     add add=3.00
     ''' % (sig,sig) )

Flow('vs',     'zero',
     '''
     math output="exp(-( (x1-0.65)*(x1-0.65)+(x2-0.65)*(x2-0.65))/(2*%g*%g))" |
     scale rscale=-0.5 |
     add add=1.50
     ''' % (sig,sig) )

Flow('ro',None,
     '''
     spike nsp=1 mag=1.0
     n1=%(nz)d o1=%(oz)g d1=%(dz)g k1=250 l1=%(nz)d
     n2=%(nx)d o2=%(ox)g d2=%(dx)g |
     add add=1.0 |
     put label1=%(lz)s label2=%(lx)s unit1=%(uz)s unit2=%(ux)s
     ''' % par)

Flow('epsilon','zero',
     '''
     add add=+0.25 |
     window f1=%d |
     pad beg1=%d |
     smooth rect1=100
     ''' % (par['nz']/2,par['nz']/2))
Flow('delta','zero',
     '''
     add add=-0.29 |
     window f2=%d |
     pad beg2=%d |
     smooth rect2=100
     ''' % (par['nx']/2,par['nx']/2))

Flow('nu',     'zero','add add=0.00')


# ------------------------------------------------------------
Plot('vp',fdmod.cgrey('bias=2.00 allpos=y labelsz=12 xll=2.2',par))
Plot('vs',fdmod.cgrey('bias=1.00 allpos=y labelsz=12 xll=2.2',par))
Plot('ro',fdmod.cgrey('bias=1.00 allpos=y labelsz=12 xll=2.2',par))
Plot('epsilon',fdmod.cgrey(' labelsz=12 xll=2.2',par))
Plot('delta',fdmod.cgrey('bias=-0.29 labelsz=12 xll=2.2',par))

for k in (['vp','vs','ro','epsilon','delta']):
    Result(k,[k,'ss'],'Overlay')

stiff.iso2d('cI','vp','vs','ro',par)
stiff.tti2d('cA','vp','vs','ro','epsilon','delta','nu',par)

# ------------------------------------------------------------
# elastic source
# ------------------------------------------------------------
fdmod.wavelet('wav_',par['frq'],par)
Flow('ver','wav_','math output="+1*input"')
Flow('hor','wav_','math output="0*input"')
Flow('wave',['ver','hor'],
     '''
     cat axis=2 space=n ${SOURCES[1:2]} |
     transp plane=12 |
     transp plane=23 |
     transp plane=12
     ''')
fdmod.ewavelet('wave','',par)

# ------------------------------------------------------------
# modeling
# ------------------------------------------------------------
# EWE modeling: output displacements
fdmod.ewefd2('dIu','uI','wave','cI','ro','ss','rr','ssou=n opot=n',par)
fdmod.ewefd2('dAu','uA','wave','cA','ro','ss','rr','ssou=n opot=n',par)

fdmod.emovie('uImovie','uI',6,'pclip=98',2,par,0.75,0.75,-8.25)
fdmod.emovie('uAmovie','uA',6,'pclip=98',2,par,0.75,0.75,-8.25)

# ------------------------------------------------------------
# derivative operators

fdd.derivatives(par)

spk.delt('spk',64,64,par)
order=8


fdd.separatorD('dzI','dxI','spk','cI','y','x','sine',1.0,order,4,1,par)
#M stationary
fdd.separatorD('dzM','dxM','spk','cA','y','x','sine',1.0,order,25,25,par)
#A non-sationary
fdd.separatorD('dzA','dxA','spk','cA','n','x','sine',1.0,order,25,25,par)


# impulse responses
fdd.oneirST('cop','dzC','dxC',7,7,'',par)
fdd.oneirST('iop','dzI','dxI',7,7,'',par)
fdd.oneirST('mop','dzM','dxM',25,25,'pclip=99.9',par)
fdd.oneirNS('aop','dzA','dxA',7,7,'labelsz=20',par)

# ------------------------------------------------------------
# isotropic displacements/potentials
#pot.displacements('uI','uIz','uIx',nframe,'',par)
#pot.potentials(   'qI','uIz','uIx','dzC','dxC','y','','',par)
#pot.potentials(   'pI','uIz','uIx','dzI','dxI','y','','',par)
#
## ------------------------------------------------------------
## anisotropic displacements/potentials
#pot.displacements('uA','uAz','uAx',nframe,'',par)
#pot.potentials(   'qA','uAz','uAx','dzC','dxC','y','','q',par)
#pot.potentials(   'pM','uAz','uAx','dzM','dxM','y','','q',par)
#pot.potentials(   'pA','uAz','uAx','dzA','dxA','n','','q',par)

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

oframe=6
nframe=1
jframe=1
for iframe in range(oframe,oframe+nframe,jframe):
    tag = "-f%02d"%iframe

    # ------------------------------------------------------------
    # isotropic displacements/potentials
    pot.displacementsC('uI'+tag,'uI','uIz'+tag,'uIx'+tag,iframe,'pclip=99.5',par)
    pot.cliptogether( 'uI'+tag,'uIz'+tag,'uIx'+tag,'"u\_z"','"u\_x"',1,'',par)

    pot.potentials(   'qI'+tag,'uIz'+tag,'uIx'+tag,'dzC','dxC','y','pclip=99.5','',par)
    pot.cliptogether( 'qI'+tag,'qI'+tag+'p', 'qI'+tag+'s',"qP","qS",1,'',par)

    pot.potentials(   'pI'+tag,'uIz'+tag,'uIx'+tag,'dzI','dxI','y','pclip=99.5','',par)
    pot.cliptogether( 'pI'+tag,'pI'+tag+'p', 'pI'+tag+'s',"qP","qS",1,'',par)
    # ------------------------------------------------------------
    # anisotropic displacements/potentials
    pot.displacementsC('uA'+tag,'uA','uAz'+tag,'uAx'+tag,iframe,'pclip=99.5',par)
    pot.cliptogether( 'uA'+tag,'uAz'+tag,'uAx'+tag,'"u\_z"','"u\_x"',1,'',par)

    pot.potentials(   'qA'+tag,'uAz'+tag,'uAx'+tag,'dzC','dxC','y','pclip=99.5','q',par)
    pot.cliptogether( 'qA'+tag,'qA'+tag+'p', 'qA'+tag+'s',"qP","qS",1,'',par)

    pot.potentials(   'pA'+tag,'uAz'+tag,'uAx'+tag,'dzA','dxA','n','pclip=99.5','q',par)
    pot.cliptogether( 'pA'+tag,'pA'+tag+'p', 'pA'+tag+'s',"qP","qS",1,'',par)


for     jx in range(3):
        fx = (jx+1) * par['nx']/4 + 1
        for jz in range(3):
            fz = (jz+1) * par['nz']/4 + 1
            tag = str(jx)+str(jz)

fdmod.boxarray('aoppos',3,(par['nx']/4+1)*par['dx'],par['nx']/4*par['dx'],
                        3,(par['nz']/4+1)*par['dz'],par['nz']/4*par['dz'],par)
Plot('aoppos',fdmod.qqplot('symbol=o plotcol=3 plotfat=10 xll=2.2',par))
Result('aoppos',['vp','aoppos'],'Overlay')


End()