Merge pull request #1154 from CliMA/gb/retoa

Sbozzolo · web-flow · commit bb1f2d8c0884 · 2025-04-23T11:31:15.000-07:00
Remove radiative_energy_toa
diff --git a/experiments/ClimaCore/test/coupled_sims.jl b/experiments/ClimaCore/test/coupled_sims.jl
@@ -1,6 +1,6 @@
 import Test: @test, @testset, @test_throws
 import ClimaComms
-@static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
+ClimaComms.@import_required_backends
 import ClimaCore as CC
 
 # Load file to test
diff --git a/experiments/ClimaEarth/setup_run.jl b/experiments/ClimaEarth/setup_run.jl
@@ -25,7 +25,7 @@ import DelimitedFiles
 # ## ClimaESM packages
 import ClimaAtmos as CA
 import ClimaComms
-@static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
+ClimaComms.@import_required_backends
 import ClimaCore as CC
 
 # ## Coupler specific imports
@@ -392,7 +392,7 @@ function CoupledSimulation(config_dict::AbstractDict)
         Interfacer.add_coupler_fields!(coupler_field_names, sim)
     end
     # add coupler fields required to track conservation, if specified
-    energy_check && push!(coupler_field_names, :radiative_energy_flux_toa, :P_net)
+    energy_check && push!(coupler_field_names, :P_net)
 
     # allocate space for the coupler fields
     coupler_fields = Interfacer.init_coupler_fields(FT, coupler_field_names, boundary_space)
diff --git a/experiments/ClimaEarth/test/component_model_tests/climaatmos_standalone/atmos_driver.jl b/experiments/ClimaEarth/test/component_model_tests/climaatmos_standalone/atmos_driver.jl
@@ -1,5 +1,5 @@
 import ClimaComms
-@static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
+ClimaComms.@import_required_backends
 import ArgParse
 import ClimaCoupler
 import ClimaCoupler: Utilities
diff --git a/experiments/ClimaEarth/test/component_model_tests/climaland_tests.jl b/experiments/ClimaEarth/test/component_model_tests/climaland_tests.jl
@@ -5,7 +5,7 @@ import ClimaCoupler
 import ClimaCoupler: FluxCalculator, Interfacer
 import Dates
 import ClimaComms
-@static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
+ClimaComms.@import_required_backends
 
 exp_dir = joinpath(pkgdir(ClimaCoupler), "experiments", "ClimaEarth")
 
diff --git a/experiments/ClimaEarth/test/debug_plots_tests.jl b/experiments/ClimaEarth/test/debug_plots_tests.jl
@@ -3,7 +3,7 @@ import Test: @test, @testset, @test_logs
 import ClimaCore as CC
 import ClimaCoupler: Interfacer
 import ClimaComms
-@static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
+ClimaComms.@import_required_backends
 
 # Prevent GKS headless operation mode warning
 ENV["GKSwstype"] = "nul"
@@ -55,7 +55,6 @@ plot_field_names(sim::Interfacer.SurfaceStub) = (:stub_field,)
         :beta,
         :z0b_sfc,
         :z0m_sfc,
-        :radiative_energy_flux_toa,
     ]
     atmos_names = (:atmos_field,)
     surface_names = (:surface_field,)
diff --git a/experiments/ClimaEarth/test/runtests.jl b/experiments/ClimaEarth/test/runtests.jl
@@ -1,6 +1,6 @@
 using SafeTestsets
 import ClimaComms
-@static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
+ClimaComms.@import_required_backends
 
 gpu_broken = ClimaComms.device() isa ClimaComms.CUDADevice
 
diff --git a/src/ConservationChecker.jl b/src/ConservationChecker.jl
@@ -6,6 +6,7 @@ This module contains functions that check global conservation of energy and wate
 module ConservationChecker
 
 import ..Interfacer, ..Utilities
+import ..Utilities: integral
 import ClimaCore as CC
 
 export AbstractConservationCheck, EnergyConservationCheck, WaterConservationCheck, check_conservation!
@@ -83,15 +84,13 @@ function check_conservation!(
     for sim in model_sims
         sim_name = Symbol(Interfacer.name(sim))
         if sim isa Interfacer.AtmosModelSimulation
-            radiative_energy_flux_toa = coupler_sim.fields.radiative_energy_flux_toa
-            # save radiation source
-            parent(radiative_energy_flux_toa) .= parent(Interfacer.get_field(sim, Val(:radiative_energy_flux_toa)))
+            radiative_energy_flux_toa = Interfacer.get_field(sim, Val(:radiative_energy_flux_toa))
 
             if isempty(ccs.toa_net_source)
-                radiation_sources_accum = sum(radiative_energy_flux_toa .* FT(float(coupler_sim.Δt_cpl))) # ∫ J / m^2 dA
+                radiation_sources_accum = integral(radiative_energy_flux_toa) * FT(float(coupler_sim.Δt_cpl)) # ∫ J / m^2 dA
             else
                 radiation_sources_accum =
-                    sum(radiative_energy_flux_toa .* FT(float(coupler_sim.Δt_cpl))) .+ ccs.toa_net_source[end] # ∫ J / m^2 dA
+                    integral(radiative_energy_flux_toa) * FT(float(coupler_sim.Δt_cpl)) + ccs.toa_net_source[end] # ∫ J / m^2 dA
             end
             push!(ccs.toa_net_source, radiation_sources_accum)
 
@@ -101,7 +100,6 @@ function check_conservation!(
 
             push!(previous, current)
             total += current + radiation_sources_accum
-
         elseif sim isa Interfacer.SurfaceModelSimulation
             # save surfaces
             area_fraction = Interfacer.get_field(sim, Val(:area_fraction))
@@ -110,7 +108,7 @@ function check_conservation!(
                 current = FT(0)
             else
                 previous = getproperty(ccs, sim_name)
-                current = sum(Interfacer.get_field(sim, Val(:energy)) .* area_fraction) # # ∫ J / m^3 dV
+                current = integral(Interfacer.get_field(sim, Val(:energy)) .* area_fraction) # # ∫ J / m^3 dV
             end
             push!(previous, current)
             total += current
@@ -120,7 +118,7 @@ function check_conservation!(
     push!(ccs.total, total)
 
     if runtime_check
-        @assert abs((total[end] - total[1]) / total[end]) < 1e-4
+        @assert abs((ccs.total[end] - ccs.total[1]) / ccs.total[end]) < 1e-4
     end
     return total
 
@@ -162,19 +160,19 @@ function check_conservation!(
 
             # save atmos
             previous = getproperty(ccs, sim_name)
-            current = sum(Interfacer.get_field(sim, Val(:water))) # kg (∫kg of water / m^3 dV)
+            current = integral(Interfacer.get_field(sim, Val(:water))) # kg (∫kg of water / m^3 dV)
             push!(previous, current)
 
         elseif sim isa Interfacer.SurfaceModelSimulation
             # save surfaces
             area_fraction = Interfacer.get_field(sim, Val(:area_fraction))
             if isnothing(Interfacer.get_field(sim, Val(:water)))
                 previous = getproperty(ccs, sim_name)
-                current = sum(PE_net .* area_fraction) # kg (∫kg of water / m^3 dV)
+                current = integral(PE_net .* area_fraction) # kg (∫kg of water / m^3 dV)
                 push!(previous, current)
             else
                 previous = getproperty(ccs, sim_name)
-                current = sum(Interfacer.get_field(sim, Val(:water)) .* area_fraction) # kg (∫kg of water / m^3 dV)
+                current = integral(Interfacer.get_field(sim, Val(:water)) .* area_fraction) # kg (∫kg of water / m^3 dV)
                 push!(previous, current)
             end
         end
@@ -185,7 +183,7 @@ function check_conservation!(
     push!(ccs.total, total)
 
     if runtime_check
-        @assert abs((total[end] - total[1]) / total[end]) < 1e-4
+        @assert abs((ccs.total[end] - ccs.total[1]) / ccs.total[end]) < 1e-4
     end
 
     return total
@@ -194,7 +192,7 @@ end
 """
     surface_water_gain_from_rates(cs::Interfacer.CoupledSimulation)
 
-Determines the total water content gain/loss of a surface from the begining of the simulation based on evaporation and precipitation rates.
+Determines the total water content gain/loss of a surface from the beginning of the simulation based on evaporation and precipitation rates.
 """
 function surface_water_gain_from_rates(cs::Interfacer.CoupledSimulation)
     evaporation = cs.fields.F_turb_moisture # kg / m^2 / s / layer depth
diff --git a/src/Interfacer.jl b/src/Interfacer.jl
@@ -190,7 +190,6 @@ get_field(
         Val{:height_sfc},
         Val{:liquid_precipitation},
         Val{:radiative_energy_flux_sfc},
-        Val{:radiative_energy_flux_toa},
         Val{:snow_precipitation},
         Val{:turblent_energy_flux},
         Val{:turbulent_moisture_flux},
diff --git a/src/Utilities.jl b/src/Utilities.jl
@@ -11,7 +11,7 @@ import ClimaCore as CC
 import Logging
 import ClimaUtilities.OutputPathGenerator: generate_output_path
 
-export swap_space!, get_device, get_comms_context, show_memory_usage, setup_output_dirs, time_to_seconds
+export swap_space!, get_device, get_comms_context, show_memory_usage, setup_output_dirs, time_to_seconds, integral
 
 """
     swap_space!(space_out::CC.Spaces.AbstractSpace, field_in::CC.Fields.Field)
@@ -188,4 +188,34 @@ function time_to_seconds(s::String)
 end
 
 
+"""
+    integral(field)
+
+Return the integral (a scalar) for `field` along its spatial dimensions.
+"""
+function integral(field::CC.Fields.Field)
+    if axes(field) isa CC.Spaces.SpectralElementSpace2D
+        # ClimaCore #1578, if field comes from Fields.level
+        # TODO: This should be fixed in ClimaCore
+        rad_grid = CC.Spaces.grid(axes(field))
+        if rad_grid isa CC.Grids.LevelGrid
+            if rad_grid.level isa CC.Utilities.PlusHalf
+                # FaceSpace
+                return sum(field ./ CC.Fields.Δz_field(field) .* 2)
+            else
+                # Center
+                return sum(field ./ CC.Fields.Δz_field(field))
+            end
+        else
+            return sum(field)
+        end
+    else
+        return sum(field)
+    end
+end
+
+function integral(field::AbstractArray)
+    return sum(field)
+end
+
 end # module
diff --git a/test/checkpointer_tests.jl b/test/checkpointer_tests.jl
@@ -1,6 +1,6 @@
 import Test: @test, @testset
 import ClimaComms
-@static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
+ClimaComms.@import_required_backends
 import ClimaCore as CC
 import ClimaCoupler: Checkpointer, Interfacer
 
diff --git a/test/conservation_checker_tests.jl b/test/conservation_checker_tests.jl
@@ -3,11 +3,10 @@
 =#
 import Test: @test, @testset
 import ClimaComms
-@static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
+ClimaComms.@import_required_backends
 import ClimaCore as CC
 import ClimaCoupler: ConservationChecker, Interfacer
-
-REGRID_DIR = @isdefined(REGRID_DIR) ? REGRID_DIR : joinpath("", "regrid_tmp/")
+import ClimaCoupler.Utilities: integral
 
 get_slab_energy(slab_sim, T_sfc) =
     slab_sim.integrator.p.params.ρ .* slab_sim.integrator.p.params.c .* T_sfc .* slab_sim.integrator.p.params.h
@@ -17,8 +16,8 @@ struct TestAtmos{I} <: Interfacer.AtmosModelSimulation
 end
 Interfacer.name(s::TestAtmos) = "TestAtmos"
 Interfacer.get_field(s::TestAtmos, ::Val{:radiative_energy_flux_toa}) = ones(s.i.space) .* 200
-Interfacer.get_field(s::TestAtmos, ::Val{:water}) = ones(s.i.space)
-Interfacer.get_field(s::TestAtmos, ::Val{:energy}) = ones(s.i.space) .* CC.Spaces.undertype(s.i.space)(1e6)
+Interfacer.get_field(s::TestAtmos, ::Val{:water}) = s.i.water
+Interfacer.get_field(s::TestAtmos, ::Val{:energy}) = s.i.energy
 
 struct TestOcean{I} <: Interfacer.SurfaceModelSimulation
     i::I
@@ -41,7 +40,10 @@ for FT in (Float32, Float64)
         space = CC.CommonSpaces.CubedSphereSpace(FT; radius = FT(6371e3), n_quad_points = 4, h_elem = 4)
 
         # set up model simulations
-        atmos = TestAtmos((; space = space))
+        initial_energy = ones(space) .* CC.Spaces.undertype(space)(1e6)
+        initial_water = ones(space) .* CC.Spaces.undertype(space)(1e5)
+
+        atmos = TestAtmos((; space = space, energy = initial_energy, water = initial_water))
         land = TestOcean((; space = space))
         ocean = TestLand((; space = space))
         ice = Interfacer.SurfaceStub((; area_fraction = CC.Fields.ones(space) .* FT(0.5)))
@@ -55,13 +57,11 @@ for FT in (Float32, Float64)
 
         # coupler fields
         cf = (;
-            radiative_energy_flux_toa = CC.Fields.ones(space),
             P_net = CC.Fields.zeros(space),
             P_liq = CC.Fields.zeros(space),
             P_snow = CC.Fields.zeros(space),
             F_turb_moisture = CC.Fields.zeros(space),
         )
-        @. cf.radiative_energy_flux_toa = 200
         @. cf.P_liq = -100
 
         # init
@@ -82,17 +82,32 @@ for FT in (Float32, Float64)
         )
 
         # set non-zero radiation and precipitation
-        F_r = cf.radiative_energy_flux_toa
+        F_r = 200 .* CC.Fields.ones(space)
         P = cf.P_liq
         Δt = float(cs.Δt_cpl)
 
+        volume = integral(ones(space))
+
+        area_fraction_scaling =
+            Interfacer.get_field(land, Val(:area_fraction)) .+ Interfacer.get_field(ocean, Val(:area_fraction))
+        water_from_precipitation = integral(P .* area_fraction_scaling) .* FT(Δt)
+        energy_from_radiation = integral(F_r) .* FT(Δt)
+        energy_per_unit_cell = CC.Fields.ones(space) .* energy_from_radiation ./ volume
+        water_per_unit_cell = CC.Fields.ones(space) .* water_from_precipitation ./ volume
+
         # analytical solution
-        tot_energy_an = sum(FT.(F_r .* 3Δt .+ 1e6 .* 1.25))
-        tot_water_an = sum(FT.(.-P .* 3Δt .* 0.5 .+ CC.Fields.ones(space)))
+        # Only ocean and atmos have energy
+        area_fraction_scaling = CC.Fields.ones(space) .+ Interfacer.get_field(ocean, Val(:area_fraction))
+        tot_energy_an = integral(area_fraction_scaling .* FT.(initial_energy)) + energy_from_radiation
+        tot_water_an = integral(FT.(initial_water)) - water_from_precipitation
 
         # run check_conservation!
         ConservationChecker.check_conservation!(cs, runtime_check = true)
+        atmos.i.energy .-= energy_per_unit_cell
+        atmos.i.water .+= water_per_unit_cell
         ConservationChecker.check_conservation!(cs, runtime_check = true)
+        atmos.i.energy .-= energy_per_unit_cell
+        atmos.i.water .+= water_per_unit_cell
         ConservationChecker.check_conservation!(cs, runtime_check = true)
 
         total_energy = cs.conservation_checks.energy.sums.total
@@ -120,7 +135,10 @@ for FT in (Float32, Float64)
         space = CC.CommonSpaces.CubedSphereSpace(FT; radius = FT(6371e3), n_quad_points = 4, h_elem = 4)
 
         # set up model simulations
-        atmos = TestAtmos((; space = space))
+        initial_energy = ones(space) .* CC.Spaces.undertype(space)(1e6)
+        initial_water = ones(space) .* CC.Spaces.undertype(space)(1e5)
+
+        atmos = TestAtmos((; space = space, energy = initial_energy, water = initial_water))
         land = TestOcean((; space = space))
         ocean = TestLand((; space = space))
         ice = Interfacer.SurfaceStub((; area_fraction = CC.Fields.ones(space) .* FT(0.5)))
@@ -134,13 +152,11 @@ for FT in (Float32, Float64)
 
         # coupler fields
         cf = (;
-            radiative_energy_flux_toa = CC.Fields.ones(space),
             P_net = CC.Fields.zeros(space),
             P_liq = CC.Fields.zeros(space),
             P_snow = CC.Fields.zeros(space),
             F_turb_moisture = CC.Fields.zeros(space),
         )
-        @. cf.radiative_energy_flux_toa = 200
         @. cf.P_liq = -100
 
         # init
diff --git a/test/interfacer_tests.jl b/test/interfacer_tests.jl
@@ -2,7 +2,7 @@ import Test: @test, @testset, @test_throws, @test_logs
 import ClimaCore as CC
 import ClimaParams as CP
 import ClimaComms
-@static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
+ClimaComms.@import_required_backends
 import Dates
 import Thermodynamics as TD
 import Thermodynamics.Parameters as TDP
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -1,6 +1,6 @@
 import SafeTestsets: @safetestset
 import ClimaComms
-@static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
+ClimaComms.@import_required_backends
 import Pkg, Artifacts
 
 gpu_broken = ClimaComms.device() isa ClimaComms.CUDADevice
diff --git a/test/utilities_tests.jl b/test/utilities_tests.jl
@@ -3,7 +3,7 @@
 =#
 import Test: @testset, @test
 import ClimaComms
-@static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
+ClimaComms.@import_required_backends
 import ClimaCoupler: Utilities
 import ClimaCore as CC
 
@@ -55,4 +55,24 @@ for FT in (Float32, Float64)
         @test typeof(Utilities.get_comms_context(parsed_args)) ==
               typeof(ClimaComms.context(ClimaComms.CPUSingleThreaded()))
     end
+
+    @testset "integral" begin
+        space2d = CC.CommonSpaces.CubedSphereSpace(FT; radius = FT(6371e3), n_quad_points = 4, h_elem = 4)
+        ones2d = ones(space2d)
+
+        space3d = CC.CommonSpaces.ExtrudedCubedSphereSpace(
+            FT;
+            z_elem = 10,
+            z_min = 0,
+            z_max = 1,
+            radius = FT(6371e3),
+            h_elem = 10,
+            n_quad_points = 4,
+            staggering = CC.CommonSpaces.CellCenter(),
+        )
+        ones3d_level = CC.Fields.level(ones(space3d), 1)
+
+        @test isapprox(Utilities.integral(ones3d_level), Utilities.integral(ones2d), rtol = 1e-5)
+        @test Utilities.integral(ones(space3d)) == sum(ones(space3d))
+    end
 end
