From e6758b335da8c081d66bcbd77b62172916823dff Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 8 Oct 2024 12:51:29 -0700 Subject: [PATCH 01/25] Delete a lot of comments --- ext/OceananigansEnzymeExt.jl | 308 ----------------------------------- 1 file changed, 308 deletions(-) diff --git a/ext/OceananigansEnzymeExt.jl b/ext/OceananigansEnzymeExt.jl index e1b7d74a1f..4c2fe35f63 100644 --- a/ext/OceananigansEnzymeExt.jl +++ b/ext/OceananigansEnzymeExt.jl @@ -7,7 +7,6 @@ using Oceananigans.Utils: contiguousrange using KernelAbstractions -#isdefined(Base, :get_extension) ? (import EnzymeCore) : (import ..EnzymeCore) isdefined(Base, :get_extension) ? (import Enzyme) : (import ..Enzyme) using Enzyme: EnzymeCore @@ -24,313 +23,6 @@ EnzymeCore.EnzymeRules.inactive(::typeof(Oceananigans.BoundaryConditions.parent_ @inline batch(::Val{1}, ::Type{T}) where T = T @inline batch(::Val{N}, ::Type{T}) where {T, N} = NTuple{N, T} -# function EnzymeCore.EnzymeRules.augmented_primal(config, -# func::EnzymeCore.Const{Type{Field}}, -# ::Type{<:EnzymeCore.Annotation{RT}}, -# loc::Union{EnzymeCore.Const{<:Tuple}, -# EnzymeCore.Duplicated{<:Tuple}}, -# grid::EnzymeCore.Annotation{<:Oceananigans.Grids.AbstractGrid}, -# T::EnzymeCore.Const{<:DataType}; kw...) where RT -# -# primal = if EnzymeCore.EnzymeRules.needs_primal(config) -# func.val(loc.val, grid.val, T.val; kw...) -# else -# nothing -# end -# -# if haskey(kw, :a) -# # copy zeroing -# kw[:data] = copy(kw[:data]) -# end -# -# shadow = if EnzymeCore.EnzymeRules.width(config) == 1 -# func.val(loc.val, grid.val, T.val; kw...) -# else -# ntuple(Val(EnzymeCore.EnzymeRules.width(config))) do i -# Base.@_inline_meta -# func.val(loc.val, grid.val, T.val; kw...) -# end -# end -# -# P = EnzymeCore.EnzymeRules.needs_primal(config) ? RT : Nothing -# B = batch(Val(EnzymeCore.EnzymeRules.width(config)), RT) -# return EnzymeCore.EnzymeRules.AugmentedReturn{P, B, Nothing}(primal, shadow, nothing) -# end -# -# ##### -# ##### Field -# ##### -# -# function EnzymeCore.EnzymeRules.reverse(config::EnzymeCore.EnzymeRules.ConfigWidth{1}, -# func::EnzymeCore.Const{Type{Field}}, -# ::RT, -# tape, -# loc::Union{EnzymeCore.Const{<:Tuple}, EnzymeCore.Duplicated{<:Tuple}}, -# grid::EnzymeCore.Const{<:Oceananigans.Grids.AbstractGrid}, -# T::EnzymeCore.Const{<:DataType}; kw...) where RT -# return (nothing, nothing, nothing) -# end -# -# function EnzymeCore.EnzymeRules.reverse(config::EnzymeCore.EnzymeRules.ConfigWidth{1}, -# func::EnzymeCore.Const{Type{Field}}, -# ::RT, -# tape, -# loc::Union{EnzymeCore.Const{<:Tuple}, EnzymeCore.Duplicated{<:Tuple}}, -# grid::EnzymeCore.Active{<:Oceananigans.Grids.AbstractGrid}, -# T::EnzymeCore.Const{<:DataType}; kw...) where RT -# return (nothing, EnzymeCore.make_ero(grid), nothing) -# end - -##### -##### FunctionField -##### - -# @inline FunctionField(L::Tuple, func, grid) = FunctionField{L[1], L[2], L[3]}(func, grid) - -# function EnzymeCore.EnzymeRules.augmented_primal(config, -# enzyme_func::Union{EnzymeCore.Const{<:Type{<:FunctionField}}, EnzymeCore.Const{Type{FT2}}}, -# ::Type{<:EnzymeCore.Annotation{RT}}, -# function_field_func, -# grid; -# clock = nothing, -# parameters = nothing) where {RT, FT2 <: FunctionField} -# -# FunctionFieldType = enzyme_func.val -# -# primal = if EnzymeCore.EnzymeRules.needs_primal(config) -# FunctionFieldType(function_field_func.val, grid.val; clock, parameters) -# else -# nothing -# end -# -# # function_field_func can be Active, Const (inactive), Duplicated (active but mutable) -# function_field_is_active = function_field_func isa Active -# # @show function_field_func -# -# # Support batched differentiation! -# config_width = EnzymeCore.EnzymeRules.width(config) -# -# dactives = if function_field_is_active -# if config_width == 1 -# Ref(EnzymeCore.make_zero(function_field_func.val)) -# else -# ntuple(Val(config_width)) do i -# Base.@_inline_meta -# Ref(EnzymeCore.make_zero(function_field_func.val)) -# end -# end -# else -# nothing -# end -# -# shadow = if config_width == 1 -# dfunction_field_func = if function_field_is_active -# dactives[] -# else -# function_field_func.dval -# end -# -# FunctionFieldType(dfunction_field_func, grid.val; clock, parameters) -# else -# ntuple(Val(config_width)) do i -# Base.@_inline_meta -# -# dfunction_field_func = if function_field_is_active -# dactives[i][] -# else -# function_field_func.dval[i] -# end -# -# FunctionFieldType(dfunction_field_func, grid.val; clock, parameters) -# end -# end -# -# P = EnzymeCore.EnzymeRules.needs_primal(config) ? RT : Nothing -# B = batch(Val(EnzymeCore.EnzymeRules.width(config)), RT) -# D = typeof(dactives) -# -# return EnzymeCore.EnzymeRules.AugmentedReturn{P, B, D}(primal, shadow, dactives) -# end -# -# function EnzymeCore.EnzymeRules.reverse(config, -# enzyme_func::Union{EnzymeCore.Const{<:Type{<:FunctionField}}, EnzymeCore.Const{Type{FT2}}}, -# ::RT, -# tape, -# function_field_func, -# grid; -# clock = nothing, -# parameters = nothing) where {RT, FT2 <: FunctionField} -# -# dactives = if function_field_func isa Active -# if EnzymeCore.EnzymeRules.width(config) == 1 -# tape[] -# else -# ntuple(Val(EnzymeCore.EnzymeRules.width(config))) do i -# Base.@_inline_meta -# tape[i][] -# end -# end -# else -# nothing -# end -# -# # return (dactives, grid (nothing)) -# return (dactives, nothing) -# end - -##### -##### launch! -##### - -# function EnzymeCore.EnzymeRules.augmented_primal(config, -# func::EnzymeCore.Const{typeof(Oceananigans.Models.flattened_unique_values)}, -# ::Type{<:EnzymeCore.Annotation{RT}}, -# a) where RT -# -# sprimal = if EnzymeCore.EnzymeRules.needs_primal(config) || EnzymeCore.EnzymeRules.needs_shadow(config) -# func.val(a.val) -# else -# nothing -# end -# -# shadow = if EnzymeCore.EnzymeRules.needs_shadow(config) -# if EnzymeCore.EnzymeRules.width(config) == 1 -# (typeof(a) <: Const) ? EnzymeCore.make_zero(sprimal)::RT : func.val(a.dval) -# else -# ntuple(Val(EnzymeCore.EnzymeRules.width(config))) do i -# Base.@_inline_meta -# (typeof(a) <: Const) ? EnzymeCore.make_zero(sprimal)::RT : func.val(a.dval[i]) -# end -# end -# else -# nothing -# end -# -# primal = if EnzymeCore.EnzymeRules.needs_primal(config) -# sprimal -# else -# nothing -# end -# -# P = EnzymeCore.EnzymeRules.needs_primal(config) ? RT : Nothing -# B = EnzymeCore.EnzymeRules.needs_primal(config) ? batch(Val(EnzymeCore.EnzymeRules.width(config)), RT) : Nothing -# -# return EnzymeCore.EnzymeRules.AugmentedReturn{P, B, Nothing}(primal, shadow, nothing) -# end -# -# function EnzymeCore.EnzymeRules.reverse(config, -# func::EnzymeCore.Const{typeof(Oceananigans.Models.flattened_unique_values)}, -# ::Type{<:EnzymeCore.Annotation{RT}}, -# tape, -# a) where RT -# -# return (nothing,) -# end - -##### -##### launch! -##### - -# function EnzymeCore.EnzymeRules.augmented_primal(config, -# func::EnzymeCore.Const{typeof(Oceananigans.Utils.launch!)}, -# ::Type{EnzymeCore.Const{Nothing}}, -# arch, -# grid, -# workspec, -# kernel!, -# kernel_args::Vararg{Any,N}; -# include_right_boundaries = false, -# reduced_dimensions = (), -# location = nothing, -# active_cells_map = nothing, -# kwargs...) where N -# -# -# workgroup, worksize = Oceananigans.Utils.work_layout(grid.val, workspec.val; -# include_right_boundaries, -# reduced_dimensions, -# location) -# -# offset = Oceananigans.Utils.offsets(workspec.val) -# -# if !isnothing(active_cells_map) -# workgroup, worksize = Oceananigans.Utils.active_cells_work_layout(workgroup, worksize, active_cells_map, grid.val) -# offset = nothing -# end -# -# if worksize != 0 -# -# # We can only launch offset kernels with Static sizes!!!! -# -# if isnothing(offset) -# loop! = kernel!.val(Oceananigans.Architectures.device(arch.val), workgroup, worksize) -# dloop! = (typeof(kernel!) <: EnzymeCore.Const) ? nothing : kernel!.dval(Oceananigans.Architectures.device(arch.val), workgroup, worksize) -# else -# loop! = kernel!.val(Oceananigans.Architectures.device(arch.val), KernelAbstractions.StaticSize(workgroup), Oceananigans.Utils.OffsetStaticSize(contiguousrange(worksize, offset))) -# dloop! = (typeof(kernel!) <: EnzymeCore.Const) ? nothing : kernel!.val(Oceananigans.Architectures.device(arch.val), KernelAbstractions.StaticSize(workgroup), Oceananigans.Utils.OffsetStaticSize(contiguousrange(worksize, offset))) -# end -# -# @debug "Launching kernel $kernel! with worksize $worksize and offsets $offset from $workspec.val" -# -# -# duploop = (typeof(kernel!) <: EnzymeCore.Const) ? EnzymeCore.Const(loop!) : EnzymeCore.Duplicated(loop!, dloop!) -# -# config2 = EnzymeCore.EnzymeRules.Config{#=needsprimal=#false, #=needsshadow=#false, #=width=#EnzymeCore.EnzymeRules.width(config), EnzymeCore.EnzymeRules.overwritten(config)[5:end]}() -# subtape = EnzymeCore.EnzymeRules.augmented_primal(config2, duploop, EnzymeCore.Const{Nothing}, kernel_args...).tape -# -# tape = (duploop, subtape) -# else -# tape = nothing -# end -# -# return EnzymeCore.EnzymeRules.AugmentedReturn{Nothing, Nothing, Any}(nothing, nothing, tape) -# end -# -# @inline arg_elem_type(::Type{T}, ::Val{i}) where {T<:Tuple, i} = eltype(T.parameters[i]) -# -# function EnzymeCore.EnzymeRules.reverse(config::EnzymeCore.EnzymeRules.ConfigWidth{1}, -# func::EnzymeCore.Const{typeof(Oceananigans.Utils.launch!)}, -# ::Type{EnzymeCore.Const{Nothing}}, -# tape, -# arch, -# grid, -# workspec, -# kernel!, -# kernel_args::Vararg{Any,N}; -# include_right_boundaries = false, -# reduced_dimensions = (), -# location = nothing, -# active_cells_map = nothing, -# kwargs...) where N -# -# subrets = if tape !== nothing -# duploop, subtape = tape -# config2 = EnzymeCore.EnzymeRules.Config{#=needsprimal=#false, #=needsshadow=#false, #=width=#EnzymeCore.EnzymeRules.width(config), EnzymeCore.EnzymeRules.overwritten(config)[5:end]}() -# EnzymeCore.EnzymeRules.reverse(config2, duploop, EnzymeCore.Const{Nothing}, subtape, kernel_args...) -# else -# res2 = ntuple(Val(N)) do i -# Base.@_inline_meta -# if kernel_args[i] isa Active -# EnzymeCore.make_zero(kernel_args[i].val) -# else -# nothing -# end -# end -# end -# -# subrets2 = ntuple(Val(N)) do i -# Base.@_inline_meta -# if kernel_args[i] isa Active -# subrets[i]::arg_elem_type(typeof(kernel_args), Val(i)) -# else -# nothing -# end -# end -# -# return (nothing, nothing, nothing, nothing, subrets2...) -# -# end - ##### ##### update_model_field_time_series! ##### From c6dcefe64ba356a5efe7a5a65f216491c3e3f546 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 8 Oct 2024 13:42:27 -0700 Subject: [PATCH 02/25] Add parameter estimation test --- test/test_enzyme.jl | 73 +++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 68 insertions(+), 5 deletions(-) diff --git a/test/test_enzyme.jl b/test/test_enzyme.jl index 00ff0b90ba..ec4bb61716 100644 --- a/test/test_enzyme.jl +++ b/test/test_enzyme.jl @@ -1,5 +1,7 @@ include("dependencies_for_runtests.jl") +using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity + # Required presently Enzyme.API.looseTypeAnalysis!(true) Enzyme.API.maxtypeoffset!(2032) @@ -7,6 +9,7 @@ Enzyme.API.maxtypeoffset!(2032) # OceananigansLogger doesn't work here -- not sure why Logging.global_logger(TestLogger()) +#= f(grid) = CenterField(grid) const maximum_diffusivity = 100 @@ -107,14 +110,10 @@ end @testset "Enzyme + Oceananigans Initialization Broadcast Kernel" begin - Nx = Ny = 64 - Nz = 8 - x = y = (-π, π) z = (-0.5, 0.5) topology = (Periodic, Periodic, Bounded) - - grid = RectilinearGrid(size=(Nx, Ny, Nz); x, y, z, topology) + grid = RectilinearGrid(size=(64, 64, 8); x, y, z, topology) model = HydrostaticFreeSurfaceModel(; grid, tracers=:c) model_tracer = model.tracers.c @@ -235,4 +234,68 @@ end @test rel_error < tol end +=# + +function buoyancy_variance!(model, e_min, Δt=10.0) + new_closure = CATKEVerticalDiffusivity(minimum_tke=e_min) + model.closure = new_closure + model.clock.time = 0 + model.clock.iteration = 0 + + bᵢ(z) = 1e-5 * z + set!(model, b=bᵢ) + + for n = 1:10 + time_step!(model, Δt) + end + + b = model.tracers.b + + # Another way to compute it + Nx, Ny, Nz = size(model.grid) + sum_b² = 0.0 + for k = 1:Nz, j = 1:Ny, i = 1:Nx + sum_b² += b[i, j, k]^2 + end + + return sum_b² +end + +@testset "Enzyme with CATKEVerticalDiffusivity" begin + for arch in archs + grid = RectilinearGrid(arch, size=50, z=(-200, 0), topology=(Flat, Flat, Bounded)) + closure = CATKEVerticalDiffusivity() + b_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(1e-7)) + tracers = (:b, :e) + buoyancy = BuoyancyTracer() + model = HydrostaticFreeSurfaceModel(; grid, closure, tracers, buoyancy, + boundary_conditions=(; b_bcs)) + + # Compute derivative by hand + e₁, e₂ = 1e-5, 2e-5 + b²₁ = buoyancy_variance!(model, e₁) + b²₂ = buoyancy_variance!(model, e₂) + db²_de_fd = (b²₂ - b²₁) / (e₂ - e₁) + + # Now for real + e = 1e-5 + dmodel = Enzyme.make_zero(model) + buoyancy_variance!(dmodel, 0) + + db²_de = autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse), + buoyancy_variance!, + Duplicated(model, dmodel), + Active(e)) + + @info """ \n + Enzyme computed $db²_de + Finite differences computed $db²_de_fd + """ + + tol = 0.01 + rel_error = abs(db²_de[1][3] - db²_de_fd) / abs(db²_de_fd) + @show db²_de, db²_de_fd + @test rel_error < tol + end +end From f9b989605f9ef6aaa750b603a93a2c8b832232f8 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 8 Oct 2024 14:22:48 -0700 Subject: [PATCH 03/25] Introduce custom DiffusivityFields to control fill_halo_regions behavior --- .../TKEBasedVerticalDiffusivities.jl | 3 +- .../catke_vertical_diffusivity.jl | 39 ++++++++++++++++--- .../tke_dissipation_vertical_diffusivity.jl | 38 +++++++++++++++++- 3 files changed, 72 insertions(+), 8 deletions(-) diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/TKEBasedVerticalDiffusivities.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/TKEBasedVerticalDiffusivities.jl index 77d0d465fb..9553c5c4dd 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/TKEBasedVerticalDiffusivities.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/TKEBasedVerticalDiffusivities.jl @@ -1,5 +1,6 @@ module TKEBasedVerticalDiffusivities +using Adapt using Adapt using CUDA using KernelAbstractions: @kernel, @index @@ -27,7 +28,7 @@ using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization, VerticalFormulation -import Oceananigans.BoundaryConditions: getbc +import Oceananigans.BoundaryConditions: getbc, fill_halo_regions! import Oceananigans.Utils: with_tracers import Oceananigans.TurbulenceClosures: validate_closure, diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl index 2adf895c82..6eca29a4f5 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl @@ -1,4 +1,3 @@ - struct CATKEVerticalDiffusivity{TD, CL, FT, DT, TKE} <: AbstractScalarDiffusivity{TD, VerticalFormulation, 2} mixing_length :: CL turbulent_kinetic_energy_equation :: TKE @@ -145,6 +144,37 @@ catke_first(catke1::FlavorOfCATKE, catke2::FlavorOfCATKE) = error("Can't have tw ##### Diffusivities and diffusivity fields utilities ##### +struct CATKEDiffusivityFields{K, L, J, T, U, KC, LC} + κu :: K + κc :: K + κe :: K + Le :: L + Jᵇ :: J + previous_compute_time :: T + previous_velocities :: U + _tupled_tracer_diffusivities :: KC + _tupled_implicit_linear_coefficients :: LC +end + +Adapt.adapt_structure(to, catke_diffusivity_fields::CATKEDiffusivityFields) = + CATKEDiffusivityFields(adapt(to, catke_diffusivity_fields.κu), + adapt(to, catke_diffusivity_fields.κc), + adapt(to, catke_diffusivity_fields.κe), + adapt(to, catke_diffusivity_fields.Le), + adapt(to, catke_diffusivity_fields.Jᵇ), + previous_compute_time[], + adapt(to, previous_velocities), + adapt(to, _tupled_tracer_diffusivities), + adapt(to, _tupled_implicit_linear_coefficients)) + +function fill_halo_regions!(catke_diffusivity_fields::CATKEDiffusivityFields, args...; kw...) + fields_with_halos_to_fill = (catke_diffusivity_fields.κu, + catke_diffusivity_fields.κc, + catke_diffusivity_fields.κe) + + return fill_halo_regions!(fields_with_halos_to_fill, args...; kw...) +end + function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfCATKE) default_diffusivity_bcs = (κu = FieldBoundaryConditions(grid, (Center, Center, Face)), @@ -170,16 +200,15 @@ function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfCATKE) _tupled_tracer_diffusivities = NamedTuple(name => name === :e ? κe : κc for name in tracer_names) _tupled_implicit_linear_coefficients = NamedTuple(name => name === :e ? Le : ZeroField() for name in tracer_names) - return (; κu, κc, κe, Le, Jᵇ, - previous_compute_time, previous_velocities, - _tupled_tracer_diffusivities, _tupled_implicit_linear_coefficients) + return CATKEDiffusivityFields(κu, κc, κe, Le, Jᵇ, + previous_compute_time, previous_velocities, + _tupled_tracer_diffusivities, _tupled_implicit_linear_coefficients) end @inline viscosity_location(::FlavorOfCATKE) = (c, c, f) @inline diffusivity_location(::FlavorOfCATKE) = (c, c, f) function compute_diffusivities!(diffusivities, closure::FlavorOfCATKE, model; parameters = :xyz) - arch = model.architecture grid = model.grid velocities = model.velocities diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/tke_dissipation_vertical_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/tke_dissipation_vertical_diffusivity.jl index d94035e507..cf8eb52ba5 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/tke_dissipation_vertical_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/tke_dissipation_vertical_diffusivity.jl @@ -149,6 +149,38 @@ end ##### Diffusivities and diffusivity fields utilities ##### +struct TKEDissipationDiffusivityFields{K, L, U, KC, LC} + κu :: K + κc :: K + κe :: K + κϵ :: K + Le :: L + Lϵ :: L + previous_velocities :: U + _tupled_tracer_diffusivities :: KC + _tupled_implicit_linear_coefficients :: LC +end + +Adapt.adapt_structure(to, tke_dissipation_diffusivity_fields::TKEDissipationDiffusivityFields) = + TKEDissipationDiffusivityFields(adapt(to, tke_dissipation_diffusivity_fields.κu), + adapt(to, tke_dissipation_diffusivity_fields.κc), + adapt(to, tke_dissipation_diffusivity_fields.κe), + adapt(to, tke_dissipation_diffusivity_fields.κϵ), + adapt(to, tke_dissipation_diffusivity_fields.Le), + adapt(to, tke_dissipation_diffusivity_fields.Lϵ), + adapt(to, previous_velocities), + adapt(to, _tupled_tracer_diffusivities), + adapt(to, _tupled_implicit_linear_coefficients)) + +function fill_halo_regions!(catke_diffusivity_fields::TKEDissipationDiffusivityFields, args...; kw...) + fields_with_halos_to_fill = (tke_dissipation_diffusivity_fields.κu, + tke_dissipation_diffusivity_fields.κc, + tke_dissipation_diffusivity_fields.κe, + tke_dissipation_diffusivity_fields.κϵ) + + return fill_halo_regions!(fields_with_halos_to_fill, args...; kw...) +end + function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfTD) default_diffusivity_bcs = (κu = FieldBoundaryConditions(grid, (Center, Center, Face)), @@ -184,8 +216,10 @@ function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfTD) _tupled_implicit_linear_coefficients = NamedTuple(name => _tupled_implicit_linear_coefficients[name] for name in tracer_names) - return (; κu, κc, κe, κϵ, Le, Lϵ, previous_velocities, - _tupled_tracer_diffusivities, _tupled_implicit_linear_coefficients) + return TKEDissipationDiffusivityFields(κu, κc, κe, κϵ, Le, Lϵ, + previous_velocities, + _tupled_tracer_diffusivities, + _tupled_implicit_linear_coefficients) end @inline viscosity_location(::FlavorOfTD) = (c, c, f) From af61d4f7ee1883b4716179e682199bffcf108b7c Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 8 Oct 2024 14:23:12 -0700 Subject: [PATCH 04/25] Some cosmetic changes --- .../HydrostaticFreeSurfaceModels/single_column_model_mode.jl | 1 - test/test_enzyme.jl | 4 ++-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl index d4d3dec8d8..4ad0225287 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl @@ -70,7 +70,6 @@ function update_state!(model::HydrostaticFreeSurfaceModel, grid::SingleColumnGri # Calculate diffusivities compute_diffusivities!(model.diffusivity_fields, model.closure, model) - fill_halo_regions!(model.diffusivity_fields, model.clock, fields(model)) for callback in callbacks diff --git a/test/test_enzyme.jl b/test/test_enzyme.jl index ec4bb61716..f1b071fc8f 100644 --- a/test/test_enzyme.jl +++ b/test/test_enzyme.jl @@ -236,13 +236,13 @@ end end =# +bᵢ(z) = 1e-5 * z + function buoyancy_variance!(model, e_min, Δt=10.0) new_closure = CATKEVerticalDiffusivity(minimum_tke=e_min) model.closure = new_closure model.clock.time = 0 model.clock.iteration = 0 - - bᵢ(z) = 1e-5 * z set!(model, b=bᵢ) for n = 1:10 From 38fbd921392aba74e2cc8ae0308c3f6c07b61bcf Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 9 Oct 2024 08:45:01 -0700 Subject: [PATCH 05/25] Even more specific fill halo regions for CATKE --- .../catke_vertical_diffusivity.jl | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl index 6eca29a4f5..29789badd3 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl @@ -168,11 +168,10 @@ Adapt.adapt_structure(to, catke_diffusivity_fields::CATKEDiffusivityFields) = adapt(to, _tupled_implicit_linear_coefficients)) function fill_halo_regions!(catke_diffusivity_fields::CATKEDiffusivityFields, args...; kw...) - fields_with_halos_to_fill = (catke_diffusivity_fields.κu, - catke_diffusivity_fields.κc, - catke_diffusivity_fields.κe) - - return fill_halo_regions!(fields_with_halos_to_fill, args...; kw...) + fill_halo_regions!(catke_diffusivity_fields.κu, args...; kw...) + fill_halo_regions!(catke_diffusivity_fields.κc, args...; kw...) + fill_halo_regions!(catke_diffusivity_fields.κe, args...; kw...) + return nothing end function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfCATKE) From 20e41659165ccf3567a4224df8e3e21cf0c4b4c4 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 24 Oct 2024 15:50:09 -0600 Subject: [PATCH 06/25] Add test for hydrostatic turbulence --- test/test_enzyme.jl | 94 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 92 insertions(+), 2 deletions(-) diff --git a/test/test_enzyme.jl b/test/test_enzyme.jl index 45faab48ff..896fefe222 100644 --- a/test/test_enzyme.jl +++ b/test/test_enzyme.jl @@ -106,7 +106,6 @@ function set_initial_condition_via_launch!(model_tracer, amplitude) end @testset "Enzyme + Oceananigans Initialization Broadcast Kernel" begin - Nx = Ny = 64 Nz = 8 @@ -157,7 +156,7 @@ end end end -@testset "Enzyme on advection and diffusion" begin +@testset "Enzyme autodifferentiation of advection and diffusion" begin Nx = Ny = 64 Nz = 8 @@ -234,3 +233,94 @@ end @test rel_error < tol end +function viscous_hydrostatic_turbulence(ν, model, u_init, v_init, Δt, u_truth, v_truth) + # Initialize the model + model.clock.iteration = 0 + model.clock.time = 0 + new_closure = ScalarDiffusivity(; ν, κ=NamedTuple()) + model.closure = new_closure + set!(model, u=u_init, v=v_init) + + # Step it forward + for n = 1:10 + time_step!(model, Δt) + end + + # Compute the sum square error + u, v, w = model.velocities + Nx, Ny, Nz = size(model.grid) + err = 0.0 + for j = 1:Ny, i = 1:Nx + err += @inbounds (u[i, j, 3] - u_truth[i, j, 3])^2 + + (v[i, j, 3] - v_truth[i, j, 3])^2 + end + + return err::Float64 +end + +@testset "Enzyme autodifferentiation of hydrostatic turbulence" begin + Random.seed!(123) + arch = CPU() + Nx = Ny = 32 + Nz = 1 + x = y = (0, 2π) + z = 1 + ν₀ = 1e-2 + + grid = RectilinearGrid(arch, size=(Nx, Ny, 1); x, y, z, topology=(Periodic, Periodic, Bounded)) + closure = ScalarDiffusivity(ν=ν₀) + momentum_advection = Centered(order=2) + + g = 4^2 + c = sqrt(g) + free_surface = ExplicitFreeSurface(gravitational_acceleration=g) + model = HydrostaticFreeSurfaceModel(; grid, momentum_advection, free_surface, closure) + + ϵ(x, y, z) = 2randn() - 1 + set!(model, u=ϵ, v=ϵ) + + u_init = deepcopy(model.velocities.u) + v_init = deepcopy(model.velocities.v) + + Δx = minimum_xspacing(grid) + Δt = 0.01 * Δx / c + for n = 1:10 + time_step!(model, Δt) + end + + u_truth = deepcopy(model.velocities.u) + v_truth = deepcopy(model.velocities.v) + + # Use a manual finite difference to compute a gradient + Δν = 1e-6 + ν1 = 2e-2 + ν2 = ν1 + Δν + e1 = viscous_hydrostatic_turbulence(ν1, model, u_init, v_init, Δt, u_truth, v_truth) + e2 = viscous_hydrostatic_turbulence(ν2, model, u_init, v_init, Δt, u_truth, v_truth) + ΔeΔν = (e2 - e1) / Δν + + @info "Finite difference computed: $ΔeΔν" + + @info "Now with autodiff..." + start_time = time_ns() + + # Use autodiff to compute a gradient + dmodel = Enzyme.make_zero(model) + dedν = autodiff(set_runtime_activity(Enzyme.Reverse), + viscous_hydrostatic_turbulence, + Active(ν1), + Duplicated(model, dmodel), + Const(u_init), + Const(v_init), + Const(Δt), + Const(u_truth), + Const(v_truth)) + + @info "Automatically computed: $dedν." + @info "Elapsed time: " * prettytime(1e-9 * (time_ns() - start_time)) + + tol = 1e-1 + rel_error = abs(dedν[1][3] - ΔeΔν) / abs(ΔeΔν) + @test rel_error < tol +end + From 92b314cd6c70f2b55e907a8831940b3237dc8a71 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 24 Oct 2024 17:02:54 -0600 Subject: [PATCH 07/25] Bypass fill_halo_regions for tuples --- .../update_hydrostatic_free_surface_model_state.jl | 7 +++++-- test/test_enzyme.jl | 11 +++++++++-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl index cf6b20eade..599bb56453 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl +++ b/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,10 @@ 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) + fill_halo_regions!(model.free_surface.η, model.clock, fields(model), async=true) + prognostic_3d_fields = (model.velocities.u, model.velocities.v, model.tracers...) + tupled_fill_halo_regions!(prognostic_3d_fields, model.grid, model.clock, fields(model); async=true) + @apply_regionally replace_horizontal_vector_halos!(model.velocities, model.grid) @apply_regionally compute_auxiliaries!(model) diff --git a/test/test_enzyme.jl b/test/test_enzyme.jl index 896fefe222..51e0442730 100644 --- a/test/test_enzyme.jl +++ b/test/test_enzyme.jl @@ -233,12 +233,19 @@ end @test rel_error < tol end +function set_viscosity!(model, viscosity) + new_closure = ScalarDiffusivity(ν=viscosity) + names = () + new_closure = with_tracers(names, new_closure) + model.closure = new_closure + return nothing +end + function viscous_hydrostatic_turbulence(ν, model, u_init, v_init, Δt, u_truth, v_truth) # Initialize the model model.clock.iteration = 0 model.clock.time = 0 - new_closure = ScalarDiffusivity(; ν, κ=NamedTuple()) - model.closure = new_closure + set_viscosity!(model, ν) set!(model, u=u_init, v=v_init) # Step it forward From b25c2b4ee7d7758620e29fab7eea4de7479a0927 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 24 Oct 2024 17:28:55 -0600 Subject: [PATCH 08/25] Cleanup to interface --- src/Fields/field_tuples.jl | 50 ++++++++----------- ...te_hydrostatic_free_surface_model_state.jl | 4 +- 2 files changed, 21 insertions(+), 33 deletions(-) diff --git a/src/Fields/field_tuples.jl b/src/Fields/field_tuples.jl index debac6aab9..9d8e02913e 100644 --- a/src/Fields/field_tuples.jl +++ b/src/Fields/field_tuples.jl @@ -55,45 +55,35 @@ 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) + return tupled_fill_halo_regions!(flattened, args...; kwargs...) +end - for field in reduced_fields - fill_halo_regions!(field, args...; kwargs...) - 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) - - # Fill halo regions for reduced and windowed fields - for field in windowed_fields - fill_halo_regions!(field, args...; kwargs...) +function tupled_fill_halo_regions!(fields, args...; kwargs...) + + ordinary_fields = Field[] + for fields in fields + if !isnothing(boundary_conditions(f)) + if field isa ReducedField || !(field isa FullField) + # Windowed and reduced fields + fill_halo_regions!(field, args...; kwargs...) + else + push!(ordinary_fields, field) + end + end end - # Fill the rest - if !isempty(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, fields), + map(boundary_conditions, fields), + default_indices(3), + map(instantiated_location, fields), + grid, args...; kwargs...) end return nothing end -function tupled_fill_halo_regions!(fields, grid, args...; kwargs...) - - # We cannot group windowed fields together, the indices must be (:, :, :)! - indices = default_indices(3) - - return fill_halo_regions!(map(data, fields), - map(boundary_conditions, fields), - indices, - map(instantiated_location, fields), - grid, args...; kwargs...) -end - ##### ##### Tracer names ##### diff --git a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl index 599bb56453..f8ca935c96 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl @@ -38,9 +38,7 @@ function update_state!(model::HydrostaticFreeSurfaceModel, grid, callbacks; comp # Update the boundary conditions @apply_regionally update_boundary_condition!(fields(model), model) - fill_halo_regions!(model.free_surface.η, model.clock, fields(model), async=true) - prognostic_3d_fields = (model.velocities.u, model.velocities.v, model.tracers...) - tupled_fill_halo_regions!(prognostic_3d_fields, model.grid, model.clock, fields(model); async=true) + tupled_fill_halo_regions!(prognostic_fields(model), model.clock, fields(model), async=true) @apply_regionally replace_horizontal_vector_halos!(model.velocities, model.grid) @apply_regionally compute_auxiliaries!(model) From fd3031cc62fdb8000d946a54e5602bc29410b737 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 24 Oct 2024 17:38:57 -0600 Subject: [PATCH 09/25] Make the test better --- test/test_enzyme.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/test/test_enzyme.jl b/test/test_enzyme.jl index 51e0442730..687ba04b8b 100644 --- a/test/test_enzyme.jl +++ b/test/test_enzyme.jl @@ -245,8 +245,9 @@ function viscous_hydrostatic_turbulence(ν, model, u_init, v_init, Δt, u_truth, # Initialize the model model.clock.iteration = 0 model.clock.time = 0 + model.clock.last_Δt = Inf set_viscosity!(model, ν) - set!(model, u=u_init, v=v_init) + set!(model, u=u_init, v=v_init, η=0) # Step it forward for n = 1:10 @@ -300,7 +301,7 @@ end # Use a manual finite difference to compute a gradient Δν = 1e-6 - ν1 = 2e-2 + ν1 = ν₀ ν2 = ν1 + Δν e1 = viscous_hydrostatic_turbulence(ν1, model, u_init, v_init, Δt, u_truth, v_truth) e2 = viscous_hydrostatic_turbulence(ν2, model, u_init, v_init, Δt, u_truth, v_truth) From 0047d2f702944e0c5fa9140372694115c43bf1be Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 25 Oct 2024 08:01:54 -0600 Subject: [PATCH 10/25] Updates --- src/Fields/field.jl | 4 +- src/TimeSteppers/quasi_adams_bashforth_2.jl | 72 +++++++++++++-------- 2 files changed, 47 insertions(+), 29 deletions(-) diff --git a/src/Fields/field.jl b/src/Fields/field.jl index ee806db849..c17c120623 100644 --- a/src/Fields/field.jl +++ b/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) diff --git a/src/TimeSteppers/quasi_adams_bashforth_2.jl b/src/TimeSteppers/quasi_adams_bashforth_2.jl index ed9f55750b..7836662afa 100644 --- a/src/TimeSteppers/quasi_adams_bashforth_2.jl +++ b/src/TimeSteppers/quasi_adams_bashforth_2.jl @@ -52,6 +52,9 @@ end reset!(timestepper::QuasiAdamsBashforth2TimeStepper) = nothing +@inline zero!(::Nothing) = nothing +@inline zero!(field) = fill!(parent(field), 0) + ##### ##### Time steppping ##### @@ -79,39 +82,23 @@ function time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt Δt == 0 && @warn "Δt == 0 may cause model blowup!" # Be paranoid and update state at iteration 0 - model.clock.iteration == 0 && update_state!(model, callbacks) - - ab2_timestepper = model.timestepper + model.clock.iteration == 0 && update_state!(model, callbacks; compute_tendencies=true) - # Change the default χ if necessary, which occurs if: + # Take an euler step if: # * We detect that the time-step size has changed. # * We detect that this is the "first" time-step, which means we # need to take an euler step. Note that model.clock.last_Δt is # initialized as Inf # * The user has passed euler=true to time_step! euler = euler || (Δt != model.clock.last_Δt) - - # If euler, then set χ = -0.5 - minus_point_five = convert(eltype(model.grid), -0.5) - χ = ifelse(euler, minus_point_five, ab2_timestepper.χ) - - # Set time-stepper χ (this is used in ab2_step!, but may also be used elsewhere) - χ₀ = ab2_timestepper.χ # Save initial value - ab2_timestepper.χ = χ + euler && @debug "Taking a forward Euler step." - # Ensure zeroing out all previous tendency fields to avoid errors in - # case G⁻ includes NaNs. See https://github.com/CliMA/Oceananigans.jl/issues/2259 + # Full step for tracers, fractional step for velocities. if euler - @debug "Taking a forward Euler step." - for field in ab2_timestepper.G⁻ - !isnothing(field) && @apply_regionally fill!(field, 0) - end + euler_step!(model, Δt) + else + ab2_step!(model, Δt) end - - # Be paranoid and update state at iteration 0 - model.clock.iteration == 0 && update_state!(model, callbacks; compute_tendencies=true) - - ab2_step!(model, Δt) # full step for tracers, fractional step for velocities. tick!(model.clock, Δt) model.clock.last_Δt = Δt @@ -123,9 +110,6 @@ function time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt update_state!(model, callbacks; compute_tendencies=true) step_lagrangian_particles!(model, Δt) - # Return χ to initial value - ab2_timestepper.χ = χ₀ - return nothing end @@ -175,7 +159,7 @@ Time step velocity fields via the 2nd-order quasi Adams-Bashforth method @kernel function ab2_step_field!(u, Δt, χ, Gⁿ, G⁻) i, j, k = @index(Global, NTuple) - FT = eltype(χ) + FT = typeof(χ) one_point_five = convert(FT, 1.5) oh_point_five = convert(FT, 0.5) @@ -184,3 +168,37 @@ end @kernel ab2_step_field!(::FunctionField, Δt, χ, Gⁿ, G⁻) = nothing +function euler_step!(model, Δt) + grid = model.grid + arch = architecture(grid) + model_fields = prognostic_fields(model) + + for (i, field) in enumerate(model_fields) + kernel_args = (field, Δt, model.timestepper.Gⁿ[i]) + launch!(arch, grid, :xyz, euler_step_field!, kernel_args...; exclude_periphery=true) + + # TODO: function tracer_index(model, field_index) = field_index - 3, etc... + tracer_index = Val(i - 3) # assumption + + implicit_step!(field, + model.timestepper.implicit_solver, + model.closure, + model.diffusivity_fields, + tracer_index, + model.clock, + Δt) + end + + return nothing +end + +@kernel function euler_step_field!(u, Δt, Gⁿ) + i, j, k = @index(Global, NTuple) + FT = eltype(u) + @inbounds u[i, j, k] += convert(FT, Δt) * Gⁿ[i, j, k] +end + +@kernel euler_step_field!(::FunctionField, Δt, Gⁿ) = nothing + + + From 409cbb29936dd1366d5b70494e18a7fe1e6a7efe Mon Sep 17 00:00:00 2001 From: jlk9 Date: Fri, 25 Oct 2024 12:42:15 -0500 Subject: [PATCH 11/25] Figuring out nan entries --- test/test_enzyme.jl | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/test/test_enzyme.jl b/test/test_enzyme.jl index bad9d9e27b..8f77b4811e 100644 --- a/test/test_enzyme.jl +++ b/test/test_enzyme.jl @@ -278,8 +278,13 @@ end # Now for real e = 1e-5 dmodel = Enzyme.make_zero(model) + + @show dmodel + buoyancy_variance!(dmodel, 0) + @show dmodel + db²_de = autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse), buoyancy_variance!, Duplicated(model, dmodel), @@ -290,10 +295,13 @@ end Finite differences computed $db²_de_fd """ - tol = 0.01 - rel_error = abs(db²_de[1][3] - db²_de_fd) / abs(db²_de_fd) - @show db²_de, db²_de_fd - @test rel_error < tol + @show db²_de_fd + @show db²_de + + #tol = 0.01 + #rel_error = abs(db²_de[1][3] - db²_de_fd) / abs(db²_de_fd) + #@show db²_de, db²_de_fd + #@test rel_error < tol end end From 1522ae47996a4df9092a27e0d32bfffc1bf53fad Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 25 Oct 2024 14:16:09 -0600 Subject: [PATCH 12/25] Fix z-coord bug --- test/test_enzyme.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_enzyme.jl b/test/test_enzyme.jl index 687ba04b8b..58c619d601 100644 --- a/test/test_enzyme.jl +++ b/test/test_enzyme.jl @@ -272,7 +272,7 @@ end Nx = Ny = 32 Nz = 1 x = y = (0, 2π) - z = 1 + z = (0, 1) ν₀ = 1e-2 grid = RectilinearGrid(arch, size=(Nx, Ny, 1); x, y, z, topology=(Periodic, Periodic, Bounded)) From b97190b66b5fec6206bb42de524412c75049a9d6 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 25 Oct 2024 14:30:37 -0600 Subject: [PATCH 13/25] Few more bugs --- test/test_enzyme.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/test_enzyme.jl b/test/test_enzyme.jl index 58c619d601..fa09e5a1f7 100644 --- a/test/test_enzyme.jl +++ b/test/test_enzyme.jl @@ -259,8 +259,8 @@ function viscous_hydrostatic_turbulence(ν, model, u_init, v_init, Δt, u_truth, Nx, Ny, Nz = size(model.grid) err = 0.0 for j = 1:Ny, i = 1:Nx - err += @inbounds (u[i, j, 3] - u_truth[i, j, 3])^2 + - (v[i, j, 3] - v_truth[i, j, 3])^2 + err += @inbounds (u[i, j, 1] - u_truth[i, j, 1])^2 + + (v[i, j, 1] - v_truth[i, j, 1])^2 end return err::Float64 @@ -328,7 +328,7 @@ end @info "Elapsed time: " * prettytime(1e-9 * (time_ns() - start_time)) tol = 1e-1 - rel_error = abs(dedν[1][3] - ΔeΔν) / abs(ΔeΔν) + rel_error = abs(dedν[1][1] - ΔeΔν) / abs(ΔeΔν) @test rel_error < tol end From 08292d2400652b4676e097816fcb2a7f9c7dcdcc Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 25 Oct 2024 15:03:13 -0600 Subject: [PATCH 14/25] Fix bug in field_tuples --- src/Fields/field_tuples.jl | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/Fields/field_tuples.jl b/src/Fields/field_tuples.jl index 9d8e02913e..048a89011e 100644 --- a/src/Fields/field_tuples.jl +++ b/src/Fields/field_tuples.jl @@ -61,8 +61,8 @@ end function tupled_fill_halo_regions!(fields, args...; kwargs...) ordinary_fields = Field[] - for fields in fields - if !isnothing(boundary_conditions(f)) + for field in fields + if !isnothing(boundary_conditions(field)) if field isa ReducedField || !(field isa FullField) # Windowed and reduced fields fill_halo_regions!(field, args...; kwargs...) @@ -72,6 +72,8 @@ function tupled_fill_halo_regions!(fields, args...; kwargs...) end end + 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, fields), From 0f487733d25e489a937a6e2d62bd5759c1765ed6 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 25 Oct 2024 15:04:57 -0600 Subject: [PATCH 15/25] Tiny cleanup ; --- src/TimeSteppers/quasi_adams_bashforth_2.jl | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/src/TimeSteppers/quasi_adams_bashforth_2.jl b/src/TimeSteppers/quasi_adams_bashforth_2.jl index 7836662afa..04097790fd 100644 --- a/src/TimeSteppers/quasi_adams_bashforth_2.jl +++ b/src/TimeSteppers/quasi_adams_bashforth_2.jl @@ -52,9 +52,6 @@ end reset!(timestepper::QuasiAdamsBashforth2TimeStepper) = nothing -@inline zero!(::Nothing) = nothing -@inline zero!(field) = fill!(parent(field), 0) - ##### ##### Time steppping ##### @@ -99,7 +96,7 @@ function time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt else ab2_step!(model, Δt) end - + tick!(model.clock, Δt) model.clock.last_Δt = Δt model.clock.last_stage_Δt = Δt # just one stage @@ -200,5 +197,3 @@ end @kernel euler_step_field!(::FunctionField, Δt, Gⁿ) = nothing - - From 9b9575f9b35ef3eff5fe4c8e34e514858ba4dfc6 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 25 Oct 2024 15:08:10 -0600 Subject: [PATCH 16/25] Another bug --- src/Fields/field_tuples.jl | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/Fields/field_tuples.jl b/src/Fields/field_tuples.jl index 048a89011e..39d3122d53 100644 --- a/src/Fields/field_tuples.jl +++ b/src/Fields/field_tuples.jl @@ -61,13 +61,13 @@ end function tupled_fill_halo_regions!(fields, args...; kwargs...) ordinary_fields = Field[] - for field in fields - if !isnothing(boundary_conditions(field)) - if field isa ReducedField || !(field isa FullField) + for f in fields + if !isnothing(boundary_conditions(f)) + if f isa ReducedField || !(f isa FullField) # Windowed and reduced fields - fill_halo_regions!(field, args...; kwargs...) + fill_halo_regions!(f, args...; kwargs...) else - push!(ordinary_fields, field) + push!(ordinary_fields, f) end end end @@ -76,10 +76,10 @@ function tupled_fill_halo_regions!(fields, args...; kwargs...) if !isempty(ordinary_fields) # ie not reduced, and with default_indices grid = first(ordinary_fields).grid - fill_halo_regions!(map(data, fields), - map(boundary_conditions, fields), + fill_halo_regions!(map(data, ordinary_fields), + map(boundary_conditions, ordinary_fields), default_indices(3), - map(instantiated_location, fields), + map(instantiated_location, ordinary_fields), grid, args...; kwargs...) end From 423281bdacfd05062f7eaed6c46bf869ede03b6d Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 25 Oct 2024 16:44:35 -0600 Subject: [PATCH 17/25] Back to pure ab2 --- .../explicit_free_surface.jl | 11 ++- .../implicit_free_surface.jl | 4 +- .../split_explicit_free_surface_kernels.jl | 72 +++++++++++-------- src/TimeSteppers/quasi_adams_bashforth_2.jl | 56 +++++---------- 4 files changed, 69 insertions(+), 74 deletions(-) diff --git a/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl index 56fd8d0049..7d44fe9b44 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl +++ b/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 + diff --git a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl index e54ad5c49c..99d0d4def6 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl @@ -124,9 +124,9 @@ build_implicit_step_solver(::Val{:Default}, grid, settings, gravitational_accele Implicitly step forward η. """ ab2_step_free_surface!(free_surface::ImplicitFreeSurface, model, Δt, χ) = - implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model, Δt, χ) + implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model, Δt) -function implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model, Δt, χ) +function implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model, Δt) η = free_surface.η g = free_surface.gravitational_acceleration rhs = free_surface.implicit_step_solver.right_hand_side diff --git a/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl b/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl index 91ffed889f..9fb98656ce 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl @@ -293,14 +293,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) @@ -313,14 +313,14 @@ function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurfac settings = free_surface.settings Nsubsteps = calculate_substeps(settings.substepping, Δt) - # barotropic time step as fraction of baroclinic step and averaging weights + # 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 + # Barotropic time step in seconds Δτᴮ = fractional_Δt * Δt - # reset free surface averages + # Reset free surface averages @apply_regionally begin initialize_free_surface_state!(free_surface.state, free_surface.η, settings.timestepper) @@ -334,7 +334,7 @@ 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 + # Prepare velocities for the barotropic correction @apply_regionally begin mask_immersed_field!(model.velocities.u) mask_immersed_field!(model.velocities.v) @@ -415,40 +415,54 @@ 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ⁿ, χ) - end + 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 - @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 + 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 -@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, χ) @@ -464,17 +478,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 diff --git a/src/TimeSteppers/quasi_adams_bashforth_2.jl b/src/TimeSteppers/quasi_adams_bashforth_2.jl index 04097790fd..3471132ec3 100644 --- a/src/TimeSteppers/quasi_adams_bashforth_2.jl +++ b/src/TimeSteppers/quasi_adams_bashforth_2.jl @@ -90,12 +90,15 @@ function time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt euler = euler || (Δt != model.clock.last_Δt) euler && @debug "Taking a forward Euler step." + # If euler, then set χ = -0.5 + minus_point_five = convert(eltype(model.grid), -0.5) + ab2_timestepper = model.timestepper + χ = ifelse(euler, minus_point_five, ab2_timestepper.χ) + χ₀ = ab2_timestepper.χ # Save initial value + ab2_timestepper.χ = χ + # Full step for tracers, fractional step for velocities. - if euler - euler_step!(model, Δt) - else - ab2_step!(model, Δt) - end + ab2_step!(model, Δt) tick!(model.clock, Δt) model.clock.last_Δt = Δt @@ -107,6 +110,9 @@ function time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt update_state!(model, callbacks; compute_tendencies=true) step_lagrangian_particles!(model, Δt) + # Return χ to initial value + ab2_timestepper.χ = χ₀ + return nothing end @@ -122,7 +128,6 @@ end """ Generic implementation. """ function ab2_step!(model, Δt) - grid = model.grid arch = architecture(grid) model_fields = prognostic_fields(model) @@ -157,43 +162,16 @@ Time step velocity fields via the 2nd-order quasi Adams-Bashforth method i, j, k = @index(Global, NTuple) FT = typeof(χ) + Δt = convert(FT, Δt) one_point_five = convert(FT, 1.5) oh_point_five = convert(FT, 0.5) + not_euler = χ != convert(FT, -0.5) # use to prevent corruption by leftover NaNs in G⁻ - @inbounds u[i, j, k] += convert(FT, Δt) * ((one_point_five + χ) * Gⁿ[i, j, k] - (oh_point_five + χ) * G⁻[i, j, k]) -end - -@kernel ab2_step_field!(::FunctionField, Δt, χ, Gⁿ, G⁻) = nothing - -function euler_step!(model, Δt) - grid = model.grid - arch = architecture(grid) - model_fields = prognostic_fields(model) - - for (i, field) in enumerate(model_fields) - kernel_args = (field, Δt, model.timestepper.Gⁿ[i]) - launch!(arch, grid, :xyz, euler_step_field!, kernel_args...; exclude_periphery=true) - - # TODO: function tracer_index(model, field_index) = field_index - 3, etc... - tracer_index = Val(i - 3) # assumption - - implicit_step!(field, - model.timestepper.implicit_solver, - model.closure, - model.diffusivity_fields, - tracer_index, - model.clock, - Δt) + @inbounds begin + Gu = (one_point_five + χ) * Gⁿ[i, j, k] - (oh_point_five + χ) * G⁻[i, j, k] * not_euler + u[i, j, k] += Δt * Gu end - - return nothing end -@kernel function euler_step_field!(u, Δt, Gⁿ) - i, j, k = @index(Global, NTuple) - FT = eltype(u) - @inbounds u[i, j, k] += convert(FT, Δt) * Gⁿ[i, j, k] -end - -@kernel euler_step_field!(::FunctionField, Δt, Gⁿ) = nothing +@kernel ab2_step_field!(::FunctionField, Δt, χ, Gⁿ, G⁻) = nothing From fe0f9cfc36c21e58dc4558de54b9a86f3f0e8e62 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 25 Oct 2024 16:50:21 -0600 Subject: [PATCH 18/25] Make set! more differentiable --- src/Fields/set!.jl | 1 + .../set_hydrostatic_free_surface_model.jl | 5 +++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index e4b8a29cb6..bbdc65c3e7 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -33,6 +33,7 @@ end 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) +@inline set!(u::Field, a::Number) = fill!(u, a) function set!(u::Field, v) u .= v # fallback diff --git a/src/Models/HydrostaticFreeSurfaceModels/set_hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/set_hydrostatic_free_surface_model.jl index f216d312c7..a9a2c26006 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/set_hydrostatic_free_surface_model.jl +++ b/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 + From e66af6438363971c73b1a5f4e71cfba7a0b4dd30 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 25 Oct 2024 19:18:11 -0600 Subject: [PATCH 19/25] Restore syntax for some reason --- .../HydrostaticFreeSurfaceModels/implicit_free_surface.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl index 99d0d4def6..e54ad5c49c 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl @@ -124,9 +124,9 @@ build_implicit_step_solver(::Val{:Default}, grid, settings, gravitational_accele Implicitly step forward η. """ ab2_step_free_surface!(free_surface::ImplicitFreeSurface, model, Δt, χ) = - implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model, Δt) + implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model, Δt, χ) -function implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model, Δt) +function implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model, Δt, χ) η = free_surface.η g = free_surface.gravitational_acceleration rhs = free_surface.implicit_step_solver.right_hand_side From 15b582a1f7d93c159cded91de2cd48d2d2231c96 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 25 Oct 2024 19:20:07 -0600 Subject: [PATCH 20/25] More appropriate test w tolerance --- test/test_enzyme.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_enzyme.jl b/test/test_enzyme.jl index fa09e5a1f7..7136181b1e 100644 --- a/test/test_enzyme.jl +++ b/test/test_enzyme.jl @@ -301,7 +301,7 @@ end # Use a manual finite difference to compute a gradient Δν = 1e-6 - ν1 = ν₀ + ν1 = ν₀ + Δν ν2 = ν1 + Δν e1 = viscous_hydrostatic_turbulence(ν1, model, u_init, v_init, Δt, u_truth, v_truth) e2 = viscous_hydrostatic_turbulence(ν2, model, u_init, v_init, Δt, u_truth, v_truth) From eced1e51f704345fe6dc72a7f803cd985964ae3e Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 31 Oct 2024 16:03:15 -0600 Subject: [PATCH 21/25] hopefully fix type stability issue with set!(u, ::Number) --- src/Fields/set!.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index bbdc65c3e7..a9cff7a0b0 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -33,7 +33,11 @@ end 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) -@inline set!(u::Field, a::Number) = fill!(u, a) + +function set!(u::Field, a::Number) + fill!(parent(u), a) + return u # return u, not parent(u), for type-stability +end function set!(u::Field, v) u .= v # fallback From 5110a5732dd0581219077b9d5741ca374c7dffb6 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 31 Oct 2024 19:55:24 -0600 Subject: [PATCH 22/25] Starting to work on Simulation --- ext/OceananigansEnzymeExt.jl | 1 + src/Simulations/simulation.jl | 2 +- src/Utils/prettytime.jl | 63 ++++++++++++++++++----------------- 3 files changed, 34 insertions(+), 32 deletions(-) diff --git a/ext/OceananigansEnzymeExt.jl b/ext/OceananigansEnzymeExt.jl index 4c2fe35f63..5d03233ac2 100644 --- a/ext/OceananigansEnzymeExt.jl +++ b/ext/OceananigansEnzymeExt.jl @@ -16,6 +16,7 @@ EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Base.:(==)), ::Oceananigans.Abstr EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.AbstractOperations.validate_grid), x...) = nothing EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.AbstractOperations.metric_function), x...) = nothing EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.Utils.flatten_reduced_dimensions), x...) = nothing +EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.Utils.prettytime), x...) = nothing EnzymeCore.EnzymeRules.inactive(::typeof(Oceananigans.Grids.total_size), x...) = nothing EnzymeCore.EnzymeRules.inactive(::typeof(Oceananigans.BoundaryConditions.parent_size_and_offset), x...) = nothing @inline EnzymeCore.EnzymeRules.inactive_type(v::Type{Oceananigans.Utils.KernelParameters}) = true diff --git a/src/Simulations/simulation.jl b/src/Simulations/simulation.jl index 3b73d85702..3c8e4dbb6b 100644 --- a/src/Simulations/simulation.jl +++ b/src/Simulations/simulation.jl @@ -144,7 +144,7 @@ iteration(sim::Simulation) = iteration(sim.model) Return `sim.model.clock.time` as a prettily formatted string." """ -prettytime(sim::Simulation, longform=true) = prettytime(time(sim)) +prettytime(sim::Simulation, longform=true) = prettytime(time(sim); longform) """ run_wall_time(sim::Simulation) diff --git a/src/Utils/prettytime.jl b/src/Utils/prettytime.jl index eb7fa14154..ad1ae172dd 100644 --- a/src/Utils/prettytime.jl +++ b/src/Utils/prettytime.jl @@ -19,48 +19,49 @@ minutes, and hours. """ function prettytime(t, longform=true) # Modified from: https://github.com/JuliaCI/BenchmarkTools.jl/blob/master/src/trials.jl - - # Some shortcuts - s = longform ? "seconds" : "s" - iszero(t) && return "0 $s" - t < 1e-9 && return @sprintf("%.3e %s", t, s) # yah that's small - - t = maybe_int(t) value, units = prettytimeunits(t, longform) - if isinteger(value) - return @sprintf("%d %s", value, units) + if iszero(value) + msg = "0 $units" + elseif value < 1e-9 + msg = @sprintf("%.3e %s", value, units) # yah that's small + elseif isinteger(value) + msg = @sprintf("%d %s", value, units) else - return @sprintf("%.3f %s", value, units) + msg = @sprintf("%.3f %s", value, units) end + + return msg::String end -function prettytimeunits(t, longform=true) - t < 1e-9 && return t, "" # just _forget_ picoseconds! - t < 1e-6 && return t * 1e9, "ns" - t < 1e-3 && return t * 1e6, "μs" - t < 1 && return t * 1e3, "ms" - if t < minute +function prettytimeunits(t::T, longform=true) where T + if t < 1e-9 # just _forget_ picoseconds! value = t - !longform && return value, "s" - units = value == 1 ? "second" : "seconds" - return value, units + units = !longform ? "s" : "seconds" + elseif t < 1e-6 + value = t * 1e9 + units = "ns" + elseif t < 1e-3 + value = t * 1e6 + units = "μs" + elseif t < 1 + value = t * 1e3 + units = "ms" + elseif t < minute + value = t + units = !longform ? "s" : value==1 ? "second" : "seconds" elseif t < hour - value = maybe_int(t / minute) - !longform && return value, "m" - units = value == 1 ? "minute" : "minutes" - return value, units + value = t / minute + units = !longform ? "m" : value==1 ? "minute" : "minutes" elseif t < day - value = maybe_int(t / hour) - units = value == 1 ? (longform ? "hour" : "hr") : - (longform ? "hours" : "hrs") - return value, units + value = t / hour + units = !longform ? "hr" : value==1 ? "hour" : "hours" else - value = maybe_int(t / day) - !longform && return value, "d" - units = value == 1 ? "day" : "days" - return value, units + value = t / day + units = !longform ? "d" : value==1 ? "day" : "days" end + + return convert(T, value), units::String end prettytime(dt::AbstractTime) = "$dt" From afd8fdd1ca65cdfbc20ac2920309eeb6e77d6697 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 31 Oct 2024 21:31:16 -0600 Subject: [PATCH 23/25] Implement get_top_tracer_bcs utility --- .../TKEBasedVerticalDiffusivities.jl | 10 +++++++++- .../catke_vertical_diffusivity.jl | 4 +++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/TKEBasedVerticalDiffusivities.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/TKEBasedVerticalDiffusivities.jl index 9553c5c4dd..acf28a0a87 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/TKEBasedVerticalDiffusivities.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/TKEBasedVerticalDiffusivities.jl @@ -19,6 +19,8 @@ using Oceananigans.BoundaryConditions: default_prognostic_bc, DefaultBoundaryCon using Oceananigans.BoundaryConditions: BoundaryCondition, FieldBoundaryConditions using Oceananigans.BoundaryConditions: DiscreteBoundaryFunction, FluxBoundaryCondition using Oceananigans.BuoyancyModels: ∂z_b, top_buoyancy_flux +using Oceananigans.BuoyancyModels: BuoyancyTracer, SeawaterBuoyancy +using Oceananigans.BuoyancyModels: TemperatureSeawaterBuoyancy, SalinitySeawaterBuoyancy using Oceananigans.Grids: inactive_cell using Oceananigans.TurbulenceClosures: @@ -160,8 +162,14 @@ function get_time_step(closure_array::AbstractArray) return get_time_step(closure) end -include("tke_top_boundary_condition.jl") +get_top_tracer_bcs(::Nothing, tracers) = NamedTuple() +get_top_tracer_bcs(::BuoyancyTracer, tracers) = (; b=tracers.b.boundary_conditions.top) +get_top_tracer_bcs(::SeawaterBuoyancy, tracers) = (T = tracers.T.boundary_conditions.top, + S = tracers.S.boundary_conditions.top) +get_top_tracer_bcs(::TemperatureSeawaterBuoyancy, tracers) = (; T = tracers.T.boundary_conditions.top) +get_top_tracer_bcs(::SalinitySeawaterBuoyancy, tracers) = (; S = tracers.S.boundary_conditions.top) +include("tke_top_boundary_condition.jl") include("catke_vertical_diffusivity.jl") include("catke_mixing_length.jl") include("catke_equation.jl") diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl index 10f5e3c6d0..9bcf287483 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl @@ -1,3 +1,5 @@ +using Oceananigans.BuoyancyModels: BuoyancyTracer, SeawaterBuoyancy + struct CATKEVerticalDiffusivity{TD, CL, FT, DT, TKE} <: AbstractScalarDiffusivity{TD, VerticalFormulation, 2} mixing_length :: CL turbulent_kinetic_energy_equation :: TKE @@ -214,7 +216,7 @@ function compute_diffusivities!(diffusivities, closure::FlavorOfCATKE, model; pa tracers = model.tracers buoyancy = model.buoyancy clock = model.clock - top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + top_tracer_bcs = get_top_tracer_bcs(model.buoyancy.model, tracers) Δt = model.clock.time - diffusivities.previous_compute_time[] diffusivities.previous_compute_time[] = model.clock.time From 3559b8da235e8df55cc3c571b3fd5ed46294c328 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 31 Oct 2024 21:31:26 -0600 Subject: [PATCH 24/25] Cleanup --- src/BuoyancyModels/seawater_buoyancy.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/BuoyancyModels/seawater_buoyancy.jl b/src/BuoyancyModels/seawater_buoyancy.jl index ba7fa68390..7472fc5615 100644 --- a/src/BuoyancyModels/seawater_buoyancy.jl +++ b/src/BuoyancyModels/seawater_buoyancy.jl @@ -9,15 +9,15 @@ temperature and salinity are active, or of type `FT` if temperature or salinity are constant, respectively. """ struct SeawaterBuoyancy{FT, EOS, T, S} <: AbstractBuoyancyModel{EOS} - equation_of_state :: EOS + equation_of_state :: EOS gravitational_acceleration :: FT - constant_temperature :: T - constant_salinity :: S + constant_temperature :: T + constant_salinity :: S end required_tracers(::SeawaterBuoyancy) = (:T, :S) -required_tracers(::SeawaterBuoyancy{FT, EOS, <:Nothing, <:Number}) where {FT, EOS} = (:T,) # active temperature only -required_tracers(::SeawaterBuoyancy{FT, EOS, <:Number, <:Nothing}) where {FT, EOS} = (:S,) # active salinity only +required_tracers(::SeawaterBuoyancy{FT, EOS, <:Nothing, <:Number}) where {FT, EOS} = tuple(:T) # active temperature only +required_tracers(::SeawaterBuoyancy{FT, EOS, <:Number, <:Nothing}) where {FT, EOS} = tuple(:S) # active salinity only Base.nameof(::Type{SeawaterBuoyancy}) = "SeawaterBuoyancy" Base.summary(b::SeawaterBuoyancy) = string(nameof(typeof(b)), " with g=", prettysummary(b.gravitational_acceleration), From 3a2d4a092c421063021f108953a7e14c65c94444 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 31 Oct 2024 21:31:42 -0600 Subject: [PATCH 25/25] Type-stabilize single column mode --- .../HydrostaticFreeSurfaceModels/single_column_model_mode.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl index 4ad0225287..7d0d854a18 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl @@ -63,7 +63,7 @@ compute_free_surface_tendency!(::SingleColumnGrid, ::SplitExplicitFreeSurfaceHFS function update_state!(model::HydrostaticFreeSurfaceModel, grid::SingleColumnGrid, callbacks; compute_tendencies = true) - fill_halo_regions!(prognostic_fields(model), model.clock, fields(model)) + tupled_fill_halo_regions!(prognostic_fields(model), model.clock, fields(model)) # Compute auxiliaries compute_auxiliary_fields!(model.auxiliary_fields)