Merge pull request #101 from CaNS-World/stag-cmpt-rho_av

Author: p-costa

src/main.f90

@@ -84,13 +84,14 @@ program cans
 #endif
   use mod_timer          , only: timer_tic,timer_toc,timer_print
   use mod_updatep        , only: updatep
-  use mod_utils          , only: bulk_mean,bulk_mean_12
+  use mod_utils          , only: bulk_mean_12_stag
   !@acc use mod_utils    , only: device_memory_footprint
   use mod_types
   implicit none
   integer , dimension(3) :: lo,hi,n,n_x_fft,n_y_fft,lo_z,hi_z,n_z
   real(rp), allocatable, dimension(:,:,:) :: u,v,w,p,pp,pn,po
-  real(rp) :: rho_av
+  real(rp), dimension(3) :: rho_av
+  logical , dimension(3) :: is_cmpt_rho_av
 #if !defined(_OPENACC)
   type(C_PTR), dimension(2,2) :: arrplanp
 #else
@@ -425,10 +426,14 @@ program cans
         !$acc update device(u,v,w,p) async(1)
         call bounduvw(cbcvel,n,bcvel,nb,is_bound,.false.,dl,dzc,dzf,u,v,w)
       else
-        rho_av = 0.
-        if(any(abs(gacc(:))>0. .and. cbcpre(0,:)//cbcpre(1,:) == 'PP')) then
-          call bulk_mean_12(n,grid_vol_ratio_c,psi,rho12,rho_av)
-        end if
+        rho_av(:) = 0.
+        is_cmpt_rho_av(:) = (abs(gacc(:)) > 0.) .and. cbcpre(0,:)//cbcpre(1,:) == 'PP'
+        if(is_cmpt_rho_av(1)) &
+          call bulk_mean_12_stag(n,1,grid_vol_ratio_c,psi,rho12,rho_av(1))
+        if(is_cmpt_rho_av(2)) &
+          call bulk_mean_12_stag(n,2,grid_vol_ratio_c,psi,rho12,rho_av(2))
+        if(is_cmpt_rho_av(3)) &
+          call bulk_mean_12_stag(n,3,grid_vol_ratio_f,psi,rho12,rho_av(3))
         call tm(tm_coeff,n,dli,dzci,dzfi,dt,dt_r, &
                 bforce,gacc,sigma,rho_av,rho12,mu12,beta12,rho0,psi,kappa,p,pn,po,s, &
                 psio,psiflx_x,psiflx_y,psiflx_z,u,v,w)

src/mom.f90

@@ -249,11 +249,11 @@ subroutine momz_d(nx,ny,nz,dxi,dyi,dzci,dzfi,mu12,psi,u,v,w,dwdt)
     end do
   end subroutine momz_d
   !
-  subroutine momx_p(nx,ny,nz,dxi,dt_r,bforce,gacc,rho0,rho_av,rho12,psi,p,pp,dudt)
+  subroutine momx_p(nx,ny,nz,dxi,dt_r,bforce,gacc,rho0,rhox_av,rho12,psi,p,pp,dudt)
     implicit none
     integer , intent(in) :: nx,ny,nz
     real(rp), intent(in) :: dxi,dt_r
-    real(rp), intent(in) :: bforce,gacc,rho0,rho_av,rho12(2)
+    real(rp), intent(in) :: bforce,gacc,rho0,rhox_av,rho12(2)
     real(rp), dimension(0:,0:,0:), intent(in   ) :: psi,p,pp
     real(rp), dimension( :, :, :), intent(inout) :: dudt
     real(rp) :: rhop,dpdl
@@ -269,7 +269,7 @@ subroutine momx_p(nx,ny,nz,dxi,dt_r,bforce,gacc,rho0,rho_av,rho12,psi,p,pp,dudt)
           rhop = rho + drho*0.5*(psi(i+1,j,k)+psi(i,j,k))
           dpdl = (p(i+1,j,k)-p(i,j,k))*dxi
           !
-          dudt(i,j,k) = dudt(i,j,k) + bforce + gacc*(rhop-rho_av) &
+          dudt(i,j,k) = dudt(i,j,k) + bforce + gacc*(rhop-rhox_av) &
 #if defined(_CONSTANT_COEFFS_POISSON)
                         -dpdl*rhop/rho0 - (1.-rhop/rho0)*( ((1.+dt_r)*p(i+1,j,k)-dt_r*(pp(i+1,j,k))) - &
                                                            ((1.+dt_r)*p(i  ,j,k)-dt_r*(pp(i  ,j,k))) &
@@ -282,11 +282,11 @@ subroutine momx_p(nx,ny,nz,dxi,dt_r,bforce,gacc,rho0,rho_av,rho12,psi,p,pp,dudt)
     end do
   end subroutine momx_p
   !
-  subroutine momy_p(nx,ny,nz,dyi,dt_r,bforce,gacc,rho0,rho_av,rho12,psi,p,pp,dvdt)
+  subroutine momy_p(nx,ny,nz,dyi,dt_r,bforce,gacc,rho0,rhoy_av,rho12,psi,p,pp,dvdt)
     implicit none
     integer , intent(in) :: nx,ny,nz
     real(rp), intent(in) :: dyi,dt_r
-    real(rp), intent(in) :: bforce,gacc,rho0,rho_av,rho12(2)
+    real(rp), intent(in) :: bforce,gacc,rho0,rhoy_av,rho12(2)
     real(rp), dimension(0:,0:,0:), intent(in   ) :: psi,p,pp
     real(rp), dimension( :, :, :), intent(inout) :: dvdt
     integer :: i,j,k
@@ -302,7 +302,7 @@ subroutine momy_p(nx,ny,nz,dyi,dt_r,bforce,gacc,rho0,rho_av,rho12,psi,p,pp,dvdt)
           rhop = rho + drho*0.5*(psi(i,j+1,k)+psi(i,j,k))
           dpdl = (p(i,j+1,k)-p(i,j,k))*dyi
           !
-          dvdt(i,j,k) = dvdt(i,j,k) + bforce + gacc*(rhop-rho_av) &
+          dvdt(i,j,k) = dvdt(i,j,k) + bforce + gacc*(rhop-rhoy_av) &
 #if defined(_CONSTANT_COEFFS_POISSON)
                         -dpdl*rhop/rho0 - (1.-1.*rhop/rho0)*( ((1.+dt_r)*p(i,j+1,k)-dt_r*(pp(i,j+1,k))) - &
                                                               ((1.+dt_r)*p(i,j  ,k)-dt_r*(pp(i,j  ,k))) &
@@ -315,12 +315,12 @@ subroutine momy_p(nx,ny,nz,dyi,dt_r,bforce,gacc,rho0,rho_av,rho12,psi,p,pp,dvdt)
     end do
   end subroutine momy_p
   !
-  subroutine momz_p(nx,ny,nz,dzci,dt_r,bforce,gacc,rho0,rho_av,rho12,psi,p,pp,dwdt)
+  subroutine momz_p(nx,ny,nz,dzci,dt_r,bforce,gacc,rho0,rhoz_av,rho12,psi,p,pp,dwdt)
     implicit none
     integer , intent(in) :: nx,ny,nz
     real(rp), intent(in), dimension(0:) :: dzci
     real(rp), intent(in) :: dt_r
-    real(rp), intent(in) :: bforce,gacc,rho0,rho_av,rho12(2)
+    real(rp), intent(in) :: bforce,gacc,rho0,rhoz_av,rho12(2)
     real(rp), dimension(0:,0:,0:), intent(in   ) :: psi,p,pp
     real(rp), dimension( :, :, :), intent(inout) :: dwdt
     real(rp) :: rhop,dpdl
@@ -336,7 +336,7 @@ subroutine momz_p(nx,ny,nz,dzci,dt_r,bforce,gacc,rho0,rho_av,rho12,psi,p,pp,dwdt
           rhop = rho + drho*0.5*(psi(i,j,k+1)+psi(i,j,k))
           dpdl = (p(i,j,k+1)-p(i,j,k))*dzci(k)
           !
-          dwdt(i,j,k) = dwdt(i,j,k) + bforce + gacc*(rhop-rho_av) &
+          dwdt(i,j,k) = dwdt(i,j,k) + bforce + gacc*(rhop-rhoz_av) &
 #if defined(_CONSTANT_COEFFS_POISSON)
                         -dpdl*rhop/rho0 - (1.-1.*rhop/rho0)*( ((1.+dt_r)*p(i,j,k+1)-dt_r*(pp(i,j,k+1))) - &
                                                               ((1.+dt_r)*p(i,j,k  )-dt_r*(pp(i,j,k  ))) &
@@ -860,7 +860,7 @@ subroutine mom_xyz_oth(n,dli,dzci,dzfi,dt_r,rho12,beta12,bforce,gacc,sigma,rho0,
     real(rp), intent(in   ), dimension(2) :: rho12,beta12
     real(rp), intent(in   ), dimension(3) :: bforce,gacc
     real(rp), intent(in   ) :: sigma
-    real(rp), intent(in   ) :: rho0,rho_av
+    real(rp), intent(in   ) :: rho0,rho_av(3)
     real(rp), intent(in   ), dimension(0:,0:,0:)           :: p,psi,kappa
     real(rp), intent(in   ), dimension(0:,0:,0:), optional :: s,pn,po
     real(rp), intent(inout), dimension( :, :, :)           :: dudt,dvdt,dwdt
@@ -872,7 +872,7 @@ subroutine mom_xyz_oth(n,dli,dzci,dzfi,dt_r,rho12,beta12,bforce,gacc,sigma,rho0,
                 q_ccc,q_pcc,q_cpc,q_ccp, &
                 s_ccc,s_pcc,s_cpc,s_ccp, &
                 k_ccc,k_pcc,k_cpc,k_ccp, &
-                bforcex,bforcey,bforcez,gaccx,gaccy,gaccz, &
+                bforcex,bforcey,bforcez,gaccx,gaccy,gaccz,rhox_av,rhoy_av,rhoz_av, &
                 dzci_c,dzci_m,dzfi_c,dzfi_p, &
                 psixp,psiyp,psizp
     real(rp) :: rhoxp,rhoyp,rhozp, &
@@ -889,6 +889,7 @@ subroutine mom_xyz_oth(n,dli,dzci,dzfi,dt_r,rho12,beta12,bforce,gacc,sigma,rho0,
     dxi = dli(1)
     dyi = dli(2)
     surf_factor = sigma*2./(rho12(1)+rho12(2))
+    rhox_av = rho_av(1); rhoy_av = rho_av(2); rhoz_av = rho_av(3)
     !
     ! making an exception for this kernel -- private variables not explicitly mentioned for the sake of conciseness
     !                                        all scalars should be firstprivate/private
@@ -948,9 +949,9 @@ subroutine mom_xyz_oth(n,dli,dzci,dzfi,dt_r,rho12,beta12,bforce,gacc,sigma,rho0,
           dpdx = (p_pcc-p_ccc)*dxi
           dpdy = (p_cpc-p_ccc)*dyi
           dpdz = (p_ccp-p_ccc)*dzci_c
-          dudt_aux = bforcex + (rhoxp-rho_av)*gaccx
-          dvdt_aux = bforcey + (rhoyp-rho_av)*gaccy
-          dwdt_aux = bforcez + (rhozp-rho_av)*gaccz
+          dudt_aux = bforcex + (rhoxp-rhox_av)*gaccx
+          dvdt_aux = bforcey + (rhoyp-rhoy_av)*gaccy
+          dwdt_aux = bforcez + (rhozp-rhoz_av)*gaccz
           !
           skappaxp = 0.5*(k_pcc+k_ccc)*surf_factor
           skappayp = 0.5*(k_cpc+k_ccc)*surf_factor

src/rk.f90

@@ -26,7 +26,7 @@ subroutine rk(rkpar,n,dli,dzci,dzfi,dt,dt_r, &
     real(rp), intent(in) :: dt,dt_r
     real(rp), intent(in   ), dimension(0:) :: dzci,dzfi
     real(rp), intent(in   ), dimension(3)  :: bforce,gacc
-    real(rp), intent(in   )                :: sigma,rho_av
+    real(rp), intent(in   )                :: sigma,rho_av(3)
     real(rp), intent(in   ), dimension(2)  :: rho12,mu12,beta12
     real(rp), intent(in   )                :: rho0
     real(rp), intent(in   ), dimension(0:,0:,0:)           :: psi,kappa

src/utils.f90

@@ -7,7 +7,7 @@
 module mod_utils
   implicit none
   private
-  public bulk_mean,bulk_mean_12,f_sizeof,swap
+  public bulk_mean,bulk_mean_12,bulk_mean_12_stag,f_sizeof,swap
   !@acc public device_memory_footprint
 contains
   subroutine bulk_mean(n,grid_vol_ratio,p,mean)
@@ -65,6 +65,75 @@ subroutine bulk_mean_12(n,grid_vol_ratio,psi,p12,mean)
     !$acc wait(1)
     call MPI_ALLREDUCE(MPI_IN_PLACE,mean,1,MPI_REAL_RP,MPI_SUM,MPI_COMM_WORLD,ierr)
   end subroutine bulk_mean_12
+  subroutine bulk_mean_12_stag(n,idir,grid_vol_ratio,psi,p12,mean)
+    !
+    ! compute the mean value of an observable over the entire domain
+    !
+    use mpi
+    use mod_types
+    implicit none
+    integer , intent(in), dimension(3) :: n
+    integer , intent(in)               :: idir
+    real(rp), intent(in), dimension(0:) :: grid_vol_ratio
+    real(rp), intent(in), dimension(0:,0:,0:) :: psi
+    real(rp), intent(in), dimension(2)        :: p12
+    real(rp), intent(out) :: mean
+    real(rp) :: psi_loc
+    integer  :: i,j,k
+    integer  :: ierr
+    mean = 0.
+    select case(idir)
+    case(1)
+      !$acc data copy(mean) async(1)
+      !$acc parallel loop collapse(3) default(present) private(psi_loc) reduction(+:mean) async(1)
+      do k=1,n(3)
+        do j=1,n(2)
+          do i=1,n(1)
+            psi_loc = 0.5*(psi(i+1,j,k)+psi(i,j,k))
+            mean = mean + (psi_loc*p12(1)+(1.-psi_loc)*p12(2))*grid_vol_ratio(k)
+          end do
+        end do
+      end do
+      !$acc end data
+    case(2)
+      !$acc data copy(mean) async(1)
+      !$acc parallel loop collapse(3) default(present) private(psi_loc) reduction(+:mean) async(1)
+      do k=1,n(3)
+        do j=1,n(2)
+          do i=1,n(1)
+            psi_loc = 0.5*(psi(i,j+1,k)+psi(i,j,k))
+            mean = mean + (psi_loc*p12(1)+(1.-psi_loc)*p12(2))*grid_vol_ratio(k)
+          end do
+        end do
+      end do
+      !$acc end data
+    case(3)
+      !$acc data copy(mean) async(1)
+      !$acc parallel loop collapse(3) default(present) private(psi_loc) reduction(+:mean) async(1)
+      do k=1,n(3)
+        do j=1,n(2)
+          do i=1,n(1)
+            psi_loc = 0.5*(psi(i,j,k+1)+psi(i,j,k))
+            mean = mean + (psi_loc*p12(1)+(1.-psi_loc)*p12(2))*grid_vol_ratio(k)
+          end do
+        end do
+      end do
+      !$acc end data
+    case default 
+      !$acc data copy(mean) async(1)
+      !$acc parallel loop collapse(3) default(present) reduction(+:mean) async(1)
+      do k=1,n(3)
+        do j=1,n(2)
+          do i=1,n(1)
+            mean = mean + (psi(i,j,k)*p12(1)+(1.-psi(i,j,k))*p12(2))*grid_vol_ratio(k)
+          end do
+        end do
+      end do
+      !$acc end data
+    end select
+    !$acc wait(1)
+    call MPI_ALLREDUCE(MPI_IN_PLACE,mean,1,MPI_REAL_RP,MPI_SUM,MPI_COMM_WORLD,ierr)
+  end subroutine bulk_mean_12_stag
   pure integer function f_sizeof(val) result(isize)
     !
     ! returns storage size of the scalar argument val in bytes