From 5c5a238d353c79f74c4269b023dc3ec81eb44da9 Mon Sep 17 00:00:00 2001 From: Andrew Shao Date: Wed, 29 Jan 2025 15:43:43 -0800 Subject: [PATCH] Refactor buffers based on Marshall's suggestions --- src/framework/MOM_diag_buffers.F90 | 198 +++++++++++++++++----------- src/framework/MOM_diag_mediator.F90 | 28 ++-- 2 files changed, 132 insertions(+), 94 deletions(-) diff --git a/src/framework/MOM_diag_buffers.F90 b/src/framework/MOM_diag_buffers.F90 index d5cde56e15..20d6419a12 100644 --- a/src/framework/MOM_diag_buffers.F90 +++ b/src/framework/MOM_diag_buffers.F90 @@ -2,7 +2,7 @@ !! diagnostics which need to store intermediate or partial states of state variables module MOM_diag_buffers -use iso_fortran_env, only : stdout => output_unit, stderr => error_unit +use MOM_io, only : stdout, stderr ! This file is part of MOM6. See LICENSE.md for the license. @@ -10,13 +10,26 @@ module MOM_diag_buffers public :: diag_buffer_unit_tests_2d, diag_buffer_unit_tests_3d +type, abstract :: buffer_base +end type buffer_base + +!> Holds a 2d field +type, extends(buffer_base) :: buffer_2d + real, dimension(:,:), allocatable :: field +end type buffer_2d + +!> Holds a 3d field +type, extends(buffer_base) :: buffer_3d + real, dimension(:,:,:), allocatable :: field +end type buffer_3d + !> The base class for the diagnostic buffers in this module type, abstract :: diag_buffer_base ; private integer :: is !< The start slot of the array i-direction integer :: js !< The start slot of the array j-direction integer :: ie !< The end slot of the array i-direction integer :: je !< The end slot of the array j-direction - real :: fill_value = 0. !< Set the fill value to use when growing the buffer + real :: fill_value = 0. !< Set the fill value to use when growing the buffer [arbitrary] integer, allocatable, dimension(:) :: ids !< List of diagnostic ids whose slot corresponds to the row in the buffer integer :: length = 0 !< The number of slots in the buffer @@ -33,8 +46,8 @@ module MOM_diag_buffers end type diag_buffer_base !> Dynamically growing buffer for 2D arrays. -type, extends(diag_buffer_base), public :: diag_buffer_2d - real, public, allocatable, dimension(:,:,:) :: buffer !< The actual buffer to store data [arbitrary] +type, extends(diag_buffer_base), public :: diag_buffer_2d; private + type(buffer_2d), public, dimension(:), allocatable :: buffer !< The actual 2D buffer which will dynamically grow contains @@ -44,9 +57,9 @@ module MOM_diag_buffers !> Dynamically growing buffer for 3D arrays. type, extends(diag_buffer_base), public :: diag_buffer_3d ; private + type(buffer_3d), public, dimension(:), allocatable :: buffer !< The actual 2D buffer which will dynamically grow integer :: ks !< The start slot in the k-dimension integer :: ke !< The last slot in the k-dimension - real, public, allocatable, dimension(:,:,:,:) :: buffer !< The actual buffer to store data [arbitrary] contains @@ -152,21 +165,29 @@ subroutine set_vertical_extent(this, ks, ke) this%ks = ks; this%ke = ke end subroutine set_vertical_extent -!> Grow a buffer for 2D arrays +!> Grow a 2d diagnostic buffer subroutine grow_2d(this) - class(diag_buffer_2d), intent(inout) :: this !< This 2d buffer + class(diag_buffer_2d), intent(inout) :: this - integer :: n - real, allocatable, dimension(:,:,:) :: temp ! Temporary array to hold the contents of the buffer - ! prior to allocating new memory [arbitrary] + integer :: i, n + integer :: is, ie, js, je + type(buffer_2d), dimension(:), allocatable :: new_buffer + ! Grow the ID array call this%grow_ids() + is = this%is; ie=this%ie; js=this%js; je=this%je n = this%length - allocate(temp(n+1, this%is:this%ie, this%js:this%je), source=this%fill_value) - if (n>0) temp(1:n,:,:) = this%buffer(:,:,:) - call move_alloc(temp, this%buffer) - this%length = this%length + 1 + + allocate(new_buffer(n+1)) + do i=1,n + allocate(new_buffer(i)%field(is:ie,js:je)) + new_buffer(i)%field(:,:) = this%buffer(i)%field(:,:) + enddo + allocate(new_buffer(n+1)%field(is:ie,js:je), source=this%fill_value) + call move_alloc(new_buffer, this%buffer) + this%length = n+1 + end subroutine grow_2d !> Store a 2D array into this buffer @@ -178,24 +199,32 @@ subroutine store_2d(this, data, id) integer :: slot slot = this%check_capacity_by_id(id) - this%buffer(slot,:,:) = data(:,:) + this%buffer(slot)%field(:,:) = data(:,:) end subroutine store_2d -!> Grow a buffer for 3d arrays +!> Grow a 2d diagnostic buffer subroutine grow_3d(this) - class(diag_buffer_3d), intent(inout) :: this !< This 3d buffer + class(diag_buffer_3d), intent(inout) :: this - integer :: n - real, allocatable, dimension(:,:,:,:) :: temp ! Temporary array to hold the contents of the buffer - ! prior to allocating new memory [arbitrary] + integer :: i, n + integer :: is, ie, js, je, ks, ke + type(buffer_3d), dimension(:), allocatable :: new_buffer + ! Grow the ID array call this%grow_ids() + is = this%is; ie=this%ie; js=this%js; je=this%je; ks=this%ks; ke=this%ke n = this%length - allocate(temp(n+1, this%is:this%ie, this%js:this%je, this%ks:this%ke), source=this%fill_value) - if (n>0) temp(1:n,:,:,:) = this%buffer(:,:,:,:) - call move_alloc(temp, this%buffer) - this%length = this%length + 1 + + allocate(new_buffer(n+1)) + do i=1,n + allocate(new_buffer(i)%field(is:ie,js:je,ks:ke)) + new_buffer(i)%field(:,:,:) = this%buffer(i)%field(:,:,:) + enddo + allocate(new_buffer(n+1)%field(is:ie,js:je,ks:ke), source=this%fill_value) + call move_alloc(new_buffer, this%buffer) + this%length = n+1 + end subroutine grow_3d !> Store a 3d array into this buffer @@ -208,7 +237,7 @@ subroutine store_3d(this, data, id) ! Find the first slot in the ids array that is 0, i.e. this is a portion of the buffer that can be reused slot = this%check_capacity_by_id(id) - this%buffer(slot,:,:,:) = data(:,:,:) + this%buffer(slot)%field(:,:,:) = data(:,:,:) end subroutine store_3d !> Unit tests for the 2d version of the diag buffer @@ -227,55 +256,57 @@ function diag_buffer_unit_tests_2d(verbose) result(fail) !> Ensure properties of a newly initialized buffer function new_buffer_2d() result(local_fail) - type(diag_buffer_2d) :: buffer_2d + type(diag_buffer_2d) :: buffer logical :: local_fail !< True if any of the unit tests fail local_fail = .false. - local_fail = local_fail .or. allocated(buffer_2d%buffer) - local_fail = local_fail .or. allocated(buffer_2d%ids) - local_fail = local_fail .or. buffer_2d%length /= 0 + local_fail = local_fail .or. allocated(buffer%buffer) + if (verbose) write(stdout,*) "new_buffer_2d: ", local_fail + local_fail = local_fail .or. allocated(buffer%ids) + if (verbose) write(stdout,*) "new_buffer_2d: ", local_fail + local_fail = local_fail .or. buffer%length /= 0 if (verbose) write(stdout,*) "new_buffer_2d: ", local_fail end function new_buffer_2d !> Test the growing of a buffer function grow_buffer_2d() result(local_fail) - type(diag_buffer_2d) :: buffer_2d + type(diag_buffer_2d) :: buffer logical :: local_fail !< True if any of the unit tests fail integer, parameter :: is=1, ie=2, js=3, je=6 integer :: i local_fail = .false. - buffer_2d = diag_buffer_2d(is=is, ie=ie, js=js, je=je) + call buffer%set_horizontal_extents(is=is, ie=ie, js=js, je=je) ! Grow the buffer 3 times do i=1,3 - call buffer_2d%grow() - local_fail = local_fail .or. (buffer_2d%length /= i) - local_fail = local_fail .or. (size(buffer_2d%buffer, 1) /= i) - local_fail = local_fail .or. (lbound(buffer_2d%buffer, 2) /= is) - local_fail = local_fail .or. (ubound(buffer_2d%buffer, 2) /= ie) - local_fail = local_fail .or. (lbound(buffer_2d%buffer, 3) /= js) - local_fail = local_fail .or. (ubound(buffer_2d%buffer, 3) /= je) + call buffer%grow() + local_fail = local_fail .or. (buffer%length /= i) + local_fail = local_fail .or. (lbound(buffer%buffer(i)%field, 1) /= is) + local_fail = local_fail .or. (ubound(buffer%buffer(i)%field, 1) /= ie) + local_fail = local_fail .or. (lbound(buffer%buffer(i)%field, 2) /= js) + local_fail = local_fail .or. (ubound(buffer%buffer(i)%field, 2) /= je) enddo if (verbose) write(stdout,*) "grow_buffer_2d: ", local_fail end function grow_buffer_2d !> Test storing a buffer based on a unique id function store_buffer_2d() result(local_fail) - type(diag_buffer_2d) :: buffer_2d + type(diag_buffer_2d) :: buffer logical :: local_fail !< True if any of the unit tests fail integer, parameter :: is=1, ie=2, js=3, je=6, nlen=3 - integer :: i + integer :: i, slot real, allocatable, dimension(:,:,:) :: test_2d allocate(test_2d(nlen, is:ie, js:je)) call random_number(test_2d) - buffer_2d = diag_buffer_2d(is=is, ie=ie, js=js, je=je) + buffer = diag_buffer_2d(is=is, ie=ie, js=js, je=je) do i=1,nlen - call buffer_2d%store(test_2d(i,:,:), i*3) + call buffer%store(test_2d(i,:,:), i*3) + slot = buffer%find_buffer_slot(i*3) + local_fail = local_fail .or. ANY(buffer%buffer(slot)%field(:,:) /= test_2d(i,:,:)) enddo - local_fail = ANY(buffer_2d%buffer /= test_2d) if (verbose) write(stdout,*) "store_buffer_2d: ", local_fail end function store_buffer_2d @@ -284,7 +315,7 @@ end function store_buffer_2d !! loop through again, but use the slots of the buffer in the following !! order: 2, 1, 3 function reuse_buffer_2d() result(local_fail) - type(diag_buffer_2d) :: buffer_2d + type(diag_buffer_2d) :: buffer logical :: local_fail !< True if any of the unit tests fail integer, parameter :: is=1, ie=2, js=3, je=6, nlen=3 @@ -296,10 +327,10 @@ function reuse_buffer_2d() result(local_fail) call random_number(test_2d_first) call random_number(test_2d_second) - buffer_2d = diag_buffer_2d(is=is, ie=ie, js=js, je=je) + call buffer%set_horizontal_extents(is=is, ie=ie, js=js, je=je) do i=1,nlen - call buffer_2d%store(test_2d_first(i,:,:), id=i*3) + call buffer%store(test_2d_first(i,:,:), id=i*3) enddo do i=1,nlen @@ -307,13 +338,15 @@ function reuse_buffer_2d() result(local_fail) ! id and new_id are multiplied by primes to make sure they are unique id = reorder(i)*3 new_id = i*7 - call buffer_2d%mark_available(id=reorder(i)*3) - call buffer_2d%store(test_2d_second(i,:,:), id=new_id) - local_fail = local_fail .or. buffer_2d%find_buffer_slot(new_id) /= new_i + call buffer%mark_available(id=reorder(i)*3) + call buffer%store(test_2d_second(i,:,:), id=new_id) + local_fail = local_fail .or. buffer%find_buffer_slot(new_id) /= new_i test_2d_first(new_i,:,:) = test_2d_second(i,:,:) enddo - local_fail = local_fail .or. any(buffer_2d%ids /= [14, 7, 21]) - local_fail = local_fail .or. any(buffer_2d%buffer /= test_2d_first) + local_fail = local_fail .or. any(buffer%ids /= [14, 7, 21]) + do i=1,nlen + local_fail = local_fail .or. any(buffer%buffer(i)%field(:,:) /= test_2d_first(i,:,:)) + enddo if (verbose) write(stdout,*) "reuse_buffer_2d: ", local_fail end function reuse_buffer_2d @@ -335,56 +368,59 @@ function diag_buffer_unit_tests_3d(verbose) result(fail) !> Ensure properties of a newly initialized buffer function new_buffer_3d() result(local_fail) - type(diag_buffer_3d) :: buffer_3d + type(diag_buffer_3d) :: buffer logical :: local_fail !< True if any of the unit tests fail local_fail = .false. - local_fail = local_fail .or. allocated(buffer_3d%buffer) - local_fail = local_fail .or. allocated(buffer_3d%ids) - local_fail = local_fail .or. buffer_3d%length /= 0 + local_fail = local_fail .or. allocated(buffer%buffer) + local_fail = local_fail .or. allocated(buffer%ids) + local_fail = local_fail .or. buffer%length /= 0 if (verbose) write(stdout,*) "new_buffer_3d: ", local_fail end function new_buffer_3d !> Test the growing of a buffer function grow_buffer_3d() result(local_fail) - type(diag_buffer_3d) :: buffer_3d + type(diag_buffer_3d) :: buffer logical :: local_fail !< True if any of the unit tests fail integer, parameter :: is=1, ie=2, js=3, je=6, ks=1, ke=10 integer :: i local_fail = .false. - buffer_3d = diag_buffer_3d(is=is, ie=ie, js=js, je=je, ks=ks, ke=ke) + call buffer%set_horizontal_extents(is=is, ie=ie, js=js, je=je) + call buffer%set_vertical_extent(ks=ks, ke=ke) ! Grow the buffer 3 times do i=1,3 - call buffer_3d%grow() - local_fail = local_fail .or. (buffer_3d%length /= i) - local_fail = local_fail .or. (size(buffer_3d%buffer, 1) /= i) - local_fail = local_fail .or. (lbound(buffer_3d%buffer, 2) /= is) - local_fail = local_fail .or. (ubound(buffer_3d%buffer, 2) /= ie) - local_fail = local_fail .or. (lbound(buffer_3d%buffer, 3) /= js) - local_fail = local_fail .or. (ubound(buffer_3d%buffer, 3) /= je) - local_fail = local_fail .or. (lbound(buffer_3d%buffer, 4) /= ks) - local_fail = local_fail .or. (ubound(buffer_3d%buffer, 4) /= ke) + call buffer%grow() + local_fail = local_fail .or. (buffer%length /= i) + local_fail = local_fail .or. (lbound(buffer%buffer(i)%field, 1) /= is) + local_fail = local_fail .or. (ubound(buffer%buffer(i)%field, 1) /= ie) + local_fail = local_fail .or. (lbound(buffer%buffer(i)%field, 2) /= js) + local_fail = local_fail .or. (ubound(buffer%buffer(i)%field, 2) /= je) + local_fail = local_fail .or. (lbound(buffer%buffer(i)%field, 3) /= ks) + local_fail = local_fail .or. (ubound(buffer%buffer(i)%field, 3) /= ke) + if (verbose) write(stdout,*) "grow_buffer_3d: ", local_fail enddo if (verbose) write(stdout,*) "grow_buffer_3d: ", local_fail end function grow_buffer_3d !> Test storing a buffer based on a unique id function store_buffer_3d() result(local_fail) - type(diag_buffer_3d) :: buffer_3d + type(diag_buffer_3d) :: buffer logical :: local_fail !< True if any of the unit tests fail integer, parameter :: is=1, ie=2, js=3, je=6, ks=1, ke=10, nlen=3 - integer :: i - real, dimension(nlen, is:ie, js:je, ks:ke) :: test_3d + integer :: i, slot + real, dimension(nlen,is:ie,js:je,ks:ke) :: test_3d + local_fail = .false. call random_number(test_3d) - buffer_3d = diag_buffer_3d(is=is, ie=ie, js=js, je=je, ks=1, ke=10) + buffer = diag_buffer_3d(is=is, ie=ie, js=js, je=je, ks=1, ke=10) do i=1,nlen - call buffer_3d%store(test_3d(i,:,:,:), i*3) + call buffer%store(test_3d(i,:,:,:), i*3) + slot = buffer%find_buffer_slot(i*3) + local_fail = local_fail .or. ANY(buffer%buffer(slot)%field(:,:,:) /= test_3d(i,:,:,:)) enddo - local_fail = ANY(buffer_3d%buffer /= test_3d) if (verbose) write(stdout,*) "store_buffer_3d: ", local_fail end function store_buffer_3d @@ -393,7 +429,7 @@ end function store_buffer_3d !! loop through again, but use the slots of the buffer in the following !! order: 2, 1, 3 function reuse_buffer_3d() result(local_fail) - type(diag_buffer_3d) :: buffer_3d + type(diag_buffer_3d) :: buffer logical :: local_fail !< True if any of the unit tests fail integer, parameter :: is=1, ie=2, js=3, je=6, ks=1, ke=10, nlen=3 @@ -405,10 +441,10 @@ function reuse_buffer_3d() result(local_fail) call random_number(test_3d_first) call random_number(test_3d_second) - buffer_3d = diag_buffer_3d(is=is, ie=ie, js=js, je=je, ks=ks, ke=ke) + buffer = diag_buffer_3d(is=is, ie=ie, js=js, je=je, ks=ks, ke=ke) do i=1,nlen - call buffer_3d%store(test_3d_first(i,:,:,:), id=i*3) + call buffer%store(test_3d_first(i,:,:,:), id=i*3) enddo do i=1,nlen @@ -416,13 +452,15 @@ function reuse_buffer_3d() result(local_fail) ! id and new_id are multiplied by primes to make sure they are unique id = reorder(i)*3 new_id = i*7 - call buffer_3d%mark_available(id=reorder(i)*3) - call buffer_3d%store(test_3d_second(i,:,:,:), id=new_id) - local_fail = local_fail .or. buffer_3d%find_buffer_slot(new_id) /= new_i + call buffer%mark_available(id=reorder(i)*3) + call buffer%store(test_3d_second(i,:,:,:), id=new_id) + local_fail = local_fail .or. buffer%find_buffer_slot(new_id) /= new_i test_3d_first(new_i,:,:,:) = test_3d_second(i,:,:,:) enddo - local_fail = local_fail .or. any(buffer_3d%ids /= [14, 7, 21]) - local_fail = local_fail .or. any(buffer_3d%buffer /= test_3d_first) + local_fail = local_fail .or. any(buffer%ids /= [14, 7, 21]) + do i=1,nlen + local_fail = local_fail .or. any(buffer%buffer(i)%field(:,:,:) /= test_3d_first(i,:,:,:)) + enddo if (verbose) write(stdout,*) "reuse_buffer_3d: ", local_fail end function reuse_buffer_3d diff --git a/src/framework/MOM_diag_mediator.F90 b/src/framework/MOM_diag_mediator.F90 index 690519a7c1..1726805ca9 100644 --- a/src/framework/MOM_diag_mediator.F90 +++ b/src/framework/MOM_diag_mediator.F90 @@ -137,8 +137,8 @@ module MOM_diag_mediator type(diag_dsamp), dimension(2:MAX_DSAMP_LEV) :: dsamp !< Downsample container ! For diagnostics posted piecemeal - type(diag_buffer_2d) :: buffer_2d !< A dynamically reallocated buffer for 2d piecemeal diagnostics - type(diag_buffer_3d) :: buffer_3d !< A dynamically reallocated buffer for 3d piecemeal diagnostics + type(diag_buffer_2d) :: piecemeal_2d !< A dynamically reallocated buffer for 2d piecemeal diagnostics + type(diag_buffer_3d) :: piecemeal_3d !< A dynamically reallocated buffer for 3d piecemeal diagnostics end type axes_grp !> Contains an array to store a diagnostic target grid @@ -1107,9 +1107,9 @@ subroutine define_axes_group(diag_cs, handles, axes, nz, vertical_coordinate_num if (axes%is_u_point) axes%mask2d => diag_cs%mask2dCu if (axes%is_v_point) axes%mask2d => diag_cs%mask2dCv if (axes%is_q_point) axes%mask2d => diag_cs%mask2dBu - call axes%buffer_2d%set_horizontal_extents(lbound(axes%mask2d,1), ubound(axes%mask2d,1), & + call axes%piecemeal_2d%set_horizontal_extents(lbound(axes%mask2d,1), ubound(axes%mask2d,1), & lbound(axes%mask2d,2), ubound(axes%mask2d,2)) - call axes%buffer_2d%set_fill_value(diag_cs%missing_value) + call axes%piecemeal_2d%set_fill_value(diag_cs%missing_value) endif ! A static 3d mask for non-native coordinates can only be setup when a grid is available axes%mask3d => null() @@ -1126,10 +1126,10 @@ subroutine define_axes_group(diag_cs, handles, axes, nz, vertical_coordinate_num if (axes%is_v_point) axes%mask3d => diag_cs%mask3dCvi if (axes%is_q_point) axes%mask3d => diag_cs%mask3dBi endif - call axes%buffer_3d%set_horizontal_extents(is=lbound(axes%mask3d,1), ie=ubound(axes%mask3d,1), & + call axes%piecemeal_3d%set_horizontal_extents(is=lbound(axes%mask3d,1), ie=ubound(axes%mask3d,1), & js=lbound(axes%mask3d,2), je=ubound(axes%mask3d,2)) - call axes%buffer_3d%set_vertical_extent(ks=lbound(axes%mask3d,3), ke=ubound(axes%mask3d,3)) - call axes%buffer_3d%set_fill_value(diag_cs%missing_value) + call axes%piecemeal_3d%set_vertical_extent(ks=lbound(axes%mask3d,3), ke=ubound(axes%mask3d,3)) + call axes%piecemeal_3d%set_fill_value(diag_cs%missing_value) endif @@ -1913,8 +1913,8 @@ subroutine post_data_3d_by_column(diag_field_id, field, diag_cs, i, j) integer :: buffer_slot diag => diag_cs%diags(diag_field_id) - buffer_slot = diag%axes%buffer_3d%check_capacity_by_id(diag_field_id) - diag%axes%buffer_3d%buffer(buffer_slot,i,j,:) = field + buffer_slot = diag%axes%piecemeal_3d%check_capacity_by_id(diag_field_id) + diag%axes%piecemeal_3d%buffer(buffer_slot)%field(i,j,:) = field(:) end subroutine post_data_3d_by_column !> Put data into the buffer for a diagnostic one point at a time @@ -1932,8 +1932,8 @@ subroutine post_data_3d_by_point(diag_field_id, field, diag_cs, i, j, k) integer :: buffer_slot diag => diag_cs%diags(diag_field_id) - buffer_slot = diag%axes%buffer_3d%find_buffer_slot(diag_field_id) - diag%axes%buffer_3d%buffer(buffer_slot,i,j,k) = field + buffer_slot = diag%axes%piecemeal_3d%find_buffer_slot(diag_field_id) + diag%axes%piecemeal_3d%buffer(buffer_slot)%field(i,j,k) = field end subroutine post_data_3d_by_point !> Post the final buffer using the standard post_data interface @@ -1946,9 +1946,9 @@ subroutine post_data_3d_final(diag_field_id, diag_cs) integer :: buffer_slot diag => diag_cs%diags(diag_field_id) - buffer_slot = diag%axes%buffer_3d%find_buffer_slot(diag_field_id) - call post_data(diag_field_id, diag%axes%buffer_3d%buffer(buffer_slot,:,:,:), diag_CS) - call diag%axes%buffer_3d%mark_available(diag_field_id) + buffer_slot = diag%axes%piecemeal_3d%find_buffer_slot(diag_field_id) + call post_data(diag_field_id, diag%axes%piecemeal_3d%buffer(buffer_slot)%field(:,:,:), diag_CS) + call diag%axes%piecemeal_3d%mark_available(diag_field_id) end subroutine post_data_3d_final !> Calculate and write out diagnostics that are the product of two 3-d arrays at u-points