Skip to content

Commit

Permalink
Move some SHT SAL related allocations to init
Browse files Browse the repository at this point in the history 
* In MOM_tidal_forcing module, spherical harmonic coefficients (for SAL)
are now parts of tidal_forcing_CS to avoid repeated allocations. The
same applies to the Love number scaling factors.

* Allocations for arrays used for reproducing sums are moved to
subroutine spherical_harmonics_init in the MOM_spherical_harmonics
module.
  • Loading branch information
@herrwang0 @Hallberg-NOAA
herrwang0 authored and Hallberg-NOAA committed Oct 13, 2022
1 parent 1c83b2d commit a2e883a
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 24 deletions.
31 changes: 18 additions & 13 deletions src/parameterizations/lateral/MOM_spherical_harmonics.F90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ module MOM_spherical_harmonics
real, allocatable :: complexFactorRe(:,:,:), complexFactorIm(:,:,:), & !< Precomputed exponential factors
complexExpRe(:,:,:), complexExpIm(:,:,:) !! at the t-cells [nondim].
real, allocatable :: aRecurrenceCoeff(:,:), bRecurrenceCoeff(:,:) !< Precomputed recurrennce coefficients [nondim].
real, allocatable :: SnmRe_reproSum(:,:,:), SnmIm_reproSum(:,:,:)
logical :: reprod_sum !< True if use reproducable global sums
end type sht_CS

Expand All @@ -36,9 +37,9 @@ module MOM_spherical_harmonics
contains
subroutine spherical_harmonics_forward(G, CS, var, SnmRe, SnmIm, Nd)
type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure.
type(sht_CS), intent(in) :: CS !< Control structure for SHT
type(sht_CS), intent(inout) :: CS !< Control structure for SHT
real, dimension(SZI_(G),SZJ_(G)), &
intent(in) :: var(:,:) !< Input 2-D variable
intent(in) :: var !< Input 2-D variable
real, intent(out) :: SnmRe(:), & !< Output real and imaginary SHT coefficients
SnmIm(:) !! [nondim]
integer, optional, intent(in) :: Nd !< Maximum degree of the spherical harmonics
Expand All @@ -50,10 +51,10 @@ subroutine spherical_harmonics_forward(G, CS, var, SnmRe, SnmIm, Nd)
pmn, & ! Current associated Legendre polynomials of degree n and order m
pmnm1, & ! Associated Legendre polynomials of degree n-1 and order m
pmnm2 ! Associated Legendre polynomials of degree n-2 and order m
real, allocatable :: SnmRe_reproSum(:,:,:), SnmIm_reproSum(:,:,:)
integer :: i, j, k
integer :: is, ie, js, je
integer :: m, n, l
! real, dimension(SZI_(G),SZJ_(G), CS%lmax) :: SnmRe_reproSum, SnmIm_reproSum

if (.not.CS%initialized) call MOM_error(FATAL, "MOM_spherical_harmonics " // &
"spherical_harmonics_forward: Module must be initialized before it is used.")
Expand All @@ -69,23 +70,20 @@ subroutine spherical_harmonics_forward(G, CS, var, SnmRe, SnmIm, Nd)
do l=1,Ltot ; SnmRe(l) = 0.0; SnmIm(l) = 0.0 ; enddo

if (CS%reprod_sum) then
allocate(SnmRe_reproSum(is:ie, js:je, CS%lmax)); SnmRe_reproSum = 0.0
allocate(SnmIm_reproSum(is:ie, js:je, CS%lmax)); SnmIm_reproSum = 0.0

do m=0,Nmax
l = order2index(m, Nmax)

do j=js,je ; do i=is,ie
SnmRe_reproSum(i,j,l) = var(i,j) * CS%Pmm(i,j,m+1) * CS%complexFactorRe(i,j,m+1)
SnmIm_reproSum(i,j,l) = var(i,j) * CS%Pmm(i,j,m+1) * CS%complexFactorIm(i,j,m+1)
CS%SnmRe_reproSum(i,j,l) = var(i,j) * CS%Pmm(i,j,m+1) * CS%complexFactorRe(i,j,m+1)
CS%SnmIm_reproSum(i,j,l) = var(i,j) * CS%Pmm(i,j,m+1) * CS%complexFactorIm(i,j,m+1)
pmnm2(i,j) = 0.0
pmnm1(i,j) = CS%Pmm(i,j,m+1)
enddo ; enddo

do n = m+1, Nmax ; do j=js,je ; do i=is,ie
pmn(i,j) = CS%aRecurrenceCoeff(n+1,m+1) * CS%cosCoLatT(i,j) * pmnm1(i,j) - CS%bRecurrenceCoeff(n+1,m+1) * pmnm2(i,j)
SnmRe_reproSum(i,j,l+n-m) = var(i,j) * pmn(i,j) * CS%complexFactorRe(i,j,m+1)
SnmIm_reproSum(i,j,l+n-m) = var(i,j) * pmn(i,j) * CS%complexFactorIm(i,j,m+1)
CS%SnmRe_reproSum(i,j,l+n-m) = var(i,j) * pmn(i,j) * CS%complexFactorRe(i,j,m+1)
CS%SnmIm_reproSum(i,j,l+n-m) = var(i,j) * pmn(i,j) * CS%complexFactorIm(i,j,m+1)
pmnm2(i,j) = pmnm1(i,j)
pmnm1(i,j) = pmn(i,j)
enddo ; enddo ; enddo
Expand Down Expand Up @@ -115,8 +113,8 @@ subroutine spherical_harmonics_forward(G, CS, var, SnmRe, SnmIm, Nd)

if (CS%reprod_sum) then
do l=1,Ltot
SnmRe(l) = reproducing_sum(SnmRe_reproSum(:,:,l))
SnmIm(l) = reproducing_sum(SnmIm_reproSum(:,:,l))
SnmRe(l) = reproducing_sum(CS%SnmRe_reproSum(:,:,l))
SnmIm(l) = reproducing_sum(CS%SnmIm_reproSum(:,:,l))
enddo
else
call sum_across_PEs(SnmRe, Ltot)
Expand All @@ -134,7 +132,7 @@ subroutine spherical_harmonics_inverse(G, CS, SnmRe, SnmIm, var, Nd)
real, intent(in) :: SnmRe(:), & !< Real and imaginary SHT coefficients with
SnmIm(:) !! any scaling factors such as Love numbers [nondim]
real, dimension(SZI_(G),SZJ_(G)), &
intent(out) :: var(:,:) !< Output 2-D variable
intent(out) :: var !< Output 2-D variable
integer, optional, intent(in) :: Nd !< Maximum degree of the spherical harmonics
!! overriding nOrder in the CS [nondim]
! local variables
Expand Down Expand Up @@ -268,6 +266,11 @@ subroutine spherical_harmonics_init(G, param_file, CS)
enddo ; enddo
enddo

if (CS%reprod_sum) then
allocate(CS%SnmRe_reproSum(is:ie, js:je, CS%lmax)); CS%SnmRe_reproSum = 0.0
allocate(CS%SnmIm_reproSum(is:ie, js:je, CS%lmax)); CS%SnmIm_reproSum = 0.0
endif

id_clock_sht = cpu_clock_id('(Ocean spherical harmonics)', grain=CLOCK_MODULE)
id_clock_sht_forward = cpu_clock_id('(Ocean SHT forward)', grain=CLOCK_ROUTINE)
id_clock_sht_inverse = cpu_clock_id('(Ocean SHT inverse)', grain=CLOCK_ROUTINE)
Expand All @@ -282,6 +285,8 @@ subroutine spherical_harmonics_end(CS)
deallocate(CS%Pmm)
deallocate(CS%complexFactorRe, CS%complexFactorIm, CS%complexExpRe, CS%complexExpIm)
deallocate(CS%aRecurrenceCoeff, CS%bRecurrenceCoeff)
if (CS%reprod_sum) &
deallocate(CS%SnmRe_reproSum, CS%SnmIm_reproSum)
end subroutine spherical_harmonics_end

!> The function calc_lmax returns the number of real elements (cosine) of the spherical harmonics,
Expand Down
23 changes: 12 additions & 11 deletions src/parameterizations/lateral/MOM_tidal_forcing.F90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@ module MOM_tidal_forcing
type(sht_CS) :: sht
integer :: sal_sht_Nd
real, allocatable :: LoveScaling(:)
real, allocatable :: SnmRe(:), SnmIm(:)
end type tidal_forcing_CS

integer :: id_clock_tides !< CPU clock for tides
Expand Down Expand Up @@ -538,6 +539,9 @@ subroutine tidal_forcing_init(Time, G, US, param_file, CS)
"calculating the self-attraction and loading term for tides.", &
default=0, do_not_log=.not. CS%tidal_sal_sht)
lmax = calc_lmax(CS%sal_sht_Nd)
allocate(CS%SnmRe(lmax)); CS%SnmRe = 0.0
allocate(CS%SnmIm(lmax)); CS%SnmIm = 0.0

allocate(CS%LoveScaling(lmax))
call calc_love_scaling(CS%sal_sht_Nd, CS%LoveScaling)
call spherical_harmonics_init(G, param_file, CS%sht)
Expand Down Expand Up @@ -730,30 +734,25 @@ subroutine calc_SAL_sht(eta, eta_sal, G, CS)
real, dimension(SZI_(G),SZJ_(G)), intent(out) :: eta_sal !< The sea surface height anomaly from
!! a time-mean geoid [Z ~> m].
! type(sht_CS), intent(in) :: sht
type(tidal_forcing_CS), intent(in) :: CS !< Tidal forcing control struct
real, allocatable :: SnmRe(:), SnmIm(:)
type(tidal_forcing_CS), intent(inout) :: CS !< Tidal forcing control struct
real, allocatable :: LoveScaling(:)

integer :: n, m, l
integer :: lmax
lmax = calc_lmax(CS%sal_sht_Nd)

call cpu_clock_begin(id_clock_SAL)

allocate(SnmRe(lmax)); SnmRe = 0.0
allocate(SnmIm(lmax)); SnmIm = 0.0

call spherical_harmonics_forward(G, CS%sht, eta, SnmRe, SnmIm, CS%sal_sht_Nd)
call spherical_harmonics_forward(G, CS%sht, eta, CS%SnmRe, CS%SnmIm, CS%sal_sht_Nd)

! Multiply scaling to each mode
do m = 0,CS%sal_sht_Nd
l = order2index(m,CS%sal_sht_Nd)
do n = m,CS%sal_sht_Nd
SnmRe(l+n-m) = SnmRe(l+n-m)*CS%LoveScaling(l+n-m)
SnmIm(l+n-m) = SnmIm(l+n-m)*CS%LoveScaling(l+n-m)
CS%SnmRe(l+n-m) = CS%SnmRe(l+n-m)*CS%LoveScaling(l+n-m)
CS%SnmIm(l+n-m) = CS%SnmIm(l+n-m)*CS%LoveScaling(l+n-m)
enddo
enddo

call spherical_harmonics_inverse(G, CS%sht, SnmRe, SnmIm, eta_sal, CS%sal_sht_Nd)
call spherical_harmonics_inverse(G, CS%sht, CS%SnmRe, CS%SnmIm, eta_sal, CS%sal_sht_Nd)

call cpu_clock_end(id_clock_SAL)
end subroutine calc_SAL_sht
Expand All @@ -776,6 +775,8 @@ subroutine tidal_forcing_end(CS)

if (CS%tidal_sal_sht) then
if (allocated(CS%LoveScaling)) deallocate(CS%LoveScaling)
if (allocated(CS%SnmRe)) deallocate(CS%SnmRe)
if (allocated(CS%SnmIm)) deallocate(CS%SnmIm)
call spherical_harmonics_end(CS%sht)
endif
end subroutine tidal_forcing_end
Expand Down

0 comments on commit a2e883a

Please sign in to comment.