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


class rayleigh(pyOM):
    
   def set_parameter(self):
     """set main parameter
     """
     M=self.fortran.pyom_module   
     fac=2.0
     M.nx=2*32*fac;
     M.nz=20*fac;
     M.ny=3
     M.dx=0.5/fac; 
     M.dz=0.5/fac; 
     M.dt=0.25/fac
     M.eps2d_sor = 1e-6
     M.eps3d_sor = 1e-6
      
     M.enable_hydrostatic           = 0
     M.enable_quicker_advection     = 0 
     M.enable_quicker_mom_advection = 0
     M.enable_biharmonic_friction   = 1 
     M.enable_vert_biha_friction    = 1
     M.enable_biharmonic_diffusion  = 1
     M.enable_vert_biha_diffusion   = 1
     M.enable_cyclic_y              = 1
     M.enable_cyclic_x              = 0
      
     M.ahbi = (10e-3/fac**2)*M.dx**2
     M.khbi = M.ahbi
     M.avbi = (10e-3/fac**2)*M.dx**2
     M.kvbi = M.avbi
     return

   
   def initial_conditions(self):
     """ setup all initial conditions
     """
     M=self.fortran.pyom_module   
     for k in range(M.nx): M.b[k,:,:,:] = 4e-5*sin(M.xt[k]/(20*M.dx)*pi)
     for k in [0,1,2]:     M.b[:,:,:,k] = M.b[:,:,:,k]*M.maskt
     return
 
   def boundary_conditions(self):
     """ setup all forcing
         surface and bottom boundary conditions
         might be variable in time, called every time step
     """
     M=self.fortran.pyom_module   
     rho_0=1024.0; g=9.801
     M.surface_flux[:,:]= 175 *0.2/rho_0*g /4185.5 
     M.bottom_flux[:,:] = 175 *0.2/rho_0*g /4185.5 
     return


   def make_plot(self):
     """ diagnose the model variables, could be replaced by other version
     """
     M=self.fortran.pyom_module         # fortran module with model variables
     self.figure.clf()
     ax=self.figure.add_subplot(211)
     ax.contour(M.xt[1:-1],M.zt[1:-1],M.b[1:-1,1,1:-1,M.tau-1].transpose()*1e2,arange(-1.2,1.2,0.1))
     ax.set_title('buoyancy')
     ax.set_xlabel('x [m]')
     ax.set_ylabel('z [m]')
     ax.axis('tight')
     return

if __name__ == "__main__": rayleigh(snapint=20).mainloop()

