From 0ae6c46b70184629b5967de65dc457980001b761 Mon Sep 17 00:00:00 2001 From: kmdeck Date: Tue, 5 Sep 2023 14:16:30 -0700 Subject: [PATCH] use dss in ClimaODEfunction; update docs Update richards_equation.jl remove spurious comma --- docs/tutorials/Bucket/bucket_tutorial.jl | 22 ++++++-- docs/tutorials/Canopy/canopy_tutorial.jl | 24 +++++--- docs/tutorials/Canopy/soil_canopy_tutorial.jl | 56 ++++++++++--------- docs/tutorials/Soil/freezing_front.jl | 28 +++++++--- docs/tutorials/Soil/richards_equation.jl | 55 +++++++++--------- docs/tutorials/Soil/soil_energy_hydrology.jl | 36 +++++++----- experiments/integrated/ozark/ozark.jl | 2 +- .../standalone/Biogeochemistry/experiment.jl | 2 +- 8 files changed, 138 insertions(+), 87 deletions(-) diff --git a/docs/tutorials/Bucket/bucket_tutorial.jl b/docs/tutorials/Bucket/bucket_tutorial.jl index 1bf85e5def..9fe27c38cc 100644 --- a/docs/tutorials/Bucket/bucket_tutorial.jl +++ b/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); diff --git a/docs/tutorials/Canopy/canopy_tutorial.jl b/docs/tutorials/Canopy/canopy_tutorial.jl index fd262c6ab0..2e87be8a4e 100644 --- a/docs/tutorials/Canopy/canopy_tutorial.jl +++ b/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, diff --git a/docs/tutorials/Canopy/soil_canopy_tutorial.jl b/docs/tutorials/Canopy/soil_canopy_tutorial.jl index 06f40f4ee9..69155c1f62 100644 --- a/docs/tutorials/Canopy/soil_canopy_tutorial.jl +++ b/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; diff --git a/docs/tutorials/Soil/freezing_front.jl b/docs/tutorials/Soil/freezing_front.jl index b36091be17..913bda390b 100644 --- a/docs/tutorials/Soil/freezing_front.jl +++ b/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); diff --git a/docs/tutorials/Soil/richards_equation.jl b/docs/tutorials/Soil/richards_equation.jl index 14f337dd2e..a6d1e9fcfc 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: @@ -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. diff --git a/docs/tutorials/Soil/soil_energy_hydrology.jl b/docs/tutorials/Soil/soil_energy_hydrology.jl index 27690982e5..72076227ea 100644 --- a/docs/tutorials/Soil/soil_energy_hydrology.jl +++ b/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) diff --git a/experiments/integrated/ozark/ozark.jl b/experiments/integrated/ozark/ozark.jl index c6848b0c55..4f382f684e 100644 --- a/experiments/integrated/ozark/ozark.jl +++ b/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, diff --git a/experiments/standalone/Biogeochemistry/experiment.jl b/experiments/standalone/Biogeochemistry/experiment.jl index 5c95874ae3..47f524ef08 100644 --- a/experiments/standalone/Biogeochemistry/experiment.jl +++ b/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,