Skip to content

[WIP] Oceananigans integration tests#2775

Open
glwagner wants to merge 3 commits intoEnzymeAD:mainfrom
glwagner:oceananigans-integration
Open

[WIP] Oceananigans integration tests#2775
glwagner wants to merge 3 commits intoEnzymeAD:mainfrom
glwagner:oceananigans-integration

Conversation

@glwagner
Copy link
Collaborator

@glwagner glwagner commented Nov 12, 2025

This PR starts adding some integration tests. I noticed that the first thing I tried failed, so we'll have to work our way up.

Working on this with @giordano and @jtbuch

@github-actions
Copy link
Contributor

github-actions bot commented Nov 12, 2025

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.
diff --git a/test/integration/Oceananigans/runtests.jl b/test/integration/Oceananigans/runtests.jl
index dab12a36..41a64e6e 100644
--- a/test/integration/Oceananigans/runtests.jl
+++ b/test/integration/Oceananigans/runtests.jl
@@ -3,14 +3,14 @@ using GPUArraysCore: @allowscalar
 using Oceananigans
 using Test
 
-function kinetic_energy(model, ν, K, dt=1)
+function kinetic_energy(model, ν, K, dt = 1)
     vitd = VerticallyImplicitTimeDiscretization()
     new_closure = ScalarDiffusivity(vitd; ν)
     tracer_names = keys(model.tracers)
     new_closure = Oceananigans.TurbulenceClosures.with_tracers(tracer_names, new_closure)
     model.closure = new_closure
 
-    for n = 1:10
+    for n in 1:10
         time_step!(model, dt)
     end
 
@@ -20,17 +20,19 @@ function kinetic_energy(model, ν, K, dt=1)
 end
 
 @testset "Column model with ScalarDiffusivity" begin
-    ν = 1e-3
-    grid = RectilinearGrid(size=128, z=(-64, 64), topology=(Flat, Flat, Bounded))
+    ν = 1.0e-3
+    grid = RectilinearGrid(size = 128, z = (-64, 64), topology = (Flat, Flat, Bounded))
     vitd = VerticallyImplicitTimeDiscretization()
     closure = ScalarDiffusivity(vitd; ν)
-    m = HydrostaticFreeSurfaceModel(; grid, closure, coriolis=FPlane(f=1e-4),
-        tracers=:b, buoyancy=BuoyancyTracer())
+    m = HydrostaticFreeSurfaceModel(;
+        grid, closure, coriolis = FPlane(f = 1.0e-4),
+        tracers = :b, buoyancy = BuoyancyTracer()
+    )
 
-    N² = 1e-6
+    N² = 1.0e-6
     bᵢ(z) = N² * z
     uᵢ(z) = exp(-z^2 / 10)
-    set!(m, b=bᵢ, u=uᵢ)
+    set!(m, b = bᵢ, u = uᵢ)
 
     u, v, w = m.velocities
     ke = (u^2 + v^2) / 2
@@ -39,11 +41,13 @@ end
     dm = Enzyme.make_zero(m)
     dK = Enzyme.make_zero(K)
 
-    dKdν = autodiff(set_strong_zero(Enzyme.ReverseWithPrimal),
-                    kinetic_energy, Active,
-                    Duplicated(m, dm),
-                    Active(ν),
-                    Duplicated(K, dK))
+    dKdν = autodiff(
+        set_strong_zero(Enzyme.ReverseWithPrimal),
+        kinetic_energy, Active,
+        Duplicated(m, dm),
+        Active(ν),
+        Duplicated(K, dK)
+    )
 
     @test dKdν[1][1] != 0
 end

@glwagner
Copy link
Collaborator Author

Looks like I need to be a little less ambitious probably


compute!(K)

return @allowscalar first(K)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@giordano @glwagner so this is failing atm because allowscalar isn't currently supported in enzyme.jl natively [we should add cc @vchuravy ]

@wsmoses wsmoses force-pushed the oceananigans-integration branch from 9b8a278 to 727981b Compare January 19, 2026 07:26
@jtbuch
Copy link

jtbuch commented Jan 19, 2026

This might be relevant for you @dkytezab

@wsmoses
Copy link
Member

wsmoses commented Jan 20, 2026

1.10 just needs to be marked runtime activity.

1.11 hits a weird issue in an error handler, @vchuravy was looking at a while back but I don't think had cycles to look into for completion

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants