Js/postprocessor

Project.toml

@@ -4,8 +4,10 @@ authors = ["CliMA Contributors <clima-software@caltech.edu>"]
 version = "0.1.2"
 
 [deps]
+ClimaAnalysis = "29b5916a-a76c-4e73-9657-3c8fd22e65e6"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
+ClimaDiagnostics = "1ecacbb8-0713-4841-9a07-eb5aa8a2d53f"
 ClimaUtilities = "b3f4f4ca-9299-4f7f-bd9b-81e1242a7513"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
@@ -15,9 +17,19 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 SurfaceFluxes = "49b00bb7-8bd4-4f2b-b78c-51cd0450215f"
 Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
 
+[weakdeps]
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+ClimaCoreMakie = "908f55d8-4145-4867-9c14-5dad1a479e4d"
+Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+
+[extensions]
+ClimaCouplerPlotsExt = ["CairoMakie", "ClimaCoreMakie", "Makie"]
+
 [compat]
+ClimaAnalysis = "0.5.20"
 ClimaComms = "0.6.2"
 ClimaCore = "0.14.25"
+ClimaDiagnostics = "0.3"
 ClimaUtilities = "0.1.22"
 Dates = "1"
 JLD2 = "0.5, 0.6"

experiments/ClimaEarth/components/atmosphere/climaatmos.jl

@@ -734,3 +734,31 @@ function climaatmos_restart_path(output_dir_root, t)
     end
     error("Restart file for time $t not found")
 end
+
+###### Additional fields for debugging
+function Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:w})
+    w_c = ones(CC.Spaces.horizontal_space(sim.domain.face_space))
+    parent(w_c) .= parent(
+        CC.Fields.level(
+            CC.Geometry.WVector.(sim.integrator.u.f.u₃),
+            5 .+ CC.Utilities.half,
+        ),
+    )
+    return w_c
+end
+
+function specific_humidity(::CA.DryModel, integrator)
+    return [eltype(integrator.u)(0)]
+end
+function specific_humidity(::Union{CA.EquilMoistModel, CA.NonEquilMoistModel}, integrator)
+    return integrator.u.c.ρq_tot
+end
+
+Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:ρq_tot}) =
+    specific_humidity(sim.integrator.p.atmos.moisture_model, sim.integrator)
+
+Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:ρe_tot}) = sim.integrator.u.c.ρe_tot
+
+function plot_field_names(sim::ClimaAtmosSimulation)
+    return (:w, :ρq_tot, :ρe_tot, :liquid_precipitation, :snow_precipitation)
+end

experiments/ClimaEarth/components/land/climaland_bucket.jl

@@ -417,3 +417,12 @@ function Checkpointer.restore_cache!(sim::BucketSimulation, new_cache)
         ignore = Set([:rc, :params, :dss_buffer_2d, :dss_buffer_3d, :graph_context]),
     )
 end
+
+###### Additional fields for debugging
+Interfacer.get_field(sim::BucketSimulation, ::Val{:σS}) = sim.integrator.u.bucket.σS
+Interfacer.get_field(sim::BucketSimulation, ::Val{:Ws}) = sim.integrator.u.bucket.Ws
+Interfacer.get_field(sim::BucketSimulation, ::Val{:W}) = sim.integrator.u.bucket.W
+
+function plot_field_names(sim::BucketSimulation)
+    return (:area_fraction, :surface_temperature, :σS, :Ws, :W)
+end

experiments/ClimaEarth/components/land/climaland_integrated.jl

@@ -717,3 +717,37 @@ function Interfacer.set_cache!(sim::ClimaLandSimulation, csf)
     Interfacer.remap!(sim.model.snow.boundary_conditions.atmos.h, csf.height_delta)
     return nothing
 end
+
+###### Additional fields for debugging
+Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:soil_water}) =
+    sim.integrator.u.soil.ϑ_l
+Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:soil_ice}) = sim.integrator.u.soil.θ_i
+Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:soil_energy}) =
+    sim.integrator.u.soil.ρe_int
+Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:canopy_temp}) =
+    sim.integrator.u.canopy.energy.T
+Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:canopy_water}) =
+    sim.integrator.u.canopy.hydraulics.ϑ_l.:1
+Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:snow_energy}) =
+    sim.integrator.u.snow.U
+Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:snow_water_equiv}) =
+    sim.integrator.u.snow.S
+Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:snow_liquid_water}) =
+    sim.integrator.u.snow.S_l
+
+function plot_field_names(sim::ClimaLandSimulation)
+    return (
+        :area_fraction,
+        :surface_direct_albedo,
+        :surface_diffuse_albedo,
+        :surface_temperature,
+        :soil_water,
+        :soil_ice,
+        :soil_energy,
+        :canopy_temp,
+        :canopy_water,
+        :snow_energy,
+        :snow_water_equiv,
+        :snow_liquid_water,
+    )
+end

experiments/ClimaEarth/setup_run.jl

@@ -39,7 +39,9 @@ import ClimaCoupler:
     FluxCalculator,
     Interfacer,
     TimeManager,
-    Utilities
+    Utilities,
+    Postprocessor
+import ClimaCoupler: ClimaCouplerPlotsExt
 import ClimaCoupler.Interfacer:
     AbstractSlabplanetSimulationMode,
     AMIPMode,
@@ -91,7 +93,6 @@ We can additionally pass the configuration dictionary to the component model ini
 include("cli_options.jl")
 include("user_io/arg_parsing.jl")
 include("user_io/postprocessing.jl")
-include("user_io/coupler_diagnostics.jl")
 
 """
     CoupledSimulation(config_file)
@@ -510,7 +511,7 @@ function CoupledSimulation(config_dict::AbstractDict)
     =#
     if use_coupler_diagnostics
         @info "Using default coupler diagnostics"
-        diags_handler = coupler_diagnostics_setup(
+        diags_handler = Postprocessor.diagnostics_setup(
             coupler_fields,
             dir_paths.coupler_output_dir,
             start_date,

experiments/ClimaEarth/test/debug_plots_tests.jl

@@ -23,7 +23,9 @@ struct ClimaLandSimulation{C} <: Interfacer.SurfaceModelSimulation
     cache::C
 end
 
-include("../user_io/debug_plots.jl")
+# Import the extension module (Makie must be loaded first to trigger the extension)
+import Makie
+import ClimaCoupler: ClimaCouplerPlotsExt
 
 Interfacer.get_field(sim::BucketSimulation, ::Val{:surface_field}) = sim.cache.surface_field
 Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:surface_field}) =
@@ -96,7 +98,7 @@ plot_field_names(sim::Interfacer.SurfaceStub) = (:stub_field,)
     )
 
     output_plots = "test_debug"
-    @test_logs (:info, "plotting debug in test_debug") match_mode = :any debug(
+    @test_logs (:info, "plotting debug in test_debug") match_mode = :any ClimaCouplerPlotsExt.debug(
         cs,
         output_plots,
     )

experiments/ClimaEarth/user_io/postprocessing.jl

@@ -1,6 +1,4 @@
 ## helpers for user-specified IO
-include("debug_plots.jl")
-include("diagnostics_plots.jl")
 include("../leaderboard/leaderboard.jl")
 include("../leaderboard/test_rmses.jl")
 
@@ -24,13 +22,30 @@ function postprocess_sim(cs, postprocessing_vars)
 
     # Plot generic diagnostics
     @info "Plotting diagnostics for coupler, atmos, land, and ocean"
-    make_diagnostics_plots(coupler_output_dir, artifacts_dir, output_prefix = "coupler_")
-    make_diagnostics_plots(atmos_output_dir, artifacts_dir, output_prefix = "atmos_")
-    make_diagnostics_plots(land_output_dir, artifacts_dir, output_prefix = "land_")
-    make_ocean_diagnostics_plots(ocean_output_dir, artifacts_dir, output_prefix = "ocean_")
+    ClimaCouplerPlotsExt.make_diagnostics_plots(
+        coupler_output_dir,
+        artifacts_dir,
+        output_prefix = "coupler_",
+    )
+    ClimaCouplerPlotsExt.make_diagnostics_plots(
+        atmos_output_dir,
+        artifacts_dir,
+        output_prefix = "atmos_",
+    )
+    ClimaCouplerPlotsExt.make_diagnostics_plots(
+        land_output_dir,
+        artifacts_dir,
+        output_prefix = "land_",
+    )
+    ClimaCouplerPlotsExt.make_ocean_diagnostics_plots(
+        ocean_output_dir,
+        artifacts_dir,
+        output_prefix = "ocean_",
+    )
 
     # Plot all model states and coupler fields (useful for debugging)
-    ClimaComms.context(cs) isa ClimaComms.SingletonCommsContext && debug(cs, artifacts_dir)
+    ClimaComms.context(cs) isa ClimaComms.SingletonCommsContext &&
+        ClimaCouplerPlotsExt.debug(cs, artifacts_dir)
 
     # If we have enough data (in time, but also enough variables), plot the leaderboard.
     # We need pressure to compute the leaderboard.
@@ -51,14 +66,14 @@ function postprocess_sim(cs, postprocessing_vars)
     # Perform conservation checks if they exist
     if !isnothing(cs.conservation_checks)
         @info "Conservation Check Plots"
-        plot_global_conservation(
+        ConservationChecker.plot_global_conservation(
             cs.conservation_checks.energy,
             cs,
             conservation_softfail,
             figname1 = joinpath(artifacts_dir, "total_energy_bucket.png"),
             figname2 = joinpath(artifacts_dir, "total_energy_log_bucket.png"),
         )
-        plot_global_conservation(
+        ConservationChecker.plot_global_conservation(
             cs.conservation_checks.water,
             cs,
             conservation_softfail,

ext/ClimaCouplerPlotsExt.jl

@@ -0,0 +1,17 @@
+"""
+    ClimaCouplerPlotsExt
+
+Extension module for ClimaCoupler that provides plotting functionality when Makie is available.
+This extension loads the plotting submodules into the Postprocessor module.
+"""
+module ClimaCouplerPlotsExt
+
+using ..ClimaCoupler
+using Makie
+using CairoMakie
+using ClimaCoreMakie
+
+include("ClimaCouplerPlotsExt/diagnostics_plots.jl")
+include("ClimaCouplerPlotsExt/debug_plots.jl")
+
+end # module ClimaCouplerPlotsExt

experiments/ClimaEarth/user_io/debug_plots.jl → ext/ClimaCouplerPlotsExt/debug_plots.jl

@@ -8,89 +8,8 @@ import ClimaAtmos as CA
 import Oceananigans as OC
 import StaticArrays
 
-"""
-    plot_global_conservation(
-        cc::AbstractConservationCheck,
-        coupler_sim::Interfacer.CoupledSimulation,
-        softfail = false;
-        figname1 = "total.png",
-        figname2 = "total_log.png",
-    )
+export debug
 
-Creates two plots of the globally integrated quantity (energy, ``\\rho e``):
-1. global quantity of each model component as a function of time,
-relative to the initial value;
-2. fractional change in the sum of all components over time on a log scale.
-"""
-function plot_global_conservation(
-    cc::ConservationChecker.AbstractConservationCheck,
-    coupler_sim::Interfacer.CoupledSimulation,
-    softfail = false;
-    figname1 = "total.png",
-    figname2 = "total_log.png",
-)
-
-    model_sims = coupler_sim.model_sims
-    ccs = cc.sums
-
-    days = collect(1:length(ccs[1])) * float(coupler_sim.Δt_cpl) / 86400
-
-    # evolution of energy of each component relative to initial value
-    total = ccs.total  # total
-
-    var_name = nameof(cc)
-    cum_total = [0.0]
-    f = Makie.Figure()
-    ax = Makie.Axis(f[1, 1], xlabel = "time [days]", ylabel = "$var_name: (t) - (t=0)")
-    Makie.lines!(ax, days, total .- total[1], label = "total"; linewidth = 3)
-    for sim in model_sims
-        sim_name = nameof(sim)
-        global_field = getproperty(ccs, Symbol(sim_name))
-        diff_global_field = (global_field .- global_field[1])
-        Makie.lines!(ax, days, diff_global_field[1:length(days)], label = sim_name)
-        cum_total .+= abs.(global_field[end])
-    end
-    if cc isa ConservationChecker.EnergyConservationCheck
-        global_field = ccs.toa_net_source
-        diff_global_field = (global_field .- global_field[1])
-        Makie.lines!(ax, days, diff_global_field[1:length(days)], label = "toa_net")
-        cum_total .+= abs.(global_field[end])
-    end
-    Makie.axislegend(ax, position = :lb)
-    Makie.save(figname1, f)
-
-    # use the cumulative global sum at the final time step as a reference for the error calculation
-    rse = abs.((total .- total[1]) ./ cum_total)
-    l_rse = log.(rse)
-    # evolution of log error of total
-    lp = Makie.lines(days, l_rse, label = "rs error")
-    lp.axis.xlabel = "time [days]"
-    lp.axis.ylabel = "log( |x(t) - x(t=0)| / Σx(t=T) )"
-    l_rse_valid = filter(x -> !isinf(x) && !isnan(x), l_rse)
-    if !isempty(l_rse_valid)
-        y_min = minimum(l_rse_valid)
-        y_max = maximum(l_rse_valid)
-        if y_min != y_max
-            Makie.ylims!(y_min, y_max)
-        else
-            # If all values are the same, add a small padding to avoid Makie error
-            padding = max(abs(y_min) * 0.01, 0.1)
-            Makie.ylims!(y_min - padding, y_max + padding)
-        end
-    end
-    Makie.axislegend(position = :lt)
-    Makie.save(figname2, lp)
-
-    # check that the relative error is small (TODO: reduce this to sqrt(eps(FT)))
-    if !softfail
-        @info typeof(cc)
-        @info rse[end]
-        @assert rse[end] < 0.035
-    end
-end
-
-
-# plotting functions for the coupled simulation
 """
     debug(cs::Interfacer.CoupledSimulation, dir = "debug", cs_fields_ref = nothing)
 
@@ -250,66 +169,6 @@ function print_extrema(field::OC.Field)
     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 Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:w})
-    w_c = ones(CC.Spaces.horizontal_space(sim.domain.face_space))
-    parent(w_c) .= parent(
-        CC.Fields.level(
-            CC.Geometry.WVector.(sim.integrator.u.f.u₃),
-            5 .+ CC.Utilities.half,
-        ),
-    )
-    return w_c
-end
-specific_humidity(::CA.DryModel, integrator) = [eltype(integrator.u)(0)]
-specific_humidity(::Union{CA.EquilMoistModel, CA.NonEquilMoistModel}, integrator) =
-    integrator.u.c.ρq_tot
-Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:ρq_tot}) =
-    specific_humidity(sim.integrator.p.atmos.moisture_model, sim.integrator)
-Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:ρe_tot}) = sim.integrator.u.c.ρe_tot
-
-# additional BucketSimulation debug fields
-Interfacer.get_field(sim::BucketSimulation, ::Val{:σS}) = sim.integrator.u.bucket.σS
-Interfacer.get_field(sim::BucketSimulation, ::Val{:Ws}) = sim.integrator.u.bucket.Ws
-Interfacer.get_field(sim::BucketSimulation, ::Val{:W}) = sim.integrator.u.bucket.W
-
-# additional ClimaLand model debug fields
-Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:soil_water}) =
-    sim.integrator.u.soil.ϑ_l
-Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:soil_ice}) = sim.integrator.u.soil.θ_i
-Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:soil_energy}) =
-    sim.integrator.u.soil.ρe_int
-Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:canopy_temp}) =
-    sim.integrator.u.canopy.energy.T
-Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:canopy_water}) =
-    sim.integrator.u.canopy.hydraulics.ϑ_l.:1
-Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:snow_energy}) =
-    sim.integrator.u.snow.U
-Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:snow_water_equiv}) =
-    sim.integrator.u.snow.S
-Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:snow_liquid_water}) =
-    sim.integrator.u.snow.S_l
-
 # currently selected plot fields
 plot_field_names(sim::Interfacer.SurfaceModelSimulation) =
     (:area_fraction, :surface_temperature)
-plot_field_names(sim::ClimaLandSimulation) = (
-    :area_fraction,
-    :surface_direct_albedo,
-    :surface_diffuse_albedo,
-    :surface_temperature,
-    :soil_water,
-    :soil_ice,
-    :soil_energy,
-    :canopy_temp,
-    :canopy_water,
-    :snow_energy,
-    :snow_water_equiv,
-    :snow_liquid_water,
-)
-plot_field_names(sim::BucketSimulation) =
-    (:area_fraction, :surface_temperature, :σS, :Ws, :W)
-plot_field_names(sim::ClimaAtmosSimulation) =
-    (:w, :ρq_tot, :ρe_tot, :liquid_precipitation, :snow_precipitation)