(0.93.2) Update Adapt.jl compat and fix Float32 CATKE on GPU (#3876)

Co-authored-by: Gregory L. Wagner <[email protected]>
CliMA · Oct 30, 2024 · f2a8fb3 · f2a8fb3 · ali-ramadhan · Oct 30, 2024
1 parent ea7eb5f
commit f2a8fb3
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 25 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.93.1"
+version = "0.93.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -43,7 +43,7 @@ OceananigansEnzymeExt = "Enzyme"
 OceananigansMakieExt = ["MakieCore", "Makie"]
 
 [compat]
-Adapt = "3, 4"
+Adapt = "^4.1.1"
 CUDA = "4.1.1, 5"
 Crayons = "4"
 CubedSphere = "0.2, 0.3"

diff --git a/src/Grids/grid_generation.jl b/src/Grids/grid_generation.jl
@@ -129,4 +129,3 @@ generate_coordinate(FT, ::Flat, N, H, c::Number, coordinate_name, arch) =
 
 generate_coordinate(FT, ::Flat, N, H, ::Nothing, coordinate_name, arch) =
     FT(1), nothing, nothing, FT(1), FT(1)
-
diff --git a/...s/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_mixing_length.jl b/...s/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_mixing_length.jl
@@ -16,7 +16,7 @@ Base.@kwdef struct CATKEMixingLength{FT}
     Cˢ   :: FT = 1.131  # Surface distance coefficient for shear length scale
     Cᵇ   :: FT = Inf    # Bottom distance coefficient for shear length scale
     Cˢᵖ  :: FT = 0.505  # Sheared convective plume coefficient
-    CRiᵟ :: FT = 1.02   # Stability function width 
+    CRiᵟ :: FT = 1.02   # Stability function width
     CRi⁰ :: FT = 0.254  # Stability function lower Ri
     Cʰⁱu :: FT = 0.242  # Shear mixing length coefficient for momentum at high Ri
     Cˡᵒu :: FT = 0.361  # Shear mixing length coefficient for momentum at low Ri
@@ -131,7 +131,7 @@ end
     ϵˢᵖ = 1 - Cˢᵖ * Riᶠ                    # ϵ = Sheared convection factor
     ℓᵉ = clip(ϵˢᵖ * ℓᵉ)
     =#
-    
+
     # Figure out which mixing length applies
     convecting = (Jᵇ > Jᵇᵋ) & (N² < 0)
     entraining = (Jᵇ > Jᵇᵋ) & (N² > 0) & (N²_above < 0)
@@ -299,4 +299,3 @@ Base.show(io::IO, ml::CATKEMixingLength) =
               " ├── Cˢᵖ:  ", ml.Cˢᵖ, '\n',
               " ├── CRiᵟ: ", ml.CRiᵟ, '\n',
               " └── CRi⁰: ", ml.CRi⁰)
-
diff --git a/...lence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl b/...lence_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
@@ -18,7 +17,7 @@ function CATKEVerticalDiffusivity{TD}(mixing_length::CL,
                                       maximum_viscosity::FT,
                                       minimum_tke::FT,
                                       minimum_convective_buoyancy_flux::FT,
-                                      negative_tke_damping_time_scale::FT, 
+                                      negative_tke_damping_time_scale::FT,
                                       tke_time_step::DT) where {TD, CL, FT, DT, TKE}
 
     return CATKEVerticalDiffusivity{TD, CL, FT, DT, TKE}(mixing_length,
@@ -173,7 +172,7 @@ function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfCATKE)
     return (; κu, κc, κe, Le, Jᵇ,
             previous_compute_time, previous_velocities,
             _tupled_tracer_diffusivities, _tupled_implicit_linear_coefficients)
-end        
+end
 
 @inline viscosity_location(::FlavorOfCATKE) = (c, c, f)
 @inline diffusivity_location(::FlavorOfCATKE) = (c, c, f)
@@ -228,7 +227,7 @@ end
     Jᵇᵋ = closure.minimum_convective_buoyancy_flux
     Jᵇᵢⱼ = @inbounds Jᵇ[i, j, 1]
     Jᵇ⁺ = max(Jᵇᵋ, Jᵇᵢⱼ, Jᵇ★) # selects fastest (dominant) time-scale
-    t★ = (ℓᴰ^2 / Jᵇ⁺)^(1/3)
+    t★ = cbrt(ℓᴰ^2 / Jᵇ⁺)
     ϵ = Δt / t★
 
     @inbounds Jᵇ[i, j, 1] = (Jᵇᵢⱼ + ϵ * Jᵇ★) / (1 + ϵ)
@@ -263,28 +262,34 @@ end
     ℓu = momentum_mixing_lengthᶜᶜᶠ(i, j, k, grid, closure, velocities, tracers, buoyancy, surface_buoyancy_flux)
     κu = ℓu * w★
     κu_max = closure.maximum_viscosity
-    return min(κu, κu_max)
+    κu★ = min(κu, κu_max)
+    FT = eltype(grid)
+    return κu★::FT
 end
 
 @inline function κcᶜᶜᶠ(i, j, k, grid, closure, velocities, tracers, buoyancy, surface_buoyancy_flux)
     w★ = ℑzᵃᵃᶠ(i, j, k, grid, turbulent_velocityᶜᶜᶜ, closure, tracers.e)
     ℓc = tracer_mixing_lengthᶜᶜᶠ(i, j, k, grid, closure, velocities, tracers, buoyancy, surface_buoyancy_flux)
     κc = ℓc * w★
     κc_max = closure.maximum_tracer_diffusivity
-    return min(κc, κc_max)
+    κc★ = min(κc, κc_max)
+    FT = eltype(grid)
+    return κc★::FT
 end
 
 @inline function κeᶜᶜᶠ(i, j, k, grid, closure, velocities, tracers, buoyancy, surface_buoyancy_flux)
     w★ = ℑzᵃᵃᶠ(i, j, k, grid, turbulent_velocityᶜᶜᶜ, closure, tracers.e)
     ℓe = TKE_mixing_lengthᶜᶜᶠ(i, j, k, grid, closure, velocities, tracers, buoyancy, surface_buoyancy_flux)
     κe = ℓe * w★
     κe_max = closure.maximum_tke_diffusivity
-    return min(κe, κe_max)
+    κe★ = min(κe, κe_max)
+    FT = eltype(grid)
+    return κe★::FT
 end
 
 @inline viscosity(::FlavorOfCATKE, diffusivities) = diffusivities.κu
 @inline diffusivity(::FlavorOfCATKE, diffusivities, ::Val{id}) where id = diffusivities._tupled_tracer_diffusivities[id]
-    
+
 #####
 ##### Show
 #####
@@ -334,4 +339,3 @@ function Base.show(io::IO, clo::CATKEVD)
               "    ├── CᵂwΔ: ", prettysummary(clo.turbulent_kinetic_energy_equation.CᵂwΔ), '\n',
               "    └── Cᵂϵ:  ", prettysummary(clo.turbulent_kinetic_energy_equation.Cᵂϵ))
 end
-
diff --git a/...bulence_closure_implementations/TKEBasedVerticalDiffusivities/time_step_catke_equation.jl b/...bulence_closure_implementations/TKEBasedVerticalDiffusivities/time_step_catke_equation.jl
@@ -117,7 +117,7 @@ end
     # Then the contribution of Jᵉ to the implicit flux is
     #
     #       Lᵂ = - Cᵂϵ * √e / Δz.
-    
+
     on_bottom = !inactive_cell(i, j, k, grid) & inactive_cell(i, j, k-1, grid)
     Δz = Δzᶜᶜᶜ(i, j, k, grid)
     Cᵂϵ = closure_ij.turbulent_kinetic_energy_equation.Cᵂϵ
@@ -190,7 +190,7 @@ function tracer_tendency_kernel_function(model::HFSM, ::Val{:e}, closures::Tuple
     else
         catke_closure = closures[catke_index]
         catke_diffusivity_fields = diffusivity_fields[catke_index]
-        return compute_hydrostatic_free_surface_Ge!, catke_closure, catke_diffusivity_fields 
+        return compute_hydrostatic_free_surface_Ge!, catke_closure, catke_diffusivity_fields
     end
 end
 
@@ -202,4 +202,3 @@ end
     return NamedTuple{tuple(names...)}(tuple(values...))
 end
 =#
-
diff --git a/test/test_hydrostatic_free_surface_models.jl b/test/test_hydrostatic_free_surface_models.jl
@@ -54,6 +54,31 @@ function hydrostatic_free_surface_model_tracers_and_forcings_work(arch)
     return nothing
 end
 
+function time_step_hydrostatic_model_with_catke_works(arch, FT)
+    grid = LatitudeLongitudeGrid(
+        arch,
+        FT,
+        topology = (Bounded, Bounded, Bounded),
+        size = (8, 8, 8),
+        longitude = (0, 1),
+        latitude = (0, 1),
+        z = (-100, 0)
+    )
+
+    model = HydrostaticFreeSurfaceModel(;
+        grid,
+        buoyancy = BuoyancyTracer(),
+        tracers = (:b, :e),
+        closure = CATKEVerticalDiffusivity(FT)
+    )
+
+    simulation = Simulation(model, Δt=1.0, stop_iteration=1)
+
+    run!(simulation)
+
+    return model.clock.iteration == 1
+end
+
 topo_1d = (Flat, Flat, Bounded)
 
 topos_2d = ((Periodic, Flat, Bounded),
@@ -66,7 +91,7 @@ topos_3d = ((Periodic, Periodic, Bounded),
 
 @testset "Hydrostatic free surface Models" begin
     @info "Testing hydrostatic free surface models..."
-      
+
     @testset "$topo_1d model construction" begin
         @info "  Testing $topo_1d model construction..."
         for arch in archs, FT in [Float64] #float_types
@@ -80,7 +105,7 @@ topos_3d = ((Periodic, Periodic, Bounded),
             @test !(:η ∈ keys(fields(model))) # doesn't include free surface
         end
     end
-    
+
     for topo in topos_2d
         @testset "$topo model construction" begin
             @info "  Testing $topo model construction..."
@@ -92,7 +117,7 @@ topos_3d = ((Periodic, Periodic, Bounded),
             end
         end
     end
-    
+
     for topo in topos_3d
         @testset "$topo model construction" begin
             @info "  Testing $topo model construction..."
@@ -180,8 +205,8 @@ topos_3d = ((Periodic, Periodic, Bounded),
         precompute_metrics = true
         lat_lon_sector_grid = LatitudeLongitudeGrid(arch; size=(H, H, H), longitude=(0, 60), latitude=(15, 75), z=(-1, 0), precompute_metrics, halo)
         lat_lon_strip_grid  = LatitudeLongitudeGrid(arch; size=(H, H, H), longitude=(-180, 180), latitude=(15, 75), z=(-1, 0), precompute_metrics, halo)
-        
-        z = z_face_generator() 
+
+        z = z_face_generator()
         lat_lon_sector_grid_stretched = LatitudeLongitudeGrid(arch; size=(H, H, H), longitude=(0, 60), latitude=(15, 75), z, precompute_metrics, halo)
         lat_lon_strip_grid_stretched  = LatitudeLongitudeGrid(arch; size=(H, H, H), longitude=(-180, 180), latitude=(15, 75), z, precompute_metrics, halo)
 
@@ -196,7 +221,7 @@ topos_3d = ((Periodic, Periodic, Bounded),
                 topo = topology(grid)
                 grid_type = typeof(grid).name.wrapper
                 free_surface_type = typeof(free_surface).name.wrapper
-                test_label = "[$arch, $grid_type, $topo, $free_surface_type]"                
+                test_label = "[$arch, $grid_type, $topo, $free_surface_type]"
                 @testset "Time-stepping HydrostaticFreeSurfaceModels with various grids $test_label" begin
                     @info "  Testing time-stepping HydrostaticFreeSurfaceModels with various grids $test_label..."
                     @test time_step_hydrostatic_model_works(grid; free_surface)
@@ -278,7 +303,7 @@ topos_3d = ((Periodic, Periodic, Bounded),
 
             @test time_step_hydrostatic_model_works(rectilinear_grid, momentum_advection  = nothing, velocities = velocities)
             @test time_step_hydrostatic_model_works(lat_lon_sector_grid, momentum_advection = nothing, velocities = velocities)
-                                            
+
             parameters = (U=1, m=0.1, W=0.001)
             u(x, y, z, t, p) = p.U
             v(x, y, z, t, p) = exp(p.m * z)
@@ -294,5 +319,11 @@ topos_3d = ((Periodic, Periodic, Bounded),
             @info "  Testing HydrostaticFreeSurfaceModel with tracers and forcings [$arch]..."
             hydrostatic_free_surface_model_tracers_and_forcings_work(arch)
         end
+
+        # See: https://github.com/CliMA/Oceananigans.jl/issues/3870
+        @testset "HydrostaticFreeSurfaceModel with Float32 CATKE [$arch]" begin
+            @info "  Testing HydrostaticFreeSurfaceModel with Float32 CATKE [$arch]..."
+            @test time_step_hydrostatic_model_with_catke_works(arch, Float32)
+        end
     end
 end
Original file line number	Diff line number	Diff line change
Expand Up		@@ -129,4 +129,3 @@ generate_coordinate(FT, ::Flat, N, H, c::Number, coordinate_name, arch) =

		generate_coordinate(FT, ::Flat, N, H, ::Nothing, coordinate_name, arch) =
		FT(1), nothing, nothing, FT(1), FT(1)