Merge pull request #3476 from CliMA/zs/conservation_diagnostics

add mass and energy conservation diagnostics

NEWS.md

@@ -3,7 +3,13 @@ ClimaAtmos.jl Release Notes
 
 Main
 -------
+### Features
+
+### Write diagnostics to text files
 
+Added functionality to write diagnostics in DictWriter to text files.
+This is useful for outputing scalar diagnostics, such as total mass of
+the atmosphere. PR [3476](https://github.com/CliMA/ClimaAtmos.jl/pull/3476)
 
 v0.27.9
 -------

config/model_configs/gpu_aquaplanet_dyamond.yml

@@ -24,3 +24,7 @@ prescribed_aerosols: ["CB1", "CB2", "DST01", "DST02", "DST03", "DST04", "OC1", "
 prescribe_clouds_in_radiation: true
 radiation_reset_rng_seed: true
 toml: [toml/longrun_aquaplanet.toml]
+diagnostics:
+  - short_name: [massa, energya]
+    period: 1hours
+    writer: dict

config/model_configs/sphere_baroclinic_wave_conservation.yml

@@ -4,3 +4,7 @@ t_end: "5days"
 moist: "equil"
 dt_save_state_to_disk: "5days"
 check_conservation: true
+diagnostics:
+  - short_name: [massa, energya]
+    period: 1days
+    writer: dict

config/model_configs/sphere_baroclinic_wave_conservation_ft64.yml

@@ -5,3 +5,7 @@ t_end: "5days"
 moist: "equil"
 dt_save_state_to_disk: "5days"
 check_conservation: true
+diagnostics:
+  - short_name: [massa, energya]
+    period: 1days
+    writer: dict

config/model_configs/sphere_baroclinic_wave_conservation_source.yml

@@ -9,3 +9,7 @@ precip_model: "0M"
 vert_diff: true
 dt_save_state_to_disk: "5days"
 check_conservation: true
+diagnostics:
+  - short_name: [massa, energya, watera, energyo, watero]
+    period: 1days
+    writer: dict

docs/Manifest.toml

@@ -342,10 +342,10 @@ version = "0.14.20"
     Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
 
 [[deps.ClimaDiagnostics]]
-deps = ["Accessors", "ClimaComms", "ClimaCore", "Dates", "NCDatasets", "SciMLBase"]
-git-tree-sha1 = "ace174fe5e4ae04c50a7835683baecd92e49c0d4"
+deps = ["Accessors", "ClimaComms", "ClimaCore", "Dates", "NCDatasets", "OrderedCollections", "SciMLBase"]
+git-tree-sha1 = "e9ac94af815dcae2a2ab24e54b53e76cca6258b7"
 uuid = "1ecacbb8-0713-4841-9a07-eb5aa8a2d53f"
-version = "0.2.10"
+version = "0.2.11"
 
 [[deps.ClimaParams]]
 deps = ["TOML"]

examples/Manifest.toml

@@ -393,10 +393,10 @@ uuid = "c8b6d40d-e815-466f-95ae-c48aefa668fa"
 version = "0.7.6"
 
 [[deps.ClimaDiagnostics]]
-deps = ["Accessors", "ClimaComms", "ClimaCore", "Dates", "NCDatasets", "SciMLBase"]
-git-tree-sha1 = "ace174fe5e4ae04c50a7835683baecd92e49c0d4"
+deps = ["Accessors", "ClimaComms", "ClimaCore", "Dates", "NCDatasets", "OrderedCollections", "SciMLBase"]
+git-tree-sha1 = "e9ac94af815dcae2a2ab24e54b53e76cca6258b7"
 uuid = "1ecacbb8-0713-4841-9a07-eb5aa8a2d53f"
-version = "0.2.10"
+version = "0.2.11"
 
 [[deps.ClimaParams]]
 deps = ["TOML"]

examples/hybrid/driver.jl

@@ -132,6 +132,9 @@ if config.parsed_args["check_conservation"]
     @test water_conservation ≈ 0 atol = 100 * eps(FT)
 end
 
+# Write diagnostics that are in DictWriter to text files
+CA.write_diagnostics_as_txt(simulation)
+
 # Visualize the solution
 if ClimaComms.iamroot(config.comms_ctx)
     include(

perf/Manifest.toml

@@ -404,10 +404,10 @@ uuid = "c8b6d40d-e815-466f-95ae-c48aefa668fa"
 version = "0.7.6"
 
 [[deps.ClimaDiagnostics]]
-deps = ["Accessors", "ClimaComms", "ClimaCore", "Dates", "NCDatasets", "SciMLBase"]
-git-tree-sha1 = "ace174fe5e4ae04c50a7835683baecd92e49c0d4"
+deps = ["Accessors", "ClimaComms", "ClimaCore", "Dates", "NCDatasets", "OrderedCollections", "SciMLBase"]
+git-tree-sha1 = "e9ac94af815dcae2a2ab24e54b53e76cca6258b7"
 uuid = "1ecacbb8-0713-4841-9a07-eb5aa8a2d53f"
-version = "0.2.10"
+version = "0.2.11"
 
 [[deps.ClimaParams]]
 deps = ["TOML"]

src/callbacks/get_callbacks.jl

@@ -33,6 +33,7 @@ function get_diagnostics(
         "average" => ((+), CAD.average_pre_output_hook!),
     )
 
+    dict_writer = CAD.DictWriter()
     hdf5_writer = CAD.HDF5Writer(output_dir)
 
     if !isnothing(parsed_args["netcdf_interpolation_num_points"])
@@ -61,11 +62,12 @@ function get_diagnostics(
         sync_schedule = CAD.EveryStepSchedule(),
         maybe_add_start_date...,
     )
-    writers = (hdf5_writer, netcdf_writer)
+    writers = (dict_writer, hdf5_writer, netcdf_writer)
 
-    # The default writer is HDF5
+    # The default writer is netcdf
     ALLOWED_WRITERS = Dict(
         "nothing" => netcdf_writer,
+        "dict" => dict_writer,
         "h5" => hdf5_writer,
         "hdf5" => hdf5_writer,
         "nc" => netcdf_writer,

src/diagnostics/Diagnostics.jl

@@ -40,9 +40,13 @@ import ..FriersonDiffusion
 import ..PrognosticEDMFX
 import ..DiagnosticEDMFX
 
+# surface_model
+import ..PrognosticSurfaceTemperature
+
 # functions used to calculate diagnostics
 import ..draft_area
 import ..compute_gm_mixing_length!
+import ..horizontal_integral_at_boundary
 
 # We need the abbreviations for symbols like curl, grad, and so on
 include(joinpath("..", "utils", "abbreviations.jl"))
@@ -58,6 +62,7 @@ import ClimaDiagnostics.Schedules:
     EveryStepSchedule, EveryDtSchedule, EveryCalendarDtSchedule
 
 import ClimaDiagnostics.Writers:
+    DictWriter,
     HDF5Writer,
     NetCDFWriter,
     write_field!,

src/diagnostics/conservation_diagnostics.jl

@@ -0,0 +1,144 @@
+# This file is included in Diagnostics.jl
+
+# Diagnostics for computing mass and energy conservation
+# All the diagnostics in this file are scalars, we need to
+# turn them into vectors as the diagnostics need to be mutable
+# Related issue: https://github.com/CliMA/ClimaDiagnostics.jl/issues/100
+
+###
+# Total mass of the air (scalar)
+###
+add_diagnostic_variable!(
+    short_name = "massa",
+    long_name = "Total Mass of the Air",
+    units = "kg",
+    compute! = (out, state, cache, time) -> begin
+        if isnothing(out)
+            return [sum(state.c.ρ)]
+        else
+            out .= [sum(state.c.ρ)]
+        end
+    end,
+)
+
+###
+# Total energy of air (scalar)
+###
+add_diagnostic_variable!(
+    short_name = "energya",
+    long_name = "Total Energy of the Air",
+    units = "J",
+    compute! = (out, state, cache, time) -> begin
+        if isnothing(out)
+            return [sum(state.c.ρe_tot)]
+        else
+            out .= [sum(state.c.ρe_tot)]
+        end
+    end,
+)
+
+###
+# Total water of the air (scalar)
+###
+compute_watera!(out, state, cache, time) =
+    compute_watera!(out, state, cache, time, cache.atmos.moisture_model)
+compute_watera!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("watera", model)
+
+function compute_watera!(
+    out,
+    state,
+    cache,
+    time,
+    moisture_model::T,
+) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
+    if isnothing(out)
+        return [sum(state.c.ρq_tot)]
+    else
+        out .= [sum(state.c.ρq_tot)]
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "watera",
+    long_name = "Total Water of the Air",
+    units = "kg",
+    compute! = compute_watera!,
+)
+
+###
+# Total energy of the slab ocean (scalar)
+###
+compute_energyo!(out, state, cache, time) =
+    compute_energyo!(out, state, cache, time, cache.atmos.surface_model)
+compute_energyo!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("energyo", model)
+
+function compute_energyo!(
+    out,
+    state,
+    cache,
+    time,
+    surface_model::T,
+) where {T <: PrognosticSurfaceTemperature}
+    sfc_cρh =
+        surface_model.ρ_ocean *
+        surface_model.cp_ocean *
+        surface_model.depth_ocean
+    if isnothing(out)
+        return [horizontal_integral_at_boundary(state.sfc.T .* sfc_cρh)]
+    else
+        out .= [horizontal_integral_at_boundary(state.sfc.T .* sfc_cρh)]
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "energyo",
+    long_name = "Total Energy of the Ocean",
+    units = "J",
+    compute! = compute_energyo!,
+)
+
+###
+# Total water of the slab ocean (scalar)
+###
+compute_watero!(out, state, cache, time) = compute_watero!(
+    out,
+    state,
+    cache,
+    time,
+    cache.atmos.moisture_model,
+    cache.atmos.surface_model,
+)
+compute_watero!(
+    _,
+    _,
+    _,
+    _,
+    moisture_model::T1,
+    surface_model::T2,
+) where {T1, T2} = error_diagnostic_variable(
+    "Can only compute total water of the ocean with a moist model and with PrognosticSurfaceTemperature",
+)
+
+function compute_watero!(
+    out,
+    state,
+    cache,
+    time,
+    moisture_model::Union{EquilMoistModel, NonEquilMoistModel},
+    surface_model::PrognosticSurfaceTemperature,
+)
+    if isnothing(out)
+        return [horizontal_integral_at_boundary(state.sfc.water)]
+    else
+        out .= [horizontal_integral_at_boundary(state.sfc.water)]
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "watero",
+    long_name = "Total Water of the Ocean",
+    units = "kg",
+    compute! = compute_watero!,
+)

src/diagnostics/diagnostic.jl

@@ -98,6 +98,7 @@ include("core_diagnostics.jl")
 include("radiation_diagnostics.jl")
 include("edmfx_diagnostics.jl")
 include("tracer_diagnostics.jl")
+include("conservation_diagnostics.jl")
 
 # Default diagnostics and higher level interfaces
 include("default_diagnostics.jl")

src/diagnostics/edmfx_diagnostics.jl

@@ -223,7 +223,7 @@ compute_husup!(
     moisture_model::T1,
     turbconv_model::T2,
 ) where {T1, T2} = error_diagnostic_variable(
-    "Can only compute updraft specific humidity and with a moist model and with EDMFX",
+    "Can only compute updraft specific humidity with a moist model and with EDMFX",
 )
 
 function compute_husup!(

src/solver/solve.jl

@@ -222,7 +222,6 @@ Return:
 - `energy_conservation = energy_net / energy_total`
 - `mass_conservation = (mass(t_end) - mass(t_0)) / mass(t_0)`
 - `water_conservation = (water_atmos + water_surface) / water_total`
-
 """
 function check_conservation(sol)
     # energy
@@ -266,3 +265,39 @@ function check_conservation(sol)
 
     return (; energy_conservation, mass_conservation, water_conservation)
 end
+
+function write_diagnostics_as_txt(simulation::AtmosSimulation)
+    foreach(
+        w -> write_diagnostics_as_txt(w, simulation.output_dir),
+        filter(w -> w isa CAD.DictWriter, simulation.output_writers),
+    )
+    return nothing
+end
+
+
+"""
+    write_diagnostics_as_txt(writer, output_dir)
+
+Write diagnostics in DictWriter to text files. This function is
+added because currently we do not support writing scalars to netcdf.
+It only supports diagnostics that are 1-element vectors.
+Related issue: https://github.com/CliMA/ClimaDiagnostics.jl/issues/100
+"""
+function write_diagnostics_as_txt(
+    writer::ClimaDiagnostics.Writers.DictWriter,
+    output_dir,
+)
+    @info "Writing diagnostics to text files"
+    for diagnostic in keys(writer.dict)
+        first(values(writer[diagnostic])) isa Vector ||
+            "write_diagnostics_as_txt is not supported for diagnostics that are not vectors"
+        filename = joinpath(output_dir, diagnostic * ".txt")
+        times = collect(keys(writer[diagnostic]))
+        values_all = getindex.(collect(values(writer[diagnostic])), 1)
+        open(filename, "w") do io
+            for (ti, vi) in zip(times, values_all)
+                println(io, "$ti $vi")
+            end
+        end
+    end
+end