from rsf.proj import *

Flow('spike',None,'spike n1=256 k1=128 d1=50 label1= unit1= | ricker1 freq=0.2')

v = 1000.0
dt = 0.01
twopi = 6.28318530718

#Flow('cspike','spike','rtoc')
Flow('cpike','spike','envelope hilb=y  | scale dscale=1')
Flow('cspike','spike cpike','cmplx ${SOURCES[1]}')

# Complex Fourier transform back and forth
##################################

Flow('rand',None,'spike n1=256 d1=50 | noise seed=2009 var=0.1 | rtoc')

# FFT forward

Flow('fft','rand','fft3 axis=1 pad=1')

# Slow inverse FT
Flow('fourier','fft',
     '''
     spray axis=2 n=256 d=50 o=0 |
     math output="exp(%g*I*x1*x2)/(256)"
     ''' % twopi)

Flow('ifourier','fft',
     '''
     spray axis=2 n=256 d=50 o=0 | transp |
     math output="exp(-%g*I*x1*x2)"
     ''' % twopi)
Flow('iden','ifourier fourier','cmatrix B=${SOURCES[1]}')
Flow('propiden','prop iden','cmatrix B=${SOURCES[1]}')

Flow('ift','fft fourier','cmatmult mat=${SOURCES[1]}')

Flow('dif','ift rand','add scale=1,-1 ${SOURCES[1]}')

Result('ift','ift rand','cat axis=2 ${SOURCES[1]} | sfreal | graph wanttitle=n')

# Variable-velocity wave propagation using one-step
####################################

Flow('vel1',None,'spike n1=156 d1=50 mag=1000 | math output="input+0.1*x1" ')
Flow('vel2',None,'spike n1=100 d1=50 mag=2500 | math output="input+0.1*(x1+156)" ')
Flow('vel','vel1 vel2','cat ${SOURCES[1]} axis=1 | put label1=Distance label2=Velocity unit1="m" unit2="m/s" | smooth rect1=4 repeat=0')
#Flow('vel','spike','math output=1000+0.1*x1 | put label1=Distance label2=Velocity unit1="m" unit2="m/s"')
Result('vel1d','vel','graph title="Velocity Profile" ')

# propagator matrix
Flow('prop','vel',
     '''
     spray axis=1 n=256 d=7.8125e-05 o=-0.01 | rtoc |
     math output="exp(-I*abs(%g*x1)*input*%g)"
     ''' % (twopi,dt))
#     math output="exp(-I*abs(%g*x1)*input*%g)-1"

Flow('prop2','prop fourier','mul ${SOURCES[1]}')

Result('propr','prop',
       '''
       real |
       grey title="One-step propagator matrix - real part" color=j scalebar=y
       label1=Wavenumber unit1=1/m label2=Distance unit2=m
       ''')

Result('propi','prop',
       '''
       imag |
       grey title="One-step propagator matrix - imaginary part" color=j scalebar=y
       label1=Wavenumber unit1=1/m label2=Distance unit2=m
       ''')

# Exact matrix multiplication
Flow('wave2','cspike prop2','cfftwave1d prop=${SOURCES[1]} sub=n nt=600 dt=%g cmplx=y' % dt)
Result('wave2','real | grey title="Wave" transp=n label1=Distance unit1=m clip=0.14')

# Lowrank decomposition of a complex matrix
Flow('prod1 left1 right1','prop',
     'transp | clrmatrix seed=2010 left=${TARGETS[1]} right=${TARGETS[2]} outputs=2 | transp')
Result('prod1r','prod1',
       '''
       real |
       grey title="Lowrank Propagator Matrix" color=j scalebar=y 
       label1=Wavenumber unit1=1/m label2=Distance unit2=m
       ''')
Result('prod1i','prod1',
       '''
       imag |
       grey title="Lowrank Propagator Matrix" color=j scalebar=y 
       label1=Wavenumber unit1=1/m label2=Distance unit2=m
       ''')

# Decomposition error
Result('proderr1r','prod1 prop',
       '''
       math other=${SOURCES[1]} output="(input-other)" | real | 
       grey title="Lowrank error - real part" color=j scalebar=y 
       label1=Wavenumber unit1=1/m label2=Distance unit2=m
       ''')

Result('proderr1i','prod1 prop',
       '''
       math other=${SOURCES[1]} output="(input-other)" | imag | 
       grey title="Lowrank error - imaginary part" color=j scalebar=y 
       label1=Wavenumber unit1=1/m label2=Distance unit2=m
       ''')

# Lowrank wave propagation
Flow('lwave1','cspike right1 left1',
     'cfftwave1d right=${SOURCES[1]} left=${SOURCES[2]} nt=600 dt=%g sub=n cmplx=y' % (dt))
Result('lwave1','real | grey title="Lowrank Wave" transp=n label1=Distance unit1=m clip=0.14')

# Wave propagation movie
Plot('lwave1',
     '''
     window j2=5 | scale axis=2 | spray axis=2 n=1 |
     graph min2=-1 max2=1.5 title="One-step (seperated)" wantaxis2=n
     ''')

Flow('error','lwave1 wave2','add ${SOURCES[1]} scale=1,-1')
Result('error','real | grey title="Error" transp=n label1=Distance unit1=m clip=0.14')

End()