Skip to content

Commit

Permalink
init plots
Browse files Browse the repository at this point in the history
clean

clean

add unit test
  • Loading branch information
LenkaNovak committed Oct 3, 2023
1 parent a8a6c58 commit 7f0921f
Show file tree
Hide file tree
Showing 5 changed files with 172 additions and 1 deletion.
1 change: 1 addition & 0 deletions experiments/AMIP/modular/coupler_driver_modular.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down
88 changes: 88 additions & 0 deletions experiments/AMIP/modular/user_io/debug_plots.jl
Original file line number Diff line number Diff line change
@@ -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)
2 changes: 2 additions & 0 deletions test/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -17,6 +18,7 @@ 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"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
Expand Down
78 changes: 78 additions & 0 deletions test/debug/debug_amip_plots.jl
Original file line number Diff line number Diff line change
@@ -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
4 changes: 3 additions & 1 deletion test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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")

0 comments on commit 7f0921f

Please sign in to comment.