Skip to content

Add diag manager test to register var with different axes permutations #1833

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

Draft
abrooks1085 wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Add diag manager test to register var with different axes permutations #1833

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
4 changes: 3 additions & 1 deletion test_fms/diag_manager/Makefile.am
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ check_PROGRAMS = test_diag_manager test_diag_manager_time \
check_time_pow check_time_rms check_subregional test_cell_measures test_var_masks \
check_var_masks test_multiple_send_data test_diag_out_yaml test_output_every_freq \
test_dm_weights test_prepend_date test_ens_runs test_diag_multi_file test_diag_attribute_add \
check_new_file_freq test_zbounds_limits test_multiple_zbounds
check_new_file_freq test_zbounds_limits test_multiple_zbounds test_generalized_indicies check_generalized_indices

# This is the source code for the test.
test_output_every_freq_SOURCES = test_output_every_freq.F90
Expand Down Expand Up @@ -71,6 +71,8 @@ test_diag_attribute_add_SOURCES = test_diag_attribute_add.F90
check_new_file_freq_SOURCES = check_new_file_freq.F90
test_zbounds_limits_SOURCES = test_zbounds_limits.F90
test_multiple_zbounds_SOURCES = test_multiple_zbounds.F90
test_generalized_indicies_SOURCES = testing_utils.F90 test_generalized_indicies.F90
check_generalized_indices_SOURCES = testing_utils.F90 check_generalized_indices.F90

TEST_EXTENSIONS = .sh
SH_LOG_DRIVER = env AM_TAP_AWK='$(AWK)' $(SHELL) \
Expand Down
133 changes: 133 additions & 0 deletions test_fms/diag_manager/check_generalized_indices.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,133 @@
!***********************************************************************
!* Apache License 2.0
!*
!* This file is part of the GFDL Flexible Modeling System (FMS).
!*
!* Licensed under the Apache License, Version 2.0 (the "License");
!* you may not use this file except in compliance with the License.
!* You may obtain a copy of the License at
!*
!* http://www.apache.org/licenses/LICENSE-2.0
!*
!* FMS is distributed in the hope that it will be useful, but WITHOUT
!* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied;
!* without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
!* PARTICULAR PURPOSE. See the License for the specific language
!* governing permissions and limitations under the License.
!***********************************************************************

!> @brief Checker for test_generalized_indicies output.
!! Verifies swapped-axis variables match identity variables under transpose:
!! var2_id(x,y) == var2_swap(y,x)
!! var3_id(x,y,z) == var3_swap(y,x,z)
program check_generalized_indices
use fms_mod, only: fms_init, fms_end, string
use fms2_io_mod, only: FmsNetcdfFile_t, read_data, open_file, close_file, get_global_attribute
use mpp_mod, only: mpp_error, FATAL, mpp_pe
use platform_mod, only: r4_kind

implicit none

type(FmsNetcdfFile_t) :: fileobj
integer :: nx, ny, nz
integer :: i

real(kind=r4_kind), allocatable :: var2_id(:,:) ! (x,y)
real(kind=r4_kind), allocatable :: var2_swap(:,:) ! (y,x)
real(kind=r4_kind), allocatable :: var3_id(:,:,:) ! (x,y,z)
real(kind=r4_kind), allocatable :: var3_swap(:,:,:) ! (y,x,z)

call fms_init()

nx = 96
ny = 96
nz = 5

if (.not. open_file(fileobj, "test_gen.nc", "read")) &
call mpp_error(FATAL, "unable to open test_gen.nc")

call check_global_attribute(fileobj, "test_generalized_indices")

allocate(var2_id(nx,ny), var2_swap(ny,nx))
allocate(var3_id(nx,ny,nz), var3_swap(ny,nx,nz))

! Output every 6 hours over 48 hours => 8 records
do i = 1, 8
var2_id = -999._r4_kind
var2_swap = -999._r4_kind
var3_id = -999._r4_kind
var3_swap = -999._r4_kind

print *, "Checking var2_swap vs var2_id - time_level:", string(i)
call read_data(fileobj, "var2_id", var2_id, unlim_dim_level=i)
call read_data(fileobj, "var2_swap", var2_swap, unlim_dim_level=i)
call check_var2_relation(var2_id, var2_swap)

print *, "Checking var3_swap vs var3_id - time_level:", string(i)
call read_data(fileobj, "var3_id", var3_id, unlim_dim_level=i)
call read_data(fileobj, "var3_swap", var3_swap, unlim_dim_level=i)
call check_var3_relation(var3_id, var3_swap)
enddo

call close_file(fileobj)
call fms_end()

contains

subroutine check_global_attribute(fileobj, expected_title)
type(FmsNetcdfFile_t), intent(in) :: fileobj
character(len=*), intent(in) :: expected_title

character(len=100) :: attribute_value

call get_global_attribute(fileobj, "title", attribute_value)
if (trim(attribute_value) .ne. trim(expected_title)) then
call mpp_error(FATAL, "Global attribute 'title' not expected value.")
endif
end subroutine check_global_attribute

subroutine check_var2_relation(v_id, v_sw)
real(kind=r4_kind), intent(in) :: v_id(:,:) ! (x,y)
real(kind=r4_kind), intent(in) :: v_sw(:,:) ! (y,x)

integer :: x, y

if (size(v_id,1) /= size(v_sw,2) .or. size(v_id,2) /= size(v_sw,1)) then
call mpp_error(FATAL, "check_var2_relation: dimension mismatch between var2_id and var2_swap")
endif

do x = 1, size(v_id,1)
do y = 1, size(v_id,2)
if (abs(v_id(x,y) - v_sw(y,x)) > 0) then
print *, mpp_pe(), "var2 mismatch at (x,y)=", x, y, " id=", v_id(x,y), " swap(y,x)=", v_sw(y,x)
call mpp_error(FATAL, "check_var2_relation: var2_swap != transpose(var2_id)")
endif
enddo
enddo
end subroutine check_var2_relation

subroutine check_var3_relation(v_id, v_sw)
real(kind=r4_kind), intent(in) :: v_id(:,:,:) ! (x,y,z)
real(kind=r4_kind), intent(in) :: v_sw(:,:,:) ! (y,x,z)

integer :: x, y, z

if (size(v_id,1) /= size(v_sw,2) .or. size(v_id,2) /= size(v_sw,1) .or. size(v_id,3) /= size(v_sw,3)) then
call mpp_error(FATAL, "check_var3_relation: dimension mismatch between var3_id and var3_swap")
endif

do x = 1, size(v_id,1)
do y = 1, size(v_id,2)
do z = 1, size(v_id,3)
if (abs(v_id(x,y,z) - v_sw(y,x,z)) > 0) then
print *, mpp_pe(), "var3 mismatch at (x,y,z)=", x, y, z, &
" id=", v_id(x,y,z), " swap(y,x,z)=", v_sw(y,x,z)
call mpp_error(FATAL, "check_var3_relation: var3_swap != var3_id with x/y swapped")
endif
enddo
enddo
enddo
end subroutine check_var3_relation

end program check_generalized_indices

228 changes: 228 additions & 0 deletions test_fms/diag_manager/test_generalized_indicies.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,228 @@
!***********************************************************************
!* Apache License 2.0
!*
!* This file is part of the GFDL Flexible Modeling System (FMS).
!*
!* Licensed under the Apache License, Version 2.0 (the "License");
!* you may not use this file except in compliance with the License.
!* You may obtain a copy of the License at
!*
!* http://www.apache.org/licenses/LICENSE-2.0
!*
!* FMS is distributed in the hope that it will be useful, but WITHOUT
!* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied;
!* without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
!* PARTICULAR PURPOSE. See the License for the specific language
!* governing permissions and limitations under the License.
!***********************************************************************

!> @brief Test generalized axis permutations (x/y for now) for send_data
!! Assumes default configuration parameters: test_normal + no_mask.
program test_generalized_indices
use fms_mod, only: fms_init, fms_end
use testing_utils, only: allocate_buffer
use platform_mod, only: r8_kind
use mpp_mod, only: FATAL, mpp_error, mpp_npes, mpp_pe, mpp_root_pe
use time_manager_mod, only: time_type, set_calendar_type, set_date, JULIAN, set_time, OPERATOR(+)
use diag_manager_mod, only: diag_manager_init, diag_manager_end, diag_axis_init, register_diag_field, &
diag_send_complete, diag_manager_set_time_end, send_data
use mpp_domains_mod, only: domain2d, mpp_define_domains, mpp_define_io_domain, mpp_get_compute_domain

implicit none

integer :: nx, ny, nz, nw
integer :: layout(2), io_layout(2)
type(domain2d) :: Domain
integer :: isc, iec, jsc, jec
integer :: nhalox, nhaloy
integer :: ntimes, i
type(time_type) :: Time, Time_step
real(r8_kind) :: missing_value

! Axes
integer :: id_x, id_y, id_z, id_w
integer :: axis(4)

! Data
real(r8_kind), allocatable :: cdata(:,:,:,:) ! canonical storage: (x,y,z,w)

! Permutation test
integer :: p_id(3), p_swap(3)
integer :: id_var2_id, id_var2_swap
integer :: id_var3_id, id_var3_swap

call fms_init
call set_calendar_type(JULIAN)
call diag_manager_init

Time = set_date(2,1,1,0,0,0)
Time_step = set_time(3600,0) ! 1 hour

nx = 96
ny = 96
nz = 5
nw = 2
ntimes = 48

nhalox = 2
nhaloy = 2
layout = (/1, mpp_npes()/)
io_layout = (/1, 1/)

! Domain
call mpp_define_domains((/1,nx,1,ny/), layout, Domain, name='2D domain', symmetry=.true., &
xhalo=nhalox, yhalo=nhaloy)
call mpp_define_io_domain(Domain, io_layout)
call mpp_get_compute_domain(Domain, isc, iec, jsc, jec)

! Data allocation + init (canonical x,y storage)
cdata = allocate_buffer(isc, iec, jsc, jec, nz, nw)
call init_buffer(cdata, isc, iec, jsc, jec, 0)

! Axes
id_x = diag_axis_init('x', real((/(i, i=1,nx)/), kind=r8_kind), 'point_E', 'x', long_name='point_E', Domain2=Domain)
id_y = diag_axis_init('y', real((/(i, i=1,ny)/), kind=r8_kind), 'point_N', 'y', long_name='point_N', Domain2=Domain)
id_z = diag_axis_init('z', real((/(i, i=1,nz)/), kind=r8_kind), 'point_Z', 'z', long_name='point_Z')
id_w = diag_axis_init('w', real((/(i, i=1,nw)/), kind=r8_kind), 'point_W', 'n', long_name='point_W')
axis = [id_x, id_y, id_z, id_w]

missing_value = -666._r8_kind

! Define permutations: identity (x,y,z) and swap (y,x,z)
p_id = [1,2,3]
p_swap = [2,1,3]

! Register permuted diagnostic fields ONCE
id_var2_id = register_diag_field('ocn_mod', 'var2_id', (/axis(p_id(1)), axis(p_id(2))/), Time, 'Var2d id', &
'mullions', missing_value=missing_value)
id_var2_swap = register_diag_field('ocn_mod', 'var2_swap', (/axis(p_swap(1)), axis(p_swap(2))/), Time, 'Var2d swap', &
'mullions', missing_value=missing_value)

id_var3_id = register_diag_field('ocn_mod', 'var3_id', (/axis(p_id(1)), axis(p_id(2)), axis(p_id(3))/), Time, &
'Var3d id', 'mullions', missing_value=missing_value)
id_var3_swap = register_diag_field('ocn_mod', 'var3_swap', (/axis(p_swap(1)), axis(p_swap(2)), axis(p_swap(3))/), Time, &
'Var3d swap', 'mullions', missing_value=missing_value)

if (mpp_pe() == mpp_root_pe()) then
print *, "Testing generalized indices in default mode (test_normal + no_mask)"
print *, " canonical storage is (x,y,z,w)"
print *, " sending:"
print *, " var2_id with axes (x,y)"
print *, " var2_swap with axes (y,x)"
print *, " var3_id with axes (x,y,z)"
print *, " var3_swap with axes (y,x,z)"
end if

call diag_manager_set_time_end(set_date(2,1,3,0,0,0))

do i = 1, ntimes
Time = Time + Time_step
call set_buffer(cdata, i)

! Identity: axes (x,y) / (x,y,z) with canonical storage
call send_var2_perm(id_var2_id, cdata, p_id, Time)
call send_var3_perm(id_var3_id, cdata, p_id, Time)

! Swap: axes (y,x) / (y,x,z) while canonical storage remains (x,y,...) -> pack to temp and send
call send_var2_perm(id_var2_swap, cdata, p_swap, Time)
call send_var3_perm(id_var3_swap, cdata, p_swap, Time)

call diag_send_complete(Time_step)
call diag_send_complete(Time_step)
end do

call diag_manager_end(Time)
call fms_end

contains

subroutine send_var2_perm(id_field, buf, p, Time_in)
integer, intent(in) :: id_field
real(r8_kind), intent(in) :: buf(:,:,:,:) ! canonical (x,y,z,w)
integer, intent(in) :: p(3)
type(time_type), intent(in) :: Time_in

logical :: used_local
real(r8_kind), allocatable :: tmp2(:,:)

! Support only identity (1,2,*) and xy-swap (2,1,*) for 2D
if (p(1)==1 .and. p(2)==2) then
used_local = send_data(id_field, buf(:,:,1,1), Time_in)
else if (p(1)==2 .and. p(2)==1) then
allocate(tmp2(size(buf,2), size(buf,1)))
tmp2 = transpose(buf(:,:,1,1))
used_local = send_data(id_field, tmp2, Time_in)
deallocate(tmp2)
else
call mpp_error(FATAL, 'send_var2_perm: only p=(1,2,*) or (2,1,*) implemented')
end if
end subroutine send_var2_perm


subroutine send_var3_perm(id_field, buf, p, Time_in)
integer, intent(in) :: id_field
real(r8_kind), intent(in) :: buf(:,:,:,:) ! canonical (x,y,z,w)
integer, intent(in) :: p(3)
type(time_type), intent(in) :: Time_in

logical :: used_local
integer :: nxloc, nyloc, nzloc, k
real(r8_kind), allocatable :: tmp3(:,:,:)

! For now, support only keeping z as z
if (p(3) /= 3) call mpp_error(FATAL, 'send_var3_perm: only permutations with p(3)=3 implemented')

if (p(1)==1 .and. p(2)==2) then
used_local = send_data(id_field, buf(:,:,:,1), Time_in)

else if (p(1)==2 .and. p(2)==1) then
nxloc = size(buf,1)
nyloc = size(buf,2)
nzloc = size(buf,3)

allocate(tmp3(nyloc, nxloc, nzloc))
do k = 1, nzloc
tmp3(:,:,k) = transpose(buf(:,:,k,1))
end do

used_local = send_data(id_field, tmp3, Time_in)
deallocate(tmp3)

else
call mpp_error(FATAL, 'send_var3_perm: only p=(1,2,3) or (2,1,3) implemented')
end if
end subroutine send_var3_perm


!> @brief initialized the buffer based on the starting/ending indices
subroutine init_buffer(buffer, is, ie, js, je, nhalo)
real(r8_kind), intent(inout) :: buffer(:,:,:,:)
integer, intent(in) :: is, ie, js, je
integer, intent(in) :: nhalo

integer :: ii, j, k, l

do ii = is, ie
do j = js, je
do k = 1, size(buffer, 3)
do l = 1, size(buffer, 4)
buffer(ii-is+1+nhalo, j-js+1+nhalo, k, l) = real(ii, kind=r8_kind)*1000._r8_kind + &
real(j, kind=r8_kind)* 10._r8_kind + &
real(k, kind=r8_kind)
end do
end do
end do
end do
end subroutine init_buffer


!> @brief Set the buffer based on the time_index
subroutine set_buffer(buffer, time_index)
real(r8_kind), intent(inout) :: buffer(:,:,:,:)
integer, intent(in) :: time_index

buffer = nint(buffer) + real(time_index, kind=r8_kind)/100._r8_kind
end subroutine set_buffer

end program test_generalized_indices

Loading
Loading