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

#H0  = 200   # water depth
#N0  = 0.001 # stability freq.
#EPS = 1e-6
#U0  = 0.15  # velocity mag.


H0  = 2000   # water depth
N0  = 0.004 # stability freq.
U0  = 0.5  # velocity mag.
EPS = 1e-5

f0  = 1e-4  # coriolis freq.
beta= 2e-11 # beta 
HY  = 0.0e-3# slope of topography

class eady2(pyOM):
   """ Eady (1941) solution
   """
   def set_parameter(self):
     """set main parameter
     """
     M=self.fortran.pyom_module   
     M.nx    = 32
     M.nz    = 15
     M.ny    = M.nx
     M.dz    = H0/(M.nz-2)
     M.eps2d_sor = EPS
     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

     # print some numbers
     Lr    = N0*H0/f0          # Rossby radius
     delta = H0/Lr             # aspect ratio
     Ro    = U0/(f0*Lr)       # Rossby number
     Ri    = N0**2*H0**2/U0**2 # Richardson number
     print
     print  ' L  = %f km'%(Lr/1e3)
     print  ' Ro = %f '%Ro
     print  ' Ri = %f '%Ri
     print  ' delta = %f '%delta
     print  ' ell = %f '%(Lr/6400e3)
     print
     
     # solve linear stability problem first
     #ksx=linspace(0,3.2,10);    kx=ksx/Lr
     #ksy=linspace(-3.2,3.2,20); ky=ksy/Lr
     ksx=linspace(0,3.2,40);    kx=ksx/Lr
     ky = array( [0./Lr] )
     zw=arange(M.nz)*M.dz
     zt=zw-M.dz/2.0
     U=U0/2+U0*zt[1:-1]/H0
     V=U*0
     B=N0**2*zt[1:-1]
     if 1:
      om_max,om,kmax,lmax,u,v,w,b,p=lin_stab.qg(U,V,B,M.dz,kx,ky,beta,f0,0,HY)
      print ' Max. growth rate QG %f 1/days ' % (-imag(om)*86400)
      print ' k_max = %f Lr , l_max = %f Lr' % (kmax*Lr,lmax*Lr)
     if 1:
      om_max,om,kmax,lmax,u,v,w,b,p=lin_stab.pe(U,V,B,M.dz,\
                                                kx,ky,0.,beta,0.,f0,0.,HY)
      print ' Max. growth rate PE %f 1/days ' % (-imag(om)*86400)
      print ' k_max = %f Lr , l_max = %f Lr' % (kmax*Lr,lmax*Lr)
     print
     
     self.lin_stab_kmax = kmax
     self.lin_stab_b = b
     self.lin_stab_p = p
     self.lin_stab_u = u
     self.lin_stab_v = v
     self.lin_stab_w = w
     L = 2*2*pi/kmax   
     M.dx  =   L/(M.nx-2) 
     M.dt    = M.dx*1200/20e3   # c = 20e3/1200.0,  dt =dx/c
     print "dx=%f km, dt= %f s "%(M.dx/1e3,M.dt)
     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   
     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   
 
     # background velocity
     for k in range(M.nz):
        M.u0[:,:,k]=(U0/2+U0*M.zt[k]/H0)*M.masku[:,:,k]

     # background buoyancy
     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]-N0**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(1,M.nz-1):
       ky=pi/((M.ny-2)*M.dx)
       kx = 2.0*self.lin_stab_kmax
       phase = cos(kx*M.xt) +complex(0,1)*sin(kx*M.xt)
       phase = phase*sin(ky*M.yt[j])
       M.b[:,j,k,0] =0.5*real(phase*self.lin_stab_b[k-1])*M.maskt[:,j,k]
       #M.u[:,j,k,0] =real(phase*self.lin_stab_u[k-1])*M.masku[:,j,k]
       #M.v[:,j,k,0] =real(phase*self.lin_stab_v[k-1])*M.maskv[:,j,k]
       #M.w[:,j,k,0] =real(phase*self.lin_stab_w[k-1])*M.maskw[:,j,k]
       #M.p_full[:,j,k,0] =real(phase*self.lin_stab_p[k-1])*M.maskt[:,j,k]

     for a in M.b,M.u,M.v,M.w,M.p_full:
       a[:,:,:,1] = a[:,:,:,0]    
       a[:,:,:,2] = a[:,:,:,0]
     return


   def topography(self):
     """ setup topography
     """
     M=self.fortran.pyom_module
     for j in range(M.ny): 
        z0 = M.yt[j]*HY
        for k in range(M.nz):
          if z0 > M.zt[k]-M.zt[0] :
            M.maskt[:,j,k]=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
       k2=M.nz/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,1:-1,M.tau-1] + M.back[1:-1,1:-1,1:-1,M.tau-1]
       b=M.maskt[1:-1,1:-1,1:-1]
       a=sum(a*b,axis=0)/sum(b,axis=0)
       co=ax.contourf(y,z,a.transpose())
       a=M.u[1:-1,1:-1,1:-1,M.tau-1] + M.u0[1:-1,1:-1,1:-1]
       b=M.masku[1:-1,1:-1,1:-1]
       a=sum(a*b,axis=0)/sum(b,axis=0)
       ax.contour(y,z,a.transpose(),10,colors='black')
       ax.set_ylabel('y [km]')

       self.figure.colorbar(co)
       ax.set_title("total buoyancy at x=%4.1f km"%x[i]);  
       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 perturb. at z=%4.1f m"%z[k]); 
       self.figure.colorbar(co)
       ax.axis('tight')
       
       ax=self.figure.add_subplot(223)
       if 1:
        co=ax.contourf(x,y,M.b[1:-1,1:-1,k2,M.tau-1].transpose()*1e2)
        a=M.b[1:-1,1:-1,k2,M.tau-1] + M.back[1:-1,1:-1,k2,M.tau-1]
        ax.contour(x,y,a.transpose(),10,colors='black')
        ax.set_title("buoyancy perturb. at z=%4.1f m"%z[k2]); 
        a=M.u[1:-1:2,1:-1:2,k,M.tau-1] + M.u0[1:-1:2,1:-1:2,k]
        b=M.v[1:-1:2,1:-1:2,k,M.tau-1] 
        ax.quiver(x[::2],y[::2],a.transpose(),b.transpose() )
        self.figure.colorbar(co)
        ax.axis('tight')
          
       if 0:
        #co=ax.contourf(x,y,M.w[1:-1,1:-1,k,M.tau-1].transpose())
        co=ax.contourf(x,y,M.p_full[1:-1,1:-1,k,M.tau-1].transpose())
        self.figure.colorbar(co)
        a=M.u[1:-1:2,1:-1:2,k,M.tau-1] + M.u0[1:-1:2,1:-1:2,k]
        b=M.v[1:-1:2,1:-1:2,k,M.tau-1] 
        ax.quiver(x[::2],y[::2],a.transpose(),b.transpose() )
        #a=M.p_full[1:-1,1:-1,k,M.tau-1] 
        #ax.contour(x,y,a.transpose(),10,colors='black')
        ax.axis('tight')
        ax.set_title("pressure pertub. at z=%4.1f m"%z[k]);
        
       if 0:
        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 perturb. at y=%4.1f km"%y[j]); 
       ax.set_xlabel('x [km]'); 
       self.figure.colorbar(co)
       ax.axis('tight')
     return

if __name__ == "__main__": eady2(snapint  = 3*24).mainloop()


