3D Waterwave Simulation using Python

I used Numpy  Matplotlib with Animation and 3d Plot module on my OS X Yosemite.

The code is still messy and clearly not efficient (there's slow loop here and there) but it works, :)
Here The Result
The Code

import numpy as np

n = 8;
g = 9.8;
dt = 0.02;
dx = 1.0;
dy = 1.0;

h = np.ones((n+2,n+2))
u = np.zeros((n+2,n+2))
v = np.zeros((n+2,n+2))

hx = np.zeros((n+1,n+1))
ux = np.zeros((n+1,n+1))
vx = np.zeros((n+1,n+1))

hy = np.zeros((n+1,n+1))
uy = np.zeros((n+1,n+1))
vy = np.zeros((n+1,n+1))

nsteps = 0
h[1,1] = .5;

def reflective():
    h[:,0] = h[:,1]
    h[:,n+1] = h[:,n]
    h[0,:] = h[1,:]
    h[n+1,:] = h[n,:]
    u[:,0] = u[:,1]
    u[:,n+1] = u[:,n]
    u[0,:] = -u[1,:]
    u[n+1,:] = -u[n,:]
    v[:,0] = -v[:,1]
    v[:,n+1] = -v[:,n]
    v[0,:] = v[1,:]
    v[n+1,:] = v[n,:]

def proses():
    #hx = (h[1:,:]+h[:-1,:])/2-dt/(2*dx)*(u[1:,:]-u[:-1,:])
    for i in range (n+1):
        for j in range(n):
            hx[i,j] = (h[i+1,j+1]+h[i,j+1])/2 - dt/(2*dx)*(u[i+1,j+1]-u[i,j+1])
            ux[i,j] = (u[i+1,j+1]+u[i,j+1])/2- dt/(2*dx)*((pow(u[i+1,j+1],2)/h[i+1,j+1]+ g/2*pow(h[i+1,j+1],2))- (pow(u[i,j+1],2)/h[i,j+1]+ g/2*pow(h[i,j+1],2)))
            vx[i,j] = (v[i+1,j+1]+v[i,j+1])/2 - dt/(2*dx)*((u[i+1,j+1]*v[i+1,j+1]/h[i+1,j+1]) - (u[i,j+1]*v[i,j+1]/h[i,j+1]))

    for i in range (n):
        for j in range(n+1):
            hy[i,j] = (h[i+1,j+1]+h[i+1,j])/2 - dt/(2*dy)*(v[i+1,j+1]-v[i+1,j])
            uy[i,j] = (u[i+1,j+1]+u[i+1,j])/2 - dt/(2*dy)*((v[i+1,j+1]*u[i+1,j+1]/h[i+1,j+1]) - (v[i+1,j]*u[i+1,j]/h[i+1,j]))
            vy[i,j] = (v[i+1,j+1]+v[i+1,j])/2 - dt/(2*dy)*((pow(v[i+1,j+1],2)/h[i+1,j+1] + g/2*pow(h[i+1,j+1],2)) - (pow(v[i+1,j],2)/h[i+1,j] + g/2*pow(h[i+1,j],2)))
    
    for i in range (1,n+1):
        for j in range(1,n+1):
            h[i,j] = h[i,j] - (dt/dx)*(ux[i,j-1]-ux[i-1,j-1]) - (dt/dy)*(vy[i-1,j]-vy[i-1,j-1])
            u[i,j] = u[i,j] - (dt/dx)*((pow(ux[i,j-1],2)/hx[i,j-1] + g/2*pow(hx[i,j-1],2)) - (pow(ux[i-1,j-1],2)/hx[i-1,j-1] + g/2*pow(hx[i-1,j-1],2))) - (dt/dy)*((vy[i-1,j]*uy[i-1,j]/hy[i-1,j]) - (vy[i-1,j-1]*uy[i-1,j-1]/hy[i-1,j-1]))
            v[i,j] = v[i,j] - (dt/dx)*((ux[i,j-1]*vx[i,j-1]/hx[i,j-1]) - (ux[i-1,j-1]*vx[i-1,j-1]/hx[i-1,j-1])) - (dt/dy)*((pow(vy[i-1,j],2)/hy[i-1,j] + g/2*pow(hy[i-1,j],2)) - (pow(vy[i-1,j-1],2)/hy[i-1,j-1] + g/2*pow(hy[i-1,j-1],2)))

#dh = dt/dt*(ux[1:,:]-ux[:-1,:])+ dt/dy*(vy[:,1:]-vy[:,:-1])
    reflective()
    return h,u,v
'''
for i in range (17):
    #print h
    proses(1)
'''

import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
from mpl_toolkits.mplot3d import Axes3D
a = n
x = np.arange(n+2)
y = np.arange(n+2)
x,y = np.meshgrid(x,y)

fig = plt.figure()

ax = fig.add_subplot(111, projection='3d')

def plotset():
    ax.set_xlim3d(0, a)
    ax.set_ylim3d(0, a)
    ax.set_zlim3d(0.5,1.5)
    ax.set_autoscalez_on(False)
    ax.zaxis.set_major_locator(LinearLocator(10))
    ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
    cset = ax.contour(x, y, h, zdir='x', offset=0 , cmap=cm.coolwarm)
    cset = ax.contour(x, y, h, zdir='y', offset=n , cmap=cm.coolwarm)
    cset = ax.contour(x, y, h, zdir='z', offset=.5, cmap=cm.coolwarm)

plotset()

surf = ax.plot_surface(x, y, h,rstride=1, cstride=1,cmap=cm.coolwarm,linewidth=0,antialiased=False, alpha=0.7)

fig.colorbar(surf, shrink=0.5, aspect=5)


from matplotlib import animation


def data(k,h,surf):
    proses()
    ax.clear()
    plotset()
    surf = ax.plot_surface(x, y, h,rstride=1, cstride=1,cmap=cm.coolwarm,linewidth=0,antialiased=False, alpha=0.7)
    return surf,

ani = animation.FuncAnimation(fig, data, fargs=(h,surf), interval=10, blit=False)
#ani.save('air.mp4', bitrate=512)
plt.show()

and the snapshot