From 66f898cdd5d93ebfcd4610b43e779df05cf309c9 Mon Sep 17 00:00:00 2001 From: Pavel Perezhogin Date: Wed, 18 Jun 2025 16:38:10 -0400 Subject: [PATCH 1/3] Apply array_sio ANN inference for computation of momentum fluxes --- .../lateral/MOM_Zanna_Bolton.F90 | 57 ++++++++++++------- 1 file changed, 37 insertions(+), 20 deletions(-) diff --git a/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 b/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 index ebbac944ac..55e5e4c36e 100644 --- a/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 +++ b/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 @@ -15,7 +15,7 @@ module MOM_Zanna_Bolton use MOM_domains, only : pass_var, 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, ANN_end, ANN_CS implicit none ; private @@ -664,13 +664,18 @@ 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 + ! number of horizontal grid points x ! (sh_xy, sh_xx, vort_xy) on a stencil [nondim] - real :: y(3) ! Vector of nondimensional - ! output features (Txy,Txx,Tyy) [nondim] + real, allocatable :: y(:,:) ! Vector of nondimensional + ! output features number of horizontal grid points x + ! (Txy,Txx,Tyy) [nondim] + real :: yy(3) ! Vector of dimensional + ! output features (Txy,Txx,Tyy) [L2 T-2 ~> m2 s-2] real :: input_norm ! Norm of input features [T-1 ~> s-1] real :: tmp ! Temporal value of squared norm [T-2 ~> s-2] integer :: offset ! Half the stencil size. Used for selection @@ -690,6 +695,11 @@ subroutine compute_stress_ANN_collocated(G, GV, CS) 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 + ! Number of horizontal grid points in ANN inference loop below + nij = (ie - is + 5) * (je - js + 5) + allocate(x(nij, 3 * CS%stencil_size**2)) + allocate(y(nij, 3)) + sh_xy_h = 0. vort_xy_h = 0. norm_h = 0. @@ -728,28 +738,32 @@ subroutine compute_stress_ANN_collocated(G, GV, CS) call pass_var(norm_h, G%Domain, clock=CS%id_clock_mpi) do k=1,nz + m = 0 do j=js-2,je+2 ; do i=is-2,ie+2 - x(1:stencil_points) = & - RESHAPE(sh_xy_h(i-offset:i+offset, & + m = m + 1 + x(m,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, & + x(m,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, & + x(m,2*stencil_points+1:3*stencil_points) = & + RESHAPE(vort_xy_h(i-offset:i+offset, & j-offset:j+offset,k), (/stencil_points/)) + + x(m,:) = x(m,:) / (norm_h(i,j,k) + CS%subroundoff_shear) + enddo; enddo - input_norm = norm_h(i,j,k) - - x(:) = x(:) / (input_norm + CS%subroundoff_shear) - - call ANN_apply(x, y, CS%ann_Tall) + call ANN_apply_array_sio(nij, x, y, CS%ann_Tall) - y(:) = y(:) * input_norm * input_norm * CS%kappa_h(i,j) + m = 0 + do j=js-2,je+2 ; do i=is-2,ie+2 + m = m+1 + yy(:) = y(m, :) * norm_h(i,j,k) * norm_h(i,j,k) * CS%kappa_h(i,j) - Txy(i,j) = y(1) - CS%Txx(i,j,k) = y(2) - CS%Tyy(i,j,k) = y(3) + Txy(i,j) = yy(1) + CS%Txx(i,j,k) = yy(2) + CS%Tyy(i,j,k) = yy(3) enddo ; enddo do J=Jsq-1,Jeq+1 ; do I=Isq-1,Ieq+1 @@ -763,6 +777,9 @@ subroutine compute_stress_ANN_collocated(G, GV, CS) call pass_var(CS%Txx, G%Domain, clock=CS%id_clock_mpi) call pass_var(CS%Tyy, G%Domain, clock=CS%id_clock_mpi) + deallocate(x) + deallocate(y) + call cpu_clock_end(CS%id_clock_stress_ANN) end subroutine compute_stress_ANN_collocated From 3017a50cbed6d7ee0819c609c939860b2c9f0f1a Mon Sep 17 00:00:00 2001 From: Pavel Perezhogin Date: Wed, 18 Jun 2025 16:55:00 -0400 Subject: [PATCH 2/3] remove trailing space --- src/parameterizations/lateral/MOM_Zanna_Bolton.F90 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 b/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 index 55e5e4c36e..1f77f17537 100644 --- a/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 +++ b/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 @@ -669,7 +669,7 @@ subroutine compute_stress_ANN_collocated(G, GV, CS) integer :: nij real, allocatable :: x(:,:) ! Vector of non-dimensional input features - ! number of horizontal grid points x + ! number of horizontal grid points x ! (sh_xy, sh_xx, vort_xy) on a stencil [nondim] real, allocatable :: y(:,:) ! Vector of nondimensional ! output features number of horizontal grid points x From 8dc1fe92a5996b849882e2b82da9777494ad60a8 Mon Sep 17 00:00:00 2001 From: Pavel Perezhogin Date: Fri, 18 Jul 2025 17:39:02 -0400 Subject: [PATCH 3/3] Initial commit --- src/parameterizations/lateral/MOM_Zanna_Bolton.F90 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 b/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 index 1f77f17537..792d89f837 100644 --- a/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 +++ b/src/parameterizations/lateral/MOM_Zanna_Bolton.F90 @@ -750,7 +750,7 @@ subroutine compute_stress_ANN_collocated(G, GV, CS) x(m,2*stencil_points+1:3*stencil_points) = & RESHAPE(vort_xy_h(i-offset:i+offset, & j-offset:j+offset,k), (/stencil_points/)) - + x(m,:) = x(m,:) / (norm_h(i,j,k) + CS%subroundoff_shear) enddo; enddo