Skip to content

Commit

Permalink
use dss in ClimaODEfunction; update docs
Browse files Browse the repository at this point in the history
  • Loading branch information
kmdeck committed Sep 5, 2023
1 parent ad81a2f commit 09716b2
Show file tree
Hide file tree
Showing 8 changed files with 151 additions and 101 deletions.
22 changes: 16 additions & 6 deletions docs/tutorials/Bucket/bucket_tutorial.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
29 changes: 16 additions & 13 deletions docs/tutorials/Canopy/canopy_tutorial.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
53 changes: 31 additions & 22 deletions docs/tutorials/Canopy/soil_canopy_tutorial.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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;
Expand Down
30 changes: 19 additions & 11 deletions docs/tutorials/Soil/freezing_front.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
62 changes: 33 additions & 29 deletions docs/tutorials/Soil/richards_equation.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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.
Expand All @@ -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);


Expand Down
46 changes: 29 additions & 17 deletions docs/tutorials/Soil/soil_energy_hydrology.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand All @@ -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);
Expand Down
8 changes: 6 additions & 2 deletions experiments/integrated/ozark/ozark.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
2 changes: 1 addition & 1 deletion experiments/standalone/Biogeochemistry/experiment.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit 09716b2

Please sign in to comment.