diff --git a/.github/workflows/CompileOrRun.yml b/.github/workflows/CompileOrRun.yml
index 0d3285c3..20d23781 100644
--- a/.github/workflows/CompileOrRun.yml
+++ b/.github/workflows/CompileOrRun.yml
@@ -48,6 +48,8 @@ jobs:
   simulation:
     name: Julia ${{ inputs.julia_version }} - ${{ inputs.sim_type }} - ${{ inputs.os }} - ${{ inputs.xla_runtime }} - ${{ inputs.compile_or_run }}
     runs-on: ${{ inputs.os }}
+    env:
+      JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
     strategy:
       fail-fast: false
     steps:
@@ -131,7 +133,7 @@ jobs:
         run: |
           OK="true"
           for OP in all-to-all all-gather all-reduce; do
-              COLLECTIVES=$(find . -name 'optimised_${{ inputs.sim_type }}*.xla' -exec grep --with-filename --line-number --extended-regexp "${OP}" '{}' \;)
+              COLLECTIVES=$(find . -name 'optimised_${{ inputs.sim_type }}*.xla' -exec grep --with-filename --line-number --extended-regexp "${OP}\(" '{}' \;)
               NUM_COLLECTIVES=$(echo "${COLLECTIVES}" | wc -l | xargs)
               if [[ -n "${COLLECTIVES}" ]]; then
                   MSG="There are ${NUM_COLLECTIVES} remaining ${OP} operations"
diff --git a/Project.toml b/Project.toml
index d6e38807..9dad234b 100644
--- a/Project.toml
+++ b/Project.toml
@@ -8,17 +8,19 @@ ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
 BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
-ClimaOcean = "0376089a-ecfe-4b0e-a64f-9c555d74d754"
+ClimaSeaIce = "6ba0ff68-24e6-4315-936c-2e99227c95a4"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 HDF5_jll = "0234f1f7-429e-5d53-9886-15a909be8d59"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
+NumericalEarth = "904d977b-046a-4731-8b86-9235c0d1ef02"
 Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433"
 Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
 SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
@@ -27,18 +29,28 @@ Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 ArgParse = "1.2.0"
 BFloat16s = "0.6.1"
 CUDA = "5.7.1"
-ClimaOcean = "0.5.10"
+ClimaSeaIce = "=0.4.5"
 Dates = "1"
 HDF5_jll = "=1.14.3"
 KernelAbstractions = "0.9.34"
 MPI = "0.20.22"
-Oceananigans = "=0.96.26"
+NumericalEarth = "0.2.2"
+Oceananigans = "0.106.1"
 PrecompileTools = "1.2.1"
 Printf = "1"
 Profile = "1"
+ReactantCore = "0.1.18"
 Reactant = "0.2.247"
 SeawaterPolynomials = "0.3.9"
 Serialization = "1"
 
+[sources.Oceananigans]
+url = "https://github.com/CliMA/Oceananigans.jl.git"
+rev = "pb/step-fix"
+
+[sources.NumericalEarth]
+url = "https://github.com/NumericalEarth/NumericalEarth.jl.git"
+rev = "glw/manual-iterate"
+
 [preferences.Reactant]
 xla_runtime = "IFRT"
diff --git a/correctness/correctness_baroclinic_instability_simulation_run.jl b/correctness/correctness_baroclinic_instability_simulation_run.jl
index 543a8684..83691ad0 100644
--- a/correctness/correctness_baroclinic_instability_simulation_run.jl
+++ b/correctness/correctness_baroclinic_instability_simulation_run.jl
@@ -9,6 +9,7 @@ const parsed_args = GordonBell25.parse_baroclinic_instability_args(;
 
 default_float_type = GordonBell25.float_type_from_args(parsed_args)
 Oceananigans.defaults.FloatType = default_float_type
+using CUDA
 using Reactant
 
 throw_error = true
@@ -40,6 +41,8 @@ vmodel = GordonBell25.baroclinic_instability_model(varch, Nx, Ny, Nz; model_kw..
 ui = 1e-3 .* rand(size(vmodel.velocities.u)...)
 vi = 1e-3 .* rand(size(vmodel.velocities.v)...)
 set!(vmodel, u=ui, v=vi)
+GordonBell25.zero_tendencies!(rmodel)
+GordonBell25.zero_tendencies!(vmodel)
 GordonBell25.sync_states!(rmodel, vmodel)
 
 @info "At the beginning:"
@@ -57,6 +60,9 @@ Oceananigans.TimeSteppers.update_state!(vmodel)
 GordonBell25.compare_states(rmodel, vmodel; include_halos, throw_error, rtol, atol)
 
 GordonBell25.sync_states!(rmodel, vmodel)
+GordonBell25.zero_tendencies!(rmodel)
+GordonBell25.zero_tendencies!(vmodel)
+# compile_options = CompileOptions(; sync=true, raise=true, strip_llvm_debuginfo=true, strip=:all)
 rfirst! = @compile compile_options=compile_options GordonBell25.first_time_step!(rmodel)
 @showtime rfirst!(rmodel)
 @showtime GordonBell25.first_time_step!(vmodel)
diff --git a/correctness/correctness_sharded_baroclinic_instability_simulation_run.jl b/correctness/correctness_sharded_baroclinic_instability_simulation_run.jl
index fe302685..4b27e8fc 100644
--- a/correctness/correctness_sharded_baroclinic_instability_simulation_run.jl
+++ b/correctness/correctness_sharded_baroclinic_instability_simulation_run.jl
@@ -1,5 +1,6 @@
 using GordonBell25
 using Oceananigans
+using Reactant
 
 const parsed_args = GordonBell25.parse_baroclinic_instability_args(;
     grid_x_default = 64,
@@ -9,7 +10,6 @@ const parsed_args = GordonBell25.parse_baroclinic_instability_args(;
 
 default_float_type = GordonBell25.float_type_from_args(parsed_args)
 Oceananigans.defaults.FloatType = default_float_type
-using Reactant
 
 if !GordonBell25.is_distributed_env_present()
     using MPI
@@ -64,6 +64,8 @@ end
 ui = 1e-3 .* rand(size(vmodel.velocities.u)...)
 vi = 1e-3 .* rand(size(vmodel.velocities.v)...)
 set!(vmodel, u=ui, v=vi)
+GordonBell25.zero_tendencies!(rmodel)
+GordonBell25.zero_tendencies!(vmodel)
 GordonBell25.sync_states!(rmodel, vmodel)
 
 compile_options = CompileOptions(; sync=true, raise=true, strip_llvm_debuginfo=true, strip=["enzymexla.kernel_call", "(::Reactant.Compiler.LLVMFunc", "ka_with_reactant", "(::KernelAbstractions.Kernel", "var\"#_launch!;_launch!"], multifloat=GordonBell25.multifloat_from_args(parsed_args))
@@ -81,6 +83,8 @@ Oceananigans.TimeSteppers.update_state!(vmodel)
 GordonBell25.compare_states(rmodel, vmodel; include_halos, throw_error, rtol, atol)
 
 GordonBell25.sync_states!(rmodel, vmodel)
+GordonBell25.zero_tendencies!(rmodel)
+GordonBell25.zero_tendencies!(vmodel)
 rfirst! = @compile compile_options=compile_options GordonBell25.first_time_step!(rmodel)
 @showtime rfirst!(rmodel)
 @showtime GordonBell25.first_time_step!(vmodel)
diff --git a/ext/PrecompileAB2StepF32.jl b/ext/PrecompileAB2StepF32.jl
index 87aedb24..05a9258c 100644
--- a/ext/PrecompileAB2StepF32.jl
+++ b/ext/PrecompileAB2StepF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin
diff --git a/ext/PrecompileCachePreviousTendenciesF32.jl b/ext/PrecompileCachePreviousTendenciesF32.jl
index 784f785a..1f3b2cb2 100644
--- a/ext/PrecompileCachePreviousTendenciesF32.jl
+++ b/ext/PrecompileCachePreviousTendenciesF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin
@@ -39,8 +38,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels:
 
     @compile_workload begin
         model = GordonBell25.baroclinic_instability_model(arch; resolution=4, Δt=60, Nz=4, grid_type=:simple_lat_lon)
-        Δt = ConcreteRNumber(60.0)
-        compiled! = @compile GordonBell25.correct_velocities_and_cache_previous_tendencies_workload!(model, Δt)
+        compiled! = @compile GordonBell25.cache_previous_tendencies_workload!(model)
     end
 
     # Reset float type
diff --git a/ext/PrecompileComputeAuxiliariesF32.jl b/ext/PrecompileComputeAuxiliariesF32.jl
index fcdd7686..b919c407 100644
--- a/ext/PrecompileComputeAuxiliariesF32.jl
+++ b/ext/PrecompileComputeAuxiliariesF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin
diff --git a/ext/PrecompileComputeBoundaryTendenciesF32.jl b/ext/PrecompileComputeBoundaryTendenciesF32.jl
index f32bbb19..5eb64382 100644
--- a/ext/PrecompileComputeBoundaryTendenciesF32.jl
+++ b/ext/PrecompileComputeBoundaryTendenciesF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin
@@ -39,7 +38,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels:
 
     @compile_workload begin
         model = GordonBell25.baroclinic_instability_model(arch; resolution=4, Δt=60, Nz=4, grid_type=:simple_lat_lon)
-        compiled! = @compile GordonBell25.compute_boundary_tendencies_workload!(model)
+        compiled! = @compile GordonBell25.compute_interior_tracer_tendencies_workload!(model)
     end
 
     # Reset float type
diff --git a/ext/PrecompileComputeInteriorMomentumTendenciesF32.jl b/ext/PrecompileComputeInteriorMomentumTendenciesF32.jl
index 16187ea0..a7f82c0a 100644
--- a/ext/PrecompileComputeInteriorMomentumTendenciesF32.jl
+++ b/ext/PrecompileComputeInteriorMomentumTendenciesF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin
diff --git a/ext/PrecompileComputeInteriorTracerTendenciesF32.jl b/ext/PrecompileComputeInteriorTracerTendenciesF32.jl
index 4f984061..db03385b 100644
--- a/ext/PrecompileComputeInteriorTracerTendenciesF32.jl
+++ b/ext/PrecompileComputeInteriorTracerTendenciesF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin
diff --git a/ext/PrecompileComputeTendenciesF32.jl b/ext/PrecompileComputeTendenciesF32.jl
index 6d9415e3..90d4197e 100644
--- a/ext/PrecompileComputeTendenciesF32.jl
+++ b/ext/PrecompileComputeTendenciesF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin
diff --git a/ext/PrecompileFillHaloRegionsF32.jl b/ext/PrecompileFillHaloRegionsF32.jl
index 308ff61c..ec60b363 100644
--- a/ext/PrecompileFillHaloRegionsF32.jl
+++ b/ext/PrecompileFillHaloRegionsF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin
diff --git a/ext/PrecompileTupledFillHaloRegionsF32.jl b/ext/PrecompileTupledFillHaloRegionsF32.jl
index 32d8a88e..8f81adb6 100644
--- a/ext/PrecompileTupledFillHaloRegionsF32.jl
+++ b/ext/PrecompileTupledFillHaloRegionsF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin
@@ -39,7 +38,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels:
 
     @compile_workload begin
         model = GordonBell25.baroclinic_instability_model(arch; resolution=4, Δt=60, Nz=4, grid_type=:simple_lat_lon)
-        compiled! = @compile GordonBell25.tupled_fill_halo_regions_workload!(model)
+        compiled! = @compile GordonBell25.fill_halo_regions_prognostic_workload!(model)
     end
 
     # Reset float type
diff --git a/sharding/less_simple_sharding_problem.jl b/sharding/less_simple_sharding_problem.jl
index 5da73e72..bbfb4963 100644
--- a/sharding/less_simple_sharding_problem.jl
+++ b/sharding/less_simple_sharding_problem.jl
@@ -84,11 +84,11 @@ grid = LatitudeLongitudeGrid(arch, size=(Nx, Ny, Nz), halo=(H, H, H), z=(-4000,
 free_surface = SplitExplicitFreeSurface(substeps=32)
 momentum_advection = WENOVectorInvariant(order=5)
 tracer_advection = WENO(order=5)
-tracers = (:T, :S, :e)
+tracers = (:T, :S)
 equation_of_state = TEOS10EquationOfState()
 buoyancy = SeawaterBuoyancy(; equation_of_state)
 closure = Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivity()
-model = HydrostaticFreeSurfaceModel(; grid, tracers, free_surface,
+model = HydrostaticFreeSurfaceModel(grid; tracers, free_surface,
                                     momentum_advection, tracer_advection,
                                     buoyancy, closure)
 
diff --git a/sharding/sharded_baroclinic_instability_simulation_compile.jl b/sharding/sharded_baroclinic_instability_simulation_compile.jl
index eeab5823..1c9578d6 100644
--- a/sharding/sharded_baroclinic_instability_simulation_compile.jl
+++ b/sharding/sharded_baroclinic_instability_simulation_compile.jl
@@ -5,6 +5,7 @@ using Reactant
 using Oceananigans
 using Oceananigans.Architectures: ReactantState
 Reactant.Compiler.WHILE_CONCAT[] = true
+# Reactant.Compiler.AGGRESSIVE_PROPAGATION[] = true
 
 const parsed_args = GordonBell25.parse_baroclinic_instability_args(;
     grid_x_default = 1536,
@@ -52,7 +53,7 @@ model = GordonBell25.baroclinic_instability_model(arch, Nx, Ny, Nz; halo=(H, H,
 GC.gc(true); GC.gc(false); GC.gc(true)
 
 TRY_COMPILE_FAILED[] = false
-Ninner = ConcreteRNumber(2)
+Ninner = ConcreteRNumber(1)
 
 for optimize in (:before_raise, false, :before_jit, true), code_type in (:hlo, :xla)
     # We only want the optimised XLA code
diff --git a/sharding/sharded_baroclinic_instability_simulation_run.jl b/sharding/sharded_baroclinic_instability_simulation_run.jl
index de25feae..717dc42b 100644
--- a/sharding/sharded_baroclinic_instability_simulation_run.jl
+++ b/sharding/sharded_baroclinic_instability_simulation_run.jl
@@ -20,6 +20,7 @@ using Oceananigans.Units
 using Oceananigans.Architectures: ReactantState
 using Random
 using Printf
+using CUDA
 using Reactant
 
 if !is_distributed_env_present()
@@ -42,16 +43,16 @@ Reactant.MLIR.IR.DUMP_MLIR_DIR[] = joinpath(@__DIR__, "mlir_dumps", jobid_procid
 Reactant.Compiler.DEBUG_DISABLE_RESHARDING[] = true
 # Reactant.Compiler.DEBUG_PRINT_CODEGEN[] = true
 Reactant.Compiler.WHILE_CONCAT[] = true
+# Reactant.Compiler.AGGRESSIVE_PROPAGATION[] = true
 # Reactant.Compiler.DUS_TO_CONCAT[] = false
 # Reactant.Compiler.SUM_TO_REDUCEWINDOW[] = true
 # Reactant.Compiler.AGGRESSIVE_SUM_TO_CONV[] = true
 
 GordonBell25.initialize(; single_gpu_per_process=false)
 
-# devarch = Oceananigans.GPU()
-devarch = Oceananigans.ReactantState()
-
-arch = devarch
+# local_arch = Oceananigans.GPU()
+local_arch = Oceananigans.ReactantState()
+arch = local_arch
 
 Ndev = if arch isa Oceananigans.ReactantState
    length(Reactant.devices())
@@ -66,12 +67,9 @@ Ry, Rx = factors(Ndev)
 if Ndev == 1
     rank = 0
 else
-    arch = Oceananigans.Distributed(
-	arch;
-        partition = Partition(Rx, Ry, 1)
-    )
-    rank = if devarch isa Oceananigans.ReactantState
-	Reactant.Distributed.local_rank()
+    arch = Oceananigans.Distributed(arch; partition = Partition(Rx, Ry, 1))
+    rank = if local_arch isa Oceananigans.ReactantState
+	    Reactant.Distributed.local_rank()
     else
        comm = MPI.COMM_WORLD
        MPI.Comm_rank(comm)
@@ -95,12 +93,13 @@ model = GordonBell25.baroclinic_instability_model(arch, Nx, Ny, Nz; halo=(H, H,
 
 Ninner = 256
 
-if devarch isa Oceananigans.ReactantState
-   Ninner = if Ndev == 1
-	 ConcreteRNumber(Ninner)
-   else
-   	ConcreteRNumber(Ninner; sharding=Sharding.NamedSharding(arch.connectivity, ()))
-   end
+if local_arch isa Oceananigans.ReactantState
+    Ninner = if Ndev == 1
+        ConcreteRNumber(Ninner)
+    else
+        sharding = Sharding.NamedSharding(arch.connectivity, ())
+   	    ConcreteRNumber(Ninner; sharding)
+    end
 end
 
 @info "[$rank] Compiling first_time_step!..." now(UTC)
@@ -114,7 +113,7 @@ end
 @info "[$rank] allocations" GordonBell25.allocatorstats()
 @info "[$rank] Compiling loop..." now(UTC)
 
-compiled_loop! = if devarch isa Oceananigans.ReactantState
+compiled_loop! = if local_arch isa Oceananigans.ReactantState
      @compile compile_options=compile_options loop!(model, Ninner)
 else
      loop!
diff --git a/sharding/simple_sharding_problem.jl b/sharding/simple_sharding_problem.jl
index 658e13d3..cce855a2 100644
--- a/sharding/simple_sharding_problem.jl
+++ b/sharding/simple_sharding_problem.jl
@@ -81,8 +81,8 @@ grid = LatitudeLongitudeGrid(arch, size=(Nx, Ny, Nz), halo=(H, H, H), z=(-4000,
 @info "[$(process_id)] allocations" GordonBell25.allocatorstats()
 
 free_surface = ExplicitFreeSurface()
-model = HydrostaticFreeSurfaceModel(; grid, tracers=:c, free_surface)
-# model = HydrostaticFreeSurfaceModel(; grid)
+model = HydrostaticFreeSurfaceModel(grid; tracers=:c, free_surface)
+# model = HydrostaticFreeSurfaceModel(grid)
 
 @show size(parent(model.velocities.u))
 @assert size(parent(model.velocities.u)) == size(parent(model.tracers.c))
diff --git a/simulations/baroclinic_instability_simulation_compile.jl b/simulations/baroclinic_instability_simulation_compile.jl
index 4b9f54e4..76b0f8fb 100644
--- a/simulations/baroclinic_instability_simulation_compile.jl
+++ b/simulations/baroclinic_instability_simulation_compile.jl
@@ -1,5 +1,6 @@
 using GordonBell25: first_time_step!, loop!, try_compile_code, preamble, TRY_COMPILE_FAILED
 using GordonBell25: baroclinic_instability_model, PROFILE
+using CUDA
 using Reactant
 using Oceananigans
 using Oceananigans.Architectures: ReactantState
diff --git a/simulations/baroclinic_instability_simulation_run.jl b/simulations/baroclinic_instability_simulation_run.jl
index d7e43931..939ebc1d 100644
--- a/simulations/baroclinic_instability_simulation_run.jl
+++ b/simulations/baroclinic_instability_simulation_run.jl
@@ -2,13 +2,14 @@ using GordonBell25: first_time_step!, time_step!, loop!, preamble
 using GordonBell25: baroclinic_instability_model
 using Oceananigans
 using Oceananigans.Architectures: ReactantState
+using CUDA
 using Reactant
 
 # Reactant.Compiler.SROA_ATTRIBUTOR[] = false
 # Reactant.MLIR.IR.DUMP_MLIR_ALWAYS[] = true
 
 preamble()
-
+Reactant.Compiler.DUMP_FAILED_LOCKSTEP[] = true
 Ninner = ConcreteRNumber(3)
 Oceananigans.defaults.FloatType = Float32
 
diff --git a/simulations/ocean_climate_simulation.jl b/simulations/ocean_climate_simulation.jl
index cc2523e3..ed4fd7ca 100644
--- a/simulations/ocean_climate_simulation.jl
+++ b/simulations/ocean_climate_simulation.jl
@@ -1,12 +1,13 @@
 using Oceananigans
 using Oceananigans.Units
 using Oceananigans.Architectures
+using CUDA
 using Reactant
 
-using ClimaOcean
-using ClimaOcean.DataWrangling: ECCO4Monthly
-using ClimaOcean.OceanSeaIceModels.InterfaceComputations: FixedIterations, ComponentInterfaces
-using OrthogonalSphericalShellGrids: TripolarGrid
+using NumericalEarth
+using NumericalEarth: EN4Monthly
+using NumericalEarth.EarthSystemModels.InterfaceComputations: TenUnrolledIterations, ComponentInterfaces
+# using OrthogonalSphericalShellGrids: TripolarGrid
 
 using Dates
 using Printf
@@ -56,37 +57,38 @@ Nz = 20 # eventually we want to increase this to between 100-600
 stop_time = 10days
 
 # Grid setup
-z_faces = exponential_z_faces(; Nz, depth=6000, h=30) # may need changing for very large Nz
-underlying_grid = TripolarGrid(arch; size=(Nx, Ny, Nz), halo=(7, 7, 7), z=z_faces)
+z_faces = ExponentialDiscretization(Nz, -6000, 0; scale=30) # may need changing for very large Nz
+underlying_grid = LatitudeLongitudeGrid(arch; size=(Nx, Ny, Nz), halo=(7, 7, 7), z=z_faces,
+                                        latitude=(-80, 80), longitude=(0, 360))
 bottom_height = regrid_bathymetry(underlying_grid) # adds Earth bathymetry from ETOPO1
 grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height))
 
 # Polar restoring setup
 dates = DateTimeProlepticGregorian(1993, 1, 1) : Month(1) : DateTimeProlepticGregorian(1993, 12, 1)
-temperature = ECCOMetadata(:temperature, dates, ECCO4Monthly())
-salinity = ECCOMetadata(:salinity, dates, ECCO4Monthly())
+temperature = Metadatum(:temperature, dates=dates, dataset=EN4Monthly())
+salinity = Metadatum(:salinity, dates=dates, dataset=EN4Monthly())
 
 restoring_rate  = 1/7days
 mask = LinearlyTaperedPolarMask(southern=(-80, -70), northern=(70, 90))
-FT = ECCORestoring(temperature, grid; mask, rate=restoring_rate)
-FS = ECCORestoring(salinity, grid; mask, rate=restoring_rate)
+FT = DatasetRestoring(temperature, grid; mask, rate=restoring_rate)
+FS = DatasetRestoring(salinity, grid; mask, rate=restoring_rate)
 
-# Ocean simulation with defaults from ClimaOcean
-ocean = ocean_simulation(grid; forcing=(T=FT, S=FT))
+# Ocean simulation with defaults from NumericalEarth
+ocean = ocean_simulation(grid; forcing=(T=FT, S=FS))
 
-# Initial ocean state from ECCO state estimate
-set!(ocean.model, T=ECCOMetadata(:temperature; dates=first(dates)),
-                  S=ECCOMetadata(:salinity; dates=first(dates)))
+# Initial ocean state from EN4
+set!(ocean.model, T=Metadatum(:temperature; dates=first(dates)),
+                  S=Metadatum(:salinity; dates=first(dates)))
 
 # Atmospheric model
 radiation  = Radiation(arch)
 atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(41))
 
 # Coupled model and simulation
-solver_stop_criteria = FixedIterations(5) # note: more iterations = more accurate
-atmosphere_ocean_flux_formulation = SimilarityTheoryFluxes(; solver_stop_criteria)
-interfaces = ComponentInterfaces(atmosphere, ocean; radiation, atmosphere_ocean_flux_formulation)
-coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation, interfaces)
+solver_stop_criteria = TenUnrolledIterations() # note: more iterations = more accurate
+atmosphere_ocean_fluxes = SimilarityTheoryFluxes(; solver_stop_criteria)
+interfaces = ComponentInterfaces(atmosphere, ocean; radiation, atmosphere_ocean_fluxes)
+coupled_model = OceanOnlyModel(ocean; atmosphere, radiation, interfaces)
 simulation = Simulation(coupled_model; Δt, stop_time)
 
 # Utility for printing progress to the terminal
@@ -106,7 +108,7 @@ function progress(sim)
                    prettytime(sim), iteration(sim), prettytime(sim.Δt),
                    umax..., Tmax, Tmin, prettytime(step_time))
 
-    ClimaOcean.@root @info(msg)
+    Oceananigans.DistributedComputations.@root @info(msg)
 
     wall_time[] = time_ns()
 
diff --git a/simulations/ocean_climate_simulation_compile.jl b/simulations/ocean_climate_simulation_compile.jl
index d3976a44..f457a70c 100644
--- a/simulations/ocean_climate_simulation_compile.jl
+++ b/simulations/ocean_climate_simulation_compile.jl
@@ -1,5 +1,6 @@
 using GordonBell25: first_time_step!, loop!, try_compile_code, preamble, TRY_COMPILE_FAILED
 using GordonBell25: data_free_ocean_climate_model_init, PROFILE
+using CUDA
 using Reactant
 using Oceananigans
 using Oceananigans.Architectures: ReactantState
diff --git a/simulations/ocean_climate_simulation_run.jl b/simulations/ocean_climate_simulation_run.jl
index 6e2745fe..a553a792 100644
--- a/simulations/ocean_climate_simulation_run.jl
+++ b/simulations/ocean_climate_simulation_run.jl
@@ -1,6 +1,7 @@
 using GordonBell25: first_time_step!, time_step!, loop!, preamble
 using GordonBell25: data_free_ocean_climate_model_init
 using Oceananigans.Architectures: ReactantState
+using CUDA
 using Reactant
 
 # Reactant.Compiler.SROA_ATTRIBUTOR[] = false
@@ -24,9 +25,9 @@ rloop! = @compile raise=true sync=true loop!(model, Ninner)
 Reactant.with_profiler("./") do
     rfirst!(model)
 end
-Reactant.with_profiler("./") do
-    rstep!(model)
-end
+# Reactant.with_profiler("./") do
+#     rstep!(model)
+# end
 Reactant.with_profiler("./") do
     rloop!(model, Ninner)
 end
diff --git a/src/baroclinic_instability_model.jl b/src/baroclinic_instability_model.jl
index 3515f436..ba74da1a 100644
--- a/src/baroclinic_instability_model.jl
+++ b/src/baroclinic_instability_model.jl
@@ -8,7 +8,6 @@
     return N² * z + Δb * μ + 1e-2 * Δb * randn()
 end
 
-
 function baroclinic_instability_model(arch; resolution, Nz, kw...)
     Nx, Ny = resolution_to_points(resolution)
     return baroclinic_instability_model(arch, Nx, Ny, Nz; kw...)
@@ -33,6 +32,9 @@ function baroclinic_instability_model(arch, Nx, Ny, Nz; Δt,
     # Coriolis forces for a rotating Earth
     coriolis = HydrostaticSphericalCoriolis(),
 
+    # Use the simplest timestepper
+    timestepper = :QuasiAdamsBashforth2,
+
     # Simple momentum advection schemes. May need to be reconsidered
     # due to Float32.
     momentum_advection = WENOVectorInvariant(order=5),
@@ -40,22 +42,13 @@ function baroclinic_instability_model(arch, Nx, Ny, Nz; Δt,
     )
 
     tracers = if buoyancy isa BuoyancyTracer
-        [:b]
+        (:b,)
     elseif buoyancy isa SeawaterBuoyancy
-        [:T, :S]
+        (:T, :S)
     else
-        Symbol[]
+        ()
     end
 
-    if closure isa Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivity
-        push!(tracers, :e)
-    elseif closure isa Oceananigans.TurbulenceClosures.TKEDissipationVerticalDiffusivity
-        push!(tracers, :e)
-        push!(tracers, :ϵ)
-    end
-
-    tracers = tuple(tracers...)
-
     grid = if grid_type === :gaussian_islands
         gaussian_islands_tripolar_grid(arch, Nx, Ny, Nz; halo)
     elseif grid_type === :simple_lat_lon
@@ -64,8 +57,8 @@ function baroclinic_instability_model(arch, Nx, Ny, Nz; Δt,
         error("grid_type=$grid_type must be :gaussian_islands or :simple_lat_lon.")
     end
 
-    model = HydrostaticFreeSurfaceModel(;
-        grid, free_surface, closure, buoyancy, tracers,
+    model = HydrostaticFreeSurfaceModel(grid;
+        free_surface, closure, buoyancy, tracers, timestepper,
         coriolis, momentum_advection, tracer_advection,
     )
 
diff --git a/src/correctness.jl b/src/correctness.jl
index f32a432a..b700f79c 100644
--- a/src/correctness.jl
+++ b/src/correctness.jl
@@ -40,27 +40,27 @@ function compare_states(m1, m2; rtol=sqrt(eps(eltype(m1.grid))), atol=0,
         if !(name ∈ (:w, :η))
             Gⁿ1 = m1.timestepper.Gⁿ
             Gⁿ2 = m2.timestepper.Gⁿ
-            approx_equal *= compare_fields("Gⁿ.$name", Gⁿ1[name], Gⁿ2[name]; rtol, atol)
+            approx_equal *= compare_interior("Gⁿ.$name", Gⁿ1[name], Gⁿ2[name]; rtol, atol)
 
             G⁻1 = m1.timestepper.G⁻
             G⁻2 = m2.timestepper.G⁻
-            approx_equal *= compare_fields("G⁻.$name", G⁻1[name], G⁻2[name]; rtol, atol)
+            approx_equal *= compare_interior("G⁻.$name", G⁻1[name], G⁻2[name]; rtol, atol)
         end
     end
 
     if m1.free_surface isa Oceananigans.SplitExplicitFreeSurface
-        names = (:U, :V, :η)
-        Φ1 = NamedTuple(name => getproperty(m1.free_surface.filtered_state, name) for name in names)
-        Φ2 = NamedTuple(name => getproperty(m2.free_surface.filtered_state, name) for name in names)
-        for name in keys(Φ1)
+        names = keys(m1.free_surface.filtered_state)
+        Φ1 = m1.free_surface.filtered_state
+        Φ2 = m2.free_surface.filtered_state
+        for name in names
             approx_equal *= compare_fields(name, Φ1[name], Φ2[name]; rtol, atol)
         end
     end
 
     if m1.closure isa Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivity
         names = (:κu, :κc, :κe, :Le, :Jᵇ)
-        Φ1 = NamedTuple(name => getproperty(m1.diffusivity_fields, name) for name in names)
-        Φ2 = NamedTuple(name => getproperty(m2.diffusivity_fields, name) for name in names)
+        Φ1 = NamedTuple(name => getproperty(m1.closure_fields, name) for name in names)
+        Φ2 = NamedTuple(name => getproperty(m2.closure_fields, name) for name in names)
         for name in keys(Φ1)
             approx_equal *= compare_fields(name, Φ1[name], Φ2[name]; rtol, atol)
         end
@@ -68,8 +68,8 @@ function compare_states(m1, m2; rtol=sqrt(eps(eltype(m1.grid))), atol=0,
 
     if m1.closure isa Oceananigans.TurbulenceClosures.TKEDissipationVerticalDiffusivity
         names = (:κu, :κc, :κe, :κϵ, :Le, :Lϵ)
-        Φ1 = NamedTuple(name => getproperty(m1.diffusivity_fields, name) for name in names)
-        Φ2 = NamedTuple(name => getproperty(m2.diffusivity_fields, name) for name in names)
+        Φ1 = NamedTuple(name => getproperty(m1.closure_fields, name) for name in names)
+        Φ2 = NamedTuple(name => getproperty(m2.closure_fields, name) for name in names)
         for name in keys(Φ1)
             approx_equal *= compare_fields(name, Φ1[name], Φ2[name]; rtol, atol)
         end
@@ -89,6 +89,14 @@ function compare_states(m1, m2; rtol=sqrt(eps(eltype(m1.grid))), atol=0,
     return nothing
 end
 
+function zero_tendencies!(model)
+    for name in keys(model.timestepper.Gⁿ)
+        parent(model.timestepper.Gⁿ[name]) .= 0
+        parent(model.timestepper.G⁻[name]) .= 0
+    end
+    return nothing
+end
+
 function sync_states!(m1, m2)
     Ψ1 = Oceananigans.fields(m1)
     Ψ2 = Oceananigans.fields(m2)
diff --git a/src/data_free_ocean_climate_model.jl b/src/data_free_ocean_climate_model.jl
index bd27d5c5..b606f459 100644
--- a/src/data_free_ocean_climate_model.jl
+++ b/src/data_free_ocean_climate_model.jl
@@ -18,6 +18,7 @@ function data_free_ocean_climate_model_init(
     )
 
     grid = gaussian_islands_tripolar_grid(arch, resolution, Nz)
+    # grid = simple_latitude_longitude_grid(arch, resolution, Nz)
 
     # See visualize_ocean_climate_simulation.jl for information about how to
     # visualize the results of this run.
@@ -29,11 +30,11 @@ function data_free_ocean_climate_model_init(
     # Set up an atmosphere
     atmos_times = range(0, 1days, length=24)
 
-    atmos_grid = LatitudeLongitudeGrid(arch,
+    topology = (Oceananigans.Grids.Periodic, Oceananigans.Grids.Bounded, Oceananigans.Grids.Flat)
+    atmos_grid = LatitudeLongitudeGrid(arch; topology,
                                        size = (360, 180),
                                        longitude = (0, 360),
-                                       latitude = (-90, 90),
-                                       topology = (Periodic, Bounded, Flat))
+                                       latitude = (-90, 90))
 
     atmosphere = PrescribedAtmosphere(atmos_grid, atmos_times)
 
@@ -61,10 +62,10 @@ function data_free_ocean_climate_model_init(
     radiation = Radiation(arch)
 
     # Coupled model
-    solver_stop_criteria = FixedIterations(5) # note: more iterations = more accurate
-    atmosphere_ocean_flux_formulation = SimilarityTheoryFluxes(; solver_stop_criteria)
-    interfaces = ComponentInterfaces(atmosphere, ocean; radiation, atmosphere_ocean_flux_formulation)
-    coupled_model = @gbprofile "OceanSeaIceModel" OceanSeaIceModel(ocean; atmosphere, radiation, interfaces)
+    solver_stop_criteria = TenUnrolledIterations() # note: more iterations = more accurate
+    atmosphere_ocean_fluxes = SimilarityTheoryFluxes(; solver_stop_criteria)
+    interfaces = ComponentInterfaces(atmosphere, ocean; radiation, atmosphere_ocean_fluxes)
+    coupled_model = @gbprofile "OceanOnlyModel" OceanOnlyModel(ocean; atmosphere, radiation, interfaces)
 
     return coupled_model
 end # data_free_ocean_climate_model_init
diff --git a/src/model_utils.jl b/src/model_utils.jl
index 24423c48..01cc92d6 100644
--- a/src/model_utils.jl
+++ b/src/model_utils.jl
@@ -5,8 +5,8 @@ using Oceananigans.Units
 using Oceananigans.Architectures: Architectures
 using SeawaterPolynomials
 
-using ClimaOcean
-using ClimaOcean.OceanSeaIceModels.InterfaceComputations: FixedIterations, ComponentInterfaces
+using NumericalEarth
+using NumericalEarth.EarthSystemModels.InterfaceComputations: TenUnrolledIterations, ComponentInterfaces
 
 using Dates
 using Printf
@@ -51,7 +51,7 @@ function simple_latitude_longitude_grid(arch, resolution, Nz)
 end
 
 function simple_latitude_longitude_grid(arch, Nx, Ny, Nz; halo=(8, 8, 8))
-    z = exponential_z_faces(; Nz, depth=4000, h=30) # may need changing for very large Nz
+    z = ExponentialDiscretization(Nz, -4000, 0; scale=1000)
 
     grid = LatitudeLongitudeGrid(arch; size=(Nx, Ny, Nz), halo, z,
         latitude = (-80, 80),
@@ -130,7 +130,7 @@ end
 
 function gaussian_islands_tripolar_grid(arch::Architectures.AbstractArchitecture, Nx, Ny, Nz; halo=(8, 8, 8))
     # Grid setup
-    z = exponential_z_faces(; Nz, depth=4000, h=30) # may need changing for very large Nz
+    z = ExponentialDiscretization(Nz, -4000, 0; scale=1000)
     underlying_grid = TripolarGrid(arch; size=(Nx, Ny, Nz), halo, z)
 
     zb = z[1]
diff --git a/src/precompile.jl b/src/precompile.jl
index 3202a2b5..d87afddd 100644
--- a/src/precompile.jl
+++ b/src/precompile.jl
@@ -7,57 +7,44 @@ using Oceananigans.Architectures:
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.TimeSteppers:
     ab2_step!,
-    correct_velocities_and_cache_previous_tendencies!
+    cache_previous_tendencies!
 
-using Oceananigans.ImmersedBoundaries:
+using Oceananigans.Grids:
     get_active_cells_map
 
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
+    compute_closure_fields!,
     compute_hydrostatic_momentum_tendencies!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!,
     interior_tendency_kernel_parameters,
-    compute_hydrostatic_boundary_tendency_contributions!,
     compute_hydrostatic_free_surface_Gc!,
-    immersed_boundary_condition,
-    complete_communication_and_compute_buffer!,
-    compute_tendencies!
+    immersed_boundary_condition
 
 #=
-# For reference
-
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+# For reference (update_state! + ab2_step! flow)
+
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
-function tupled_fill_halo_regions_workload!(model)
-    tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
+function fill_halo_regions_prognostic_workload!(model)
+    fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
 end
 
 function compute_tendencies_workload!(model)
-    compute_tendencies!(model, [])
-end
-
-function compute_boundary_tendencies_workload!(model)
-    compute_hydrostatic_boundary_tendency_contributions!(model.timestepper.Gⁿ,
-                                                         model.architecture,
-                                                         model.velocities,
-                                                         model.tracers,
-                                                         model.clock,
-                                                         fields(model),
-                                                         model.closure,
-                                                         model.buoyancy)
+    compute_momentum_tendencies!(model, [])
+    compute_tracer_tendencies!(model)
 end
 
 function compute_interior_momentum_tendencies_workload!(model)
@@ -69,7 +56,6 @@ function compute_interior_momentum_tendencies_workload!(model)
     #compute_hydrostatic_free_surface_tendency_contributions!(model, kernel_parameters; active_cells_map)
     
     compute_hydrostatic_momentum_tendencies!(model, model.velocities, kernel_parameters; active_cells_map)
-    complete_communication_and_compute_buffer!(model, grid, arch)
 end
 
 function compute_interior_tracer_tendencies_workload!(model)
@@ -92,10 +78,10 @@ function compute_interior_tracer_tendencies_workload!(model)
                      c_immersed_bc,
                      model.buoyancy,
                      model.biogeochemistry,
-                     model.velocities,
+                     model.transport_velocities,
                      model.free_surface,
                      model.tracers,
-                     model.diffusivity_fields,
+                     model.closure_fields,
                      model.auxiliary_fields,
                      model.clock,
                      c_forcing)
@@ -104,25 +90,24 @@ function compute_interior_tracer_tendencies_workload!(model)
                 compute_hydrostatic_free_surface_Gc!,
                 c_tendency,
                 grid,
-                active_cells_map,
                 args;
                 active_cells_map)
     end
 end
 
 function compute_auxiliaries_workload!(model)
-    compute_auxiliaries!(model)
+    compute_closure_fields!(model.closure_fields, model.closure, model)
 end
 
 function fill_halo_regions_workload!(model)
-    fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
+    fill_halo_regions!(model.closure_fields; only_local_halos=true)
 end
 
 function ab2_step_workload!(model, Δt)
     ab2_step!(model, Δt)
 end
 
-function correct_velocities_and_cache_previous_tendencies_workload!(model, Δt)
-    correct_velocities_and_cache_previous_tendencies!(model, Δt)
+function cache_previous_tendencies_workload!(model)
+    cache_previous_tendencies!(model)
 end
 
diff --git a/src/timestepping_utils.jl b/src/timestepping_utils.jl
index 4241f640..3b4ddf98 100644
--- a/src/timestepping_utils.jl
+++ b/src/timestepping_utils.jl
@@ -1,6 +1,8 @@
 using Reactant
 using Oceananigans
-import Oceananigans.TimeSteppers: first_time_step!, time_step!
+import Oceananigans.TimeSteppers: time_step!, update_state!
+using Oceananigans.Models: AbstractModel
+using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper
 using Reactant_jll: libReactantExtra
 
 const TRY_COMPILE_FAILED = Ref(false)
@@ -18,14 +20,27 @@ function try_compile_code(f)
     end
 end
 
+function first_time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper})
+    Reactant.Profiler.annotate("first_time_step") do
+        Oceananigans.initialize!(model)
+        Oceananigans.TimeSteppers.update_state!(model)
+        Δt = model.clock.last_Δt + 0
+        Oceananigans.TimeSteppers.time_step!(model, Δt; euler=true)
+    end
+    return nothing
+end
+
 function first_time_step!(model)
     Reactant.Profiler.annotate("first_time_step") do
-        Δt = model.clock.last_Δt
-        Oceananigans.TimeSteppers.first_time_step!(model, Δt)
+        Oceananigans.initialize!(model)
+        Oceananigans.TimeSteppers.update_state!(model)
+        Δt = model.clock.last_Δt + 0
+        Oceananigans.TimeSteppers.time_step!(model, Δt)
     end
     return nothing
 end
 
+
 function time_step!(model)
     Reactant.Profiler.annotate("time_step") do
         Δt = model.clock.last_Δt + 0