From a6599a5661fe0beef0f1425e9e346542f54584d6 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        | 33 +++++++++++++-------
 test/Operators/finitedifference/linsolve.jl  |  8 ++---
 3 files changed, 28 insertions(+), 17 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 c07e5347e2..59122ffbee 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
 
@@ -448,7 +450,9 @@ function set_initial_condition(
     return Y
 end
 
-function rhs!(dYdt, y, parameters, t)
+function rhs!(_dYdt, _y, parameters, t)
+    dYdt = _dYdt.c
+    y = _y.c
     NVTX.@range "rhs!" begin
         (; f, h_s, ghost_buffer, params) = parameters
         (; D₄, g) = params
@@ -487,7 +491,7 @@ function rhs!(dYdt, y, parameters, t)
             NVTX.@range "dss" Spaces.weighted_dss2!(dYdt, ghost_buffer)
         end
     end
-    return dYdt
+    return _dYdt
 end
 
 function shallow_water_driver_cuda(ARGS, ::Type{FT}) where {FT}
@@ -523,12 +527,13 @@ function shallow_water_driver_cuda(ARGS, ::Type{FT}) where {FT}
     space = Spaces.SpectralElementSpace2D(grid_topology, quad)
     f = set_coriolis_parameter(space, test)
     h_s = surface_topography(space, test)
-    Y = set_initial_condition(space, test)
+    Y = Fields.FieldVector(; c = set_initial_condition(space, test))
+    @show eltype(Y)
 
-    ghost_buffer = Spaces.create_dss_buffer(Y)
-    Spaces.weighted_dss_start!(Y, ghost_buffer)
-    Spaces.weighted_dss_internal!(Y, ghost_buffer)
-    Spaces.weighted_dss_ghost!(Y, ghost_buffer)
+    ghost_buffer = Spaces.create_dss_buffer(Y.c)
+    Spaces.weighted_dss_start!(Y.c, ghost_buffer)
+    Spaces.weighted_dss_internal!(Y.c, ghost_buffer)
+    Spaces.weighted_dss_ghost!(Y.c, ghost_buffer)
 
     parameters =
         (; f = f, h_s = h_s, ghost_buffer = ghost_buffer, params = test.params)
@@ -562,12 +567,16 @@ 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(
+    @show (0.0, T)
+    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,
diff --git a/test/Operators/finitedifference/linsolve.jl b/test/Operators/finitedifference/linsolve.jl
index cf6041391b..5ab3e6549d 100644
--- a/test/Operators/finitedifference/linsolve.jl
+++ b/test/Operators/finitedifference/linsolve.jl
@@ -35,7 +35,7 @@ face_space = Spaces.FaceExtrudedFiniteDifferenceSpace(center_space)
 =#
 face_space = Spaces.FaceFiniteDifferenceSpace(center_space)
 
-function _linsolve!(x, A, b, update_matrix = false; kwargs...)
+function test_linsolve!(x, A, b, update_matrix = false; kwargs...)
 
     FT = Spaces.undertype(axes(x.c))
 
@@ -88,11 +88,11 @@ W = SchurComplementW(Y, use_transform, jacobi_flags)
 
 using JET
 using Test
-@time _linsolve!(Y, W, b)
-@time _linsolve!(Y, W, b)
+@time test_linsolve!(Y, W, b)
+@time test_linsolve!(Y, W, b)
 
 @testset "JET test for `apply` in linsolve! kernel" begin
-    @test_opt _linsolve!(Y, W, b)
+    @test_opt test_linsolve!(Y, W, b)
 end
 
 ClimaCore.Operators.allow_mismatched_fd_spaces() = false