Create the diag buffer type for 2d arrays

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.