# grid
nx=np.int32(5000)
h=np.float32(0.005)

# time
nt=np.int32(5000)
dt=np.float32(0.001)

# velocity
v=np.zeros(nx,dtype=np.float32)
v[:]=4.0

# source position
delta=np.zeros_like(v)
delta[nx/2]=1.

# source wavelet
w=np.fromfile("wavelet.bin",dtype=np.float32)

plot(np.linspace(0.,nt*dt,nt),w,'k-')
xlabel("Time (s)",fontsize='large')
ylabel('Amplitude',fontsize='large')
xlim([0,5])
savefig('wavelet.eps')

#import numpy as np
def model1d(nx,nt,h,dt,v,delta,w):
    um=np.zeros_like(v)
    uo=np.zeros_like(um)
    ud=np.zeros_like(um)
    up=np.zeros_like(um)

    h2=h*h
    dt2=dt*dt
    seismo=np.zeros((nt,nx),dtype=np.float32)
    
    for it in xrange(nt):
        for ix in xrange(1,nx-1):
            ud[ix]=(uo[ix-1]+uo[ix+1]-2.*uo[ix])/h2
        for ix in xrange(nx):
            up[ix]=(ud[ix]+delta[ix]*w[it])*v[ix]*v[ix]*dt2+2.*uo[ix]-um[ix]
        um,uo,up=uo,up,um
        for ix in xrange(nx):
            seismo[it,ix]=uo[ix]
    return seismo

def model1d_vec(nx,nt,h,dt,v,delta,w):
    um=np.zeros_like(v)
    uo=np.zeros_like(um)
    ud=np.zeros_like(um)
    
    h2=h*h
    dt2=dt*dt
    seismo=np.zeros((nt,nx),dtype=np.float32)
    
    for it in xrange(nt):
        ud[1:-1]=(uo[:-2]+uo[2:]-2.*uo[1:-1])/h2
        up=(ud+delta*w[it])*v*v*dt2 + 2.*uo - um
        um,uo=uo,up
        seismo[it,:]=uo[:]
    return seismo

seismo=model1d_vec(nx,nt,h,dt,v,delta,w)
seismo=seismo.clip(seismo.min()*0.01,seismo.max()*0.01)
plt.imshow(seismo,aspect=3,extent=[0,(nx-1)*h,(nt-1)*dt,0],cmap=cm.gray_r)
xlabel('Distance (km)',fontsize='large')
ylabel('Time (s)',fontsize='large')
savefig('seismo.eps')

seismo=model1d(nx,nt,h,dt,v,delta,w)
seismo=seismo.clip(seismo.min()*0.01,seismo.max()*0.01)
plt.imshow(seismo,aspect=1.0,cmap=cm.gray_r)

%timeit -r 3 -n 1 model1d_vec(nx,nt,h,dt,v,delta,w)

%timeit -r 3 -n 1 model1d(nx,nt,h,dt,v,delta,w)

import numba
model1d_nb=numba.autojit(model1d)

%timeit -r 3 -n 1 model1d_nb(nx,nt,h,dt,v,delta,w)

%load_ext cythonmagic

%%cython
import cython
import numpy as np
cimport numpy as np
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
def model1d_cy(int nx,int nt,float h,float dt,np.ndarray[float] v,np.ndarray[float] delta,np.ndarray[float] w):
    cdef np.ndarray[float] um=np.zeros_like(v)
    cdef np.ndarray[float] uo=np.zeros_like(um)
    cdef np.ndarray[float] ud=np.zeros_like(um)
    cdef np.ndarray[float] up=np.zeros_like(um)

    cdef float h2=h*h
    cdef float dt2=dt*dt
    cdef np.ndarray[float,ndim=2] seismo=np.zeros((nt,nx),dtype=np.float32)
    cdef int it,ix
    
    for it in xrange(nt):
        for ix in xrange(1,nx-1):
            ud[ix]=(uo[ix-1]+uo[ix+1]-2.*uo[ix])/h2
        for ix in xrange(nx):
            up[ix]=(ud[ix]+delta[ix]*w[it])*v[ix]*v[ix]*dt2+2.*uo[ix]-um[ix]
        um,uo,up=uo,up,um
        for ix in xrange(nx):
            seismo[it,ix]=uo[ix]
    return seismo

seismo=model1d_cy(nx,nt,h,dt,v,delta,w)
seismo=seismo.clip(seismo.min()*0.01,seismo.max()*0.01)
plt.imshow(seismo,aspect=1,cmap=cm.gray_r)

%timeit -r 3 -n 1 model1d_cy(nx,nt,h,dt,v,delta,w)

%%file model1d_f.f90

subroutine model1d_vec(nx,nt,h,dt,v,delta,w,seismo)
implicit none
integer,intent(in):: nx,nt
real,intent(in):: h,dt,v(nx),delta(nx),w(nt)
real,intent(out):: seismo(nx,nt)
real,dimension(nx):: um,uo,up,ud
real:: h2,dt2
integer:: it

um=0.
uo=0.
ud=0.
seismo=0.

h2=h*h
dt2=dt*dt

do it=1,nt
    ud(2:nx-1)=(uo(1:nx-2)+uo(3:nx)-2.*uo(2:nx-1))/h2
    up=(ud+delta*w(it))*v**2*dt2 + 2.*uo-um
    um=uo
    uo=up
    seismo(:,it)=uo(:)
enddo
end subroutine
                                                
subroutine model1d(nx,nt,h,dt,v,delta,w,seismo)
implicit none
integer,intent(in):: nx,nt
real,intent(in):: h,dt,v(nx),delta(nx),w(nt)
real,intent(out):: seismo(nx,nt)
real,dimension(nx):: um,uo,up,ud
real:: h2,dt2
integer it,ix
um=0.
uo=0.
ud=0.
seismo=0.

h2=h*h
dt2=dt*dt

do it=1,nt
    do ix=2,nx-1
        ud(ix)=(uo(ix-1)+uo(ix+1)-2.*uo(ix))/h2
    enddo
    do ix=1,nx
        up(ix)=(ud(ix)+delta(ix)*w(it))*v(ix)**2*dt2&
        + 2.*uo(ix)-um(ix)
    enddo
    um=uo
    uo=up
    do ix=1,nx
        seismo(ix,it)=uo(ix)
    enddo
enddo
end subroutine

!f2py -c -m model1d_f model1d_f.f90 --f90exec=/opt/local/bin/gfortran-mp-4.9 > log.txt

import model1d_f
print model1d_f.__doc__

vtr=v.transpose()
dtr=delta.transpose()
seismo=model1d_f.model1d(h,dt,vtr,dtr,w)
seismo=seismo.clip(seismo.min()*0.01,seismo.max()*0.01)
plt.imshow(seismo.transpose(),aspect=1,cmap=cm.gray_r)

%timeit -r 3 -n 1 model1d_f.model1d(h,dt,vtr,dtr,w)

%timeit -r 3 -n 1 model1d_f.model1d_vec(h,dt,vtr,dtr,w)