use imp tend in canopy alone, update diag, harmonic mean at faces (#958)

kmdeck · web-flow · commit 1f8b44b385bd · 2024-12-16T15:52:08.000-08:00
diff --git a/docs/tutorials/standalone/Canopy/canopy_tutorial.jl b/docs/tutorials/standalone/Canopy/canopy_tutorial.jl
@@ -262,6 +262,10 @@ canopy = ClimaLand.Canopy.CanopyModel{FT}(;
 
 Y, p, coords = ClimaLand.initialize(canopy)
 exp_tendency! = make_exp_tendency(canopy);
+imp_tendency! = make_imp_tendency(canopy)
+jacobian! = make_jacobian(canopy);
+jac_kwargs =
+    (; jac_prototype = ClimaLand.ImplicitEquationJacobian(Y), Wfact = jacobian!);
 
 # Provide initial conditions for the canopy hydraulics model
 
@@ -317,12 +321,21 @@ cb = SciMLBase.CallbackSet(driver_cb, saving_cb);
 
 
 # Select a timestepping algorithm and setup the ODE problem.
-
-timestepper = CTS.RK4();
-ode_algo = CTS.ExplicitAlgorithm(timestepper)
+timestepper = CTS.ARS111();
+ode_algo = CTS.IMEXAlgorithm(
+    timestepper,
+    CTS.NewtonsMethod(
+        max_iters = 6,
+        update_j = CTS.UpdateEvery(CTS.NewNewtonIteration),
+    ),
+);
 
 prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLand.dss!),
+    CTS.ClimaODEFunction(
+        T_exp! = exp_tendency!,
+        T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
+        dss! = ClimaLand.dss!,
+    ),
     Y,
     (t0, tf),
     p,
diff --git a/experiments/integrated/fluxnet/run_fluxnet.jl b/experiments/integrated/fluxnet/run_fluxnet.jl
@@ -30,11 +30,11 @@ include(joinpath(climaland_dir, "experiments/integrated/fluxnet/data_tools.jl"))
 include(joinpath(climaland_dir, "experiments/integrated/fluxnet/plot_utils.jl"))
 
 # Read in the site to be run from the command line
-if length(ARGS) < 1
-    error("Must provide site ID on command line")
-end
+#if length(ARGS) < 1
+#    error("Must provide site ID on command line")
+#end
 
-site_ID = ARGS[1]
+site_ID = "US-MOz"#ARGS[1]
 
 # Read all site-specific domain parameters from the simulation file for the site
 include(
diff --git a/experiments/standalone/Vegetation/no_vegetation.jl b/experiments/standalone/Vegetation/no_vegetation.jl
@@ -162,7 +162,10 @@ canopy = ClimaLand.Canopy.CanopyModel{FT}(;
 
 Y, p, coords = ClimaLand.initialize(canopy)
 exp_tendency! = make_exp_tendency(canopy);
-
+imp_tendency! = make_imp_tendency(canopy)
+jacobian! = make_jacobian(canopy);
+jac_kwargs =
+    (; jac_prototype = ClimaLand.ImplicitEquationJacobian(Y), Wfact = jacobian!);
 
 ψ_leaf_0 = FT(-2e5 / 9800)
 
@@ -195,11 +198,22 @@ updatefunc = ClimaLand.make_update_drivers(drivers)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 cb = SciMLBase.CallbackSet(driver_cb, saving_cb);
 
-timestepper = CTS.RK4();
-ode_algo = CTS.ExplicitAlgorithm(timestepper)
+# Set up timestepper
+timestepper = CTS.ARS111();
+ode_algo = CTS.IMEXAlgorithm(
+    timestepper,
+    CTS.NewtonsMethod(
+        max_iters = 6,
+        update_j = CTS.UpdateEvery(CTS.NewNewtonIteration),
+    ),
+);
 
 prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLand.dss!),
+    CTS.ClimaODEFunction(
+        T_exp! = exp_tendency!,
+        T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
+        dss! = ClimaLand.dss!,
+    ),
     Y,
     (t0, tf),
     p,
diff --git a/experiments/standalone/Vegetation/varying_lai.jl b/experiments/standalone/Vegetation/varying_lai.jl
@@ -162,6 +162,10 @@ canopy = ClimaLand.Canopy.CanopyModel{FT}(;
 
 Y, p, coords = ClimaLand.initialize(canopy)
 exp_tendency! = make_exp_tendency(canopy);
+imp_tendency! = make_imp_tendency(canopy)
+jacobian! = make_jacobian(canopy);
+jac_kwargs =
+    (; jac_prototype = ClimaLand.ImplicitEquationJacobian(Y), Wfact = jacobian!);
 
 
 ψ_leaf_0 = FT(-2e5 / 9800)
@@ -195,11 +199,22 @@ updatefunc = ClimaLand.make_update_drivers(drivers)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 cb = SciMLBase.CallbackSet(driver_cb, saving_cb);
 
-timestepper = CTS.RK4();
-ode_algo = CTS.ExplicitAlgorithm(timestepper)
+# Set up timestepper
+timestepper = CTS.ARS111();
+ode_algo = CTS.IMEXAlgorithm(
+    timestepper,
+    CTS.NewtonsMethod(
+        max_iters = 6,
+        update_j = CTS.UpdateEvery(CTS.NewNewtonIteration),
+    ),
+);
 
 prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLand.dss!),
+    CTS.ClimaODEFunction(
+        T_exp! = exp_tendency!,
+        T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
+        dss! = ClimaLand.dss!,
+    ),
     Y,
     (t0, tf),
     p,
diff --git a/experiments/standalone/Vegetation/varying_lai_with_stem.jl b/experiments/standalone/Vegetation/varying_lai_with_stem.jl
@@ -161,7 +161,10 @@ canopy = ClimaLand.Canopy.CanopyModel{FT}(;
 
 Y, p, coords = ClimaLand.initialize(canopy)
 exp_tendency! = make_exp_tendency(canopy);
-
+imp_tendency! = make_imp_tendency(canopy)
+jacobian! = make_jacobian(canopy);
+jac_kwargs =
+    (; jac_prototype = ClimaLand.ImplicitEquationJacobian(Y), Wfact = jacobian!);
 
 ψ_leaf_0 = FT(-2e5 / 9800)
 ψ_stem_0 = FT(-1e5 / 9800)
@@ -202,11 +205,22 @@ updatefunc = ClimaLand.make_update_drivers(drivers)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 cb = SciMLBase.CallbackSet(driver_cb, saving_cb);
 
-timestepper = CTS.RK4();
-ode_algo = CTS.ExplicitAlgorithm(timestepper)
+# Set up timestepper
+timestepper = CTS.ARS111();
+ode_algo = CTS.IMEXAlgorithm(
+    timestepper,
+    CTS.NewtonsMethod(
+        max_iters = 6,
+        update_j = CTS.UpdateEvery(CTS.NewNewtonIteration),
+    ),
+);
 
 prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLand.dss!),
+    CTS.ClimaODEFunction(
+        T_exp! = exp_tendency!,
+        T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
+        dss! = ClimaLand.dss!,
+    ),
     Y,
     (t0, tf),
     p,
diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl
@@ -117,7 +117,7 @@ function default_diagnostics(
             "trans",
             "clhf",
             "cshf",
-            # "lwp", # last(p.canopy.hydraulics.ψ) errors
+            "lwp",
             # "fa", # return a Tuple
             "far",
             "lai",
@@ -264,7 +264,7 @@ function default_diagnostics(
             "trans",
             "clhf",
             "cshf",
-            # "lwp", # last(p.canopy.hydraulics.ψ) errors
+            "lwp",
             # "fa", # return a Tuple
             "far",
             "lai",
diff --git a/src/diagnostics/define_diagnostics.jl b/src/diagnostics/define_diagnostics.jl
@@ -224,7 +224,7 @@ function define_diagnostics!(land_model)
     ### Canopy - Hydraulics
     # Leaf water potential
 
-    #=
+
     add_diagnostic_variable!(
         short_name = "lwp",
         long_name = "Leaf Water Potential",
@@ -234,18 +234,18 @@ function define_diagnostics!(land_model)
         compute! = (out, Y, p, t) ->
             compute_leaf_water_potential!(out, Y, p, t, land_model),
     )
-
-    # Flux per ground area
-    add_diagnostic_variable!(
-        short_name = "fa",
-        long_name = "Flux Per Ground Area",
-        standard_name = "flux_per_ground_area",
-        units = "m s^-1",
-        comments = "Flux of water volume per m^2 of plant per second, multiplied by the area index (plant area/ground area).",
-        compute! = (out, Y, p, t) ->
-            compute_flux_per_ground_area!(out, Y, p, t, land_model),
-    )
-    =#
+    #=
+        # Flux per ground area
+        add_diagnostic_variable!(
+            short_name = "fa",
+            long_name = "Flux Per Ground Area",
+            standard_name = "flux_per_ground_area",
+            units = "m s^-1",
+            comments = "Flux of water volume per m^2 of plant per second, multiplied by the area index (plant area/ground area).",
+            compute! = (out, Y, p, t) ->
+                compute_flux_per_ground_area!(out, Y, p, t, land_model),
+        )
+        =#
 
     # Root flux per ground area
     add_diagnostic_variable!(
diff --git a/src/diagnostics/land_compute_methods.jl b/src/diagnostics/land_compute_methods.jl
@@ -88,9 +88,24 @@ end
 } p.canopy.energy.turbulent_fluxes.shf
 
 # Canopy - Hydraulics
-#@diagnostic_compute "leaf_water_potential" Union{SoilCanopyModel, LandModel} last(
-#    p.canopy.hydraulics.ψ,
-#)
+function compute_leaf_water_potential!(
+    out,
+    Y,
+    p,
+    t,
+    land_model::Union{SoilCanopyModel, LandModel},
+)
+    hydraulics = land_model.canopy.hydraulics
+    n_stem = hydraulics.n_stem
+    n_leaf = hydraulics.n_leaf
+    n = n_stem + n_leaf
+    if isnothing(out)
+        return p.canopy.hydraulics.ψ.:($n)
+    else
+        out .= p.canopy.hydraulics.ψ.:($n)
+    end
+end
+
 # @diagnostic_compute "flux_per_ground_area" Union{SoilCanopyModel, LandModel} p.canopy.hydraulics.fa # return a Tuple
 @diagnostic_compute "root_flux_per_ground_area" Union{
     SoilCanopyModel,
diff --git a/src/integrated/soil_canopy_root_interactions.jl b/src/integrated/soil_canopy_root_interactions.jl
@@ -9,9 +9,15 @@ function update_root_extraction!(p, Y, t, land)
     z = land.soil.domain.fields.z
     (; conductivity_model) = land.canopy.hydraulics.parameters
     area_index = p.canopy.hydraulics.area_index
-
+    # allocates
     above_ground_area_index =
-        getproperty(area_index, land.canopy.hydraulics.compartment_labels[1])
+        PlantHydraulics.harmonic_mean.(
+            getproperty(
+                area_index,
+                land.canopy.hydraulics.compartment_labels[1],
+            ),
+            getproperty(area_index, :root),
+        )
     # Note that we model the flux between each soil layer and the canopy as:
     # Flux = -K_eff x [(ψ_canopy - ψ_soil)/(z_canopy - z_soil) + 1], where
     # K_eff = K_soil K_canopy /(K_canopy + K_soil)
diff --git a/src/standalone/Vegetation/Canopy.jl b/src/standalone/Vegetation/Canopy.jl
@@ -539,7 +539,7 @@ function ClimaLand.make_update_aux(
                         conductivity_model,
                         ψ.:($$ip1),
                     ),
-                ) * areaip1
+                ) * PlantHydraulics.harmonic_mean(areaip1, areai)
         end
         # We update the fa[n_stem+n_leaf] element once we have computed transpiration, below
         # update photosynthesis and conductance terms
diff --git a/src/standalone/Vegetation/PlantHydraulics.jl b/src/standalone/Vegetation/PlantHydraulics.jl
@@ -350,6 +350,15 @@ function ClimaLand.Canopy.set_canopy_prescribed_field!(
     @. p.canopy.hydraulics.area_index.root = RAI
 end
 
+"""
+    harmonic_mean(x::FT,y::FT) where {FT} 
+
+Computes the harmonic mean of x >=0 and y >=0; returns zero if both
+x and y are zero.
+    
+"""
+harmonic_mean(x::FT, y::FT) where {FT} = x * y / max(x + y, eps(FT))
+
 
 """
     water_flux(
@@ -378,7 +387,7 @@ F = -k_eff * (ψ1 +z1 - ψ2 - z2) (m/s).
 This currently assumes the path lengths are equal.
 """
 function water_flux(z1::FT, z2::FT, ψ1::FT, ψ2::FT, K1::FT, K2::FT) where {FT}
-    K_eff = K1 * K2 / max(K1 + K2, eps(FT))
+    K_eff = harmonic_mean(K1, K2)
     flux = -K_eff * ((ψ2 - ψ1) / (z2 - z1) + 1)
     return flux # (m/s)
 end
@@ -595,22 +604,15 @@ function make_compute_exp_tendency(
         @inbounds for i in 1:(n_stem + n_leaf)
             im1 = i - 1
             ip1 = i + 1
-            # To prevent dividing by zero, change AI/(AI x dz)" to
-            # "AI/max(AI x dz, eps(FT))"
-            AIdz =
-                max.(
-                    getproperty(area_index, model.compartment_labels[i]) .* (
-                        model.compartment_surfaces[ip1] -
-                        model.compartment_surfaces[i]
-                    ),
-                    eps(FT),
-                )
+            # To prevent dividing by zero, change AI/(AI x dz)" to AI/max(AI x dz, eps(FT))"
+            AI = getproperty(area_index, model.compartment_labels[i]) # this is a field; should not allocate here
+            dz = model.compartment_surfaces[ip1] - model.compartment_surfaces[i] # currently this is a scalar. in the future, this will be a field.
             if i == 1
                 @inbounds @. dY.canopy.hydraulics.ϑ_l.:($$i) =
-                    1 / AIdz * (fa_roots - fa.:($$i))
+                    1 / max(AI * dz, eps(FT)) * (fa_roots - fa.:($$i))
             else
                 @inbounds @. dY.canopy.hydraulics.ϑ_l.:($$i) =
-                    1 / AIdz * (fa.:($$im1) - fa.:($$i))
+                    1 / max(AI * dz, eps(FT)) * (fa.:($$im1) - fa.:($$i))
             end
         end
     end
@@ -662,7 +664,10 @@ function root_water_flux_per_ground_area!(
     fa .= FT(0.0)
     @inbounds for i in 1:n_root_layers
         above_ground_area_index =
-            getproperty(area_index, model.compartment_labels[1])
+            harmonic_mean.(
+                getproperty(area_index, model.compartment_labels[1]),
+                getproperty(area_index, :root),
+            )
 
         if i != n_root_layers
             @. fa +=
