Cannot time step a non-hydrostatic model with default PISCES

[ali-ramadhan]

Huge push with #178!

I know PISCES is still in early development but I was curious to play around with it.

Is it not supposed to work out of the box with the default, e.g. just biogeochemistry = PISCES(; grid)? I'm guessing one of the root solvers does not like zero carbon or calcite concentrations.

MWE:

using Oceananigans
using OceanBioME

grid = RectilinearGrid(CPU(), size=(16, 16, 16), extent=(1, 1, 1))

model = NonhydrostaticModel(; grid, biogeochemistry=PISCES(; grid))

time_step!(model, 1)

Error:

┌ Warning: This implementation of PISCES is in early development and has not yet been validated against the operational version
└ @ OceanBioME.Models.PISCESModel ~/atdepth/OceanBioME.jl/src/Models/AdvectedPopulations/PISCES/PISCES.jl:346
ERROR: LoadError: ArgumentError: The interval [a,b] is not a bracketing interval.
You need f(a) and f(b) to have different signs (f(a) * f(b) < 0).
Consider a different bracket or try fzero(f, c) with an initial guess c.


Stacktrace:
  [1] assert_bracket
    @ ~/.julia/packages/Roots/KNVCY/src/Bracketing/bracketing.jl:52 [inlined]
  [2] init_state(::Roots.Bisection, F::Roots.Callable_Function{Val{…}, Val{…}, typeof(OceanBioME.Models.CarbonChemistryModel.alkalinity_residual), @NamedTuple{…}}, x₀::Float64, x₁::Float64, fx₀::Float64, fx₁::Float64; m::Float64, fm::Float64)
    @ Roots ~/.julia/packages/Roots/KNVCY/src/Bracketing/bisection.jl:50
  [3] init_state(::Roots.Bisection, F::Roots.Callable_Function{Val{…}, Val{…}, typeof(OceanBioME.Models.CarbonChemistryModel.alkalinity_residual), @NamedTuple{…}}, x₀::Float64, x₁::Float64, fx₀::Float64, fx₁::Float64)
    @ Roots ~/.julia/packages/Roots/KNVCY/src/Bracketing/bisection.jl:34
  [4] init_state(M::Roots.Bisection, F::Roots.Callable_Function{Val{…}, Val{…}, typeof(OceanBioME.Models.CarbonChemistryModel.alkalinity_residual), @NamedTuple{…}}, x::Tuple{Float64, Float64})
    @ Roots ~/.julia/packages/Roots/KNVCY/src/Bracketing/bracketing.jl:6
  [5] #init#43
    @ ~/.julia/packages/Roots/KNVCY/src/find_zero.jl:299 [inlined]
  [6] init
    @ ~/.julia/packages/Roots/KNVCY/src/find_zero.jl:289 [inlined]
  [7] solve(𝑭𝑿::Roots.ZeroProblem{…}, M::Roots.Bisection, p::@NamedTuple{…}; verbose::Bool, kwargs::@Kwargs{…})
    @ Roots ~/.julia/packages/Roots/KNVCY/src/find_zero.jl:491
  [8] solve
    @ ~/.julia/packages/Roots/KNVCY/src/find_zero.jl:484 [inlined]
  [9] #find_zero#40
    @ ~/.julia/packages/Roots/KNVCY/src/find_zero.jl:220 [inlined]
 [10] find_zero (repeats 2 times)
    @ ~/.julia/packages/Roots/KNVCY/src/find_zero.jl:210 [inlined]
 [11] solve_for_H
    @ ~/atdepth/OceanBioME.jl/src/Models/CarbonChemistry/carbon_chemistry.jl:186 [inlined]
 [12] carbonate_concentration(cc::CarbonChemistry{…}; DIC::Float64, T::Float64, S::Float64, Alk::Float64, pH::Nothing, P::Float64, lon::Int64, lat::Int64, boron::Float64, sulfate::Float64, fluoride::Float64, silicate::Float64, phosphate::Int64, upper_pH_bound::Int64, lower_pH_bound::Int64)
    @ OceanBioME.Models.CarbonChemistryModel ~/atdepth/OceanBioME.jl/src/Models/CarbonChemistry/calcite_concentration.jl:46
 [13] carbonate_concentration
    @ ~/atdepth/OceanBioME.jl/src/Models/CarbonChemistry/calcite_concentration.jl:1 [inlined]
 [14] calcite_saturation(cc::CarbonChemistry{…}; DIC::Float64, T::Float64, S::Float64, Alk::Float64, pH::Nothing, P::Float64, boron::Float64, sulfate::Float64, fluoride::Float64, calcium_ion_concentration::Float64, silicate::Float64, phosphate::Int64, upper_pH_bound::Int64, lower_pH_bound::Int64)
    @ OceanBioME.Models.CarbonChemistryModel ~/atdepth/OceanBioME.jl/src/Models/CarbonChemistry/calcite_concentration.jl:67
 [15] macro expansion
    @ ~/atdepth/OceanBioME.jl/src/Models/AdvectedPopulations/PISCES/compute_calcite_saturation.jl:34 [inlined]
 [16] cpu__compute_calcite_saturation!
    @ ~/.julia/packages/KernelAbstractions/xRXlv/src/macros.jl:291 [inlined]
 [17] cpu__compute_calcite_saturation!(__ctx__::KernelAbstractions.CompilerMetadata{…}, carbon_chemistry::CarbonChemistry{…}, calcite_saturation::Field{…}, grid::RectilinearGrid{…}, model_fields::@NamedTuple{…})
    @ OceanBioME.Models.PISCESModel ./none:0
 [18] __thread_run(tid::Int64, len::Int64, rem::Int64, obj::KernelAbstractions.Kernel{…}, ndrange::Nothing, iterspace::KernelAbstractions.NDIteration.NDRange{…}, args::Tuple{…}, dynamic::KernelAbstractions.NDIteration.NoDynamicCheck)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/xRXlv/src/cpu.jl:144
 [19] __run(obj::KernelAbstractions.Kernel{…}, ndrange::Nothing, iterspace::KernelAbstractions.NDIteration.NDRange{…}, args::Tuple{…}, dynamic::KernelAbstractions.NDIteration.NoDynamicCheck, static_threads::Bool)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/xRXlv/src/cpu.jl:111
 [20] (::KernelAbstractions.Kernel{…})(::CarbonChemistry{…}, ::Vararg{…}; ndrange::Nothing, workgroupsize::Nothing)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/xRXlv/src/cpu.jl:46
 [21] (::KernelAbstractions.Kernel{…})(::CarbonChemistry{…}, ::Vararg{…})
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/xRXlv/src/cpu.jl:39
 [22] launch!(::CPU, ::RectilinearGrid{…}, ::Symbol, ::typeof(OceanBioME.Models.PISCESModel._compute_calcite_saturation!), ::CarbonChemistry{…}, ::Vararg{…}; include_right_boundaries::Bool, reduced_dimensions::Tuple{}, location::Nothing, active_cells_map::Nothing, kwargs::@Kwargs{})
    @ Oceananigans.Utils ~/.julia/packages/Oceananigans/jUuNd/src/Utils/kernel_launching.jl:168
 [23] launch!
    @ ~/.julia/packages/Oceananigans/jUuNd/src/Utils/kernel_launching.jl:154 [inlined]
 [24] compute_calcite_saturation!
    @ ~/atdepth/OceanBioME.jl/src/Models/AdvectedPopulations/PISCES/compute_calcite_saturation.jl:14 [inlined]
 [25] update_biogeochemical_state!(model::NonhydrostaticModel{…}, bgc::PISCES{…})
    @ OceanBioME.Models.PISCESModel ~/atdepth/OceanBioME.jl/src/Models/AdvectedPopulations/PISCES/update_state.jl:13
 [26] update_biogeochemical_state!
    @ ~/atdepth/OceanBioME.jl/src/OceanBioME.jl:156 [inlined]
 [27] update_state!(model::NonhydrostaticModel{…}, callbacks::Vector{…}; compute_tendencies::Bool)
    @ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/jUuNd/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl:50
 [28] update_state!
    @ ~/.julia/packages/Oceananigans/jUuNd/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl:20 [inlined]
 [29] time_step!(model::NonhydrostaticModel{…}, Δt::Int64; callbacks::Vector{…})
    @ Oceananigans.TimeSteppers ~/.julia/packages/Oceananigans/jUuNd/src/TimeSteppers/runge_kutta_3.jl:114
 [30] time_step!(model::NonhydrostaticModel{…}, Δt::Int64)
    @ Oceananigans.TimeSteppers ~/.julia/packages/Oceananigans/jUuNd/src/TimeSteppers/runge_kutta_3.jl:81
 [31] top-level scope
    @ ~/atdepth/OceanBioME.jl/time_step_pisces.jl:8
 [32] include(fname::String)
    @ Base.MainInclude ./client.jl:489
 [33] top-level scope
    @ REPL[24]:1
in expression starting at /home/alir/atdepth/OceanBioME.jl/time_step_pisces.jl:8
Some type information was truncated. Use `show(err)` to see complete types.

[ali-ramadhan]

Borrowing some code from test_PISCES.jl I was able to time step it! Also works with GPU() 🎉

I'll leave this issue open in case it makes sense to make the default PISCES(; grid) work. But if not, then please feel free to close this issue.

---

using Oceananigans
using OceanBioME

using Oceananigans.Fields: ConstantField, FunctionField

grid = RectilinearGrid(CPU(), size=(16, 16, 16), extent=(100, 100, 100))

@inline total_light(x, y, z) = 3light(x, y, z)
@inline light(x, y, z) = ifelse(z <= 0, exp(z/10), 2-exp(-z/10))

PAR₁ = PAR₂ = PAR₃ = FunctionField{Center, Center, Center}(light, grid)
PAR  = FunctionField{Center, Center, Center}(total_light, grid)

mixed_layer_depth = ConstantField(-25)

light_attenuation = PrescribedPhotosyntheticallyActiveRadiation((; PAR, PAR₁, PAR₂, PAR₃))

biogeochemistry = PISCES(; grid, light_attenuation, mixed_layer_depth)

model = NonhydrostaticModel(; grid, biogeochemistry)

time_step!(model, 1)

[glwagner]

Should also change the API to PISCES(grid)

[jagoosw]

I came across this a few times too. I think this is from when it tries to solve for the calcite saturation which fails when DIC=Alk=0.

We've discussed it a little before but I'm not sure how best to handle that error (especially on GPU since it produces a really long confusing error).

We should make the default work too.

[jagoosw]

Also, regarding the warning about it being in development: we think that it's correct and should be able to do what the default version (i.e. not quota/simple/etc) of PISCES does, but I haven't been able to test that it exactly replicates the results of the existing implementations.

Our intention is to run a simple configuration of the NEMO version to compare to directly but I haven't had time (and struggled with the FORTRAN stuff when I did try before).

I'm going to make an issue for this so

[ali-ramadhan]

I came across this a few times too. I think this is from when it tries to solve for the calcite saturation which fails when DIC=Alk=0.

We've discussed it a little before but I'm not sure how best to handle that error (especially on GPU since it produces a really long confusing error).

We should make the default work too.

Yeah I can see it being tricky to handle... If it's just erroring when solving for calcite saturation, can the calculation be skipped with a if DIC == 0 && Alk == 0 clause?

[jagoosw]

I came across this a few times too. I think this is from when it tries to solve for the calcite saturation which fails when DIC=Alk=0.
We've discussed it a little before but I'm not sure how best to handle that error (especially on GPU since it produces a really long confusing error).
We should make the default work too.

Yeah I can see it being tricky to handle... If it's just erroring when solving for calcite saturation, can the calculation be skipped with a if DIC == 0 && Alk == 0 clause?

It could be, but my understanding is that this might affect performance on GPU. Since in almost all cases the outcome of the if would be the same on all threads the branching might not be a problem since they would all take the same time, but I don't have a good enough grasp on exactly how branching causes slowdown.

[glwagner]

There are a few ways, you can use clipping to ensure that you don't get a DomainError, but then mark the kernel as failed somehow at the same time, either by inserting a NaN or with something more specific:

# Eg store in `Field{Center, Center, Center}(grid, Bool)`
negative_found[i, j, k] = (DIC < 0) | (Alk < 0)

# clip and move on
DIC = max(zero(grid), DIC)
Alk = max(zero(grid), Alk)
# etc

Outside the kernel you can check negative_found and throw an error.

Or you can just set DIC and Alk to NaN.

Probably other solutions if you can be creative. Short-circuiting logic should not be necessary.