from rsf.proj import *

Flow('spike0',None,'spike n1=256 k1=128 d1=50 label1= unit1= | bandpass fhi=0.002')
Flow('spike',None,'spike n1=256 k1=128 d1=50 label1= unit1= | ricker1 freq=0.1')
Flow('zero','spike','math output=0')

v = 1000.0
dt = 0.01
twopi = 6.28318530718

# Constant-velocity wave propagation
####################################

prev = 'zero'
curr = 'spike'

# Time steps
steps = ['spike']
for it in range(500):
    next = 'step%d' % it

    Flow(next,[curr,prev],
         '''
         fft1 |
         math output="2*(cos(abs(%g*x1)*%g)-1)*input" |
         fft1 inv=y |
         add scale=1,-1,2 ${SOURCES[1]} ${SOURCES[0]}
         ''' % (twopi,v*dt))

    steps.append(next)

    prev = curr
    curr = next

Flow('wave',steps,'cat axis=2 ${SOURCES[1:%d]}' % len(steps))
Result('wave','grey title=Wave transp=n')

prev = 'zero'
curr = 'spike'

# Time steps
steps = ['spike']
for it in range(500):
    next = 'astep%d' % it

    Flow(next,[curr,prev],
         '''
         fft1 |
         math output="2*cos(abs(%g*x1)*%g)*input" |
         fft1 inv=y |
         add scale=1,-1 ${SOURCES[1]} 
         ''' % (twopi,v*dt))

    steps.append(next)

    prev = curr
    curr = next

Flow('wave1',steps,'cat axis=2 ${SOURCES[1:%d]}' % len(steps))
Result('wave1','grey title=Wave transp=n')



# Constant-velocity wave propagation (cosine FT)
################################################

prev = 'zero'
curr = 'spike'

# Time steps
steps = ['spike']
for it in range(500):
    next = 'cstep%d' % it

    Flow(next,[curr,prev],
         '''
         cosft sign1=1 |
         math output="2*cos(abs(%g*x1)*%g)*input" |
         cosft sign1=-1 |
         add scale=1,-1 ${SOURCES[1]}
         ''' % (twopi,v*dt))

    steps.append(next)

    prev = curr
    curr = next

Flow('cwave',steps,'cat axis=2 ${SOURCES[1:%d]}' % len(steps))
Result('cwave','grey title="Wave (Cosine FT)" transp=n')

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

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

# FFT forward

Flow('fft','rand','fft1')

# Slow inverse FT

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

Flow('fourier0','fourier','window n1=1 squeeze=n | scale dscale=0.5')
Flow('fouriern','fourier','window n1=1 f1=-1 squeeze=n | scale dscale=0.5')
Flow('fourierm','fourier','window n1=127 f1=1')

Flow('fourier2','fourier0 fourierm fouriern','cat axis=1 ${SOURCES[1:3]}')

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

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

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

# Variable-velocity wave propagation
####################################

Flow('vel','spike','math output=1000+0.1*x1')
Result('vel','graph title=Velocity')

Flow('cvel','spike','pad end1=1 | math output=1000+0.1*x1')

# propagator matrix
Flow('prop','vel',
     '''
     spray axis=1 n=129 d=0.000078125 o=0 |
     math output="2*(cos(abs(%g*x1)*input*%g)-1)" 
     ''' % (twopi,dt))

# propagator matrix - cosine FT
Flow('cprop','cvel',
     '''
     spray axis=1 n=257 d=0.0000390625 o=0 |
     math output="2*(cos(abs(%g*x1)*input*%g)-1)" 
     ''' % (twopi,dt))

# propagator matrix - large dt
Flow('prop1','vel',
     '''
     spray axis=1 n=129 d=0.000078125 o=0 |
     math output="2*(cos(abs(%g*x1)*input*%g)-1)" 
     ''' % (twopi,20*dt))

Flow('prop2','prop fourier2','rtoc | mul ${SOURCES[1]}')
Flow('prop10','prop1 fourier2','rtoc | mul ${SOURCES[1]}')
Flow('propo2','propo fourier2','rtoc | mul ${SOURCES[1]}')

Result('prop','grey title="Propagator Matrix" color=j bias=-1 scalebar=y label1=Wavenumber unit1=1/m label2=Distance unit2=m')


# propagator matrix (without -1)
Flow('propo','vel',
     '''
     spray axis=1 n=129 d=0.000078125 o=0 |
     math output="2*cos(abs(%g*x1)*input*%g)" 
     ''' % (twopi,dt))

Result('propo','grey title="Propagator Matrix" color=j scalebar=y label1=Wavenumber unit1=1/m label2=Distance unit2=m')


# Lowrank decomposition
#######################

targets = ','.join(map(lambda x: '\'${TARGETS[%d]}\'' % x, range(6)))

Flow('prod left mid right','prop',
     'lrmatrix seed=2010 left=${TARGETS[1]} mid=${TARGETS[2]} right=${TARGETS[3]}')
Flow('prod2 left2 right2','prop',
     'lrmatrix seed=2010 left=${TARGETS[1]} right=${TARGETS[2]} outputs=2')

Flow('cprod cleft cmid cright','cprop',
     'lrmatrix seed=2012 left=${TARGETS[1]} mid=${TARGETS[2]} right=${TARGETS[3]}')
Flow('tprod tleft tmid tright','prop1',
     'lrmatrix seed=2011 left=${TARGETS[1]} mid=${TARGETS[2]} right=${TARGETS[3]}')
Flow('prodo lefto mido righto','propo',
     'lrmatrix seed=2011 left=${TARGETS[1]} mid=${TARGETS[2]} right=${TARGETS[3]}')



Result('prod','grey title="Lowrank Propagator Matrix" color=j bias=-1 scalebar=y label1=Wavenumber unit1=1/m label2=Distance unit2=m')
Result('prodo','grey title="Lowrank Propagator Matrix" color=j scalebar=y label1=Wavenumber unit1=1/m label2=Distance unit2=m')

Result('proderr','prod prop',
       'add scale=1,-1 ${SOURCES[1]} | grey title="Lowrank Approximation Error" color=j scalebar=y label1=Wavenumber unit1=1/m label2=Distance unit2=m')

# Exact matrix multiplication
Flow('wave2','spike prop2','fftwave1 prop=${SOURCES[1]} nt=500 dt=%g' % dt)
Result('wave2','window max2=3 | grey title="Wave" transp=n clip=0.5 label1=Distance unit1=m')

Flow('waveo','spike propo2','fftwave1 sub=n prop=${SOURCES[1]} nt=500 dt=%g' % dt)
Result('waveo','window max2=3 | grey title="Wave" transp=n clip=0.5 label1=Distance unit1=m')


Flow('wave10','spike prop10','fftwave1 prop=${SOURCES[1]} nt=50 dt=%g' % (20*dt))
Result('wave10','window max2=3 | grey title="Wave" transp=n label1=Distance unit1=m')


# Propagating a spike
Flow('wave0','spike0 prop2','fftwave1 prop=${SOURCES[1]} nt=500 dt=%g' % dt)
Result('wave0','window max2=3 | grey title="Wave" transp=n clip=0.5 label1=Distance unit1=m')

# FT in time
Flow('helm','wave0','window n2=300 | transp | fft1')

# Instantanous phase

Flow('dreal','helm','real | deriv')
Flow('dimag','helm','imag | deriv')

Flow('iphase','dreal dimag helm',
     '''
     cmplx ${SOURCES[1]} | math f=${SOURCES[2]} output=input/f | imag |
     window max1=50
     ''')

#Result('iphase','grey title="Instantaneous Phase" ')

# Lowrank propagation
Flow('awave2','spike left2 right2',
     'fftwave1 prop=${SOURCES[1]} right=${SOURCES[2]} nt=500 dt=%g' % dt)
Result('awave2',
       '''
       window max2=3 |
       grey title="Lowrank Wave" transp=n clip=0.5 label1=Distance unit1=m
       ''')

Flow('awaveo','spike mido lefto righto','fftwave1 sub=n prop=${SOURCES[1]} left=${SOURCES[2]} right=${SOURCES[3]} nt=500 dt=%g' % dt)
Result('awaveo','window max2=3 | grey title="Lowrank Wave" transp=n clip=0.5 label1=Distance unit1=m')

Flow('cwave2','spike cmid cleft cright',
     'pad end1=1 | cosftwave1 prop=${SOURCES[1]} left=${SOURCES[2]} right=${SOURCES[3]} nt=500 dt=%g | window n1=256' % dt)
Result('cwave2','window max2=3 | grey title="Lowrank Wave (Cosine FT)" transp=n clip=0.5 label1=Distance unit1=m')

Flow('awave10','spike tmid tleft tright','fftwave1 prop=${SOURCES[1]} left=${SOURCES[2]} right=${SOURCES[3]} nt=50 dt=%g' % (20*dt))
Result('awave10','window max2=3 | grey title="Lowrank Wave" transp=n label1=Distance unit1=m')

Result('waverr','awave2 wave2','add scale=1,-1 ${SOURCES[1]} | window max2=3 | grey title="Lowrank Wave Error (x 10)" transp=n clip=0.05 label1=Distance unit1=m')



End()