SoilCanopyModel uses implicit solver

docs/tutorials/integrated/soil_canopy_tutorial.jl

@@ -82,7 +82,7 @@ land_domain = Column(; zlim = (zmin, zmax), nelements = nelements);
 #   the simulation with atmospheric and radiative flux models. We also
 # read in the observed LAI and let that vary in time in a prescribed manner.
 
-# Use the data tools for reading FLUXNET data sets 
+# Use the data tools for reading FLUXNET data sets
 include(
     joinpath(pkgdir(ClimaLand), "experiments/integrated/fluxnet/data_tools.jl"),
 );
@@ -326,6 +326,12 @@ land = SoilCanopyModel{FT}(;
 
 Y, p, coords = initialize(land);
 exp_tendency! = make_exp_tendency(land);
+imp_tendency! = make_imp_tendency(land);
+tendency_jacobian! = make_tendency_jacobian(land);
+jac_kwargs = (;
+    jac_prototype = ClimaLand.ImplicitEquationJacobian(Y),
+    Wfact = tendency_jacobian!,
+);
 
 # We need to provide initial conditions for the soil and canopy hydraulics
 # models:
@@ -370,8 +376,14 @@ dt = Float64(30)
 n = 120
 saveat = Array(t0:(n * dt):tf)
 
-timestepper = CTS.RK4()
-ode_algo = CTS.ExplicitAlgorithm(timestepper);
+timestepper = CTS.ARS343()
+ode_algo = CTS.IMEXAlgorithm(
+    timestepper,
+    CTS.NewtonsMethod(
+        max_iters = 1,
+        update_j = CTS.UpdateEvery(CTS.NewNewtonIteration),
+    ),
+);
 
 # Now set the initial values for the cache variables for the combined soil and plant model.
 
@@ -394,7 +406,11 @@ cb = SciMLBase.CallbackSet(driver_cb, saving_cb);
 
 # Carry out the simulation
 prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!),
+    CTS.ClimaODEFunction(
+        T_exp! = exp_tendency!,
+        T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
+        dss! = ClimaLand.dss!,
+    ),
     Y,
     (t0, tf),
     p,

experiments/integrated/fluxnet/US-Var/US-Var_simulation.jl

@@ -1,5 +1,5 @@
 """This file contains simulation variables for running Clima Land on the US-Var
-fluxtower site. This includes both the domain variables and timestepping 
+fluxtower site. This includes both the domain variables and timestepping
 variables for running the simulation."""
 
 # DOMAIN SETUP:
@@ -20,7 +20,7 @@ h_stem = FT(0) # m
 t0 = Float64(21 * 3600 * 24)# start day 21 of the year
 
 # Time step size:
-dt = Float64(40)
+dt = Float64(20)
 
 # Number of timesteps between saving output:
 n = 45

experiments/integrated/fluxnet/fluxnet_simulation.jl

@@ -1,12 +1,19 @@
-"""This file contains the site-generic time variables for running ClimaLand on 
-fluxtower sites. These work in tandem with the site-specific timing parameters 
+"""This file contains the site-generic time variables for running ClimaLand on
+fluxtower sites. These work in tandem with the site-specific timing parameters
 found in the {site-ID}_simulation.jl files in each site directory."""
 
 N_spinup_days = 30
 N_days = N_spinup_days + 30
 tf = Float64(t0 + 3600 * 24 * N_days)
 t_spinup = Float64(t0 + N_spinup_days * 3600 * 24)
 saveat = Array(t_spinup:(n * dt):tf)
-timestepper = CTS.RK4()
+
 # Set up timestepper
-ode_algo = CTS.ExplicitAlgorithm(timestepper)
+timestepper = CTS.ARS343();
+ode_algo = CTS.IMEXAlgorithm(
+    timestepper,
+    CTS.NewtonsMethod(
+        max_iters = 1,
+        update_j = CTS.UpdateEvery(CTS.NewNewtonIteration),
+    ),
+);

experiments/integrated/fluxnet/ozark_pft.jl

@@ -1,5 +1,5 @@
 """
-This experiment tests running the Ozark site (US-MOz) using plant parameters 
+This experiment tests running the Ozark site (US-MOz) using plant parameters
 defined by plant functional types instead of fully site-specific parameters.
 """
 
@@ -252,6 +252,12 @@ land = SoilCanopyModel{FT}(;
 )
 Y, p, cds = initialize(land)
 exp_tendency! = make_exp_tendency(land)
+imp_tendency! = make_imp_tendency(land);
+tendency_jacobian! = make_tendency_jacobian(land);
+jac_kwargs = (;
+    jac_prototype = ClimaLand.ImplicitEquationJacobian(Y),
+    Wfact = tendency_jacobian!,
+);
 
 #Initial conditions
 Y.soil.ϑ_l =
@@ -274,7 +280,7 @@ Y.soil.ρe_int =
         Ref(land.soil.parameters),
     )
 Y.soilco2.C .= FT(0.000412) # set to atmospheric co2, mol co2 per mol air
-ψ_stem_0 = FT(-1e5 / 9800) # pressure in the leaf divided by rho_liquid*gravitational acceleration [m] 
+ψ_stem_0 = FT(-1e5 / 9800) # pressure in the leaf divided by rho_liquid*gravitational acceleration [m]
 ψ_leaf_0 = FT(-2e5 / 9800)
 ψ_comps = n_stem > 0 ? [ψ_stem_0, ψ_leaf_0] : ψ_leaf_0
 
@@ -302,10 +308,15 @@ saving_cb = ClimaLand.NonInterpSavingCallback(sv, saveat)
 updateat = Array(t0:DATA_DT:tf)
 updatefunc = ClimaLand.make_update_drivers(atmos, radiation)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
-cb = SciMLBase.CallbackSet(driver_cb, saving_cb)
+cb = SciMLBase.CallbackSet(driver_cb, saving_cb);
 
+# Problem definition and solve
 prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction((T_exp!) = exp_tendency!),
+    CTS.ClimaODEFunction(
+        T_exp! = exp_tendency!,
+        T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
+        dss! = ClimaLand.dss!,
+    ),
     Y,
     (t0, tf),
     p,

experiments/integrated/fluxnet/run_fluxnet.jl

@@ -212,6 +212,12 @@ land = SoilCanopyModel{FT}(;
 )
 Y, p, cds = initialize(land)
 exp_tendency! = make_exp_tendency(land)
+imp_tendency! = make_imp_tendency(land)
+tendency_jacobian! = make_tendency_jacobian(land);
+jac_kwargs = (;
+    jac_prototype = ClimaLand.ImplicitEquationJacobian(Y),
+    Wfact = tendency_jacobian!,
+);
 
 #Initial conditions
 Y.soil.ϑ_l =
@@ -234,7 +240,7 @@ Y.soil.ρe_int =
         Ref(land.soil.parameters),
     )
 Y.soilco2.C .= FT(0.000412) # set to atmospheric co2, mol co2 per mol air
-ψ_stem_0 = FT(-1e5 / 9800) # pressure in the leaf divided by rho_liquid*gravitational acceleration [m] 
+ψ_stem_0 = FT(-1e5 / 9800) # pressure in the leaf divided by rho_liquid*gravitational acceleration [m]
 ψ_leaf_0 = FT(-2e5 / 9800)
 ψ_comps = n_stem > 0 ? [ψ_stem_0, ψ_leaf_0] : ψ_leaf_0
 
@@ -265,7 +271,11 @@ driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 cb = SciMLBase.CallbackSet(driver_cb, saving_cb)
 
 prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction((T_exp!) = exp_tendency!),
+    CTS.ClimaODEFunction(
+        T_exp! = exp_tendency!,
+        T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
+        dss! = ClimaLand.dss!,
+    ),
     Y,
     (t0, tf),
     p,

experiments/integrated/performance/conservation/ozark_conservation.jl

@@ -46,8 +46,8 @@ for float_type in (Float32, Float64)
     prob = SciMLBase.ODEProblem(
         CTS.ClimaODEFunction(
             T_exp! = exp_tendency!,
+            T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
             dss! = ClimaLand.dss!,
-            T_imp! = nothing,
         ),
         Y,
         (t0, tf),

experiments/integrated/performance/conservation/ozark_conservation_setup.jl

@@ -192,8 +192,15 @@ land = SoilCanopyModel{FT}(;
     canopy_model_args = canopy_model_args,
 )
 exp_tendency! = make_exp_tendency(land)
+imp_tendency! = make_imp_tendency(land);
+tendency_jacobian! = make_tendency_jacobian(land);
 set_initial_cache! = make_set_initial_cache(land)
 Y, p, cds = initialize(land)
+jac_kwargs = (;
+    jac_prototype = ClimaLand.ImplicitEquationJacobian(Y),
+    Wfact = tendency_jacobian!,
+);
+
 #Initial conditions
 Y.soil.ϑ_l = drivers.SWC.values[1 + Int(round(t0 / 1800))] # Get soil water content at t0
 # recalling that the data is in intervals of 1800 seconds. Both the data

experiments/integrated/performance/profile_allocations.jl

@@ -325,26 +325,43 @@ land = SoilCanopyModel{FT}(;
     canopy_component_args = canopy_component_args,
     canopy_model_args = canopy_model_args,
 )
+
+# Define explicit and implicit tendencies, and the jacobian
 exp_tendency! = make_exp_tendency(land)
+imp_tendency! = make_imp_tendency(land);
+tendency_jacobian! = make_tendency_jacobian(land);
+
 # Set up timestepping and simulation callbacks
 dt = Float64(150)
 tf = Float64(t0 + 100dt)
 saveat = Array(t0:dt:tf)
 updateat = Array(t0:dt:tf)
-timestepper = CTS.RK4()
-ode_algo = CTS.ExplicitAlgorithm(timestepper)
+timestepper = CTS.ARS343()
+ode_algo = CTS.IMEXAlgorithm(
+    timestepper,
+    CTS.NewtonsMethod(
+        max_iters = 1,
+        update_j = CTS.UpdateEvery(CTS.NewNewtonIteration),
+    ),
+);
 
 updatefunc = ClimaLand.make_update_drivers(atmos, radiation)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 # Set initial conditions
 Y, p = set_initial_conditions(land, t0)
 
+# Set up jacobian info
+jac_kwargs = (;
+    jac_prototype = ClimaLand.ImplicitEquationJacobian(Y),
+    Wfact = tendency_jacobian!,
+);
+
 # Solve simulation
 prob = SciMLBase.ODEProblem(
     CTS.ClimaODEFunction(
         T_exp! = exp_tendency!,
+        T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
         dss! = ClimaLand.dss!,
-        T_imp! = nothing,
     ),
     Y,
     (t0, tf),

lib/ClimaLandSimulations/src/utilities/make_timestepper.jl

@@ -4,20 +4,33 @@ export make_timestepper
     make_timestepper(site_setup_out;
         N_spinup_days = 30,
         N_days_sim = 30,
-        timestepper = CTS.RK4(),
-        ode_algo = CTS.ExplicitAlgorith(timestepper)
-        )
+        timestepper = CTS.ARS343(),
+        ode_algo = CTS.IMEXAlgorithm(
+            timestepper,
+            CTS.NewtonsMethod(
+                max_iters = 1,
+                update_j = CTS.UpdateEvery(CTS.NewNewtonIteration),
+            ),
+        ),
+    )
 
-Define the setup for the simulation timestepper. 
-the default timestepper is 4th order Runge Kutta method,
- other timesteppers from ClimaTimeSteppers.jl can be used. 
+Define the setup for the simulation timestepper.
+The default timestepper (ARS343) is an IMEX ARK algorithm with
+3 implicit stages, 4 implicit stages, and 3rd order accuracy.
+Other IMEX timesteppers from ClimaTimeSteppers.jl can be used.
 """
 function make_timestepper(
     site_setup_out;
     N_spinup_days = 30,
     N_days_sim = 30,
-    timestepper = CTS.RK4(),
-    ode_algo = CTS.ExplicitAlgorithm(timestepper),
+    timestepper = CTS.ARS343(),
+    ode_algo = CTS.IMEXAlgorithm(
+        timestepper,
+        CTS.NewtonsMethod(
+            max_iters = 1,
+            update_j = CTS.UpdateEvery(CTS.NewNewtonIteration),
+        ),
+    ),
 )
     N_days = N_spinup_days + N_days_sim
     tf = Float64(site_setup_out.t0 + 3600 * 24 * N_days)

src/ClimaLand.jl

@@ -23,7 +23,7 @@ include("shared_utilities/models.jl")
 include("shared_utilities/drivers.jl")
 include("shared_utilities/boundary_conditions.jl")
 include("shared_utilities/sources.jl")
-include("shared_utilities/implicit_tendencies.jl")
+include("shared_utilities/implicit_timestepping.jl")
 include("standalone/Bucket/Bucket.jl")
 
 """
@@ -134,27 +134,18 @@ end
 
 function make_imp_tendency(land::AbstractLandModel)
     components = land_components(land)
-
-    # If all component models are stepped explicitly, do nothing in imp_tendency!
-    if all(c -> typeof(c) .<: AbstractExpModel, components)
-        function do_nothing_imp_tendency!(dY, Y, p, t) end
-        return do_nothing_imp_tendency!
-    else
-        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
+    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
-        return imp_tendency!
     end
+    return imp_tendency!
 end
 
 function make_exp_tendency(land::AbstractLandModel)
@@ -197,6 +188,18 @@ function make_update_boundary_fluxes(land::AbstractLandModel)
     return update_boundary_fluxes!
 end
 
+function make_update_jacobian(land::AbstractLandModel)
+    components = land_components(land)
+    update_jacobian_function_list =
+        map(x -> make_update_jacobian(getproperty(land, x)), components)
+    function update_jacobian!(jacobian, Y, p, dtγ, t)
+        for f! in update_jacobian_function_list
+            f!(jacobian, Y, p, dtγ, t)
+        end
+    end
+    return update_jacobian!
+end
+
 """
     land_components(land::AbstractLandModel)