Use ClimaTimeSteppers in CUDA shallow water example

CliMA · Jul 28, 2023 · 93324c1 · 93324c1
1 parent 85ea29c
commit 93324c1
Showing 1 changed file with 10 additions and 5 deletions.
diff --git 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,