diff --git a/.github/workflows/Integration.yml b/.github/workflows/Integration.yml index 496c6578a2..57c1c20c6d 100644 --- a/.github/workflows/Integration.yml +++ b/.github/workflows/Integration.yml @@ -53,6 +53,7 @@ jobs: - DynamicExpressions - SymbolicRegression - Lux + - Oceananigans - SciML - KernelAbstractions - MatrixAlgebraKit diff --git a/test/integration/Oceananigans/Project.toml b/test/integration/Oceananigans/Project.toml new file mode 100644 index 0000000000..adf453c5ca --- /dev/null +++ b/test/integration/Oceananigans/Project.toml @@ -0,0 +1,12 @@ +[deps] +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" +Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[compat] +GPUArraysCore = "=0.2.0" +Oceananigans = "=0.101.0" + +[sources] +Enzyme = { path = "../../.." } diff --git a/test/integration/Oceananigans/runtests.jl b/test/integration/Oceananigans/runtests.jl new file mode 100644 index 0000000000..45f388a40d --- /dev/null +++ b/test/integration/Oceananigans/runtests.jl @@ -0,0 +1,50 @@ +using Enzyme +using GPUArraysCore: @allowscalar +using Oceananigans +using Test + +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 + time_step!(model, dt) + end + + compute!(K) + + # TODO return @allowscalar first(K) + return first(K) +end + +@testset "Column model with ScalarDiffusivity" begin + ν = 1e-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()) + + N² = 1e-6 + bᵢ(z) = N² * z + uᵢ(z) = exp(-z^2 / 10) + set!(m, b=bᵢ, u=uᵢ) + + u, v, w = m.velocities + ke = (u^2 + v^2) / 2 + K = Field(Integral(ke)) + + 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)) + + @test dKdν[1][1] != 0 +end