From c524c17a1084b3bcf8995733afaeda1dee912a1c Mon Sep 17 00:00:00 2001
From: LenkaNovak <lenka@caltech.edu>
Date: Mon, 2 Oct 2023 16:11:25 -0700
Subject: [PATCH] clean

---
 .../AMIP/modular/user_io/debug_plots.jl       | 88 ++++++++++++-------
 1 file changed, 56 insertions(+), 32 deletions(-)

diff --git a/experiments/AMIP/modular/user_io/debug_plots.jl b/experiments/AMIP/modular/user_io/debug_plots.jl
index ae97f2e509..d96bffadb9 100644
--- a/experiments/AMIP/modular/user_io/debug_plots.jl
+++ b/experiments/AMIP/modular/user_io/debug_plots.jl
@@ -1,34 +1,72 @@
 using Plots
 using ClimaCorePlots
 using Printf
+using ClimaCoupler.Interfacer: ComponentModelSimulation, SurfaceModelSimulation
 
-# plotting functions for the coupler
+# 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)
     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)) )
+        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
 
-# plotting functions for each model
-function debug(sim, dir)
+"""
+    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)) )
+        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(boundary_space)
@@ -37,37 +75,23 @@ function get_field(sim::ClimaAtmosSimulation, ::Val{:w})
 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
-get_field(sim::ClimaAtmosSimulation, ::Val{:col_integrated_rain}) = sim.integrator.p.col_integrated_rain
-get_field(sim::ClimaAtmosSimulation, ::Val{:dif_flux_energy_bc}) = Geometry.WVector.(sim.integrator.p.sfc_conditions.ρ_flux_h_tot).components.data.:1
-get_field(sim::ClimaAtmosSimulation, ::Val{:dif_flux_ρq_tot_bc}) = Geometry.WVector.(sim.integrator.p.sfc_conditions.ρ_flux_q_tot).components.data.:1
-get_field(sim::ClimaAtmosSimulation, ::Val{:radiation_flux}) = ones(boundary_space)
 
 # 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
-get_field(sim::BucketSimulation, ::Val{:T}) = sim.integrator.u.bucket.T
-get_field(sim::BucketSimulation, ::Val{:P_snow}) = sim.integrator.p.bucket.P_snow
-get_field(sim::BucketSimulation, ::Val{:P_liq}) = sim.integrator.p.bucket.P_liq
 
-# chosen plot fields
+# currently selected plot fields
 plot_field_names(sim::SurfaceModelSimulation) = (:area_fraction, :surface_temperature, :surface_humidity, :air_density)
-plot_field_names(sim::BucketSimulation) = (:area_fraction, :surface_temperature, :surface_humidity, :air_density, :σS, :Ws, :W, :T, :P_snow, :P_liq)
-plot_field_names(sim::ClimaAtmosSimulation) = (:w, :ρq_tot, :ρe_tot, :col_integrated_rain, :dif_flux_energy_bc, :dif_flux_ρq_tot_bc, :radiation_flux)
-
-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
-
-using Printf
-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
+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,
+    :turbulent_energy_flux,
+    :turbulent_moisture_flux,
+    :radiative_energy_flux,
+)