Add test for autodifferentiating hydrostatic turbulence simulation

Project.toml

@@ -50,7 +50,7 @@ CubedSphere = "0.2, 0.3"
 Dates = "1.9"
 Distances = "0.10"
 DocStringExtensions = "0.8, 0.9"
-Enzyme = "0.13.3"
+Enzyme = "0.13.13"
 FFTW = "1"
 Glob = "1.3"
 IncompleteLU = "0.2"

src/Fields/field.jl

@@ -404,8 +404,8 @@ Base.checkbounds(f::Field, I...) = Base.checkbounds(f.data, I...)
 @propagate_inbounds Base.lastindex(f::Field) = lastindex(f.data)
 @propagate_inbounds Base.lastindex(f::Field, dim) = lastindex(f.data, dim)
 
-Base.fill!(f::Field, val) = fill!(parent(f), val)
-Base.parent(f::Field) = parent(f.data)
+@inline Base.fill!(f::Field, val) = fill!(parent(f), val)
+@inline Base.parent(f::Field) = parent(f.data)
 Adapt.adapt_structure(to, f::Field) = Adapt.adapt(to, f.data)
 
 total_size(f::Field) = total_size(f.grid, location(f), f.indices)

src/Fields/field_tuples.jl

@@ -56,42 +56,73 @@ Fill halo regions for all `fields`. The algorithm:
 function fill_halo_regions!(maybe_nested_tuple::Union{NamedTuple, Tuple}, args...; kwargs...)
     flattened = flattened_unique_values(maybe_nested_tuple)
 
-    # Sort fields into ReducedField and Field with non-nothing boundary conditions
-    fields_with_bcs = filter(f -> !isnothing(boundary_conditions(f)), flattened)
-    reduced_fields  = filter(f -> f isa ReducedField, fields_with_bcs)
-
-    for field in reduced_fields
-        fill_halo_regions!(field, args...; kwargs...)
+    # Look for grid within the flattened field tuple:
+    for f in flattened
+        if isdefined(f, :grid)
+            grid = f.grid
+            return tupled_fill_halo_regions!(flattened, grid, args...; kwargs...)
+        end
     end
 
-    # MultiRegion fields are considered windowed_fields (indices isa MultiRegionObject))
-    windowed_fields = filter(f -> !(f isa FullField), fields_with_bcs)
-    ordinary_fields = filter(f -> (f isa FullField) && !(f isa ReducedField), fields_with_bcs)
+    return tupled_fill_halo_regions!(flattened, args...; kwargs...)
+end
 
-    # Fill halo regions for reduced and windowed fields
-    for field in windowed_fields
-        fill_halo_regions!(field, args...; kwargs...)
+# Version where we find grid amongst ordinary fields:
+function tupled_fill_halo_regions!(fields, args...; kwargs...)
+
+    ordinary_fields = Field[]
+    for f in fields
+        if !isnothing(boundary_conditions(f))
+            if f isa ReducedField || !(f isa FullField)
+                # Windowed and reduced fields
+                fill_halo_regions!(f, args...; kwargs...)
+            else
+                push!(ordinary_fields, f)
+            end
+        end
     end
 
-    # Fill the rest
-    if !isempty(ordinary_fields)
+    ordinary_fields = tuple(ordinary_fields...)
+
+    if !isempty(ordinary_fields) # ie not reduced, and with default_indices
         grid = first(ordinary_fields).grid
-        tupled_fill_halo_regions!(ordinary_fields, grid, args...; kwargs...)
+        fill_halo_regions!(map(data, ordinary_fields),
+                           map(boundary_conditions, ordinary_fields),
+                           default_indices(3),
+                           map(instantiated_location, ordinary_fields),
+                           grid, args...; kwargs...)
     end
 
     return nothing
 end
-
+
+# Version where grid is provided:
 function tupled_fill_halo_regions!(fields, grid, args...; kwargs...)
 
-    # We cannot group windowed fields together, the indices must be (:, :, :)!
-    indices = default_indices(3)        
+    ordinary_fields = Field[]
+    for f in fields
+        if !isnothing(boundary_conditions(f))
+            if f isa ReducedField || !(f isa FullField)
+                # Windowed and reduced fields
+                fill_halo_regions!(f, args...; kwargs...)
+            else
+                push!(ordinary_fields, f)
+            end
+        end
+    end
 
-    return fill_halo_regions!(map(data, fields),
-                              map(boundary_conditions, fields),
-                              indices,
-                              map(instantiated_location, fields),
-                              grid, args...; kwargs...)
+    ordinary_fields = tuple(ordinary_fields...)
+
+    if !isempty(ordinary_fields) # ie not reduced, and with default_indices
+        grid = first(ordinary_fields).grid
+        fill_halo_regions!(map(data, ordinary_fields),
+                           map(boundary_conditions, ordinary_fields),
+                           default_indices(3),
+                           map(instantiated_location, ordinary_fields),
+                           grid, args...; kwargs...)
+    end
+
+    return nothing
 end
 
 #####

src/Fields/set!.jl

@@ -34,6 +34,11 @@ set!(u::Field, f::Function) = set_to_function!(u, f)
 set!(u::Field, a::Union{Array, CuArray, OffsetArray}) = set_to_array!(u, a)
 set!(u::Field, v::Field) = set_to_field!(u, v)
 
+function set!(u::Field, a::Number)
+    fill!(interior(u), a) # note all other set! only change interior
+    return u # return u, not parent(u), for type-stability
+end
+
 function set!(u::Field, v)
     u .= v # fallback
     return u

src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl

@@ -32,7 +32,6 @@ on_architecture(to, free_surface::ExplicitFreeSurface) =
 function materialize_free_surface(free_surface::ExplicitFreeSurface{Nothing}, velocities, grid)
     η = free_surface_displacement_field(velocities, free_surface, grid)
     g = convert(eltype(grid), free_surface.gravitational_acceleration)
-
     return ExplicitFreeSurface(η, g)
 end
 
@@ -62,10 +61,16 @@ explicit_ab2_step_free_surface!(free_surface, model, Δt, χ) =
 ##### Kernel
 #####
 
-@kernel function _explicit_ab2_step_free_surface!(η, Δt, χ::FT, Gηⁿ, Gη⁻, Nz) where FT
+@kernel function _explicit_ab2_step_free_surface!(η, Δt, χ, Gηⁿ, Gη⁻, Nz)
     i, j = @index(Global, NTuple)
+    FT = typeof(χ)
+    one_point_five = convert(FT, 1.5)
+    oh_point_five = convert(FT, 0.5)
+    not_euler = χ != convert(FT, -0.5)
 
     @inbounds begin
-        η[i, j, Nz+1] += Δt * ((FT(1.5) + χ) * Gηⁿ[i, j, Nz+1] - (FT(0.5) + χ) * Gη⁻[i, j, Nz+1])
+        Gη = (one_point_five + χ) * Gηⁿ[i, j, Nz+1] - (oh_point_five  + χ) * Gη⁻[i, j, Nz+1] * not_euler
+        η[i, j, Nz+1] += Δt * Gη
     end
 end
+

src/Models/HydrostaticFreeSurfaceModels/set_hydrostatic_free_surface_model.jl

@@ -33,7 +33,7 @@ T₀[T₀ .< 0.5] .= 0
 
 set!(model, u=u₀, v=v₀, T=T₀)
 
-model.velocities.u
+@model.velocities.u
 
 # output
 
@@ -61,7 +61,8 @@ model.velocities.u
     end
 
     initialize!(model)
-    update_state!(model; compute_tendencies = false)
+    update_state!(model; compute_tendencies=false)
 
     return nothing
 end
+

src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl

@@ -168,7 +168,6 @@ end
 end
 
 function initialize_free_surface_state!(state, η, timestepper)
-
     parent(state.U) .= parent(state.U̅)
     parent(state.V) .= parent(state.V̅)
 
@@ -213,11 +212,9 @@ function barotropic_split_explicit_corrector!(u, v, free_surface, grid)
     Hᶠᶜ, Hᶜᶠ   = free_surface.auxiliary.Hᶠᶜ, free_surface.auxiliary.Hᶜᶠ
     arch       = architecture(grid)
 
-
-    # take out "bad" barotropic mode,
+    # Remove incorrect barotropic mode and add correction
     # !!!! reusing U and V for this storage since last timestep doesn't matter
     compute_barotropic_mode!(U, V, grid, u, v)
-    # add in "good" barotropic mode
     launch!(arch, grid, :xyz, _barotropic_split_explicit_corrector!,
             u, v, U̅, V̅, U, V, Hᶠᶜ, Hᶜᶠ, grid)
 
@@ -228,14 +225,14 @@ end
 Explicitly step forward η in substeps.
 """
 ab2_step_free_surface!(free_surface::SplitExplicitFreeSurface, model, Δt, χ) =
-    split_explicit_free_surface_step!(free_surface, model, Δt, χ)
+    split_explicit_free_surface_step!(free_surface, model, Δt)
 
 function initialize_free_surface!(sefs::SplitExplicitFreeSurface, grid, velocities)
     @apply_regionally compute_barotropic_mode!(sefs.state.U̅, sefs.state.V̅, grid, velocities.u, velocities.v)
     fill_halo_regions!((sefs.state.U̅, sefs.state.V̅, sefs.η))
 end
 
-function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurface, model, Δt, χ)
+function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurface, model, Δt)
 
     # Note: free_surface.η.grid != model.grid for DistributedSplitExplicitFreeSurface
     # since halo_size(free_surface.η.grid) != halo_size(model.grid)
@@ -247,16 +244,16 @@ function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurfac
     # Calculate the substepping parameterers
     settings = free_surface.settings
     Nsubsteps = calculate_substeps(settings.substepping, Δt)
-
-    # barotropic time step as fraction of baroclinic step and averaging weights
-    fractional_Δt, weights = calculate_adaptive_settings(settings.substepping, Nsubsteps)
+    
+    # Barotropic time step as fraction of baroclinic step and averaging weights
+    fractional_Δt, weights = calculate_adaptive_settings(settings.substepping, Nsubsteps) 
     Nsubsteps = length(weights)
 
-    # barotropic time step in seconds
-    Δτᴮ = fractional_Δt * Δt
-
-    # reset free surface averages
-    @apply_regionally begin
+    # Barotropic time step in time units
+    Δτᴮ = fractional_Δt * Δt  
+    
+    # Reset free surface averages
+    @apply_regionally begin 
         initialize_free_surface_state!(free_surface.state, free_surface.η, settings.timestepper)
 
         # Solve for the free surface at tⁿ⁺¹
@@ -269,8 +266,8 @@ function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurfac
     fields_to_fill = (free_surface.state.U̅, free_surface.state.V̅)
     fill_halo_regions!(fields_to_fill; async = true)
 
-    # Preparing velocities for the barotropic correction
-    @apply_regionally begin
+    # Prepare velocities for the barotropic correction
+    @apply_regionally begin 
         mask_immersed_field!(model.velocities.u)
         mask_immersed_field!(model.velocities.v)
     end
@@ -350,40 +347,55 @@ function iterate_split_explicit!(free_surface, grid, Δτᴮ, weights, ::Val{Nsu
 end
 
 # Calculate RHS for the barotropic time step.
-@kernel function _compute_integrated_ab2_tendencies!(Gᵁ, Gⱽ, grid, ::Nothing, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
+@kernel function _compute_integrated_tendencies!(Gᵁ, Gⱽ, grid, ::Nothing, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
     i, j  = @index(Global, NTuple)
     k_top = grid.Nz + 1
 
-    @inbounds Gᵁ[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * ab2_step_Gu(i, j, 1, grid, Gu⁻, Guⁿ, χ)
-    @inbounds Gⱽ[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * ab2_step_Gv(i, j, 1, grid, Gv⁻, Gvⁿ, χ)
+    @inbounds begin
+        Gᵁ[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * u_tendency_increment(i, j, 1, grid, Gu⁻, Guⁿ, χ)
+        Gⱽ[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * v_tendency_increment(i, j, 1, grid, Gv⁻, Gvⁿ, χ)
 
-    for k in 2:grid.Nz
-        @inbounds Gᵁ[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * ab2_step_Gu(i, j, k, grid, Gu⁻, Guⁿ, χ)
-        @inbounds Gⱽ[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * ab2_step_Gv(i, j, k, grid, Gv⁻, Gvⁿ, χ)
+        for k in 2:grid.Nz	
+            Gᵁ[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * u_tendency_increment(i, j, k, grid, Gu⁻, Guⁿ, χ)
+            Gⱽ[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * v_tendency_increment(i, j, k, grid, Gv⁻, Gvⁿ, χ)
+        end	
     end
 end
 
 # Calculate RHS for the barotropic time step.q
-@kernel function _compute_integrated_ab2_tendencies!(Gᵁ, Gⱽ, grid, active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
+@kernel function _compute_integrated_tendencies!(Gᵁ, Gⱽ, grid, active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
     idx = @index(Global, Linear)
     i, j = active_linear_index_to_tuple(idx, active_cells_map)
     k_top = grid.Nz+1
 
-    @inbounds Gᵁ[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * ab2_step_Gu(i, j, 1, grid, Gu⁻, Guⁿ, χ)
-    @inbounds Gⱽ[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * ab2_step_Gv(i, j, 1, grid, Gv⁻, Gvⁿ, χ)
+    @inbounds begin
+        Gᵁ[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * u_tendency_increment(i, j, 1, grid, Gu⁻, Guⁿ, χ)
+        Gⱽ[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * v_tendency_increment(i, j, 1, grid, Gv⁻, Gvⁿ, χ)
+
+        for k in 2:grid.Nz	
+            Gᵁ[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * u_tendency_increment(i, j, k, grid, Gu⁻, Guⁿ, χ)
+            Gⱽ[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * v_tendency_increment(i, j, k, grid, Gv⁻, Gvⁿ, χ)
+        end	
 
-    for k in 2:grid.Nz
-        @inbounds Gᵁ[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * ab2_step_Gu(i, j, k, grid, Gu⁻, Guⁿ, χ)
-        @inbounds Gⱽ[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * ab2_step_Gv(i, j, k, grid, Gv⁻, Gvⁿ, χ)
     end
 end
 
-@inline ab2_step_Gu(i, j, k, grid, G⁻, Gⁿ, χ::FT) where FT =
-    @inbounds ifelse(peripheral_node(i, j, k, grid, f, c, c), zero(grid), (convert(FT, 1.5) + χ) *  Gⁿ[i, j, k] - G⁻[i, j, k] * (convert(FT, 0.5) + χ))
-
-@inline ab2_step_Gv(i, j, k, grid, G⁻, Gⁿ, χ::FT) where FT =
-    @inbounds ifelse(peripheral_node(i, j, k, grid, c, f, c), zero(grid), (convert(FT, 1.5) + χ) *  Gⁿ[i, j, k] - G⁻[i, j, k] * (convert(FT, 0.5) + χ))
+@inline function u_tendency_increment(i, j, k, grid, G⁻, Gⁿ, χ)
+    FT = typeof(χ)
+    peripheral = peripheral_node(i, j, k, grid, f, c, c)
+    not_euler = χ != convert(FT, -0.5)
+    ΔGu = @inbounds (convert(FT, 1.5) + χ) * Gⁿ[i, j, k] - (convert(FT, 0.5) + χ) * G⁻[i, j, k] * not_euler
+    return ifelse(peripheral, zero(grid), ΔGu)
+end
 
+@inline function v_tendency_increment(i, j, k, grid, G⁻, Gⁿ, χ)
+    FT = typeof(χ)
+    peripheral = peripheral_node(i, j, k, grid, c, f, c)
+    not_euler = χ != convert(FT, -0.5)
+    ΔGv = @inbounds (convert(FT, 1.5) + χ) *  Gⁿ[i, j, k] - (convert(FT, 0.5) + χ) * G⁻[i, j, k] * not_euler
+    return ifelse(peripheral, zero(grid), ΔGv)
+end
+
 # Setting up the RHS for the barotropic step (tendencies of the barotropic velocity components)
 # This function is called after `calculate_tendency` and before `ab2_step_velocities!`
 function setup_free_surface!(model, free_surface::SplitExplicitFreeSurface, χ)
@@ -399,17 +411,15 @@ function setup_free_surface!(model, free_surface::SplitExplicitFreeSurface, χ)
     @apply_regionally setup_split_explicit_tendency!(auxiliary, model.grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
 
     fields_to_fill = (auxiliary.Gᵁ, auxiliary.Gⱽ)
-    fill_halo_regions!(fields_to_fill; async = true)
+    fill_halo_regions!(fields_to_fill; async=true)
 
     return nothing
 end
 
 @inline function setup_split_explicit_tendency!(auxiliary, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
     active_cells_map = retrieve_surface_active_cells_map(grid)
-
-    launch!(architecture(grid), grid, :xy, _compute_integrated_ab2_tendencies!, auxiliary.Gᵁ, auxiliary.Gⱽ, grid,
-            active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ; active_cells_map)
-
+    kernel_args = (auxiliary.Gᵁ, auxiliary.Gⱽ, grid, active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
+    launch!(architecture(grid), grid, :xy, _compute_integrated_tendencies!, kernel_args...; active_cells_map)
     return nothing
 end
 

src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl

@@ -8,7 +8,7 @@ using Oceananigans.TurbulenceClosures: compute_diffusivities!
 using Oceananigans.ImmersedBoundaries: mask_immersed_field!, mask_immersed_field_xy!, inactive_node
 using Oceananigans.Models: update_model_field_time_series!
 using Oceananigans.Models.NonhydrostaticModels: update_hydrostatic_pressure!, p_kernel_parameters
-using Oceananigans.Fields: replace_horizontal_vector_halos!
+using Oceananigans.Fields: replace_horizontal_vector_halos!, tupled_fill_halo_regions!
 
 import Oceananigans.Models.NonhydrostaticModels: compute_auxiliaries!
 import Oceananigans.TimeSteppers: update_state!
@@ -38,7 +38,8 @@ function update_state!(model::HydrostaticFreeSurfaceModel, grid, callbacks; comp
     # Update the boundary conditions
     @apply_regionally update_boundary_condition!(fields(model), model)
 
-    fill_halo_regions!(prognostic_fields(model), model.clock, fields(model); async = true)
+    tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model), async=true)
+
     @apply_regionally replace_horizontal_vector_halos!(model.velocities, model.grid)
     @apply_regionally compute_auxiliaries!(model)