From 111f4843469a5d6f9a7e499190e148dd7d379702 Mon Sep 17 00:00:00 2001 From: Jago Strong-Wright Date: Fri, 12 Jul 2024 15:23:37 +0100 Subject: [PATCH 1/6] made the box model ~2x faster --- benchmark/box_model.jl | 55 ++++++++++++++++++++++++++++++++++++ src/BoxModel/boxmodel.jl | 13 ++++++--- src/BoxModel/timesteppers.jl | 52 +++++++++++++++++++++------------- 3 files changed, 97 insertions(+), 23 deletions(-) create mode 100644 benchmark/box_model.jl diff --git a/benchmark/box_model.jl b/benchmark/box_model.jl new file mode 100644 index 000000000..0c012a76a --- /dev/null +++ b/benchmark/box_model.jl @@ -0,0 +1,55 @@ +# # Calibrating a biogeochemical model with `EnsembleKalmanProcesses` +# +# In this example we calibrate some of the parameters for the [NPZD](@ref NPZD) model +# in a simple box model setup using a data assimilation package [EnsembleKalmanProcesses](https://github.com/CliMA/EnsembleKalmanProcesses.jl). +# First we setup the model and generate synthetic data with "true" parameters. We then +# define priors and setup an EKP to solve. +# +# While this is a very simple situation it illustrates the ease of integration with +# data assimilation tools. Examples given in the EnsembleKalmanProcesses docs illustrate +# how the package can be used to solve more complex forward models. + +# ## Install dependencies +# First we ensure we have the required dependencies installed +# ```julia +# using Pkg +# pkg "add OceanBioME, Oceananigans, CairoMakie, EnsembleKalmanProcesses, Distributions" +# ``` + +using OceanBioME, EnsembleKalmanProcesses, JLD2, CairoMakie, Oceananigans.Units, Oceananigans, BenchmarkTools + +const year = years = 365day + +# Setup the forward model + +@inline PAR⁰(t) = 60 * (1 - cos((t + 15days) * 2π / year)) * (1 / (1 + 0.2 * exp(-((mod(t, year) - 200days) / 50days)^2))) + 2 + +z = -10 # nominal depth of the box for the PAR profile +@inline PAR(t)::Float64 = PAR⁰(t) * exp(0.2z) # Modify the PAR based on the nominal depth and exponential decay + +function run_box_simulation() + biogeochemistry = NutrientPhytoplanktonZooplanktonDetritus(; grid = BoxModelGrid, + light_attenuation_model = nothing) + + model = BoxModel(; biogeochemistry, forcing = (; PAR)) + + set!(model, N = 10.0, P = 0.1, Z = 0.01) + + simulation = Simulation(model; Δt = 20minutes, stop_iteration = 1000, verbose = false) + + simulation.output_writers[:fields] = JLD2OutputWriter(model, model.fields; filename = "box_benchmarking.jld2", schedule = IterationInterval(20), overwrite_existing = true) + + @info "Running the model..." + run!(simulation) +end + +##### +##### results +##### + +# Config: 1000 iterations with output every 8 hours and 20min steps +# origional @btime 317.5ms (1483607 allocations: 243.03 MiB) +# removing kernel launching from store tendencies: 291.546 ms (1293607 allocations: 187.95 MiB) +# removed kernel launching from rk3 substepping: 265.823 ms (1000607 allocations: 79.11 MiB) +# removed broadcasting in update state: 120.605 ms (619379 allocations: 63.98 MiB) +# no outputs: 23.523 ms (370344 allocations: 30.31 MiB) \ No newline at end of file diff --git a/src/BoxModel/boxmodel.jl b/src/BoxModel/boxmodel.jl index 6123c22b8..1d4263e91 100644 --- a/src/BoxModel/boxmodel.jl +++ b/src/BoxModel/boxmodel.jl @@ -30,10 +30,11 @@ import Base: show, summary @inline no_func(args...) = 0.0 -mutable struct BoxModel{G, B, F, FO, TS, C, PF} <: AbstractModel{TS} +mutable struct BoxModel{G, B, F, FV, FO, TS, C, PF} <: AbstractModel{TS} grid :: G # here so that simualtion can be built biogeochemistry :: B fields :: F + field_values :: FV forcing :: FO timestepper :: TS clock :: C @@ -66,22 +67,26 @@ function BoxModel(; biogeochemistry::B, variables = (required_biogeochemical_tracers(biogeochemistry)..., required_biogeochemical_auxiliary_fields(biogeochemistry)...) fields = NamedTuple{variables}([CenterField(BoxModelGrid) for var in eachindex(variables)]) + field_values = zeros(length(variables)) forcing = NamedTuple{variables}([variable in keys(forcing) ? forcing[variable] : no_func for variable in variables]) timestepper = BoxModelTimeStepper(timestepper, BoxModelGrid, keys(fields)) F = typeof(fields) + FV = typeof(field_values) FO = typeof(forcing) TS = typeof(timestepper) - return BoxModel{typeof(BoxModelGrid), B, F, FO, TS, C, PF}(BoxModelGrid, biogeochemistry, fields, forcing, timestepper, clock, prescribed_fields) + return BoxModel{typeof(BoxModelGrid), B, F, FV, FO, TS, C, PF}(BoxModelGrid, biogeochemistry, fields, field_values, forcing, timestepper, clock, prescribed_fields) end function update_state!(model::BoxModel, callbacks=[]; compute_tendencies = true) t = model.clock.time - @inbounds for field in model.prescribed_fields - (field in keys(model.fields)) && (model.fields[field] .= model.forcing[field](t)) + for field in model.prescribed_fields + if field in keys(model.fields) + @inbounds model.fields[field][1, 1, 1] = @inbounds model.forcing[field](t) + end end for callback in callbacks diff --git a/src/BoxModel/timesteppers.jl b/src/BoxModel/timesteppers.jl index 8bcd8590c..97bd5591c 100644 --- a/src/BoxModel/timesteppers.jl +++ b/src/BoxModel/timesteppers.jl @@ -18,10 +18,8 @@ end function store_tendencies!(model::BoxModel) model_fields = prognostic_fields(model) - for field_name in keys(model_fields) - launch!(architecture(model), model.grid, :xyz, store_field_tendencies!, - model.timestepper.G⁻[field_name], - model.timestepper.Gⁿ[field_name]) + @inbounds for field_name in keys(model_fields) + model.timestepper.G⁻[field_name][1, 1, 1] = model.timestepper.Gⁿ[field_name][1, 1, 1] end return nothing @@ -30,18 +28,14 @@ end function compute_tendencies!(model::BoxModel, callbacks) Gⁿ = model.timestepper.Gⁿ - model_fields = @inbounds [field[1] for field in model.fields] - - @inbounds for tracer in required_biogeochemical_tracers(model.biogeochemistry) - if !(tracer == :T) - getproperty(Gⁿ, tracer)[1] = model.biogeochemistry(Val(tracer), 0.0, 0.0, 0.0, model.clock.time, model_fields...) + model.forcing[tracer](model.clock.time, model_fields...) - end + @inbounds for (i, field) in enumerate(model.fields) + model.field_values[i] = field[1, 1, 1] end - @inbounds for variable in required_biogeochemical_auxiliary_fields(model.biogeochemistry) - if !(variable == :PAR) - getproperty(Gⁿ, variable)[1] = model.forcing[variable](model.clock.time, model_fields...) - end + for tracer in required_biogeochemical_tracers(model.biogeochemistry) + forcing = @inbounds model.forcing[tracer] + + @inbounds Gⁿ[tracer][1, 1, 1] = tracer_tendency(Val(tracer), model.biogeochemistry, forcing, model.clock.time, model.field_values) end for callback in callbacks @@ -53,15 +47,35 @@ function compute_tendencies!(model::BoxModel, callbacks) return nothing end +@inline tracer_tendency(val_name, biogeochemistry, forcing, time, model_fields) = + biogeochemistry(val_name, 0, 0, 0, time, model_fields...) + forcing(time, model_fields...) + +@inline tracer_tendency(::Val{:T}, biogeochemistry, forcing, time, model_fields) = 0 + function rk3_substep!(model::BoxModel, Δt, γⁿ, ζⁿ) - workgroup, worksize = work_layout(model.grid, :xyz) - substep_field_kernel! = rk3_substep_field!(device(architecture(model)), workgroup, worksize) model_fields = prognostic_fields(model) for (i, field) in enumerate(model_fields) - substep_field_kernel!(field, Δt, γⁿ, ζⁿ, - model.timestepper.Gⁿ[i], - model.timestepper.G⁻[i]) + Gⁿ = @inbounds model.timestepper.Gⁿ[i] + G⁻ = @inbounds model.timestepper.G⁻[i] + + rk3_substep!(field, Δt, γⁿ, ζⁿ, Gⁿ, G⁻) + end + + return nothing +end + +function rk3_substep!(U, Δt, γⁿ::FT, ζⁿ, Gⁿ, G⁻) where FT + @inbounds begin + U[1, 1, 1] += convert(FT, Δt) * (γⁿ * Gⁿ[1, 1, 1] + ζⁿ * G⁻[1, 1, 1]) + end + + return nothing +end + +function rk3_substep!(U, Δt, γ¹::FT, ::Nothing, G¹, G⁰) where FT + @inbounds begin + U[1, 1, 1] += convert(FT, Δt) * γ¹ * G¹[1, 1, 1] end return nothing From 48af01339149fb44257ec355557d5939a03d20b0 Mon Sep 17 00:00:00 2001 From: Jago Strong-Wright Date: Fri, 12 Jul 2024 15:43:21 +0100 Subject: [PATCH 2/6] made speedy output writer (~10x faster than normal writer) --- benchmark/box_model.jl | 25 ++++++++++++++++++++++--- src/BoxModel/boxmodel.jl | 3 ++- src/BoxModel/output_writer.jl | 27 +++++++++++++++++++++++++++ src/OceanBioME.jl | 2 +- 4 files changed, 52 insertions(+), 5 deletions(-) create mode 100644 src/BoxModel/output_writer.jl diff --git a/benchmark/box_model.jl b/benchmark/box_model.jl index 0c012a76a..21dfebc22 100644 --- a/benchmark/box_model.jl +++ b/benchmark/box_model.jl @@ -16,7 +16,7 @@ # pkg "add OceanBioME, Oceananigans, CairoMakie, EnsembleKalmanProcesses, Distributions" # ``` -using OceanBioME, EnsembleKalmanProcesses, JLD2, CairoMakie, Oceananigans.Units, Oceananigans, BenchmarkTools +using OceanBioME, Oceananigans.Units, Oceananigans, BenchmarkTools const year = years = 365day @@ -37,12 +37,29 @@ function run_box_simulation() simulation = Simulation(model; Δt = 20minutes, stop_iteration = 1000, verbose = false) - simulation.output_writers[:fields] = JLD2OutputWriter(model, model.fields; filename = "box_benchmarking.jld2", schedule = IterationInterval(20), overwrite_existing = true) + #simulation.output_writers[:fields] = JLD2OutputWriter(model, model.fields; filename = "box_benchmarking.jld2", schedule = IterationInterval(20), overwrite_existing = true) + + fast_output = SpeedyOutput("box_benchmarking.jld2") + + simulation.callbacks[:output] = Callback(fast_output, IterationInterval(20);) @info "Running the model..." run!(simulation) end +function fast_output(sim, fname) + file = jldopen(fname, "w+") + + model = sim.model + + t = time(sim) + + file["fields/$t"] = model.field_values + + close(file) + + return nothing +end ##### ##### results ##### @@ -52,4 +69,6 @@ end # removing kernel launching from store tendencies: 291.546 ms (1293607 allocations: 187.95 MiB) # removed kernel launching from rk3 substepping: 265.823 ms (1000607 allocations: 79.11 MiB) # removed broadcasting in update state: 120.605 ms (619379 allocations: 63.98 MiB) -# no outputs: 23.523 ms (370344 allocations: 30.31 MiB) \ No newline at end of file +# no outputs: 23.523 ms (370344 allocations: 30.31 MiB) +# outputting every 20 steps: 1.181s +# outputting every 20 steps with `SpeedyOutput`: 34ms \ No newline at end of file diff --git a/src/BoxModel/boxmodel.jl b/src/BoxModel/boxmodel.jl index 1d4263e91..70d9db1db 100644 --- a/src/BoxModel/boxmodel.jl +++ b/src/BoxModel/boxmodel.jl @@ -3,7 +3,7 @@ Integrate biogeochemical models on a single point """ module BoxModels -export BoxModel, run!, set!, SaveBoxModel +export BoxModel, run!, set!, SpeedyOutput using Oceananigans: Clock, prettytime using Oceananigans.Biogeochemistry: @@ -137,6 +137,7 @@ end default_included_properties(::BoxModel) = [:grid] include("timesteppers.jl") +include("output_writer.jl") summary(::BoxModel{B, V, F, TS, C}) where {B, V, F, TS, C} = string("Biogeochemical box model") show(io::IO, model::BoxModel{B, V, F, TS, C}) where {B, V, F, TS, C} = diff --git a/src/BoxModel/output_writer.jl b/src/BoxModel/output_writer.jl new file mode 100644 index 000000000..519faccf9 --- /dev/null +++ b/src/BoxModel/output_writer.jl @@ -0,0 +1,27 @@ +struct SpeedyOutput{FN, B} # I will make this an `AbstractOutputWriter` at some point but this will do for now + filename :: FN + overwrite_existing :: B + + function SpeedyOutput(filename::FN; overwrite_existing::B = true) where {FN, B} + return new{FN, B}(filename, overwrite_existing) + end +end + +function (save::SpeedyOutput)(simulation) + model = simulation.model + + t = time(simulation) + iter = model.clock.iteration + + file = jldopen(save.filename, ifelse(save.overwrite_existing, "w+", "w")) + + file["timeseries/t/$iter"] = t + + for (i, name) in enumerate(keys(model.fields)) + file["timeseries/$name/$iter"] = model.field_values[i] + end + + close(file) + + return nothing +end \ No newline at end of file diff --git a/src/OceanBioME.jl b/src/OceanBioME.jl index 7bb9c5d6c..9d06d57da 100644 --- a/src/OceanBioME.jl +++ b/src/OceanBioME.jl @@ -11,7 +11,7 @@ export Biogeochemistry, LOBSTER, NutrientPhytoplanktonZooplanktonDetritus, NPZD, export SLatissima # Box model -export BoxModel, BoxModelGrid, SaveBoxModel, run!, set! +export BoxModel, BoxModelGrid, SpeedyOutput, run!, set! # Particles export Particles From 3195c770d5aa18d724b435f5d32ed06b4e46a4d0 Mon Sep 17 00:00:00 2001 From: Jago Strong-Wright Date: Sat, 13 Jul 2024 12:43:06 +0100 Subject: [PATCH 3/6] removed npzd box model example (not included in docs before anyway) --- examples/box_npzd.jl | 73 -------------------------------------------- 1 file changed, 73 deletions(-) delete mode 100644 examples/box_npzd.jl diff --git a/examples/box_npzd.jl b/examples/box_npzd.jl deleted file mode 100644 index b0a9ca53e..000000000 --- a/examples/box_npzd.jl +++ /dev/null @@ -1,73 +0,0 @@ -# # Box model -# In this example we setup a simple NPZD biogeochemical model in a single box configuration. -# This example demonstrates: -# - How to setup OceanBioME's biogeochemical models as a stand-alone box model - -# ## Install dependencies -# First we check we have the dependencies installed -# ```julia -# using Pkg -# pkg"add OceanBioME, JLD2, CairoMakie" -# ``` - -# ## Model setup -# Load the packages and setup the initial and forcing conditions -using OceanBioME - -minute = minutes = 60 -hour = hours = 60minutes -day = days = 24hours -year = years = 365day -nothing #hide - -# This is forced by a prescribed time-dependent photosynthetically available radiation (PAR) -PAR⁰(t) = 50 * (1 - cos((t + 15days) * 2π / (365days))) * (1 / (1 + 0.2 * exp(-((mod(t, 365days) - 200days) / 50days)^2))) + 10 - -z = 10 # specify the nominal depth of the box for the PAR profile -PAR(t) = PAR⁰(t) * exp(-0.1z) # Modify the PAR based on the nominal depth and exponential decay - -T(t) = 5 * (1 - cos((t + 30days) * 2π / (365days))) / 2 + 15 -nothing #hide - -# Set up the model. Here, first specify the biogeochemical model, followed by initial conditions and the start and end times -model = BoxModel(biogeochemistry = NutrientPhytoplanktonZooplanktonDetritus(grid = BoxModelGrid()), forcing = (; PAR, T)) -model.Δt = 2day -model.stop_time = 10years - -set!(model, N = 7.0, P = 0.01, Z = 0.05) - -# ## Run the model (should only take a few seconds) -@info "Running the model..." -run!(model, save_interval = 1, save = SaveBoxModel("box_npzd.jld2")) - - -# ## Load the output -using JLD2 - -file = jldopen("box_npzd.jld2") -vars = (:N, :P, :Z, :D, :T, :PAR) -times = parse.(Float64, keys(file["values"])) - -timeseries = NamedTuple{vars}(ntuple(t -> zeros(length(times)), length(vars))) - -for (idx, time) in enumerate(times) - values = file["values/$time"] - for tracer in vars - getproperty(timeseries, tracer)[idx] = values[tracer] - end -end - -close(file) - -# ## And plot -using CairoMakie - -fig = Figure(size = (800, 1200), fontsize = 24) - -axs = [] -for (idx, tracer) in enumerate(vars) - push!(axs, Axis(fig[idx, 1], ylabel = "$tracer", xlabel = "Year", xticks=(0:10))) - lines!(axs[end], times / year, timeseries[tracer], linewidth = 3) -end - -fig From f5a83773b45c7704c08dae3f73807cce1fb13f5d Mon Sep 17 00:00:00 2001 From: Jago Strong-Wright Date: Fri, 19 Jul 2024 18:34:38 +0100 Subject: [PATCH 4/6] Fairly sure this isn't breaking --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 2cfd1e0cb..7781e1e0b 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "OceanBioME" uuid = "a49af516-9db8-4be4-be45-1dad61c5a376" authors = ["Jago Strong-Wright and contributors"] -version = "0.10.3" +version = "0.10.4" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" From 116fa193c2d0cf24bdebcf8290399c175d49dcc2 Mon Sep 17 00:00:00 2001 From: Jago Strong-Wright Date: Wed, 24 Jul 2024 14:15:55 +0100 Subject: [PATCH 5/6] made prescribed par --- examples/box.jl | 14 +++++++++++--- src/BoxModel/boxmodel.jl | 22 ++++++++++----------- src/BoxModel/timesteppers.jl | 6 +++++- src/Light/Light.jl | 3 ++- src/Light/prescribed.jl | 37 ++++++++++++++++++++++++++++++++++++ src/OceanBioME.jl | 2 +- 6 files changed, 67 insertions(+), 17 deletions(-) create mode 100644 src/Light/prescribed.jl diff --git a/examples/box.jl b/examples/box.jl index f07bc8702..4574c0641 100644 --- a/examples/box.jl +++ b/examples/box.jl @@ -13,19 +13,27 @@ # ## Model setup # Load the packages and setup the initial and forcing conditions using OceanBioME, Oceananigans, Oceananigans.Units +using Oceananigans.Fields: FunctionField const year = years = 365day + +grid = BoxModelGrid +clock = Clock(time = zero(grid)) + nothing #hide # This is forced by a prescribed time-dependent photosynthetically available radiation (PAR) PAR⁰(t) = 60 * (1 - cos((t + 15days) * 2π / year)) * (1 / (1 + 0.2 * exp(-((mod(t, year) - 200days) / 50days)^2))) + 2 -z = -10 # specify the nominal depth of the box for the PAR profile -PAR(t) = PAR⁰(t) * exp(0.2z) # Modify the PAR based on the nominal depth and exponential decay +const z = -10 # specify the nominal depth of the box for the PAR profile +PAR_func(t) = PAR⁰(t) * exp(0.2z) # Modify the PAR based on the nominal depth and exponential decay + +PAR = FunctionField{Center, Center, Center}(PAR_func, grid; clock) nothing #hide # Set up the model. Here, first specify the biogeochemical model, followed by initial conditions and the start and end times -model = BoxModel(biogeochemistry = LOBSTER(grid = BoxModelGrid, light_attenuation_model = nothing), forcing = (; PAR)) +model = BoxModel(; biogeochemistry = LOBSTER(; grid, light_attenuation_model = PrescribedPhotosyntheticallyActiveRadiation(PAR)), + clock) set!(model, NO₃ = 10.0, NH₄ = 0.1, P = 0.1, Z = 0.01) diff --git a/src/BoxModel/boxmodel.jl b/src/BoxModel/boxmodel.jl index 70d9db1db..1a6546461 100644 --- a/src/BoxModel/boxmodel.jl +++ b/src/BoxModel/boxmodel.jl @@ -30,7 +30,7 @@ import Base: show, summary @inline no_func(args...) = 0.0 -mutable struct BoxModel{G, B, F, FV, FO, TS, C, PF} <: AbstractModel{TS} +mutable struct BoxModel{G, B, F, FV, FO, TS, C, PT} <: AbstractModel{TS} grid :: G # here so that simualtion can be built biogeochemistry :: B fields :: F @@ -38,15 +38,15 @@ mutable struct BoxModel{G, B, F, FV, FO, TS, C, PF} <: AbstractModel{TS} forcing :: FO timestepper :: TS clock :: C - prescribed_fields :: PF + prescribed_tracers :: PT end """ - BoxModel(; biogeochemistry::B, + BoxModel(; biogeochemistry, forcing = NamedTuple(), timestepper = :RungeKutta3, - clock::C = Clock(; time = 0.0), - prescribed_fields::PF = (:T, :PAR)) + clock = Clock(; time = 0.0), + prescribed_tracers = (:T, )) Constructs a box model of a `biogeochemistry` model. Once this has been constructed you can set initial condiitons by `set!(model, X=1.0...)` and then `run!(model)`. @@ -57,17 +57,17 @@ Keyword Arguments - `forcing`: NamedTuple of additional forcing functions for the biogeochemical tracers to be integrated - `timestepper`: Timestepper to integrate model - `clock`: Oceananigans clock to keep track of time -- `prescribed_fields`: Tuple of fields names (Symbols) which are not integrated but provided in `forcing` as a function of time with signature `f(t)` +- `prescribed_tracers`: Tuple of fields names (Symbols) which are not integrated but provided in `forcing` as a function of time with signature `f(t)` """ function BoxModel(; biogeochemistry::B, forcing = NamedTuple(), timestepper = :RungeKutta3, clock::C = Clock(; time = 0.0), - prescribed_fields::PF = (:T, :PAR)) where {B, C, PF} + prescribed_tracers::PT = (T = (t) -> 0, )) where {B, C, PT} - variables = (required_biogeochemical_tracers(biogeochemistry)..., required_biogeochemical_auxiliary_fields(biogeochemistry)...) + variables = required_biogeochemical_tracers(biogeochemistry) fields = NamedTuple{variables}([CenterField(BoxModelGrid) for var in eachindex(variables)]) - field_values = zeros(length(variables)) + field_values = zeros(length(variables)+length(required_biogeochemical_auxiliary_fields(biogeochemistry))) forcing = NamedTuple{variables}([variable in keys(forcing) ? forcing[variable] : no_func for variable in variables]) timestepper = BoxModelTimeStepper(timestepper, BoxModelGrid, keys(fields)) @@ -77,13 +77,13 @@ function BoxModel(; biogeochemistry::B, FO = typeof(forcing) TS = typeof(timestepper) - return BoxModel{typeof(BoxModelGrid), B, F, FV, FO, TS, C, PF}(BoxModelGrid, biogeochemistry, fields, field_values, forcing, timestepper, clock, prescribed_fields) + return BoxModel{typeof(BoxModelGrid), B, F, FV, FO, TS, C, PT}(BoxModelGrid, biogeochemistry, fields, field_values, forcing, timestepper, clock, prescribed_tracers) end function update_state!(model::BoxModel, callbacks=[]; compute_tendencies = true) t = model.clock.time - for field in model.prescribed_fields + for field in model.prescribed_tracers if field in keys(model.fields) @inbounds model.fields[field][1, 1, 1] = @inbounds model.forcing[field](t) end diff --git a/src/BoxModel/timesteppers.jl b/src/BoxModel/timesteppers.jl index 97bd5591c..ac30aff7b 100644 --- a/src/BoxModel/timesteppers.jl +++ b/src/BoxModel/timesteppers.jl @@ -1,5 +1,5 @@ using Oceananigans.Architectures: device -using Oceananigans.Biogeochemistry: update_tendencies! +using Oceananigans.Biogeochemistry: update_tendencies!, biogeochemical_auxiliary_fields using Oceananigans.TimeSteppers: rk3_substep_field!, store_field_tendencies!, RungeKutta3TimeStepper, QuasiAdamsBashforth2TimeStepper using Oceananigans.Utils: work_layout, launch! @@ -32,6 +32,10 @@ function compute_tendencies!(model::BoxModel, callbacks) model.field_values[i] = field[1, 1, 1] end + @inbounds for (i, field) in enumerate(values(biogeochemical_auxiliary_fields(model.biogeochemistry))) + model.field_values[i+length(model.fields)] = field[1, 1, 1] + end + for tracer in required_biogeochemical_tracers(model.biogeochemistry) forcing = @inbounds model.forcing[tracer] diff --git a/src/Light/Light.jl b/src/Light/Light.jl index 7b400dae0..c9e649579 100644 --- a/src/Light/Light.jl +++ b/src/Light/Light.jl @@ -3,7 +3,7 @@ Light attenuation by chlorophyll as described by [Karleskind2011](@citet) (imple """ module Light -export TwoBandPhotosyntheticallyActiveRadiation, update_PAR! +export TwoBandPhotosyntheticallyActiveRadiation, update_PAR!, PrescribedPhotosyntheticallyActiveRadiation using KernelAbstractions, Oceananigans.Units using Oceananigans.Architectures: device, architecture @@ -29,6 +29,7 @@ import Oceananigans.BoundaryConditions: _fill_top_halo! include("2band.jl") include("morel.jl") +include("prescribed.jl") default_surface_PAR(x, y, t) = default_surface_PAR(t) default_surface_PAR(x_or_y, t) = default_surface_PAR(t) diff --git a/src/Light/prescribed.jl b/src/Light/prescribed.jl new file mode 100644 index 000000000..c2f1d64f1 --- /dev/null +++ b/src/Light/prescribed.jl @@ -0,0 +1,37 @@ +using Oceananigans.Fields: compute! + +maybe_named_fields(field) = ((:PAR, ), (field, )) +maybe_named_fields(fields::NamedTuple) = (keys(fields), values(fields)) + +struct PrescribedPhotosyntheticallyActiveRadiation{F, FN} + fields :: F + field_names :: FN + + PrescribedPhotosyntheticallyActiveRadiation(fields::F, names::FN) where {F, FN} = + new{F, FN}(fields, names) + + function PrescribedPhotosyntheticallyActiveRadiation(fields) + names, values = maybe_named_fields(fields) + + F = typeof(values) + FN = typeof(names) + + return new{F, FN}(values, names) + end +end + +function update_biogeochemical_state!(model, PAR::PrescribedPhotosyntheticallyActiveRadiation) + for field in PAR.fields + compute!(field) + end + + return nothing +end + +summary(::PrescribedPhotosyntheticallyActiveRadiation{FT}) where {FT} = string("Prescribed PAR") +show(io::IO, model::PrescribedPhotosyntheticallyActiveRadiation{FT}) where {FT} = print(io, summary(model)) + +biogeochemical_auxiliary_fields(par::PrescribedPhotosyntheticallyActiveRadiation) = NamedTuple{par.field_names}(par.fields) + +adapt_structure(to, par::PrescribedPhotosyntheticallyActiveRadiation) = PrescribedPhotosyntheticallyActiveRadiation(adapt(to, par.fields), + adapt(to, par.field_names)) diff --git a/src/OceanBioME.jl b/src/OceanBioME.jl index 9d06d57da..1df55b2a1 100644 --- a/src/OceanBioME.jl +++ b/src/OceanBioME.jl @@ -17,7 +17,7 @@ export BoxModel, BoxModelGrid, SpeedyOutput, run!, set! export Particles # Light models -export TwoBandPhotosyntheticallyActiveRadiation +export TwoBandPhotosyntheticallyActiveRadiation, PrescribedPhotosyntheticallyActiveRadiation # Boundaries export Boundaries, Sediments, GasExchange, FlatSediment From 9066a7ab8cfeaee73b714748002b144c0558aab7 Mon Sep 17 00:00:00 2001 From: Jago Strong-Wright Date: Wed, 24 Jul 2024 14:38:40 +0100 Subject: [PATCH 6/6] added docstring --- src/Light/prescribed.jl | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/src/Light/prescribed.jl b/src/Light/prescribed.jl index c2f1d64f1..44c4d385c 100644 --- a/src/Light/prescribed.jl +++ b/src/Light/prescribed.jl @@ -1,8 +1,26 @@ -using Oceananigans.Fields: compute! +using Oceananigans.Fields: compute!, AbstractField + +function maybe_named_fields(field) + + isa(field, AbstractField) || @warn "fields: $field is not an `AbstractField" + + return ((:PAR, ), (field, )) +end -maybe_named_fields(field) = ((:PAR, ), (field, )) maybe_named_fields(fields::NamedTuple) = (keys(fields), values(fields)) +""" + PrescribedPhotosyntheticallyActiveRadiation(fields) + +`PrescribedPhotosyntheticallyActiveRadiation` returns "prescribed" PAR +fields which are user specified, e.g. they may be `FunctionField`s or +`ConstantField`s. + +`fields` may either be an `AbstractField` or a `NamedTuple` of names and +fields which will be returned in `biogeochemical_auxiliary_fields`, if only +one field is present the field will be named `PAR`. +""" + struct PrescribedPhotosyntheticallyActiveRadiation{F, FN} fields :: F field_names :: FN