

GUI = False

import sys; sys.path.append('../py_src')
if GUI:
   from pyOM_gui import pyOM_gui as pyOM
else:
   from pyOM_cdf import pyOM_cdf as pyOM
from numpy import *

class eady1(pyOM):
   """ Eady (1941) solution
   """
   def set_parameter(self):
     """set main parameter
     """
     M=self.fortran.pyom_module   
     M.nx    = 32
     M.nz    = 20
     M.ny    = 32
     M.dx    = 20e3
     M.dz    = 100.0 
     M.dt    = 1200.0 
     M.eps2d_sor = 1e-6
     M.a_v   = 10.e-4 
     M.k_v   = 1.0e-4 
     M.enable_hydrostatic           = 1
     M.enable_cyclic_x              = 1
     M.enable_quicker_advection     = 1
     M.enable_quicker_mom_advection = 1
     M.enable_back_state            = 1
     M.enable_back_zonal_flow       = 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 = 1e-4; beta=0e-11
     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   
     U_0 = 0.5
     N_0 = 0.004
     f=M.coriolis_t[M.ny/2]
     h = (M.nz-2)*M.dz
     kx=1.6*f/(N_0*h)
     ky=pi/((M.ny-2)*M.dx)
     d=f/N_0/(kx**2+ky**2)**0.5

     fxa=(exp(h/d)+exp(-h/d))/(exp(h/d)-exp(-h/d))
     c1= (1+0.25*(h/d)**2-h/d*fxa )*complex(1,0)
     c1=(sqrt(c1)*d/h+0.5)*U_0
     A=(U_0-c1)/U_0*h/d
     
     # background velocity
     for k in range(M.nz):
        M.u0[:,:,k]=(U_0/2+U_0*M.zt[k]/(M.nz*M.dz))*M.masku[:,:,k]
        #M.u0[:,:,k]=(U_0/2+U_0*M.zt[k]/(M.nz*M.dz))*M.masku[:,:,k]\
        #             *exp( -(M.yt-M.yt[M.ny/2])**2/200e3**2)
     # background buoyancy
     #for k in range(M.nz):
     #  for j in range(M.ny-1):
     #   fxa = (M.coriolis_t[j]+M.coriolis_t[j+1])/2.0
     #   M.back[:,j+1,k,0]=M.back[:,j,k,0]+\
     #                     M.dx*M.maskt[:,j,k]*U_0/(M.nz*M.dz)*fxa
     
     for k in range(1,M.nz-1):
       for j in range(M.ny-1):
        fxa = (M.coriolis_t[j]+M.coriolis_t[j+1])/2.0
        uz1 = (M.u0[:,j,k+1]-M.u0[:,j,k])/M.dz*M.masku[:,j,k+1]
        uz2 = (M.u0[:,j,k]-M.u0[:,j,k-1])/M.dz*M.masku[:,j,k-1]
        uz3 = (M.u0[:,j+1,k+1]-M.u0[:,j+1,k])/M.dz*M.masku[:,j+1,k+1]
        uz4 = (M.u0[:,j+1,k]-M.u0[:,j+1,k-1])/M.dz*M.masku[:,j+1,k-1]
        uz = (uz1+uz2+uz3+uz4)/(M.masku[:,j,k+1]+M.masku[:,j,k-1] \
                          +M.masku[:,j+1,k+1]+M.masku[:,j+1,k-1]+1e-12)
        M.back[:,j+1,k,0]=M.back[:,j,k,0]+M.dx*M.maskt[:,j,k]*uz*fxa

     for k in range(M.nz):
        M.back[:,:,k,0]=(M.back[:,:,k,0]-N_0**2*M.zt[k])*M.maskt[:,:,k]
     M.back[:,:,:,1] = M.back[:,:,:,0]    
     M.back[:,:,:,2] = M.back[:,:,:,0]

     # perturbation buoyancy
     for j in range(M.ny):
      for k in range(M.nz):
       phiz=A/d*sinh(M.zt[k]/d)+cosh(M.zt[k]/d)/d
       M.b[:,j,k,0] =0.1*sin(kx*M.xt)*sin(ky*M.yt[j])*abs(phiz)*M.maskt[:,j,k]
       #M.b[:,j,k,0] =1e-5*random.randn(M.nx)*M.maskt[:,j,k]

       
     M.b[:,:,:,1] = M.b[:,:,:,0]    
     M.b[:,:,:,2] = M.b[:,:,:,0]
     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
       k=M.nz*3/4
       i=int(M.nx/2)
       j=int(M.nx/2)
       x=M.xt[1:-1]/1e3
       y=M.yt[1:-1]/1e3
       z=M.zt[1:-1]

       self.figure.clf()
       ax=self.figure.add_subplot(221)
       #a=M.b[1:-1,1:-1,k,M.tau-1] + M.back[1:-1,1:-1,k,M.tau-1]
       #co=ax.contourf(x,y,a.transpose()*1e2)
       #ax.set_ylabel('y [km]')
       a=M.b[i,1:-1,1:-1,M.tau-1] + M.back[i,1:-1,1:-1,M.tau-1]
       co=ax.contourf(y,z,a.transpose())
       a=M.u[i,1:-1,1:-1,M.tau-1] + M.u0[i,1:-1,1:-1]
       ax.contour(y,z,a.transpose(),10,colors='black')
       ax.set_ylabel('y [km]')

       self.figure.colorbar(co)
       ax.set_title('total buoyancy');  
       ax.axis('tight') 
         
       ax=self.figure.add_subplot(222)
       co=ax.contourf(x,y,M.b[1:-1,1:-1,k,M.tau-1].transpose()*1e2)
       a=M.b[1:-1,1:-1,k,M.tau-1] + M.back[1:-1,1:-1,k,M.tau-1]
       ax.contour(x,y,a.transpose(),10,colors='black')
       ax.set_title('buoyancy perturbation'); 
       self.figure.colorbar(co)
       ax.axis('tight')
       
       ax=self.figure.add_subplot(223)
       #a=M.b[i,1:-1,1:-1,M.tau-1] + M.back[i,1:-1,1:-1,M.tau-1]
       a=M.p_full[i,1:-1,1:-1,M.tau-1] 
       co=ax.contourf(y,z,a.transpose())
       self.figure.colorbar(co)
       a=M.u[i,1:-1,1:-1,M.tau-1] #+ M.u0[i,1:-1,1:-1]
       if a[0,0] != a.all():
           ax.contour(y,z,a.transpose(),10,colors='k')
       #ax.set_title('total buoyancy'); 
       ax.set_title('pressure perturbation'); 
       ax.set_xlabel('y [km]'); ax.set_ylabel('z [m]')
       ax.axis('tight')
       
       ax=self.figure.add_subplot(224)
       a=M.p_full[1:-1,j,1:-1,M.tau-1] 
       co=ax.contourf(x,z,a.transpose())
       ax.set_title('pressure perturbation'); 
       ax.set_xlabel('x [km]'); 
       self.figure.colorbar(co)
       ax.axis('tight')
     return

if __name__ == "__main__": 
 
  if GUI:
     eady1(snapint  = 3*24).mainloop()
  else:
     model=eady1()
     model.run(snapint = 3*24, runlen = 30*86400 )




