# Plot sum of cosines to make square wave
# Try changing n to see what happens with wave
n = 5
k = 1.0
x = linspace(0,10,1000)
# fac allows us to calculate factors of -1 or +1 easily
fac = 1
# Set initial value
standingwave = 0.5 + 2*fac*cos(k*x)/pi
# Iterate over odd numbers (even numbers have zero coefficients)
for i in range(3,int(n+1),2):
    fac *= -1
    standingwave += (2*fac*cos(i*k*x))/(i*pi)
plot(x,standingwave)

# Show 3D plot ?
# Import necessary 3D plotting
from mpl_toolkits.mplot3d.axes3d import Axes3D
# Create figure
fig = plt.figure(figsize=(9,6))
axm = fig.add_subplot(1, 1, 1, projection='3d')
axm.set_xlabel('x')
axm.set_ylabel('t')
axm.set_zlabel(r'$\psi$')
# Specify x and t and put onto mesh
phix = linspace(0,4*pi,100)
phit = linspace(0,4*pi,100)
x,t = meshgrid(phix, phit)
# Choose k and omega
k = 0.5
omega = 0.5
Gamma = -0.5
c = omega/k
# Create variable
zm = exp(0.5*Gamma*x/c)*cos(k*x - omega*t)
# Plot surface
pm = axm.plot_surface(x, t, zm, linewidth=0,rstride=2,cstride=2,cmap=cm.coolwarm)
# Set angles
axm.view_init(50, 30)
# Add colour bar
fig.colorbar(pm, shrink=0.5)


# Plot water dispersion
g = 9.8 # acceleration due to gravity
rho = 1000 # density
gamma = 0.073 # surface tension
# Ripples are found for lambda < 0.017m (surface tension dominated)
# Shallow water is kh << 1
# Deep water is kh >> 1
fig = figure(figsize=[12,3])
# Plot vs x
ax_shallow = fig.add_subplot(121)
ax_deep = fig.add_subplot(122)
# Set h
h = 10.0
k_shallow = linspace(0,1/h,1000)
k_deep = linspace(1/h,100*h,1000)
#ax.axis([0,4,-1.5*A,1.5*A])
print "Depth, h, is ",h
# Plot shallow water wave dispersion
omega_shallow = sqrt((g*k_shallow + gamma*k_shallow*k_shallow*k_shallow/rho)*tanh(k_shallow*h))
ax_shallow.plot(k_shallow,omega_shallow)
# Limiting behaviour 
ax_shallow.plot(k_shallow,sqrt(g*h)*k_shallow)
# Plot deep water wave dispersion
omega_deep = sqrt((g*k_deep + gamma*k_deep*k_deep*k_deep/rho)*tanh(k_deep*h))
ax_deep.plot(k_deep,omega_deep)
# Plot limiting behaviour: ripples (lambda<0.017 i.e. large k) SURFACE TENSION
ax_deep.plot(k_deep,sqrt(gamma/rho)*k_deep**1.5)
# Long wavelength (small k) behaviour GRAVITY
ax_deep.plot(k_deep,sqrt(g*k_deep))

# Plot motion of water wave particles for deep water
k = 1.0
omega = 1.0
t = linspace(0,2*pi/k,1000)
x = 0.0
for i in range(5):
    y=-0.2*i
    psi_x_d = exp(k*y)*sin(omega*t)
    psi_y_d = exp(k*y)*cos(omega*t)
    plot(psi_x_d,psi_y_d+y)
axis('scaled')

# Plot motion of water wave particles for shallow water
k = 2*pi/10.0 # Set k by a nominal wavelength
omega = 1.0
h = 1.0 # Depth in m
if k*h < 1:
    t = linspace(0,2*pi/k,1000)
    x = 0.0
    for i in range(5):
        y=-0.2*i
        psi_x_d = sin(omega*t)/(k*h)
        psi_y_d = (1+y/h)*cos(omega*t)
        plot(psi_x_d,psi_y_d+y)
    axis('scaled')
    axis([-2,2,-2,2])
else:
    print 'kh is not less than one - this is deep water ! ',k*h