cleanup

docs/src/utilities.md

@@ -10,4 +10,5 @@ modules in the coupler.
     ClimaCoupler.Utilities.CouplerSimulation
     ClimaCoupler.Utilities.heaviside
     ClimaCoupler.Utilities.swap_space!
+    ClimaCoupler.Utilities.trigger_callback
 ```

experiments/AMIP/moist_mpi_earth/coupler_driver_modular.jl

@@ -61,6 +61,7 @@ using UnPack
 using ClimaCore.Utilities: half, PlusHalf
 using ClimaCore: InputOutput, Fields
 
+
 include("cli_options.jl")
 (s, parsed_args) = parse_commandline()
 
@@ -93,7 +94,8 @@ import ClimaCoupler
 import ClimaCoupler.Regridder: land_sea_mask, update_masks!, combine_surfaces!, dummmy_remap!
 import ClimaCoupler.ConservationChecker:
     EnergyConservationCheck, WaterConservationCheck, check_conservation!, plot_global_conservation
-import ClimaCoupler.Utilities: CoupledSimulation, float_type, swap_space!
+import ClimaCoupler.Utilities: CoupledSimulation, float_type, swap_space!, Monthly, EveryTimestep
+import ClimaCoupler.Diagnostics: get_var, init_diagnostics, accumulate_diagnostics!, save_diagnostics, TimeMean
 
 pkg_dir = pkgdir(ClimaCoupler)
 COUPLER_OUTPUT_DIR = joinpath(pkg_dir, "experiments/AMIP/moist_mpi_earth/output", joinpath(mode_name, run_name))
@@ -113,7 +115,6 @@ mask_data = joinpath(mask_dataset_path(), "seamask.nc")
 include("coupler_utils/flux_calculator.jl")
 include("coupler_utils/calendar_timer.jl")
 include("coupler_utils/bcfile_reader.jl")
-include("../../../src/Diagnostics.jl")
 include("coupler_utils/offline_postprocessor.jl")
 
 ## user-specified diagnostics
@@ -259,7 +260,7 @@ monthly_3d_diags = init_diagnostics(
     (:T, :u, :q_tot),
     atmos_sim.domain.center_space;
     save = Monthly(),
-    operations = (; accumulate = TimeMean([0])),
+    operations = (; accumulate = TimeMean([Int(0)])),
     output_dir = COUPLER_OUTPUT_DIR,
     name_tag = "monthly_mean_3d_",
 )
@@ -268,7 +269,7 @@ monthly_2d_diags = init_diagnostics(
     (:precipitation, :toa, :T_sfc),
     boundary_space;
     save = Monthly(),
-    operations = (; accumulate = TimeMean([0])),
+    operations = (; accumulate = TimeMean([Int(0)])),
     output_dir = COUPLER_OUTPUT_DIR,
     name_tag = "monthly_mean_2d_",
 )

experiments/AMIP/moist_mpi_earth/user_diagnostics.jl

@@ -19,7 +19,8 @@ end
 
 Zonal wind (m s⁻¹).
 """
-get_var(cs::CoupledSimulation, ::Val{:u}) = ClimaCore.Geometry.UVVector.(cs.model_sims.atmos_sim.integrator.u.c.uₕ).components.data.:1
+get_var(cs::CoupledSimulation, ::Val{:u}) =
+    ClimaCore.Geometry.UVVector.(cs.model_sims.atmos_sim.integrator.u.c.uₕ).components.data.:1
 
 """
     get_var(cs, ::Val{:q_tot}) 
@@ -83,4 +84,3 @@ Combined surface temperature (K).
 get_var(cs::CoupledSimulation, ::Val{:T_sfc}) = swap_space!(cs.fields.T_S, cs.boundary_space)
 
 # land diagnotics
-

src/ClimaCoupler.jl

@@ -12,5 +12,6 @@ include("../test/TestHelper.jl")
 include("Utilities.jl")
 include("Regridder.jl")
 include("ConservationChecker.jl")
+include("Diagnostics.jl")
 
 end

src/Diagnostics.jl

@@ -3,12 +3,14 @@
 
 This module contains functions for defining, gathering and outputting online model diagnostics from the Coupler.  
 =#
-#module Diagnostics
+module Diagnostics
 
-using ClimaCore: Spaces, Fields#, Domains, Meshes, Topologies
-using ClimaCoupler.Utilities: AbstractFrequency, Monthly, EveryTimestep
+using ClimaCore: Spaces, Fields
+using ClimaCoupler.Utilities: AbstractFrequency, Monthly, EveryTimestep, CoupledSimulation
+using Dates
+using ClimaComms
 
-export get_var
+export get_var, init_diagnostics, accumulate_diagnostics!, save_diagnostics, TimeMean
 
 abstract type AbstractOutputGroup end
 
@@ -32,12 +34,17 @@ end
 
 Initializes diagnostics groups. 
 """
-function init_diagnostics(names::Tuple, space::Spaces.AbstractSpace; save = EveryTimestep(), operations = (;), output_dir = "", name_tag = "")
-    data = NamedTuple{names}(
-        ntuple(i -> Fields.zeros(space), length(names)),
-    )
+function init_diagnostics(
+    names::Tuple,
+    space::Spaces.AbstractSpace;
+    save = EveryTimestep(),
+    operations = (;),
+    output_dir = "",
+    name_tag = "",
+)
+    data = NamedTuple{names}(ntuple(i -> Fields.zeros(space), length(names)))
     return DiagnosticsGroup(Fields.FieldVector(; data...), operations, save, output_dir, name_tag)
-end 
+end
 
 """
     get_var(cs::CoupledSimulation, x)
@@ -49,7 +56,7 @@ Example:
 get_var(cs, ::Val{:T_sfc}) = cs.fields.T_S
 
 """
-get_var(::CoupledSimulation, x ) = @warn "Variable $x is not defined." 
+get_var(::CoupledSimulation, x) = @warn "Variable $x is not defined."
 
 """
     accumulate_diagnostics!(cs::CoupledSimulation)
@@ -67,7 +74,7 @@ end
 """
     collect_diags(cs::CoupledSimulation, dg::DiagnosticsGroup)
 
-Collects diagnostics in diags names
+Collects diagnostics in diags names.
 """
 function collect_diags(cs::CoupledSimulation, dg::DiagnosticsGroup)
 
@@ -89,12 +96,12 @@ Applies iteratively all specified diagnostics operations.
 
 function iterate_operations(cs::CoupledSimulation, dg::DiagnosticsGroup, new_diags::Fields.FieldVector)
     if isempty(dg.operations)
-        operation(cs, dg, new_diags, nothing) 
+        operation(cs, dg, new_diags, nothing)
     else
         for op in dg.operations
-            operation(cs, dg, new_diags, op) 
+            operation(cs, dg, new_diags, op)
         end
-    end 
+    end
 end
 
 """
@@ -104,7 +111,7 @@ Accumulates in time all entries in `new_diags` and saves the result in `dg.field
 """
 function operation(::CoupledSimulation, dg::DiagnosticsGroup, new_diags::Fields.FieldVector, ::TimeMean)
     dg.field_vector .+= new_diags
-    dg.operations.accumulate.ct[1] += FT(1)
+    dg.operations.accumulate.ct[1] += Int(1)
 
     return nothing
 end
@@ -128,12 +135,9 @@ Saves all entries in `dg` in separate HDF5 files per variable in `output_dir`.
 """
 function save_diagnostics(cs::CoupledSimulation)
     for dg in cs.diagnostics
-        if trigger_callback(cs, dg.save) 
+        if trigger_callback(cs, dg.save)
             pre_save(dg.operations.accumulate, cs, dg)
-            save_diagnostics(
-                    cs,
-                    dg,
-                )
+            save_diagnostics(cs, dg)
             post_save(dg.operations.accumulate, cs, dg)
         end
     end
@@ -153,7 +157,7 @@ function save_diagnostics(cs::CoupledSimulation, dg::DiagnosticsGroup)
     @info "Saving coupler diagnostics: $tag $diag_save"
 
     for (name, values) in zip(diag_names, diag_values)
-        output_file = joinpath(output_dir,  "$name.$diag_save.$tag." * string(date) * ".hdf5")
+        output_file = joinpath(output_dir, "$name.$diag_save.$tag." * string(date) * ".hdf5")
         hdfwriter = InputOutput.HDF5Writer(output_file, cs.comms_ctx)
         InputOutput.HDF5.write_attribute(hdfwriter.file, "unix time", save_time_format(date, diag_save))
         InputOutput.write!(hdfwriter, values, string(name))
@@ -191,18 +195,19 @@ end
 
 Collects variables and performs all specified operations before saving the snapshot diagnostics.
 """
-pre_save(::Nothing, cs::CoupledSimulation, dg::DiagnosticsGroup) = iterate_operations(cs, dg, collect_diags(cs, dg)) 
+pre_save(::Nothing, cs::CoupledSimulation, dg::DiagnosticsGroup) = iterate_operations(cs, dg, collect_diags(cs, dg))
 
 """
     post_save(::TimeMean, cs::CoupledSimulation, dg::DiagnosticsGroup)
 
 Resets accumulating fields and counts after saving the diagnostics. 
 """
 function post_save(::TimeMean, cs::CoupledSimulation, dg::DiagnosticsGroup)
-    dg.field_vector .= (0.0)
+    FT = eltype(dg.field_vector)
+    dg.field_vector .= FT(0.0)
     dg.operations.accumulate.ct .= FT(0)
 end
 
 post_save(::Nothing, cs::CoupledSimulation, dg::DiagnosticsGroup) = nothing
 
-# end #module
+end # module

test/conservation_checker_tests.jl

@@ -64,7 +64,7 @@ function coupler_sim_from_file(
     FT = eltype(land_mask)
 
     Utilities.CoupledSimulation(
-        ClimaComms.SingletonCommsContext(), 
+        ClimaComms.SingletonCommsContext(),
         tspan,
         dates,
         boundary_space,

test/diagnostics_tests.jl

@@ -2,26 +2,28 @@
     Unit tests for ClimaCoupler Diagnostics module
 =#
 using Test
-using ClimaCoupler.Utilities
+using Dates
+using ClimaComms
+using ClimaCoupler: Utilities, TestHelper
 
 FT = Float64
-get_var(cs::CoupledSimulation, ::Val{:x}) = FT(1) 
+get_var(cs::CoupledSimulation, ::Val{:x}) = FT(1)
 
 @testset "init_diagnostics" begin
     names = (:x, :y)
     space = create_space(FT)
     dg = init_diagnostics(names, space)
     @test typeof(dg) == DiagnosticsGroup{EveryTimestep}
-end 
+end
 
 @testset "accumulate_diagnostics!, collect_diags, iterate_operations, operation{accumulation{TimeMean, Nothing}}" begin
-    cases = (nothing, TimeMean([0]))
+    cases = (nothing, TimeMean([Int(0)]))
     expected_results = (FT(2), FT(3))
     for (c_i, case) in enumerate(cases)
         names = (:x,)
         space = create_space(FT)
-        dg_2d = init_diagnostics(names, space,  save = EveryTimestep(), operations =  (; accumulate = case))
-        dg_2d.field_vector .= FT(2) 
+        dg_2d = init_diagnostics(names, space, save = EveryTimestep(), operations = (; accumulate = case))
+        dg_2d.field_vector .= FT(2)
         cs = Utilities.CoupledSimulation(
             nothing, # comms_ctx
             nothing, # tspan
@@ -40,18 +42,24 @@ end
         accumulate_diagnostics!(cs)
         @test cs.diagnostics[1].field_vector[1] == expected_results[c_i]
 
-    end 
+    end
 end
 
 @testset "save_diagnostics" begin
     test_dir = "diag_test_dir"
     nmes = (:x,)
     space = create_space(FT)
-    dg_2d = init_diagnostics(nmes, space,  save = EveryTimestep(), operations = (; accumulate = TimeMean([0])), output_dir = test_dir) # or use accumulate = nothing for snapshop save
+    dg_2d = init_diagnostics(
+        nmes,
+        space,
+        save = EveryTimestep(),
+        operations = (; accumulate = TimeMean([Int(0)])),
+        output_dir = test_dir,
+    ) # or use accumulate = nothing for snapshop save
     cs = Utilities.CoupledSimulation(
         ClimaComms.SingletonCommsContext(), # comms_ctx
         nothing,# tspan
-        (date = [DateTime(0,2)], date1 = [DateTime(0,1)]), # dates
+        (date = [DateTime(0, 2)], date1 = [DateTime(0, 1)]), # dates
         nothing, # boundary_space
         nothing, # fields
         nothing, # parsed_args
@@ -63,28 +71,28 @@ end
         (;), # mode
         (dg_2d,), # diagnostics
     )
-    save_diagnostics(cs,cs.diagnostics[1])
-    file = filter(x->endswith(x, ".hdf5"), readdir(test_dir))
+    save_diagnostics(cs, cs.diagnostics[1])
+    file = filter(x -> endswith(x, ".hdf5"), readdir(test_dir))
     @test !isempty(file)
     rm(test_dir; recursive = true, force = true)
-    
+
 end
 
 @testset "save_time_format" begin
-    date = DateTime(1970,2,1,0,1)
+    date = DateTime(1970, 2, 1, 0, 1)
     unix = save_time_format(date, Monthly())
     @test unix == 0
 end
 
 @testset "pre_save{TimeMean, Nothing}, post_save" begin
-    cases = (nothing, TimeMean([0]),)
-    expected_results = ((FT(3), FT(1), FT(1)), (FT(4), FT(2.5), FT(0)),)
+    cases = (nothing, TimeMean([Int(0)]))
+    expected_results = ((FT(3), FT(1), FT(1)), (FT(4), FT(2.5), FT(0)))
 
     for (c_i, case) in enumerate(cases)
         names = (:x,)
         space = create_space(FT)
-        dg_2d = init_diagnostics(nmes, space,  save = EveryTimestep(), operations = (; accumulate = case))
-        dg_2d.field_vector .= FT(3) 
+        dg_2d = init_diagnostics(nmes, space, save = EveryTimestep(), operations = (; accumulate = case))
+        dg_2d.field_vector .= FT(3)
         cs = Utilities.CoupledSimulation(
             nothing, # comms_ctx
             nothing, # tspan
@@ -110,4 +118,3 @@ end
         @test cs.diagnostics[1].field_vector[1] == expected_results[c_i][3]
     end
 end
-