From 7d18e7e22810ba29e64fb0c4cef0429f71322609 Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Tue, 11 Jul 2023 08:58:12 -0700
Subject: [PATCH] Fixes

---
 examples/hybrid/sphere/balanced_flow_rhoe.jl |  4 +++-
 examples/sphere/shallow_water_cuda.jl        | 15 ++++++++++-----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/examples/hybrid/sphere/balanced_flow_rhoe.jl b/examples/hybrid/sphere/balanced_flow_rhoe.jl
index f0b8fc4768..aa5ffa2e74 100644
--- 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/sphere/shallow_water_cuda.jl b/examples/sphere/shallow_water_cuda.jl
index 71fd90637b..20d2dcd8fd 100644
--- a/examples/sphere/shallow_water_cuda.jl
+++ b/examples/sphere/shallow_water_cuda.jl
@@ -3,6 +3,8 @@ using ClimaComms
 using DocStringExtensions
 using LinearAlgebra
 using ClimaTimeSteppers, DiffEqBase
+import OrdinaryDiffEq as ODE
+import ClimaTimeSteppers as CTS
 using DiffEqCallbacks
 using NVTX
 
@@ -564,12 +566,15 @@ function shallow_water_driver_cuda(ARGS, ::Type{FT}) where {FT}
         end
         return nothing
     end
-
-    prob =
-        ODEProblem(ClimaODEFunction(; T_exp! = rhs!), Y, (0.0, T), parameters)
-    integrator = DiffEqBase.init(
+    prob = ODE.ODEProblem(
+        CTS.ClimaODEFunction(; T_exp! = rhs!),
+        Y,
+        (FT(0), T),
+        parameters,
+    )
+    integrator = ODE.init(
         prob,
-        ExplicitAlgorithm(SSP33ShuOsher()),
+        CTS.ExplicitAlgorithm(CTS.SSP33ShuOsher()),
         dt = dt,
         saveat = [],
         progress = true,