use dss in ClimaODEfunction; update docs

docs/tutorials/Bucket/bucket_tutorial.jl

@@ -294,13 +294,23 @@ 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)
-
-# Then we can set up the simulation and solve it:
-prob =
-    ODE.ODEProblem(CTS.ClimaODEFunction(T_exp! = exp_tendency!), Y, (t0, tf), p);
+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!)
+prob = ODE.ODEProblem(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.

docs/tutorials/Canopy/canopy_tutorial.jl

@@ -249,12 +249,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
 
@@ -301,17 +297,24 @@ 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)
 
-prob = ODE.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLSM.dss!),
-    Y,
-    (t0, tf),
-    p,
-);
+# 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 = ODE.ODEProblem(clima_ode_function, Y, (t0, tf), p);
 
 # Now, we can solve the problem and store the model data in the saveat array, 
 # using [`OrdinaryDiffEq.jl`](https://github.com/SciML/OrdinaryDiffEq.jl) and

docs/tutorials/Canopy/soil_canopy_tutorial.jl

@@ -317,13 +317,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:
@@ -355,35 +350,49 @@ for i in 1:2
         augmented_liquid_fraction.(plant_ν, S_l_ini[i])
 end;
 
-# Now initialize 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
+# 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 initialize 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 = ODE.ODEProblem(clima_ode_function, Y, (t0, tf), p);
 
-prob =
-    ODE.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 = ODE.solve(
     prob,
     ode_algo;

docs/tutorials/Soil/freezing_front.jl

@@ -200,21 +200,29 @@ end
 
 init_soil!(Y, coords.z, soil.parameters);
 
-# Create the tendency function, and choose a timestep, integration timespan, etc:
-exp_tendency! = make_exp_tendency(soil)
+# 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)
 dt = FT(60)
 tf = FT(3600 * 50)
 
-# Choose a timestepper and set up the ODE problem:
-timestepper = CTS.RK4();
-ode_algo = CTS.ExplicitAlgorithm(timestepper)
-prob = ODE.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLSM.dss!),
-    Y,
-    (t0, tf),
-    p,
-);
+
+# 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 = ODE.ODEProblem(clima_ode_function, Y, (t0, tf), p);
 
 # Now we can solve the problem.
 sol = ODE.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:
@@ -153,14 +146,41 @@ coords |> propertynames
 # conditions.
 Y.soil.ϑ_l .= FT(0.494);
 
-# Next, we turn to timestepping. We choose the initial and final times, as well as a 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.
-t0 = FT(0)
+# 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!,
+    ),
+
+    # Next, we turn to timestepping. We choose the initial and final times, as well as a 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.
+    t0 = FT(0)
 tf = FT(60 * 60 * 24 * 36)
 dt = FT(1e3);
 
-# Now, we choose the timestepping algorithm we want to use.
+# Now, we choose the imex 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.
@@ -173,23 +193,7 @@ 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
-# [OrdinaryDiffEq.jl](https://github.com/SciML/OrdinaryDiffEq.jl) and
-# [ClimaTimeSteppers.jl](https://github.com/CliMA/ClimaTimeSteppers.jl).
-prob = ODE.ODEProblem(
-    CTS.ClimaODEFunction(
-        T_exp! = exp_tendency!,
-        T_imp! = ODE.ODEFunction(imp_tendency!; jac_kwargs...),
-        dss! = ClimaLSM.dss!,
-    ),
-    Y,
-    (t0, tf),
-    p,
-);
+prob = ODE.ODEProblem(clima_ode_function, Y, (t0, tf), p);
 sol = ODE.solve(prob, ode_algo; dt = dt, adaptive = false);
 
 

docs/tutorials/Soil/soil_energy_hydrology.jl

@@ -202,7 +202,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:
@@ -243,18 +242,40 @@ function init_soil!(Y, z, params)
 end
 
 init_soil!(Y, coords.z, soil.parameters);
-# Updating the auxiliary to the initial state is not necessary, because
-# the first step of the tendency function is to update the auxiliary state.
-# However, we want to plot the initial state in this tutorial, so
-# we update it here.
-update_aux! = make_update_aux(soil)
-update_aux!(p, Y, FT(0.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);
+
+# Choose a timestep, integration timespan:
 t0 = FT(0)
 tf = FT(60 * 60 * 72)
 dt = FT(30.0);
 
-# We use [OrdinaryDiffEq.jl](https://github.com/SciML/OrdinaryDiffEq.jl) for carrying out the time integration. By default, it
+# Updating the auxiliary to the initial state is not necessary, because
+# the first step of the tendency function is to update the auxiliary state.
+# However, we want to plot the initial state in this tutorial, so
+# we update it here.
+update_aux! = make_update_aux(soil)
+update_aux!(p, Y, t0);
+
+# 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 = ODE.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)
@@ -264,15 +285,6 @@ saved_values = (;
 );
 cb = ClimaLSM.NonInterpSavingCallback(saved_values, saveat);
 
-# Choose a timestepper and set up the ODE problem:
-timestepper = CTS.RK4();
-ode_algo = CTS.ExplicitAlgorithm(timestepper)
-prob = ODE.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLSM.dss!),
-    Y,
-    (t0, tf),
-    p,
-);
 
 # Now we can solve the problem.
 sol = ODE.solve(prob, ode_algo; dt = dt, saveat = saveat, callback = cb);

experiments/integrated/ozark/ozark.jl

@@ -211,8 +211,12 @@ sv = (;
 )
 cb = ClimaLSM.NonInterpSavingCallback(sv, saveat)
 
-prob =
-    ODE.ODEProblem(CTS.ClimaODEFunction(T_exp! = exp_tendency!), Y, (t0, tf), p);
+prob = ODE.ODEProblem(
+    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLSM.dss!),
+    Y,
+    (t0, tf),
+    p,
+);
 sol = ODE.solve(
     prob,
     ode_algo;

experiments/standalone/Biogeochemistry/experiment.jl

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