use imp solver for all SoilCanopyModel runs

CliMA · Jun 12, 2024 · cf46982 · cf46982
1 parent 61cbdeb
commit cf46982
Show file tree

Hide file tree

Showing 6 changed files with 69 additions and 18 deletions.
diff --git a/docs/tutorials/integrated/soil_canopy_tutorial.jl b/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,

diff --git a/experiments/integrated/fluxnet/US-Var/US-Var_simulation.jl b/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
diff --git a/experiments/integrated/performance/conservation/ozark_conservation.jl b/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),

diff --git a/experiments/integrated/performance/conservation/ozark_conservation_setup.jl b/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

diff --git a/experiments/integrated/performance/profile_allocations.jl b/experiments/integrated/performance/profile_allocations.jl
@@ -325,14 +325,29 @@ 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);
+jac_kwargs = (;
+    jac_prototype = ClimaLand.ImplicitEquationJacobian(Y),
+    Wfact = tendency_jacobian!,
+);
+
 # 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)
@@ -343,8 +358,8 @@ Y, p = set_initial_conditions(land, t0)
 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),

diff --git a/lib/ClimaLandSimulations/src/utilities/make_timestepper.jl b/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 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)