use dss in ClimaODEfunction; update docs

Update richards_equation.jl

remove spurious comma

docs/tutorials/Bucket/bucket_tutorial.jl

@@ -301,18 +301,32 @@ set_initial_aux_state!(p, Y, t0);
 # of ordinary differential equations:
 exp_tendency! = make_exp_tendency(model);
 
-# Now we choose our timestepping algorithm.
+# Now we choose our (explicit) timestepping algorithm.
 timestepper = CTS.RK4()
-ode_algo = CTS.ExplicitAlgorithm(timestepper)
-
+ode_algo = CTS.ExplicitAlgorithm(timestepper);
+
+
+# We use the
+# [ClimaTimeSteppers.jl](https://github.com/CliMA/ClimaTimeSteppers.jl)
+# interface for handling the specification of implicitly and explicitly
+# treated terms.
+# To set up the ClimaODEFunction, we must specify:
+# - the ODE function/tendency which is treated explicitly in time
+# - the ODE function/tendency which is treated implicitly in time (none here),
+# - the ClimaLSM.dss! function, which does nothing for single column
+#   domains but carries out the dss step needed for domains with spectral
+#   element discretization (employed by Clima in the horizontal directions)
+clima_ode_function =
+    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLSM.dss!);
 # Then we can set up the simulation and solve it:
 prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!),
+    clima_ode_function,
     Y,
     (t0, tf),
     p,
 );
 
+
 # We need a callback to get and store the auxiliary fields, as they
 # are not stored by default.
 saveat = collect(t0:Δt:tf);

docs/tutorials/Canopy/canopy_tutorial.jl

@@ -248,12 +248,8 @@ canopy = ClimaLSM.Canopy.CanopyModel{FT}(;
     radiation = radiation,
 );
 
-# Initialize the state vectors and obtain the model coordinates, then get the
-# explicit time stepping tendency that updates auxiliary and prognostic
-# variables that are stepped explicitly.
-
+# Initialize the state vectors and obtain the model coordinates
 Y, p, coords = ClimaLSM.initialize(canopy)
-exp_tendency! = make_exp_tendency(canopy);
 
 # Provide initial conditions for the canopy hydraulics model
 
@@ -300,13 +296,25 @@ sv = (;
 )
 cb = ClimaLSM.NonInterpSavingCallback(sv, saveat);
 
-# Select a timestepping algorithm and setup the ODE problem.
-
+# Select an (explicit) timestepping algorithm and setup the ODE problem.
 timestepper = CTS.RK4();
 ode_algo = CTS.ExplicitAlgorithm(timestepper)
 
+# We use the
+# [ClimaTimeSteppers.jl](https://github.com/CliMA/ClimaTimeSteppers.jl)
+# interface for handling the specification of implicitly and explicitly
+# treated terms.
+# To set up the ClimaODEFunction, we must specify:
+# - the ODE function/tendency which is treated explicitly in time
+# - the ODE function/tendency which is treated implicitly in time (none here),
+# - the ClimaLSM.dss! function, which does nothing for single column
+#   domains but carries out the dss step needed for domains with spectral
+#   element discretization (employed by Clima in the horizontal directions)
+exp_tendency! = make_exp_tendency(canopy);
+clima_ode_function =
+    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLSM.dss!);
 prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLSM.dss!),
+    clima_ode_function,
     Y,
     (t0, tf),
     p,

docs/tutorials/Canopy/soil_canopy_tutorial.jl

@@ -316,13 +316,8 @@ land = SoilCanopyModel{FT}(;
     canopy_model_args = canopy_model_args,
 );
 
-# Now we can initialize the state vectors and model coordinates, and initialize 
-# the explicit/implicit tendencies as usual. The Richard's equation time 
-# stepping is done implicitly, while the canopy model may be explicitly stepped, 
-# so we use an IMEX (implicit-explicit) scheme for the combined model.
-
+# Now we can initialize the state vectors and model coordinates. 
 Y, p, coords = initialize(land);
-exp_tendency! = make_exp_tendency(land);
 
 # We need to provide initial conditions for the soil and canopy hydraulics 
 # models:
@@ -354,39 +349,48 @@ for i in 1:2
         augmented_liquid_fraction.(plant_ν, S_l_ini[i])
 end;
 
-# Now set the initial conditions for the auxiliary variables for the combined soil and plant model.
-
-t0 = FT(0)
-set_initial_aux_state! = make_set_initial_aux_state(land)
-set_initial_aux_state!(p, Y, t0);
-
 # Select the timestepper and solvers needed for the specific problem. Specify the time range and dt
-# value over which to perform the simulation.
-
 t0 = FT(150 * 3600 * 24)# start mid year
 N_days = 100
 tf = t0 + FT(3600 * 24 * N_days)
-dt = FT(30)
-n = 120
-saveat = Array(t0:(n * dt):tf)
 
+# Now set the initial conditions for the auxiliary variables for the combined soil and plant model.
+set_initial_aux_state! = make_set_initial_aux_state(land)
+set_initial_aux_state!(p, Y, t0);
+
+# Pick the timestepper, timestep. Here we choose an explicit algorithm,
+# which is appropriate for this model.
+dt = FT(30)
 timestepper = CTS.RK4()
 ode_algo = CTS.ExplicitAlgorithm(timestepper);
 
-# And now perform the simulation as always.
-
+# By default, only the state Y is saved. We'd like to save `p` as well,
+# so we can define a callback which does so here:
+n = 120
+saveat = Array(t0:(n * dt):tf)
 sv = (;
     t = Array{FT}(undef, length(saveat)),
     saveval = Array{ClimaCore.Fields.NamedTuple}(undef, length(saveat)),
 )
-cb = ClimaLSM.NonInterpSavingCallback(sv, saveat)
+cb = ClimaLSM.NonInterpSavingCallback(sv, saveat);
+
+# We use the
+# [ClimaTimeSteppers.jl](https://github.com/CliMA/ClimaTimeSteppers.jl)
+# interface for handling the specification of implicitly and explicitly
+# treated terms.
+# To set up the ClimaODEFunction, we must specify:
+# - the ODE function/tendency which is treated explicitly in time
+# - the ODE function/tendency which is treated implicitly in time (none here),
+# - the ClimaLSM.dss! function, which does nothing for single column
+#   domains but carries out the dss step needed for domains with spectral
+#   element discretization (employed by Clima in the horizontal directions)
+exp_tendency! = make_exp_tendency(land);
+clima_ode_function =
+    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLSM.dss!)
+prob = SciMLBase.ODEProblem(clima_ode_function, Y, (t0, tf), p);
 
-prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!),
-    Y,
-    (t0, tf),
-    p,
-);
+# Now, wrap the problem, algorithm, timestep, and callback together
+# to carry out the simulation.
 sol = SciMLBase.solve(
     prob,
     ode_algo;

docs/tutorials/Soil/freezing_front.jl

@@ -200,26 +200,40 @@ end
 
 init_soil!(Y, coords.z, soil.parameters);
 
-# We choose the initial and final simulation times:
+# Here we choose the timestepping algorithm. For the energy & hydrology
+# soil model, we currently only support explicit timestepping.
+# Here, we use the explicit timestepping algorithm RK4, which is suitable
+# for this model because it has no defined implicit tendendency function.
+timestepper = CTS.RK4();
+ode_algo = CTS.ExplicitAlgorithm(timestepper);
+
+# Choose a timestep, integration timespan:
 t0 = FT(0)
-tf = FT(60 * 60 * 50);
+tf = FT(60 * 60 * 50)
+dt = FT(60);
 
 # We set the aux state corresponding to the initial conditions
 # of the state Y:
 set_initial_aux_state! = make_set_initial_aux_state(soil);
 set_initial_aux_state!(p, Y, t0);
-# Create the tendency function, and choose a timestep, and timestepper:
+# To set up the ClimaODEFunction which will be executed to step the
+# system explicitly in time, we must specify:
+# - the ODE function/tendency which is treated explicitly in time
+# - the ClimaLSM.dss! function, which does nothing for single column
+#   domains but carries out the dss step needed for domains with spectral
+#   element discretization (employed by Clima in the horizontal directions)
 exp_tendency! = make_exp_tendency(soil)
-dt = FT(60)
-timestepper = CTS.RK4()
-ode_algo = CTS.ExplicitAlgorithm(timestepper)
+clima_ode_function =
+    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLSM.dss!);
+
 prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLSM.dss!),
+    clima_ode_function,
     Y,
     (t0, tf),
     p,
 );
 
+
 # Now we can solve the problem.
 sol = SciMLBase.solve(prob, ode_algo; dt = dt, saveat = 0:3600:tf);
 

docs/tutorials/Soil/richards_equation.jl

@@ -131,13 +131,6 @@ soil = Soil.RichardsModel{FT}(;
     sources = sources,
 );
 
-# Here we create the explicit and implicit tendencies, which update prognostic
-# variable components that are stepped explicitly and implicitly, respectively.
-# We also create the function which is used to update our Jacobian.
-exp_tendency! = make_exp_tendency(soil);
-imp_tendency! = ClimaLSM.make_imp_tendency(soil);
-update_jacobian! = ClimaLSM.make_update_jacobian(soil);
-
 # # Set up the simulation
 # We can now initialize the prognostic and auxiliary variable vectors, and take
 # a peek at what those variables are:
@@ -162,12 +155,38 @@ tf = FT(60 * 60 * 24 * 36);
 set_initial_aux_state! = make_set_initial_aux_state(soil);
 set_initial_aux_state!(p, Y, t0);
 
-# Next, we turn to timestepping.
+# For solving Richards equation, the `RichardsModel` is set up to
+# treat the vertical transport of water implicitly in time,
+# and the horizontal transport (if applicable) explicitly in time.
+# We use the
+# [ClimaTimeSteppers.jl](https://github.com/CliMA/ClimaTimeSteppers.jl)
+# interface for handling the specification of implicitly and explicitly
+# treated terms.
+# To set up the ClimaODEFunction, we must specify:
+# - the ODE function/tendency which is treated explicitly in time
+# - the ODE function/tendency which is treated implicitly in time,
+#   along with information about the Jacobian of this function
+# - the ClimaLSM.dss! function, which does nothing for single column
+#   domains but carries out the dss step needed for domains with spectral
+#   element discretization (employed by Clima in the horizontal directions)
+# Here we set up the information used for the Jacobian of the implicit
+exp_tendency! = ClimaLSM.make_exp_tendency(soil);
+imp_tendency! = ClimaLSM.make_imp_tendency(soil);
+update_jacobian! = ClimaLSM.make_update_jacobian(soil);
+jac_kwargs =
+    (; jac_prototype = RichardsTridiagonalW(Y), Wfact = update_jacobian!);
+clima_ode_function =
+    CTS.ClimaODEFunction(
+        T_exp! = exp_tendency!,
+        T_imp! = ODE.ODEFunction(imp_tendency!; jac_kwargs...),
+        dss! = ClimaLSM.dss!,
+    )
+
+
+# Now, we choose the imex timestepping algorithm we want to use and the timestep.
 # As usual, your timestep depends on the problem you are solving, the accuracy
 # of the solution required, and the timestepping algorithm you are using.
 dt = FT(1e3);
-
-# Now, we choose the timestepping algorithm we want to use.
 # We'll use the ARS111 algorithm with 1 Newton iteration per timestep;
 # you can also specify a convergence criterion and a maximum number
 # of Newton iterations.
@@ -180,26 +199,12 @@ ode_algo = CTS.IMEXAlgorithm(
     ),
 );
 
-# Here we set up the information used for our Jacobian.
-jac_kwargs =
-    (; jac_prototype = RichardsTridiagonalW(Y), Wfact = update_jacobian!);
-
 # And then we can solve the system of equations, using
 # [SciMLBase.jl](https://github.com/SciML/SciMLBase.jl) and
 # [ClimaTimeSteppers.jl](https://github.com/CliMA/ClimaTimeSteppers.jl).
-prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(
-        T_exp! = exp_tendency!,
-        T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
-        dss! = ClimaLSM.dss!,
-    ),
-    Y,
-    (t0, tf),
-    p,
-);
+prob = SciMLBase.ODEProblem(clima_ode_function, Y, (t0, tf), p);
 sol = SciMLBase.solve(prob, ode_algo; dt = dt, adaptive = false);
 
-
 # # Create some plots
 # We'll plot the moisture content vs depth in the soil, as well as
 # the expected profile of `ϑ_l` in hydrostatic equilibrium.

docs/tutorials/Soil/soil_energy_hydrology.jl

@@ -201,7 +201,6 @@ soil = Soil.EnergyHydrology{FT}(;
     sources = sources,
 );
 
-exp_tendency! = make_exp_tendency(soil);
 # # Set up the simulation
 # We can now initialize the prognostic and auxiliary variable vectors, and take
 # a peek at what those variables are:
@@ -242,6 +241,7 @@ function init_soil!(Y, z, params)
 end
 
 init_soil!(Y, coords.z, soil.parameters);
+<<<<<<< HEAD
 
 # We choose the initial and final simulation times:
 t0 = FT(0)
@@ -252,21 +252,27 @@ tf = FT(60 * 60 * 72);
 set_initial_aux_state! = make_set_initial_aux_state(soil);
 set_initial_aux_state!(p, Y, t0);
 
-# We use [ClimaTimesteppers.jl](https://github.com/CliMA/ClimaTimesteppers.jl) for carrying out the time integration.
-
-# Choose a timestepper and set up the ODE problem:
-dt = FT(30.0);
+# Next we choose the timestepping algorithm. For the energy & hydrology
+# soil model, we currently only support explicit timestepping.
+# Here, we use the explicit timestepping algorithm RK4, which is suitable
+# for this model because it has no defined implicit tendendency function.
+# We use [ClimaTimeSteppers.jl](https://github.com/CliMA/ClimaTimeSteppers.jl)
+# for timestepping.
 timestepper = CTS.RK4();
-ode_algo = CTS.ExplicitAlgorithm(timestepper)
-prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLSM.dss!),
-    Y,
-    (t0, tf),
-    p,
-);
-
-
-#By default, it
+ode_algo = CTS.ExplicitAlgorithm(timestepper);
+dt = FT(30.0);
+# To set up the ClimaODEFunction which will be executed to step the
+# system explicitly in time, we must specify:
+# - the ODE function/tendency which is treated explicitly in time
+# - the ClimaLSM.dss! function, which does nothing for single column
+#   domains but carries out the dss step needed for domains with spectral
+#   element discretization (employed by Clima in the horizontal directions)
+exp_tendency! = make_exp_tendency(soil)
+clima_ode_function =
+    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLSM.dss!);
+prob = SciMLBase.ODEProblem(clima_ode_function, Y, (t0, tf), p);
+
+# By default, the `solve` command
 # only returns Y and t at each time we request output (`saveat`, below). We use
 # a callback in order to also get the auxiliary vector `p` back:
 saveat = collect(t0:FT(1000 * dt):tf)

experiments/integrated/ozark/ozark.jl

@@ -210,7 +210,7 @@ sv = (;
 cb = ClimaLSM.NonInterpSavingCallback(sv, saveat)
 
 prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!),
+    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLSM.dss!),
     Y,
     (t0, tf),
     p,

experiments/standalone/Biogeochemistry/experiment.jl

@@ -153,7 +153,7 @@ saved_values = (;
 cb = ClimaLSM.NonInterpSavingCallback(saved_values, saveat)
 
 prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = Soil_bio_exp_tendency!),
+    CTS.ClimaODEFunction(T_exp! = Soil_bio_exp_tendency!, dss! = ClimaLSM.dss!),
     Y,
     (t0, tf),
     p,