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

BETA=2e-11

class enso1(pyOM):
   """ Enso response
   """
   def set_parameter(self):
     """set main parameter
     """
     M=self.fortran.pyom_module   
     M.nx    = 64
     M.nz    = 3
     M.ny    = 64
     M.dx    = 200e3
     M.dz = (500e3**2*BETA) **2 /9.81
     M.dt    = 3600.0 /2.
     M.eps2d_sor = 1e-22
     M.enable_hydrostatic          = 1
     M.enable_cyclic_x             = 0
     M.enable_quicker_advection    = 1
     M.enable_quicker_mom_advection= 1
     M.enable_free_surface         = 1
     return

   def set_coriolis(self):
     """ vertical and horizontal Coriolis parameter on yt grid
         routine is called after initialization of grid
     """
     M=self.fortran.pyom_module   
     f0=-M.ny*M.dx/2.0*BETA   # equatorial beta plane
     M.coriolis_t[:]   =  f0+BETA*M.yt[:]
     M.coriolis_hor[:] =  0.
     return



   def initial_conditions(self):
     """ setup all initial conditions
     """
     M=self.fortran.pyom_module   
     cn =  (M.dz*9.81)**0.5  
     hn=cn**2/9.81
     Re = (cn /BETA)**0.5  
     y0=M.ny*M.dx*0.5
     g=9.81
     A = .1
     for i in range(M.nx):
       for j in range(M.ny):
         M.eta[i,j,:]=0.1*exp( -(M.xt[i]-y0*0.3)**2/(2*Re)**2 \
                               -(M.yt[j]-y0)**2/(2*Re)**2 )
     return

   def make_plot(self):
     """ make a plot using methods of self.figure
     """
     if hasattr(self,'figure'):
       M=self.fortran.pyom_module         # fortran module with model variables
       x=M.xt[1:-1]/1e3
       y=M.yt[1:-1]/1e3

       self.figure.clf()
       ax=self.figure.add_subplot(111)
       a=M.eta[1:-1,1:-1,M.tau-1] 
       co=ax.contourf(x,y,a.transpose(),15)
       ax.axis('tight') 
       
       a=M.u[1:-1:2,1:-1:2,1,M.tau-1] 
       b=M.v[1:-1:2,1:-1:2,1,M.tau-1] 
       ax.quiver(x[::2],y[::2],a.transpose(),b.transpose() )
       self.figure.colorbar(co)
     return

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