from rsf.proj import *
import math

# get the data

Fetch('pptrace.rsf','attr')
Flow('pptrace1','pptrace.rsf',
     'dd form=native | bandpass flo=3 | scale dscale=0.1 | put label1=Time unit1=s')

graph = '''
        window max1=6 |
        graph pad=n crowd1=0.75 label2=Amplitude unit2= label1=Time unit1=s crowd2=0.3 wanttitle=n
        min2=-10 max2=10
        '''

Plot('trace','pptrace1',
     '''
     graph title="Seismic Trace" label1= unit1= label2=Amplitude
     labelsz=15 titlesz=18 wheretitle=t wherexlabel=b min2=-10 max2=10
     ''')

Result('trace','pptrace1',graph)

# local frequency (local attributes)

Plot('freq','pptrace1',
     '''
     iphase rect1=30 hertz=y |
     graph title="Local Frequency"
     min2=0 max2=50 plotcol=5
     label1=Time unit1=s label2=Frequency unit2=Hz 
     labelsz=15 titlesz=18 wheretitle=t wherexlabel=b wantaxis1=n
     ''')

Plot('seis','trace freq','OverUnderAniso')
Result('seis','Overlay',vppen='yscale=0.5')

# EMD
Flow('emd','pptrace1','emd')

# apply shifts in time

d1 = 0.004 # time sampling

Flow('shift','pptrace1','shift1 ns=5')

# find adaptive PEF

Flow('pef3 pre3','shift pptrace1',
     'window n2=2 | lpf match=${SOURCES[1]} rect1=30 pred=${TARGETS[1]}')
Flow('err3','pre3 pptrace1','add scale=-1,1 ${SOURCES[1]}')

Plot('err3',
     '''
     graph title="Nonstationary Deconvolution" label1= unit1= label2=Amplitude
     labelsz=15 titlesz=18 wheretitle=t wherexlabel=b 
     ''')

Plot('decon3','trace err3','OverUnderAniso')

# analytical trace
Flow('itrace','pptrace1','envelope hilb=y')
Flow('ctrace','pptrace1 itrace','cmplx ${SOURCES[1]}')

Flow('ishift','shift','envelope hilb=y')
Flow('cshift','shift ishift','cmplx ${SOURCES[1]} | put o1=0')

Flow('cpef3 cpre3','cshift ctrace',
     'window n2=1 | clpf match=${SOURCES[1]} rect1=30 pred=${TARGETS[1]}')
Flow('cerr3','cpre3 ctrace','add scale=-1,1 ${SOURCES[1]}')

Plot('cerr3',
     '''
     real |
     graph title="Nonstationary Deconvolution" label1= unit1= label2=Amplitude
     labelsz=15 titlesz=18 wheretitle=t wherexlabel=b min2=-10 max2=10
     ''')

Plot('cdecon3','trace cerr3','OverUnderAniso')
Result('cdecon3','Overlay',vppen='yscale=0.5')

# find local frequency from adaptive PEF

Flow('atten','pef3','window n2=1 f2=1 | math output="sqrt(-input)" ')
Flow('ifreq','pef3 atten',
     '''
     window n2=1 | math atten=${SOURCES[1]} output="input/(2*atten)" | 
     clip2 lower=-1 upper=1 |
     math output="acos(input)/%g"
     ''' % (2*math.pi*d1))

Plot('ifreq',
     '''
     graph title="Local Frequency (from NAR)"
     min2=0 max2=50 plotcol=5
     label1=Time unit1=s label2=Frequency unit2=Hz
     labelsz=15 titlesz=18 wheretitle=t wherexlabel=b wantaxis1=n
     ''')

Plot('iseis','trace ifreq','OverUnderAniso')
Result('iseis','Overlay',vppen='yscale=0.5')

# find local frequency from complex-valued adaptive PEF

Flow('icfreq','cpef3','math output="atan(arg(input))/%g" | real' % (2*math.pi*d1))

Plot('icfreq',
     '''
     graph title="Local Frequency (from CNAR)"
     min2=0 max2=50 plotcol=5
     label1=Time unit1=s label2=Frequency unit2=Hz
     labelsz=15 titlesz=18 wheretitle=t wherexlabel=b wantaxis1=n
     ''')

Plot('icseis','trace icfreq','OverUnderAniso')
Result('icseis','Overlay',vppen='yscale=0.5')

# find complex adaptive PEF

Flow('cpef cpre','cshift ctrace',
     'clpf match=${SOURCES[1]} rect1=75 pred=${TARGETS[1]}')
Flow('cerr','cpre ctrace','add scale=-1,1 ${SOURCES[1]}')

Plot('cerr',
     '''
     real |
     graph title="Nonstationary Deconvolution" label1= unit1= label2=Amplitude
     labelsz=15 titlesz=18 wheretitle=t wherexlabel=b min2=-10 max2=10
     ''')

Result('gerr','cerr','real | ' + graph)

Plot('cdecon','trace cerr','OverUnderAniso')

Result('cdecon','Overlay',vppen='yscale=0.5')

Result('cerr','cerr ctrace',
       'cat axis=2 ${SOURCES[1]} | real | window max1=6 | dots gaineach=n')

# Find complex polynomial roots

Flow('cpoly','cpef','window n2=1 | math output=-1 | cat axis=2 $SOURCE')
Flow('croots','cpoly','transp | roots verb=n niter=100 sort=p | transp')

# Frequency components

Flow('group','croots','math output="-arg(input)/%g" | real' % (2*math.pi*d1))

Result('group4','group',
       '''
       window max1=6 |
       graph title="Instantaneous Frequencies" label1=Time unit1=s label2=Frequency unit2=Hz
       plotfat=3 min2=0
       ''')

Flow('freqs','group',
     '''
     causint | math output="input*%g/(x1+%g)" 
     ''' % (d1,d1))

Result('freqs','graph title="Local Frequencies" min2=0')

Flow('comps','freqs','rtoc | math output="exp(I*input*x1*%g)" ' % (2*math.pi) )

Flow('cwht cfit','comps ctrace',
     'clpf match=${SOURCES[1]} rect1=6 pred=${TARGETS[1]}')

Flow('cdif','cfit ctrace','add scale=1,-1 ${SOURCES[1]}')

Result('gfit','cfit','real | ' + graph)

Result('cfit','cdif cfit ctrace',
       '''
       cat axis=2 ${SOURCES[1:3]} | real | window max1=6 |
       dots labels=Residual:Fit:Data gaineach=n label1=Time unit1=s
       ''')

Flow('csign','comps cwht','math other=${SOURCES[1]} output="input*other"')

Result('csign','csign cdif',
       '''
       cat axis=2 ${SOURCES[1]} | real | window max1=6 n2=4 |
       dots gaineach=n yreverse=y label1=Time unit1=s 
       ''')

for case in range(4):
     sign = 'sign%d' % case
     Result(sign,'csign',
            '''
            window n2=1 f2=%d | real | %s plotcol=%d min2=-5 max2=5 crowd2=0.2 parallel2=n
            ''' % (case,graph,(2,3,4,2)[case]))

# mask low amplitudes

Flow('mask','cwht','math output="abs(input)" | real | mask min=0.2')

group = []
for case in range(4):
    mask = 'mask%d' % case
    Flow(mask,'mask','window n2=1 f2=%d' % case)
    grup = 'group%d' % case
    Flow(grup,['group',mask],
         '''
         window n2=1 f2=%d | 
         rtoc | math output="x1+I*input" | 
         transp plane=12 | headerwindow mask=${SOURCES[1]} | window
         ''' % case)
    Plot(grup,
         '''
         graph  title="Instantaneous Frequencies" label1=Time unit1=s label2=Frequency unit2=Hz
         plotfat=3 min2=0 plotcol=%d min2=0 max2=80 min1=0 max1=6
         ''' % (2,3,4,1)[case])
    group.append(grup)

Result('tgroup',group,'Overlay')

End()