    

!=======================================================================
!       Bolus velocity and isopycnal diffusion
!       linked in the code in pyOM.f90
!=======================================================================


 module bolus_module
!-----------------------------------------------------------------------
!     module for bolus velocities
!-----------------------------------------------------------------------
      implicit none
      real*8, allocatable :: ub(:,:,:)
      real*8, allocatable :: vb(:,:,:)
      real*8, allocatable :: wb(:,:,:)
      real*8, allocatable :: psix(:,:,:)
      real*8, allocatable :: psiy(:,:,:)
      real*8, allocatable :: ahthk(:,:,:)
      real*8, parameter :: epsln = 1.0d-20 ! a small parameter
      logical :: enable_bolus_diag                     = .false. ! write output to netcdf file
 end module bolus_module
 

 subroutine init_bolus
!-----------------------------------------------------------------------
!     initialize module
!-----------------------------------------------------------------------
      use pyOM_module   
      use bolus_module
      implicit none
      real*8 :: fxa

      if (my_pe==0) print*,' Initializing bolus velocity module '
      allocate( ub(nx ,ny ,nz) ); ub=0.
      allocate( vb(nx ,ny ,nz) ); vb=0.
      allocate( wb(nx ,ny ,nz) ); wb=0.
      allocate( psix(nx ,ny ,nz) ); psix=0.
      allocate( psiy(nx ,ny ,nz) ); psiy=0.
      allocate( ahthk(nx ,ny ,nz) ); ahthk=0.
      ahthk = K_gm

 if (enable_bolus_diag) then
      if (my_pe==0) call init_bolus_diag
 else
      if (my_pe==0) print*,' no diagnostic output '
 endif
!
!-----------------------------------------------------------------------
!     Compute the grid factors which set the maximum slopes available
!     for the mixing schemes. 
!-----------------------------------------------------------------------
!
      fxa = dz*dx/(4*K_gm*dt+epsln)
      if (my_pe==0) then
       print*,' steepest slope possible for linear stab = ',fxa
       print*,' critical slope specified as               ',iso_slopec
      print*,' K_gm  = ',k_gm,' m^2/s'
      endif
      if (my_pe==0) print*,' done initializing GM module '
 end subroutine init_bolus




 subroutine bolus_velocity
!-----------------------------------------------------------------------
!     Verteiler funktion
!-----------------------------------------------------------------------
      use pyOM_module   
      use bolus_module
      use fcontrol_module
      implicit none
      integer :: n

      call calculate_bolus
      call add_bolus(b)
      if (enable_back_state) call add_bolus_back

      if (enable_diag_tracer ) then
       do n=1,nt
        call add_bolus(tr(:,:,:,:,n))
       enddo
      endif
      if (enable_bolus_diag .and. mod(itt,int(snapint/dt))  == 0)  then
        call diag_bolus
      endif
 end subroutine bolus_velocity


! subroutine bolus_set_ahthk(diff)
!-----------------------------------------------------------------------
!     set thickness diffusivity from other module
!-----------------------------------------------------------------------
!      use pyOM_module   
!      use bolus_module
!      implicit none
!      real*8 :: diff(nx ,ny ,nz)
!      call border_exchg3D(nx,ny,nz,ahthk,1); 
!      call setcyclic3D(nx,ny,nz,ahthk)
! end subroutine bolus_set_ahthk



 subroutine calculate_bolus
!-----------------------------------------------------------------------
!     calculate bolus velocities using diffusivity ahthk
!-----------------------------------------------------------------------
      use pyOM_module   
      use bolus_module
      implicit none
      integer :: i,j,k,js,je
      real*8 :: fxa,fxb,sx(nx ,ny ,nz),bz(nx ,ny ,nz),sy(nx ,ny ,nz),a(nx ,ny ,nz)

      js=max(2,js_pe); je = min(je_pe,ny -1)
!-----------------------------------------------------------------------
!     vertical derivative of buoyancy
!-----------------------------------------------------------------------
      bz(:,js_pe:je_pe,:)=0
      do k=1,nz-1
       do j=js,je
        do i=2,nx -1
         bz(i,j,k)= (b(i,j,k+1,taum1)-b(i,j,k,taum1))/dz*maskW(i,j,k) 
        enddo
       enddo
      enddo
      if (enable_back_state) then  
       do k=1,nz-1
        do j=js,je
         do i=2,nx -1
          bz(i,j,k)= bz(i,j,k)+ (back(i,j,k+1,taum1)-back(i,j,k,taum1))/dz*maskW(i,j,k) 
         enddo
        enddo
       enddo
      endif
      call border_exchg3D(nx,ny,nz,bz,1); call setcyclic3D(nx,ny,nz,bz)
!-----------------------------------------------------------------------
!     zonal component of streamfunction
!-----------------------------------------------------------------------
      a(:,js_pe:je_pe,:)=0
      sx(:,js_pe:je_pe,:)=0
      psix(:,js_pe:je_pe,:)=0
      do k=2,nz-1
       do j=js,je
        do i=2,nx -1
         a(i,j,k)=(b(i+1,j,k,taum1)-b(i,j,k,taum1))/dx*maskU(i,j,k)
        enddo
       enddo
      enddo
      if (enable_back_state) then  
       do k=2,nz-1
        do j=js,je
         do i=2,nx -1
          a(i,j,k)=a(i,j,k)+ (back(i+1,j,k,taum1)-back(i,j,k,taum1))/dx*maskU(i,j,k)
         enddo
        enddo
       enddo
      endif
      do k=2,nz-1
       do j=js,je
        do i=2,nx -1
         fxa=(a(i,j,k+1)+a(i,j,k))/(maskU(i,j,k)+maskU(i,j,k+1)+epsln)
         fxb=(bz(i+1,j,k)+bz(i,j,k))/(maskW(i+1,j,k)+maskW(i,j,k)+epsln)
         fxa=fxa*sign(dble(1.),fxb)/(abs(fxb)+epsln)
!         sx(i,j,k)=fxa*0.5*(1.+tanh((-abs(fxa)+iso_slopec)/iso_dslope))   
         sx(i,j,k)= max(min(fxa,iso_slopec),-iso_slopec)
        enddo
       enddo
      enddo
      do k=2,nz-1
       do j=js,je
        do i=2,nx -1
         fxa=ahthk(i,j,k)+ahthk(i+1,j,k)+ahthk(i,j,k+1)+ahthk(i+1,j,k+1)
         psix(i,j,k)=fxa/4.*sx(i,j,k)*maskW(i,j,k)*maskU(i,j,k)  *maskU(i,j,k+1)*maskW(i+1,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,psix,1); call setcyclic3D(nx,ny,nz,psix)
!-----------------------------------------------------------------------
!     meridional component of streamfunction
!-----------------------------------------------------------------------
      a(:,js_pe:je_pe,:)=0
      sy(:,js_pe:je_pe,:)=0
      psiy(:,js_pe:je_pe,:)=0
      do k=2,nz-1
       do j=js,je
        do i=2,nx -1
         a(i,j,k)=(b(i,j+1,k,taum1)-b(i,j,k,taum1))/dx*maskV(i,j,k)
        enddo
       enddo
      enddo
      if (enable_back_state) then  
       do k=2,nz-1
        do j=js,je
         do i=2,nx -1
          a(i,j,k)=a(i,j,k)+(back(i,j+1,k,taum1)-back(i,j,k,taum1))/dx*maskV(i,j,k)
         enddo
        enddo
       enddo
      endif
      do k=2,nz-1
       do j=js,je
        do i=2,nx -1
         fxa=(a(i,j,k+1)+a(i,j,k))/(maskV(i,j,k)+maskV(i,j,k+1)+epsln)
         fxb=(bz(i,j+1,k)+bz(i,j,k))/(maskW(i,j+1,k)+maskW(i,j,k)+epsln)
         fxa=fxa*sign(dble(1.),fxb)/(abs(fxb)+epsln)
!         sy(i,j,k)=fxa*0.5*(1.+tanh((-abs(fxa)+iso_slopec)/iso_dslope))   
         sy(i,j,k)= max(min(fxa,iso_slopec),-iso_slopec)
        enddo
       enddo
      enddo
      do k=2,nz-1
       do j=js,je
        do i=2,nx -1
         fxa=ahthk(i,j,k)+ahthk(i,j+1,k)+ahthk(i,j,k+1)+ahthk(i,j+1,k+1)
         psiy(i,j,k)=fxa/4.0*sy(i,j,k)*maskW(i,j,k)*maskV(i,j,k)  *maskV(i,j,k+1)*maskW(i,j+1,k)
        enddo
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,psiy,1); call setcyclic3D(nx,ny,nz,psiy)
!-----------------------------------------------------------------------
!        bolus velocities
!-----------------------------------------------------------------------
      do k=2,nz-1
       do j=js,je
        do i=2,nx -1
         ub(i,j,k)=(psix(i,j,k)-psix(i,j,k-1))/dz*maskU(i,j,k)
         vb(i,j,k)=(psiy(i,j,k)-psiy(i,j,k-1))/dz*maskV(i,j,k)
         wb(i,j,k)=(-(psiy(i,j,k)-psiy(i,j-1,k))/dx  -(psix(i,j,k)-psix(i-1,j,k))/dx)*maskW(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,ub,1); call setcyclic3D(nx,ny,nz,ub)
      call border_exchg3D(nx,ny,nz,vb,1); call setcyclic3D(nx,ny,nz,vb)
      call border_exchg3D(nx,ny,nz,wb,1); call setcyclic3D(nx,ny,nz,wb)
 end subroutine calculate_bolus


 subroutine add_bolus(tr_)
!-----------------------------------------------------------------------
!     add divergence of advective eddy fluxes to tracer tr_
!-----------------------------------------------------------------------
      use pyOM_module   
      use bolus_module
      implicit none
      real*8 :: tr_(nx ,ny ,nz,3)
      real*8 :: adv_fe(nx ,ny ,nz), adv_ft(nx ,ny ,nz)
      real*8 :: adv_fn(nx ,ny ,nz)
      integer :: i,j,k,js,je


      js=max(2,js_pe); je = min(je_pe,ny -1)
      do k=2,nz-1
       do j=js,je
        do i=1,nx -1
         adv_fe(i,j,k)= 0.5*(tr_(i,j,k,tau)+tr_(i+1,j,k,tau) )*ub(i,j,k)*maskU(i,j,k)
        enddo
       enddo
      enddo
      adv_fn(:,1,:)=0
      do k=2,nz-1
       do j=js,je
        do i=2,nx -1
         adv_fn(i,j,k)=0.5*(tr_(i,j,k,tau)+tr_(i,j+1,k,tau) )*vb(i,j,k)*maskV(i,j,k)
        enddo
       enddo
      enddo
      do k=1,nz-1
       do j=js,je
        do i=2,nx -1
         adv_ft(i,j,k)=0.5*(tr_(i,j,k,tau)+tr_(i,j,k+1,tau) )*wb(i,j,k)*maskW(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,adv_fn,1)
      call setcyclic3D(nx,ny,nz,adv_fe)
      call setcyclic3D(nx,ny,nz,adv_fn)
      do k=2,nz-1
       do j=js,je
        do i=2,nx -1
         tr_(i,j,k,taup1)=tr_(i,j,k,taup1)+maskT(i,j,k)*c2dt*(    &
                          -(adv_fe(i,j,k)-adv_fe(i-1,j,k))/dx   &
                          -(adv_fn(i,j,k)-adv_fn(i,j-1,k))/dx   &
                          -(adv_ft(i,j,k)-adv_ft(i,j,k-1))/dz )
        enddo
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,tr_(:,:,:,taup1),2);
      call setcyclic3D(nx,ny,nz,tr_(:,:,:,taup1) )
 end subroutine add_bolus




 subroutine add_bolus_back
!-----------------------------------------------------------------------
!     add advection of background density
!-----------------------------------------------------------------------
      use pyOM_module   
      use bolus_module
      implicit none
      real*8 :: adv_fe(nx ,ny ,nz), adv_ft(nx ,ny ,nz), adv_fn(nx ,ny ,nz)
      integer :: i,j,k,js,je

      js=max(2,js_pe); je = min(je_pe,ny -1)
      do k=2,nz-1
       do j=js,je
        do i=1,nx -1
         adv_fe(i,j,k)= 0.5*(back(i,j,k,tau)+back(i+1,j,k,tau) )*ub(i,j,k)*maskU(i,j,k)
        enddo
       enddo
      enddo
      adv_fn(:,1,:)=0
      do k=2,nz-1
       do j=js,je
        do i=2,nx -1
         adv_fn(i,j,k)=0.5*(back(i,j,k,tau)+back(i,j+1,k,tau) )*vb(i,j,k)*maskV(i,j,k)
        enddo
       enddo
      enddo
      do k=1,nz-1
       do j=js,je
        do i=2,nx -1
         adv_ft(i,j,k)=0.5*(back(i,j,k,tau)+back(i,j,k+1,tau) )*wb(i,j,k)*maskW(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,adv_fn,1)
      call setcyclic3D(nx,ny,nz,adv_fe)
      call setcyclic3D(nx,ny,nz,adv_fn)
      do k=2,nz-1
       do j=js,je
        do i=2,nx -1
         b(i,j,k,taup1)=b(i,j,k,taup1)+maskT(i,j,k)*c2dt*(    &
                          -(adv_fe(i,j,k)-adv_fe(i-1,j,k))/dx   &
                          -(adv_fn(i,j,k)-adv_fn(i,j-1,k))/dx   &
                          -(adv_ft(i,j,k)-adv_ft(i,j,k-1))/dz )
        enddo
       enddo
      enddo
      call border_exchg3D(nx,ny,nz,b(:,:,:,taup1),2);
      call setcyclic3D(nx,ny,nz,b(:,:,:,taup1) )
 end subroutine add_bolus_back







 subroutine init_bolus_diag
!-----------------------------------------------------------------------
!     initialize NetCDF snapshot file
!-----------------------------------------------------------------------
      use pyOM_module   
      use bolus_module
      implicit none
      include "netcdf.inc"
      integer :: ncid,iret
      integer :: lon_tdim,lon_udim,z_tdim,z_udim,itimedim
      integer :: lat_tdim,lat_udim,varid
      integer :: dims(4)
      character :: name*24, unit*16
      real*8, parameter :: spval = -1.0d33

      call def_grid_cdf('bolus.cdf')
      iret=nf_open('bolus.cdf',NF_WRITE,ncid)
      iret=nf_set_fill(ncid, NF_NOFILL, iret)
      call ncredf(ncid, iret)
      iret=nf_inq_dimid(ncid,'xt',lon_tdim)
      iret=nf_inq_dimid(ncid,'xu',lon_udim)
      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)

!     2 dim variables on t grid
      dims = (/Lon_tdim,lat_udim, z_udim, iTimedim/)
      varid  = ncvdef (ncid,'psiy', NCFLOAT,4,dims,iret)
      name = 'Meridional streamfct'; unit = 'm^2/s'
      call dvcdf(ncid,varid,name,24,unit,16,spval)

      dims = (/Lon_udim,lat_tdim, z_udim, iTimedim/)
      varid  = ncvdef (ncid,'psix', NCFLOAT,4,dims,iret)
      name = 'Zonal streamfct'; unit = 'm^2/s'
      call dvcdf(ncid,varid,name,24,unit,16,spval)

      dims = (/Lon_tdim,lat_tdim, z_tdim, iTimedim/)
      varid  = ncvdef (ncid,'ahthk', NCFLOAT,4,dims,iret)
      name = 'thickness diffusivity'; unit = 'm^2/s'
      call dvcdf(ncid,varid,name,24,unit,16,spval)

      call ncclos (ncid, iret)
 end subroutine init_bolus_diag


 subroutine diag_bolus
!-----------------------------------------------------------------------
!     write to NetCDF snapshot file
!-----------------------------------------------------------------------
      use pyOM_module   
      use bolus_module
      use fcontrol_module
      implicit none
      include "netcdf.inc"
      integer :: ncid,iret,npe, corner(4), edges(4)
      real*8 :: a(nx ,js_pe:je_pe,nz)
      integer :: itdimid,ilen,rid,itimeid
      integer :: js,je
      real*8 :: fxa,time
      real*8, parameter :: spval = -1.0d33

      js=max(2,js_pe); je = min(je_pe,ny -1)

      do npe=0,n_pes
       if (my_pe==npe) then

        iret=nf_open('bolus.cdf',NF_WRITE,ncid)
        iret=nf_set_fill(ncid, NF_NOFILL, iret)
        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 = itt*dt
         fxa = time/86400.0
         iret= nf_put_vara_double(ncid,itimeid,ilen,1,fxa)
        endif

        Corner = (/1,js_pe,1,ilen/); 
        edges  = (/nx ,je_pe-js_pe+1,nz,1/)
        iret=nf_inq_varid(ncid,'psiy',rid)
        a=psiy(:,js_pe:je_pe,:)
        iret= nf_put_vara_double(ncid,rid,corner,edges,a)

        iret=nf_inq_varid(ncid,'psix',rid)
        a=psix(:,js_pe:je_pe,:)
        iret= nf_put_vara_double(ncid,rid,corner,edges,a)

        iret=nf_inq_varid(ncid,'ahthk',rid)
        a=ahthk(:,js_pe:je_pe,:)
        iret= nf_put_vara_double(ncid,rid,corner,edges,a)

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


    
