diff --git a/docs/tutorials/Bucket/bucket_tutorial.jl b/docs/tutorials/Bucket/bucket_tutorial.jl index 12bfd82bac..9f9d0c9cce 100644 --- a/docs/tutorials/Bucket/bucket_tutorial.jl +++ b/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. diff --git a/docs/tutorials/Canopy/canopy_tutorial.jl b/docs/tutorials/Canopy/canopy_tutorial.jl index 9defdd761b..2b165ff4c4 100644 --- a/docs/tutorials/Canopy/canopy_tutorial.jl +++ b/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 diff --git a/docs/tutorials/Canopy/soil_canopy_tutorial.jl b/docs/tutorials/Canopy/soil_canopy_tutorial.jl index 067732ccf3..4fd0e1faa7 100644 --- a/docs/tutorials/Canopy/soil_canopy_tutorial.jl +++ b/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; diff --git a/docs/tutorials/Soil/freezing_front.jl b/docs/tutorials/Soil/freezing_front.jl index 1ee1d5dee2..0462415d59 100644 --- a/docs/tutorials/Soil/freezing_front.jl +++ b/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); diff --git a/docs/tutorials/Soil/richards_equation.jl b/docs/tutorials/Soil/richards_equation.jl index 19b8013e0d..6dbaf4953d 100644 --- a/docs/tutorials/Soil/richards_equation.jl +++ b/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); diff --git a/docs/tutorials/Soil/soil_energy_hydrology.jl b/docs/tutorials/Soil/soil_energy_hydrology.jl index 0dd267c897..504b2d988d 100644 --- a/docs/tutorials/Soil/soil_energy_hydrology.jl +++ b/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); diff --git a/experiments/integrated/ozark/ozark.jl b/experiments/integrated/ozark/ozark.jl index 8a79db2162..3478691308 100644 --- a/experiments/integrated/ozark/ozark.jl +++ b/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; diff --git a/experiments/standalone/Biogeochemistry/experiment.jl b/experiments/standalone/Biogeochemistry/experiment.jl index 2806af0c26..d02d652dad 100644 --- a/experiments/standalone/Biogeochemistry/experiment.jl +++ b/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,