more fix of MPI writing of timers

moment_kinetics/src/file_io.jl

@@ -2585,8 +2585,8 @@ function write_timing_data(io_moments, t_idx, dfns=false)
             push!(times, timer.accumulated_data.time)
             push!(ncalls, timer.accumulated_data.ncalls)
             push!(allocs, timer.accumulated_data.allocs)
-            for (sub_dict,sub_timer) ∈ zip(values(names_dict), values(timer.inner_timers))
-                walk_through_timers(sub_dict, sub_timer)
+            for (sub_name, sub_dict) ∈ pairs(names_dict)
+                walk_through_timers(sub_dict, timer[sub_name])
             end
         end
         walk_through_timers(sorted_timer_names["moment_kinetics"],
@@ -2705,9 +2705,9 @@ function write_timing_data(io_moments, t_idx, dfns=false)
         allocs_data = Dict{mk_int,Vector{mk_int}}()
         times_data[0], ncalls_data[0], allocs_data[0] = get_data_from_timers()
         for irank ∈ 1:block_size[]-1
-            times_data[irank] = zeros(mk_int, timer_names_per_rank[irank][2])
-            ncalls_data[irank] = zeros(mk_int, timer_names_per_rank[irank][2])
-            allocs_data[irank] = zeros(mk_int, timer_names_per_rank[irank][2])
+            times_data[irank] = zeros(mk_int, timer_names_per_rank[irank][2][])
+            ncalls_data[irank] = zeros(mk_int, timer_names_per_rank[irank][2][])
+            allocs_data[irank] = zeros(mk_int, timer_names_per_rank[irank][2][])
         end
         times_reqs = MPI.Request[MPI.Irecv!(times_data[irank], comm_block[]; source=irank)
                                  for irank ∈ 1:block_size[]-1]
@@ -2740,7 +2740,12 @@ function write_timing_data(io_moments, t_idx, dfns=false)
                 io_ncalls = io_group["ncalls:" * this_name]
                 io_allocs = io_group["allocs:" * this_name]
                 if t_idx < 0
-                    this_t_idx = length(io_time)
+                    # "time:moment_kinetics" is the top-level timer and is the first one
+                    # created, so it should always exist, and be written at each timestep.
+                    # If we got the length of `time:$this_name`, the variable might have
+                    # the wrong length (e.g. if it has only just been created and has
+                    # length 1).
+                    this_t_idx = length(io_group["time:moment_kinetics"])
                 else
                     this_t_idx = t_idx
                 end
@@ -2774,17 +2779,24 @@ function write_timing_data(io_moments, t_idx, dfns=false)
                 io_ncalls = io_group["ncalls:" * this_name]
                 io_allocs = io_group["allocs:" * this_name]
                 if t_idx < 0
-                    this_t_idx = length(io_time)
+                    # "time:moment_kinetics" is the top-level timer and is the first one
+                    # created, so it should always exist, and be written at each timestep.
+                    # If we got the length of `time:$this_name`, the variable might have
+                    # the wrong length (e.g. if it has only just been created and has
+                    # length 1).
+                    this_t_idx = length(io_group["time:moment_kinetics"])
                 else
                     this_t_idx = t_idx
                 end
                 if write_from_this_rank
-                    append_to_dynamic_var(io_time, times_data[counter], this_t_idx,
-                                          parallel_io, write_from_this_rank=true)
-                    append_to_dynamic_var(io_ncalls, ncalls_data[counter], this_t_idx,
-                                          parallel_io, write_from_this_rank=true)
-                    append_to_dynamic_var(io_allocs, allocs_data[counter], this_t_idx,
+                    append_to_dynamic_var(io_time, times_data[irank][counter], this_t_idx,
                                           parallel_io, write_from_this_rank=true)
+                    append_to_dynamic_var(io_ncalls, ncalls_data[irank][counter],
+                                          this_t_idx, parallel_io,
+                                          write_from_this_rank=true)
+                    append_to_dynamic_var(io_allocs, allocs_data[irank][counter],
+                                          this_t_idx, parallel_io,
+                                          write_from_this_rank=true)
                 else
                     append_to_dynamic_var(io_time, nothing, this_t_idx,
                                           parallel_io, write_from_this_rank=false)
@@ -2827,7 +2839,7 @@ properly.
 """
 function write_final_timing_data_to_binary(io_or_file_info_moments, io_or_file_info_dfns)
     io_moments = nothing
-    io_dfns = nothing
+    io_dfns_moments = nothing
     @serial_region begin
         # Only read/write from first process in each 'block'
 
@@ -2846,10 +2858,11 @@ function write_final_timing_data_to_binary(io_or_file_info_moments, io_or_file_i
             io_dfns = reopen_dfns_io(io_or_file_info_dfns)
             closefile = true
         end
+        io_dfns_moments = io_dfns.io_moments
     end
 
     write_timing_data(io_moments, -1)
-    write_timing_data(io_dfns.io_moments, -1, true)
+    write_timing_data(io_dfns_moments, -1, true)
 
     @serial_region begin
         closefile && close(io_moments.fid)