Merge pull request #3123 from CliMA/gb/check_conservation

Move conservation check to src

examples/hybrid/driver.jl

@@ -116,55 +116,24 @@ end
 
 # Conservation checks
 if config.parsed_args["check_conservation"]
-    @info "Checking conservation"
     FT = Spaces.undertype(axes(sol.u[end].c.ρ))
+    @info "Checking conservation"
+    (; energy_conservation, mass_conservation, water_conservation) =
+        CA.check_conservation(sol)
+
+    @info "    Net energy change / total energy: $energy_conservation"
+    @info "    Net mass change / total mass: $mass_conservation"
 
-    # energy
-    energy_total = sum(sol.u[end].c.ρe_tot)
-    energy_atmos_change = sum(sol.u[end].c.ρe_tot) - sum(sol.u[1].c.ρe_tot)
     sfc = p.atmos.surface_model
-    if sfc isa CA.PrognosticSurfaceTemperature
-        sfc_cρh = sfc.ρ_ocean * sfc.cp_ocean * sfc.depth_ocean
-        energy_total +=
-            CA.horizontal_integral_at_boundary(sol.u[end].sfc.T .* sfc_cρh)
-        energy_surface_change =
-            CA.horizontal_integral_at_boundary(
-                sol.u[end].sfc.T .- sol.u[1].sfc.T,
-            ) * sfc_cρh
-    else
-        energy_surface_change = -p.net_energy_flux_sfc[][]
-    end
-    energy_radiation_input = -p.net_energy_flux_toa[][]
-
-    energy_net =
-        abs(
-            energy_atmos_change + energy_surface_change -
-            energy_radiation_input,
-        ) / energy_total
-    @info "    Net energy change: $energy_net"
-    if CA.has_no_source_or_sink(config.parsed_args)
-        @test (energy_net / energy_total) ≈ 0 atol = 50 * eps(FT)
-    else
-        @test (energy_net / energy_total) ≈ 0 atol = sqrt(eps(FT))
-    end
 
     if CA.has_no_source_or_sink(config.parsed_args)
-        # mass
-        @test sum(sol.u[1].c.ρ) ≈ sum(sol.u[end].c.ρ) rtol = 50 * eps(FT)
+        @test energy_conservation ≈ 0 atol = 50 * eps(FT)
+        @test mass_conservation ≈ 0 atol = 50 * eps(FT)
     else
+        @test energy_conservation ≈ 0 atol = sqrt(eps(FT))
         if sfc isa CA.PrognosticSurfaceTemperature
-            # water
-            water_total = sum(sol.u[end].c.ρq_tot)
-            water_atmos_change =
-                sum(sol.u[end].c.ρq_tot) - sum(sol.u[1].c.ρq_tot)
-            water_surface_change = CA.horizontal_integral_at_boundary(
-                sol.u[end].sfc.water .- sol.u[1].sfc.water,
-            )
-
-            water_net =
-                abs(water_atmos_change + water_surface_change) / water_total
-            @info "    Net water change: $water_net"
-            @test water_net ≈ 0 atol = 100 * sqrt(eps(FT))
+            @info "    Net water change: $water_conservation"
+            @test water_conservation ≈ 0 atol = 100 * sqrt(eps(FT))
         end
     end
 end

src/solver/solve.jl

@@ -196,3 +196,67 @@ function precompile_callbacks(integrator)
     @assert B
     return nothing
 end
+
+
+check_conservation(simulation::AtmosSimulation) =
+    check_conservation(simulation.integrator.sol)
+check_conservation(integrator::CTS.DistributedODEIntegrator) =
+    check_conservation(integrator.sol)
+check_conservation(atmos_sol::AtmosSolveResults) =
+    check_conservation(atmos_sol.sol)
+
+"""
+    check_conservation(sol)
+    check_conservation(simulation)
+    check_conservation(integrator)
+
+Return:
+- `energy_conservation = energy_net / energy_total`
+- `mass_conservation = (mass(t_end) - mass(t_0)) / mass(t_0)`
+- `water_conservation = (water_atmos + water_surface) / water_total`
+
+"""
+function check_conservation(sol)
+    # energy
+    energy_total = sum(sol.u[end].c.ρe_tot)
+    energy_atmos_change = sum(sol.u[end].c.ρe_tot) - sum(sol.u[1].c.ρe_tot)
+    p = sol.prob.p
+    sfc = p.atmos.surface_model
+    if sfc isa PrognosticSurfaceTemperature
+        sfc_cρh = sfc.ρ_ocean * sfc.cp_ocean * sfc.depth_ocean
+        energy_total +=
+            horizontal_integral_at_boundary(sol.u[end].sfc.T .* sfc_cρh)
+        energy_surface_change =
+            horizontal_integral_at_boundary(
+                sol.u[end].sfc.T .- sol.u[1].sfc.T,
+            ) * sfc_cρh
+    else
+        energy_surface_change = -p.net_energy_flux_sfc[][]
+    end
+    energy_radiation_input = -p.net_energy_flux_toa[][]
+
+    energy_conservation =
+        abs(
+            energy_atmos_change + energy_surface_change -
+            energy_radiation_input,
+        ) / energy_total
+
+    mass_conservation =
+        (sum(sol.u[end].c.ρ) - sum(sol.u[1].c.ρ)) / sum(sol.u[1].c.ρ)
+
+    water_conservation = zero(Spaces.undertype(axes(sol.u[end].c.ρ)))
+
+    if sfc isa PrognosticSurfaceTemperature
+        # water
+        water_total = sum(sol.u[end].c.ρq_tot)
+        water_atmos_change = sum(sol.u[end].c.ρq_tot) - sum(sol.u[1].c.ρq_tot)
+        water_surface_change = horizontal_integral_at_boundary(
+            sol.u[end].sfc.water .- sol.u[1].sfc.water,
+        )
+
+        water_conservation =
+            abs(water_atmos_change + water_surface_change) / water_total
+    end
+
+    return (; energy_conservation, mass_conservation, water_conservation)
+end