
import sys; sys.path.append('../py_src')
from pyOM_gui import pyOM_gui as pyOM
from numpy import *

class kelvin_helm1(pyOM):
   
   """ vertical shear instability
   """
   
   def set_parameter(self):
     """set main parameter
     """
     M=self.fortran.pyom_module   
     fac=0.5
     M.nx=3*64*fac;
     M.nz=40*fac;
     M.ny=3
     
     M.dx=0.25/fac 
     M.dz=0.25/fac 
     M.dt=0.025/fac
     M.eps2d_sor = 1e-6
     M.eps3d_sor = 1e-6
     
     M.enable_hydrostatic          = 0
     M.enable_cyclic_x             = 1
     M.enable_cyclic_y             = 1
     
     if 1:
       M.enable_biharmonic_friction  = 1
       M.enable_vert_biha_friction   = 1
       M.enable_quicker_advection    = 1
       M.enable_quicker_mom_advection= 1
       M.enable_biharmonic_diffusion = 0
       M.enable_vert_biha_diffusion  = 0
       M.ahbi = (5e-3/fac**2)*M.dx**2
       M.khbi = M.ahbi
       M.avbi = (5e-3/fac**2)*M.dx**2
       M.kvbi = M.avbi
     else: 
       M.a_h = 20e-3/fac**2; M.k_h = M.a_h; M.a_v = M.a_h; M.k_v = M.a_h
     return

   def initial_conditions(self):
     """ setup all initial conditions
     """
     M=self.fortran.pyom_module   
     self.ini_b = zeros((M.nx,M.ny,M.nz),'d')
     self.ini_bs= zeros((M.nx,M.ny,M.nz),'d')
     self.ini_u = zeros((M.nx,M.ny,M.nz),'d')
     for k in range(M.nz):
       for j in range(M.ny):
         fxa=1e-3*M.zt[0]*sin( M.xt/(20*M.dx)*pi)
         self.ini_b[:,j,k]  = 1e-2*tanh( (    M.zt[k]-1.0*M.zt[M.nz/2] ) /M.zt[0]*100 )*M.maskt[:,j,k]
         self.ini_bs[:,j,k] = 1e-2*tanh( (fxa+M.zt[k]-1.0*M.zt[M.nz/2] ) /M.zt[0]*100 )*M.maskt[:,j,k]
         self.ini_u[:,j,k]  = ( 2+tanh( (M.zt[k]-1.0*M.zt[M.nz/2])/M.zt[0]*100) )*M.masku[:,j,k]
     for k in range(3): M.b[:,:,:,k] = self.ini_bs
     for k in range(3): M.u[:,:,:,k] = self.ini_u

     self.T_rest = ones((M.nx,M.ny,M.nz),'d') / (15.*M.dt) 
     self.T_rest[(M.nx/8-1):,:,:]=0
     return


   def restoring_zones(self):
     """  add here restoring zones
     """
     M=self.fortran.pyom_module   
     M.b[:,:,:,M.taup1-1]=M.b[:,:,:,M.taup1-1] + M.maskt*2*M.dt*self.T_rest*(self.ini_b-M.b[:,:,:,M.taum1-1])
     return

   def momentum_restoring_zones(self):
     """  add here momentum restoring zones
     """
     M=self.fortran.pyom_module   
     M.fu[:,:,:]=M.fu[:,:,:] + M.masku*self.T_rest*(self.ini_u-M.u[:,:,:,M.taum1-1])
     return

   def diagnose2(self):
     """ diagnose the model variables, could be replaced by other version
     """
     import matplotlib
     matplotlib.use('TKAgg')
     import matplotlib.pyplot as plt
     pyOM.diagnose(self)
     M=self.fortran.pyom_module         # fortran module with model variables
     if M.my_pe==0:
         i=M.nx/8-1
         plt.ion()
         plt.clf()
         plt.subplot(311)
         plt.contour(M.xt[i:],M.zt[1:-1],M.b[i:,1,1:-1,M.tau-1].transpose()*1e2,arange(-1.2,1.2,0.25) )
         plt.title('buoyancy')
         plt.xlabel('x [m]')
         plt.ylabel('z [m]')
         plt.draw()
     self.barrier()
     return
  
   def make_plot(self):
       M=self.fortran.pyom_module         # fortran module with model variables
       i=M.nx/8-1
       self.figure.clf()
       ax=self.figure.add_subplot(311)
       ax.contour(M.xt[i:],M.zt[1:-1],M.b[i:,1,1:-1,M.tau-1].transpose()*1e2,arange(-1.2,1.2,0.25) )
       ax.set_title('buoyancy')
       ax.set_xlabel('x [m]')
       ax.set_ylabel('z [m]')
       ax.axis('tight')
       return

  
if __name__ == "__main__": kelvin_helm1(snapint=5).mainloop()

    #model  = kelvin_helm1()
    #model.run(runlen = 360., snapint  = 0.25)
