Skip to content

Commit

Permalink
Optimize the reading of ensembles and setup for global multiscale runs (
Browse files Browse the repository at this point in the history
NOAA-EMC#594)

This update improves the efficiency of the GSI, especially for
multiscale runs. Details can be found in issue NOAA-EMC#585
  • Loading branch information
jderber-NOAA committed Sep 22, 2023
1 parent aa656c6 commit 2f4e7fe
Show file tree
Hide file tree
Showing 13 changed files with 790 additions and 546 deletions.
152 changes: 100 additions & 52 deletions src/gsi/apply_scaledepwgts.f90
Original file line number Diff line number Diff line change
Expand Up @@ -16,18 +16,18 @@ function fwgtofwvlen (rvlft,rvrgt,rcons,rlen,rinput)
!
!$$$ end documentation block

use kinds, only: r_kind,i_kind,r_single
use kinds, only: r_kind
implicit none

real(r_kind),intent(in) :: rvlft,rvrgt,rcons,rlen,rinput
real(r_kind) :: fwgtofwvlen
real(r_kind) :: rlen1,rtem1,rconshalf

rlen1=rlen/10.0_r_kind ! rlen corresponds to a (-5,5) region
rconshalf=0.5_r_kind*rcons
if(rinput > rvlft .and. rinput < rvrgt) then
fwgtofwvlen=rcons
else
rlen1=rlen/10.0_r_kind ! rlen corresponds to a (-5,5) region
rconshalf=0.5_r_kind*rcons
rtem1=min(abs(rinput-rvlft),abs(rinput-rvrgt))
fwgtofwvlen=rconshalf*(1.0_r_kind+tanh(5.0_r_kind-rtem1/rlen1))
endif
Expand All @@ -41,23 +41,21 @@ subroutine init_mult_spc_wgts(jcap_in)
!
!$$$ end documentation block

use kinds, only: r_kind,i_kind,r_single
use constants, only: zero,half,one,two,three,rearth,pi,tiny_r_kind
use kinds, only: r_kind,i_kind
use constants, only: zero,half,one,rearth,pi,tiny_r_kind
use mpimod, only: mype
use general_sub2grid_mod, only: general_sub2grid_create_info
use egrid2agrid_mod,only: g_create_egrid2agrid
use general_sub2grid_mod, only: sub2grid_info
use hybrid_ensemble_parameters, only: nsclgrp
use hybrid_ensemble_parameters, only: spc_multwgt,spcwgt_params,r_ensloccov4scl
implicit none

integer(i_kind),intent(in ) :: jcap_in

integer(i_kind) i
integer(i_kind) i,l,ks
integer(i_kind) ig
real(r_kind) rwv0,rtem1,rtem2
real (r_kind):: fwgtofwvlen
real(r_kind) :: rwv0,rtem1,rtem2
real(r_kind) :: fwgtofwvlen
real(r_kind) :: totwvlength
real(r_kind),dimension(0:jcap_in,nsclgrp) :: spcwgt
logical :: l_sum_spc_weights

! Spectral scale decomposition is differernt between SDL-cross and SDL-nocross
Expand All @@ -67,9 +65,9 @@ subroutine init_mult_spc_wgts(jcap_in)
l_sum_spc_weights = .true.
end if

spc_multwgt(0,1)=one
spcwgt(0,1)=one
do ig=2,nsclgrp
spc_multwgt(0,ig)=zero
spcwgt(0,ig)=zero
end do


Expand All @@ -79,34 +77,66 @@ subroutine init_mult_spc_wgts(jcap_in)
rtem1=zero
do ig=1,nsclgrp
if(ig /= 2) then
spc_multwgt(i,ig)=fwgtofwvlen(spcwgt_params(1,ig),spcwgt_params(2,ig),&
spcwgt_params(3,ig),spcwgt_params(4,ig),totwvlength)
spc_multwgt(i,ig)=min(max(spc_multwgt(i,ig),zero),one)
spcwgt(i,ig)=fwgtofwvlen(spcwgt_params(1,ig),spcwgt_params(2,ig),&
spcwgt_params(3,ig),spcwgt_params(4,ig),totwvlength)
spcwgt(i,ig)=min(max(spcwgt(i,ig),zero),one)
if(l_sum_spc_weights) then
rtem1=rtem1+spc_multwgt(i,ig)
rtem1=rtem1+spcwgt(i,ig)
else
rtem1=rtem1+spc_multwgt(i,ig)*spc_multwgt(i,ig)
rtem1=rtem1+spcwgt(i,ig)*spcwgt(i,ig)
endif
endif
enddo
rtem2 =1.0_r_kind - rtem1
if(rtem2 >= zero) then

if(l_sum_spc_weights) then
spc_multwgt(i,2)=rtem2
spcwgt(i,2)=rtem2
else
spc_multwgt(i,2)=sqrt(rtem2)
spcwgt(i,2)=sqrt(rtem2)
endif
else
if(mype == 0)write(6,*) ' rtem2 < zero ',i,rtem2,(spc_multwgt(i,ig),ig=1,nsclgrp)
spc_multwgt(i,2)=zero
if(mype == 0)write(6,*) ' rtem2 < zero ',i,rtem2,(spcwgt(i,ig),ig=1,nsclgrp)
spcwgt(i,2)=zero
endif
enddo
!! Code borrowed from spvar in splib

spc_multwgt = zero
do ig=1,nsclgrp
do i=0,jcap_in
ks=2*i
spc_multwgt(ks+1,ig)=spcwgt(i,ig)
end do
do i=0,jcap_in
do l=MAX(1,i-jcap_in),MIN(i,jcap_in)
ks=l*(2*jcap_in+1-l)+2*i
spc_multwgt(ks+1,ig) = spcwgt(i,ig)
spc_multwgt(ks+2,ig) = spcwgt(i,ig)
end do
end do
end do


return
end subroutine init_mult_spc_wgts
subroutine destroy_mult_spc_wgts
!$$$ subprogram documentation block
!
! subprogram: destroy_mult_spc_wgts
!
!$$$ end documentation block

use hybrid_ensemble_parameters, only: spc_multwgt,spcwgt_params
implicit none

deallocate(spc_multwgt,spcwgt_params)

subroutine apply_scaledepwgts(grd_in,sp_in,wbundle,spwgts,wbundle2)
return
end subroutine destroy_mult_spc_wgts


subroutine apply_scaledepwgts(m,grd_in,sp_in)
!
! Program history log:
! 2017-03-30 J. Kay, X. Wang - copied from Kleist's apply_scaledepwgts and
Expand All @@ -115,49 +145,67 @@ subroutine apply_scaledepwgts(grd_in,sp_in,wbundle,spwgts,wbundle2)
!
use constants, only: one
use control_vectors, only: control_vector
use kinds, only: r_kind,i_kind
use kinds, only: r_single
use general_specmod, only: general_spec_multwgt
use kinds, only: r_kind,i_kind,r_single
use gsi_bundlemod, only: gsi_bundle
use general_sub2grid_mod, only: general_sub2grid,general_grid2sub
use general_specmod, only: spec_vars
use general_sub2grid_mod, only: sub2grid_info
use mpimod, only: mpi_comm_world,mype
use hybrid_ensemble_parameters, only: spc_multwgt,en_perts,nsclgrp,n_ens
use mpimod, only: mype
implicit none

! Declare passed variables
type(gsi_bundle),intent(in) :: wbundle
type(gsi_bundle),intent(inout) :: wbundle2
integer,intent(in) :: m
type(spec_vars),intent (in):: sp_in
type(sub2grid_info),intent(in)::grd_in
real(r_kind),dimension(0:sp_in%jcap),intent(in):: spwgts

! Declare local variables
integer(i_kind) kk
integer(i_kind) kk,ig,n,ig2,i,j

real(r_kind),dimension(grd_in%nlat*grd_in%nlon*grd_in%nlevs_alloc) :: hwork
real(r_kind),dimension(grd_in%nlat,grd_in%nlon,grd_in%nlevs_alloc) :: work
real(r_kind),dimension(sp_in%nc):: spc1

! Beta1 first
! Get from subdomains to
call general_sub2grid(grd_in,wbundle%values,hwork)
work=reshape(hwork,(/grd_in%nlat,grd_in%nlon,grd_in%nlevs_alloc/))

do kk=1,grd_in%nlevs_alloc
! Transform from physical space to spectral space
call general_g2s0(grd_in,sp_in,spc1,work(:,:,kk))

! Apply spectral weights
call general_spec_multwgt(sp_in,spc1,spwgts)
! Transform back to physical space
call general_s2g0(grd_in,sp_in,spc1,work(:,:,kk))
real(r_single),dimension(grd_in%nlat,grd_in%nlon,grd_in%nlevs_alloc,nsclgrp) :: hwork2
real(r_kind),dimension(grd_in%nlat,grd_in%nlon) :: work
real(r_kind),dimension(sp_in%nc,grd_in%nlevs_alloc):: spc1
real(r_kind),dimension(sp_in%nc):: spc2

do n=1,n_ens
! Get from subdomains to full grid
call general_sub2grid(grd_in,en_perts(n,1,m)%valuesr4(:),hwork2(:,:,:,1))

!$omp parallel do schedule(static,1) private(i,j,kk,work)
do kk=1,grd_in%nlevs_loc
do j=1,grd_in%nlon
do i=1,grd_in%nlat
work(i,j)=hwork2(i,j,kk,1)
end do
end do
! Transform from physical space to spectral space
call general_g2s0(grd_in,sp_in,spc1(1,kk),work)

end do
!$omp parallel do schedule(static,1) private(kk,ig,ig2,i,j,work,spc2)
do ig2=1,nsclgrp*grd_in%nlevs_loc
ig=(ig2-1)/grd_in%nlevs_loc+1
kk=ig2-(ig-1)*grd_in%nlevs_loc

do i=1,sp_in%nc
spc2(i)=spc1(i,kk)*spc_multwgt(i,ig)
end do
! Apply spectral weights
! Transform back to physical space
call general_s2g0(grd_in,sp_in,spc2,work)

do j=1,grd_in%nlon
do i=1,grd_in%nlat
hwork2(i,j,kk,ig)=work(i,j)
end do
end do
end do
do ig=1,nsclgrp

! Transfer work back to subdomains
call general_grid2sub(grd_in,hwork2(:,:,:,ig),en_perts(n,ig,m)%valuesr4(:))
end do
end do

! Transfer work back to subdomains
hwork=reshape(work,(/grd_in%nlat*grd_in%nlon*grd_in%nlevs_alloc/))
call general_grid2sub(grd_in,hwork,wbundle2%values)

return
end subroutine apply_scaledepwgts
8 changes: 5 additions & 3 deletions src/gsi/control_vectors.f90
Original file line number Diff line number Diff line change
Expand Up @@ -897,21 +897,23 @@ real(r_quad) function qdot_prod_sub(xcv,ycv)
itot=max(m3d,0)+max(m2d,0)
if(l_hyb_ens)itot=itot+n_ens*naensgrp
allocate(partsum(itot))
partsum=zero_quad
do ii=1,nsubwin
!$omp parallel do schedule(dynamic,1) private(i)
do i = 1,m3d
partsum(i) = dplevs(xcv%step(ii)%r3(i)%q,ycv%step(ii)%r3(i)%q,ihalo=1)
partsum(i) = partsum(i)+dplevs(xcv%step(ii)%r3(i)%q,ycv%step(ii)%r3(i)%q,ihalo=1)
enddo
!$omp parallel do schedule(dynamic,1) private(i)
do i = 1,m2d
partsum(m3d+i) = dplevs(xcv%step(ii)%r2(i)%q,ycv%step(ii)%r2(i)%q,ihalo=1)
partsum(m3d+i) = partsum(m3d+i)+dplevs(xcv%step(ii)%r2(i)%q,ycv%step(ii)%r2(i)%q,ihalo=1)
enddo
if(l_hyb_ens) then
do ig=1,naensgrp
nigtmp=n_ens*(ig-1)
!$omp parallel do schedule(dynamic,1) private(i)
do i = 1,n_ens
partsum(m3d+m2d+nigtmp+i) = dplevs(xcv%aens(ii,ig,i)%r3(1)%q,ycv%aens(ii,ig,i)%r3(1)%q,ihalo=1)
partsum(m3d+m2d+nigtmp+i) = partsum(m3d+m2d+nigtmp+i) + &
dplevs(xcv%aens(ii,ig,i)%r3(1)%q,ycv%aens(ii,ig,i)%r3(1)%q,ihalo=1)
end do
end do
end if
Expand Down
Loading

0 comments on commit 2f4e7fe

Please sign in to comment.