Merge #436

436: Revamp ConservationChecker r=LenkaNovak a=LenkaNovak



Co-authored-by: LenkaNovak <lenka@caltech.edu>

config/model_configs/coarse_mpi_n2.yml

@@ -13,5 +13,5 @@ h_elem: 6
 dt_save_restart: "5days"
 precip_model: "0M"
 run_name: "coarse_mpi_n2"
-apply_limiter: true
+apply_limiter: false
 job_id: "coarse_mpi_n2"

config/model_configs/coarse_single_modular_ft64.yml

@@ -13,6 +13,6 @@ mono_surface: true
 h_elem: 6
 dt_save_restart: "10days"
 precip_model: "0M"
-apply_limiter: true
+apply_limiter: false
 job_id: "coarse_single_modular_ft64"
 

config/model_configs/coarse_single_modular_ft64_monthly_checkpoints.yml

@@ -14,5 +14,5 @@ mono_surface: false
 h_elem: 6
 dt_save_restart: "10days"
 precip_model: "0M"
-apply_limiter: true
+apply_limiter: false
 job_id: "coarse_single_modular_ft64_monthly_checkpoints"

config/model_configs/default_modular_unthreaded.yml

@@ -12,5 +12,5 @@ mono_surface: true
 h_elem: 4
 precip_model: "0M"
 anim: true
-apply_limiter: true
+apply_limiter: false
 job_id: "default_modular_unthreaded"

config/model_configs/default_mono.yml

@@ -11,5 +11,5 @@ mono_surface: true
 h_elem: 4
 precip_model: "0M"
 run_name: "default_mono"
-apply_limiter: true
+apply_limiter: false
 job_id: "default_mono"

config/model_configs/default_notmono.yml

@@ -11,5 +11,5 @@ mono_surface: false
 h_elem: 4
 precip_model: "0M"
 run_name: "default_notmono"
-apply_limiter: true
+apply_limiter: false
 job_id: "default_notmono"

config/model_configs/interactive_debug.yml

@@ -4,9 +4,9 @@ rad: "gray"
 energy_check: true
 mode_name: "slabplanet"
 t_end: "10days"
-dt_save_to_sol: "3600secs"
-dt_cpl: 200
-dt: "200secs"
+dt_save_to_sol: "2days"
+dt_cpl: 400
+dt: "400secs"
 mono_surface: true
 turb_flux_partition: "CombinedStateFluxes"
 h_elem: 4

config/model_configs/slabplanet_albedo_function.yml

@@ -12,5 +12,5 @@ h_elem: 4
 precip_model: "0M"
 land_albedo_type: "map_static"
 run_name: "slabplanet_albedo_function"
-apply_limiter: true
+apply_limiter: false
 job_id: "slabplanet_albedo_function"

config/model_configs/slabplanet_albedo_static_map.yml

@@ -12,5 +12,5 @@ h_elem: 4
 precip_model: "0M"
 land_albedo_type: "map_static"
 run_name: "slabplanet_albedo_static_map"
-apply_limiter: true
+apply_limiter: false
 job_id: "slabplanet_albedo_static_map"

config/model_configs/slabplanet_albedo_temporal_map.yml

@@ -12,5 +12,5 @@ h_elem: 4
 precip_model: "0M"
 land_albedo_type: "map_temporal"
 run_name: "slabplanet_albedo_temporal_map"
-apply_limiter: true
+apply_limiter: false
 job_id: "slabplanet_albedo_temporal_map"

config/model_configs/slabplanet_default.yml

@@ -12,5 +12,5 @@ mono_surface: true
 h_elem: 4
 precip_model: "0M"
 anim: true
-apply_limiter: true
+apply_limiter: false
 job_id: "slabplanet_default"

config/model_configs/slabplanet_nonmono.yml

@@ -12,5 +12,5 @@ mono_surface: false
 h_elem: 4
 precip_model: "0M"
 anim: true
-apply_limiter: true
+apply_limiter: false
 job_id: "slabplanet_nonmono"

config/model_configs/slabplanet_partitioned_fluxes.yml

@@ -13,5 +13,5 @@ mono_surface: true
 h_elem: 4
 precip_model: "0M"
 anim: true
-apply_limiter: true
+apply_limiter: false
 job_id: "slabplanet_partitioned_fluxes"

config/model_configs/target_amip_n1_shortrun.yml

@@ -20,5 +20,5 @@ t_end: "0.01days"
 dt_save_to_sol: "1days"
 mono_surface: false
 precip_model: "0M"
-apply_limiter: true
+apply_limiter: false
 job_id: "target_amip_n1_shortrun"

config/model_configs/target_params_in_slab.yml

@@ -17,5 +17,5 @@ mono_surface: true
 h_elem: 6
 precip_model: "0M"
 albedo_from_file: false
-apply_limiter: true
+apply_limiter: false
 job_id: "target_params_in_slab"

config/model_configs/target_params_in_slab_test2.yml

@@ -16,5 +16,5 @@ mono_surface: true
 h_elem: 6
 precip_model: "0M"
 run_name: "target_params_in_slab_test2"
-apply_limiter: true
+apply_limiter: false
 job_id: "target_params_in_slab_test2"

config/model_configs/target_params_in_slab_test3.yml

@@ -16,5 +16,5 @@ mono_surface: true
 h_elem: 6
 precip_model: "0M"
 run_name: "target_params_in_slab_test3"
-apply_limiter: true
+apply_limiter: false
 job_id: "target_params_in_slab_test3"

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

@@ -71,9 +71,14 @@ end
 # extensions required by the Interfacer
 get_field(sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux}) =
     Fields.level(sim.integrator.p.ᶠradiation_flux, half)
-get_field(sim::ClimaAtmosSimulation, ::Val{:liquid_precipitation}) =
-    sim.integrator.p.col_integrated_rain .+ sim.integrator.p.col_integrated_snow # all fallen snow melts for now
-get_field(sim::ClimaAtmosSimulation, ::Val{:snow_precipitation}) = sim.integrator.p.col_integrated_snow .* FT(0)
+function get_field(sim::ClimaAtmosSimulation, ::Val{:liquid_precipitation})
+    ρ_liq = CAP.ρ_cloud_liq(sim.integrator.p.params)
+    sim.integrator.p.col_integrated_rain .* ρ_liq # kg/m^2/s
+end
+function get_field(sim::ClimaAtmosSimulation, ::Val{:snow_precipitation})
+    ρ_liq = CAP.ρ_cloud_liq(sim.integrator.p.params)
+    sim.integrator.p.col_integrated_snow .* ρ_liq  # kg/m^2/s
+end
 
 get_field(sim::ClimaAtmosSimulation, ::Val{:turbulent_energy_flux}) =
     Geometry.WVector.(sim.integrator.p.sfc_conditions.ρ_flux_h_tot)
@@ -252,3 +257,46 @@ Extension of Checkpointer.get_model_state_vector to get the model state.
 function get_model_state_vector(sim::ClimaAtmosSimulation)
     return sim.integrator.u
 end
+
+"""
+    get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:F_radiative_TOA})
+
+Extension of Interfacer.get_field to get the net TOA radiation, which is a sum of the
+upward and downward longwave and shortwave radiation.
+"""
+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}) = atmos_sim.integrator.u.c.ρe_tot
+
+get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:water}) = atmos_sim.integrator.u.c.ρq_tot

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

@@ -1,35 +1,3 @@
-
-"""
-    get_bucket_energy(bucket_sim)
-
-Returns the volumetric internal energy of the bucket land model.
-"""
-function get_land_energy(bucket_sim::BucketSimulation, e_per_area)
-    # required by ConservationChecker
-    e_per_area .= zeros(axes(bucket_sim.integrator.u.bucket.W))
-    soil_depth = FT = eltype(bucket_sim.integrator.u.bucket.W)
-    ClimaCore.Fields.bycolumn(axes(bucket_sim.integrator.u.bucket.T)) do colidx
-        e_per_area[colidx] .=
-            bucket_sim.model.parameters.ρc_soil .* mean(bucket_sim.integrator.u.bucket.T[colidx]) .*
-            bucket_sim.domain.soil_depth
-    end
-
-    e_per_area .+=
-        -LSMP.LH_f0(bucket_sim.model.parameters.earth_param_set) .*
-        LSMP.ρ_cloud_liq(bucket_sim.model.parameters.earth_param_set) .* bucket_sim.integrator.u.bucket.σS
-    return e_per_area
-end
-
-"""
-    get_land_temp_from_state(land_sim, u)
-Returns the surface temperature of the earth, computed
-from the state u.
-"""
-function get_land_temp_from_state(land_sim, u)
-    # required by viz_explorer.jl
-    return ClimaLSM.surface_temperature(land_sim.model, u, land_sim.integrator.p, land_sim.integrator.t)
-end
-
 """
     make_lsm_domain(
         atmos_boundary_space::ClimaCore.Spaces.SpectralElementSpace2D,
@@ -105,10 +73,12 @@ function update_field!(sim::BucketSimulation, ::Val{:radiative_energy_flux}, fie
     parent(sim.integrator.p.bucket.R_n) .= parent(field)
 end
 function update_field!(sim::BucketSimulation, ::Val{:liquid_precipitation}, field)
-    parent(sim.integrator.p.bucket.P_liq) .= parent(field)
+    ρ_liq = (LSMP.ρ_cloud_liq(sim.model.parameters.earth_param_set))
+    parent(sim.integrator.p.bucket.P_liq) .= parent(field ./ ρ_liq)
 end
 function update_field!(sim::BucketSimulation, ::Val{:snow_precipitation}, field)
-    parent(sim.integrator.p.bucket.P_snow) .= parent(field)
+    ρ_liq = (LSMP.ρ_cloud_liq(sim.model.parameters.earth_param_set))
+    parent(sim.integrator.p.bucket.P_snow) .= parent(field ./ ρ_liq)
 end
 
 function update_field!(sim::BucketSimulation, ::Val{:air_density}, field)
@@ -153,3 +123,44 @@ Extension of Checkpointer.get_model_state_vector to get the model state.
 function get_model_state_vector(sim::BucketSimulation)
     return sim.integrator.u.bucket
 end
+
+"""
+    get_field(bucket_sim::BucketSimulation, ::Val{:energy})
+
+Extension of Interfacer.get_field that provides the total energy contained in the bucket, including the latent heat due to snow melt.
+"""
+function get_field(bucket_sim::BucketSimulation, ::Val{:energy})
+    # required by ConservationChecker
+    e_per_area = zeros(axes(bucket_sim.integrator.u.bucket.W))
+    ClimaCore.Fields.bycolumn(axes(bucket_sim.integrator.u.bucket.T)) do colidx
+        e_per_area[colidx] .=
+            bucket_sim.model.parameters.ρc_soil .* mean(bucket_sim.integrator.u.bucket.T[colidx]) .*
+            bucket_sim.domain.soil_depth
+    end
+
+    e_per_area .+=
+        -LSMP.LH_f0(bucket_sim.model.parameters.earth_param_set) .*
+        LSMP.ρ_cloud_liq(bucket_sim.model.parameters.earth_param_set) .* bucket_sim.integrator.u.bucket.σS
+    return e_per_area
+end
+
+"""
+    get_field(bucket_sim::BucketSimulation, ::Val{:water})
+
+Extension of Interfacer.get_field that provides the total water contained in the bucket, including the liquid water in snow.
+"""
+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)
+Returns the surface temperature of the earth, computed from the state u.
+"""
+function get_land_temp_from_state(land_sim, u)
+    # required by viz_explorer.jl
+    return ClimaLSM.surface_temperature(land_sim.model, u, land_sim.integrator.p, land_sim.integrator.t)
+end

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

@@ -157,3 +157,14 @@ 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})
+
+Extension of Interfacer.get_field to get the energy of the ocean.
+It multiplies the the slab temperature by the heat capacity, density, and depth.
+"""
+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
+
+get_field(sim::SlabOceanSimulation, ::Val{:water}) = nothing

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

@@ -167,3 +167,14 @@ 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})
+
+Extension of Interfacer.get_field to get the energy of the ocean.
+It multiplies the the slab temperature by the heat capacity, density, and depth.
+"""
+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
+
+get_field(sim::PrescribedIceSimulation, ::Val{:water}) = nothing

experiments/AMIP/modular/coupler_driver_modular.jl

@@ -387,8 +387,7 @@ if energy_check
         mode_name == "slabplanet" && !CA.is_distributed(ClimaComms.context(boundary_space)),
         "Only non-distributed slabplanet allowable for energy_check"
     )
-    conservation_checks =
-        (; energy = EnergyConservationCheck([], [], [], [], [], []), water = WaterConservationCheck([], [], [], []))
+    conservation_checks = (; energy = EnergyConservationCheck(model_sims), water = WaterConservationCheck(model_sims))
 end
 
 ## coupler simulation
@@ -516,7 +515,7 @@ function solve_coupler!(cs)
         end
 
         ## compute global energy
-        !isnothing(cs.conservation_checks) ? check_conservation!(cs, get_slab_energy, get_land_energy) : nothing
+        !isnothing(cs.conservation_checks) ? check_conservation!(cs) : nothing
 
         ## run component models sequentially for one coupling timestep (Δt_cpl)
         ClimaComms.barrier(comms_ctx)