NP model

[iuryt]

Hi,

I am trying to setup a NP model with Oceananigans.
Once I have it ready I will announce it here. We can even have it as an example for those who would like to run biogeochemical models on Oceananigans.

The model has two nutrients and the following equations:

Based on the Convecting Plankton example, I know that I could setup most of the extra terms using Forcing.

But there are three parts that I would like to know your opinion.

1. Equation 2 is averaged over the mixed layer.
2. Equation 2 has a term that depends on the vertical cumulative integral of the chlorophyll that depends on P concentration
3. There is a sinking term for phytoplankton

Any idea on how to obtain and use the mixed-layer depth in a way that I could set up different density increment criteria?
What about the vertical cumulative integral of the chlorophyll that depends on P concentration?
How can I create a sinking term for P?

Status (tasks)

This implementation is currently in the following repository:
https://github.com/iuryt/Bioceananigans.jl

Based on the discussion below:

• Implement the biogeochemical relationships
• Implement the light profile function
• Average the light over the mixed layer
• Implement the shading of the phytoplankton
• Implement a time-varying light
• Implement the sinking velocity

[glwagner]

Any idea on how to obtain and use the mixed-layer depth in a way that I could set up different density increment criteria?

Perhaps you want something like

h = Field{Center, Center, Nothing}(grid) # defines a field that's reduced in the vertical direction
# ... much later ...

function compute_mixed_layer_depth!(sim)
    # compute h (referencing h as a global variable...)
    return nothing
end

simulation.callbacks[:compute_mld] = Callback(compute_mixed_layer_depth!) # computes mixed layer depth every time-step

You can also use model.auxiliary_fields, but this would require using multiple dispatch to define an appropriate compute! for your mixed layer depth field. Eminently doable, but perhaps a little more complicated.

What about the vertical cumulative integral of the the chlorophyll that depends on P concentration?

This should also be computed with a callback; the tricky part is that we might need trial and error to actually write out the computation. There is cumsum!:

help?> cumsum!
search: cumsum! cumsum

  cumsum!(B, A; dims::Integer)

  Cumulative sum of A along the dimension dims, storing the result in B. See also cumsum.

The good news is that this works on the GPU too (I think, I just did a quick test). But the bad news is that we actually need the sum in the reverse direction (ie from k=Nz down rather than k=1 up). As usual there are quite a few solutions to this, but I'm not sure the best one:

• Write our own kernel (this wouldn't be so different to what we do for the hydrostatic pressure:

  Oceananigans.jl/src/Models/NonhydrostaticModels/update_hydrostatic_pressure.jl

  Lines 9 to 17 in c74dfc6

  	@kernel function _update_hydrostatic_pressure!(pHY′, grid, buoyancy, C)
  	i, j = @index(Global, NTuple)
  	@inbounds pHY′[i, j, grid.Nz] = - ℑzᵃᵃᶠ(i, j, grid.Nz+1, grid, z_dot_g_b, buoyancy, C) * Δzᶜᶜᶠ(i, j, grid.Nz+1, grid)
  	@unroll for k in grid.Nz-1 : -1 : 1
  	@inbounds pHY′[i, j, k] = pHY′[i, j, k+1] - ℑzᵃᵃᶠ(i, j, k+1, grid, z_dot_g_b, buoyancy, C) * Δzᶜᶜᶠ(i, j, k+1, grid)
  	end
  	end

• Somehow figure out how to use cumsum! with Iterators.reverse (idea from https://discourse.julialang.org/t/cumsum-reversed/69355/3):

help?> Iterators.reverse
  Iterators.reverse(itr)

  Given an iterator itr, then reverse(itr) is an iterator over the same collection but in the reverse order.

  This iterator is "lazy" in that it does not make a copy of the collection in order to reverse it; see Base.reverse for an eager implementation.

  Not all iterator types T support reverse-order iteration. If T doesn't, then iterating over Iterators.reverse(itr::T) will throw a MethodError because of the missing iterate methods for Iterators.Reverse{T}.
  (To implement these methods, the original iterator itr::T can be obtained from r = Iterators.reverse(itr) by r.itr.)

  Examples
  ≡≡≡≡≡≡≡≡≡≡

  julia> foreach(println, Iterators.reverse(1:5))
  5
  4
  3
  2
  1

Finally note that a CumulativeIntegral might be appropriate for src implementation alongside Average and Integral:

https://github.com/CliMA/Oceananigans.jl/blob/main/src/AbstractOperations/metric_field_reductions.jl

However since clearly the "reversed" CumulativeIntegral is important, we'd need to figure out how to support that at the same time (ie Field(CumulativeIntegral(C, Down(), dims=3)) or something)

How can I create a sinking term for P?

Again a few ways:

1. Create a 3D array that represents this forcing term, then compute it in a Callback using, for example, AbstractOperations or
2. Use discrete_form=true and write out the flux divergence using Oceanangians operators. There are a few ways to do this; you can either pull the advective flux divergence operator from the source code, or you can compute this advective term with second-order differences:

using Oceananigans.Operators

# Implement a sinking term with second-order advection scheme
@inline w_Pᶜᶜᶠ(i, j, k, grid, w, P) = w * ℑzᵃᵃᶠ(i, j, k, grid, P)
@inline sinking(i, j, k, grid, w_sinking, P) = ∂zᶜᶜᶜ(i, j, k, grid, w_Pᶜᶜᶠ, w_sinking, P)

# Use Oceananigans advection operator (but then w_sinking must be field / array)
using Oceananigans.Advection: div_Uc
const scheme = WENO5()
@inline sinking(i, j, k, grid, w_sinking, P) = div_Uc(i, j, k, grid, scheme, (u=ZeroField(), v=ZeroField(), w=w_sinking), P)

The above I hard-coded WENO5() advection scheme, but any of them could be used. div_Uc is defined here:

Oceananigans.jl/src/Advection/tracer_advection_operators.jl

Line 27 in c74dfc6

@inline function div_Uc(i, j, k, grid, advection, U, c)

This is a pretty cool project! I think the model is complicated enough that it could warrant inclusion in the source. It'd be fun to have a Bioceananigans module too. But maybe we should see how the user script turns out first, then we can decide what the best course of action would be to add such things to the source code more permanently.

[glwagner]

I looked a little into cumsum! and my guess is that a cumulative interval field reduction will be most easily implemented if we write our own kernels, similar to how we calculate hydrostatic pressure.

If anyone is interested to learn more, I found that cumsum! calls Base.accumulate! which can accumulate any operator (like + for cumsum!):

https://github.com/JuliaLang/julia/blob/9d31f6a8d94bc63c5f39bdc7d2593616643490f4/base/accumulate.jl#L337

help?> accumulate!
search: accumulate! accumulate

  accumulate!(op, B, A; [dims], [init])

  Cumulative operation op on A along the dimension dims, storing the result in B. Providing dims is optional for vectors. If the keyword argument init is given, its value is used to instantiate the accumulation.
  See also accumulate.

  Examples
  ≡≡≡≡≡≡≡≡≡≡

  julia> x = [1, 0, 2, 0, 3];

  julia> y = [0, 0, 0, 0, 0];

  julia> accumulate!(+, y, x);

  julia> y
  5-element Vector{Int64}:
   1
   1
   3
   3
   6

  julia> A = [1 2; 3 4];

  julia> B = [0 0; 0 0];

  julia> accumulate!(-, B, A, dims=1);

  julia> B
  2×2 Matrix{Int64}:
    1   2
   -2  -2

  julia> accumulate!(-, B, A, dims=2);

  julia> B
  2×2 Matrix{Int64}:
   1  -1
   3  -1

accumulate! is also supported by CuArray. However, I am not sure that we can accumulate! an AbstractOperation due to the type restriction to AnyCuArray:

https://github.com/JuliaGPU/CUDA.jl/blob/e1fb0f47ac8ec4f08e042caae316ff6da417d8c3/src/accumulate.jl#L204

I don't know if that type restriction is necessary or perhaps can be relaxed.

[iuryt]

Thanks @glwagner
I am planning to have something to show by this week.

using Oceananigans.Advection: div_Uc
@inline sinking(i, j, k, grid, w_sinking, P) = div_Uc(i, j, k, grid, advection, (u=ZeroField(), v=ZeroField(), w=w_sinking), P)

Using the advection scheme from the model seems a better option for me. Now I have another question, how can I add multiple forcings for Phytoplankton? I mean, some will be in discrete form (as "sinking") and others will not be discrete, for instance:

@inline growing_and_grazing(x, y, z, t, P, params) = (params.μ₀ * exp(z / params.λ) - params.m) * P
plankton_dynamics_parameters = (μ₀ = 1/day,   # surface growth rate
                                 λ = 5,       # sunlight attenuation length scale (m)
                                 m = 0.1/day) # mortality rate due to virus and zooplankton grazing

plankton_dynamics = Forcing(growing_and_grazing, field_dependencies = :P,
                            parameters = plankton_dynamics_parameters)

While creating the model object, can I merge multiple forcings?

model = NonhydrostaticModel(grid = grid,
                            advection = advection,
                            closure=vertical_closure,
                            tracers = (:b, :P),
                            buoyancy = BuoyancyTracer(),
                            forcing = (P=plankton_dynamics,))

[glwagner]

At the moment you will have to assemble the total forcing manually. We could support auto-assembly (eg with a forcing tuple, similar to a closure tuple). I think it's particularly an issue for adding continuous form forcings together, since there's no easy way to do that from the user side.

For this application you should convert growing_and_grazing to the discrete_form. Note that the two are interchangeable provided you know that Pis located at(Center, Center, Center)`. Then

• x = xnode(Center, Center, Center, i, j, k, grid)
• y = ynode(Center, Center, Center, i, j, k, grid)
• z = znode(Center, Center, Center, i, j, k, grid)
• you can also write x, y, z = node(Center, Center, Center, i, j, k, grid)
• t = clock.time
• all prognostic fields are in the discrete argument fields, so P -> fields.P[i, j, k].

The *node functions come from Oceananigans.Grids, and the discrete_form also accepts parameters.

[iuryt]

Now I noticed you mentioned about *node functions here.

[iuryt]

Just to give some idea on what I am doing, up to now, what I am trying to do is something like this:

using Oceananigans
using Oceananigans.Units

# define the size and max depth of the simulation
const Nx = 2
const Ny = 2
const Nz = 48 # number of points in z
const H = 1000 # maximum depth

# create the grid of the model
grid = RectilinearGrid(GPU(),
    size=(Nx,Ny,Nz),
    x=(-Nx/2,Nx/2),
    y=(-Ny/2,Ny/2),
    z=(H * cos.(LinRange(π/2,0,Nz+1)) .- H)meters,
    topology=(Periodic, Periodic, Bounded)
)

# define the turbulence closure of the model
closure = ScalarDiffusivity(ν=1e-5, κ=1e-5)

# constants for the NP model
const μ₀ = 1/day   # surface growth rate
const m = 0.015/day # mortality rate due to virus and zooplankton grazing

const Kw = 0.059 # meter^-1
const Kc = 0.041 # m^2 mg^-1
const kn = 0.75
const kr = 0.5
const α = 0.0538/day

# create the mld field that will be updated at every timestep
h = Field{Center, Center, Nothing}(grid) 

#light = Field{Center, Center, Center}(grid)
#@inline light_growth(light) = (light*α)/sqrt(μ₀^2 + (light*α)^2)

const L0 = 100
# evolution of the available light at the surface
@inline L(z) = L0*exp.(z*Kw)
# light profile
@inline light_growth(z,t) = μ₀ * (L(z)*α)/sqrt(μ₀^2 + (L(z)*α)^2)

# nitrate and ammonium limiting
@inline N_lim(N,Nr) = (N/(N+kn)) * (kr/(Nr+kr))
@inline Nr_lim(Nr) =  (Nr/(Nr+kr))

# functions for the NP model
P_forcing(x, y, z, t, P, N, Nr)  =   light_growth(z,t) * (N_lim(N,Nr) + Nr_lim(Nr)) * P - m * P^2
N_forcing(x, y, z, t, P, N, Nr)  = - light_growth(z,t) * N_lim(N,Nr) * P
Nr_forcing(x, y, z, t, P, N, Nr) = - light_growth(z,t) * Nr_lim(Nr) * P + m * P^2

# using the functions to determine the forcing
P_dynamics = Forcing(P_forcing, field_dependencies = (:P,:N,:Nr))
N_dynamics = Forcing(N_forcing, field_dependencies = (:P,:N,:Nr))
Nr_dynamics = Forcing(Nr_forcing, field_dependencies = (:P,:N,:Nr))

# create the model
model = NonhydrostaticModel(grid = grid,
                            closure=closure,
                            tracers = (:b, :P, :N, :Nr),
                            buoyancy = BuoyancyTracer(),
                            forcing = (P=P_dynamics,N=N_dynamics,Nr=Nr_dynamics))

# mld
const cz = -200
# gravity
const g = 9.82
# reference density
const ρₒ = 1026

# initial buoyancy profile based on Argo data
@inline B(x, y, z) = -(g/ρₒ)*0.25*tanh(0.0027*(-653.3-z))-6.8*z/1e5+1027.56
# initial phytoplankton profile
@inline P(x, y, z) = ifelse(z>cz,0.4,0)

# setting the initial conditions
set!(model;b=B,P=P,N=13,Nr=0)

# create a simulation
simulation = Simulation(model, Δt = 1minutes, stop_time = 20days)

# function to compute the mld
function compute_mixed_layer_depth!(sim)
    # get the buoyancy from the simulation
    b = sim.model.tracers.b
    # extract the model nodes
    x,y,z = nodes((Center,Center,Center),sim.model.grid)
    
    # loop by x-axis
    for i=1:length(x)
        # loop by y axis
        for j=1:length(y)
            # get the density profile at (i,j)
            ρⁿ = -(ρₒ/g)*Array(interior(b, i, j, :))
            # surface density
            ρs = ρⁿ[end]
            # criterion for density increment
            dρ = 0.01
            # loop backwards (surface->bottom)
            for k=0:length(z)-1
                # if the density satisfies the criterium
                if ρⁿ[end-k]>=ρs+dρ
                    # update mld values
                    h[i,j]=z[end-k]
                    # break the loop
                    break
                end
            end
        end
    end
    return nothing
end

# add the function to the callbacks of the simulation
simulation.callbacks[:compute_mld] = Callback(compute_mixed_layer_depth!)

# function to update the light profile (not working)
function compute_light_profile!(sim)
    x,y,z = nodes((Center,Center,Center),sim.model.grid)
    
    dz = z[2:end]-z[1:end-1]
    dz = [dz[1];dz]
    
    for i=1:length(x)
        for j=1:length(y)
            Li = L0(time(sim))*exp.(z*Kw)
            
            inmld = z.>h[i,j]
            Li[inmld] .= sum(Li[inmld].*dz[inmld])/abs(h[i,j])
            
            for k=1:length(z)
                light[i,j,k] = Li[k]
            end
        end
    end
    return nothing
end
# simulation.callbacks[:compute_light] = Callback(compute_light_profile!)

# writing the output
simulation.output_writers[:fields] =
    NetCDFOutputWriter(model, merge(model.velocities, model.tracers), filepath = "data/NP_output.nc",
                     schedule=TimeInterval(1hour))


using Printf

function print_progress(simulation)
    b, P, N, Nr = simulation.model.tracers

    # Print a progress message
    msg = @sprintf("i: %04d, t: %s, Δt: %s, P_max = %.1e, N_min = %.1e, Nr_max = %.1e, wall time: %s\n",
                   iteration(simulation),
                   prettytime(time(simulation)),
                   prettytime(simulation.Δt),
                   maximum(P), minimum(N), maximum(Nr),
                   prettytime(simulation.run_wall_time))

    @info msg

    return nothing
end

simulation.callbacks[:progress] = Callback(print_progress, TimeInterval(1hour))

# run the simulation
run!(simulation)

But this is not yet working for the average light and I am not using the mld yet. I don't know what to do with the halo regions and I am pretty sure there is a better way than compute extract and compute dz at every call. Any suggestions?

I created a public repo for that https://github.com/iuryt/bioceananigans.jl

Also, I am defining light and mld as global fields... what if I want to save them with the output data?

[glwagner]

I edited the code in a way that I hope makes it work but it needs to be tested. Please report back with the final debugged version!

[glwagner]

I also got confused because in you case 1 is at the top, but sometimes the first index is at the bottom on Oceananigans..

If you're referring to

@inbounds h[i, j, 1] = - z_below_mixed_layer_ij

then note that h is a two-dimensional field and only has one vertical index, k=1. The loop is written to start at the top:

@unroll for k in grid.Nz - 1 : -1 : 1 # scroll to point just above bottom

and decrease k until it hits the bottom. There very well could be another error or bug so let me know if there is a specific piece of the code that you think has an issue!

[glwagner]

I couldn't find a reference for this znode method and znodes doesn't have that many args. Should I define it separately?

znode is defined here:

Oceananigans.jl/src/Grids/grid_utils.jl

Line 227 in b7485ec

@inline znode(LX, LY, LZ, i, j, k, grid) = znode(LZ, k, grid)

znodes is different; it returns an array. znode extracts the znode at index k.

The fallback is valid for on a rectilinear grid where z only varies with k. For znode it's unlikely you would care about 3D curvilinear grids any time soon. For xnode and ynode however it's recommended to specify all three locations, so that your code is also valid on the sphere.

[iuryt]

Thanks again for the patience to explain that to me. You don't have to edit the code for that. I can do by myself. I will debug it with a case I know the answer and let you know. Once this is done, I will use it as an example for the light.

[glwagner]

Could be fun to add a simple version of this model as a validation case in #2389

[iuryt]

@glwagner and others

Just to update, I worked on the code for obtaining the mixed layer depth fixing some bugs and this is also now interpolating linearly.

Sorry for the lack of information in the plots. The color is N2.

Before interpolation:

After interpolation:

If you want to check how the code is right now, click here.

[francispoulin]

Much smoother and looks nicer

[iuryt]

Hey guys,

I have just finished computing the average light in the mixed layer and I am stuck in a problem now.
I am trying to use the light field that is updated at every timestep in the forcing functions and this is not working.
Any idea how to make it work?

Code:

# create the mld field that will be updated at every timestep
h = Field{Center, Center, Nothing}(grid) 
light = Field{Center, Center, Center}(grid)

const L0 = 100
# evolution of the available light at the surface
@inline L(z) = L0*exp.(z*Kw)
# light profile
@inline light_growth(z) = μ₀ * (L(z)*α)/sqrt(μ₀^2 + (L(z)*α)^2)

# nitrate and ammonium limiting
@inline N_lim(N,Nr) = (N/(N+kn)) * (kr/(Nr+kr))
@inline Nr_lim(Nr) =  (Nr/(Nr+kr))

# functions for the NP model
P_forcing(x, y, z, t, P, N, Nr)  =   light * (N_lim(N,Nr) + Nr_lim(Nr)) * P - m * P^2
N_forcing(x, y, z, t, P, N, Nr)  = - light * N_lim(N,Nr) * P
Nr_forcing(x, y, z, t, P, N, Nr) = - light * Nr_lim(Nr) * P + m * P^2

# using the functions to determine the forcing
P_dynamics = Forcing(P_forcing, field_dependencies = (:P,:N,:Nr))
N_dynamics = Forcing(N_forcing, field_dependencies = (:P,:N,:Nr))
Nr_dynamics = Forcing(Nr_forcing, field_dependencies = (:P,:N,:Nr))

Full code

Error:

InvalidIRError: compiling kernel gpu_calculate_Gc!(Cassette.Context{nametype(CUDACtx), Nothing, Nothing, KernelAbstractions.var"##PassType#257", Nothing, Cassette.DisableHooks}, typeof(Oceananigans.Models.NonhydrostaticModels.gpu_calculate_Gc!), KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.StaticSize{(1, 100, 48)}, KernelAbstractions.NDIteration.DynamicCheck, Nothing, Nothing, KernelAbstractions.NDIteration.NDRange{3, KernelAbstractions.NDIteration.StaticSize{(1, 1, 48)}, KernelAbstractions.NDIteration.StaticSize{(1, 100, 1)}, Nothing, Nothing}}, OffsetArrays.OffsetArray{Float64, 3, CUDA.CuDeviceArray{Float64, 3, 1}}, RectilinearGrid{Float64, Flat, Periodic, Bounded, Float64, Float64, OffsetArrays.OffsetVector{Float64, CUDA.CuDeviceVector{Float64, 1}}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}}, OffsetArrays.OffsetVector{Float64, CUDA.CuDeviceVector{Float64, 1}}, Nothing}, Val{2}, CenteredSecondOrder, Tuple{ScalarDiffusivity{Oceananigans.TurbulenceClosures.ExplicitTimeDiscretization, Oceananigans.TurbulenceClosures.ThreeDimensionalFormulation, Float64, NamedTuple{(:b, :P, :N, :Nr), NTuple{4, Float64}}}, ScalarDiffusivity{Oceananigans.TurbulenceClosures.ExplicitTimeDiscretization, Oceananigans.TurbulenceClosures.ThreeDimensionalFormulation, Float64, NamedTuple{(:b, :P, :N, :Nr), NTuple{4, Float64}}}}, Buoyancy{BuoyancyTracer, Oceananigans.Grids.ZDirection}, NamedTuple{(:velocities, :tracers), Tuple{NamedTuple{(:u, :v, :w), Tuple{Oceananigans.Fields.ZeroField{Int64, 3}, Oceananigans.Fields.ZeroField{Int64, 3}, Oceananigans.Fields.ZeroField{Int64, 3}}}, NamedTuple{(:b, :P, :N, :Nr), NTuple{4, Oceananigans.Fields.ZeroField{Int64, 3}}}}}, NamedTuple{(:u, :v, :w), Tuple{OffsetArrays.OffsetArray{Float64, 3, CUDA.CuDeviceArray{Float64, 3, 1}}, OffsetArrays.OffsetArray{Float64, 3, CUDA.CuDeviceArray{Float64, 3, 1}}, OffsetArrays.OffsetArray{Float64, 3, CUDA.CuDeviceArray{Float64, 3, 1}}}}, NamedTuple{(:b, :P, :N, :Nr), NTuple{4, OffsetArrays.OffsetArray{Float64, 3, CUDA.CuDeviceArray{Float64, 3, 1}}}}, Tuple{Nothing, Nothing}, Oceananigans.Forcings.ContinuousForcing{Center, Center, Center, Nothing, typeof(P_forcing), Nothing, Tuple{Nothing, Int64, Int64, Int64}, Tuple{typeof(Oceananigans.Operators.identity1), typeof(Oceananigans.Operators.identity2), typeof(Oceananigans.Operators.identity3), typeof(Oceananigans.Operators.identity4)}}, NamedTuple{(:time, :iteration, :stage), Tuple{Float64, Int64, Int64}}) resulted in invalid LLVM IR
Reason: unsupported dynamic function invocation (call to field_arguments)
Stacktrace:
  [1] call
    @ ~/.julia/packages/Cassette/34vIw/src/context.jl:456
  [2] fallback
    @ ~/.julia/packages/Cassette/34vIw/src/context.jl:454
  [3] _overdub_fallback(::Any, ::Vararg{Any, N} where N)
    @ ~/.julia/packages/Cassette/34vIw/src/overdub.jl:586
  [4] overdub
    @ ~/.julia/packages/Cassette/34vIw/src/overdub.jl:586
  [5] overdub
    @ ~/.julia/packages/Oceananigans/ynvn3/src/Utils/user_function_arguments.jl:22
  [6] (::Oceananigans.Forcings.ContinuousForcing{Center, Center, Center, Nothing, typeof(P_forcing), Nothing, Tuple{Nothing, Int64, Int64, Int64}, Tuple{typeof(Oceananigans.Operators.identity1), typeof(Oceananigans.Operators.identity2), typeof(Oceananigans.Operators.identity3), typeof(Oceananigans.Operators.identity4)}})(::Int64, ::Int64, ::Int64, ::RectilinearGrid{Float64, Flat, Periodic, Bounded, Float64, Float64, OffsetArrays.OffsetVector{Float64, CUDA.CuDeviceVector{Float64, 1}}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}}, OffsetArrays.OffsetVector{Float64, CUDA.CuDeviceVector{Float64, 1}}, Nothing}, ::NamedTuple{(:time, :iteration, :stage), Tuple{Float64, Int64, Int64}}, ::NamedTuple{(:u, :v, :w, :b, :P, :N, :Nr), NTuple{7, OffsetArrays.OffsetArray{Float64, 3, CUDA.CuDeviceArray{Float64, 3, 1}}}})
    @ ~/.julia/packages/Oceananigans/ynvn3/src/Forcings/continuous_forcing.jl:118
  [7] overdub
    @ ~/.julia/packages/Oceananigans/ynvn3/src/Forcings/continuous_forcing.jl:118
  [8] overdub
    @ ~/.julia/packages/Oceananigans/ynvn3/src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl:211
  [9] macro expansion
    @ ~/.julia/packages/Oceananigans/ynvn3/src/Models/NonhydrostaticModels/calculate_nonhydrostatic_tendencies.jl:137
 [10] overdub
    @ ~/.julia/packages/KernelAbstractions/3ZHln/src/macros.jl:80
 [11] overdub
    @ ~/.julia/packages/Cassette/34vIw/src/overdub.jl:0

[glwagner]

light is a field (referenced as a global variable in the forcing functions as if it were a constant). I also can't see where m is defined. But here you will need the discrete_form because light is not a prognostic variable and thus not available via field_dependencies (at least right now...). You may need

# functions for the NP model
@inline P_forcing(i, j, k, grid, clock, fields, p)  = @inbounds   p.light[i, j, k] * (N_lim(fields.N[i, j, k], fields.Nr[i, j, k]) + Nr_lim(fields.Nr[i, j, k])) * fields.P[i, j, k] - p.m * P^2
@inline N_forcing(i, j, k, grid, clock, fields, p)  = @inbounds - p.light[i, j, k] * N_lim(fields.N[i, j, k], fields.Nr[i, j, k]) * fields.P
@inline Nr_forcing(i, j, k, grid, clock, fields, p) = @inbounds - p.light[i, j, k] * Nr_lim(fields.Nr[i, j, k]) * fields.P[i, j, k] + p.m * fields.P[i, j, k]^2

# using the functions to determine the forcing
P_dynamics = Forcing(P_forcing, discrete_form=true, parameters=(; m, light))
N_dynamics = Forcing(N_forcing, discrete_form=true, parameters=(; m, light))
Nr_dynamics = Forcing(Nr_forcing, discrete_form=true, parameters=(; m, light))

From my brain:

1. Test on the CPU before running on GPU to find "ordinary bugs". You'll get much better error messages on the CPU. Then if there's still a problem on the GPU, there's a relatively narrower range of possible issues.
2. It might make sense to define intermediate functions that take in the fields evaluated at i, j, k and parameters like m, eg, functions that look like what you had without the arguments x, y, z, t.
3. Always put spaces after commas to show love to people reading your code ❤️

[iuryt]

Sorry. I had put the link for the full code, but should have had put the full code here directly.
I will follow your suggestions. Thanks!
And thanks for the programming style advice as well, haha

[iuryt]

It worked!

That's what I am doing, based on your suggestion

# constants for the NP model
const μ₀ = 1/day   # surface growth rate
const m = 0.015/day # mortality rate due to virus and zooplankton grazing
const Kw = 0.059 # meter^-1
const Kc = 0.041 # m^2 mg^-1
const kn = 0.75
const kr = 0.5
const α = 0.0538/day
const L0 = 100


# create the mld field that will be updated at every timestep
h = Field{Center, Center, Nothing}(grid) 
light = Field{Center, Center, Center}(grid)

# evolution of the available light at the surface
@inline L(z) = L0*exp.(z*Kw)
# light profile
@inline light_growth(z) = μ₀ * (L(z)*α)/sqrt(μ₀^2 + (L(z)*α)^2)

# nitrate and ammonium limiting
@inline N_lim(N, Nr) = (N/(N+kn)) * (kr/(Nr+kr))
@inline Nr_lim(Nr) =  (Nr/(Nr+kr))

# functions for the NP model
@inline P_forcing(light, P, N, Nr)  =   light * (N_lim(N, Nr) + Nr_lim(Nr)) * P - m * P^2
@inline N_forcing(light, P, N, Nr)  = - light * N_lim(N, Nr) * P
@inline Nr_forcing(light, P, N, Nr) = - light * Nr_lim(Nr) * P + m * P^2

# functions for the NP model
@inline P_forcing(i, j, k, grid, clock, fields, p)  = @inbounds P_forcing(p.light[i, j, k], fields.P[i, j, k], fields.N[i, j, k], fields.Nr[i, j, k])
@inline N_forcing(i, j, k, grid, clock, fields, p)  = @inbounds N_forcing(p.light[i, j, k], fields.P[i, j, k], fields.N[i, j, k], fields.Nr[i, j, k])
@inline Nr_forcing(i, j, k, grid, clock, fields, p) = @inbounds Nr_forcing(p.light[i, j, k], fields.P[i, j, k], fields.N[i, j, k], fields.Nr[i, j, k])

# using the functions to determine the forcing
P_dynamics = Forcing(P_forcing, discrete_form=true, parameters=(; light))
N_dynamics = Forcing(N_forcing, discrete_form=true, parameters=(; light))
Nr_dynamics = Forcing(Nr_forcing, discrete_form=true, parameters=(; light))

[glwagner]

Perfect! Nicely done!

[iuryt]

@glwagner

I am now working on the vertical integration of chlorophyll. Based on the hydrostatic pressure code,

Oceananigans.jl/src/Models/NonhydrostaticModels/update_hydrostatic_pressure.jl

Lines 9 to 17 in c74dfc6

	@kernel function _update_hydrostatic_pressure!(pHY′, grid, buoyancy, C)
	i, j = @index(Global, NTuple)
	@inbounds pHY′[i, j, grid.Nz] = - ℑzᵃᵃᶠ(i, j, grid.Nz+1, grid, z_dot_g_b, buoyancy, C) * Δzᶜᶜᶠ(i, j, grid.Nz+1, grid)
	@unroll for k in grid.Nz-1 : -1 : 1
	@inbounds pHY′[i, j, k] = pHY′[i, j, k+1] - ℑzᵃᵃᶠ(i, j, k+1, grid, z_dot_g_b, buoyancy, C) * Δzᶜᶜᶠ(i, j, k+1, grid)
	end
	end

I have some questions:

• You are interpolating to face because this is used for w? If so, in my case I could simply

 @inbounds P_integ[i, j, grid.Nz] = P[i, j, grid.Nz] * Δzᶜᶜᶜ(i, j, grid.Nz, grid)

• Why are you using Δzᶜᶜᶠ instead of Δzᶜᶜᶜ, as tracers, phytoplankton and buoyancy are centered, right?

 @unroll for k in grid.Nz-1 : -1 : 1 
     @inbounds pHY′[i, j, k] = pHY′[i, j, k+1] - ℑzᵃᵃᶠ(i, j, k+1, grid, z_dot_g_b, buoyancy, C) * Δzᶜᶜᶠ(i, j, k+1, grid) 
 end 

Why rolling from grid.Nz-1 to 1 if you sum 1 to k? Wouldn't be better to simply

 @unroll for k in grid.Nz : -1 : 2 
     @inbounds pHY′[i, j, k] = pHY′[i, j, k] - ℑzᵃᵃᶠ(i, j, k+1, grid, z_dot_g_b, buoyancy, C) * Δzᶜᶜᶠ(i, j, k, grid) 
 end 

[glwagner]

This integral is derived from the hydrostatic pressure gradient

(p[k] - p[k-1]) / Δzᶜᶜᶠ[k]

which is located at cell interfaces in the vertical. Another way to think about this is that the expression arises in the vertical momentum equation, which is located at cell interfaces in the vertical.

as tracers, phytoplankton and buoyancy are centered, right?

Yes. That's why we have to interpolate buoyancy to cell interfaces in the vertical with:

ℑzᵃᵃᶠ(i, j, k+1, grid, z_dot_g_b, buoyancy, C)

Why rolling from grid.Nz-1 to 1 if you sum 1 to k? Wouldn't be better to simply

How do we determine pHY′ at k=1 then?

[iuryt]

Got it!

[iuryt]

How to organize different options for the package

I would like to ask something to everyone who is watching this discussion.

My goal is to make a package based on Oceanostics.jl for Bioceananigans.jl. For now, this is going to basically have a kernel for computing the mixed-layer depth and different approaches for the light-limiting growth:

1. Average the light over the mixed layer
2. Average the light-limiting growth over the mixed layer
3. Don't average the light over the mixed layer

We could also add an option for using or not the negative feedback that phytoplankton can cause on the light for deeper levels (aka phytoplankton shading). The number of phyto/zooplankton species, types of nutrients and the equations are provided by the user and we could give some examples to help with the configuration for different models.

The problem

With the on and off options for shading and the three approaches for the light-limiting growth, we have at least 6 options for the users.
I need help to think about the best way to handle that.

My solution (for now)

Give two key arguments with the following options:

1. average
   • nothing (default): do not average light or light-limiting growth.
   • light: average the light over the mixed layer
   • growth: average the light-limiting growth over the mixed layer
2. shading
   • false (default): do not compute and use the phytoplankton shading
   • true: compute and use the phytoplankton shading

I was trying to avoid having string options for key arguments in the kernels, but that was the only solution I found for this without having to duplicate code. If this is a good solution, we will need to check the selected options somehow before calling the function, right?