Merge pull request #134 from OceanBioME/jsw/river-input

(0.6.0) Refactor models for more abstraction (and bumped compats)

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.5.0"
+version = "0.6.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -12,12 +12,12 @@ Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 
 [compat]
-Adapt = "3.4"
-JLD2 = "0.4.31"
+Adapt = "3"
+JLD2 = "0.4"
 KernelAbstractions = "0.9"
-Oceananigans = "0.84.1, 0.85"
+Oceananigans = "0.84.1, 0.85, 0.86, 0.87"
 Roots = "2"
-StructArrays = "0.6.15"
+StructArrays = "0.4, 0.5, 0.6"
 julia = "1.9"
 
 [extras]

docs/Project.toml

@@ -17,4 +17,4 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 CairoMakie = "0.10"
 DocumenterCitations = "1"
 Literate = "≥2.9.0"
-Oceananigans = "0.84.1, 0.85"
+Oceananigans = "0.84.1, 0.85, 0.86, 0.87"

docs/src/model_components/biogeochemical/LOBSTER.md

@@ -13,15 +13,16 @@ julia> grid = RectilinearGrid(size=(3, 3, 30), extent=(10, 10, 200));
 
 julia> bgc_model = LOBSTER(; grid, carbonates = true)
 Lodyc-DAMTP Ocean Biogeochemical Simulation Tools for Ecosystem and Resources (LOBSTER) model (Float64) 
- Light Attenuation Model: 
-    └── Two-band light attenuation model (Float64)
  Optional components:
     ├── Carbonates ✅ 
     ├── Oxygen ❌ 
     └── Variable Redfield Ratio ❌
  Sinking Velocities:
     ├── sPOM: 0.0 to -3.47e-5 m/s 
-    └── bPOM: 0.0 to -0.0023148148148148147 m/s
+    └── bPOM: 0.0 to -0.0023148148148148147 m/s 
+ Light attenuation: Two-band light attenuation model (Float64)
+ Sediment: Nothing
+ Particles: Nothing
 ```
 
 ## Model equations

docs/src/model_components/biogeochemical/index.md

@@ -5,7 +5,7 @@ Biogeochemical (BGC) models can be used within the [Oceananigans biogeochemistry
 For details of the BGC models currently implemented please see the following pages.
 
 ## Oceananigans setup
-At the simplest level all that is required to setup a BGC model is to pass it to the Oceananigans model setup:
+At the simplest level all that is required to setup an existing OceanBioME BGC model is to pass it to the Oceananigans model setup:
 ```julia
 model = NonhydrostaticModel(; grid,
                               ...,
@@ -19,3 +19,7 @@ MODEL_NAME(; grid, growth_rate = 10.0)
 This will set up the required tracers and auxiliary fields, and you may also set boundary conditions or additional forcing through the normal Oceananigans setup. 
 
 Models usually have a default [light attenuation model](@ref light) specified, these may be substituted easily by passing different models as parameters as above.
+
+Our models are implemented in an abstract framework `Biogeochemistry` which contains `underlying_biogeochemistry`, `light_attenuation`, `sediment`, and `modifiers`.
+This is automatically set up for existing BGC models, but may also be used to couple any BGC model with light attenuation and sediments.
+See the [implementation page](@ref model_implementation) for some more information on how to couple other models.

docs/src/model_components/individuals/index.md

@@ -43,7 +43,7 @@ function update_lagrangian_particle_properties!(particles::GrowingParticles, mod
     return nothing
 end
 
-nothing # hide
+nothing #hide
 ```
 
 In this example the particles will not move around, and are only integrated on a single thread. For a more comprehensive example see the [Sugar Kelp](@ref SLatissima) implementation. We then need to update the tracer tendencies to match the nutrients' uptake:
@@ -67,7 +67,7 @@ function update_tendencies!(bgc, particles::GrowingParticles, model)
     return nothing
 end
 
-nothing # hide
+nothing #hide
 ```
 
 Now we can just plug this into any biogeochemical model setup to have particles (currently [NPZD](@ref NPZD) and [LOBSTER](@ref LOBSTER)):

docs/src/model_components/utils.md

@@ -22,20 +22,17 @@ simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))
 ## Negative tracer detection
 As a temporary measure we have implemented a callback to either detect negative tracers and either scale a conserved group, force them back to zero, or throw an error. Please see the numerical implementations' page for details. This can be set up by:
 ```julia
-negativity_protection! = ScaleNegativeTracers(tracers = (:P, :Z, :N))
-simulation.callbacks[:neg] = Callback(negativity_protection!; callsite = UpdateStateCallsite())
+negativity_protection = ScaleNegativeTracers((:P, :Z, :N))
+biogeochemistry = Biogeochemistry(...; modifiers = negativity_protection)
 ```
 You may also pass a scale factor for each component (e.g. in case they have different redfield ratios):
 ```julia
-negativity_protection! = ScaleNegativeTracers(tracers = (:P, :Z, :D), scalefactors = (P = 1, Z = 1, D = 2))
-simulation.callbacks[:neg] = Callback(negativity_protection!; callsite = UpdateStateCallsite())
+negativity_protection = ScaleNegativeTracers((:P, :Z, :N); scalefactors = (1, 1, 2))
+biogeochemistry = Biogeochemistry(...; modifiers = negativity_protection)
 ```
 Here you should carefully consider which tracers form a conserved group (if at all). Alternatively, force to zero by:
 ```julia
-simulation.callbacks[:neg] = Callback(OceanBioME.no_negative_tracers!, callsite = UpdateStateCallsite())
+negativity_protection = ZeroNegativeTracers()
+biogeochemistry = Biogeochemistry(...; modifiers = negativity_protection)
 ```
-or throw an error:
-```julia
-simulation.callbacks[:neg] = Callback(OceanBioME.error_on_neg!, callsite = UpdateStateCallsite())
-```
-The latter two both optionally take a named tuple of parameters which may include `exclude` which can be a tuple of tracer names (Symbols) which are allowed to be negative.
+The latter optionally takes a named tuple of parameters that may include `exclude`, which can be a tuple of tracer names (Symbols) which are allowed to be negative.

docs/src/model_implementation.md

@@ -18,20 +18,19 @@ For this example we are going to implement the simple Nutrient-Phytoplankton mod
 The first step is to import the abstract type from OceanBioME, some units from Oceananigans (for ease of parameter definition), and [`import`](https://stackoverflow.com/questions/27086159/what-is-the-difference-between-using-and-import-in-julia-when-building-a-mod) some functions from Oceananigans in order to add methods to:
 
 ```@example implementing
-using OceanBioME: ContinuousFormBiogeochemistry
+using OceanBioME: Biogeochemistry
+using Oceananigans.Biogeochemistry: AbstractContinuousFormBiogeochemistry
 using Oceananigans.Units
 
 import Oceananigans.Biogeochemistry: required_biogeochemical_tracers,
                                      required_biogeochemical_auxiliary_fields,
-                                     update_biogeochemical_state!,
-                                     biogeochemical_drift_velocity,
-                                     biogeochemical_auxiliary_fields
+                                     biogeochemical_drift_velocity
 ```
 
 We then define our `struct` with the model parameters, as well as slots for the particles, light attenuation, and sediment models:
 
 ```@example implementing
-@kwdef struct NutrientPhytoplankton{FT, LA, S, W, P} <:ContinuousFormBiogeochemistry{LA, S, P}
+@kwdef struct NutrientPhytoplankton{FT, W} <: AbstractContinuousFormBiogeochemistry
             base_growth_rate :: FT = 1.27 / day              # 1 / seconds
     nutrient_half_saturation :: FT = 0.025 * 1000 / 14       # mmol N / m³
        light_half_saturation :: FT = 300.0                   # micro einstein / m² / s
@@ -41,10 +40,7 @@ We then define our `struct` with the model parameters, as well as slots for the
               mortality_rate :: FT = 0.15 / day              # 1 / seconds
      crowding_mortality_rate :: FT = 0.004 / day / 1000 * 14 # 1 / seconds / mmol N / m³
 
-     light_attenuation_model :: LA = nothing
-              sediment_model :: S  = nothing
             sinking_velocity :: W  = 2 / day
-                   particles :: P  = nothing
 end
 ```
 
@@ -54,8 +50,6 @@ Here, we use descriptive names for the parameters. Below, each of these paramete
 required_biogeochemical_tracers(::NutrientPhytoplankton) = (:N, :P, :T)
 
 required_biogeochemical_auxiliary_fields(::NutrientPhytoplankton) = (:PAR, )
-
-biogeochemical_auxiliary_fields(bgc::NutrientPhytoplankton) = biogeochemical_auxiliary_fields(bgc.light_attenuation_model)
 ```
 
 Next, we define the functions that specify how the phytoplankton ``P`` evolve. In the absence of advection and diffusion (both of which are handled by Oceananigans), we want the phytoplankton to evolve at the rate given by:
@@ -126,10 +120,6 @@ Finally, we need to define some functions to allow us to update the time-depende
 using OceanBioME: BoxModel
 import OceanBioME.BoxModels: update_boxmodel_state!
 
-function update_biogeochemical_state!(bgc::NutrientPhytoplankton, model)
-    update_PAR!(model, bgc.light_attenuation_model)
-end
-
 function update_boxmodel_state!(model::BoxModel{<:NutrientPhytoplankton, <:Any, <:Any, <:Any, <:Any, <:Any})
     getproperty(model.values, :PAR) .= model.forcing.PAR(model.clock.time)
     getproperty(model.values, :T) .= model.forcing.T(model.clock.time)
@@ -228,7 +218,7 @@ import OceanBioME.Boundaries.Sediments: nitrogen_flux, carbon_flux, remineralisa
 @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(bgc::NutrientPhytoplankton, tracers, i, j) = nitrogen_flux(i, j, k, grid, advection, bgc, tracers) * 6.56
+@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
 ```
@@ -254,17 +244,22 @@ grid = RectilinearGrid(topology = (Flat, Flat, Bounded), size = (32, ), extent =
 
 # setup the biogeochemical model
 
-light_attenuation_model = TwoBandPhotosyntheticallyActiveRadiation(; grid, surface_PAR)
+light_attenuation = TwoBandPhotosyntheticallyActiveRadiation(; grid, surface_PAR)
 
-sediment_model = InstantRemineralisation(; grid)
+sediment = InstantRemineralisation(; grid)
 
 sinking_velocity = ZFaceField(grid)
 
 w_sink(x, y, z) = 2 / day * tanh(z / 5)
 
 set!(sinking_velocity, w_sink)
 
-biogeochemistry = NutrientPhytoplankton(; light_attenuation_model, sinking_velocity, sediment_model) 
+negative_tracer_scaling = ScaleNegativeTracers((:N, :P))
+
+biogeochemistry = Biogeochemistry(NutrientPhytoplankton(; sinking_velocity);
+                                  light_attenuation,
+                                  sediment,
+                                  modifiers = negative_tracer_scaling) 
 
 # put the model together
 
@@ -284,15 +279,12 @@ simulation.output_writers[:tracers] = JLD2OutputWriter(model, model.tracers,
                                                        schedule = TimeInterval(1day),
                                                        overwrite_existing = true)
 
-simulation.output_writers[:sediment] = JLD2OutputWriter(model, model.biogeochemistry.sediment_model.fields,
+simulation.output_writers[:sediment] = JLD2OutputWriter(model, model.biogeochemistry.sediment.fields,
                                                         indices = (:, :, 1),
                                                         filename = "column_np_sediment.jld2",
                                                         schedule = TimeInterval(1day),
                                                         overwrite_existing = true)
 
-scale_negative_tracers = ScaleNegativeTracers(; model, tracers = (:N, :P))
-simulation.callbacks[:nan_tendencies] = Callback(remove_NaN_tendencies!; callsite = TendencyCallsite())
-
 run!(simulation)
 ```
 

docs/src/quick_start.md

@@ -3,7 +3,7 @@
 OceanBioME provides biogeochemical models to plug into [Oceananigans](https://github.com/CliMA/Oceananigans.jl), for example this code will run one month of a single column, 7 variable (P, Z, sPOM, bPOM, DOM, NO₃, NH₄) biogeochemical situation with constant forcing.
 
 First we need to check we have the required dependencies:
-```@example quickstart
+```julia
 using Pkg
 Pkg.add(["OceanBioME", "Oceananigans"])
 ```
@@ -30,11 +30,15 @@ simulation.output_writers[:profiles] = JLD2OutputWriter(model, model.tracers,
 run!(simulation)
 ```
 
-This isn't quite as simple as it could be as it records the output so that we can visualize it:
+We can then visualize it, first check the required packages are installed:
 
-```@example quickstart
+```julia
 Pkg.add("CairoMakie")
+```
 
+and then load the data and plot:
+
+```@example quickstart
 using CairoMakie
 
 phytoplankton = FieldTimeSeries("quickstart.jld2", "P")

examples/column.jl

@@ -40,14 +40,22 @@ nothing #hide
 # Define the grid.
 grid = RectilinearGrid(size=(1, 1, 50), extent=(20meters, 20meters, 200meters))
 
-# Specify the boundary conditions for DIC and O₂ based on the air-sea CO₂ and O₂ flux
+
+# ## Model
+# First we define the biogeochemical model including carbonate chemistry 
+# and scaling of negative tracers(see discussion in the [positivity preservation](@ref pos-preservation))
+# and then setup the Oceananigans model with the boundary condition for the DIC based on the air-sea CO₂ flux.
+
+biogeochemistry = LOBSTER(; grid,
+                            surface_phytosynthetically_active_radiation = PAR⁰,
+                            carbonates = true,
+                            scale_negatives = true)
+
 CO₂_flux = GasExchange(; gas = :CO₂, temperature = temp, salinity = (args...) -> 35)
 
 model = NonhydrostaticModel(; grid,
-                              closure = ScalarDiffusivity(ν = κₜ, κ = κₜ), 
-                              biogeochemistry = LOBSTER(; grid,
-                                                          surface_phytosynthetically_active_radiation = PAR⁰,
-                                                          carbonates = true),
+                              closure = ScalarDiffusivity(ν = κₜ, κ = κₜ),
+                              biogeochemistry,
                               boundary_conditions = (DIC = FieldBoundaryConditions(top = CO₂_flux), ))
 
 set!(model, P = 0.03, Z = 0.03, NO₃ = 4.0, NH₄ = 0.05, DIC = 2239.8, Alk = 2409.0)
@@ -56,15 +64,14 @@ set!(model, P = 0.03, Z = 0.03, NO₃ = 4.0, NH₄ = 0.05, DIC = 2239.8, Alk = 2
 # Next we setup a simulation and add some callbacks that:
 # - Show the progress of the simulation
 # - Store the model and particles output
-# - Prevent the tracers from going negative from numerical error (see discussion in the [positivity preservation](@ref pos-preservation) implementation section)
 
 simulation = Simulation(model, Δt = 3minutes, stop_time = 100days)
 
 progress_message(sim) = @printf("Iteration: %04d, time: %s, Δt: %s, wall time: %s\n",
-                                                        iteration(sim),
-                                                        prettytime(sim),
-                                                        prettytime(sim.Δt),
-                                                        prettytime(sim.run_wall_time))      
+                                iteration(sim),
+                                prettytime(sim),
+                                prettytime(sim.Δt),
+                                prettytime(sim.run_wall_time))
 
 simulation.callbacks[:progress] = Callback(progress_message, TimeInterval(10days))
 
@@ -73,9 +80,6 @@ simulation.output_writers[:profiles] = JLD2OutputWriter(model, merge(model.trace
                                                         filename = "$filename.jld2",
                                                         schedule = TimeInterval(1day),
                                                         overwrite_existing = true)
-
-scale_negative_tracers = ScaleNegativeTracers(; model, tracers = (:NO₃, :NH₄, :P, :Z, :sPOM, :bPOM, :DOM))
-simulation.callbacks[:neg] = Callback(scale_negative_tracers; callsite = UpdateStateCallsite())
 nothing #hide
 
 # ## Run!
@@ -95,17 +99,20 @@ bPOM = FieldTimeSeries("$filename.jld2", "bPOM")
 
 x, y, z = nodes(P)
 times = P.times
+nothing #hide
 
 # We compute the  air-sea CO₂ flux at the surface (corresponding to vertical index `k = grid.Nz`) and
 # the carbon export by computing how much carbon sinks below some arbirtrary depth; here we use depth 
 # that corresponds to `k = grid.Nz - 20`.
 air_sea_CO₂_flux = zeros(length(times))
 carbon_export = zeros(length(times))
 
+using Oceananigans.Biogeochemistry: biogeochemical_drift_velocity
+
 for (i, t) in enumerate(times)
     air_sea_CO₂_flux[i] = CO₂_flux.condition.parameters(0.0, 0.0, t, DIC[1, 1, grid.Nz, i], Alk[1, 1, grid.Nz, i], temp(1, 1, 0, t), 35)
-    carbon_export[i] = (sPOM[1, 1, grid.Nz-20, i] * model.biogeochemistry.sinking_velocities.sPOM.w[1, 1, grid.Nz-20] +
-                        bPOM[1, 1, grid.Nz-20, i] * model.biogeochemistry.sinking_velocities.bPOM.w[1, 1, grid.Nz-20]) * model.biogeochemistry.organic_redfield
+    carbon_export[i] = (sPOM[1, 1, grid.Nz-20, i] * biogeochemical_drift_velocity(model.biogeochemistry, Val(:sPOM)).w[1, 1, grid.Nz-20] +
+                        bPOM[1, 1, grid.Nz-20, i] * biogeochemical_drift_velocity(model.biogeochemistry, Val(:bPOM)).w[1, 1, grid.Nz-20]) * redfield(Val(:sPOM), model.biogeochemistry)
 end
 
 # Both `air_sea_CO₂_flux` and `carbon_export` are in units `mmol CO₂ / (m² s)`.