updated experiments

CliMA · Sep 13, 2023 · b80ca0f · b80ca0f
1 parent e377256
commit b80ca0f
Show file tree

Hide file tree

Showing 7 changed files with 19 additions and 125 deletions.
diff --git a/experiments/integrated/ozark/conservation/ozark_conservation.jl b/experiments/integrated/ozark/conservation/ozark_conservation.jl
@@ -168,8 +168,6 @@ land = SoilCanopyModel{FT}(;
     canopy_model_args = canopy_model_args,
 )
 Y, p, cds = initialize(land)
-exp_tendency! = make_exp_tendency(land)
-
 #Initial conditions
 Y.soil.ϑ_l = SWC[1 + Int(round(t0 / 1800))] # Get soil water content at t0
 # recalling that the data is in intervals of 1800 seconds. Both the data
@@ -211,17 +209,8 @@ sv = (;
     saveval = Array{NamedTuple}(undef, length(saveat)),
 )
 cb = ClimaLSM.NonInterpSavingCallback(sv, saveat)
-
-prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(
-        T_exp! = exp_tendency!,
-        dss! = ClimaLSM.dss!,
-        T_imp! = nothing,
-    ),
-    Y,
-    (t0, tf),
-    p,
-);
+clima_ode_function = ClimaLSM.get_ClimaODEFunction(land)
+prob = SciMLBase.ODEProblem(clima_ode_function, Y, (t0, tf), p);
 sol = SciMLBase.solve(
     prob,
     ode_algo;

diff --git a/experiments/integrated/ozark/hydrology_only/ozark.jl b/experiments/integrated/ozark/hydrology_only/ozark.jl
@@ -158,8 +158,6 @@ land = SoilPlantHydrologyModel{FT}(;
     canopy_model_args = canopy_model_args,
 )
 Y, p, cds = initialize(land)
-exp_tendency! = make_exp_tendency(land)
-imp_tendency! = make_imp_tendency(land)
 
 #Initial conditions
 Y.soil.ϑ_l = SWC[Int(round(t0 / 1800))] # Get soil water content at t0
@@ -182,9 +180,6 @@ end
 set_initial_aux_state! = make_set_initial_aux_state(land)
 set_initial_aux_state!(p, Y, t0);
 
-# Set up timestepper and jacobian for soil
-update_jacobian! = make_update_jacobian(land.soil)
-
 ode_algo = CTS.IMEXAlgorithm(
     timestepper,
     CTS.NewtonsMethod(
@@ -194,26 +189,15 @@ ode_algo = CTS.IMEXAlgorithm(
     ),
 )
 
-W = RichardsTridiagonalW(Y)
-jac_kwargs = (; jac_prototype = W, Wfact = update_jacobian!)
 
 # Simulation
 sv = (;
     t = Array{FT}(undef, length(saveat)),
     saveval = Array{NamedTuple}(undef, length(saveat)),
 )
 cb = ClimaLSM.NonInterpSavingCallback(sv, saveat)
-
-prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(
-        T_exp! = exp_tendency!,
-        T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
-        dss! = ClimaLSM.dss!,
-    ),
-    Y,
-    (t0, tf),
-    p,
-)
+clima_ode_function = ClimaLSM.get_ClimaODEFunction(land, Y)
+prob = SciMLBase.ODEProblem(clima_ode_function, Y, (t0, tf), p)
 sol = SciMLBase.solve(
     prob,
     ode_algo;

diff --git a/experiments/integrated/ozark/ozark.jl b/experiments/integrated/ozark/ozark.jl
@@ -165,7 +165,6 @@ land = SoilCanopyModel{FT}(;
     canopy_model_args = canopy_model_args,
 )
 Y, p, cds = initialize(land)
-exp_tendency! = make_exp_tendency(land)
 
 #Initial conditions
 Y.soil.ϑ_l = SWC[1 + Int(round(t0 / 1800))] # Get soil water content at t0
@@ -208,13 +207,8 @@ sv = (;
     saveval = Array{NamedTuple}(undef, length(saveat)),
 )
 cb = ClimaLSM.NonInterpSavingCallback(sv, saveat)
-
-prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = exp_tendency!, dss! = ClimaLSM.dss!),
-    Y,
-    (t0, tf),
-    p,
-);
+clima_ode_function = ClimaLSM.get_ClimaODEFunction(land)
+prob = SciMLBase.ODEProblem(clima_ode_function, Y, (t0, tf), p);
 sol = SciMLBase.solve(
     prob,
     ode_algo;

diff --git a/experiments/standalone/Biogeochemistry/experiment.jl b/experiments/standalone/Biogeochemistry/experiment.jl
@@ -137,7 +137,7 @@ init_soil!(Y, z, model.soil.parameters)
 init_co2!(Y, z)
 t0 = FT(0.0)
 set_initial_aux_state!(p, Y, t0);
-Soil_bio_exp_tendency! = make_exp_tendency(model)
+clima_ode_function = ClimaLSM.get_ClimaODEFunction(model)
 
 tf = FT(10000)
 dt = FT(10)
@@ -152,12 +152,7 @@ saved_values = (;
 )
 cb = ClimaLSM.NonInterpSavingCallback(saved_values, saveat)
 
-prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = Soil_bio_exp_tendency!, 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, callback = cb)
 
 # Animation

diff --git a/experiments/standalone/Soil/evaporation.jl b/experiments/standalone/Soil/evaporation.jl
@@ -154,16 +154,10 @@ set_initial_aux_state!(p, Y, t0);
 
 # Timestepping:
 dt = FT(1)
-soil_exp_tendency! = make_exp_tendency(soil)
 timestepper = CTS.RK4()
 ode_algo = CTS.ExplicitAlgorithm(timestepper)
-
-prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(T_exp! = soil_exp_tendency!, dss! = ClimaLSM.dss!),
-    Y,
-    (t0, tf),
-    p,
-)
+clima_ode_function = ClimaLSM.get_ClimaODEFunction(soil)
+prob = SciMLBase.ODEProblem(clima_ode_function, Y, (t0, tf), p)
 sol = SciMLBase.solve(prob, ode_algo; dt = dt, saveat = 3600)
 
 (; ν, θ_r, d_ds) = soil.parameters

diff --git a/experiments/standalone/Soil/richards_comparison.jl b/experiments/standalone/Soil/richards_comparison.jl
@@ -50,10 +50,6 @@ include(joinpath(pkgdir(ClimaLSM), "parameters", "create_parameters.jl"))
 
     # specify ICs
     Y.soil.ϑ_l .= FT(0.24)
-    exp_tendency! = make_exp_tendency(soil)
-    imp_tendency! = ClimaLSM.make_imp_tendency(soil)
-    update_jacobian! = ClimaLSM.make_update_jacobian(soil)
-
     t0 = FT(0)
     set_initial_aux_state!(p, Y, t0)
     tf = FT(1e6)
@@ -70,21 +66,8 @@ include(joinpath(pkgdir(ClimaLSM), "parameters", "create_parameters.jl"))
             convergence_checker = conv_checker,
         ),
     )
-
-    # set up jacobian info
-    jac_kwargs =
-        (; jac_prototype = RichardsTridiagonalW(Y), Wfact = update_jacobian!)
-
-    prob = SciMLBase.ODEProblem(
-        CTS.ClimaODEFunction(
-            T_exp! = exp_tendency!,
-            T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
-            dss! = ClimaLSM.dss!,
-        ),
-        Y,
-        (t0, tf),
-        p,
-    )
+    clima_ode_function = ClimaLSM.get_ClimaODEFunction(soil, Y)
+    prob = SciMLBase.ODEProblem(clima_ode_function, Y, (t0, tf), p)
     sol = SciMLBase.solve(prob, ode_algo; dt = dt, saveat = 10000)
 
     N = length(sol.t)
@@ -144,10 +127,6 @@ end
 
     # specify ICs
     Y.soil.ϑ_l .= FT(0.1)
-    exp_tendency! = make_exp_tendency(soil)
-    imp_tendency! = ClimaLSM.make_imp_tendency(soil)
-    update_jacobian! = ClimaLSM.make_update_jacobian(soil)
-
     t0 = FT(0)
     set_initial_aux_state!(p, Y, t0)
     tf = FT(60 * 60 * 0.8)
@@ -165,20 +144,8 @@ end
             convergence_checker = conv_checker,
         ),
     )
-    # set up jacobian info
-    jac_kwargs =
-        (; jac_prototype = RichardsTridiagonalW(Y), Wfact = update_jacobian!)
-
-    prob = SciMLBase.ODEProblem(
-        CTS.ClimaODEFunction(
-            T_exp! = exp_tendency!,
-            T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
-            dss! = ClimaLSM.dss!,
-        ),
-        Y,
-        (t0, tf),
-        p,
-    )
+    clima_ode_function = ClimaLSM.get_ClimaODEFunction(soil, Y)
+    prob = SciMLBase.ODEProblem(clima_ode_function, Y, (t0, tf), p)
     sol = SciMLBase.solve(prob, ode_algo; dt = dt, saveat = 60 * dt)
 
     N = length(sol.t)

diff --git a/experiments/standalone/Soil/water_conservation.jl b/experiments/standalone/Soil/water_conservation.jl
@@ -84,11 +84,7 @@ soil = Soil.RichardsModel{FT}(;
     boundary_conditions = boundary_fluxes,
     sources = sources,
 )
-
-exp_tendency! = make_exp_tendency(soil)
 set_initial_aux_state! = make_set_initial_aux_state(soil);
-imp_tendency! = make_imp_tendency(soil)
-update_jacobian! = make_update_jacobian(soil)
 
 rmses = Array{FT}(undef, length(dts))
 mass_errors = Array{FT}(undef, length(dts))
@@ -99,19 +95,8 @@ for i in eachindex(dts)
     @. Y.soil.ϑ_l = FT(0.24)
     set_initial_aux_state!(p, Y, FT(0.0))
 
-    jac_kwargs =
-        (; jac_prototype = RichardsTridiagonalW(Y), Wfact = update_jacobian!)
-
-    prob = SciMLBase.ODEProblem(
-        CTS.ClimaODEFunction(
-            T_exp! = exp_tendency!,
-            T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
-            dss! = ClimaLSM.dss!,
-        ),
-        Y,
-        (t_start, t_end),
-        p,
-    )
+    clima_ode_function = ClimaLSM.get_ClimaODEFunction(soil, Y)
+    prob = SciMLBase.ODEProblem(clima_ode_function, Y, (t_start, t_end), p)
 
     sol = SciMLBase.solve(prob, ode_algo; dt = dt, saveat = dt)
 
@@ -187,11 +172,7 @@ soil_dirichlet = Soil.RichardsModel{FT}(;
     sources = sources,
 )
 
-exp_tendency! = make_exp_tendency(soil_dirichlet)
 set_initial_aux_state! = make_set_initial_aux_state(soil_dirichlet);
-imp_tendency! = make_imp_tendency(soil_dirichlet)
-update_jacobian! = make_update_jacobian(soil_dirichlet)
-update_aux! = make_update_aux(soil_dirichlet)
 
 rmses_dirichlet = Array{FT}(undef, length(dts))
 mass_errors_dirichlet = Array{FT}(undef, length(dts))
@@ -202,19 +183,9 @@ for i in eachindex(dts)
     @. Y.soil.ϑ_l = FT(0.24)
     set_initial_aux_state!(p, Y, FT(0.0))
 
-    jac_kwargs =
-        (; jac_prototype = RichardsTridiagonalW(Y), Wfact = update_jacobian!)
-
-    prob = SciMLBase.ODEProblem(
-        CTS.ClimaODEFunction(
-            T_exp! = exp_tendency!,
-            T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
-            dss! = ClimaLSM.dss!,
-        ),
-        Y,
-        (t_start, t_end),
-        p,
-    )
+    clima_ode_function = ClimaLSM.get_ClimaODEFunction(soiil_dirichlet, Y)
+
+    prob = SciMLBase.ODEProblem(clima_ode_function, Y, (t_start, t_end), p)
     sol = SciMLBase.solve(prob, ode_algo; dt = dt, saveat = dt)
 
     # Calculate water mass balance over entire simulation