From 35d4e866cf99202d60bb1c3bb36f433d0bc51fa9 Mon Sep 17 00:00:00 2001
From: Ryan Mulhall <35538242+rem1776@users.noreply.github.com>
Date: Wed, 15 Nov 2023 13:11:53 -0500
Subject: [PATCH] fix: diag_send_complete loops and add get_file_ids (#1407)

---
 diag_manager/fms_diag_field_object.F90 |  8 ++++
 diag_manager/fms_diag_object.F90       | 66 ++++++++++++++------------
 2 files changed, 43 insertions(+), 31 deletions(-)

diff --git a/diag_manager/fms_diag_field_object.F90 b/diag_manager/fms_diag_field_object.F90
index 1eb0221e94..ffecfc650a 100644
--- a/diag_manager/fms_diag_field_object.F90
+++ b/diag_manager/fms_diag_field_object.F90
@@ -164,6 +164,7 @@ module fms_diag_field_object_mod
      procedure :: get_math_needs_to_be_done
      procedure :: add_area_volume
      procedure :: append_time_cell_methods
+     procedure :: get_file_ids
 end type fmsDiagField_type
 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! variables !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 type(fmsDiagField_type) :: null_ob
@@ -1639,5 +1640,12 @@ function get_starting_compute_domain(axis_ids, diag_axis) &
   enddo axis_loop
 end function get_starting_compute_domain
 
+!> Get list of field ids
+pure function get_file_ids(this)
+  class(fmsDiagField_type), intent(in) :: this
+  integer, allocatable :: get_file_ids(:) !< Ids of the FMS_diag_files the variable
+  get_file_ids = this%file_ids
+end function
+
 #endif
 end module fms_diag_field_object_mod
diff --git a/diag_manager/fms_diag_object.F90 b/diag_manager/fms_diag_object.F90
index 894d8023d8..ea651e725f 100644
--- a/diag_manager/fms_diag_object.F90
+++ b/diag_manager/fms_diag_object.F90
@@ -46,6 +46,7 @@ module fms_diag_object_mod
 use omp_lib
 #endif
 use mpp_domains_mod, only: domain1d, domain2d, domainUG, null_domain2d
+use fms_string_utils_mod, only: string
 use platform_mod
 implicit none
 private
@@ -648,43 +649,46 @@ subroutine fms_diag_send_complete(this, time_step)
   class(*), pointer :: input_data_buffer(:,:,:,:)
   character(len=128) :: error_string
   type(fmsDiagIbounds_type) :: bounds
+  integer, dimension(:), allocatable :: file_ids !< Array of file IDs for a field
+  logical, parameter :: DEBUG_SC = .true. !< turn on output for debugging
 
   !< Update the current model time by adding the time_step
   this%current_model_time = this%current_model_time + time_step
 
   !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
   !! In the future, this may be parallelized for offloading
-  file_loop: do ifile = 1, size(this%FMS_diag_files)
-    diag_file => this%FMS_diag_files(ifile)
-    field_outer_if: if (size(diag_file%FMS_diag_file%get_field_ids()) .ge. 1) then
-      allocate (file_field_ids(size(diag_file%FMS_diag_file%get_field_ids() )))
-      file_field_ids = diag_file%FMS_diag_file%get_field_ids()
-      field_loop: do ifield = 1, size(file_field_ids)
-        ! If the field is not registered go away
-        if (.not. diag_file%FMS_diag_file%is_field_registered(ifield)) cycle
-
-        diag_field => this%FMS_diag_fields(file_field_ids(ifield))
-        !> Check if math needs to be done
-        math = diag_field%get_math_needs_to_be_done()
-        calling_math: if (math) then
-          input_data_buffer => diag_field%get_data_buffer()
-          call bounds%reset_bounds_from_array_4D(input_data_buffer)
-          call this%allocate_diag_field_output_buffers(input_data_buffer, file_field_ids(ifield))
-          error_string = this%fms_diag_do_reduction(input_data_buffer, file_field_ids(ifield), &
-            diag_field%get_mask(), diag_field%get_weight(), &
-            bounds, .False., Time=this%current_model_time)
-          if (trim(error_string) .ne. "") call mpp_error(FATAL, "Field:"//trim(diag_field%get_varname()//&
-            " -"//trim(error_string)))
-        endif calling_math
-        !> Clean up, clean up, everybody everywhere
-        if (associated(diag_field)) nullify(diag_field)
-      enddo field_loop
-      !> Clean up, clean up, everybody do your share
-      if (allocated(file_field_ids)) deallocate(file_field_ids)
-    endif field_outer_if
-  enddo file_loop
-
-  call this%fms_diag_do_io()
+  ! loop through each field
+  field_loop: do ifield = 1, size(this%FMS_diag_fields)
+    diag_field => this%FMS_diag_fields(ifield)
+    if(.not. diag_field%is_registered()) cycle
+    if(DEBUG_SC) call mpp_error(NOTE, "fms_diag_send_complete:: var: "//diag_field%get_varname())
+    ! get files the field is in
+    allocate (file_ids(size(diag_field%get_file_ids() )))
+    file_ids = diag_field%get_file_ids()
+    math = diag_field%get_math_needs_to_be_done()
+    ! if doing math loop through each file for given field
+    doing_math: if (size(file_ids) .ge. 1 .and. math) then
+      ! Check if buffer alloc'd
+      has_input_buff: if (diag_field%has_input_data_buffer()) then
+        input_data_buffer => diag_field%get_data_buffer()
+        ! reset bounds, allocate output buffer, and update it with reduction
+        call bounds%reset_bounds_from_array_4D(input_data_buffer)
+        call this%allocate_diag_field_output_buffers(input_data_buffer, ifield)
+        error_string = this%fms_diag_do_reduction(input_data_buffer, ifield, &
+                              diag_field%get_mask(), diag_field%get_weight(), &
+                              bounds, .False., Time=this%current_model_time)
+        if (trim(error_string) .ne. "") call mpp_error(FATAL, "Field:"//trim(diag_field%get_varname()//&
+                                                       " -"//trim(error_string)))
+      else
+        call mpp_error(FATAL, "diag_send_complete:: no input buffer allocated for field"//diag_field%get_longname())
+      endif has_input_buff
+    endif doing_math
+    !> Clean up, clean up, everybody do your share
+    if (allocated(file_ids)) deallocate(file_ids)
+    if (associated(diag_field)) nullify(diag_field)
+  enddo field_loop
+
+call this%fms_diag_do_io()
 #endif
 
 end subroutine fms_diag_send_complete