diff --git a/src/framework/MOM_diag_buffers.F90 b/src/framework/MOM_diag_buffers.F90 new file mode 100644 index 0000000000..7993d47548 --- /dev/null +++ b/src/framework/MOM_diag_buffers.F90 @@ -0,0 +1,193 @@ +!> Provides buffers that can dynamically grow as needed. These are primarily intended for the +!! 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 + +! This file is part of MOM6. See LICENSE.md for the license. + +implicit none ; private + +public :: diag_buffer_unit_tests + +type :: diag_buffer_base ; private + integer :: is !< The start index of the array i-direction + integer :: js !< The start index of the array j-direction + integer :: ie !< The end index of the array i-direction + integer :: je !< The end index of the array j-direction + + integer :: length = 0 !< The size of the buffer + integer, allocatable, dimension(:) :: ids !< List of diagnostic ids whose index corresponds to the row in the buffer + + contains + + procedure, public :: grow_ids + procedure, public :: find_buffer_index +end type diag_buffer_base + +!> Dynamically growing buffer for 2D arrays. +type, extends(diag_buffer_base), public :: diag_buffer_2d + real, allocatable, dimension(:,:,:) :: buffer !< The actual buffer to store data + + contains + + procedure :: grow => grow_2d + procedure :: store => store_2d +end type diag_buffer_2d + +!> Dynamically growing buffer for 3D arrays. +type, extends(diag_buffer_base), public :: diag_buffer_3d + real, allocatable, dimension(:,:,:,:) :: buffer !< The actual buffer to store data + + contains + +! procedure :: grow => grow_3d +! procedure :: store => store_3d +end type diag_buffer_3d + +contains + +!> Return the index of the buffer corresponding to the diagnostic id +pure function find_buffer_index(this, id) result(index) + class(diag_buffer_base), intent(in) :: this !< The diagnostic buffer + integer, intent(in) :: id !< The diagnostic id + + integer, dimension(1) :: temp + integer :: index !< The index in the buffer corresponding to the diagnostic id + + !NOTE: Alternatively could do index = SUM(findloc(...)) + temp = findloc(this%ids(:), id) + index = temp(1) + +end function find_buffer_index + +!> Mark an index in the buffer as available for use by clearing the entry in index -> id map array +subroutine mark_available(this, idx) + class(diag_buffer_2d) :: this !< This 2d buffer + integer, intent(in) :: idx !< The diagnostic id + + this%ids(idx) = 0 +end subroutine mark_available + +!> Grow the ids array +subroutine grow_ids(this) + class(diag_buffer_base), intent(inout) :: this !< This buffer + + integer, allocatable, dimension(:) :: temp + integer :: n_old, n_new + + n_old = this%length + n_new = n_old + 1 + + allocate(temp(n_new), source=0) + if(n_old>0) temp(1:n_old) = this%ids + call move_alloc(temp, this%ids) +end subroutine + +!> Grow a buffer for 2D arrays +subroutine grow_2d(this) + class(diag_buffer_2d) :: this !< This 2d buffer + + integer :: n_old, n_new + real, allocatable, dimension(:,:,:) :: temp + + n_old = this%length + n_new = n_old + 1 + + allocate(temp(n_new, this%is:this%ie, this%js:this%je), source=0.) + if(n_old>0) temp(1:n_old,:,:) = this%buffer + call move_alloc(temp, this%buffer) +end subroutine grow_2d + +!> Store a 2D array into this buffer +subroutine store_2d(this, data, id) + class(diag_buffer_2d), intent(inout) :: this !< This 2d buffer + real, dimension(:,:), intent(in) :: data !< The data to be stored in the buffer + integer, optional, intent(in) :: id !< The diagnostic id + + integer :: idx + + ! Find the first index in the ids array that is 0, i.e. this is a portion of the buffer that can be reused + idx = this%find_buffer_index(0) + if(idx==0) then + call this%grow() + idx = this%length + 1 + endif + this%buffer(idx,:,:) = data(:,:) + this%ids(idx) = id + +end subroutine store_2d + +function diag_buffer_unit_tests(verbose) result(fail) + logical, intent(in) :: verbose !< If true, write results to stdout + logical :: fail !< True if any of the unit tests fail + + fail = .false. + write(stdout,*) '==== MOM_diag_buffers: diag_buffers_unit_tests ===' + fail = fail .or. new_buffer_2d() + + contains + + !> Ensure properties of a newly initialized buffer + function new_buffer_2d() result(local_fail) + type(diag_buffer_2d) :: buffer_2d + 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 + if(verbose) write(stdout,*) "new_buffer_2d: ", local_fail + end function new_buffer_2d + + function grow_buffer_2d() result(local_fail) + type(diag_buffer_2d) :: buffer_2d + logical :: local_fail !< True if any of the unit tests fail + integer :: is, ie, js, je + integer :: i + + local_fail = .false. + is = 1; ie=2; js=3; je=6 + + buffer_2d = diag_buffer_2d(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) + enddo + if(v) write(stdout,*) "grow_buffer_2d: ", local_fail + end function grow_buffer_2d + + function store_buffer_2d() result(local_fail) + type(diag_buffer_2d) :: buffer_2d + logical :: local_fail !< True if any of the unit tests fail + + integer :: is, ie, js, je, nlen + integer :: i + real, allocatable, dimension(:,:,:) :: test_2d + + nlen = 3 + is = 1; ie = 2; js = 3; je = 6 + + 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) + + do i=1,nlen + call buffer_2d%store(test_2d(i,:,:), i) + enddo + local_fail = ANY(buffer_2d%buffer /= test_2d) + + if(v) write(stdout,*) "store_buffer_2d: ", local_fail + end function store_buffer_2d +end function diag_buffer_unit_tests + +end module MOM_diag_buffers + +!> \namespace mom_cpu_clock +!! +!! APIs are defined and implemented in mom_cpu_clock_infra.