Make BoxModels speedy again (+ introduce prescribed PAR)

Project.toml

@@ -1,7 +1,7 @@
 name = "OceanBioME"
 uuid = "a49af516-9db8-4be4-be45-1dad61c5a376"
 authors = ["Jago Strong-Wright <[email protected]> and contributors"]
-version = "0.10.3"
+version = "0.10.4"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

benchmark/box_model.jl

@@ -0,0 +1,74 @@
+# # 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, 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)
+
+    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
+#####
+
+# 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)
+# outputting every 20 steps: 1.181s 
+# outputting every 20 steps with `SpeedyOutput`: 34ms

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)
 

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: 
@@ -30,22 +30,23 @@ 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, PT} <: AbstractModel{TS}
                grid :: G # here so that simualtion can be built
     biogeochemistry :: B
              fields :: F
+       field_values :: FV
             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)`.
 
@@ -56,32 +57,36 @@ 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)+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))
 
     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, 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
 
-    @inbounds for field in model.prescribed_fields 
-        (field in keys(model.fields)) && (model.fields[field] .= model.forcing[field](t))
+    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
     end
 
     for callback in callbacks
@@ -132,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} = 

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

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!
 
@@ -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,18 @@ end
 function compute_tendencies!(model::BoxModel, callbacks)
     Gⁿ = model.timestepper.Gⁿ
 
-    model_fields = @inbounds [field[1] for field in model.fields]
+    @inbounds for (i, field) in enumerate(model.fields)
+        model.field_values[i] = field[1, 1, 1]
+    end
 
-    @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(values(biogeochemical_auxiliary_fields(model.biogeochemistry)))
+        model.field_values[i+length(model.fields)] = 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 +51,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