Merge #452

452: Add AMIP debug plots r=LenkaNovak a=LenkaNovak



Co-authored-by: LenkaNovak <lenka@caltech.edu>

experiments/AMIP/modular/coupler_driver_modular.jl

@@ -125,6 +125,7 @@ parsed_args = parse_commandline(argparse_settings())
 # modify parsed args for fast testing from REPL #hide
 pkg_dir = pkgdir(ClimaCoupler)
 if isinteractive()
+    include("user_io/debug_plots.jl")
     parsed_args["config_file"] =
         isnothing(parsed_args["config_file"]) ? joinpath(pkg_dir, "config/model_configs/interactive_debug.yml") :
         parsed_args["config_file"]

experiments/AMIP/modular/user_io/debug_plots.jl

@@ -0,0 +1,88 @@
+using Plots
+using ClimaCorePlots
+using Printf
+using ClimaCoupler.Interfacer: ComponentModelSimulation, SurfaceModelSimulation
+
+# plotting functions for the coupled simulation
+"""
+    debug(cs::CoupledSimulation, dir = "debug")
+
+Plot the fields of a coupled simulation and save plots to a directory.
+"""
+function debug(cs::CoupledSimulation, dir = "debug")
+    mkpath(dir)
+    @info "plotting debug in " * dir
+    for sim in cs.model_sims
+        debug(sim, dir)
+    end
+    debug(cs.fields, dir)
+end
+
+"""
+    debug(cs_fields::NamedTuple, dir)
+
+Plot useful coupler fields (in `field_names`) and save plots to a directory.
+"""
+function debug(cs_fields::NamedTuple, dir)
+    field_names = (:F_turb_energy, :F_turb_moisture, :P_liq, :T_S, :ρ_sfc, :q_sfc)
+    all_plots = []
+    for field_name in field_names
+        field = getproperty(cs_fields, field_name)
+        push!(all_plots, Plots.plot(field, title = string(field_name) * print_extrema(field)))
+    end
+    fig = Plots.plot(all_plots..., size = (1500, 800))
+    Plots.png(joinpath(dir, "debug_coupler"))
+end
+
+"""
+    debug(sim::ComponentModelSimulation, dir)
+
+Plot the fields of a component model simulation and save plots to a directory.
+"""
+function debug(sim::ComponentModelSimulation, dir)
+
+    field_names = plot_field_names(sim)
+
+    all_plots = []
+    for field_name in field_names
+        field = get_field(sim, Val(field_name))
+        push!(all_plots, Plots.plot(field, title = string(field_name) * print_extrema(field)))
+    end
+    fig = Plots.plot(all_plots..., size = (1500, 800))
+    Plots.png(joinpath(dir, "debug_$(name(sim))"))
+
+end
+
+"""
+    print_extrema(field::ClimaCore.Fields.Field)
+
+Return the minimum and maximum values of a field as a string.
+"""
+function print_extrema(field::ClimaCore.Fields.Field)
+    ext_vals = extrema(field)
+    min = @sprintf("%.2E", ext_vals[1])
+    max = @sprintf("%.2E", ext_vals[2])
+    return " [$min, $max]"
+end
+
+# below are additional fields specific to this experiment (ourside of the required coupler fields) that we are interested in plotting for debugging purposes
+
+# additional ClimaAtmos model debug fields
+function get_field(sim::ClimaAtmosSimulation, ::Val{:w})
+    w_c = ones(sim.domain.face_space.horizontal_space)
+    parent(w_c) .= parent(Fields.level(Geometry.WVector.(sim.integrator.u.f.u₃), 5 .+ half))
+    return w_c
+end
+get_field(sim::ClimaAtmosSimulation, ::Val{:ρq_tot}) = sim.integrator.u.c.ρq_tot
+get_field(sim::ClimaAtmosSimulation, ::Val{:ρe_tot}) = sim.integrator.u.c.ρe_tot
+
+# additional BucketSimulation debug fields
+get_field(sim::BucketSimulation, ::Val{:σS}) = sim.integrator.u.bucket.σS
+get_field(sim::BucketSimulation, ::Val{:Ws}) = sim.integrator.u.bucket.Ws
+get_field(sim::BucketSimulation, ::Val{:W}) = sim.integrator.u.bucket.W
+
+# currently selected plot fields
+plot_field_names(sim::SurfaceModelSimulation) = (:area_fraction, :surface_temperature, :surface_humidity)
+plot_field_names(sim::BucketSimulation) =
+    (:area_fraction, :surface_temperature, :surface_humidity, :air_density, :σS, :Ws, :W)
+plot_field_names(sim::ClimaAtmosSimulation) = (:w, :ρq_tot, :ρe_tot, :liquid_precipitation, :snow_precipitation)

test/Project.toml

@@ -5,6 +5,7 @@ CLIMAParameters = "6eacf6c3-8458-43b9-ae03-caf5306d3d53"
 ClimaAtmos = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
+ClimaCorePlots = "cf7c7e5a-b407-4c48-9047-11a94a308626"
 ClimaCoupler = "4ade58fe-a8da-486c-bd89-46df092ec0c7"
 ClimaLSM = "7884a58f-fab6-4fd0-82bb-ecfedb2d8430"
 ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
@@ -17,6 +18,8 @@ MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

test/debug/debug_amip_plots.jl

@@ -0,0 +1,78 @@
+# testing functions used to produce user-defined debugging plots for AMIP experiments
+
+using Test
+using ClimaCore
+using ClimaCoupler: TestHelper
+import ClimaCoupler.Interfacer:
+    update_field!, AtmosModelSimulation, SurfaceModelSimulation, SurfaceStub, get_field, update_field!, name
+using ClimaCoupler.Utilities: CoupledSimulation, CoupledSimulation
+
+FT = Float64
+
+struct ClimaAtmosSimulation{C} <: AtmosModelSimulation
+    cache::C
+end
+name(sim::ClimaAtmosSimulation) = "ClimaAtmosSimulation"
+get_field(sim::AtmosModelSimulation, ::Val{:atmos_field}) = sim.cache.atmos_field
+
+struct BucketSimulation{C} <: SurfaceModelSimulation
+    cache::C
+end
+name(sim::BucketSimulation) = "BucketSimulation"
+
+include("../../experiments/AMIP/modular/user_io/debug_plots.jl")
+
+get_field(sim::BucketSimulation, ::Val{:surface_field}) = sim.cache.surface_field
+get_field(sim::SurfaceStub, ::Val{:stub_field}) = sim.cache.stub_field
+
+plot_field_names(sim::ClimaAtmosSimulation) = (:atmos_field,)
+plot_field_names(sim::BucketSimulation) = (:surface_field,)
+plot_field_names(sim::SurfaceStub) = (:stub_field,)
+
+@testset "import_atmos_fields!" begin
+
+    boundary_space = TestHelper.create_space(FT)
+    coupler_names = (:F_turb_energy, :F_turb_moisture, :P_liq, :T_S, :ρ_sfc, :q_sfc)
+    atmos_names = (:atmos_field,)
+    surface_names = (:surface_field,)
+    stub_names = (:stub_field,)
+
+    atmos_fields = NamedTuple{atmos_names}(ntuple(i -> ClimaCore.Fields.ones(boundary_space), length(atmos_names)))
+    surface_fields =
+        NamedTuple{surface_names}(ntuple(i -> ClimaCore.Fields.ones(boundary_space), length(surface_names)))
+    stub_fields = NamedTuple{stub_names}(ntuple(i -> ClimaCore.Fields.ones(boundary_space), length(stub_names)))
+    coupler_fields =
+        NamedTuple{coupler_names}(ntuple(i -> ClimaCore.Fields.zeros(boundary_space), length(coupler_names)))
+
+    model_sims = (;
+        atmos_sim = ClimaAtmosSimulation(atmos_fields),
+        surface_sim = BucketSimulation(surface_fields),
+        ice_sim = SurfaceStub(stub_fields),
+    )
+    cs = CoupledSimulation{FT}(
+        nothing, # comms_ctx
+        nothing, # dates
+        nothing, # boundary_space
+        coupler_fields, # fields
+        nothing, # parsed_args
+        nothing, # conservation_checks
+        (Int(0), Int(1)), # tspan
+        Int(200), # t
+        Int(200), # Δt_cpl
+        (;), # surface_masks
+        model_sims, # model_sims
+        (;), # mode
+        (), # diagnostics
+    )
+
+    output_plots = "test_debug"
+    debug(cs, output_plots)
+    @test isfile("test_debug/debug_ClimaAtmosSimulation.png")
+    @test isfile("test_debug/debug_BucketSimulation.png")
+    @test isfile("test_debug/debug_SurfaceStub.png")
+    @test isfile("test_debug/debug_coupler.png")
+
+    # remove output
+    rm(output_plots; recursive = true)
+
+end

test/runtests.jl

@@ -57,6 +57,8 @@ end
 @safetestset "component model test: slab ocean" begin
     include("component_model_tests/slab_ocean_tests.jl")
 end
-
+@safetestset "debug diagnostics: amip plots" begin
+    include("debug/debug_amip_plots.jl")
+end
 
 # include("CoupledSimulations/cplsolver.jl")