
!=======================================================================
!       zonal averaging diagnostics
!       linked in the code in driver only
!=======================================================================
! 
!  Reference buoyancy profile given by N_0  has to be set
! 
!  Note that the damping time scales of the resoring zones have to 
!  be adjusted in subroutine init_zonal_averages to match 
!  the setting in  subroutine restoring_zones

module zonal_averages_module
!=======================================================================
!       Module for zonal mean diagnostics
!=======================================================================
      implicit none
      real*8, allocatable, dimension(:,:) :: ulen,vlen,wlen,blen
      real*8, allocatable, dimension(:,:) :: um,vm,wm,bm,pm
      real*8, allocatable, dimension(:,:,:)   :: vb2,wb2 ,tm
      real*8, allocatable, dimension(:,:,:,:) :: vt2,wt2 
      real*8, allocatable, dimension(:,:) :: eke,bws,epe,T2,T4u,T4v,T4
      real*8, allocatable, dimension(:,:) :: mke,mpe,bwq,rest,restm
      real*8, allocatable, dimension(:,:) :: eps_b,eps_u,eps_v
      real*8, allocatable, dimension(:,:) :: eps_kh,eps_kv
      real*8, allocatable, dimension(:,:) :: vps,wps,veke,weke,vepe,wepe
      real*8, allocatable, dimension(:,:) :: vv,vvb,ageo,vw,ageo2,uv,vvu
      real*8, allocatable, dimension(:,:) :: agev1,agev2
      real*8, allocatable, dimension(:,:) :: mpv_qg,mpv_er,mrv_er
      real*8, allocatable, dimension(:,:,:) :: vq_qg,vq_er,wq_er,wq_qg
      real*8, allocatable, dimension(:,:) :: vvq_qg,vvq_er
      real*8, allocatable, dimension(:) :: b_r, p_r

      integer :: nr_var = 5  ! highest moment to be calculated
      integer :: nr_tr  = 0  ! numbers of tracers to be considered
      real*8, allocatable, dimension(:) :: t_rest
      real*8,parameter :: epsln     = 1.0d-20
      real*8,parameter :: spval = -1.0d33
      real*8,parameter :: N_0 = 0.004

      logical :: diag_eddy_fluxes          = .true.
      logical :: diag_eddy_flux_tendency   = .false.
      logical :: diag_energy_budget        = .false.
      logical :: diag_potential_vorticity  = .false.
end module zonal_averages_module



subroutine zonal_averages
!=======================================================================
!      Verteiler function
!=======================================================================
      use pyOM_module   
      use zonal_averages_module
      use fcontrol_module
      implicit none
      if ( mod(itt,int(snapint/dt))  == 0)  then
       call zonal_averages_means
       if (diag_eddy_fluxes) call zonal_averages_eddy_fluxes
       if (diag_energy_budget) call zonal_averages_energy
       if (diag_eddy_flux_tendency) call zonal_averages_flux_tendency
       if (diag_potential_vorticity) call zonal_averages_pv_fluxes
       call zonal_averages_write
      endif
end subroutine zonal_averages



subroutine init_zonal_averages
!=======================================================================
!      Initialize zonal mean diagnostics
!=======================================================================
      use pyOM_module   
      use zonal_averages_module
      use fcontrol_module
      implicit none
      include "netcdf.inc"
      integer :: i,j,k,js,je,n
      integer :: ncid,iret
      integer :: z_tdim,z_udim,itimedim,varid
      integer :: lat_tdim,lat_udim, dims(4)
      character :: name*64, unit*32
      if (my_pe==0) then
       print*,' Initializing zonal averaging module '
       print*,' using background stability freq. ', N_0,'1/s'
      endif

      allocate( b_r(nz) ); b_r=0.0
      allocate( p_r(nz) ); p_r=0.0
      b_r = -zt*N_0**2
      p_r(nz) = 0
      do k=nz-1,1,-1
       p_r(k)=p_r(k+1)+0.5*(b_r(k+1)+b_r(k))*dz
      enddo


      if (enable_diag_tracer)   nr_tr = nt

      if (my_pe==0) then
       print*,' found ',nr_tr,' tracer '
      endif

      allocate( um(ny ,nz) ); um=0.0
      allocate( vm(ny ,nz) ); vm=0.0
      allocate( wm(ny ,nz) ); wm=0.0
      allocate( bm(ny ,nz) ); bm=0.0
      allocate( pm(ny ,nz) ); pm=0.0
      allocate( tm(ny ,nz,nr_tr) ); tm=0.0

    if (diag_eddy_fluxes) then
      allocate( vb2(ny ,nz,1:nr_var) ); vb2=0.0
      allocate( wb2(ny ,nz,1:nr_var) ); wb2=0.0
      allocate( vt2(ny ,nz,nr_tr,1:nr_var) ); vt2=0.0
      allocate( wt2(ny ,nz,nr_tr,1:nr_var) ); wt2=0.0
    endif

    if (diag_energy_budget) then
      allocate( eke(ny ,nz) ); eke=0.0
      allocate( epe(ny ,nz) ); epe=0.0
      allocate( mke(ny ,nz) ); mke=0.0
      allocate( mpe(ny ,nz) ); mpe=0.0
      allocate( bws(ny ,nz) ); bws=0.0
      allocate( bwq(ny ,nz) ); bwq=0.0
      allocate( T2 (ny ,nz) ); T2=0.0
      allocate( T4u(ny ,nz) ); T4u=0.0
      allocate( T4v(ny ,nz) ); T4v=0.0
      allocate( T4(ny ,nz) );  T4=0.0

      allocate( eps_b(ny ,nz) ); eps_b=0.0
      allocate( eps_u(ny ,nz) ); eps_u=0.0
      allocate( eps_v(ny ,nz) ); eps_v=0.0
      allocate( eps_kh(ny ,nz)); eps_kh=0.0
      allocate( eps_kv(ny ,nz)); eps_kv=0.0
      allocate( rest(ny ,nz) );  rest=0.0
      allocate( restm(ny ,nz) ); restm=0.0

      allocate( vps(ny ,nz) ); vps=0.0
      allocate( wps(ny ,nz) ); wps=0.0
      allocate( veke(ny ,nz) ); veke=0.0
      allocate( weke(ny ,nz) ); weke=0.0
      allocate( vepe(ny ,nz) ); vepe=0.0
      allocate( wepe(ny ,nz) ); wepe=0.0
     endif

     if (diag_eddy_flux_tendency) then
      allocate( vv(ny ,nz)  ); vv  =0.0
      allocate( vvb(ny ,nz) ); vvb =0.0
      allocate( ageo(ny ,nz)); ageo=0.0
      allocate( ageo2(ny ,nz));ageo2=0.0
      allocate( vw(ny ,nz)  ); vw  =0.0
      allocate( uv(ny ,nz)  ); uv  =0.0
      allocate( vvu(ny ,nz)  ); vvu  =0.0
      allocate( agev1(ny ,nz));agev1=0.0
      allocate( agev2(ny ,nz));agev2=0.0
     endif

     if (diag_potential_vorticity) then
      allocate( mpv_qg(ny ,nz)  ); mpv_qg=0.0
      allocate( mpv_er(ny ,nz)  ); mpv_er=0.0
      allocate( mrv_er(ny ,nz)  ); mrv_er=0.0
      allocate( vq_qg(ny ,nz,1:nr_var)  );  vq_qg=0.0
      allocate( vq_er(ny ,nz,1:nr_var)  );  vq_er=0.0
      allocate( wq_qg(ny ,nz,1:nr_var)  );  wq_qg=0.0
      allocate( wq_er(ny ,nz,1:nr_var)  );  wq_er=0.0
      allocate( vvq_er(ny ,nz)  ); vvq_er=0.0
      allocate( vvq_qg(ny ,nz)  ); vvq_qg=0.0
     endif

      allocate( ulen(ny ,nz) ); ulen=0.0
      allocate( vlen(ny ,nz) ); vlen=0.0
      allocate( wlen(ny ,nz) ); wlen=0.0
      allocate( blen(ny ,nz) ); blen=0.0

      js=max(2,js_pe); je = min(je_pe,ny -1)
      do j=js,je
       do i=2,nx -1
        ulen(j,:)=ulen(j,:)+dx*maskU(i,j,:)
        vlen(j,:)=vlen(j,:)+dx*maskV(i,j,:)
        wlen(j,:)=wlen(j,:)+dx*maskW(i,j,:)
        blen(j,:)=blen(j,:)+dx*maskT(i,j,:)
       enddo 
      enddo 
!       Restoring time scale for buffer zone
      allocate( t_rest(ny )); t_rest=0.
!      do j=2,3+1
!        t_rest(j)= 1./(5*85400)/(j-1.)
!      enddo
!      do j=ny -1,ny -3,-1
!        t_rest(j)= 1./(5*85400)/(-1.*(j-ny ))
!      enddo

    if (my_pe==0) then 

      call def_grid_cdf('zonal_ave.cdf')
      iret=nf_open('zonal_ave.cdf',NF_WRITE,ncid)
      iret=nf_set_fill(ncid, NF_NOFILL, iret)
      call ncredf(ncid, iret)
      iret=nf_inq_dimid(ncid,'yt',lat_tdim)
      iret=nf_inq_dimid(ncid,'yu',lat_udim)
      iret=nf_inq_dimid(ncid,'zt',z_tdim)
      iret=nf_inq_dimid(ncid,'zu',z_udim)
      iret=nf_inq_dimid(ncid,'Time',itimedim)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'um', NCFLOAT,3,dims,iret)
      name = 'Zonal mean zonal velocity'; unit = 'm/s'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vm', NCFLOAT,3,dims,iret)
      name = 'Zonal mean meridional velocity'; unit = 'm/s'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_udim, iTimedim,1/)
      varid = ncvdef (ncid,'wm', NCFLOAT,3,dims,iret)
      name = 'Zonal mean vertical velocity'; unit = 'm/s'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'bm', NCFLOAT,3,dims,iret)
      name = 'Zonal mean buoyancy'; unit = 'm/s**2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'pm', NCFLOAT,3,dims,iret)
      name = 'Zonal mean pressure'; unit = 'm**2/s**2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      do n=1,nr_tr
       dims = (/lat_tdim, z_tdim, iTimedim,1/)
       write(name, '("trm",i2)') n; call replace_space_zero(name)
       varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
       name = 'zonal mean tracer'; unit = ' '
       call dvcdf(ncid,varid,name,64,unit,32,spval)
      enddo

    if (diag_eddy_fluxes) then
      do k=1,nr_var
       write(unit, '("m**",i2,"/s**",i2)')k+1,2*k+1
       dims = (/lat_udim, z_tdim, iTimedim,1/)
       write(name, '("vb_",i2)') k; call replace_space_zero(name)
       varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
       write(name, '("Meridional buoyancy ",i2,". moment flux")')k
       call dvcdf(ncid,varid,name,64,unit,32,spval)

       dims = (/lat_tdim, z_udim, iTimedim,1/)
       write(name, '("wb_",i2)') k; call replace_space_zero(name)
       varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
       write(name, '("Vertical buoyancy ",i2,". moment flux")')k
       call dvcdf(ncid,varid,name,64,unit,32,spval)

       do n=1,nr_tr
        unit = 'm/s'
        dims = (/lat_udim, z_tdim, iTimedim,1/)
        write(name, '("vt",i2,"_",i2)')n,k;call replace_space_zero(name)
        varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
        write(name, '("Meridional ",i2,". moment tracer flux")')k
        call dvcdf(ncid,varid,name,64,unit,32,spval)

        dims = (/lat_tdim, z_udim, iTimedim,1/)
        write(name, '("wt",i2,"_",i2)')n,k;call replace_space_zero(name)
        varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
        write(name, '("Vertical ",i2,". moment tracer flux")')k
        call dvcdf(ncid,varid,name,64,unit,32,spval)
       enddo
      enddo
     endif

    if (diag_energy_budget) then
      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'eke', NCFLOAT,3,dims,iret)
      name = 'Eddy kinetic energy'; unit = 'm**2/s**2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'mke', NCFLOAT,3,dims,iret)
      name = 'Mean kinetic energy'; unit = 'm**2/s**2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'epe', NCFLOAT,3,dims,iret)
      name = 'Eddy potential energy'; unit = 'm**2/s**2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'mpe', NCFLOAT,3,dims,iret)
      name = 'Mean potential energy'; unit = 'm**2/s**2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'bws', NCFLOAT,3,dims,iret)
      name = 'EKE production (EPE->EKE)'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'bwq', NCFLOAT,3,dims,iret)
      name = 'MPE production (MPE->MKE)'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'T2', NCFLOAT,3,dims,iret)
      name = 'MPE production (EPE->MPE)'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'T4', NCFLOAT,3,dims,iret)
      name = 'MKE production (EKE->MKE)'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'eps_b', NCFLOAT,3,dims,iret)
      name = 'EPE dissipation'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'eps_kh', NCFLOAT,3,dims,iret)
      name = 'Horizontal EKE dissipation'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'eps_kv', NCFLOAT,3,dims,iret)
      name = 'Vertical EKE dissipation'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'rest', NCFLOAT,3,dims,iret)
      name = 'EPE dissipation by nudging'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'restm', NCFLOAT,3,dims,iret)
      name = 'MPE dissipation by nudging'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vps', NCFLOAT,3,dims,iret)
      name = 'Meridional wave radiation'; unit = 'm**3/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_udim, iTimedim,1/)
      varid = ncvdef (ncid,'wps', NCFLOAT,3,dims,iret)
      name = 'Vertical wave radiation'; unit = 'm**3/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'veke', NCFLOAT,3,dims,iret)
      name = 'Meridional EKE advection'; unit = 'm**3/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_udim, iTimedim,1/)
      varid = ncvdef (ncid,'weke', NCFLOAT,3,dims,iret)
      name = 'Vertical EKE advection'; unit = 'm**3/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vepe', NCFLOAT,3,dims,iret)
      name = 'Meridional EPE advection'; unit = 'm**3/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_udim, iTimedim,1/)
      varid = ncvdef (ncid,'wepe', NCFLOAT,3,dims,iret)
      name = 'Vertical EPE advection'; unit = 'm**3/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)
     endif

    if (diag_eddy_flux_tendency) then
      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vv', NCFLOAT,3,dims,iret)
      name = 'Meridional velocity variance'; unit = 'm^2/s^2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vvb', NCFLOAT,3,dims,iret)
      name = 'Triple variance flux'; unit = 'm^3/s^4'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'ageo', NCFLOAT,3,dims,iret)
      name = 'Ageostrophic terms'; unit = 'm^2/s^4'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'ageo2', NCFLOAT,3,dims,iret)
      name = 'Ageostrophic (beta) term'; unit = 'm^2/s^4'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_udim, iTimedim,1/)
      varid = ncvdef (ncid,'vw', NCFLOAT,3,dims,iret)
      name = 'Velocity correlation'; unit = 'm^2/s^4'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'uv', NCFLOAT,3,dims,iret)
      name = 'Velocity correlation'; unit = 'm^2/s^2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vvu', NCFLOAT,3,dims,iret)
      name = 'Flux of velocity correlation'; unit = 'm^3/s^3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'agev1', NCFLOAT,3,dims,iret)
      name = 'Ageostrophic term 1'; unit = 'm^2/s^3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'agev2', NCFLOAT,3,dims,iret)
      name = 'Ageostrophic term 2'; unit = 'm^2/s^3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)
     endif

     if (diag_potential_vorticity) then
      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'mpv_qg', NCFLOAT,3,dims,iret)
      name = 'Potential vorticity (QG)'; unit = '1/s'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vvq_qg', NCFLOAT,3,dims,iret)
      name = 'Triple var. flux of PV (QG)'; unit = 'm/s^2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'mpv_er', NCFLOAT,3,dims,iret)
      name = 'Potential vorticity (Ertel)'; unit = '1/s'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'mrv_er', NCFLOAT,3,dims,iret)
      name = 'Relative vorticity (Ertel)'; unit = '1/s'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vvq_er', NCFLOAT,3,dims,iret)
      name = 'Triple var. flux of PV (Ertel)'; unit = 'm^2/s^3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      do k=1,nr_var
       write(unit, '("m**",i2,"/s**",i2)')k,k+1

       dims = (/lat_udim, z_tdim, iTimedim,1/)
       write(name, '("vq_qg_",i2)') k; call replace_space_zero(name)
       varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
       write(name, '("Meridional PV (QG) ",i2,". moment flux")')k
       call dvcdf(ncid,varid,name,64,unit,32,spval)

       dims = (/lat_udim, z_tdim, iTimedim,1/)
       write(name, '("vq_er_",i2)') k; call replace_space_zero(name)
       varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
       write(name, '("Meridional PV (Ertel) ",i2,". moment flux")')k
       call dvcdf(ncid,varid,name,64,unit,32,spval)

       dims = (/lat_tdim, z_udim, iTimedim,1/)
       write(name, '("wq_qg_",i2)') k; call replace_space_zero(name)
       varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
       write(name, '("Vertical PV (QG) ",i2,". moment flux")')k
       call dvcdf(ncid,varid,name,64,unit,32,spval)

       dims = (/lat_tdim, z_udim, iTimedim,1/)
       write(name, '("wq_er_",i2)') k; call replace_space_zero(name)
       varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
       write(name, '("Vertical PV (Ertel) ",i2,". moment flux")')k
       call dvcdf(ncid,varid,name,64,unit,32,spval)

      enddo
     endif
      varid = ncvdef (ncid,'b_r', NCFLOAT,1,z_tdim,iret)
      name = 'Reference buoyancy profile'; unit = 'm/s^2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      call ncendf(ncid, iret)
      iret= nf_put_vara_double(ncid,varid,1,nz,b_r)

      call ncclos (ncid, iret)
    endif ! my_pe ==0
    if (my_pe==0) print*,' Done initializing zonal averaging module'
  end subroutine init_zonal_averages



 subroutine zonal_averages_write
!======================================================================
!       compute zonal mean diagnostics
!=======================================================================
      use pyOM_module   
      use zonal_averages_module
      use fcontrol_module
      implicit none
      include "netcdf.inc"
      integer :: ncid,iret,n,npe, corner(4), edges(4),m
      integer :: itdimid,ilen,itimeid
      integer :: trmid(nr_tr),varid
      real*8 :: fxa,time
      character :: name*64
!-----------------------------------------------------------------------
!      write to netcdf file
!-----------------------------------------------------------------------
      do npe=0,n_pes
       if (my_pe==npe) then

         iret=nf_open('zonal_ave.cdf',NF_WRITE,ncid)
         iret=nf_set_fill(ncid, NF_NOFILL, iret)
         do n=1,nr_tr
          write(name, '("trm",i2)') n; call replace_space_zero(name)
          iret=nf_inq_varid(ncid,name,trmid(n))
         enddo

         iret=nf_inq_dimid(ncid,'Time',itdimid)
         iret=nf_inq_dimlen(ncid, itdimid,ilen)
         iret=nf_inq_varid(ncid,'Time',itimeid)
         if (my_pe==0) then
          ilen=ilen+1
          time = dt*itt !current_time-initial_time
          fxa = time/86400. ! time%days + time%seconds/86400.
          iret= nf_put_vara_double(ncid,itimeid,ilen,1,fxa)
         endif

         Corner = (/js_pe,1,ilen,1/); 
         edges  = (/je_pe-js_pe+1,nz,1,1/)
         iret=nf_inq_varid(ncid,'um',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges, um(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'vm',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges, vm(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'wm',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges, wm(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'bm',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges, bm(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'pm',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges, pm(js_pe:je_pe,:))
         do n=1,nr_tr
          iret= nf_put_vara_double(ncid,trmid(n),corner,edges, tm(js_pe:je_pe,:,n))
         enddo

    if (diag_eddy_fluxes) then
         do m=1,nr_var
          write(name, '("vb_",i2)')m; call replace_space_zero(name)
          iret=nf_inq_varid(ncid,name,varid)
          iret= nf_put_vara_double(ncid,varid,corner,edges,     vb2(js_pe:je_pe,:,m))
          write(name, '("wb_",i2)')m; call replace_space_zero(name)
          iret=nf_inq_varid(ncid,name,varid)
          iret= nf_put_vara_double(ncid,varid,corner,edges,  wb2(js_pe:je_pe,:,m))
          do n=1,nr_tr
           write(name, '("vt",i2,"_",i2)')n,m;
           call replace_space_zero(name)
           iret=nf_inq_varid(ncid,name,varid)
           iret= nf_put_vara_double(ncid,varid,corner,edges,  vt2(js_pe:je_pe,:,n,m))
           write(name, '("wt",i2,"_",i2)')n,m;
           call replace_space_zero(name)
           iret=nf_inq_varid(ncid,name,varid)
           iret= nf_put_vara_double(ncid,varid,corner,edges,  wt2(js_pe:je_pe,:,n,m))
          enddo
         enddo
      endif

      if (diag_energy_budget) then
         iret=nf_inq_varid(ncid,'eke',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,   eke(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'mke',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,   mke(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'epe',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,   epe(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'mpe',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,  mpe(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'bws',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,  bws(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'bwq',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,  bwq(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'T2',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,  T2(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'T4',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges, T4u(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'T4v',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges, T4v(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'eps_b',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges, eps_b(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'eps_kh',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,  eps_kh(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'eps_kv',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,  eps_kv(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'rest',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,  rest(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'restm',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,   restm(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'vps',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,   vps(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'wps',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,  wps(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'veke',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges, veke(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'weke',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,  weke(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'vepe',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,   vepe(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'wepe',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,  wepe(js_pe:je_pe,:))
       endif

      if (diag_eddy_flux_tendency) then
         iret=nf_inq_varid(ncid,'vv',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,    vv(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'vvb',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,   vvb(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'vw',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,  vw(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'ageo',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,  ageo(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'ageo2',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,  ageo2(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'uv',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges, uv(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'vvu',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges, vvu(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'agev1',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges, agev1(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'agev2',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges, agev2(js_pe:je_pe,:))
        endif

        if (diag_potential_vorticity) then
         iret=nf_inq_varid(ncid,'mpv_qg',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,   mpv_qg(js_pe:je_pe,:))

         iret=nf_inq_varid(ncid,'mpv_er',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,   mpv_er(js_pe:je_pe,:))

         iret=nf_inq_varid(ncid,'mrv_er',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,  mrv_er(js_pe:je_pe,:))

         iret=nf_inq_varid(ncid,'vvq_er',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges, vvq_er(js_pe:je_pe,:))

         iret=nf_inq_varid(ncid,'vvq_qg',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges, vvq_qg(js_pe:je_pe,:))

         do m=1,nr_var

          write(name, '("vq_qg_",i2)')m; call replace_space_zero(name)
          iret=nf_inq_varid(ncid,name,varid)
          iret=nf_put_vara_double(ncid,varid,corner,edges, vq_qg(js_pe:je_pe,:,m))

          write(name, '("wq_qg_",i2)')m; call replace_space_zero(name)
          iret=nf_inq_varid(ncid,name,varid)
          iret=nf_put_vara_double(ncid,varid,corner,edges,  wq_qg(js_pe:je_pe,:,m))

          write(name, '("vq_er_",i2)')m; call replace_space_zero(name)
          iret=nf_inq_varid(ncid,name,varid)
          iret=nf_put_vara_double(ncid,varid,corner,edges,  vq_er(js_pe:je_pe,:,m))

          write(name, '("wq_er_",i2)')m; call replace_space_zero(name)
          iret=nf_inq_varid(ncid,name,varid)
          iret=nf_put_vara_double(ncid,varid,corner,edges, wq_er(js_pe:je_pe,:,m))

         enddo
        endif

         call ncclos (ncid, iret)
       endif
       call fortran_barrier
      enddo
      end subroutine zonal_averages_write


  subroutine zonal_averages_means
!=======================================================================
!       get zonal means
!=======================================================================
      use pyOM_module   
      use zonal_averages_module
      use fcontrol_module
      implicit none
      integer :: i,j,js,je,n

      js=max(2,js_pe); je = min(je_pe,ny -1)
      um=0.0;vm=0.0;wm=0.0;bm=0.0;tm=0.0
      do j=js,je 
        do i=2,nx -1
         um(j,:)=um(j,:)+u(i,j,:,tau)*dx*maskU(i,j,:) 
         vm(j,:)=vm(j,:)+v(i,j,:,tau)*dx*maskV(i,j,:) 
         wm(j,:)=wm(j,:)+w(i,j,:,tau)*dx*maskW(i,j,:) 
         bm(j,:)=bm(j,:)+b(i,j,:,  tau)*dx*maskT(i,j,:) 
         pm(j,:)=pm(j,:)+p_full(i,j,:,tau)*dx*maskT(i,j,:) 
        enddo 
      enddo 

    if (enable_diag_tracer) then
      do j=js,je 
        do i=2,nx -1
         do n=1,nr_tr
          tm(j,:,n)=tm(j,:,n)+tr(i,j,:,tau,n)*dx*maskT(i,j,:) 
         enddo
        enddo 
      enddo 
    endif

      where (ulen/=0); um=um/ulen; elsewhere; um=spval ; end where
      where (vlen/=0); vm=vm/vlen; elsewhere; vm=spval ; end where
      where (wlen/=0); wm=wm/wlen; elsewhere; wm=spval ; end where
      where (blen/=0); bm=bm/blen; elsewhere; bm=spval ; end where
      where (blen/=0); pm=pm/blen; elsewhere; pm=spval ; end where
      call border_exchg_merid(ny,nz,um,1)
      call border_exchg_merid(ny,nz,vm,1)
      call border_exchg_merid(ny,nz,wm,1)
      call border_exchg_merid(ny,nz,bm,1)
      call border_exchg_merid(ny,nz,pm,1)
      do n=1,nr_tr
        where (blen/=0); tm(:,:,n)=tm(:,:,n)/blen
        elsewhere; tm(:,:,n)=spval;end where
        call border_exchg_merid(ny,nz,tm(:,:,n),1)
      enddo
 end subroutine zonal_averages_means



  subroutine zonal_averages_eddy_fluxes
!=======================================================================
!       compute zonal mean eddy fluxes
!=======================================================================
      use pyOM_module   
      use zonal_averages_module
      use fcontrol_module
      implicit none
      integer :: i,j,k,js,je,m,n

      js=max(2,js_pe); je = min(je_pe,ny -1)
      vb2=0.0;wb2=0.0;vt2=0.0;wt2=0.0
      do m=1,nr_var
        do j=js,je 
         do k=1,nz
          do i=2,nx -1
           vb2(j,k,m)=vb2(j,k,m)+v(i,j,k,tau)*dx*maskV(i,j,k)   &
           *.5/m*( (b(i,j  ,k,tau)-bm(j  ,k))**m + (b(i,j+1,k,tau)-bm(j+1,k))**m  )
          enddo
         enddo
         do k=1,nz-1
          do i=2,nx -1
           wb2(j,k,m)=wb2(j,k,m)+w(i,j,k,tau)*dx*maskW(i,j,k)    &
            *.5/m*( (b(i,j,k  ,tau)-bm(j,k  ))**m + (b(i,j,k+1,tau)-bm(j,k+1))**m )
          enddo
         enddo
       enddo
       where (vlen/=0); vb2(:,:,m)=vb2(:,:,m)/vlen; 
       elsewhere; vb2(:,:,m)=spval; end where
       where (wlen/=0); wb2(:,:,m)=wb2(:,:,m)/wlen; 
       elsewhere; wb2(:,:,m)=spval; end where

 if (enable_diag_tracer) then
        do n=1,nr_tr
          do j=js,je 
           do k=1,nz
            do i=2,nx -1
             vt2(j,k,n,m)=vt2(j,k,n,m)+v(i,j,k,tau)*dx*maskV(i,j,k)   &
             *.5/m*( (tr(i,j  ,k,tau,n)-tm(j  ,k,n))**m  + (tr(i,j+1,k,tau,n)-tm(j+1,k,n))**m )
            enddo
           enddo
           do k=1,nz-1
            do i=2,nx -1
             wt2(j,k,n,m)=wt2(j,k,n,m)+w(i,j,k,tau)*dx*maskW(i,j,k)   &
             *.5/m*( (tr(i,j,k  ,tau,n)-tm(j,k  ,n))**m + (tr(i,j,k+1,tau,n)-tm(j,k+1,n))**m )
            enddo
           enddo
          enddo ! j
          where (vlen/=0); vt2(:,:,n,m)=vt2(:,:,n,m)/vlen; 
          elsewhere; vt2(:,:,n,m)=spval; end where
          where (wlen/=0); wt2(:,:,n,m)=wt2(:,:,n,m)/wlen; 
          elsewhere; wt2(:,:,n,m)=spval; end where
        enddo ! n
 endif
      enddo ! m
 end subroutine zonal_averages_eddy_fluxes





 subroutine zonal_averages_pv_fluxes
!=======================================================================
!       compute zonal mean potential vorticity
!=======================================================================
      use pyOM_module   
      use zonal_averages_module
      implicit none
      integer :: i,j,k,js,je,m
      real*8 :: pv_qg(nx ,ny ,nz),f0,y0,ps(nx ,ny ,nz)
      real*8 :: rv_qg(nx ,ny ,nz),pv_er(nx ,ny ,nz)
      real*8 :: rv(nx ,ny ,nz),fxa,beta

      f0 = coriolis_t(ny /2); y0 = yt(ny /2)
      js=max(2,js_pe); je = min(je_pe,ny -1)
!
!       relative potential vorticity
!
      rv=0.0
      do j=js,je 
       do k=2,nz-1
        do i=2,nx -1
          rv(i,j,k)=(v(i+1,j,k,tau)-v(i,j,k,tau) )/dx*maskU(i,j,k) &
                   -(u(i,j+1,k,tau)-u(i,j,k,tau) )/dx*maskV(i,j,k)
        enddo 
       enddo 
      enddo 
      call border_exchg3D(nx,ny,nz,rv,1)
!
!       Ertel potential vorticity  Q N0^2 = b_z (v_x-u_y + f) + u_z b_y - v_z b_x
!
      pv_er=0.
      do j=js,je 
       do k=2,nz-1
        do i=2,nx -1
         fxa= (rv(i,j,k)+rv(i-1,j,k)+rv(i,j-1,k)+rv(i-1,j-1,k))/4.0
         pv_er(i,j,k)=(fxa+coriolis_t(j)  )*  ((b(i,j,k+1,tau)-b(i,j,k,tau))/dz*maskW(i,j,k)   &
                  +(b(i,j,k,tau)-b(i,j,k-1,tau))/dz*maskW(i,j,k-1))/ (maskW(i,j,k)+maskW(i,j,k-1)+epsln) 
         pv_er(i,j,k)=-N_0**(-2)*pv_er(i,j,k)
        enddo 
       enddo 
      enddo 
      call border_exchg3D(nx,ny,nz,pv_er,1); call setcyclic3D(nx,ny,nz,pv_er)
!
!       QG streamfunction
!
      psi=0.0
      do k=1,nz
       do j=js,je 
        ps(:,j,k)=(p_full(:,j,k,tau)-p_r(k))/f0
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,ps,1); call setcyclic3D(nx,ny,nz,ps)
!
!       QG relative potential vorticity
!
      rv_qg=0.
      do j=js,je 
       do k=2,nz-1
        do i=2,nx -1
         rv_qg(i,j,k)= ((ps(i,j+1,k)-ps(i,j,k))/dx*maskV(i,j  ,k)-(ps(i,j,k)-ps(i,j-1,k))/dx*maskV(i,j-1,k))/dx  &
         +((ps(i+1,j,k)-ps(i,j,k))/dx*maskU(i-1,j,k) -(ps(i,j,k)-ps(i-1,j,k))/dx*maskU(i-1,j,k))/dx
        enddo 
       enddo 
      enddo 
      call border_exchg3D(nx,ny,nz,rv_qg,1); call setcyclic3D(nx,ny,nz,rv_qg)
!
!       add QG planetary and stretching term to relative vorticity
!
      pv_qg=0.
      do j=js,je 
       beta = (coriolis_t(j+1)-coriolis_t(j))/dx
       pv_qg(:,j,:)=rv_qg(:,j,:)+beta*(yt(j)-y0)
       do k=2,nz-1
         pv_qg(:,j,k)=pv_qg(:,j,k)+f0**2/N_0**2* ((ps(:,j,k+1)-ps(:,j,k))/dz*maskW(:,j,k)   &
          -(ps(:,j,k)-ps(:,j,k-1))/dz*maskW(:,j,k-1))/dz
       enddo 
       k=2
       pv_qg(:,j,k)=pv_qg(:,j,k)+f0**2/N_0**2* ((ps(:,j,k+1)-ps(:,j,k))/dz*maskW(:,j,k)  &
          -(b(:,j,k,tau)-bm(j,k))*maskT(i,j,k) )/dz
       k=nz-1
       pv_qg(:,j,k)=pv_qg(:,j,k)+f0**2/N_0**2* ((b(:,j,k,tau)-bm(j,k))*maskT(:,j,k)   &
          -(ps(:,j,k)-ps(:,j,k-1))/dz*maskW(:,j,k-1))/dz
      enddo 
      call border_exchg3D(nx,ny,nz,pv_qg,1); call setcyclic3D(nx,ny,nz,pv_qg)
!
!       mean QG and Ertel PV
!
      mpv_qg=0.
      do j=js,je 
        do i=2,nx -1
         mpv_qg(j,:)=mpv_qg(j,:)+pv_qg(i,j,:)*dx*maskT(i,j,:) 
         mpv_er(j,:)=mpv_er(j,:)+pv_er(i,j,:)*dx*maskT(i,j,:) 
         mrv_er(j,:)=mrv_er(j,:)+rv(i,j,:)*dx*maskT(i,j,:) 
        enddo 
      enddo 
      where (blen/=0); mpv_qg=mpv_qg/blen; 
      elsewhere; mpv_qg=spval ; end where
      where (blen/=0); mpv_er=mpv_er/blen; 
      elsewhere; mpv_er=spval ; end where
      where (vlen/=0); mrv_er=mrv_er/vlen; 
      elsewhere; mrv_er=spval ; end where
      call border_exchg_merid(ny,nz,mpv_qg,1)
      call border_exchg_merid(ny,nz,mpv_er,1)
      call border_exchg_merid(ny,nz,mrv_er,1)
!
!       meridional and vertical fluxes of PV
!
      vq_qg=0.;vq_er=0.; wq_qg=0.;wq_er=0.; 
      do m=1,nr_var
        do j=js,je 
         do k=1,nz
          do i=2,nx -1
           vq_qg(j,k,m)=vq_qg(j,k,m)+v(i,j,k,tau)*dx*maskV(i,j,k)   &
           *.5/m*( (pv_qg(i,j  ,k)-mpv_qg(j  ,k))**m + (pv_qg(i,j+1,k)-mpv_qg(j+1,k))**m  )
           vq_er(j,k,m)=vq_er(j,k,m)+v(i,j,k,tau)*dx*maskV(i,j,k)   &
           *.5/m*( (pv_er(i,j  ,k)-mpv_er(j  ,k))**m + (pv_er(i,j+1,k)-mpv_er(j+1,k))**m  )
          enddo
         enddo
         do k=1,nz-1
          do i=2,nx -1
           wq_qg(j,k,m)=wq_qg(j,k,m)+w(i,j,k,tau)*dx*maskW(i,j,k)    &
            *.5/m*( (pv_qg(i,j,k  )-mpv_qg(j,k  ))**m  + (pv_qg(i,j,k+1)-mpv_qg(j,k+1))**m )
           wq_er(j,k,m)=wq_er(j,k,m)+w(i,j,k,tau)*dx*maskW(i,j,k)    &
            *.5/m*( (pv_er(i,j,k  )-mpv_er(j,k  ))**m + (pv_er(i,j,k+1)-mpv_er(j,k+1))**m )
          enddo
         enddo
       enddo
       where (vlen/=0); vq_qg(:,:,m)=vq_qg(:,:,m)/vlen; 
       elsewhere; vq_qg(:,:,m)=spval; end where
       where (wlen/=0); wq_qg(:,:,m)=wq_qg(:,:,m)/wlen; 
       elsewhere; wq_qg(:,:,m)=spval; end where
       where (vlen/=0); vq_er(:,:,m)=vq_er(:,:,m)/vlen; 
       elsewhere; vq_er(:,:,m)=spval; end where
       where (wlen/=0); wq_er(:,:,m)=wq_er(:,:,m)/wlen; 
       elsewhere; wq_er(:,:,m)=spval; end where
      enddo
!
!      triple flux of PV (for eddy pv flux tendency)
!
      vvq_qg=0.;vvq_er=0.
      do j=js,je 
       do k=1,nz
        do i=2,nx -1
         vvq_qg(j,k)=vvq_qg(j,k) +(v(i,j,k,tau)-vm(j,k))**2*dx*maskV(i,j,k)   &
           *((pv_qg(i,j  ,k)-mpv_qg(j  ,k)) +(pv_qg(i,j+1,k)-mpv_qg(j+1,k)))/2.
         vvq_er(j,k)=vvq_er(j,k) +(v(i,j,k,tau)-vm(j,k))**2*dx*maskV(i,j,k)   &
           *((pv_er(i,j  ,k)-mpv_er(j  ,k)) +(pv_er(i,j+1,k)-mpv_er(j+1,k)))/2.
        enddo
       enddo
      enddo
      where (vlen/=0); vvq_qg(:,:)=vvq_qg(:,:)/vlen; 
      elsewhere; vvq_qg(:,:)=spval; end where
      where (vlen/=0); vvq_er(:,:)=vvq_er(:,:)/vlen; 
      elsewhere; vvq_er(:,:)=spval; end where
 end subroutine zonal_averages_pv_fluxes



 subroutine zonal_averages_flux_tendency
!=======================================================================
!        eddy flux balance :  (v'b')_t = - v'^2 b_y + ...
!=======================================================================
      use pyOM_module   
      use zonal_averages_module
      implicit none
      integer :: i,j,k,js,je
      real*8 :: fxa,a(nx ,ny ,nz),beta

      js=max(2,js_pe); je = min(je_pe,ny -1)
      vv=0.;vvb=0.;vw=0.0; ageo=0.0; ageo2=0.
      do j=js,je 
         beta = (coriolis_t(j+1)-coriolis_t(j))/dx
        do k=2,nz-1
         do i=2,nx -1
          vv(j,k)=vv(j,k)+maskV(i,j,k)*dx*(v(i,j,k,tau)-vm(j,k))**2
          vvb(j,k)=vvb(j,k)+maskV(i,j,k)*dx*   (v(i,j,k,tau)-vm(j,k))**2  &
              *0.5*( b(i,j,k,tau)-bm(j,k)+b(i,j+1,k,tau)-bm(j+1,k))
          vw(j,k)=vw(j,k)+maskW(i,j,k)* N_0**2*dx* 0.5*( (v(i,j  ,k,tau)-vm(j  ,k))*maskV(i,j  ,k)  &
                   +(v(i,j-1,k,tau)-vm(j-1,k))*maskV(i,j-1,k)) *(w(i,j,k,tau)-wm(j,k))
          ageo(j,k)=ageo(j,k)+maskV(i,j,k)*dx*0.5*( b(i,j,k,tau)-bm(j,k)+b(i,j+1,k,tau)-bm(j+1,k))   &
              *( (p_full(i,j+1,k,tau)-pm(j+1,k) -p_full(i,j  ,k,tau)+pm(j  ,k))/dx   &
          +(coriolis_t(j    )*(u(i-1,j  ,k,tau)-um(j,k)  +u(i  ,j  ,k,tau)-um(j,k))    &
           +coriolis_t(j+1  )*(u(i-1,j+1,k,tau)-um(j+1,k) +u(i  ,j+1,k,tau)-um(j+1,k)))/4.0  )
          ageo2(j,k)=ageo2(j,k)+maskV(i,j,k)*dx* 0.5*( b(i,j,k,tau)-bm(j,k)+b(i,j+1,k,tau)-bm(j+1,k))*(   &
          +(beta*(yt(j  )-yt(ny /2))*(u(i-1,j  ,k,tau)-um(j,k)    +u(i  ,j  ,k,tau)-um(j,k))   &
           +beta*(yt(j+1)-yt(ny /2))*(u(i-1,j+1,k,tau)-um(j+1,k)+u(i  ,j+1,k,tau)-um(j+1,k)))/4.0)
         enddo 
        enddo 
      enddo 
      where (vlen/=0); vv=vv/vlen; elsewhere; vv=spval; end where
      where (vlen/=0); vvb=vvb/vlen; elsewhere; vvb=spval; end where
      where (wlen/=0); vw=vw/wlen; elsewhere; vw=spval; end where
      where (vlen/=0); ageo=ageo/vlen;elsewhere;ageo=spval; end where
      where (vlen/=0);ageo2=ageo2/vlen;elsewhere;ageo2=spval;end where

!
!        tendency of cross correlation u'v' 
!
      uv=0.;vvu=0.
      do j=js,je 
       do k=2,nz-1
        do i=2,nx -1
         fxa=0.5*(u(i,j,k,tau)-um(j,k)+u(i,j+1,k,tau)-um(j+1,k) )   &
                 *0.5*(v(i,j,k,tau)-vm(j,k)+v(i+1,j,k,tau)-vm(j,k) )
         uv(j,k)=uv(j,k)+maskV(i,j,k)*dx*fxa
         vvu(j,k)=vvu(j,k)+maskV(i,j,k)*dx*fxa *0.5*(v(i,j,k,tau)+v(i+1,j,k,tau))
        enddo 
       enddo 
      enddo 
      where (vlen/=0);uv=uv/vlen;elsewhere;uv=spval;end where
      where (vlen/=0);vvu=vvu/vlen;elsewhere;vvu=spval;end where

      do j=js,je 
       do k=2,nz-1
        do i=2,nx -1
          a(i,j,k)= maskU(i,j,k)*(  coriolis_t(j)*(v(i,j  ,k,tau)+v(i+1,j  ,k,tau)+  &
           v(i,j-1,k,tau)+v(i+1,j-1,k,tau))/4.0 -(p_full(i+1,j,k,tau)-p_full(i,j,k,tau))/dx )
        enddo 
       enddo 
      enddo 
      call border_exchg3D(nx,ny,nz,a,1); call setcyclic3D(nx,ny,nz,a)
      agev1=0.
      do j=js,je 
       do k=2,nz-1
        do i=2,nx -1
         agev1(j,k)=agev1(j,k)+maskV(i,j,k)*dx* 0.5*(a(i,j,k)+a(i,j+1,k))*  &
          0.5*(v(i,j,k,tau)-vm(j,k)+v(i+1,j,k,tau)-vm(j,k))
        enddo 
       enddo 
      enddo 
      where (vlen/=0);agev1=agev1/vlen;elsewhere;agev1=spval;end where

      do j=js,je 
       do k=2,nz-1
        do i=2,nx -1
          a(i,j,k)= maskV(i,j,k)*( -(coriolis_t(j  )*(u(i-1,j  ,k,tau)+u(i,j  ,k,tau))  &
           +coriolis_t(j+1)*(u(i-1,j+1,k,tau)+u(i,j+1,k,tau)))/4.0  &
          -(p_full(i,j+1,k,tau)-p_full(i,j,k,tau))/dx) 
        enddo 
       enddo 
      enddo 
      call border_exchg3D(nx,ny,nz,a,1); call setcyclic3D(nx,ny,nz,a)
      agev2=0.
      do j=js,je 
       do k=2,nz-1
        do i=2,nx -1
         agev2(j,k)=agev2(j,k)+maskU(i,j,k)*dx* 0.5*(a(i,j,k)+a(i+1,j,k))*  &
          0.5*(u(i,j,k,tau)-um(j,k)+u(i,j+1,k,tau)-um(j+1,k))
        enddo 
       enddo 
      enddo 
      where (vlen/=0);agev2=agev2/vlen;elsewhere;agev2=spval;end where

 end subroutine zonal_averages_flux_tendency







 subroutine zonal_averages_energy
!=======================================================================
!       compute zonal mean energy budget
!=======================================================================
      use pyOM_module   
      use zonal_averages_module
      implicit none

      integer :: i,j,k,js,je
      real*8 :: a(ny ,nz),bs(nx ,ny ,nz),bms(ny ,nz)
      real*8 :: bss(nx ,ny ,nz),del2(nx ,ny ,nz),us(nx ,ny ,nz)

      js=max(2,js_pe); je = min(je_pe,ny -1)
      bs=0.0;bms=0.0;bss=0.  ! buoyancy deviation from reference state
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
           bs(i,j,k)=(b(i,j,k,tau)-b_r(k))*maskT(i,j,k)
         enddo 
        enddo 
      enddo 
      call border_exchg3D(nx,ny,nz,bs,1); call setcyclic3D(nx,ny,nz,bs)
      do j=js,je 
        do i=2,nx -1
         bms(j,:)=bms(j,:)+bs(i,j,:)*dx*maskT(i,j,:) 
        enddo 
      enddo 
      where (blen/=0) bms=bms/blen
      call border_exchg_merid(ny,nz,bms,1)
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
           bss(i,j,k)=(bs(i,j,k)-bms(j,k))*maskT(i,j,k)
         enddo 
        enddo 
      enddo 
      call border_exchg3D(nx,ny,nz,bss,1); call setcyclic3D(nx,ny,nz,bss)
!-----------------------------------------------------------------------
!       Mean energy
!-----------------------------------------------------------------------
      where (vlen==0); vm=0; end where
      where (wlen==0); wm=0; end where
      call border_exchg_merid(ny,nz,vm,1)
      call border_exchg_merid(ny,nz,wm,1)

      do j=js,je 
        do k=2,nz-1
         mke(j,k)=um(j,k)**2/2+(0.5*(vm(j,k)+vm(j-1,k)))**2/2.
         mpe(j,k)=bms(j,k)**2/2/N_0**2
         bwq(j,k)= 0.25*wm(j,k  )*(bms(j,k  )+bms(j,k+1)) +0.25*wm(j,k-1)*(bms(j,k-1)+bms(j,k  ))
        enddo
      enddo
      where (blen==0); mke=spval; end where
      where (blen==0); mpe=spval; end where
      where (blen==0); bwq=spval; end where

      where (vlen==0); vm=spval; end where
      where (wlen==0); wm=spval; end where
      call border_exchg_merid(ny,nz,vm,1)
      call border_exchg_merid(ny,nz,wm,1)
!-----------------------------------------------------------------------
!        Eddy energy
!-----------------------------------------------------------------------
      eke=0.0;bws=0.0;epe=0.0;
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
          eke(j,k)=eke(j,k)+dx*( ( 0.5*((u(i  ,j,k,tau)-um(j  ,k))*maskU(i  ,j,k)   &
                  +(u(i-1,j,k,tau)-um(j  ,k))*maskU(i-1,j,k) ) )**2/2   &
           +( 0.5*((v(i,j-1,k,tau)-vm(j-1,k))*maskV(i,j-1,k)   &
                  +(v(i,j-1,k,tau)-vm(j  ,k))*maskV(i,j  ,k) ) )**2/2  )
          epe(j,k)=epe(j,k)+dx*bss(i,j,k)**2/2/N_0**2

          bws(j,k)=bws(j,k)+dx*( (w(i,j,k  ,tau)-wm(j,k  ))*0.25*maskW(i,j,k  )*  &
               (bss(i,j,k  )+bss(i,j,k+1)) +(w(i,j,k-1,tau)-wm(j,k-1))*0.25*maskW(i,j,k-1)*  &
               (bss(i,j,k-1)+bss(i,j,k  ))    )
         enddo 
        enddo 
      enddo 
      where (blen/=0); eke=eke/blen; elsewhere; eke=spval; end where
      where (blen/=0); epe=epe/blen; elsewhere; epe=spval; end where
      where (blen/=0); bws=bws/blen; elsewhere; bws=spval; end where
!-----------------------------------------------------------------------
!        Baroclinic instability term
!-----------------------------------------------------------------------
      a=0.0 ! v'b' bar b_y /N_0**2 
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
          a(j,k)=a(j,k)+dx/N_0**2*(  (v(i,j,k,tau)-vm(j,k))*  &
            0.5*(bss(i,j,k)+bss(i,j+1,k))*(bms(j+1,k)-bms(j,k))/dx    )*maskV(i,j,k)
         enddo 
        enddo 
      enddo 
      where (vlen/=0) a=a/vlen
      call border_exchg_merid(ny,nz,a,1)
      do j=js,je 
        do k=2,nz-1
         T2(j,k)=0.5*( a(j-1,k)+a(j,k))
        enddo 
      enddo 

      a=0.0 ! w'b' bar b_z /N_0**2 
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
          a(j,k)=a(j,k)+dx/N_0**2*(  (w(i,j,k,tau)-wm(j,k))*   &
            0.5*(bss(i,j,k)+bss(i,j,k+1))*( bms(j,k+1)-bms(j,k))/dz     )*maskW(i,j,k)
         enddo 
        enddo 
      enddo 
      where (wlen/=0) a=a/wlen
      do j=js,je 
        do k=2,nz-1
         T2(j,k)=T2(j,k)+0.5*(a(j,k-1)+a(j,k))
        enddo 
      enddo 
      where (blen==0); T2 =spval; end where

!-----------------------------------------------------------------------
!       Shear instability terms
!-----------------------------------------------------------------------
      a=0.0; 
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1 ! bar(u' v') bar u_y    
          a(j,k)=a(j,k)+dx*(  0.5*( u(i,j,k,tau)-um(j,k) + u(i,j+1,k,tau)-um(j+1,k))   &
             *0.5*( v(i,j,k,tau)-vm(j,k) + v(i+1,j,k,tau)-vm(j,k))   &
             *(um(j+1,k)-um(j,k))/dx   )*maskU(i,j,k)*maskU(i,j+1,k)
         enddo 
        enddo 
      enddo 
      where (ulen/=0) a=a/ulen
      call border_exchg_merid(ny,nz,a,1)
      do j=js,je 
        do k=2,nz-1
         T4u(j,k)=0.5*( a(j-1,k)+a(j,k))
        enddo 
      enddo 

      a=0.0; 
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1 ! bar(u' w') bar u_z
          a(j,k)=a(j,k)+dx*(   0.5*( u(i,j,k,tau)-um(j,k) + u(i,j,k+1,tau)-um(j,k+1))   &
             *0.5*( w(i,j,k,tau)-wm(j,k) + w(i+1,j,k,tau)-wm(j,k))   &
             *(um(j,k+1)-um(j,k))/dz   )*maskU(i,j,k)*maskU(i,j,k+1)
         enddo 
        enddo 
      enddo 
      where (ulen/=0) a=a/ulen
      do j=js,je 
        do k=2,nz-1
         T4u(j,k)=T4u(j,k)+0.5*( a(j,k)+a(j,k-1))
        enddo 
      enddo 
!       where (blen==0); T4u=spval; end where

      a=0.0; 
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1 ! bar(v' v') bar v_y
          a(j,k)=a(j,k)+dx*(  0.5*( v(i,j,k,tau)-vm(j,k) + v(i,j-1,k,tau)-vm(j-1,k))   &
             *0.5*( v(i,j,k,tau)-vm(j,k) + v(i,j-1,k,tau)-vm(j-1,k))   &
             *(vm(j,k)-vm(j-1,k))/dx    )*maskV(i,j,k)*maskV(i,j-1,k)
         enddo 
        enddo 
      enddo 
      where (vlen/=0) a=a/vlen
      do j=js,je 
        do k=2,nz-1
         T4v(j,k)=a(j,k)
        enddo 
      enddo 

      a=0.0; 
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1 ! bar(v' w') bar v_z
          a(j,k)=a(j,k)+dx*(  0.5*( v(i,j,k,tau)-vm(j,k) + v(i,j,k+1,tau)-vm(j,k+1))   &
             *0.5*( w(i,j,k,tau)-wm(j,k) + w(i,j+1,k,tau)-wm(j+1,k))   &
             *(vm(j,k+1)-vm(j,k))/dz     )*maskV(i,j,k)*maskV(i,j,k+1)
         enddo 
        enddo 
      enddo 
      where (vlen/=0) a=a/vlen
      call border_exchg_merid(ny,nz,a,1)
      do j=js,je 
        do k=2,nz-1
         T4v(j,k)=T4v(j,k)+(a(j,k)+a(j,k-1)+a(j-1,k)+a(j-1,k-1))/4.0
        enddo 
      enddo 
      T4u=T4u+T4v
      where (blen==0); T4u=spval; end where
      where (blen==0); T4v=spval; end where
!-----------------------------------------------------------------------
!        Dissipation of EKE:  M'_u**2   M'_v**2
!-----------------------------------------------------------------------
      eps_kh = 0.
      if (enable_biharmonic_friction) then
        do j=js,je 
         do k=2,nz-1
          do i=2,nx -1
            us(i,j,k)=(u(i,j,k,tau)-um(j,k))*maskU(i,j,k)
          enddo 
         enddo 
        enddo 
        call border_exchg3D(nx,ny,nz,us,1); call setcyclic3D(nx,ny,nz,us)
        call biha_diss_u(us,del2)
        eps_u=0.0;
        do j=js,je 
         do k=2,nz-1
          do i=2,nx -1
           eps_u(j,k)=eps_u(j,k)+dx*del2(i,j,k)**2*maskU(i,j,k)
          enddo 
         enddo 
        enddo 
        do j=js,je 
         do k=2,nz-1
          do i=2,nx -1
            us(i,j,k)=(v(i,j,k,tau)-vm(j,k))*maskV(i,j,k)
          enddo 
         enddo 
        enddo 
        call border_exchg3D(nx,ny,nz,us,1); call setcyclic3D(nx,ny,nz,us)
        call biha_diss_v(us,del2)
        eps_v=0.0;
        do j=js,je 
         do k=2,nz-1
          do i=2,nx -1
           eps_v(j,k)=eps_v(j,k)+dx*del2(i,j,k)**2*maskV(i,j,k)
          enddo 
         enddo 
        enddo 
        call border_exchg_merid(ny,nz,eps_v,1)
        eps_kh=0
        do j=js,je 
         do k=2,nz-1
          eps_kh(j,k)=eps_u(j,k)+0.5*(eps_v(j,k)+eps_v(j-1,k))
         enddo 
        enddo 
      endif

!        add here dissipation by harmonic friction

      where (blen/=0); eps_kh=eps_kh/blen; 
      elsewhere; eps_kh=spval; end where
!-----------------------------------------------------------------------
!        Dissipation of EKE:  A_v (u'_z)**2  +A_v (v'_z)**2 
!-----------------------------------------------------------------------
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
           us(i,j,k)=(u(i,j,k,tau)-um(j,k))*maskU(i,j,k)
         enddo 
        enddo 
      enddo 
      call border_exchg3D(nx,ny,nz,us,1); call setcyclic3D(nx,ny,nz,us)
      eps_u=0.0;
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
          eps_u(j,k)=eps_u(j,k)+dx*A_v*( 0.5*(us(i,j,k+1)-us(i,j,k  ))/dz*maskW(i,j,k  )   &
           +0.5*(us(i,j,k  )-us(i,j,k-1))/dz*maskW(i,j,k-1)    )**2         
         enddo 
        enddo 
      enddo 
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
           us(i,j,k)=(v(i,j,k,tau)-vm(j,k))*maskV(i,j,k)
         enddo 
        enddo 
      enddo 
      call border_exchg3D(nx,ny,nz,us,1); call setcyclic3D(nx,ny,nz,us)
      eps_v=0.0;
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
          eps_v(j,k)=eps_v(j,k)+dx*A_v*( 0.5*(us(i,j,k+1)-us(i,j,k  ))/dz*maskW(i,j,k  )   &
           +0.5*(us(i,j,k  )-us(i,j,k-1))/dz*maskW(i,j,k-1)   )**2         
         enddo 
        enddo 
      enddo 
      call border_exchg_merid(ny,nz,eps_v,1)
      eps_kv=0
      do j=js,je 
        do k=2,nz-1
         eps_kv(j,k)=eps_u(j,k)+0.5*(eps_v(j,k)+eps_v(j-1,k))
        enddo 
      enddo 
!
!        Dissipation of EKE by bottom stress
!
      if (enable_bottom_stress) then
        eps_u=0.;eps_v=0.
        do j=js,je 
         do i=2,nx -1
          k=max(1,k_bottom_u(i,j))
          eps_u(j,k)=eps_u(j,k)    +dx*cdbot*(u(i,j,k,tau)-um(j,k))**2*masku(i,j,k)
          k=max(1,k_bottom_v(i,j))
          eps_v(j,k)=eps_v(j,k)   +dx*cdbot*(v(i,j,k,tau)-vm(j,k))**2*maskV(i,j,k)
         enddo
        enddo
        do j=js,je 
         do k=2,nz-1
          eps_kv(j,k)=eps_kv(j,k)    +eps_u(j,k)+0.5*(eps_v(j,k)+eps_v(j-1,k))
         enddo 
        enddo 
      endif
!
!        Dissipation of EKE by interior stress
!
      if (enable_interior_stress) then
        eps_u=0.;eps_v=0.
        do k=2,nz-1
         do j=js,je 
          do i=2,nx -1
           eps_u(j,k)=eps_u(j,k)   +dx*cdint*(u(i,j,k,tau)-um(j,k))**2*masku(i,j,k)
           eps_v(j,k)=eps_v(j,k)   +dx*cdint*(v(i,j,k,tau)-vm(j,k))**2*maskV(i,j,k)
          enddo
         enddo
        enddo
        do k=2,nz-1
         do j=js,je 
           eps_kv(j,k)=eps_kv(j,k)+     eps_u(j,k)+0.5*(eps_v(j,k)+eps_v(j-1,k))
         enddo 
        enddo 
      endif

      where (blen/=0); eps_kv=eps_kv/blen; 
      elsewhere; eps_kv=spval; end where
!-----------------------------------------------------------------------
!        Dissipation of EPE:  K_v (b'_z)**2 /N**2 +  M'_b**2
!-----------------------------------------------------------------------
      eps_b=0.0;
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
          eps_b(j,k)=eps_b(j,k)+dx/N_0**2*( +0.5*(K_v+K_b(i,j,k  ))   &
             *(bss(i,j,k+1)-bss(i,j,k))/dz *(b(i,j,k+1,tau)-b(i,j,k,tau))/dz*maskW(i,j,k)    &
           +0.5*(K_v+K_b(i,j,k-1)) *(bss(i,j,k)-bss(i,j,k-1))/dz   &
             *(b(i,j,k,tau)-b(i,j,k-1,tau))/dz*maskW(i,j,k-1)     )
!     &       (K_v+0.5*(K_b(i,j,k)+K_b(i,j,k-1)) )
!     &        *( (b(i,j,k+1,tau)-b(i,j,k-1,tau))/(2*dz) )**2
         enddo 
        enddo 
      enddo 
      if (enable_biharmonic_diffusion) then
        call biha_diss_b(bss,del2)
        do j=js,je 
         do k=2,nz-1
          do i=2,nx -1
           eps_b(j,k)=eps_b(j,k)+dx/N_0**2*del2(i,j,k)**2*maskT(i,j,k)
          enddo 
         enddo 
        enddo 
      endif
      where (blen/=0); eps_b=eps_b/blen; 
      elsewhere; eps_b=spval; end where
!-----------------------------------------------------------------------
!        Dissipation of EPE due to buffer zones
!-----------------------------------------------------------------------
      rest =0.0;
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
          rest (j,k)=rest (j,k)+dx/N_0**2*( bss(i,j,k)*b(i,j,k,tau)*t_rest(j)*maskT(i,j,k)  )
         enddo 
        enddo 
      enddo 
      where (blen/=0); rest =rest /blen; 
      elsewhere; rest =spval; end where

      restm =0.0;
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
          restm(j,k)=restm(j,k)+dx/N_0**2*(  bms(j,k)*b(i,j,k,tau)*t_rest(j)*maskT(i,j,k)  )
         enddo 
        enddo 
      enddo 
      where (blen/=0); restm=restm/blen; 
      elsewhere; restm=spval; end where
!-----------------------------------------------------------------------
!        EKE radiation  vp wp
!-----------------------------------------------------------------------
      vps=0.0; wps=0.0;
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
          vps(j,k)=vps(j,k)+dx*( ( v(i,j,k,tau)-vm(j,k) )*  &
             0.5*(p_full(i,j,k,tau)+p_full(i,j+1,k,tau))*maskV(i,j,k)     )
          wps(j,k)=wps(j,k)+dx*( ( w(i,j,k,tau)-wm(j,k) )*  &
             0.5*(p_full(i,j,k,tau)+p_full(i,j,k+1,tau))*maskW(i,j,k)     )
         enddo 
        enddo 
      enddo 
      where (vlen/=0); vps  =vps  /vlen; 
      elsewhere; vps  =spval; end where
      where (wlen/=0); wps  =wps  /wlen; 
      elsewhere; wps  =spval; end where
!-----------------------------------------------------------------------
!        EKE advection
!-----------------------------------------------------------------------
      del2=0 
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
           del2(i,j,k)= ( 0.5*((u(i  ,j,k,tau)-um(j  ,k))*maskU(i  ,j,k)   &
                  +(u(i-1,j,k,tau)-um(j  ,k))*maskU(i-1,j,k) ) )**2/2   &
           +( 0.5*((v(i,j-1,k,tau)-vm(j-1,k))*maskV(i,j-1,k)   &
                  +(v(i,j-1,k,tau)-vm(j  ,k))*maskV(i,j  ,k) ) )**2/2
         enddo 
        enddo 
      enddo 
      call border_exchg3D(nx,ny,nz,del2,1); call setcyclic3D(nx,ny,nz,del2)

      veke=0.0; weke=0.0;
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
          veke(j,k)=veke(j,k)+dx*(v(i,j,k,tau)*maskV(i,j,k)*   &
             0.5*(del2(i,j,k)+del2(i,j+1,k))*maskV(i,j,k)     )
          weke(j,k)=weke(j,k)+dx*(  w(i,j,k,tau)*maskW(i,j,k)*  &
             0.5*(del2(i,j,k)+del2(i,j,k+1))*maskW(i,j,k)         )
         enddo 
        enddo 
      enddo 
      where (vlen/=0); veke =veke /vlen; 
      elsewhere; veke =spval; end where
      where (wlen/=0); weke =weke /wlen; 
      elsewhere; weke =spval; end where
!-----------------------------------------------------------------------
!        EPE advection
!-----------------------------------------------------------------------
      del2=0 
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
           del2(i,j,k)= bss(i,j,k)**2/2/N_0**2
         enddo 
        enddo 
      enddo 
      call border_exchg3D(nx,ny,nz,del2,1); call setcyclic3D(nx,ny,nz,del2)

      vepe=0.0; wepe=0.0;
      do j=js,je 
        do k=2,nz-1
         do i=2,nx -1
          vepe(j,k)=vepe(j,k)+dx*(  v(i,j,k,tau)*maskV(i,j,k)*   &
             0.5*(del2(i,j,k)+del2(i,j+1,k))*maskV(i,j,k)                     )
          wepe(j,k)=wepe(j,k)+dx*( w(i,j,k,tau)*maskW(i,j,k)*   &
             0.5*(del2(i,j,k)+del2(i,j,k+1))*maskW(i,j,k)          )
         enddo 
        enddo 
      enddo 
      where (vlen/=0); vepe =vepe /vlen; 
      elsewhere; vepe =spval; end where
      where (wlen/=0); wepe =wepe /wlen; 
      elsewhere; wepe =spval; end where
 end subroutine zonal_averages_energy




 subroutine biha_diss_b(var,del2)
      use pyOM_module   
      implicit none
      integer :: i,j,k,js,je
      real*8 :: var(nx ,ny ,nz) 
      real*8 :: diff_fe(nx ,ny ,nz), diff_fn(nx ,ny ,nz) 
      real*8 :: del2(nx ,ny ,nz),diff_tx,diff_ty,diffx
      DIFF_Tx(i,j,k)=(diff_fe(i,j,k)-diff_fe(i-1,j,k))/dx
      DIFF_Ty(i,j,k)=(diff_fn(i,j,k)-diff_fn(i,j-1,k))/dx

      js=max(2,js_pe); je = min(je_pe,ny -1)
      diffx = sqrt(abs(Khbi)); diff_fe=0.0;diff_fn=0.0
      do k=1,nz
       do j=js,je
        do i=1,nx -1
         diff_fe(i,j,k)=diffx*(var(i+1,j,k)-var(i,j,k))/dx     *maskU(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,diff_fe,1); call setcyclic3D(nx,ny,nz,diff_fe)
      do k=1,nz-1
       do j=js,je
        do i=1,nx 
         diff_fn(i,j,k)=diffx*(var(i,j+1,k)-var(i,j,k))/dx       *maskV(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,diff_fn,1); call setcyclic3D(nx,ny,nz,diff_fn)
      del2=0.0
      do k=2,nz-1
       do j=js,je
        do i=2,nx -1
         del2(i,j,k) = (DIFF_Tx(i,j,k) + DIFF_Ty(i,j,k))*maskT(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,del2,1); call setcyclic3D(nx,ny,nz,del2)
  end subroutine biha_diss_b




 subroutine biha_diss_u (var,del2)
      use pyOM_module   
      implicit none
      integer :: i,j,k,js,je
      real*8 :: var(nx ,ny ,nz) 
      real*8 :: diff_fe(nx ,ny ,nz), diff_fn(nx ,ny ,nz) 
      real*8 :: del2(nx ,ny ,nz),diff_tx,diff_ty,diffx
      DIFF_Tx(i,j,k) = (diff_fe(i,j,k)-diff_fe(i-1,j,k))/dx
      DIFF_Ty(i,j,k) = (diff_fn(i,j,k)-diff_fn(i,j-1,k))/dx

      js=max(2,js_pe); je = min(je_pe,ny -1)
      diffx = sqrt(abs(ahbi)); diff_fe=0.0;diff_fn=0.0
      do k=1,nz
       do j=js,je
        do i=1,nx -1
         diff_fe(i,j,k)=diffx*(var(i+1,j,k)-var(i,j,k))/dx
        enddo
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,diff_fe,1); call setcyclic3D(nx,ny,nz,diff_fe)
      do k=1,nz-1
       do j=js,je
        do i=1,nx 
         diff_fn(i,j,k)=diffx*(var(i,j+1,k)-var(i,j,k))/dx     *maskU(i,j+1,k)*maskU(i,j,k)
        enddo
       enddo
      enddo
      if (enable_noslip) then
       do j=js-1,je
         diff_fn(:,j,:)=diff_fn(:,j,:)-2*diffx*var(:,j,:)/dx *(1-maskU(:,j+1,:))*maskU(:,j,:)   &
                                      +2*diffx*var(:,j+1,:)/dx *(1-maskU(:,j,:))*maskU(:,j+1,:)
       enddo
      endif
      call border_exchg3D(nx,ny,nz,diff_fn,1); call setcyclic3D(nx,ny,nz,diff_fn)
      del2=0.0
      do k=2,nz-1
       do j=js,je
        do i=2,nx -1
         del2(i,j,k) = (DIFF_Tx(i,j,k) + DIFF_Ty(i,j,k))*maskU(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,del2,1); call setcyclic3D(nx,ny,nz,del2)
 end subroutine biha_diss_u




 subroutine biha_diss_v (var,del2)
      use pyOM_module   
      implicit none
      integer :: i,j,k,js,je
      real*8 :: var(nx ,ny ,nz) 
      real*8 :: diff_fe(nx ,ny ,nz), diff_fn(nx ,ny ,nz) 
      real*8 :: del2(nx ,ny ,nz),diff_tx,diff_ty,diffx
      DIFF_Tx(i,j,k)=(diff_fe(i,j,k)-diff_fe(i-1,j,k))/dx
      DIFF_Ty(i,j,k)=(diff_fn(i,j,k)-diff_fn(i,j-1,k))/dx

      js=max(2,js_pe); je = min(je_pe,ny -1)
      diffx = sqrt(abs(ahbi)); diff_fe=0.0;diff_fn=0.0
      do k=1,nz
       do j=js,je
        do i=1,nx -1
         diff_fe(i,j,k)=diffx*(var(i+1,j,k)-var(i,j,k))/dx  *maskV(i,j,k)*maskV(i+1,j,k)
        enddo
       enddo
      enddo
      if (enable_noslip) then
       do i=1,nx -1
         diff_fe(i,:,:)=diff_fe(i,:,:)-2*diffx*var(i,:,:)/dx    *(1-maskV(i+1,:,:))*maskV(i,:,:)  &
                                      +2*diffx*var(i+1,:,:)/dx  *(1-maskV(i,:,:))*maskV(i+1,:,:)
       enddo
      endif
      call border_exchg3D(nx,ny,nz,diff_fe,1); call setcyclic3D(nx,ny,nz,diff_fe)
      do k=1,nz-1
       do j=js,je
        do i=1,nx 
         diff_fn(i,j,k)=diffx*(var(i,j+1,k)-var(i,j,k))/dx
        enddo
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,diff_fn,1); call setcyclic3D(nx,ny,nz,diff_fn)
      del2=0.0
      do k=2,nz-1
       do j=js,je
        do i=2,nx -1
         del2(i,j,k) = (DIFF_Tx(i,j,k) + DIFF_Ty(i,j,k))*maskV(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,del2,1); call setcyclic3D(nx,ny,nz,del2)
 end subroutine biha_diss_v


