diff --git a/config_src/drivers/timing_tests/time_MOM_ANN.F90 b/config_src/drivers/timing_tests/time_MOM_ANN.F90 index a399835d89..a6906970dd 100644 --- a/config_src/drivers/timing_tests/time_MOM_ANN.F90 +++ b/config_src/drivers/timing_tests/time_MOM_ANN.F90 @@ -5,7 +5,7 @@ program time_MOM_ANN use MOM_ANN, only : ANN_CS use MOM_ANN, only : ANN_allocate, ANN_apply, ANN_end use MOM_ANN, only : ANN_apply_vector_orig, ANN_apply_vector_oi -use MOM_ANN, only : ANN_apply_array_sio +use MOM_ANN, only : ANN_apply_array_sio, ANN_apply_array_sio_r4 use MOM_ANN, only : ANN_random implicit none @@ -72,7 +72,10 @@ program time_MOM_ANN 2, "MOM_ANN:ANN_apply_vector_oi(array)") write(*,"(',')") call time_ANN(nlayers, nin, layer_width, nout, nsamp, nits, nxy, & - 12, "MOM_ANN:ANN_apply_array_sio(array)") + 3, "MOM_ANN:ANN_apply_array_sio(array)") +write(*,"(',')") +call time_ANN(nlayers, nin, layer_width, nout, nsamp, nits, nxy, & + 4, "MOM_ANN:ANN_apply_array_sio_r4(array)") write(*,"()") write(*,'(a)') "}" @@ -101,9 +104,9 @@ subroutine time_ANN(nlayers, nin, width, nout, nsamp, nits, nxy, impl, label) real :: x_s(nin) ! Inputs (just features) [nondim] real :: y_s(nin) ! Outputs (just features) [nondim] real :: x_fs(nin,nxy) ! Inputs (feature, space) [nondim] - real :: y_fs(nin,nxy) ! Outputs (feature, space) [nondim] - real :: x_sf(nin,nxy) ! Inputs (space, feature) [nondim] - real :: y_sf(nin,nxy) ! Outputs (space, feature) [nondim] + real :: y_fs(nout,nxy) ! Outputs (feature, space) [nondim] + real :: x_sf(nxy,nin) ! Inputs (space, feature) [nondim] + real :: y_sf(nxy,nout) ! Outputs (space, feature) [nondim] integer :: iter, samp ! Loop counters integer :: ij ! Horizontal loop index real :: start, finish, timing ! CPU times [s] @@ -117,6 +120,7 @@ subroutine time_ANN(nlayers, nin, width, nout, nsamp, nits, nxy, impl, label) widths(nlayers) = nout call ANN_random(ANN, nlayers, widths) + call random_number(x_s) call random_number(x_fs) call random_number(x_sf) @@ -131,7 +135,6 @@ subroutine time_ANN(nlayers, nin, width, nout, nsamp, nits, nxy, impl, label) do samp = 1, nsamp select case (impl) case (0) - aits = nits call cpu_time(start) do iter = 1, nits ! Make many passes to reduce sampling error call ANN_apply(x_s, y_s, ANN) @@ -153,13 +156,20 @@ subroutine time_ANN(nlayers, nin, width, nout, nsamp, nits, nxy, impl, label) enddo enddo call cpu_time(finish) - case (12) + case (3) call cpu_time(start) do iter = 1, aits ! Make many passes to reduce sampling error call ANN_apply_array_sio(nxy, x_sf(:,:), y_sf(:,:), ANN) enddo call cpu_time(finish) asamp = nsamp * aits ! Account for working on whole arrays + case (4) + call cpu_time(start) + do iter = 1, aits ! Make many passes to reduce sampling error + call ANN_apply_array_sio_r4(nxy, x_sf(:,:), y_sf(:,:), ANN) + enddo + call cpu_time(finish) + asamp = nsamp * aits ! Account for working on whole arrays end select timing = ( finish - start ) / real(nits) ! Average time per call diff --git a/src/framework/MOM_ANN.F90 b/src/framework/MOM_ANN.F90 index 56cafcf7d5..881bc98700 100644 --- a/src/framework/MOM_ANN.F90 +++ b/src/framework/MOM_ANN.F90 @@ -13,6 +13,7 @@ module MOM_ANN public ANN_init, ANN_allocate, ANN_apply, ANN_end, ANN_unit_tests public ANN_apply_vector_orig, ANN_apply_vector_oi, ANN_apply_array_sio +public ANN_apply_array_sio_r4 public set_layer, set_input_normalization, set_output_normalization public ANN_random, randomize_layer @@ -34,6 +35,8 @@ module MOM_ANN real, allocatable :: A(:,:) !< Matrix in column-major order !! of size A(output_width, input_width) [nondim] real, allocatable :: b(:) !< bias vector of size output_width [nondim] + real(4), allocatable :: A_r4(:,:) !< Same as A(:,:) but in real(4) [nondim] + real(4), allocatable :: b_r4(:) !< Same as b(:) but in real(4) [nondim] end type layer_type !> Control structure/type for ANN @@ -117,10 +120,12 @@ subroutine ANN_init(CS, NNfile) fieldname = trim('A') // trim(layer_num_str) call MOM_read_data(NNfile, fieldname, CS%layers(i)%A, & (/1,1,1,1/),(/CS%layers(i)%output_width,CS%layers(i)%input_width,1,1/)) + CS%layers(i)%A_r4(:,:) = real(CS%layers(i)%A(:,:), kind=4) ! Reading bias b fieldname = trim('b') // trim(layer_num_str) call MOM_read_data(NNfile, fieldname, CS%layers(i)%b) + CS%layers(i)%b_r4(:) = real(CS%layers(i)%b(:), kind=4) enddo ! No activation function for the last layer @@ -170,6 +175,8 @@ subroutine ANN_allocate(CS, num_layers, layer_sizes) allocate( CS%layers(l)%A(CS%layers(l)%output_width, CS%layers(l)%input_width) ) allocate( CS%layers(l)%b(CS%layers(l)%output_width) ) + allocate( CS%layers(l)%A_r4(CS%layers(l)%output_width, CS%layers(l)%input_width) ) + allocate( CS%layers(l)%b_r4(CS%layers(l)%output_width) ) CS%parameters = CS%parameters & + CS%layer_sizes(l) * CS%layer_sizes(l+1) & ! For weights @@ -228,6 +235,8 @@ subroutine ANN_end(CS) do i = 1, CS%num_layers-1 deallocate(CS%layers(i)%A) deallocate(CS%layers(i)%b) + deallocate(CS%layers(i)%A_r4) + deallocate(CS%layers(i)%b_r4) enddo deallocate(CS%layers) @@ -242,6 +251,15 @@ pure elemental function activation_fn(x) result (y) end function activation_fn +!> The default activation function in real(4) precision +pure elemental function activation_fn_r4(x) result (y) + real(4), intent(in) :: x !< Scalar input value [nondim] + real(4) :: y !< Scalar output value [nondim] + + y = max(x, 0.0_4) ! ReLU activation + +end function activation_fn_r4 + !> Single application of ANN inference using vector input and output !! !! This implementation is the simplest using allocation and de-allocation @@ -440,6 +458,82 @@ subroutine layer_apply_sio(nij, x, y, layer) end subroutine layer_apply_sio end subroutine ANN_apply_array_sio +!> Same as ANN_apply_array_sio, but casts input and output +!! vectors to real(4) internally and performs ANN inference +!! in real(4) precision. On average, twice faster than original +!! ANN_apply_array_sio. +subroutine ANN_apply_array_sio_r4(nij, x, y, CS) + type(ANN_CS), intent(in) :: CS !< ANN control structure + integer, intent(in) :: nij !< Size of spatial dimension + real, intent(in) :: x(nij, CS%layer_sizes(1)) !< input [arbitrary] + real, intent(inout) :: y(nij, CS%layer_sizes(CS%num_layers)) !< output [arbitrary] + ! Local variables + real(4), allocatable :: x_1(:,:), x_2(:,:) ! intermediate states [nondim] + integer :: l, i, o ! Layer, input, output index + + allocate( x_1( nij, maxval( CS%layer_sizes(:) ) ) ) + allocate( x_2( nij, maxval( CS%layer_sizes(:) ) ) ) + + ! Normalize input + do i = 1, CS%layer_sizes(1) + x_1(:,i) = real(( x(:,i) - CS%input_means(i) ) * CS%input_norms(i), kind=4) + enddo + + ! Apply Linear layers + do l = 1, CS%num_layers-2, 2 + call layer_apply_sio(nij, x_1, x_2, CS%layers(l)) + call layer_apply_sio(nij, x_2, x_1, CS%layers(l+1)) + enddo + if (mod(CS%num_layers,2)==0) then + call layer_apply_sio(nij, x_1, x_2, CS%layers(CS%num_layers-1)) + ! Un-normalize output + do o = 1, CS%layer_sizes(CS%num_layers) + y(:,o) = real(x_2(:,o) * CS%output_norms(o) + CS%output_means(o), kind=8) + enddo + else + ! Un-normalize output + do o = 1, CS%layer_sizes(CS%num_layers) + y(:,o) = real(x_1(:,o) * CS%output_norms(o) + CS%output_means(o), kind=8) + enddo + endif + + deallocate(x_1, x_2) + + contains + + !> Applies linear layer to input data x and stores the result in y with + !! y = A*x + b with optional application of the activation function so the + !! overall operations is ReLU(A*x + b) + subroutine layer_apply_sio(nij, x, y, layer) + type(layer_type), intent(in) :: layer !< Linear layer + integer, intent(in) :: nij !< Size of spatial dimension + real(4), intent(in) :: x(nij, layer%input_width) !< Input vector [nondim] + real(4), intent(inout) :: y(nij, layer%output_width) !< Output vector [nondim] + ! Local variables + integer :: i, o ! Input, output indices + ! We introduce rescaling which gives bitwise the same answer if there is no underflow + ! We assume that overflow is unlikely because x is always on the order of one + real(4), parameter :: boost = 2.**33 ! Shifts exponent by ~ 1e+10 + real(4), parameter :: inv_boost = 2.**(-33) ! Shifts exponent back + + do o = 1, layer%output_width + ! Add bias + y(:,o) = layer%b_r4(o) * boost + ! Multiply by kernel + do i = 1, layer%input_width + y(:,o) = y(:,o) + (x(:,i) * boost) * layer%A_r4(o, i) + enddo + ! Apply activation function + if (layer%activation) then + y(:,o) = activation_fn_r4(y(:,o) * inv_boost) + else + y(:,o) = y(:,o) * inv_boost + endif + enddo + + end subroutine layer_apply_sio +end subroutine ANN_apply_array_sio_r4 + !> Sets weights and bias for a single layer subroutine set_layer(ANN, layer, weights, biases, activation) type(ANN_CS), intent(inout) :: ANN !< ANN control structure @@ -456,12 +550,14 @@ subroutine set_layer(ANN, layer, weights, biases, activation) if ( size(biases) /= size(ANN%layers(layer)%b) ) & call MOM_error(FATAL, "MOM_ANN, set_layer: mismatch in size of biases") ANN%layers(layer)%b(:) = biases(:) + ANN%layers(layer)%b_r4(:) = real(biases(:), kind=4) if ( size(weights,1) /= size(ANN%layers(layer)%A,1) ) & call MOM_error(FATAL, "MOM_ANN, set_layer: mismatch in size of weights (first dim)") if ( size(weights,2) /= size(ANN%layers(layer)%A,2) ) & call MOM_error(FATAL, "MOM_ANN, set_layer: mismatch in size of weights (second dim)") ANN%layers(layer)%A(:,:) = weights(:,:) + ANN%layers(layer)%A_r4(:,:) = real(weights(:,:), kind=4) ANN%layers(layer)%activation = activation end subroutine set_layer @@ -669,6 +765,9 @@ logical function ANN_unit_tests(verbose) ! as above with v5 of ANN_apply applied to 2d inputs, x(space,feature) call ANN_apply_array_sio(2, reshape([0.,1.,2.,3.,4.,5.,6.,7.],[2,4]), y2, ANN) call test%real_arr(2, y2, [2.,5.], 'Rectifier+summation+bias+norms 4-layer array v2') + + call ANN_apply_array_sio_r4(2, reshape([0.,1.,2.,3.,4.,5.,6.,7.],[2,4]), y2, ANN) + call test%real_arr(2, y2, [2.,5.], 'Rectifier+summation+bias+norms 4-layer array v2 real(4)') deallocate( y2 ) call ANN_end(ANN) @@ -685,6 +784,8 @@ logical function ANN_unit_tests(verbose) deallocate( y ) call ANN_random(ANN, nlay, widths) allocate( x(widths(1)), y(widths(nlay)), y_good(widths(nlay)) ) + call random_number(x) + x(:) = 2. * x(:) - 1. call ANN_apply_vector_orig(x, y_good, ANN) call ANN_apply_vector_oi(x, y, ANN) rand_res = rand_res .or. maxval( abs( y(:) - y_good(:) ) ) > 0. ! Check results from v2 = v1 @@ -695,6 +796,9 @@ logical function ANN_unit_tests(verbose) call ANN_apply_array_sio(20, x2, y2, ANN) rand_res = rand_res .or. maxval( abs( maxval(y2(:,:),1) - y_good(:) ) ) > 0. ! Check results from array v2 = v1 rand_res = rand_res .or. maxval( abs( minval(y2(:,:),1) - y_good(:) ) ) > 0. ! Check results from array v2 = v1 + call ANN_apply_array_sio_r4(20, x2, y2, ANN) + rand_res = rand_res .or. maxval( abs( maxval(y2(:,:),1) - y_good(:) ) ) > 1.e-5 ! Lower Real(4) precision + rand_res = rand_res .or. maxval( abs( minval(y2(:,:),1) - y_good(:) ) ) > 1.e-5 ! Lower Real(4) precision deallocate( x, y, y_good, x2, y2 ) call ANN_end(ANN) enddo diff --git a/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 b/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 index ebbac944ac..3118bde207 100644 --- a/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 +++ b/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 @@ -12,10 +12,10 @@ module MOM_Zanna_Bolton use MOM_domains, only : create_group_pass, do_group_pass, group_pass_type, & start_group_pass, complete_group_pass use MOM_domains, only : To_North, To_East -use MOM_domains, only : pass_var, CORNER +use MOM_domains, only : pass_var, pass_vector, CORNER use MOM_cpu_clock, only : cpu_clock_id, cpu_clock_begin, cpu_clock_end use MOM_cpu_clock, only : CLOCK_MODULE, CLOCK_ROUTINE -use MOM_ANN, only : ANN_init, ANN_apply, ANN_end, ANN_CS +use MOM_ANN, only : ANN_init, ANN_apply_array_sio_r4, ANN_end, ANN_CS implicit none ; private @@ -62,7 +62,8 @@ module MOM_Zanna_Bolton real, dimension(:,:,:), allocatable :: & Txx, & !< Subgrid stress xx component in h [L2 T-2 ~> m2 s-2] Tyy, & !< Subgrid stress yy component in h [L2 T-2 ~> m2 s-2] - Txy !< Subgrid stress xy component in q [L2 T-2 ~> m2 s-2] + Txy, & !< Subgrid stress xy component in q [L2 T-2 ~> m2 s-2] + Txy_h !< Subgrid stress xy component in h [L2 T-2 ~> m2 s-2] real, dimension(:,:), allocatable :: & kappa_h, & !< Scaling coefficient in h points [L2 ~> m2] @@ -97,7 +98,8 @@ module MOM_Zanna_Bolton integer :: id_clock_copy integer :: id_clock_cdiss integer :: id_clock_stress - integer :: id_clock_stress_ANN + integer :: id_clock_ANN_inference + integer :: id_clock_ANN_features integer :: id_clock_divergence integer :: id_clock_mpi integer :: id_clock_filter @@ -232,9 +234,10 @@ subroutine ZB2020_init(Time, G, GV, US, param_file, diag, CS, use_ZB2020) CS%id_clock_copy = cpu_clock_id('(ZB2020 copy fields)', grain=CLOCK_ROUTINE, sync=.false.) CS%id_clock_cdiss = cpu_clock_id('(ZB2020 compute c_diss)', grain=CLOCK_ROUTINE, sync=.false.) CS%id_clock_stress = cpu_clock_id('(ZB2020 compute stress)', grain=CLOCK_ROUTINE, sync=.false.) - CS%id_clock_stress_ANN = cpu_clock_id('(ZB2020 compute stress ANN)', grain=CLOCK_ROUTINE, sync=.false.) + CS%id_clock_ANN_inference = cpu_clock_id('(ZB2020 ANN inference)', grain=CLOCK_ROUTINE, sync=.false.) + CS%id_clock_ANN_features = cpu_clock_id('(ZB2020 ANN features)', grain=CLOCK_ROUTINE, sync=.false.) CS%id_clock_divergence = cpu_clock_id('(ZB2020 compute divergence)', grain=CLOCK_ROUTINE, sync=.false.) - CS%id_clock_mpi = cpu_clock_id('(ZB2020 filter MPI exchanges)', grain=CLOCK_ROUTINE, sync=.false.) + CS%id_clock_mpi = cpu_clock_id('(ZB2020 MPI exchanges)', grain=CLOCK_ROUTINE, sync=.false.) CS%id_clock_filter = cpu_clock_id('(ZB2020 filter no MPI)', grain=CLOCK_ROUTINE, sync=.false.) CS%id_clock_post = cpu_clock_id('(ZB2020 post data)', grain=CLOCK_ROUTINE, sync=.false.) CS%id_clock_source = cpu_clock_id('(ZB2020 compute energy source)', grain=CLOCK_ROUTINE, sync=.false.) @@ -253,6 +256,10 @@ subroutine ZB2020_init(Time, G, GV, US, param_file, diag, CS, use_ZB2020) allocate(CS%vort_xy(SZIB_(G),SZJB_(G),SZK_(GV)), source=0.) allocate(CS%hq(SZIB_(G),SZJB_(G),SZK_(GV))) + if (CS%use_ann) then + allocate(CS%Txy_h(SZI_(G),SZJ_(G),SZK_(GV)), source=0.) + endif + allocate(CS%Txx(SZI_(G),SZJ_(G),SZK_(GV)), source=0.) allocate(CS%Tyy(SZI_(G),SZJ_(G),SZK_(GV)), source=0.) allocate(CS%Txy(SZIB_(G),SZJB_(G),SZK_(GV)), source=0.) @@ -321,6 +328,9 @@ end subroutine ZB2020_init subroutine ZB2020_end(CS) type(ZB2020_CS), intent(inout) :: CS !< ZB2020 control structure. + if (CS%use_ann) then + deallocate(CS%Txy_h) + endif deallocate(CS%sh_xx) deallocate(CS%sh_xy) deallocate(CS%vort_xy) @@ -664,106 +674,135 @@ subroutine compute_stress_ANN_collocated(G, GV, CS) type(ZB2020_CS), intent(inout) :: CS !< ZB2020 control structure. integer :: is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz - integer :: i, j, k, n + integer :: i, j, k, n, m integer :: ii, jj + integer :: nij - real :: x(3*CS%stencil_size**2) ! Vector of non-dimensional input features + real, allocatable :: x(:,:) ! Vector of non-dimensional input features of size + ! number of horizontal grid points times ! (sh_xy, sh_xx, vort_xy) on a stencil [nondim] - real :: y(3) ! Vector of nondimensional - ! output features (Txy,Txx,Tyy) [nondim] - real :: input_norm ! Norm of input features [T-1 ~> s-1] + real, allocatable :: y(:,:) ! Vector of nondimensional output features of size + ! number of horizontal grid points times + ! (Txy,Txx,Tyy) [nondim] + real, allocatable :: input_norm(:) ! Inverse norm of input features in center points [T ~> s] + real, allocatable :: output_norm(:)! Norm of output features in center points [L2T-2 ~> m2 s-2] real :: tmp ! Temporal value of squared norm [T-2 ~> s-2] + real :: x1, x2, x3 ! Components of the velocity gradient tensor + ! (sh_xy, sh_xx, vort_xy) at a point [T-1 ~> s-1] integer :: offset ! Half the stencil size. Used for selection integer :: stencil_points ! The number of points after flattening - real, dimension(SZI_(G),SZJ_(G),SZK_(GV)) :: & - sh_xy_h, & ! sh_xy interpolated to the center [T-1 ~> s-1] - vort_xy_h, & ! vort_xy interpolated to the center [T-1 ~> s-1] - norm_h ! Norm of input feautres in center points [T-1 ~> s-1] - real, dimension(SZI_(G),SZJ_(G)) :: & - sqr_h, & ! Squared norm of velocity gradients in center points [T-2 ~> s-2] - Txy ! Predicted Txy in center points [T-1 ~> s-1] + sh_xy_h, & ! Shearing strain interpolated to h point [T-1 ~> s-1] + vort_xy_h ! Vorticity interpolated to h point [T-1 ~> s-1] - call cpu_clock_begin(CS%id_clock_stress_ANN) + type(group_pass_type) :: pass_vel_grads ! A handle used for group halo passes + type(group_pass_type) :: pass_flux ! A handle used for group halo passes is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = GV%ke Isq = G%IscB ; Ieq = G%IecB ; Jsq = G%JscB ; Jeq = G%JecB - sh_xy_h = 0. - vort_xy_h = 0. - norm_h = 0. - - call pass_var(CS%sh_xy, G%Domain, clock=CS%id_clock_mpi, position=CORNER) - call pass_var(CS%sh_xx, G%Domain, clock=CS%id_clock_mpi) - call pass_var(CS%vort_xy, G%Domain, clock=CS%id_clock_mpi, position=CORNER) + ! Number of horizontal grid points in ANN inference loop below + nij = (ie - is + 1) * (je - js + 1) + allocate(x(nij, 3 * CS%stencil_size**2)) + allocate(y(nij, 3)) + allocate(input_norm(nij)) + allocate(output_norm(nij)) + + ! If stencil_size==3, the halo required to apply the model + ! in center points is only 1. This halo is available + ! without MPI exhchange in symmetric and non-symmetric memory models + if (CS%stencil_size > 3) then + call create_group_pass(pass_vel_grads, CS%sh_xy, G%Domain, position=CORNER) + call create_group_pass(pass_vel_grads, CS%vort_xy, G%Domain, position=CORNER) + call do_group_pass(pass_vel_grads, G%Domain, clock=CS%id_clock_mpi) + call pass_var(CS%sh_xx, G%Domain, clock=CS%id_clock_mpi) + endif offset = (CS%stencil_size-1)/2 stencil_points = CS%stencil_size**2 - ! Interpolate input features do k=1,nz - do j=js-2,je+2 ; do i=is-2,ie+2 - ! It is assumed that B.C. is applied to sh_xy and vort_xy - sh_xy_h(i,j,k) = 0.25 * ( (CS%sh_xy(I-1,J-1,k) + CS%sh_xy(I,J,k)) & - + (CS%sh_xy(I-1,J,k) + CS%sh_xy(I,J-1,k)) ) - - vort_xy_h(i,j,k) = 0.25 * ( (CS%vort_xy(I-1,J-1,k) + CS%vort_xy(I,J,k)) & - + (CS%vort_xy(I-1,J,k) + CS%vort_xy(I,J-1,k)) ) - - sqr_h(i,j) = (CS%sh_xx(i,j,k)**2 + sh_xy_h(i,j,k)**2 + vort_xy_h(i,j,k)**2) * G%mask2dT(i,j) + call cpu_clock_begin(CS%id_clock_ANN_features) + ! Precompute interpolated values to efficiently reuse in the next loop. + ! Interpolation from corner to center assuming that B.C. + ! is already applied + do j=js-1,je+1 ; do i=is-1,ie+1 + sh_xy_h(i,j) = 0.25 * ( (CS%sh_xy(i-1,j-1,k) + CS%sh_xy(i,j,k)) & + + (CS%sh_xy(i-1,j,k) + CS%sh_xy(i,j-1,k)) ) + vort_xy_h(i,j) = 0.25 * ( (CS%vort_xy(i-1,j-1,k) + CS%vort_xy(i,j,k)) & + + (CS%vort_xy(i-1,j,k) + CS%vort_xy(i,j-1,k)) ) enddo; enddo + m = 0 do j=js,je ; do i=is,ie - tmp = 0.0 - do jj=j-offset,j+offset; do ii=i-offset,i+offset - tmp = tmp + sqr_h(ii,jj) - enddo; enddo - norm_h(i,j,k) = sqrt(tmp) + m = m + 1 + tmp = 0. + n = 0 + ! Fuse assembling a vector of input features + ! and computation of its norm + do jj = j-offset, j+offset + do ii = i-offset, i+offset + n = n + 1 + x1 = sh_xy_h(ii,jj) + x2 = CS%sh_xx(ii,jj,k) + x3 = vort_xy_h(ii,jj) + + x(m,n) = x1 + x(m,n+stencil_points) = x2 + x(m,n+2*stencil_points) = x3 + + tmp = tmp + (((x1*x1) + x2*x2) + x3*x3) + end do + end do + ! Momentum fluxes scale as dx^2 * |grad V|^2 + output_norm(m) = tmp * CS%kappa_h(i,j) + ! Input features are simply normalized by their norm + input_norm(m) = 1. / (sqrt(tmp) + CS%subroundoff_shear) enddo; enddo - enddo - - call pass_var(sh_xy_h, G%Domain, clock=CS%id_clock_mpi) - call pass_var(vort_xy_h, G%Domain, clock=CS%id_clock_mpi) - call pass_var(norm_h, G%Domain, clock=CS%id_clock_mpi) - - do k=1,nz - do j=js-2,je+2 ; do i=is-2,ie+2 - x(1:stencil_points) = & - RESHAPE(sh_xy_h(i-offset:i+offset, & - j-offset:j+offset,k), (/stencil_points/)) - x(stencil_points+1:2*stencil_points) = & - RESHAPE(CS%sh_xx(i-offset:i+offset, & - j-offset:j+offset,k), (/stencil_points/)) - x(2*stencil_points+1:3*stencil_points) = & - RESHAPE(vort_xy_h(i-offset:i+offset, & - j-offset:j+offset,k), (/stencil_points/)) - - input_norm = norm_h(i,j,k) - - x(:) = x(:) / (input_norm + CS%subroundoff_shear) - - call ANN_apply(x, y, CS%ann_Tall) + ! Normalize the input features using dimensional scaling + do n=1, 3*stencil_points + do m=1,nij + x(m,n) = x(m,n) * input_norm(m) + enddo + enddo + call cpu_clock_end(CS%id_clock_ANN_features) - y(:) = y(:) * input_norm * input_norm * CS%kappa_h(i,j) + call cpu_clock_begin(CS%id_clock_ANN_inference) + call ANN_apply_array_sio_r4(nij, x, y, CS%ann_Tall) + call cpu_clock_end(CS%id_clock_ANN_inference) - Txy(i,j) = y(1) - CS%Txx(i,j,k) = y(2) - CS%Tyy(i,j,k) = y(3) + call cpu_clock_begin(CS%id_clock_ANN_features) + m = 0 + do j=js,je ; do i=is,ie + m = m+1 + ! Denormalize the output features using dimensional scaling + CS%Txy_h(i,j,k) = y(m, 1) * output_norm(m) + CS%Txx(i,j,k) = y(m, 2) * output_norm(m) + CS%Tyy(i,j,k) = y(m, 3) * output_norm(m) enddo ; enddo - do J=Jsq-1,Jeq+1 ; do I=Isq-1,Ieq+1 - CS%Txy(I,J,k) = 0.25 * ( (Txy(i+1,j+1) + Txy(i,j)) & - + (Txy(i+1,j) + Txy(i,j+1))) * G%mask2dBu(I,J) - enddo; enddo - + call cpu_clock_end(CS%id_clock_ANN_features) enddo ! end of k loop - call pass_var(CS%Txy, G%Domain, clock=CS%id_clock_mpi, position=CORNER) - call pass_var(CS%Txx, G%Domain, clock=CS%id_clock_mpi) - call pass_var(CS%Tyy, G%Domain, clock=CS%id_clock_mpi) + call create_group_pass(pass_flux, CS%Txy_h, G%Domain, halo=2) + call create_group_pass(pass_flux, CS%Txx, G%Domain, halo=2) + call create_group_pass(pass_flux, CS%Tyy, G%Domain, halo=2) + call do_group_pass(pass_flux, G%Domain, clock=CS%id_clock_mpi) + + call cpu_clock_begin(CS%id_clock_ANN_features) + do k=1,nz + do J=js-2,Jeq+1 ; do I=is-2,Ieq+1 + CS%Txy(I,J,k) = 0.25 * ( (CS%Txy_h(i+1,j+1,k) + CS%Txy_h(i,j,k)) & + + (CS%Txy_h(i+1,j,k) + CS%Txy_h(i,j+1,k))) * G%mask2dBu(I,J) + enddo; enddo + enddo - call cpu_clock_end(CS%id_clock_stress_ANN) + deallocate(x) + deallocate(y) + deallocate(input_norm) + deallocate(output_norm) + call cpu_clock_end(CS%id_clock_ANN_features) end subroutine compute_stress_ANN_collocated @@ -1183,11 +1222,11 @@ subroutine compute_energy_source(u, v, h, fx, fy, G, GV, CS) intent(in) :: h !< Layer thicknesses [H ~> m or kg m-2]. real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)), & - intent(in) :: fx !< Zonal acceleration due to convergence of - !! along-coordinate stress tensor [L T-2 ~> m s-2] + intent(inout) :: fx !< Zonal acceleration due to convergence of + !! along-coordinate stress tensor [L T-2 ~> m s-2] real, dimension(SZI_(G),SZJB_(G),SZK_(GV)), & - intent(in) :: fy !< Meridional acceleration due to convergence - !! of along-coordinate stress tensor [L T-2 ~> m s-2] + intent(inout) :: fy !< Meridional acceleration due to convergence + !! of along-coordinate stress tensor [L T-2 ~> m s-2] real :: KE_term(SZI_(G),SZJ_(G),SZK_(GV)) ! A term in the kinetic energy budget ! [H L2 T-3 ~> m3 s-3 or W m-2] @@ -1201,14 +1240,14 @@ subroutine compute_energy_source(u, v, h, fx, fy, G, GV, CS) real :: vh ! Transport through meridional faces = v*h*dx, ! [H L2 T-1 ~> m3 s-1 or kg s-1]. - type(group_pass_type) :: pass_KE_uv ! A handle used for group halo passes - integer :: is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz integer :: i, j, k if (CS%id_KE_ZB2020 > 0) then call cpu_clock_begin(CS%id_clock_source) - call create_group_pass(pass_KE_uv, KE_u, KE_v, G%Domain, To_North+To_East) + if (.not. G%symmetric) then + call pass_vector(fx, fy, G%Domain, To_North+To_East, clock=CS%id_clock_mpi, halo=1) + endif is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = GV%ke Isq = G%IscB ; Ieq = G%IecB ; Jsq = G%JscB ; Jeq = G%JecB @@ -1218,17 +1257,17 @@ subroutine compute_energy_source(u, v, h, fx, fy, G, GV, CS) do k=1,nz KE_u(:,:) = 0. KE_v(:,:) = 0. - do j=js,je ; do I=Isq,Ieq + do j=js,je ; do I=is-1,Ieq uh = u(I,j,k) * 0.5 * (G%mask2dT(i,j)*h(i,j,k) + G%mask2dT(i+1,j)*h(i+1,j,k)) * & G%dyCu(I,j) KE_u(I,j) = uh * G%dxCu(I,j) * fx(I,j,k) enddo ; enddo - do J=Jsq,Jeq ; do i=is,ie + do J=js-1,Jeq ; do i=is,ie vh = v(i,J,k) * 0.5 * (G%mask2dT(i,j)*h(i,j,k) + G%mask2dT(i,j+1)*h(i,j+1,k)) * & G%dxCv(i,J) KE_v(i,J) = vh * G%dyCv(i,J) * fy(i,J,k) enddo ; enddo - call do_group_pass(pass_KE_uv, G%domain, clock=CS%id_clock_mpi) + do j=js,je ; do i=is,ie KE_term(i,j,k) = 0.5 * G%IareaT(i,j) & * ((KE_u(I,j) + KE_u(I-1,j)) + (KE_v(i,J) + KE_v(i,J-1)))