1 import sys; sys.path.append('../py_src')
   2 
   3 from pyOM_gui import pyOM_gui as pyOM
   4 #from pyOM_cdf import pyOM_cdf as pyOM
   5 from numpy import *
   6 
   7 
   8 fac = 0.5
   9 mix = 2e-5/fac**2
  10 
  11 class kelv2(pyOM):
  12    """ 
  13    """
  14    def set_parameter(self):
  15      """set main parameter
  16      """
  17      M=self.fortran.main_module   
  18 
  19      (M.nx,M.ny,M.nz)=(int(64),  1, int(40) )
  20      M.dt_mom     = 0.05/fac
  21      M.dt_tracer  = 0.05/fac
  22 
  23      M.enable_conserve_energy = 0
  24      M.coord_degree           = 0
  25      M.enable_cyclic_x        = 1
  26      M.enable_hydrostatic     = 0
  27      M.eq_of_state_type       = 1 
  28 
  29      M.congr_epsilon = 1e-12
  30      M.congr_max_iterations = 5000
  31      M.congr_epsilon_non_hydro=   1e-8
  32      M.congr_max_itts_non_hydro = 5000    
  33 
  34      M.enable_explicit_vert_friction = 1
  35      M.kappam_0 = mix/fac**2
  36      M.enable_hor_friction = 1
  37      M.a_h = mix/fac**2
  38 
  39      M.enable_superbee_advection = 1
  40      return
  41 
  42    def set_mean_state1(self,dz):
  43      # two vorticity jumps and one density jump
  44      M=self.fortran.main_module   
  45      zw=arange(M.nz)*dz+ dz
  46      zt=zw-dz/2.0
  47      zt = zt - zw[-1]
  48      zw = zw - zw[-1]
  49      alpha = self.fortran.linear_eq_of_state.linear_eq_of_state_drhodt()
  50      u=1
  51      z1=-15.0
  52      z2=-5.0
  53      S=2*u/(z2-z1)
  54      self.U=0*zt-u
  55      for k in range(M.nz):
  56          if   zt[k]>z1 and zt[k]< z2: self.U[k]=S*(zt[k]-z1)-u
  57          elif zt[k]>=z2:              self.U[k]=u
  58      T0 = 1.0
  59      z3=-10.0
  60      self.T=0*zt
  61      for k in range(M.nz):
  62            if zt[k]>z3: self.T[k]=T0
  63      self.B=-self.T*alpha*9.81/1024.
  64      return
  65 
  66 
  67    def set_mean_state2(self,dz):
  68      # two layer flow
  69      M=self.fortran.main_module   
  70      zw=arange(M.nz)*dz+ dz
  71      zt=zw-dz/2.0
  72      zt = zt - zw[-1]
  73      zw = zw - zw[-1]
  74      alpha = self.fortran.linear_eq_of_state.linear_eq_of_state_drhodt()
  75      self.T=-(9.85-6.5*tanh( (zt-zt[M.nz/2-1] ) /zt[0]*10 ))
  76      self.B=-self.T*alpha*9.81/1024.
  77      self.U=0.6+0.5*tanh( (zt-zt[M.nz/2-1])/zt[0]*10)
  78      return 
  79  
  80    def calc_lin_stab(self,dz):
  81      M=self.fortran.main_module   
  82      alpha = self.fortran.linear_eq_of_state.linear_eq_of_state_drhodt()
  83      
  84      self.set_mean_state1(dz)
  85      #self.set_mean_state2(dz)
  86      
  87      k=linspace(0,1.0,50);   
  88      om_max,om,kmax,u,v,w,b,p=pe(self.U,self.B,dz,k,0,0)
  89      print ' Max. growth rate %f 1/s ' % (-imag(om))
  90      print ' k_max = %f ' % (kmax,)
  91      self.kmax = kmax
  92      self.u_pert = 5e-2*real(u)
  93      self.w_pert = 5e-2*real(w)
  94      self.T_pert = -5e-2*real(b)/alpha*1024/9.81
  95      return 
  96  
  97 
  98    def set_grid(self):
  99      M=self.fortran.main_module   
 100      M.dzt[:]=0.25/fac
 101      self.calc_lin_stab(0.25/fac)
 102 
 103      L = 2*2*pi/self.kmax    # two waves
 104      dx  =   L/M.nx 
 105      M.dt_tracer  = dx*0.02/0.25   # c = 0.25/0.05  dt =dx/c
 106      M.dt_mom     = dx*0.02/0.25
 107      print "dx=%f m, dt= %f s "%(dx,M.dt_tracer)
 108      
 109      M.dxt[:]=dx
 110      M.dyt[:]=dx
 111      return
 112 
 113          
 114    def set_initial_conditions(self):
 115      """ setup initial conditions
 116      """
 117      M=self.fortran.main_module
 118      for k in range(M.nz):
 119        for i in range(M.xt.shape[0]):
 120         M.temp[i,:,k,:] = self.T[k]+self.T_pert[k]*sin( self.kmax*M.xt[i] ) 
 121         M.u[i,:,k,:]    = self.U[k]+self.u_pert[k]*sin( self.kmax*M.xt[i] ) 
 122         M.w[i,:,k,:]    =           self.w_pert[k]*sin( self.kmax*M.xt[i] ) 
 123      return
 124 
 125 
 126    def user_defined_signal(self):
 127        """ this routine must be called by all processors
 128        """
 129        M=self.fortran.main_module  
 130        a = zeros( (M.nx,M.ny), 'd', order = 'F')
 131        a[M.is_pe-1:M.ie_pe,0] = M.xt[2:-2]
 132        self.fortran.pe0_recv_2d(a)
 133        self.xt_gl = a[:,0].copy()
 134        
 135        self.temp_gl = zeros( (M.nx,M.ny,M.nz), 'd', order = 'F')
 136        self.u_gl = zeros( (M.nx,M.ny,M.nz), 'd', order = 'F')
 137        self.w_gl = zeros( (M.nx,M.ny,M.nz), 'd', order = 'F')
 138        for k in range(M.nz):
 139          a[M.is_pe-1:M.ie_pe,M.js_pe-1:M.je_pe] = where( M.maskt[2:-2,2:-2,k] >0,  M.temp[2:-2,2:-2,k,M.tau-1] , NaN) 
 140          self.fortran.pe0_recv_2d(a)
 141          self.temp_gl[:,:,k]=a.copy()
 142 
 143          a[M.is_pe-1:M.ie_pe,M.js_pe-1:M.je_pe] = where( M.masku[2:-2,2:-2,k] >0,  M.u[2:-2,2:-2,k,M.tau-1] , NaN) 
 144          self.fortran.pe0_recv_2d(a)
 145          self.u_gl[:,:,k]=a.copy()
 146 
 147          a[M.is_pe-1:M.ie_pe,M.js_pe-1:M.je_pe] = where( M.maskw[2:-2,2:-2,k] >0,  M.w[2:-2,2:-2,k,M.tau-1] , NaN) 
 148          self.fortran.pe0_recv_2d(a)
 149          self.w_gl[:,:,k]=a.copy()
 150 
 151        return
 152 
 153    def make_plot(self):
 154        M=self.fortran.main_module
 155        
 156        self.set_signal('user_defined') # following routine is called by all PEs
 157        self.user_defined_signal()
 158        
 159        self.figure.clf()
 160        ax=self.figure.add_subplot(211)
 161        
 162        co=ax.contourf(self.xt_gl,M.zt, self.temp_gl[:,0,:].transpose())
 163        self.figure.colorbar(co)
 164        ax.quiver(self.xt_gl[::2],M.zt[::2],self.u_gl[::2,0,::2].transpose(),self.w_gl[::2,0,::2].transpose() )
 165        ax.set_title('Temperature [deg C]')
 166        ax.set_ylabel('z [m]')
 167        #ax.axis('tight')
 168        
 169        ax=self.figure.add_subplot(212)
 170        co=ax.contourf(self.xt_gl,M.zt, self.w_gl[:,0,:].transpose())
 171        self.figure.colorbar(co)
 172        ax.set_title('vertical velocity [m/s]')
 173        ax.set_xlabel('x [m]')
 174        return
 175 
 176 
 177 
 178 def pe(U,B0,dz,kx,fh,f0):
 179   # solution for primitive equations
 180   import numpy as np
 181   import sys
 182   cs = 150 # speed of sound, artificially reduced
 183   N=U.shape[0]
 184   # derivatives of U,V and B
 185   UZ=np.zeros(N,'d');BZ=np.zeros(N,'d')
 186   for n in range(1,N-1):
 187    UZ[n]=(U[n+1]-U[n-1])/(2*dz)
 188    BZ[n]=(B0[n+1]-B0[n-1])/(2*dz)
 189   UZ[0]=UZ[1];UZ[-1]=UZ[-2];
 190   BZ[0]=BZ[1];BZ[-1]=BZ[-2];
 191   # allocate some variables
 192   I=complex(0,1)
 193   A  = np.zeros((5,5,N),'Complex64')
 194   B  = np.zeros((5,5,N),'Complex64')
 195   C  = np.zeros((5,5,N),'Complex64')
 196   AA = np.zeros((5*N,5*N),'Complex64')
 197   om_max= np.zeros((kx.shape[0],),'Complex64');
 198   omax  = complex(0,0)
 199   AAmax = np.zeros((5*N,5*N),'Complex64')
 200   kmax  = 0;
 201   # enter main loop
 202   for i in range(kx.shape[0]): # loop over zonal wavelength
 203       sys.stdout.write('\b'*21+'calculating i=%3i/%3i'%(i,kx.shape[0]) )
 204       sys.stdout.flush()
 205 
 206       Uk   = kx[i]*U # k \cdot \v U
 207       UZk  = kx[i]*UZ # k \cdot \v U'
 208 
 209       kh2 = kx[i]**2 + 1e-18
 210 
 211       for n in range(N): # loop over depth
 212         np1 = min(N-1,n+1)
 213         nm  = max(0,n-1)
 214 
 215         B[0,2,n]=  0.
 216         B[1,2,n]= -(kx[i]*fh+UZk[n])/(2*kh2)
 217         B[3,2,n]= I*BZ[n]/2
 218         B[4,2,n]= -cs**2*I/dz
 219 
 220         C[2,0,n]=  0
 221         C[2,1,n]= -kx[i]*fh/2
 222         C[2,3,n]= -I/2
 223         C[2,4,n]=  I/dz
 224 
 225         A[0,:,n]= [ Uk[n],I*f0,0,0,0 ]
 226         A[1,:,n]= [-I*f0 ,Uk[n], -(kx[i]*fh+UZk[n])/(2*kh2),0,I]
 227         A[2,:,n]= [ 0 ,-kx[i]*fh/2 ,(Uk[n]+Uk[np1])/2,-I/2,-I/dz ]
 228         A[3,:,n]= [ -f0*UZk[n],0,I*BZ[n]/2,Uk[n], 0  ]
 229         A[4,:,n]= [ 0, -I*cs**2*kh2 ,  I*cs**2/dz, 0   , 0 ]
 230 
 231       # upper boundary condition
 232       A[2,:,-1]=0; B[2,:,-1]=0; C[2,:,-1]=0;
 233       A[:,2,-1]=0
 234       C[:,2,-1]=0
 235       C[:,2,-2]=0
 236 
 237       # lower boundary condition
 238       A[:,2,0]=A[:,2,0]-B[:,2,0]
 239 
 240       # build large matrix
 241       n1 = 0; n2 = 5
 242       AA[ n1:n2,(n1+5):(n2+5) ] = C[:,:,0]
 243       AA[ n1:n2,n1:n2         ] = A[:,:,0]
 244       for n in range(1,N-1):
 245         n1 = 5*n; n2 = 5*(n+1)
 246         AA[ n1:n2, n1  :n2  ]     = A[:,:,n]
 247         AA[ n1:n2, (n1-5):(n2-5)] = B[:,:,n]
 248         AA[ n1:n2, (n1+5):(n2+5)] = C[:,:,n]
 249       n1 = 5*(N-1); n2 = 5*N
 250       AA[ n1:n2, n1    :n2     ] = A[:,:,-1];
 251       AA[ n1:n2, (n1-5):(n2-5) ] = B[:,:,-1];
 252 
 253       # eigenvalues of matrix
 254       om = np.linalg.eigvals(AA)
 255       kh=kx[i]
 256       # search minimal imaginary eigenvalue
 257       if kh>0: 
 258         om = np.extract( np.abs( np.real( om/kh )) <0.5*cs, om )
 259         n=np.argmin( np.imag(om) )
 260         om_max[i]=om[n]
 261         # look for global minimum
 262         if np.imag(om[n]) < np.imag(omax):
 263            omax=om[n]; kmax=kx[i]; AAmax[:,:] = AA
 264   sys.stdout.write('\n')
 265 
 266   #eigenvectors for global minimum
 267   om, phi=np.linalg.eig(AAmax)
 268   n=np.argmin( np.imag(om) )
 269   om=om[n]
 270   phi = phi[:,n]
 271 
 272   #complete solution
 273   str= phi[0::5]
 274   pot= phi[1::5]
 275   u= -complex(0,1)*kmax*pot
 276   v=-complex(0,1)*kmax*str
 277   w= phi[2::5]
 278   b= phi[3::5]
 279   p= phi[4::5]
 280 
 281   return om_max,om,kmax,u,v,w,b,p
 282 
 283 
 284 if __name__ == "__main__":
 285       m=kelv2()
 286       dt=m.fortran.main_module.dt_tracer
 287       m.run( snapint = dt*100 ,runlen = dt*10000)