conservation computation, diagnostics, plotting [skip ci]

Author: kmdeck

src/diagnostics/default_diagnostics.jl

@@ -270,6 +270,8 @@ function default_diagnostics(
     output_writer,
     output_vars = :short,
     average_period = :monthly,
+    conservation = false,
+    conservation_period = Day(10),
 ) where {FT}
 
     define_diagnostics!(land_model)
@@ -388,4 +390,25 @@ function default_diagnostics(
             start_date,
         )
     end
+
+    if conservation
+        additional_diags = ["epa", "epac", "wvpa", "wvpac"]
+        additional_outputs = vcat(
+            map(additional_diags) do short_name
+                output_schedule_func =
+                    conservation_period isa Period ?
+                    EveryCalendarDtSchedule(conservation_period; start_date) : EveryDtSchedule(conservation_period)
+                return ScheduledDiagnostic(
+                    variable = get_diagnostic_variable(short_name),
+                    compute_schedule_func = EveryStepSchedule(),
+                    output_schedule_func = output_schedule_func,
+                    output_writer = output_writer,
+                )
+            end...,
+        )
+    else
+        additional_outputs = []
+    end
+
+    return [default_outputs..., additional_outputs...]
 end

src/diagnostics/land_compute_methods.jl

@@ -55,6 +55,11 @@ end
 @diagnostic_compute "water_volume_per_area_change" EnergyHydrology Y.soil.∫F_vol_liq_water_dt
 @diagnostic_compute "energy_per_area_change" EnergyHydrology Y.soil.∫F_e_dt
 
+@diagnostic_compute "water_volume_per_area" LandModel p.total_water
+@diagnostic_compute "energy_per_area" LandModel p.total_energy
+@diagnostic_compute "water_volume_per_area_change" LandModel Y.∫F_vol_liq_water_dt
+@diagnostic_compute "energy_per_area_change" LandModel Y.∫F_e_dt
+
 
 ### BucketModel ###
 

src/integrated/land.jl

@@ -222,6 +222,8 @@ lsm_aux_vars(m::LandModel) = (
     :ϵ_sfc,
     :α_sfc,
     :α_ground,
+    :total_energy,
+    :total_water,
 )
 
 """
@@ -249,6 +251,8 @@ lsm_aux_types(m::LandModel{FT}) where {FT} = (
     FT,
     FT,
     NamedTuple{(:PAR, :NIR), Tuple{FT, FT}},
+    FT,
+    FT,
 )
 
 """
@@ -276,8 +280,157 @@ lsm_aux_domain_names(m::LandModel) = (
     :surface,
     :surface,
     :surface,
+    :surface,
+    :surface,
 )
 
+function initialize_prognostic(
+    model::LandModel{FT},
+    coords::NamedTuple,
+) where {FT}
+    components = land_components(model)
+    Y_state_list = map(components) do (component)
+        submodel = getproperty(model, component)
+        getproperty(initialize_prognostic(submodel, coords), component)
+    end
+    conservation_vars = (:∫Fwdt, :∫Fedt)
+    conservation_types = (FT, FT)
+    conservation_domain_names = (:surface, :surface)
+    conservation_state_list = initialize_vars(
+        conservation_vars,
+        conservation_types,
+        conservation_domain_names,
+        coords,
+        :conservation_check,
+    )
+    Y = ClimaCore.Fields.FieldVector(;
+        NamedTuple{(components..., conservation_vars...)}((
+            Y_state_list...,
+            conservation_state_list.conservation_check...,
+        ))...,
+    )
+    return Y
+end
+
+function make_update_aux(land::LandModel)
+    components = land_components(land)
+    update_aux_function_list =
+        map(x -> make_update_aux(getproperty(land, x)), components)
+    function update_aux!(p, Y, t)
+        for f! in update_aux_function_list
+            f!(p, Y, t)
+        end
+        sfc_cache = p.scratch1 .* 0
+        ClimaLand.total_liq_water_vol_per_area!(
+            p.total_water,
+            land,
+            Y,
+            p,
+            t,
+            sfc_cache,
+        )
+        ClimaLand.total_energy_per_area!(
+            p.total_energy,
+            land,
+            Y,
+            p,
+            t,
+            sfc_cache,
+        )
+    end
+    return update_aux!
+end
+
+function make_imp_tendency(land::LandModel)
+    components = land_components(land)
+    compute_imp_tendency_list =
+        map(x -> make_compute_imp_tendency(getproperty(land, x)), components)
+    update_aux! = make_update_aux(land)
+    update_boundary_fluxes! = make_update_boundary_fluxes(land)
+    function imp_tendency!(dY, Y, p, t)
+        update_aux!(p, Y, t)
+        update_boundary_fluxes!(p, Y, t)
+        for f! in compute_imp_tendency_list
+            f!(dY, Y, p, t)
+        end
+        # Soil and canopy implicit terms
+        # soil contributions were computed by soil and account for snow cover fraction
+        @. dY.∫Fedt = -(
+            (1 - p.snow.snow_cover_fraction) * (
+                p.soil.R_n +
+                p.soil.turbulent_fluxes.lhf +
+                p.soil.turbulent_fluxes.shf
+            ) +
+            p.canopy.turbulent_fluxes.shf +
+            p.canopy.turbulent_fluxes.lhf -
+            p.canopy.radiative_transfer.SW_n -
+            p.canopy.radiative_transfer.LW_n
+        )
+        @. dY.∫Fwdt = -(
+            (1 - p.snow.snow_cover_fraction) *
+            (p.drivers.P_liq + p.soil.turbulent_fluxes.vapor_flux_liq) +
+            p.soil.R_s - p.soil.bottom_bc.water
+        )
+
+    end
+    return imp_tendency!
+end
+
+function make_exp_tendency(land::LandModel{FT}) where {FT}
+    components = land_components(land)
+    compute_exp_tendency_list =
+        map(x -> make_compute_exp_tendency(getproperty(land, x)), components)
+    update_aux! = make_update_aux(land)
+    update_boundary_fluxes! = make_update_boundary_fluxes(land)
+    function exp_tendency!(dY, Y, p, t)
+        update_aux!(p, Y, t)
+        # Treat the following terms explicitly!
+
+        # update root extraction
+        update_root_extraction!(p, Y, t, land) # defined in src/integrated/soil_canopy_root_interactions.jl
+
+        # Compute the ground heat flux in place:
+        update_soil_snow_ground_heat_flux!(
+            p,
+            Y,
+            land.soil.parameters,
+            land.snow.parameters,
+            land.soil.domain,
+            FT,
+        )
+        # Treat the following terms explicitly!
+        update_boundary_fluxes!(p, Y, t)
+        for f! in compute_exp_tendency_list
+            f!(dY, Y, p, t)
+        end
+        _LH_f0 = LP.LH_f0(land.soil.parameters.earth_param_set)
+        _ρ_liq = LP.ρ_cloud_liq(land.soil.parameters.earth_param_set)
+        ρe_falling_snow = -_LH_f0 * _ρ_liq # per unit vol of liquid water
+        # Explicit source terms for soil, and snow fluxes
+        # soil contributions were computed by soil and account for snow cover fraction
+        @. dY.∫Fedt = -(
+            p.drivers.P_snow * ρe_falling_snow +
+            (
+                p.snow.turbulent_fluxes.lhf +
+                p.snow.turbulent_fluxes.shf +
+                p.snow.R_n - p.snow.energy_runoff
+            ) * p.snow.snow_cover_fraction
+        )
+        # Explicit snow, canopy, source terms for soil (including sublimation, subsurface runoff)
+        @. dY.∫Fwdt =
+            -p.soil.R_ss -
+            p.soil.turbulent_fluxes.vapor_flux_ice *
+            (1 - p.snow.snow_cover_fraction) - (
+                p.drivers.P_snow +
+                (p.drivers.P_liq + p.snow.turbulent_fluxes.vapor_flux) *
+                p.snow.snow_cover_fraction +
+                p.canopy.turbulent_fluxes.transpiration
+            )
+    end
+    return exp_tendency!
+end
+
+
 """
     make_update_boundary_fluxes(
         land::LandModel{FT, MM, SM, RM, SnM},
@@ -313,9 +466,6 @@ function make_update_boundary_fluxes(
 
     function update_boundary_fluxes!(p, Y, t)
         earth_param_set = land.soil.parameters.earth_param_set
-        # update root extraction
-        update_root_extraction!(p, Y, t, land) # defined in src/integrated/soil_canopy_root_interactions.jl
-
         # Radiation - updates Rn for soil and snow also
         lsm_radiant_energy_fluxes!(
             p,
@@ -330,16 +480,6 @@ function make_update_boundary_fluxes(
             p,
             land.soil.parameters.earth_param_set,
         )
-
-        # Compute the ground heat flux in place:
-        update_soil_snow_ground_heat_flux!(
-            p,
-            Y,
-            land.soil.parameters,
-            land.snow.parameters,
-            land.soil.domain,
-            FT,
-        )
         #Now update snow boundary conditions, which rely on the ground heat flux
         update_snow_bf!(p, Y, t)
 
@@ -527,7 +667,6 @@ function soil_boundary_fluxes!(
         p.excess_heat_flux +
         p.snow.snow_cover_fraction * p.ground_heat_flux +
         infiltration_energy_flux
-
     return nothing
 end
 
@@ -615,7 +754,7 @@ function snow_boundary_fluxes!(
     e_flux_falling_rain =
         Snow.energy_flux_falling_rain(bc.atmos, p, model.parameters)
 
-    # positive fluxes are TOWARDS atmos, but R_n positive if snow absorbs energy
+    # positive fluxes are TOWARDS atmos
     @. p.snow.total_energy_flux =
         e_flux_falling_snow +
         (

src/integrated/soil_canopy_root_interactions.jl

@@ -42,6 +42,7 @@ function update_root_extraction!(p, Y, t, land)
             z,
             land.canopy.hydraulics.parameters.rooting_depth,
         )
+
     @. p.root_energy_extraction =
         p.root_extraction * ClimaLand.Soil.volumetric_internal_energy_liq(
             p.soil.T,
@@ -163,4 +164,5 @@ function ClimaLand.source!(
 
     ClimaCore.Operators.column_integral_definite!(p.scratch1, p.root_extraction)
     @. dY.soil.∫F_vol_liq_water_dt += p.scratch1
+    # We don't update ∫Sdz here because it cancels out by design with the bottom flux into the canopy (within the land model)
 end

src/integrated/soil_snow_model.jl

@@ -359,7 +359,6 @@ function soil_boundary_fluxes!(
         p.excess_heat_flux +
         p.snow.snow_cover_fraction * p.ground_heat_flux +
         infiltration_energy_flux
-
     return nothing
 end
 

src/shared_utilities/drivers.jl

@@ -1227,6 +1227,7 @@ function make_update_drivers(a::PrescribedAtmosphere{FT}) where {FT}
     function update_drivers!(p, t)
         evaluate!(p.drivers.P_liq, a.liquid_precip, t)
         evaluate!(p.drivers.P_snow, a.snow_precip, t)
+        p.drivers.P_snow .*= 0
         evaluate!(p.drivers.T, a.T, t)
         evaluate!(p.drivers.P, a.P, t)
         evaluate!(p.drivers.u, a.u, t)

src/shared_utilities/implicit_timestepping.jl

@@ -102,6 +102,8 @@ function FieldMatrixWithSolver(Y::ClimaCore.Fields.FieldVector)
     explicit_vars = (
         @name(soil.∫F_vol_liq_water_dt),
         @name(soil.∫F_e_dt),
+        @name(∫F_vol_liq_water_dt),
+        @name(∫F_e_dt),
         @name(soilco2.C),
         @name(soil.θ_i),
         @name(canopy.hydraulics.ϑ_l),

src/simulations/initial_conditions.jl

@@ -166,13 +166,14 @@ function set_snow_initial_conditions!(
     regridder_type = regridder_type,
     extrapolation_bc = extrapolation_bc,
 )
-    Y.snow.S .= SpaceVaryingInput(
-        snow_ic_path,
-        "swe",
-        surface_space;
-        regridder_type,
-        regridder_kwargs = (; extrapolation_bc,),
-    )
+    Y.snow.S .=
+        SpaceVaryingInput(
+            snow_ic_path,
+            "swe",
+            surface_space;
+            regridder_type,
+            regridder_kwargs = (; extrapolation_bc,),
+        ) .* 0
     Y.snow.S_l .= 0
     p.snow.T .= enforce_snow_temperature_constraint.(Y.snow.S, p.snow.T)
     Y.snow.U .=

src/standalone/Vegetation/PlantHydraulics.jl

@@ -342,11 +342,11 @@ function ClimaLand.Canopy.set_canopy_prescribed_field!(
     t,
 ) where {FT}
     (; LAIfunction, SAI, RAI) = component.parameters.ai_parameterization
-    evaluate!(p.canopy.hydraulics.area_index.leaf, LAIfunction, t)
-    p.canopy.hydraulics.area_index.leaf .=
-        clip.(p.canopy.hydraulics.area_index.leaf, FT(0.05))
-    @. p.canopy.hydraulics.area_index.stem = SAI
-    @. p.canopy.hydraulics.area_index.root = RAI
+    #evaluate!(p.canopy.hydraulics.area_index.leaf, LAIfunction, t)
+    p.canopy.hydraulics.area_index.leaf .= 0#
+    #clip.(p.canopy.hydraulics.area_index.leaf, FT(0.05))
+    @. p.canopy.hydraulics.area_index.stem = 0#SAI
+    @. p.canopy.hydraulics.area_index.root = 0#RAI
 end
 
 """