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


RESOLVE = 0.125
#RESOLVE = 1.

class acc1(pyOM):
   """ a simple global model with a Southern Ocean and Atlantic part
   """
   def set_parameter(self):
     """set main parameter
     """
     M=self.fortran.pyom_module   
     M.nx = 128*RESOLVE 
     M.nz = 20# *RESOLVE;
     M.ny = 128*RESOLVE
     M.dx = 20e3/RESOLVE
     M.dz = 50.0# /RESOLVE
     M.dt = 1200.0/RESOLVE
     M.k_h=0000.0 /RESOLVE
     M.a_h=10000.0 #/RESOLVE
     M.k_gm = 2000.0
     M.ahbi  = 5e11/RESOLVE**2
     M.khbi  = M.ahbi
     M.k_v    = 1.0e-4#/RESOLVE
     M.a_v    = 10.e-4#/RESOLVE
     M.cdbot =  1e-5#*RESOLVE
     M.eps2d_sor = 1e-9
      
     M.enable_hydrostatic          = 1
     M.enable_noslip               = 1
     M.enable_cyclic_x             = 1
     M.enable_bottom_stress        = 0
     M.enable_bottom_noslip        = 1
     M.enable_biharmonic_friction  = 0
     M.enable_biharmonic_diffusion = 0
     M.enable_quicker_advection    = 1  
     M.enable_free_surface         = 0
     M.enable_expl_free_surf       = 0
     M.enable_vert_friction_trm    = 1

     M.enable_diag_tracer = 0
     M.nt=1
     M.delimit_tracer_fluxes      = 1
     M.enable_isopycnal_diffusion = 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   
     rad  = 6370.0e3   # earth radius
     om   = pi/43082.0 # earth rotation freq. in 1/s
     phi0 = -25 /180. *pi
     M.coriolis_t[:]   =  2*om*sin(phi0)+2*om*cos(phi0)/rad*M.yt[:]
     M.coriolis_hor[:] =  0.
     return

   def b_surface(self):
     M=self.fortran.pyom_module   
     L_y = M.yt[-1]-M.yt[0]
     db = -30e-3
     bstar=db*ones((M.ny,),'d')
     bstar[:M.ny/2] = db*M.yt[:M.ny/2]/(L_y/2.)
     y2=L_y*0.75
     bstar[M.ny/4*3:] = db*( 1- (M.yt[M.ny/4*3:]-y2)/(M.yt[-1]-y2))
     return bstar

   def topography(self):
     """ setup topography
     """
     M=self.fortran.pyom_module   
     Lx = M.xt[-1]-M.xt[0]
     Ly = M.yt[-1]-M.yt[0]
     for j in range(M.ny):
       for i in range(M.nx):
         if M.yt[j]>Ly/2.0 and (M.xt[i]>Lx*.75 or M.xt[i]<Lx*.25) :
            M.maskt[i,j,:]=0
     return

   def initial_conditions(self):
     """ setup initial conditions
     """
     M=self.fortran.pyom_module   
     N0=0.004
     bstar = self.b_surface()
     for j in range(M.ny):
       for k in range(M.nz):
         M.b[:,j,k,0] = (-30e-3-N0**2*M.zt[k] )*M.maskt[:,j,k]
         M.b[:,j,k,1] = (-30e-3-N0**2*M.zt[k] )*M.maskt[:,j,k]
         M.b[:,j,k,2] = (-30e-3-N0**2*M.zt[k] )*M.maskt[:,j,k]

   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   
     L_y = M.yt[-1]-M.yt[0]
     t_rest=30*86400
     bstar = self.b_surface()
     for j in range(M.ny/2):
       y = M.yu[j]/L_y
       M.surface_taux[:,j] = 0.1e-3*sin(2*pi*y)*M.masku[:,j,-2]
     for j in range(M.ny):
       M.surface_flux[:,j]=M.dz/t_rest*(bstar[j]-M.b[:,j,-2,M.tau-1])*M.maskt[:,j,-2]
     return

  
   def tracer_sources(self):
     """ setup tracer source. Here a simple age tracer
     """
     M=self.fortran.pyom_module   
     M.tr[:,:,:,M.taup1-1,0]=M.tr[:,:,:,M.taup1-1,0]+M.maskt[:,:,:]*2*M.dt*(1.0/86400.)
     k=M.nz-2
     M.tr[:,:,k,M.taup1-1,0]=M.tr[:,:,k,M.taup1-1,0]+M.maskt[:,:,k]*2*M.dt*1./(1*86400.)*(0-M.tr[:,:,k,M.taum1-1,0] )
     return



   def make_plot(self):
     """ diagnose the model variables, could be replaced by other version
     """
     if hasattr(self,'figure'):
         M=self.fortran.pyom_module     # fortran module with model variables
         k=M.nz*3/4
         i=int(M.nx/2)
         j=int(3*M.ny/4)
        
         fxa = -1024/9.81 /0.2 
         self.figure.clf()
         ax=self.figure.add_subplot(221)
         a=M.b[1:-1,1:-1,k,M.tau-1]*fxa 
         a=where( M.maskt[1:-1,1:-1,k] >0, a, NaN)
         co=ax.contourf(M.xt[1:-1]/1e3,M.yt[1:-1]/1e3,a.transpose())
         ax.set_title('buoyancy')
         ax.set_ylabel('y [km]')
         ax.axis('tight') 
         self.figure.colorbar(co,format='%.1f')
         
         ax=self.figure.add_subplot(222)
         
         a=M.b[1:-1,1:-1,1:-1,M.tau-1]*fxa 
         b=M.maskt[1:-1,1:-1,1:-1]
         a=sum(a*b,axis=0)/sum(b,axis=0)
         co=ax.contourf(M.yt[1:-1]/1e3,M.zt[1:-1],a.transpose())

         a=M.v[1:-1,1:-1,1:-1,M.tau-1]*M.maskv[1:-1,1:-1,1:-1]
         b=M.maskv[1:-1,1:-1,1:-1]
         a=sum(a,axis=0)*M.dx
         a=cumsum(a,axis=1)*M.dz /1e6
         try: 
           ax.contour(M.yt[1:-1]/1e3,M.zt[1:-1],a.transpose(),10,colors='black')
         except: 
           pass
         ax.set_title('MOC')
         ax.axis('tight') 
         self.figure.colorbar(co,format='%.1f')
         
         ax=self.figure.add_subplot(223)
         a=M.u[1:-1,1:-1,1:-1,M.tau-1]*M.masku[1:-1,1:-1,1:-1]
         b=M.masku[1:-1,1:-1,1:-1]
         a=sum(a,axis=0)/sum(b,axis=0)
         co=ax.contourf(M.yt[1:-1]/1e3,M.zt[1:-1],a.transpose())
         ax.set_title('zonal velocity')
         ax.axis('tight') 
         ax.set_xlabel('x [km]')
         self.figure.colorbar(co,format='%.2f')
         
         if M.enable_diag_tracer :
           ax=self.figure.add_subplot(224)
           a=M.tr[1:-1,1:-1,1:-1,M.tau-1,0]*M.maskt[1:-1,1:-1,1:-1]
           b=M.maskt[1:-1,1:-1,1:-1]
           a=sum(a,axis=0)/sum(b,axis=0)
           co=ax.contourf(M.yt[1:-1]/1e3,M.zt[1:-1],a.transpose())
           ax.set_title('age tracer')
           ax.axis('tight') 
           ax.set_xlabel('x [km]')
           self.figure.colorbar(co)
     return

if __name__ == "__main__": acc1(snapint  = 24 ).mainloop()
