Merge #1401

1401: Use ClimaTimeSteppers r=charleskawczynski a=charleskawczynski This is a peel off from #1358. Co-authored-by: Charles Kawczynski <[email protected]>
CliMA · Aug 1, 2023 · f84433c · f84433c
2 parents d772077 + f0941a0
commit f84433c
Show file tree

Hide file tree

Showing 11 changed files with 21 additions and 16 deletions.
diff --git a/examples/hybrid/driver.jl b/examples/hybrid/driver.jl
@@ -24,6 +24,7 @@ postprocessing(sol, output_dir) = nothing
 
 ################################################################################
 
+import ClimaTimeSteppers as CTS
 using ClimaComms
 const comms_ctx = ClimaComms.context()
 is_distributed = comms_ctx isa ClimaComms.MPICommsContext
@@ -145,13 +146,14 @@ end
 callback =
     CallbackSet(dss_callback, save_to_disk_callback, additional_callbacks...)
 
-problem = SplitODEProblem(
-    ODEFunction(
-        implicit_tendency!;
-        jac_kwargs(ode_algo, Y, jacobian_flags)...,
-        tgrad = (∂Y∂t, Y, p, t) -> (∂Y∂t .= FT(0)),
+problem = ODE.ODEProblem(
+    CTS.ClimaODEFunction(;
+        T_imp! = ODEFunction(
+            implicit_tendency!;
+            jac_kwargs(ode_algo, Y, jacobian_flags)...,
+        ),
+        T_exp! = remaining_tendency!,
     ),
-    remaining_tendency!,
     Y,
     (t_start, t_end),
     p,

diff --git a/examples/hybrid/plane/inertial_gravity_wave.jl b/examples/hybrid/plane/inertial_gravity_wave.jl
@@ -60,10 +60,11 @@ fps = 2
 horizontal_mesh = periodic_line_mesh(; x_max, x_elem = x_elem)
 
 # Additional values required for driver
+# dt may need tweaking for is_small_scale = false
 dt = is_small_scale ? FT(1.5) : FT(20)
 t_end = is_small_scale ? FT(60 * 60 * 0.5) : FT(60 * 60 * 8)
 dt_save_to_sol = t_end / (animation_duration * fps)
-ode_algorithm = OrdinaryDiffEq.Rosenbrock23
+ode_algorithm = CTS.SSP333
 jacobian_flags = (;
     ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode = ᶜ𝔼_name == :ρe ? :no_∂ᶜp∂ᶜK : :exact,
     ∂ᶠ𝕄ₜ∂ᶜρ_mode = :exact,

diff --git a/examples/hybrid/sphere/balanced_flow_rhoe.jl b/examples/hybrid/sphere/balanced_flow_rhoe.jl
@@ -1,4 +1,5 @@
 using Test
+import ClimaTimeSteppers as CTS
 using ClimaCorePlots, Plots
 using ClimaCore.DataLayouts
 
@@ -15,7 +16,8 @@ t_end = FT(60 * 60)
 dt = FT(5)
 dt_save_to_sol = FT(50)
 dt_save_to_disk = FT(0) # 0 means don't save to disk
-ode_algorithm = OrdinaryDiffEq.SSPRK33
+ode_algorithm = CTS.SSP33ShuOsher
+jacobian_flags = (; ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode = :exact, ∂ᶠ𝕄ₜ∂ᶜρ_mode = :exact)
 
 additional_cache(ᶜlocal_geometry, ᶠlocal_geometry, dt) = merge(
     hyperdiffusion_cache(ᶜlocal_geometry, ᶠlocal_geometry; κ₄ = FT(2e17)),

diff --git a/examples/hybrid/sphere/balanced_flow_rhotheta.jl b/examples/hybrid/sphere/balanced_flow_rhotheta.jl
@@ -14,7 +14,7 @@ t_end = FT(60 * 60 * 24 * 10)
 dt = FT(400)
 dt_save_to_sol = FT(60 * 60 * 24)
 dt_save_to_disk = FT(0) # 0 means don't save to disk
-ode_algorithm = OrdinaryDiffEq.Rosenbrock23
+ode_algorithm = CTS.SSP333
 jacobian_flags = (; ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode = :exact, ∂ᶠ𝕄ₜ∂ᶜρ_mode = :exact)
 
 additional_cache(ᶜlocal_geometry, ᶠlocal_geometry, dt) = merge(

diff --git a/examples/hybrid/sphere/baroclinic_wave_rhoe.jl b/examples/hybrid/sphere/baroclinic_wave_rhoe.jl
@@ -14,7 +14,7 @@ t_end = FT(60 * 60 * 24 * 10)
 dt = FT(400)
 dt_save_to_sol = FT(60 * 60 * 24)
 dt_save_to_disk = FT(0) # 0 means don't save to disk
-ode_algorithm = OrdinaryDiffEq.Rosenbrock23
+ode_algorithm = CTS.SSP333
 jacobian_flags = (; ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode = :no_∂ᶜp∂ᶜK, ∂ᶠ𝕄ₜ∂ᶜρ_mode = :exact)
 
 additional_cache(ᶜlocal_geometry, ᶠlocal_geometry, dt) = merge(

diff --git a/examples/hybrid/sphere/baroclinic_wave_rhotheta.jl b/examples/hybrid/sphere/baroclinic_wave_rhotheta.jl
@@ -14,7 +14,7 @@ t_end = FT(60 * 60 * 24 * 10)
 dt = FT(400)
 dt_save_to_sol = FT(60 * 60 * 24)
 dt_save_to_disk = FT(0) # 0 means don't save to disk
-ode_algorithm = OrdinaryDiffEq.Rosenbrock23
+ode_algorithm = CTS.SSP333
 jacobian_flags = (; ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode = :exact, ∂ᶠ𝕄ₜ∂ᶜρ_mode = :exact)
 
 additional_cache(ᶜlocal_geometry, ᶠlocal_geometry, dt) = merge(

diff --git a/examples/hybrid/sphere/held_suarez_rhoe.jl b/examples/hybrid/sphere/held_suarez_rhoe.jl
@@ -14,7 +14,7 @@ t_end = FT(60 * 60 * 24 * 10)
 dt = FT(400)
 dt_save_to_sol = FT(60 * 60 * 24)
 dt_save_to_disk = FT(0) # 0 means don't save to disk
-ode_algorithm = OrdinaryDiffEq.Rosenbrock23
+ode_algorithm = CTS.SSP333
 jacobian_flags = (; ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode = :no_∂ᶜp∂ᶜK, ∂ᶠ𝕄ₜ∂ᶜρ_mode = :exact)
 
 # Additional values required for driver

diff --git a/examples/hybrid/sphere/held_suarez_rhoe_int.jl b/examples/hybrid/sphere/held_suarez_rhoe_int.jl
@@ -14,7 +14,7 @@ t_end = FT(60 * 60 * 24 * 10)
 dt = FT(400)
 dt_save_to_sol = FT(60 * 60 * 24)
 dt_save_to_disk = FT(0) # 0 means don't save to disk
-ode_algorithm = OrdinaryDiffEq.Rosenbrock23
+ode_algorithm = CTS.SSP333
 jacobian_flags = (; ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode = :exact, ∂ᶠ𝕄ₜ∂ᶜρ_mode = :exact)
 
 additional_cache(ᶜlocal_geometry, ᶠlocal_geometry, dt) = merge(

diff --git a/examples/hybrid/sphere/held_suarez_rhotheta.jl b/examples/hybrid/sphere/held_suarez_rhotheta.jl
@@ -15,7 +15,7 @@ t_end = FT(60 * 60 * 24 * 10)
 dt = FT(400)
 dt_save_to_sol = FT(60 * 60 * 24)
 dt_save_to_disk = FT(0) # 0 means don't save to disk
-ode_algorithm = OrdinaryDiffEq.Rosenbrock23
+ode_algorithm = CTS.SSP333
 jacobian_flags = (; ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode = :exact, ∂ᶠ𝕄ₜ∂ᶜρ_mode = :exact)
 
 additional_cache(ᶜlocal_geometry, ᶠlocal_geometry, dt) = merge(

diff --git a/examples/hybrid/sphere/held_suarez_rhotheta_scaling.jl b/examples/hybrid/sphere/held_suarez_rhotheta_scaling.jl
@@ -16,7 +16,7 @@ t_end = FT(60 * 60 * 1)
 dt = FT(100)
 dt_save_to_sol = FT(60 * 60 * 24)
 dt_save_to_disk = FT(0) # 0 means don't save to disk
-ode_algorithm = OrdinaryDiffEq.Rosenbrock23
+ode_algorithm = CTS.SSP333
 jacobian_flags = (; ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode = :exact, ∂ᶠ𝕄ₜ∂ᶜρ_mode = :exact)
 
 additional_cache(ᶜlocal_geometry, ᶠlocal_geometry, dt) = merge(

diff --git a/examples/hybrid/sphere/held_suarez_rhotheta_tempest.jl b/examples/hybrid/sphere/held_suarez_rhotheta_tempest.jl
@@ -14,7 +14,7 @@ t_end = FT(60 * 60 * 24 * 10)
 dt = FT(400)
 dt_save_to_sol = FT(60 * 60 * 24)
 dt_save_to_disk = FT(0) # 0 means don't save to disk
-ode_algorithm = OrdinaryDiffEq.Rosenbrock23
+ode_algorithm = CTS.SSP333
 jacobian_flags = (; ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode = :exact, ∂ᶠ𝕄ₜ∂ᶜρ_mode = :exact)
 
 additional_cache(ᶜlocal_geometry, ᶠlocal_geometry, dt) = merge(