source: codes/icosagcm/trunk/src/dissip/dissip_gcm.F90 @ 954

MODULE dissip_gcm_mod
  USE icosa
  USE abort_mod
  PRIVATE

  TYPE(t_field),POINTER,SAVE :: f_due_diss1(:)
  TYPE(t_field),POINTER,SAVE :: f_due_diss2(:)

  TYPE(t_field),POINTER,SAVE :: f_dtheta_diss(:)
  TYPE(t_field),POINTER,SAVE :: f_dtheta_rhodz_diss(:)
  TYPE(t_message),SAVE :: req_due, req_dtheta 
 
  INTEGER,SAVE :: nitergdiv=1
!$OMP THREADPRIVATE(nitergdiv)
  INTEGER,SAVE :: nitergrot=1
!$OMP THREADPRIVATE(nitergrot)
  INTEGER,SAVE :: niterdivgrad=1
!$OMP THREADPRIVATE(niterdivgrad)

  REAL,ALLOCATABLE,SAVE :: tau_graddiv(:)
!$OMP THREADPRIVATE(tau_graddiv)
  REAL,ALLOCATABLE,SAVE :: tau_gradrot(:)
!$OMP THREADPRIVATE(tau_gradrot)
  REAL,ALLOCATABLE,SAVE :: tau_divgrad(:)
!$OMP THREADPRIVATE(tau_divgrad)
  
  REAL,SAVE :: cgraddiv
!$OMP THREADPRIVATE(cgraddiv)
  REAL,SAVE :: cgradrot
!$OMP THREADPRIVATE(cgradrot)
  REAL,SAVE :: cdivgrad
!$OMP THREADPRIVATE(cdivgrad)

  INTEGER, SAVE :: rayleigh_friction_type, rayleigh_shear
!$OMP THREADPRIVATE(rayleigh_friction_type)
!$OMP THREADPRIVATE(rayleigh_shear)
  REAL, SAVE    :: rayleigh_tau, rayleigh_limlat
!$OMP THREADPRIVATE(rayleigh_tau)
!$OMP THREADPRIVATE(rayleigh_limlat)
  
  REAL,SAVE    :: dtdissip
!$OMP THREADPRIVATE(dtdissip)
 
  PUBLIC init_dissip, dissip
CONTAINS

  SUBROUTINE allocate_dissip
  USE icosa
  IMPLICIT NONE  
    CALL allocate_field(f_due_diss1,field_u,type_real,llm,ondevice=.TRUE.)
    CALL allocate_field(f_due_diss2,field_u,type_real,llm,ondevice=.TRUE.)
    CALL allocate_field(f_dtheta_diss,field_t,type_real,llm)
    CALL allocate_field(f_dtheta_rhodz_diss,field_t,type_real,llm,ondevice=.TRUE.)
    ALLOCATE(tau_graddiv(llm))
    ALLOCATE(tau_gradrot(llm))    
    ALLOCATE(tau_divgrad(llm))
    !$acc enter data create(tau_graddiv(:),tau_gradrot(:),tau_divgrad(:)) async
  END SUBROUTINE allocate_dissip
 
  SUBROUTINE init_dissip
  USE icosa
  USE disvert_mod
  USE mpi_mod
  USE mpipara
  USE transfert_mod
  USE time_mod
  USE transfert_omp_mod
  USE omp_para
  USE abort_mod
  IMPLICIT NONE
  
  TYPE(t_field),POINTER,SAVE  :: f_u(:)
  TYPE(t_field),POINTER,SAVE  :: f_du(:)
  REAL(rstd),POINTER     :: u(:)
  REAL(rstd),POINTER     :: du(:)
  TYPE(t_field),POINTER,SAVE  :: f_theta(:)
  TYPE(t_field),POINTER ,SAVE :: f_dtheta(:)
  REAL(rstd),POINTER    :: theta(:)
  REAL(rstd),POINTER    :: dtheta(:)
  REAL(rstd)            :: dumax,dumax1
  REAL(rstd)            :: dthetamax,dthetamax1
  REAL(rstd)            :: r
  REAL(rstd)            :: tau
  REAL(rstd)            :: zz, zvert, fact
  INTEGER               :: l
  CHARACTER(len=255)    :: rayleigh_friction_key
  REAL(rstd)            :: mintau
  
  ! New variables added for dissipation vertical profile (SF, 19/09/18)
  INTEGER    :: vert_prof_dissip
  REAL(rstd) :: dissip_factz,dissip_deltaz,dissip_zref,pseudoz
                        
  INTEGER :: i,j,ij,ind,it,iter,M

   rayleigh_friction_key='none'
   CALL getin("rayleigh_friction_type",rayleigh_friction_key)
   SELECT CASE(TRIM(rayleigh_friction_key))
   CASE('none')
      rayleigh_friction_type=0
      IF (is_master) PRINT *, 'No Rayleigh friction'
   CASE('dcmip2_schaer_noshear')
      CALL abort_acc("rayleigh_friction_type /= 'none'")
      rayleigh_friction_type=1
      rayleigh_shear=0
      IF (is_master) PRINT *, 'Rayleigh friction : Schaer-like mountain without shear DCMIP2.1'
   CASE('dcmip2_schaer_shear')
      CALL abort_acc("rayleigh_friction_type /= 'none'")
      rayleigh_shear=1
      rayleigh_friction_type=2
      IF (is_master) PRINT *, 'Rayleigh friction : Schaer-like mountain with shear DCMIP2.2'
   CASE('giant_liu_schneider') 
      CALL abort_acc("rayleigh_friction_type /= 'none'")
      rayleigh_friction_type=99
      IF (is_master) PRINT *, 'Rayleigh friction : giant planets Liu Schneider 2010'
   CASE DEFAULT
      IF (is_master) PRINT *, 'Bad selector : rayleigh_friction_type =', TRIM(rayleigh_friction_key), &
        ' in dissip_gcm.f90/init_dissip'
      STOP
   END SELECT

   IF(rayleigh_friction_type>0) THEN
      rayleigh_tau=0.
      CALL getin("rayleigh_friction_tau",rayleigh_tau)
      rayleigh_tau = rayleigh_tau / scale_factor
      IF(rayleigh_tau<=0) THEN
         IF (is_master) PRINT *, 'Forbidden : negative value for rayleigh_friction_tau =',rayleigh_tau
         STOP
      END IF
      IF(rayleigh_friction_type == 99) THEN
         rayleigh_limlat=0.
         CALL getin("rayleigh_limlat",rayleigh_limlat)
         rayleigh_limlat = rayleigh_limlat*3.14159/180.
      ENDIF
   END IF

    CALL allocate_dissip
    CALL allocate_field(f_u,field_u,type_real,ondevice=.TRUE.)
    CALL allocate_field(f_du,field_u,type_real,ondevice=.TRUE.)
    CALL allocate_field(f_theta,field_t,type_real,ondevice=.TRUE.)
    CALL allocate_field(f_dtheta,field_t,type_real,ondevice=.TRUE.)
    
    CALL init_message(f_due_diss1,req_e1_vect,req_due)
    CALL init_message(f_dtheta_diss,req_i1,req_dtheta)

    tau_graddiv(:)=5000
    CALL getin("tau_graddiv",tau)
    tau_graddiv(:)=tau/scale_factor

    CALL getin("nitergdiv",nitergdiv)
    
    tau_gradrot(:)=5000
    CALL getin("tau_gradrot",tau)
    tau_gradrot(:)=tau/scale_factor

    CALL getin("nitergrot",nitergrot)
    

    tau_divgrad(:)=5000
    CALL getin("tau_divgrad",tau)
    tau_divgrad(:)=tau/scale_factor

    CALL getin("niterdivgrad",niterdivgrad)


    cgraddiv=-1
    cdivgrad=-1
    cgradrot=-1

!$OMP BARRIER
!$OMP MASTER
    DO ind=1,ndomain
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      u=f_u(ind)

! set random seed to get reproductibility when using a different number of process
      CALL RANDOM_SEED(size=M)
      CALL RANDOM_SEED(put=(/(i,i=1,M)/))

      ! This cannot be ported on GPU due to compiler limitations
      DO j=jj_begin,jj_end
        DO i=ii_begin,ii_end
          ij=(j-1)*iim+i   
          CALL RANDOM_NUMBER(r)
          u(ij+u_right)=r-0.5
          CALL RANDOM_NUMBER(r)
          u(ij+u_lup)=r-0.5
          CALL RANDOM_NUMBER(r)
          u(ij+u_ldown)=r-0.5
        ENDDO
      ENDDO        
    ENDDO
!$OMP END MASTER  
!$OMP BARRIER

    DO it=1,20
      
      dumax=0
      DO iter=1,nitergdiv
        CALL update_device_field(f_u)
        CALL transfert_request(f_u,req_e1_vect)
         
        DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          u=f_u(ind)
          du=f_du(ind)
          CALL compute_gradiv_inplace(u,Ai,ne,le,de,1,1)
          ! This should be ported on GPU but we are running into compiler issues...
          !$acc update host(u(:)) wait
          du=u
        ENDDO
      ENDDO

      CALL update_device_field(f_du)
      CALL transfert_request(f_du,req_e1_vect)
      CALL update_host_field(f_du)

      DO ind=1,ndomain
        IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
        CALL swap_dimensions(ind)
        CALL swap_geometry(ind)
        du=f_du(ind)
         
        ! Not ported on GPU because all the other kernels cannot be ported
        DO j=jj_begin,jj_end
          DO i=ii_begin,ii_end
            ij=(j-1)*iim+i   
            if (le(ij+u_right)>1e-100) dumax=MAX(dumax,ABS(du(ij+u_right)))
            if (le(ij+u_lup)>1e-100)   dumax=MAX(dumax,ABS(du(ij+u_lup)))
            if (le(ij+u_ldown)>1e-100) dumax=MAX(dumax,ABS(du(ij+u_ldown)))
          ENDDO
        ENDDO
      ENDDO

      IF (using_mpi) THEN 
        CALL reduce_max_omp(dumax,dumax1)
!$OMP MASTER        
        CALL MPI_ALLREDUCE(dumax1,dumax,1,MPI_REAL8,MPI_MAX,comm_icosa,ierr)
!$OMP END MASTER     
        CALL bcast_omp(dumax)  
      ELSE
        CALL allreduce_max_omp(dumax,dumax1)
        dumax=dumax1
      ENDIF  
                        
      DO ind=1,ndomain
        IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
        CALL swap_dimensions(ind)
        CALL swap_geometry(ind)
        u=f_u(ind)
        du=f_du(ind)
        ! This should be ported on GPU but we are running into compiler issues...
        u=du/dumax
      ENDDO
      IF (is_master) PRINT *,"gradiv : it :",it ,": dumax",dumax

    ENDDO 
    IF (is_master) PRINT *,"gradiv : dumax",dumax
    IF (is_master) PRINT *, 'mean T-cell edge size (km)', 1.45*radius/iim_glo/1000., &
                              'effective T-cell half-edge size (km)', dumax**(-.5/nitergdiv)/1000
    IF (is_master)  PRINT *, 'Max. time step assuming c=340 m/s and Courant number=3 (ARK2.3) :', &
                               3./340.*dumax**(-.5/nitergdiv)

    cgraddiv=dumax**(-1./nitergdiv)
    IF (is_master) PRINT *,"cgraddiv : ",cgraddiv

!$OMP BARRIER
!$OMP MASTER
    DO ind=1,ndomain
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      u=f_u(ind)

! set random seed to get reproductibility when using a different number of process
      CALL RANDOM_SEED(size=M)
      CALL RANDOM_SEED(put=(/(i,i=1,M)/))
 
      ! This cannot be ported on GPU due to compiler limitations
       DO j=jj_begin,jj_end
        DO i=ii_begin,ii_end
          ij=(j-1)*iim+i   
          CALL RANDOM_NUMBER(r)
          u(ij+u_right)=r-0.5
          CALL RANDOM_NUMBER(r)
          u(ij+u_lup)=r-0.5
          CALL RANDOM_NUMBER(r)
          u(ij+u_ldown)=r-0.5
        ENDDO
      ENDDO        
    ENDDO
!$OMP END MASTER  
!$OMP BARRIER

    DO it=1,20
 
      dumax=0
      DO iter=1,nitergrot
        CALL update_device_field(f_u)
        CALL transfert_request(f_u,req_e1_vect)
        DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          u=f_u(ind)
          du=f_du(ind)
          CALL compute_gradrot_inplace(u,Av,ne,le,de,1,1)
          ! This should be ported on GPU but we are running into compiler issues...
          !$acc update host(u(:)) wait
          du=u
        ENDDO
      ENDDO
      
      CALL update_device_field(f_du)
      CALL transfert_request(f_du,req_e1_vect)
      CALL update_host_field(f_du)
      
      DO ind=1,ndomain
        IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
        CALL swap_dimensions(ind)
        CALL swap_geometry(ind)
        du=f_du(ind)
        
        ! Not ported on GPU because all the other kernels cannot be ported
        DO j=jj_begin,jj_end
          DO i=ii_begin,ii_end
            ij=(j-1)*iim+i   
            if (le(ij+u_right)>1e-100) dumax=MAX(dumax,ABS(du(ij+u_right)))
            if (le(ij+u_lup)>1e-100) dumax=MAX(dumax,ABS(du(ij+u_lup)))
            if (le(ij+u_ldown)>1e-100) dumax=MAX(dumax,ABS(du(ij+u_ldown)))
          ENDDO
        ENDDO
      ENDDO

      IF (using_mpi) THEN 
        CALL reduce_max_omp(dumax,dumax1)
!$OMP MASTER        
        CALL MPI_ALLREDUCE(dumax1,dumax,1,MPI_REAL8,MPI_MAX,comm_icosa,ierr)
!$OMP END MASTER     
        CALL bcast_omp(dumax)  
      ELSE
        CALL allreduce_max_omp(dumax,dumax1)
        dumax=dumax1
      ENDIF  

      
      DO ind=1,ndomain
        IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
        CALL swap_dimensions(ind)
        CALL swap_geometry(ind)
        u=f_u(ind)
        du=f_du(ind)
        ! This should be ported on GPU but we are running into compiler issues...
        u=du/dumax
      ENDDO
      
      IF (is_master) PRINT *,"gradrot : it :",it ,": dumax",dumax

    ENDDO 
    IF (is_master) PRINT *,"gradrot : dumax",dumax
  
    cgradrot=dumax**(-1./nitergrot) 
    IF (is_master) PRINT *,"cgradrot : ",cgradrot
   


!$OMP BARRIER
!$OMP MASTER
    DO ind=1,ndomain
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      theta=f_theta(ind)

! set random seed to get reproductibility when using a different number of process
      CALL RANDOM_SEED(size=M)
      CALL RANDOM_SEED(put=(/(i,i=1,M)/))

      ! This cannot be ported on GPU due to compiler limitations
      DO j=jj_begin,jj_end
        DO i=ii_begin,ii_end
          ij=(j-1)*iim+i   
          CALL RANDOM_NUMBER(r)
          theta(ij)=r-0.5
        ENDDO
      ENDDO  
    ENDDO
!$OMP END MASTER
!$OMP BARRIER

    DO it=1,20

      dthetamax=0
      DO iter=1,niterdivgrad
        CALL update_device_field(f_theta)
        CALL transfert_request(f_theta,req_i1)
        DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          theta=f_theta(ind)
          dtheta=f_dtheta(ind)
          CALL compute_divgrad_inplace(theta,Ai,ne,le,de,1,1)
          ! This should be ported on GPU but we are running into compiler issues...
          !$acc update host(theta(:)) wait
          dtheta=theta
        ENDDO
      ENDDO

      CALL update_device_field(f_dtheta)
      CALL transfert_request(f_dtheta,req_i1)
      CALL update_host_field(f_dtheta)
      
      DO ind=1,ndomain
        IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
        CALL swap_dimensions(ind)
        CALL swap_geometry(ind)
        dtheta=f_dtheta(ind)
        
        ! Not ported on GPU because all the other kernels cannot be ported
        DO j=jj_begin,jj_end
          DO i=ii_begin,ii_end
            ij=(j-1)*iim+i   
            dthetamax=MAX(dthetamax,ABS(dtheta(ij)))
          ENDDO
        ENDDO
      ENDDO

      IF (using_mpi) THEN 
        CALL reduce_max_omp(dthetamax ,dthetamax1)
!$OMP MASTER        
        CALL MPI_ALLREDUCE(dthetamax1,dthetamax,1,MPI_REAL8,MPI_MAX,comm_icosa,ierr)
!$OMP END MASTER     
        CALL bcast_omp(dthetamax)  
      ELSE
        CALL allreduce_max_omp(dthetamax,dthetamax1)
        dumax=dumax1
      ENDIF  
      
      IF (is_master) PRINT *,"divgrad : it :",it ,": dthetamax",dthetamax

      DO ind=1,ndomain
        IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
        CALL swap_dimensions(ind)
        CALL swap_geometry(ind)
        theta=f_theta(ind)
        dtheta=f_dtheta(ind)
        ! This should be ported on GPU but we are running into compiler issues...
        theta=dtheta/dthetamax
      ENDDO
    ENDDO 

    IF (is_master) PRINT *,"divgrad : divgrad",dthetamax
  
    cdivgrad=dthetamax**(-1./niterdivgrad) 
    IF (is_master) PRINT *,"cdivgrad : ",cdivgrad

    ! Added CMIP5 dissipation vertical profile (SF, 19/09/18)
    vert_prof_dissip = 1
    CALL getin('vert_prof_dissip',vert_prof_dissip)
    IF (vert_prof_dissip == 1) then
       dissip_zref = 30.
       dissip_deltaz = 10.
       dissip_factz = 4.
       CALL getin('dissip_factz',dissip_factz )
       CALL getin('dissip_deltaz',dissip_deltaz )
       CALL getin('dissip_zref',dissip_zref )
       DO l=1,llm
          pseudoz=8.*log(preff/presnivs(l))
          zvert=1+ &
               (tanh((pseudoz-dissip_zref)/dissip_deltaz)+1.)/2. &
               *(dissip_factz-1.)
          tau_graddiv(l) = tau_graddiv(l)/zvert
          tau_gradrot(l) = tau_gradrot(l)/zvert
          tau_divgrad(l) = tau_divgrad(l)/zvert
       ENDDO
    ELSE     
       fact=2
       DO l=1,llm
          IF(ap_bp_present) THEN ! height-dependent dissipation
             zz= 1. - preff/presnivs(l)
          ELSE
             zz = 0.
          END IF
          zvert=fact-(fact-1)/(1+zz*zz)
          tau_graddiv(l) = tau_graddiv(l)/zvert
          tau_gradrot(l) = tau_gradrot(l)/zvert
          tau_divgrad(l) = tau_divgrad(l)/zvert
       ENDDO
    ENDIF

    mintau=tau_graddiv(1)
    DO l=1,llm
      mintau=MIN(mintau,tau_graddiv(l))
      mintau=MIN(mintau,tau_gradrot(l))
      mintau=MIN(mintau,tau_divgrad(l))
    ENDDO

    IF(rayleigh_friction_type>0) mintau=MIN(mintau,rayleigh_tau)

    IF(mintau>0) THEN
       IF (itau_dissip==0) THEN
         IF (is_master) PRINT *,"init_dissip: Automatic computation of itau_dissip..."
         itau_dissip=INT(mintau/dt)
       ENDIF
    ELSE
       IF (is_master) PRINT *,"init_dissip: No dissipation time set, setting itau_dissip to 1000000000"
       itau_dissip=100000000
    END IF
    itau_dissip=MAX(1,itau_dissip)
    dtdissip=itau_dissip*dt
    IF (is_master) THEN
      PRINT *,"init_dissip: rayleigh_tau",rayleigh_tau, "mintau ",mintau
      PRINT *,"init_dissip: itau_dissip",itau_dissip," dtdissip ",dtdissip
    ENDIF

    !$acc update device(tau_graddiv(:),tau_gradrot(:),tau_divgrad(:)) async

  END SUBROUTINE init_dissip 
 
  
  SUBROUTINE dissip(f_ps,f_mass,f_phis,f_geopot,f_theta_rhodz,f_ue, f_dtheta_rhodz,f_due)
  USE icosa
  USE theta2theta_rhodz_mod
  USE pression_mod
  USE exner_mod
  USE geopotential_mod
  USE trace
  USE time_mod
  USE omp_para
  USE abort_mod
  IMPLICIT NONE
    TYPE(t_field),POINTER :: f_ps(:), f_mass(:), f_phis(:), f_geopot(:)
    TYPE(t_field),POINTER :: f_theta_rhodz(:), f_dtheta_rhodz(:)
    TYPE(t_field),POINTER :: f_ue(:), f_due(:)

    REAL(rstd),POINTER         :: due(:,:)
    REAL(rstd),POINTER         :: phi(:,:), ue(:,:)
    REAL(rstd),POINTER         :: due_diss1(:,:)
    REAL(rstd),POINTER         :: due_diss2(:,:)
    REAL(rstd),POINTER         :: dtheta_rhodz(:,:,:)
    REAL(rstd),POINTER         :: dtheta_rhodz_diss(:,:)

    INTEGER :: ind, shear
    INTEGER :: l,ij,nn

!$OMP BARRIER
    
    CALL trace_start("dissip")
    CALL gradiv(f_ue,f_due_diss1)
    CALL gradrot(f_ue,f_due_diss2)

    CALL divgrad_theta_rhodz(f_mass,f_theta_rhodz,f_dtheta_rhodz_diss)
    
    DO ind=1,ndomain
      IF (.NOT. assigned_domain(ind)) CYCLE
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      due=f_due(ind) 
      due_diss1=f_due_diss1(ind)
      due_diss2=f_due_diss2(ind)
      dtheta_rhodz=f_dtheta_rhodz(ind)
      dtheta_rhodz_diss=f_dtheta_rhodz_diss(ind)

      !$acc parallel loop collapse(2) present(due(:,:), dtheta_rhodz(:,:,:), due_diss1(:,:), due_diss2(:,:), dtheta_rhodz_diss(:,:), tau_graddiv(:), tau_gradrot(:), tau_divgrad(:)) async
      DO l=ll_begin,ll_end
!DIR$ SIMD
        DO ij=ij_begin,ij_end

            due(ij+u_right,l) = -0.5*( due_diss1(ij+u_right,l)/tau_graddiv(l) + due_diss2(ij+u_right,l)/tau_gradrot(l))*itau_dissip 
            due(ij+u_lup,l)   = -0.5*( due_diss1(ij+u_lup,l)  /tau_graddiv(l) + due_diss2(ij+u_lup,l)  /tau_gradrot(l))*itau_dissip
            due(ij+u_ldown,l) = -0.5*( due_diss1(ij+u_ldown,l)/tau_graddiv(l) + due_diss2(ij+u_ldown,l)/tau_gradrot(l))*itau_dissip

            dtheta_rhodz(ij,l,1) = -0.5*dtheta_rhodz_diss(ij,l)/tau_divgrad(l)*itau_dissip
        ENDDO
      ENDDO
      !$acc end parallel loop

!      dtheta_rhodz=0
!      due=0

      IF(rayleigh_friction_type>0) THEN
        CALL abort_acc("dissip/rayleigh_friction_type>0")
       IF(rayleigh_friction_type<99) THEN
         phi=f_geopot(ind)
         ue=f_ue(ind)
         DO l=ll_begin,ll_end
            DO ij=ij_begin,ij_end
               CALL relax(t_right, u_right)
               CALL relax(t_lup,   u_lup)
               CALL relax(t_ldown, u_ldown)
            ENDDO
         END DO
       ELSE
         ue=f_ue(ind)
            !$acc parallel loop present(ue(:,:), due(:,:), lat_e(:))
            DO ij=ij_begin,ij_end
              nn = ij+u_right
              IF (ABS(lat_e(nn)) .gt. rayleigh_limlat) THEN
                !print*, "latitude", lat_e(nn)*180./3.14159
                due(nn,ll_begin:ll_begin+1) = due(nn,ll_begin:ll_begin+1) - (ue(nn,ll_begin:ll_begin+1)/rayleigh_tau)
              ENDIF
              nn = ij+u_lup
              IF (ABS(lat_e(nn)) .gt. rayleigh_limlat) THEN
                due(nn,ll_begin:ll_begin+1) = due(nn,ll_begin:ll_begin+1) - (ue(nn,ll_begin:ll_begin+1)/rayleigh_tau)
              ENDIF
              nn = ij+u_ldown
              IF (ABS(lat_e(nn)) .gt. rayleigh_limlat) THEN
                due(nn,ll_begin:ll_begin+1) = due(nn,ll_begin:ll_begin+1) - (ue(nn,ll_begin:ll_begin+1)/rayleigh_tau)
              ENDIF
            ENDDO
            !$acc end parallel loop
       ENDIF
      END IF
   END DO
   CALL trace_end("dissip")

   !CALL write_dissip_tendencies
!$OMP BARRIER

    CONTAINS

      SUBROUTINE relax(shift_t, shift_u)
        USE dcmip_initial_conditions_test_1_2_3
        REAL(rstd) :: z, ulon,ulat, & ! input to test2_schaer_mountain
             p,hyam,hybm,w,t,phis,ps,rho,q, &   ! unused input/output to test2_schaer_mountain
             fz, u3d(3), uref
        REAL(rstd), PARAMETER :: zh=2e4,ztop=3e4  ! DCMIP values
        LOGICAL :: hybrid_eta
        INTEGER :: shift_u, shift_t, zcoords, nn
        z = (phi(ij,l)+phi(ij+shift_t,l))/(2.*g)
        IF(z>zh) THEN  ! relax only in the sponge layer z>zh

           nn = ij+shift_u
           zcoords = 1 ; hybrid_eta = .FALSE. ! use z instead of p or hyam/hybm
           CALL test2_schaer_mountain(lon_e(nn),lat_e(nn),p,z,zcoords,hybrid_eta, &
                hyam,hybm,shear,ulon,ulat,w,t,phis,ps,rho,q)
!           u3d = ulon*elon_e(nn,:) + ulat*elat_e(nn,:)
           u3d = ulon*elon_e(nn,:) ! ulat=0
           uref = sum(u3d*ep_e(nn,:))

           fz = sin((pi/2)*(z-zh)/(ztop-zh))
           fz = fz*fz/rayleigh_tau
           due(nn,l) = due(nn,l) - itau_dissip*fz*(ue(nn,l)-uref)
!           fz = 1./rayleigh_tau
!           due(nn,l) = due(nn,l) - fz*ue(nn,l)
         END IF
       END SUBROUTINE relax
      
       SUBROUTINE write_dissip_tendencies
         USE observable_mod, ONLY : f_buf_ulon, f_buf_ulat
         USE wind_mod
         USE output_field_mod

         CALL transfert_request(f_due_diss1,req_e1_vect)
         CALL un2ulonlat(f_due_diss1, f_buf_ulon, f_buf_ulat, (1./(tau_graddiv(1))))
         CALL output_field("dulon_diss1",f_buf_ulon)
         CALL output_field("dulat_diss1",f_buf_ulat)
!
         CALL transfert_request(f_due_diss2,req_e1_vect)
         CALL un2ulonlat(f_due_diss2, f_buf_ulon, f_buf_ulat, (1./(tau_graddiv(1))))
         CALL output_field("dulon_diss2",f_buf_ulon)
         CALL output_field("dulat_diss2",f_buf_ulat)
       END SUBROUTINE write_dissip_tendencies

  END SUBROUTINE dissip

  SUBROUTINE gradiv(f_ue,f_due)
  USE icosa
  USE trace
  USE omp_para
  IMPLICIT NONE
    TYPE(t_field),POINTER :: f_ue(:)
    TYPE(t_field),POINTER :: f_due(:)
    REAL(rstd),POINTER    :: ue(:,:)
    REAL(rstd),POINTER    :: due(:,:)
    INTEGER :: ind
    INTEGER :: it,l,ij
       
    CALL trace_start("gradiv")

    DO ind=1,ndomain
      IF (.NOT. assigned_domain(ind)) CYCLE
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      ue=f_ue(ind)
      due=f_due(ind) 
      !$acc parallel loop present(ue(:,:), due(:,:)) async
      DO  l = ll_begin, ll_end
        !$acc loop
!DIR$ SIMD
        DO ij=ij_begin,ij_end
             due(ij+u_right,l)=ue(ij+u_right,l)
             due(ij+u_lup,l)=ue(ij+u_lup,l)
             due(ij+u_ldown,l)=ue(ij+u_ldown,l)
        ENDDO
      ENDDO
    ENDDO

    DO it=1,nitergdiv
      CALL send_message(f_due,req_due)
      CALL wait_message(req_due)

      DO ind=1,ndomain
        IF (.NOT. assigned_domain(ind)) CYCLE
        CALL swap_dimensions(ind)
        CALL swap_geometry(ind)
        due=f_due(ind)
        CALL compute_gradiv_inplace(due,Ai,ne,le,de,ll_begin,ll_end)
      ENDDO
    ENDDO

   CALL trace_end("gradiv")

  END SUBROUTINE gradiv
  

  SUBROUTINE gradrot(f_ue,f_due)
  USE icosa
  USE trace
  USE omp_para
  IMPLICIT NONE
    TYPE(t_field),POINTER :: f_ue(:)
    TYPE(t_field),POINTER :: f_due(:)
    REAL(rstd),POINTER    :: ue(:,:)
    REAL(rstd),POINTER    :: due(:,:)
    INTEGER :: ind
    INTEGER :: it,l,ij
       
    CALL trace_start("gradrot")

    DO ind=1,ndomain
      IF (.NOT. assigned_domain(ind)) CYCLE
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      ue=f_ue(ind)
      due=f_due(ind) 
      !$acc parallel loop present(ue(:,:), due(:,:)) async
      DO  l = ll_begin, ll_end
         !$acc loop
!DIR$ SIMD
        DO ij=ij_begin,ij_end
             due(ij+u_right,l)=ue(ij+u_right,l)
             due(ij+u_lup,l)=ue(ij+u_lup,l)
             due(ij+u_ldown,l)=ue(ij+u_ldown,l)
        ENDDO
      ENDDO
    ENDDO

    DO it=1,nitergrot
      CALL send_message(f_due,req_due)
      CALL wait_message(req_due)
        
      DO ind=1,ndomain
        IF (.NOT. assigned_domain(ind)) CYCLE
        CALL swap_dimensions(ind)
        CALL swap_geometry(ind)
        due=f_due(ind) 
        CALL compute_gradrot_inplace(due,Av,ne,le,de,ll_begin,ll_end)
      ENDDO

    ENDDO

    CALL trace_end("gradrot")

  END SUBROUTINE gradrot
  
  SUBROUTINE divgrad(f_theta,f_dtheta)
  USE icosa
  USE trace
  USE omp_para
  IMPLICIT NONE
    TYPE(t_field),POINTER :: f_theta(:)
    TYPE(t_field),POINTER :: f_dtheta(:)
    REAL(rstd),POINTER    :: theta(:,:)
    REAL(rstd),POINTER    :: dtheta(:,:)
    INTEGER :: ind, l, ij
    INTEGER :: it

    CALL trace_start("divgrad")
       
    DO ind=1,ndomain
      IF (.NOT. assigned_domain(ind)) CYCLE
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      theta=f_theta(ind)
      dtheta=f_dtheta(ind) 
      ! Replace Fortran 90 construct dtheta=theta because it confuses PGI acc kernels...
      !$acc parallel loop collapse(2) present(theta(:,:), dtheta(:,:)) async
      DO l=ll_begin,ll_end
!DIR$ SIMD
        DO ij=1,iim*jjm
          dtheta(ij,l)=theta(ij,l)
        ENDDO
      ENDDO
      !$acc end parallel loop
    ENDDO

    DO it=1,niterdivgrad
      CALL transfert_request(f_dtheta,req_i1)
        
      DO ind=1,ndomain
        IF (.NOT. assigned_domain(ind)) CYCLE
        CALL swap_dimensions(ind)
        CALL swap_geometry(ind)
        dtheta=f_dtheta(ind) 
        CALL compute_divgrad_inplace(dtheta,Ai,ne,le,de,ll_begin,ll_end)
      ENDDO

    ENDDO

    CALL trace_end("divgrad")

  END SUBROUTINE divgrad
   
  SUBROUTINE divgrad_theta_rhodz(f_mass,f_theta_rhodz,f_dtheta_rhodz)
  USE icosa
  USE trace
  USE omp_para
  IMPLICIT NONE
    TYPE(t_field),POINTER :: f_mass(:)
    TYPE(t_field),POINTER :: f_theta_rhodz(:)
    TYPE(t_field),POINTER :: f_dtheta_rhodz(:)
    
    REAL(rstd),POINTER :: mass(:,:)
    REAL(rstd),POINTER :: theta_rhodz(:,:,:)
    REAL(rstd),POINTER :: dtheta_rhodz(:,:)

    INTEGER :: ind
    INTEGER :: it,l,ij

    CALL trace_start("divgrad")

    DO ind=1,ndomain
      IF (.NOT. assigned_domain(ind)) CYCLE
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      mass=f_mass(ind)
      theta_rhodz=f_theta_rhodz(ind)
      dtheta_rhodz=f_dtheta_rhodz(ind) 
      !$acc parallel loop present(mass(:,:), theta_rhodz(:,:,:), dtheta_rhodz(:,:)) async
      DO  l = ll_begin, ll_end
        !$acc loop
!DIR$ SIMD
        DO ij=ij_begin,ij_end
            dtheta_rhodz(ij,l) = theta_rhodz(ij,l,1) / mass(ij,l)
        ENDDO
      ENDDO
    ENDDO

    DO it=1,niterdivgrad
      CALL send_message(f_dtheta_rhodz,req_dtheta)
      CALL wait_message(req_dtheta)
      DO ind=1,ndomain
        IF (.NOT. assigned_domain(ind)) CYCLE
        CALL swap_dimensions(ind)
        CALL swap_geometry(ind)
        dtheta_rhodz=f_dtheta_rhodz(ind) 
        CALL compute_divgrad_inplace(dtheta_rhodz,Ai,ne,le,de,ll_begin,ll_end)
      ENDDO

    ENDDO

    DO ind=1,ndomain
      IF (.NOT. assigned_domain(ind)) CYCLE
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      dtheta_rhodz=f_dtheta_rhodz(ind) 
      mass=f_mass(ind)
    
      !$acc parallel loop collapse(2) present(mass(:,:), dtheta_rhodz(:,:)) async
      DO  l = ll_begin, ll_end
!DIR$ SIMD
        DO ij=ij_begin,ij_end
            dtheta_rhodz(ij,l) = dtheta_rhodz(ij,l) * mass(ij,l)
        ENDDO
      ENDDO
      !$acc end parallel loop
    ENDDO


    CALL trace_end("divgrad")

  END SUBROUTINE divgrad_theta_rhodz

  SUBROUTINE compute_gradiv(ue,gradivu_e,Ai,ne,le,de,llb,lle)
    INTEGER,INTENT(IN)     :: llb
    INTEGER,INTENT(IN)     :: lle
    REAL(rstd),INTENT(OUT) :: gradivu_e(iim*3*jjm,llm)
    REAL(rstd),INTENT(IN)  :: ue(iim*3*jjm,llm)
    REAL(rstd),INTENT(IN)  :: Ai(iim*jjm)
    INTEGER,INTENT(IN)     :: ne(iim*jjm,6)
    REAL(rstd),INTENT(IN)  :: le(iim*3*jjm)
    REAL(rstd),INTENT(IN)  :: de(iim*3*jjm)

    gradivu_e = ue
    CALL compute_gradiv_inplace(gradivu_e,Ai,ne,le,de,llb,lle)

  END SUBROUTINE compute_gradiv
  
  SUBROUTINE compute_gradiv_inplace(ue_gradivu_e,Ai,ne,le,de,llb,lle)
    INTEGER,INTENT(IN)     :: llb
    INTEGER,INTENT(IN)     :: lle
    REAL(rstd),INTENT(INOUT) :: ue_gradivu_e(iim*3*jjm,llm)
    REAL(rstd),INTENT(IN)  :: Ai(iim*jjm)
    INTEGER,INTENT(IN)     :: ne(iim*jjm,6)
    REAL(rstd),INTENT(IN)  :: le(iim*3*jjm)
    REAL(rstd),INTENT(IN)  :: de(iim*3*jjm)
    REAL(rstd) :: divu_i(iim*jjm,llb:lle)
    
    INTEGER :: ij,l

    ! ue and gradivu_e second dimension is not always llm so use the bounds explicitly
    !$acc data present( ue_gradivu_e(:,llb:lle), Ai(:), ne(:,:), le(:), de(:)) create(divu_i(:,:)) async

    !$acc parallel loop collapse(2) async
    DO l=llb,lle
      !DIR$ SIMD
      DO ij=ij_begin,ij_end
          divu_i(ij,l)=1./Ai(ij)*(ne(ij,right)*ue_gradivu_e(ij+u_right,l)*le(ij+u_right)  +  &
                             ne(ij,rup)*ue_gradivu_e(ij+u_rup,l)*le(ij+u_rup)        +  &  
                             ne(ij,lup)*ue_gradivu_e(ij+u_lup,l)*le(ij+u_lup)        +  &  
                             ne(ij,left)*ue_gradivu_e(ij+u_left,l)*le(ij+u_left)     +  &  
                             ne(ij,ldown)*ue_gradivu_e(ij+u_ldown,l)*le(ij+u_ldown)  +  &  
                             ne(ij,rdown)*ue_gradivu_e(ij+u_rdown,l)*le(ij+u_rdown))
      ENDDO
    ENDDO
    !$acc end parallel loop
    
    !$acc parallel loop collapse(2) async
    DO l=llb,lle
      !DIR$ SIMD
      DO ij=ij_begin,ij_end
          ue_gradivu_e(ij+u_right,l)=-1/de(ij+u_right)*(ne(ij,right)*divu_i(ij,l)+ ne(ij+t_right,left)*divu_i(ij+t_right,l) ) 
          ue_gradivu_e(ij+u_lup,l)=-1/de(ij+u_lup)*(ne(ij,lup)*divu_i(ij,l)+ ne(ij+t_lup,rdown)*divu_i(ij+t_lup,l))   
          ue_gradivu_e(ij+u_ldown,l)=-1/de(ij+u_ldown)*(ne(ij,ldown)*divu_i(ij,l)+ne(ij+t_ldown,rup)*divu_i(ij+t_ldown,l) )
      ENDDO
    ENDDO
    !$acc end parallel loop

    !$acc parallel loop collapse(2) async
    DO l=llb,lle
      !DIR$ SIMD
      DO ij=ij_begin,ij_end
          ue_gradivu_e(ij+u_right,l)=-ue_gradivu_e(ij+u_right,l)*cgraddiv
          ue_gradivu_e(ij+u_lup,l)=-ue_gradivu_e(ij+u_lup,l)*cgraddiv
          ue_gradivu_e(ij+u_ldown,l)=-ue_gradivu_e(ij+u_ldown,l)*cgraddiv
      ENDDO
    ENDDO
    !$acc end parallel loop
    !$acc end data   
  END SUBROUTINE compute_gradiv_inplace

  SUBROUTINE compute_divgrad(theta,divgrad_i,Ai,ne,le,de,llb,lle)
    INTEGER,INTENT(IN)     :: llb
    INTEGER,INTENT(IN)     :: lle
    REAL(rstd),INTENT(IN) :: theta(iim*jjm,1:lle)
    REAL(rstd),INTENT(OUT) :: divgrad_i(iim*jjm,1:lle)
    REAL(rstd),INTENT(IN)  :: Ai(iim*jjm)
    INTEGER,INTENT(IN)     :: ne(iim*jjm,6)
    REAL(rstd),INTENT(IN)  :: le(iim*3*jjm)
    REAL(rstd),INTENT(IN)  :: de(iim*3*jjm)

    divgrad_i = theta
    CALL compute_divgrad_inplace(divgrad_i,Ai,ne,le,de,llb,lle)
  END SUBROUTINE compute_divgrad
  
  SUBROUTINE compute_divgrad_inplace(theta_divgrad_i,Ai,ne,le,de,llb,lle)
    INTEGER,INTENT(IN)     :: llb
    INTEGER,INTENT(IN)     :: lle
    REAL(rstd),INTENT(INOUT) :: theta_divgrad_i(iim*jjm,1:lle)
    REAL(rstd),INTENT(IN)  :: Ai(iim*jjm)
    INTEGER,INTENT(IN)     :: ne(iim*jjm,6)
    REAL(rstd),INTENT(IN)  :: le(iim*3*jjm)
    REAL(rstd),INTENT(IN)  :: de(iim*3*jjm)
    REAL(rstd)  :: grad_e(3*iim*jjm,llb:lle)

    INTEGER :: ij,l

    ! theta and divgrad_i second dimension is not always llm so use the bounds explicitly
    !$acc data present(theta_divgrad_i(:,llb:lle), Ai(:), de(:), ne(:,:), le(:)) create(grad_e(:,:)) async
       
    !$acc parallel loop collapse(2) async
    DO l=llb,lle
      !DIR$ SIMD
      DO ij=ij_begin_ext,ij_end_ext 
          grad_e(ij+u_right,l)=-1/de(ij+u_right)*(ne(ij,right)*theta_divgrad_i(ij,l)+ ne(ij+t_right,left)*theta_divgrad_i(ij+t_right,l) )
          grad_e(ij+u_lup,l)=-1/de(ij+u_lup)*(ne(ij,lup)*theta_divgrad_i(ij,l)+ ne(ij+t_lup,rdown)*theta_divgrad_i(ij+t_lup,l ))
          grad_e(ij+u_ldown,l)=-1/de(ij+u_ldown)*(ne(ij,ldown)*theta_divgrad_i(ij,l)+ne(ij+t_ldown,rup)*theta_divgrad_i(ij+t_ldown,l) )
      ENDDO
    ENDDO
    !$acc end parallel loop
    
    
    !$acc parallel loop collapse(2) async
    DO l=llb,lle
      !DIR$ SIMD
      DO ij=ij_begin,ij_end
          theta_divgrad_i(ij,l)=1./Ai(ij)*(ne(ij,right)*grad_e(ij+u_right,l)*le(ij+u_right)  +  &
                             ne(ij,rup)*grad_e(ij+u_rup,l)*le(ij+u_rup)              +  &
                             ne(ij,lup)*grad_e(ij+u_lup,l)*le(ij+u_lup)              +  &
                             ne(ij,left)*grad_e(ij+u_left,l)*le(ij+u_left)           +  &
                             ne(ij,ldown)*grad_e(ij+u_ldown,l)*le(ij+u_ldown)        +  & 
                             ne(ij,rdown)*grad_e(ij+u_rdown,l)*le(ij+u_rdown))
      ENDDO
    ENDDO
    !$acc end parallel loop
   
    !$acc parallel loop collapse(2) async
    DO l=llb,lle
      DO ij=ij_begin,ij_end
          theta_divgrad_i(ij,l)=-theta_divgrad_i(ij,l)*cdivgrad
      ENDDO
    ENDDO
    !$acc end parallel loop
    !$acc end data
  END SUBROUTINE compute_divgrad_inplace

  SUBROUTINE compute_gradrot(ue,gradrot_e,Av,ne,le,de,llb,lle)
    INTEGER,INTENT(IN)     :: llb
    INTEGER,INTENT(IN)     :: lle
    REAL(rstd),INTENT(IN) :: ue(iim*3*jjm,lle)
    REAL(rstd),INTENT(OUT) :: gradrot_e(iim*3*jjm,lle)
    REAL(rstd),INTENT(IN)  :: Av(2*iim*jjm)
    INTEGER,INTENT(IN)     :: ne(iim*jjm,6)
    REAL(rstd),INTENT(IN)  :: le(iim*3*jjm)
    REAL(rstd),INTENT(IN)  :: de(iim*3*jjm)

    gradrot_e = ue
    CALL compute_gradrot_inplace(gradrot_e,Av,ne,le,de,llb,lle)
  END SUBROUTINE compute_gradrot

  SUBROUTINE compute_gradrot_inplace(ue_gradrot_e,Av,ne,le,de,llb,lle)
    INTEGER,INTENT(IN)     :: llb
    INTEGER,INTENT(IN)     :: lle
    REAL(rstd),INTENT(INOUT) :: ue_gradrot_e(iim*3*jjm,lle)
    REAL(rstd),INTENT(IN)  :: Av(2*iim*jjm)
    INTEGER,INTENT(IN)     :: ne(iim*jjm,6)
    REAL(rstd),INTENT(IN)  :: le(iim*3*jjm)
    REAL(rstd),INTENT(IN)  :: de(iim*3*jjm)
    REAL(rstd) :: rot_v(2*iim*jjm,llb:lle)

    INTEGER :: ij,l

    ! ue and gradrot_e second dimension is not always llm so use the bounds explicitly
    ! gradrot_e should be copyout but using copy instead allows to compare the output
    ! more easily as the code sometimes uses unintialed values
    !$acc data present(ue_gradrot_e(:,llb:lle), Av(:), ne(:,:), de(:), le(:)) create(rot_v(:,:)) async
     
    !$acc parallel loop collapse(2) async
    DO l=llb,lle
!DIR$ SIMD
      DO ij=ij_begin_ext,ij_end_ext        
          rot_v(ij+z_up,l)= 1./Av(ij+z_up)*(  ne(ij,rup)*ue_gradrot_e(ij+u_rup,l)*de(ij+u_rup) &
                                + ne(ij+t_rup,left)*ue_gradrot_e(ij+t_rup+u_left,l)*de(ij+t_rup+u_left) &
                                - ne(ij,lup)*ue_gradrot_e(ij+u_lup,l)*de(ij+u_lup) )                               
          rot_v(ij+z_down,l) = 1./Av(ij+z_down)*( ne(ij,ldown)*ue_gradrot_e(ij+u_ldown,l)*de(ij+u_ldown) &
                                     + ne(ij+t_ldown,right)*ue_gradrot_e(ij+t_ldown+u_right,l)*de(ij+t_ldown+u_right) &
                                     - ne(ij,rdown)*ue_gradrot_e(ij+u_rdown,l)*de(ij+u_rdown) )
      ENDDO
    ENDDO                              
    !$acc end parallel loop
  
    !$acc parallel loop collapse(2) async
    DO l=llb,lle
!DIR$ SIMD
      DO ij=ij_begin,ij_end
          ue_gradrot_e(ij+u_right,l)=1/le(ij+u_right)*ne(ij,right)*(rot_v(ij+z_rdown,l)-rot_v(ij+z_rup,l))       
          ue_gradrot_e(ij+u_lup,l)=1/le(ij+u_lup)*ne(ij,lup)*(rot_v(ij+z_up,l)-rot_v(ij+z_lup,l))        
          ue_gradrot_e(ij+u_ldown,l)=1/le(ij+u_ldown)*ne(ij,ldown)*(rot_v(ij+z_ldown,l)-rot_v(ij+z_down,l))
      ENDDO
    ENDDO
    !$acc end parallel loop

    !$acc parallel loop collapse(2) async
    DO l=llb,lle
!DIR$ SIMD
      DO ij=ij_begin,ij_end    
          ue_gradrot_e(ij+u_right,l)=-ue_gradrot_e(ij+u_right,l)*cgradrot       
          ue_gradrot_e(ij+u_lup,l)=-ue_gradrot_e(ij+u_lup,l)*cgradrot
          ue_gradrot_e(ij+u_ldown,l)=-ue_gradrot_e(ij+u_ldown,l)*cgradrot        
      ENDDO
    ENDDO
    !$acc end parallel loop
    !$acc end data   
  END SUBROUTINE compute_gradrot_inplace


END MODULE dissip_gcm_mod