Skip to content

Commit

Permalink
fix: fms2_io domain_read for z axis (#1620)
Browse files Browse the repository at this point in the history
  • Loading branch information
uramirez8707 authored Dec 18, 2024
1 parent d79c3b5 commit 01a3a00
Show file tree
Hide file tree
Showing 2 changed files with 86 additions and 24 deletions.
41 changes: 17 additions & 24 deletions fms2_io/include/domain_read.inc
Original file line number Diff line number Diff line change
Expand Up @@ -334,7 +334,10 @@ subroutine domain_read_3d(fileobj, variable_name, vdata, unlim_dim_level, &
xpos, ypos, isd, isc, xc_size, jsd, jsc, yc_size, buffer_includes_halos, &
msg="file:"//trim(fileobj%path)//" and variable:"//trim(variable_name))
c(:) = 1
if (present(corner)) c = corner
e(:) = shape(vdata)
if (present(edge_lengths)) e = edge_lengths
call mpp_get_global_domain(io_domain, xbegin=xgbegin, xsize=xgsize, position=xpos)
call mpp_get_global_domain(io_domain, ybegin=ygbegin, ysize=ygsize, position=ypos)
Expand Down Expand Up @@ -503,6 +506,7 @@ subroutine domain_read_4d(fileobj, variable_name, vdata, unlim_dim_level, &
msg="file:"//trim(fileobj%path)//" and variable:"//trim(variable_name))
c(:) = 1
e(:) = shape(vdata)
if (present(edge_lengths)) e = edge_lengths
!I/O root reads in the data and scatters it.
if (fileobj%is_root) then
Expand All @@ -515,6 +519,7 @@ subroutine domain_read_4d(fileobj, variable_name, vdata, unlim_dim_level, &
call mpp_get_global_domain(io_domain, xbegin=xgmin, position=xpos)
call mpp_get_global_domain(io_domain, ybegin=ygmin, position=ypos)
do i = 1, size(fileobj%pelist)
if (present(corner)) c = corner
c(xdim_index) = pe_isc(i)
c(ydim_index) = pe_jsc(i)
if (fileobj%adjust_indices) then
Expand All @@ -532,13 +537,11 @@ subroutine domain_read_4d(fileobj, variable_name, vdata, unlim_dim_level, &
corner=c, edge_lengths=e, broadcast=.false.)
if (i .eq. 1) then
!Root rank stores data directly.
c = 1
if (buffer_includes_halos) then
!Adjust if the input buffer has room for halos.
c(xdim_index) = isc - isd + 1
c(ydim_index) = jsc - jsd + 1
else
c(xdim_index) = 1
c(ydim_index) = 1
endif
call put_array_section(buf_i4_kind, vdata, c, e)
else
Expand All @@ -555,13 +558,11 @@ subroutine domain_read_4d(fileobj, variable_name, vdata, unlim_dim_level, &
corner=c, edge_lengths=e, broadcast=.false.)
if (i .eq. 1) then
!Root rank stores data directly.
c = 1
if (buffer_includes_halos) then
!Adjust if the input buffer has room for halos.
c(xdim_index) = isc - isd + 1
c(ydim_index) = jsc - jsd + 1
else
c(xdim_index) = 1
c(ydim_index) = 1
endif
call put_array_section(buf_i8_kind, vdata, c, e)
else
Expand All @@ -578,13 +579,11 @@ subroutine domain_read_4d(fileobj, variable_name, vdata, unlim_dim_level, &
corner=c, edge_lengths=e, broadcast=.false.)
if (i .eq. 1) then
!Root rank stores data directly.
c = 1
if (buffer_includes_halos) then
!Adjust if the input buffer has room for halos.
c(xdim_index) = isc - isd + 1
c(ydim_index) = jsc - jsd + 1
else
c(xdim_index) = 1
c(ydim_index) = 1
endif
call put_array_section(buf_r4_kind, vdata, c, e)
else
Expand All @@ -601,13 +600,11 @@ subroutine domain_read_4d(fileobj, variable_name, vdata, unlim_dim_level, &
corner=c, edge_lengths=e, broadcast=.false.)
if (i .eq. 1) then
!Root rank stores data directly.
c = 1
if (buffer_includes_halos) then
!Adjust if the input buffer has room for halos.
c(xdim_index) = isc - isd + 1
c(ydim_index) = jsc - jsd + 1
else
c(xdim_index) = 1
c(ydim_index) = 1
endif
call put_array_section(buf_r8_kind, vdata, c, e)
else
Expand All @@ -626,6 +623,7 @@ subroutine domain_read_4d(fileobj, variable_name, vdata, unlim_dim_level, &
deallocate(pe_jsc)
deallocate(pe_jcsize)
else
c = 1
if (buffer_includes_halos) then
c(xdim_index) = isc - isd + 1
c(ydim_index) = jsc - jsd + 1
Expand Down Expand Up @@ -724,6 +722,7 @@ subroutine domain_read_5d(fileobj, variable_name, vdata, unlim_dim_level, &
msg="file:"//trim(fileobj%path)//" and variable:"//trim(variable_name))
c(:) = 1
e(:) = shape(vdata)
if (present(edge_lengths)) e = edge_lengths
!I/O root reads in the data and scatters it.
if (fileobj%is_root) then
Expand All @@ -737,6 +736,7 @@ subroutine domain_read_5d(fileobj, variable_name, vdata, unlim_dim_level, &
call mpp_get_global_domain(io_domain, ybegin=ygmin, position=ypos)
do i = 1, size(fileobj%pelist)
!Calculate the indices of the domain-decomposed chunk relative to its position in the file.
if (present(corner)) c = corner
c(xdim_index) = pe_isc(i)
c(ydim_index) = pe_jsc(i)
if (fileobj%adjust_indices) then
Expand All @@ -755,13 +755,11 @@ subroutine domain_read_5d(fileobj, variable_name, vdata, unlim_dim_level, &
if (i .eq. 1) then
!Root rank stores data directly. Re-adjust the indicies relative
!to the input buffer vdata.
c = 1
if (buffer_includes_halos) then
!Adjust if the input buffer has room for halos.
c(xdim_index) = isc - isd + 1
c(ydim_index) = jsc - jsd + 1
else
c(xdim_index) = 1
c(ydim_index) = 1
endif
call put_array_section(buf_i4_kind, vdata, c, e)
else
Expand All @@ -778,13 +776,11 @@ subroutine domain_read_5d(fileobj, variable_name, vdata, unlim_dim_level, &
corner=c, edge_lengths=e, broadcast=.false.)
if (i .eq. 1) then
!Root rank stores data directly.
c = 1
if (buffer_includes_halos) then
!Adjust if the input buffer has room for halos.
c(xdim_index) = isc - isd + 1
c(ydim_index) = jsc - jsd + 1
else
c(xdim_index) = 1
c(ydim_index) = 1
endif
call put_array_section(buf_i8_kind, vdata, c, e)
else
Expand All @@ -801,13 +797,11 @@ subroutine domain_read_5d(fileobj, variable_name, vdata, unlim_dim_level, &
corner=c, edge_lengths=e, broadcast=.false.)
if (i .eq. 1) then
!Root rank stores data directly.
c = 1
if (buffer_includes_halos) then
!Adjust if the input buffer has room for halos.
c(xdim_index) = isc - isd + 1
c(ydim_index) = jsc - jsd + 1
else
c(xdim_index) = 1
c(ydim_index) = 1
endif
call put_array_section(buf_r4_kind, vdata, c, e)
else
Expand All @@ -824,13 +818,11 @@ subroutine domain_read_5d(fileobj, variable_name, vdata, unlim_dim_level, &
corner=c, edge_lengths=e, broadcast=.false.)
if (i .eq. 1) then
!Root rank stores data directly.
c = 1
if (buffer_includes_halos) then
!Adjust if the input buffer has room for halos.
c(xdim_index) = isc - isd + 1
c(ydim_index) = jsc - jsd + 1
else
c(xdim_index) = 1
c(ydim_index) = 1
endif
call put_array_section(buf_r8_kind, vdata, c, e)
else
Expand All @@ -849,6 +841,7 @@ subroutine domain_read_5d(fileobj, variable_name, vdata, unlim_dim_level, &
deallocate(pe_jsc)
deallocate(pe_jcsize)
else
c = 1
if (buffer_includes_halos) then
c(xdim_index) = isc - isd + 1
c(ydim_index) = jsc - jsd + 1
Expand Down
69 changes: 69 additions & 0 deletions test_fms/fms2_io/test_domain_io.F90
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,9 @@ program test_domain_read
call read_data_wrapper(fileobj, "var3", 3, var_data_out, var_data_in)
call read_data_wrapper(fileobj, "var4", 4, var_data_out, var_data_in)
call read_data_wrapper(fileobj, "var5", 5, var_data_out, var_data_in)
call read_data_wrapper(fileobj, "var3", 6, var_data_out, var_data_in)
call read_data_wrapper(fileobj, "var4", 7, var_data_out, var_data_in)
call read_data_wrapper(fileobj, "var5", 8, var_data_out, var_data_in)

call close_file(fileobj)
endif
Expand Down Expand Up @@ -295,6 +298,72 @@ subroutine read_data_wrapper(fileob, var_name, dim, var_data, ref_data)

call read_data(fileob, trim(var_name)//"_i8", var_data%var_i8(:,:,:,:,:))
call compare_var_data(mpp_chksum(var_data%var_i8(:,:,:,:,:)), mpp_chksum(ref_data%var_i8(:,:,:,:,:)), "var5_i8")
case(6)
!Only read the second third dimension (3d case)
call var_data_init(var_data)
call read_data(fileob, trim(var_name)//"_r4", var_data%var_r4(:,:,1:1,1,1), &
corner=(/1, 1, 2/), edge_lengths=(/ nx, ny, 1/))
call compare_var_data(mpp_chksum(var_data%var_r4(:,:,1:1,1,1)), mpp_chksum(ref_data%var_r4(:,:,2:2,1,1)), &
"var3_r4-slice")

call read_data(fileob, trim(var_name)//"_r8", var_data%var_r8(:,:,1:1,1,1), &
corner=(/1, 1, 2/), edge_lengths=(/ nx, ny, 1/))
call compare_var_data(mpp_chksum(var_data%var_r8(:,:,1:1,1,1)), mpp_chksum(ref_data%var_r8(:,:,2:2,1,1)), &
"var3_r8-slice")

call read_data(fileob, trim(var_name)//"_i4", var_data%var_i4(:,:,1:1,1,1), &
corner=(/1, 1, 2/), edge_lengths=(/ nx, ny, 1/))
call compare_var_data(mpp_chksum(var_data%var_i4(:,:,1:1,1,1)), mpp_chksum(ref_data%var_i4(:,:,2:2,1,1)), &
"var3_i4-slice")

call read_data(fileob, trim(var_name)//"_i8", var_data%var_i8(:,:,1:1,1,1), &
corner=(/1, 1, 2/), edge_lengths=(/ nx, ny, 1/))
call compare_var_data(mpp_chksum(var_data%var_i8(:,:,1:1,1,1)), mpp_chksum(ref_data%var_i8(:,:,2:2,1,1)), &
"var3_i8-slice")
case(7)
!Only read the second third dimension (4d case)
call var_data_init(var_data)
call read_data(fileob, trim(var_name)//"_r4", var_data%var_r4(:,:,1:1,:,1), &
corner=(/1, 1, 2, 1/), edge_lengths=(/ nx, ny, 1, ndim4/))
call compare_var_data(mpp_chksum(var_data%var_r4(:,:,1:1,:,1)), mpp_chksum(ref_data%var_r4(:,:,2:2,:,1)), &
"var4_r4-slice")

call read_data(fileob, trim(var_name)//"_r8", var_data%var_r8(:,:,1:1,:,1), &
corner=(/1, 1, 2, 1/), edge_lengths=(/ nx, ny, 1, ndim4/))
call compare_var_data(mpp_chksum(var_data%var_r8(:,:,1:1,:,1)), mpp_chksum(ref_data%var_r8(:,:,2:2,:,1)), &
"var4_r8-slice")

call read_data(fileob, trim(var_name)//"_i4", var_data%var_i4(:,:,1:1,:,1), &
corner=(/1, 1, 2, 1/), edge_lengths=(/ nx, ny, 1, ndim4/))
call compare_var_data(mpp_chksum(var_data%var_i4(:,:,1:1,:,1)), mpp_chksum(ref_data%var_i4(:,:,2:2,:,1)), &
"var4_i4-slice")

call read_data(fileob, trim(var_name)//"_i8", var_data%var_i8(:,:,1:1,:,1), &
corner=(/1, 1, 2, 1/), edge_lengths=(/ nx, ny, 1, ndim4/))
call compare_var_data(mpp_chksum(var_data%var_i8(:,:,1:1,:,1)), mpp_chksum(ref_data%var_i8(:,:,2:2,:,1)), &
"var4_i8-slice")
case(8)
!Only read the second third dimension (5d case)
call var_data_init(var_data)
call read_data(fileob, trim(var_name)//"_r4", var_data%var_r4(:,:,1:1,:,:), &
corner=(/1, 1, 2, 1, 1/), edge_lengths=(/ nx, ny, 1, ndim4, ndim5/))
call compare_var_data(mpp_chksum(var_data%var_r4(:,:,1:1,:,:)), mpp_chksum(ref_data%var_r4(:,:,2:2,:,:)), &
"var5_r4-slice")

call read_data(fileob, trim(var_name)//"_r8", var_data%var_r8(:,:,1:1,:,:), &
corner=(/1, 1, 2, 1, 1/), edge_lengths=(/ nx, ny, 1, ndim4, ndim5/))
call compare_var_data(mpp_chksum(var_data%var_r8(:,:,1:1,:,:)), mpp_chksum(ref_data%var_r8(:,:,2:2,:,:)), &
"var5_r8-slice")

call read_data(fileob, trim(var_name)//"_i4", var_data%var_i4(:,:,1:1,:,:), &
corner=(/1, 1, 2, 1, 1/), edge_lengths=(/ nx, ny, 1, ndim4, ndim5/))
call compare_var_data(mpp_chksum(var_data%var_i4(:,:,1:1,:,:)), mpp_chksum(ref_data%var_i4(:,:,2:2,:,:)), &
"var5_i4-slice")

call read_data(fileob, trim(var_name)//"_i8", var_data%var_i8(:,:,1:1,:,:), &
corner=(/1, 1, 2, 1, 1/), edge_lengths=(/ nx, ny, 1, ndim4, ndim5/))
call compare_var_data(mpp_chksum(var_data%var_i8(:,:,1:1,:,:)), mpp_chksum(ref_data%var_i8(:,:,2:2,:,:)), &
"var5_i8-slice")
end select

end subroutine read_data_wrapper
Expand Down

0 comments on commit 01a3a00

Please sign in to comment.