Add adapt_structure for ForcingKernelFunction (#4579)

* add adapt_structure for ForcingKernelFunction

* add test

* Update Project.toml

* remove stray line

* fix tests

cc @glwagner

---------

Co-authored-by: Navid C. Constantinou <[email protected]>

Author: glwagner

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.96.30"
+version = "0.96.31"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/Forcings/continuous_forcing.jl

@@ -168,4 +168,3 @@ on_architecture(to, forcing::ContinuousForcing{LX, LY, LZ}) where {LX, LY, LZ} =
                                   on_architecture(to, forcing.field_dependencies),
                                   on_architecture(to, forcing.field_dependencies_indices),
                                   on_architecture(to, forcing.field_dependencies_interp))
-

src/Models/forcing_operation.jl

@@ -1,10 +1,14 @@
+using Adapt
 
 import Oceananigans.Utils: prettysummary
 
 struct ForcingKernelFunction{F}
     forcing :: F
 end
 
+Adapt.adapt_structure(to, fkf::ForcingKernelFunction) =
+    ForcingKernelFunction(adapt(to, fkf.forcing))
+
 prettysummary(kf::ForcingKernelFunction) = "ForcingKernelFunction"
 
 @inline function (kf::ForcingKernelFunction)(i, j, k, grid, args...)

test/test_model_diagnostics.jl

@@ -1,6 +1,8 @@
 include("dependencies_for_runtests.jl")
 
-using Oceananigans.Models: BoundaryConditionOperation, BoundaryConditionField
+using Oceananigans.Forcings: ContinuousForcing
+using Oceananigans.Models: BoundaryConditionOperation, BoundaryConditionField,
+                           ForcingOperation, ForcingField, ForcingKernelFunction
 
 u_bottom_drag_continuous(x, y, t, grid, u, v, Cᴰ) = -Cᴰ * u * sqrt(u^2 + v^2)
 v_bottom_drag_discrete(i, j, grid, clock, fields, Cᴰ) = - @inbounds Cᴰ * fields.v[i, j, 1]
@@ -85,10 +87,10 @@ c_bottom_flux(i, j, grid, clock, fields, t★) = - @inbounds fields.c[i, j, 1] /
 
                 @test location(c_bottom_bc) == (Center, Center, Nothing)
                 @test location(c_top_bc)    == (Center, Center, Nothing)
-            
+
                 initial_c(x, y, z) = 1 + 1 * (z > 0)
                 set!(model, c=initial_c)
-            
+
                 c_bottom_bc_field = BoundaryConditionField(c, :bottom, model)
                 c_top_bc_field = BoundaryConditionField(c, :top, model)
 
@@ -97,11 +99,38 @@ c_bottom_flux(i, j, grid, clock, fields, t★) = - @inbounds fields.c[i, j, 1] /
 
                 @test c_bottom_bc_field isa BoundaryConditionField
                 @test c_top_bc_field isa BoundaryConditionField
-                
+
                 @test all(interior(c_bottom_bc_field) .≈ - 1 / t★)
                 @test all(interior(c_top_bc_field) .≈ + 2 / t★)
             end
         end
     end
 end
 
+damping(x, y, z, t, c, τ) = - c / τ
+
+@testset "ForcingOperation and ForcingField" begin
+    for arch in archs
+        grid = RectilinearGrid(arch,
+                               size=(4, 4, 4),
+                               x = (0, 1),
+                               y = (0, 1),
+                               z = (-1, 1),
+                               topology=(Bounded, Bounded, Bounded))
+
+
+        c_forcing = Forcing(damping, field_dependencies=:c, parameters=60)
+        model = NonhydrostaticModel(; grid, tracers=:c, forcing=(; c=c_forcing))
+        c_forcing_op = ForcingOperation(:c, model)
+        @test c_forcing_op isa KernelFunctionOperation
+        @test c_forcing_op.kernel_function isa ForcingKernelFunction
+        @test c_forcing_op.kernel_function.forcing isa ContinuousForcing
+        @test c_forcing_op.kernel_function.forcing.parameters == 60
+
+        set!(model, c=1)
+        c_forcing_field = ForcingField(:c, model)
+        compute!(c_forcing_field)
+        @test c_forcing_field isa Field
+        @test all(interior(c_forcing_field) .== - 1/60)
+    end
+end