Skip to content

Apply array_sio function in ANN inference for momentum fluxes #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
adcroft merged 3 commits into from
Jul 19, 2025
Merged

Apply array_sio function in ANN inference for momentum fluxes #5

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
66f898c
Apply array_sio ANN inference for computation of momentum fluxes
Pperezhogin Jun 18, 2025
3017a50
remove trailing space
Pperezhogin Jun 18, 2025
8dc1fe9
Initial commit
Pperezhogin Jul 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 36 additions & 19 deletions src/parameterizations/lateral/MOM_Zanna_Bolton.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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/))

input_norm = norm_h(i,j,k)

x(:) = x(:) / (input_norm + CS%subroundoff_shear)
x(m,:) = x(m,:) / (norm_h(i,j,k) + CS%subroundoff_shear)
enddo; enddo

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
Expand All @@ -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
Expand Down
Loading