from rsf.proj import *

Flow('spike',None,
     '''
     spike d1=1. d2=1. n1=100 n2=100 k1=24 k2=13 nsp=1 
     label1=x unit1=m label2=y unit2=m
     ''')
Plot('spike','grey titlesz=15 labelsz=10 crowd=0.75 pclip=100 title="(a)" ')

Fetch('schlum.HH','vsp')
Flow('schlum','schlum.HH',
     'dd form=native | put label1=time unit1=s label2=depth unit2=m')
Plot('schlum','grey titlesz=15 labelsz=10 crowd=0.75 title="(a)" ')

take = dict(spike='real',
            schlum='math output="abs(input)" | real')
clip = dict(spike='clip=1 label1="k\_x\^" label2="k\_y\^" ',
            schlum='''
                   clip=2e6 allpos=y label1=frequency label2=wavenumber
                   unit1=Hz unit2=1/m color=b
                   ''')
grid = dict(spike='d1=0.0976562 o1=-50 d2=0.0976562 o2=-50 unit1=1/m',
            schlum='d1=0.0813802 o1=-41.6667 d2=3.90625e-05 o2=-0.02')

for data in ('spike','schlum'):
    plots = [data]

    # 2-D FFT
    fft2 = data+'-fft2'
    Flow(fft2,data,
         '''
         pad n1=1024 n2=1024 | rtoc | 
         fft3 axis=1 pad=1 | fft3 axis=2 pad=1
         ''')

    # 1-D helix FFT
    fft  = data+'-fft'
    fft1 = data+'-fft1'
    left = data+'-left'
    rite = data+'-rite'
    diff = data+'-diff'

    Flow(fft,data,
         '''
         pad n1=1024 n2=1024 | put n2=1 n1=1048576 | rtoc | 
         fft3 axis=1 pad=1 | put n1=1024 n2=1024 | transp
         ''')
    Flow(left,fft,'window n2=512')
    Flow(rite,fft,'window f2=512 n1=1023 | pad beg1=1')
    Flow(fft1,[rite,left],
         '''
         cat axis=2 ${SOURCES[1]} | 
         put %s
         ''' % grid[data])
    Flow(diff,[fft1,fft2],'add scale=1,-1 ${SOURCES[1]}')

    for case in Split('fft2 fft1 diff'):
        plot='-'.join([data,case])
        Plot(plot,
             '''
             %s | 
             grey titlesz=15 labelsz=10 crowd=0.75 %s title="(%c)"
             ''' % (take[data],clip[data],
                    {'fft2':'b','fft1':'c','diff':'d'}[case]))
        plots.append(plot)

    Result(data+'spec',plots,'TwoRows')

End()