vel=np.fromfile('marm16km.drt',dtype=float32)
nx=576
nz=188
vel.shape=(nx,nz)
nw=10

def plot_vel(vel,img='Marmousi.eps',h=0.016):
    xmax=(vel.shape[0]-1)*h
    zmax=(vel.shape[1]-1)*h
    fs=14
    figure(figsize=(10,4))
    ax=imshow(vel.transpose(),extent=(0,xmax,zmax,0),cmap=cm.gray_r)
    tick_params(labelsize=fs)
    xlabel('Distance (km)',fontsize=fs)
    ylabel('Depth (km)',fontsize=fs)

    ax.axes.set_yticks(arange(0,zmax+1,1)) 
    cb=colorbar(shrink=0.68,pad=0.02,aspect=10,ticks=arange(1.5,6,1))
    clim([1.5,5.5])
    cb.set_label('km/s',fontsize=fs)
    ct=getp(cb.ax,'ymajorticklabels')
    setp(ct,fontsize=fs)
    
    savefig(img)

plot_vel(vel)

def moving_average(n1,n2,vel,nw):
    svel=np.empty_like(vel)
    for i in xrange(n1):
        for j in xrange(n2):
            s=0.
            isum=0
            for k in xrange(max(0,i-nw),min(i+nw,n1)):
                for l in xrange(max(0,j-nw),min(j+nw,n2)):
                    s+=vel[k,l]
                    isum+=1
            svel[i,j]=s/isum
    return svel

def moving_average_vec(n1,n2,vel,nw):
    svel=np.empty_like(vel)
    for i in xrange(n1):
        for j in xrange(n2):
            svel[i,j]=np.average(vel[max(0,i-nw):min(i+nw,n1),max(0,j-nw):min(j+nw,n2)])
    return svel

%timeit -r 5 -n 1 moving_average(nx,nz,vel,nw)

svel=moving_average_vec(nx,nz,vel,nw)
imshow(svel.transpose(),cmap=cm.gray_r)

%timeit -r 5 -n 1 moving_average_vec(nx,nz,vel,nw)



import numba

moving_average_nb=numba.autojit(moving_average)
moving_average_nv=numba.autojit(moving_average_vec)

svel=moving_average_nb(nx,nz,vel,nw)

plot_vel(svel,img='Marmousi_MA.eps')

%timeit -r 5 -n 1 moving_average_nb(nx,nz,vel,nw)

%timeit -r 5 -n 1 moving_average_nv(nx,nz,vel,nw)

%load_ext cythonmagic

%%cython
import cython
import numpy as np
cimport numpy as np
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
def moving_average_cy(int n1,int n2,np.ndarray[float,ndim=2] vel,int nw):
    cdef np.ndarray[float,ndim=2] svel=np.empty_like(vel)
    cdef int i,j,k,l,isum
    cdef double s
    for i in xrange(n1):
        for j in xrange(n2):
            s=0.
            isum=0
            for k in xrange(max(0,i-nw),min(i+nw,n1)):
                for l in xrange(max(0,j-nw),min(j+nw,n2)):
                    s+=vel[k,l]
                    isum+=1
            svel[i,j]=s/isum
    return svel

svel=moving_average_cy(nx,nz,vel,nw)
imshow(svel.transpose())

%timeit -r 5 -n 1 moving_average_cy(nx,nz,vel,nw)

%%file moving_average.f90
subroutine moving_average(n1,n2,vel,nw,svel)
implicit none
integer,intent(in):: n1,n2,nw
real,intent(in):: vel(n1,n2)
real,intent(out):: svel(n1,n2)
integer i,j,k,l,isum
real s
do i=1,n1
do j=1,n2
    s=0.
    isum=0
    do k=max(1,i-nw),min(n1,i+nw)
    do l=max(1,j-nw),min(n2,j+nw)
        s=s+vel(i,j)
        isum=isum+1
    enddo
    enddo
    svel(i,j)=s/isum
enddo
enddo
end subroutine

subroutine moving_average_vec(n1,n2,vel,nw,svel)
implicit none
integer,intent(in):: n1,n2,nw
real,intent(in):: vel(n1,n2)
real,intent(out):: svel(n1,n2)
integer i,j,i1,i2,j1,j2,isum
do i=1,n1
    i1=max(1,i-nw)
    i2=min(n1,i+nw)
    do j=1,n2
        j1=max(1,j-nw)
        j2=min(n2,j+nw)
        isum=(i2-i1+1)*(j2-j1+1)
        svel(i,j)=sum(vel(i1:i2,j1:j2))/isum
    enddo
enddo
end subroutine

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

import mov_avg_f
print mov_avg_f.__doc__

velt=vel.transpose()

svel=mov_avg_f.moving_average(velt,nw)
imshow(svel,cmap=cm.gray_r)

%timeit -r 5 -n 1 mov_avg_f.moving_average(velt,nw)

%timeit -r 5 -n 1 mov_avg_f.moving_average_vec(velt,nw)



def moving_average_i(n1,n2,vel,nw):
    svel=np.empty_like(vel)
    for i in xrange(n1):
        i1=max(0,i-nw)
        i2=min(i+nw,n1)
        for j in xrange(n2):
            j1=max(0,j-nw)
            j2=min(j+nw,n2)
            s=0.
            isum=0
            for k in xrange(i1,i2):
                for l in xrange(j1,j2):
                    s+=vel[k,l]
                    isum+=1
            svel[i,j]=s/isum
    return svel

def moving_average_vec_i(n1,n2,vel,nw):
    svel=np.empty_like(vel)
    for i in xrange(n1):
        i1=max(0,i-nw)
        i2=min(i+nw,n1)
        for j in xrange(n2):
            j1=max(0,j-nw)
            j2=min(j+nw,n2)
            svel[i,j]=np.average(vel[i1:i2,j1:j2])
    return svel

%timeit -r 5 -n 1 moving_average_i(nx,nz,vel,nw)

%timeit -r 5 -n 1 moving_average_vec_i(nx,nz,vel,nw)