WIP

experiments/AMIP/modular/components/atmosphere/climaatmos_init.jl

@@ -234,3 +234,41 @@ Extension of Checkpointer.get_model_state_vector to get the model state.
 function get_model_state_vector(sim::ClimaAtmosSimulation)
     return sim.integrator.u
 end
+
+
+function get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:F_radiative_TOA})
+    radiation = atmos_sim.integrator.p.radiation_model
+    FT = eltype(atmos_sim.integrator.u)
+    # save radiation source
+    if radiation != nothing
+        face_space = axes(atmos_sim.integrator.u.f)
+        z = parent(Fields.coordinate_field(face_space).z)
+        Δz_top = round(FT(0.5) * (z[end, 1, 1, 1, 1] - z[end - 1, 1, 1, 1, 1]))
+        n_faces = length(z[:, 1, 1, 1, 1])
+
+        LWd_TOA = Fields.level(
+            RRTMGPI.array2field(FT.(atmos_sim.integrator.p.radiation_model.face_lw_flux_dn), face_space),
+            n_faces - half,
+        )
+        LWu_TOA = Fields.level(
+            RRTMGPI.array2field(FT.(atmos_sim.integrator.p.radiation_model.face_lw_flux_up), face_space),
+            n_faces - half,
+        )
+        SWd_TOA = Fields.level(
+            RRTMGPI.array2field(FT.(atmos_sim.integrator.p.radiation_model.face_sw_flux_dn), face_space),
+            n_faces - half,
+        )
+        SWu_TOA = Fields.level(
+            RRTMGPI.array2field(FT.(atmos_sim.integrator.p.radiation_model.face_sw_flux_up), face_space),
+            n_faces - half,
+        )
+
+        return LWd_TOA .+ SWd_TOA .- LWu_TOA .- SWu_TOA # [W/m^2]
+    else
+        return FT(0)
+    end
+end
+
+get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:energy}) = sum(atmos_sim.integrator.u.c.ρe_tot)
+
+get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:water}) = sum(atmos_sim.integrator.u.c.ρq_tot)

experiments/AMIP/modular/components/land/bucket_utils.jl

@@ -4,7 +4,7 @@
 
 Returns the volumetric internal energy of the bucket land model.
 """
-function get_land_energy(bucket_sim::BucketSimulation, e_per_area)
+function get_field(bucket_sim::BucketSimulation, ::Val{:energy})
     # required by ConservationChecker
     e_per_area .= zeros(axes(bucket_sim.integrator.u.bucket.W))
     soil_depth = FT = eltype(bucket_sim.integrator.u.bucket.W)
@@ -20,14 +20,20 @@ function get_land_energy(bucket_sim::BucketSimulation, e_per_area)
     return e_per_area
 end
 
+function get_field(bucket_sim::BucketSimulation, ::Val{:water})
+    ρ_cloud_liq = ClimaLSM.LSMP.ρ_cloud_liq(bucket_sim.model.parameters.earth_param_set)
+    return
+        @. (bucket_sim.integrator.u.bucket.σS + bucket_sim.integrator.u.bucket.W + bucket_sim.integrator.u.bucket.Ws) * ρ_cloud_liq  # kg water / m2
+end
+
 """
-    get_land_temp_from_state(land_sim, u)
+    get_land_temp_from_state(bucket_sim, u)
 Returns the surface temperature of the earth, computed
 from the state u.
 """
-function get_land_temp_from_state(land_sim, u)
+function get_land_temp_from_state(bucket_sim, u)
     # required by viz_explorer.jl
-    return ClimaLSM.surface_temperature(land_sim.model, u, land_sim.integrator.p, land_sim.integrator.t)
+    return ClimaLSM.surface_temperature(bucket_sim.model, u, bucket_sim.integrator.p, bucket_sim.integrator.t)
 end
 
 """

experiments/AMIP/modular/components/ocean/slab_ocean_init.jl

@@ -157,3 +157,6 @@ Extension of Checkpointer.get_model_state_vector to get the model state.
 function get_model_state_vector(sim::SlabOceanSimulation)
     return sim.integrator.u
 end
+
+get_field(sim::SlabOceanSimulation, ::Val{:energy}) =
+    sim.integrator.p.params.ρ .* sim.integrator.p.params.c .* sim.integrator.u.T_sfc .* sim.integrator.p.params.h

experiments/AMIP/modular/components/ocean/slab_seaice_init.jl

@@ -167,3 +167,6 @@ Extension of Checkpointer.get_model_state_vector to get the model state.
 function get_model_state_vector(sim::PrescribedIceSimulation)
     return sim.integrator.u
 end
+
+get_field(sim::PrescribedIceSimulation, ::Val{:energy}) =
+    sim.integrator.p.params.ρ .* sim.integrator.p.params.c .* sim.integrator.u.T_sfc .* sim.integrator.p.params.h

src/ClimaCoupler.jl

@@ -13,7 +13,7 @@ include("../test/TestHelper.jl")
 include("Utilities.jl")
 include("TimeManager.jl")
 include("Regridder.jl")
-include("ConservationChecker.jl")
+# include("ConservationChecker.jl")
 include("BCReader.jl")
 include("FluxCalculator.jl")
 include("FieldExchanger.jl")

src/ConservationChecker.jl

@@ -295,10 +295,10 @@ function plot_global_conservation(
     figname2 = "total_water_log.png",
 ) where {FT}
     times = collect(1:length(cc.ρq_tot_atmos)) * coupler_sim.Δt_cpl
-    diff_ρe_tot_atmos = (cc.ρq_tot_atmos .- cc.ρq_tot_atmos[2])
-    diff_ρe_tot_slab = (cc.ρq_tot_land .- cc.ρq_tot_land[2])
-    diff_ρe_tot_slab_seaice = (cc.ρq_tot_seaice .- cc.ρq_tot_seaice[2])
-    diff_ρe_tot_slab_ocean = (cc.ρq_tot_ocean .- cc.ρq_tot_ocean[2])
+    diff_ρe_tot_atmos = (cc.ρq_tot_atmos .- cc.ρq_tot_atmos[1])
+    diff_ρe_tot_slab = (cc.ρq_tot_land .- cc.ρq_tot_land[1])
+    diff_ρe_tot_slab_seaice = (cc.ρq_tot_seaice .- cc.ρq_tot_seaice[1])
+    diff_ρe_tot_slab_ocean = (cc.ρq_tot_ocean .- cc.ρq_tot_ocean[1])
 
     times_days = times ./ (24 * 60 * 60)
     Plots.plot(times_days, diff_ρe_tot_atmos[1:length(times_days)], label = "atmos")