Fix ConversationChecker

`sum` behaves differently depending on the space, so I added `integral`.

In addition to this, the test was not correctly accounting for the
movement of water and energy.

Author: Sbozzolo

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!
@@ -85,26 +86,11 @@ function check_conservation!(
         if sim isa Interfacer.AtmosModelSimulation
             radiative_energy_flux_toa = Interfacer.get_field(sim, Val(:radiative_energy_flux_toa))
 
-            # ClimaCore #1578, if radiative_energy_flux_toa comes from Fields.level
-            # TODO: Fix this in ClimaCore
-            rad_grid = CC.Spaces.grid(axes(radiative_energy_flux_toa))
-            if rad_grid isa CC.Grids.LevelGrid
-                if rad_grid.level isa CC.Utilities.PlusHalf
-                    # FaceSpace
-                    surface_integral =
-                        sum(radiative_energy_flux_toa ./ CC.Fields.Δz_field(radiative_energy_flux_toa) .* 2)
-                else
-                    # Center
-                    surface_integral = sum(radiative_energy_flux_toa ./ CC.Fields.Δz_field(radiative_energy_flux_toa))
-                end
-            else
-                surface_integral = sum(radiative_energy_flux_toa)
-            end
-
             if isempty(ccs.toa_net_source)
-                radiation_sources_accum = surface_integral * 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 = surface_integral * FT(float(coupler_sim.Δt_cpl)) + ccs.toa_net_source[end] # ∫ J / m^2 dA
+                radiation_sources_accum =
+                    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)
 
@@ -122,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
@@ -174,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
@@ -206,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

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

test/conservation_checker_tests.jl

@@ -6,8 +6,7 @@ import ClimaComms
 @static pkgversion(ClimaComms) >= v"0.6" && 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)))
@@ -84,13 +86,28 @@ for FT in (Float32, Float64)
         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
@@ -118,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)))

test/utilities_tests.jl

@@ -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