Merge pull request #138 from OceanBioME/jsw/gpu-sediment-bug

(0.7.0) Fix all the GPU bugs that have crept in

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.9.3"
 manifest_format = "2.0"
-project_hash = "b8336a615981d764881d88f0d8193dd7f4667beb"
+project_hash = "d35eb03107a4f7f3d176d38b0533075a0ad26402"
 
 [[deps.AbstractFFTs]]
 deps = ["LinearAlgebra"]
@@ -500,9 +500,9 @@ version = "1.2.0"
 
 [[deps.Oceananigans]]
 deps = ["Adapt", "CUDA", "Crayons", "CubedSphere", "Dates", "Distances", "DocStringExtensions", "FFTW", "Glob", "IncompleteLU", "InteractiveUtils", "IterativeSolvers", "JLD2", "KernelAbstractions", "LinearAlgebra", "Logging", "MPI", "NCDatasets", "OffsetArrays", "OrderedCollections", "PencilArrays", "PencilFFTs", "Pkg", "Printf", "Random", "Rotations", "SeawaterPolynomials", "SparseArrays", "Statistics", "StructArrays"]
-git-tree-sha1 = "8746c619e15950de59160b2d99e61961cfa57fcc"
+git-tree-sha1 = "1821dce5269caa6f9dc3df3e1b97928b39f233b5"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
-version = "0.87.2"
+version = "0.88.0"
 
 [[deps.OffsetArrays]]
 deps = ["Adapt"]

Project.toml

@@ -1,7 +1,7 @@
 name = "OceanBioME"
 uuid = "a49af516-9db8-4be4-be45-1dad61c5a376"
 authors = ["Jago Strong-Wright <[email protected]>", "John R Taylor <[email protected]>", "Si Chen <[email protected]>"]
-version = "0.6.0"
+version = "0.7.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -15,10 +15,11 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 Adapt = "3"
 JLD2 = "0.4"
 KernelAbstractions = "0.9"
-Oceananigans = "0.84.1, 0.85, 0.86, 0.87"
+Oceananigans = "0.84.1, 0.85, 0.86, 0.87, 0.88"
 Roots = "2"
 StructArrays = "0.4, 0.5, 0.6"
 julia = "1.9"
+Documenter = "0.27"
 
 [extras]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

docs/Project.toml

@@ -18,3 +18,4 @@ CairoMakie = "0.10"
 DocumenterCitations = "1"
 Literate = "≥2.9.0"
 Oceananigans = "0.84.1, 0.85, 0.86, 0.87"
+Documenter = "0.27"

docs/make.jl

@@ -53,7 +53,7 @@ model_parameters = (LOBSTER(; grid = BoxModelGrid()),
                     NutrientPhytoplanktonZooplanktonDetritus(; grid = BoxModelGrid()),
                     SLatissima(),
                     TwoBandPhotosyntheticallyActiveRadiation(; grid = BoxModelGrid()),
-                    SimpleMultiG(; grid = BoxModelGrid(), depth = 1000),
+                    SimpleMultiG(; grid = BoxModelGrid()),
                     InstantRemineralisation(; grid = BoxModelGrid()),
                     OCMIP_default,
                     GasExchange(; gas = :CO₂).condition.parameters,

docs/src/model_components/utils.md

@@ -8,11 +8,6 @@ We have added a few additional utilities which extend the capabilities of Oceana
 wizard = TimeStepWizard(cfl = 0.2, diffusive_cfl = 0.2, max_change = 2.0, min_change = 0.5, cell_advection_timescale = column_advection_timescale)
 simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))
 ```
-Additionally, in a column model you may have a functional definition for the viscosity, so we define an additional diffusion timescale function:
-```julia
-wizard = TimeStepWizard(cfl = 0.2, diffusive_cfl = 0.2, max_change = 2.0, min_change = 0.5, cell_diffusion_timescale = column_diffusion_timescale, cell_advection_timescale = column_advection_timescale)
-simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))
-```
 Finally, sinking may be more limiting than the normal advective CFL conditions so, we have an additional cell advection timescale defined for 3D models:
 ```julia
 wizard = TimeStepWizard(cfl = 0.6, diffusive_cfl = 0.5, max_change = 1.5, min_change = 0., cell_advection_timescale = sinking_advection_timescale)

docs/src/model_implementation.md

@@ -213,14 +213,16 @@ Another aspect that OceanBioME includes is sediment models. Doing this varies be
 ```@example implementing
 using OceanBioME.Boundaries.Sediments: sinking_flux
 
-import OceanBioME.Boundaries.Sediments: nitrogen_flux, carbon_flux, remineralisation_receiver
+import OceanBioME.Boundaries.Sediments: nitrogen_flux, carbon_flux, remineralisation_receiver, sinking_tracers
 
 @inline nitrogen_flux(i, j, k, grid, advection, bgc::NutrientPhytoplankton, tracers) =
      sinking_flux(i, j, k, grid, advection, Val(:P), bgc, tracers)
                  
 @inline carbon_flux(i, j, k, grid, advection, bgc::NutrientPhytoplankton, tracers) = nitrogen_flux(i, j, k, grid, advection, bgc, tracers) * 6.56
 
 @inline remineralisation_receiver(::NutrientPhytoplankton) = :N
+
+@inline sinking_tracers(::NutrientPhytoplankton) = (:P, )
 ```
 
 ### Putting it together

examples/data_forced.jl

@@ -20,6 +20,12 @@
 using OceanBioME
 using Oceananigans, Random, Printf, NetCDF, Interpolations, DataDeps
 using Oceananigans.Units
+using Oceananigans.Fields: FunctionField
+
+import Oceananigans.TurbulenceClosures: maximum_numeric_diffusivity
+
+maximum_numeric_diffusivity(κ::NamedTuple) = maximum(maximum.(values(κ)))
+maximum_numeric_diffusivity(κ::FunctionField) = maximum(κ)
 
 const year = years = 365days # just for these idealised cases
 nothing #hide
@@ -48,9 +54,10 @@ t_function(x, y, z, t) = temperature_itp(mod(t, 364days))
 s_function(x, y, z, t) = salinity_itp(mod(t, 364days))
 surface_PAR(x, y, t) = PAR_itp(mod(t, 364days))
 κₜ(x, y, z, t) = 2e-2 * max(1 - (z + mld_itp(mod(t, 364days)) / 2)^2 / (mld_itp(mod(t, 364days)) / 2)^2, 0) + 1e-4
+
 nothing #hide
 
-# ## Grid and PAR field
+# ## Grid and diffusivity field
 # Define the grid (in this case a non uniform grid for better resolution near the surface) and an extra Oceananigans field for the PAR to be stored in
 Nz = 33
 Lz = 600meters
@@ -63,6 +70,10 @@ z_faces(k) = Lz * (ζ₀(k) * Σ(k) - 1)
 
 grid = RectilinearGrid(size = (1, 1, Nz), x = (0, 20meters), y = (0, 20meters), z = z_faces)
 
+clock = Clock(; time = 0.0)
+
+κ = FunctionField{Center, Center, Center}(κₜ, grid; clock)
+
 # ## Biogeochemical and Oceananigans model
 # Here we instantiate the LOBSTER model with carbonate chemistry and a surface flux of DIC (CO₂)
 
@@ -73,8 +84,8 @@ biogeochemistry = LOBSTER(; grid,
 
 CO₂_flux = GasExchange(; gas = :CO₂, temperature = t_function, salinity = s_function)
 
-model = NonhydrostaticModel(; grid,
-                              closure = ScalarDiffusivity(ν = κₜ, κ = κₜ),
+model = NonhydrostaticModel(; grid, clock,
+                              closure = ScalarDiffusivity(ν = κ, κ = κ),
                               biogeochemistry,
                               boundary_conditions = (DIC = FieldBoundaryConditions(top = CO₂_flux),))
 
@@ -106,7 +117,6 @@ simulation.output_writers[:profiles] = JLD2OutputWriter(model,
 
 wizard = TimeStepWizard(cfl = 0.2, diffusive_cfl = 0.2,
                         max_change = 1.5, min_change = 0.75,
-                        cell_diffusion_timescale = column_diffusion_timescale,
                         cell_advection_timescale = column_advection_timescale)
 
 simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))

src/Boundaries/Sediments/Sediments.jl

@@ -7,7 +7,7 @@ using KernelAbstractions
 using OceanBioME: Biogeochemistry, BoxModelGrid
 
 using Oceananigans
-using Oceananigans.Architectures: device
+using Oceananigans.Architectures: device, architecture, arch_array
 using Oceananigans.Utils: launch!
 using Oceananigans.Advection: advective_tracer_flux_z
 using Oceananigans.Units: day
@@ -35,13 +35,20 @@ function update_tendencies!(bgc, sediment::FlatSediment, model)
         launch!(arch, model.grid, :xy, store_flat_tendencies!, sediment.tendencies.G⁻[i], sediment.tendencies.Gⁿ[i])
     end
 
+    biogeochemistry = bgc.underlying_biogeochemistry
+
     launch!(arch, model.grid, :xy,
             _calculate_tendencies!,
-            sediment, bgc, model.grid, model.advection, model.tracers, model.timestepper)
+            sediment, biogeochemistry, model.grid, 
+            sinking_advection(biogeochemistry, model.advection), 
+            required_tracers(sediment, biogeochemistry, model.tracers), 
+            required_tendencies(sediment, biogeochemistry, model.timestepper.Gⁿ), 
+            sediment.tendencies.Gⁿ)
 
     return nothing
 end
 
+
 @kernel function store_flat_tendencies!(G⁻, G⁰)
     i, j = @index(Global, NTuple)
     @inbounds G⁻[i, j, 1] = G⁰[i, j, 1]
@@ -50,10 +57,12 @@ end
 @inline nitrogen_flux() = 0
 @inline carbon_flux() = 0
 @inline remineralisation_receiver() = nothing
+@inline sinking_tracers() = nothing
+@inline required_tracers() = nothing
+@inline required_tendencies() = nothing
 
-@inline nitrogen_flux(i, j, k, grid, advection, bgc::Biogeochemistry, tracers) = nitrogen_flux(i, j, k, grid, advection, bgc.underlying_biogeochemistry, tracers)
-@inline carbon_flux(i, j, k, grid, advection, bgc::Biogeochemistry, tracers) = carbon_flux(i, j, k, grid, advection, bgc.underlying_biogeochemistry, tracers)
-@inline remineralisation_receiver(bgc::Biogeochemistry) = remineralisation_receiver(bgc.underlying_biogeochemistry)
+@inline sinking_advection(bgc, advection) = advection
+@inline sinking_advection(bgc, advection::NamedTuple) = advection[sinking_tracers(bgc)]
 
 @inline advection_scheme(advection, val_tracer) = advection
 @inline advection_scheme(advection::NamedTuple, val_tracer::Val{T}) where T = advection[T]
@@ -81,7 +90,8 @@ calculate_bottom_indices(grid) = ones(Int, size(grid)[1:2]...)
 end
 
 function calculate_bottom_indices(grid::ImmersedBoundaryGrid)
-    indices = zeros(Int, size(grid)[1:2]...)
+    arch = architecture(grid)
+    indices = arch_array(arch, zeros(Int, size(grid)[1:2]...))
 
     launch!(grid.architecture, grid, :xy, find_bottom_cell, grid, indices)
 

src/Boundaries/Sediments/instant_remineralization.jl

@@ -63,7 +63,7 @@ function InstantRemineralisation(; grid,
     tendencies = (Gⁿ = NamedTuple{tracer_names}(Tuple(CenterField(grid) for tracer in tracer_names)),
                   G⁻ = NamedTuple{tracer_names}(Tuple(CenterField(grid) for tracer in tracer_names)))
 
-    bottom_indices = calculate_bottom_indices(grid)
+    bottom_indices = arch_array(architecture(grid), calculate_bottom_indices(grid))
 
     return InstantRemineralisation(burial_efficiency_constant1,
                                    burial_efficiency_constant2,
@@ -74,18 +74,22 @@ function InstantRemineralisation(; grid,
 end
 
 adapt_structure(to, sediment::InstantRemineralisation) = 
-    SimpleMultiG(sediment.burial_efficiency_constant1,
-                 sediment.burial_efficiency_constant2,
-                 sediment.burial_efficiency_half_saturaiton,
-                 adapt(to, sediment.fields),
-                 adapt(to, sediment.tendencies))
+    InstantRemineralisation(adapt(to, sediment.burial_efficiency_constant1),
+                            adapt(to, sediment.burial_efficiency_constant2),
+                            adapt(to, sediment.burial_efficiency_half_saturaiton),
+                            adapt(to, sediment.fields),
+                            nothing,
+                            adapt(to, sediment.bottom_indices))
 
 sediment_tracers(::InstantRemineralisation) = (:N_storage, )
 sediment_fields(model::InstantRemineralisation) = (N_storage = model.fields.N_storage, )
 
+@inline required_tracers(::InstantRemineralisation, bgc, tracers) = tracers[(sinking_tracers(bgc)..., remineralisation_receiver(bgc))]
+@inline required_tendencies(::InstantRemineralisation, bgc, tracers) = tracers[remineralisation_receiver(bgc)]
+
 @inline bottom_index_array(sediment::InstantRemineralisation) = sediment.bottom_indices
 
-@kernel function _calculate_tendencies!(sediment::InstantRemineralisation, bgc, grid, advection, tracers, timestepper)
+@kernel function _calculate_tendencies!(sediment::InstantRemineralisation, bgc, grid, advection, tracers, tendencies, sediment_tendencies)
     i, j = @index(Global, NTuple)
 
     k = bottom_index(i, j, sediment)
@@ -97,9 +101,9 @@ sediment_fields(model::InstantRemineralisation) = (N_storage = model.fields.N_st
     burial_efficiency = sediment.burial_efficiency_constant1 + sediment.burial_efficiency_constant2 * ((flux * 6.56) / (7 + flux * 6.56)) ^ 2
 
     # sediment evolution
-    @inbounds sediment.tendencies.Gⁿ.N_storage[i, j, 1] = burial_efficiency * flux
+    @inbounds sediment_tendencies.N_storage[i, j, 1] = burial_efficiency * flux
 
-    @inbounds timestepper.Gⁿ[remineralisation_receiver(bgc)][i, j, 1] += flux * (1 - burial_efficiency) / Δz
+    @inbounds tendencies[i, j, k] += flux * (1 - burial_efficiency) / Δz
 end
 
 summary(::InstantRemineralisation{FT}) where {FT} = string("Single-layer instant remineralisaiton ($FT)")