Merge pull request #217 from OceanBioME/jsw/rewrite-PISCES-again

Rewrites PISCES to be a discrete form model, I think over all this makes the implementation a lot more flexible

docs/src/appendix/library.md

@@ -27,6 +27,42 @@ Modules = [OceanBioME.Models.LOBSTERModel]
 Modules = [OceanBioME.Models.PISCESModel]
 ```
 
+```@autodocs
+Modules = [OceanBioME.Models.PISCESModel.DissolvedOrganicMatter]
+```
+
+```@autodocs
+Modules = [OceanBioME.Models.PISCESModel.ParticulateOrganicMatter]
+```
+
+```@autodocs
+Modules = [OceanBioME.Models.PISCESModel.Iron]
+```
+
+```@autodocs
+Modules = [OceanBioME.Models.PISCESModel.InorganicCarbons]
+```
+
+```@autodocs
+Modules = [OceanBioME.Models.PISCESModel.Zooplankton]
+```
+
+```@autodocs
+Modules = [OceanBioME.Models.PISCESModel.Phytoplankton]
+```
+
+```@autodocs
+Modules = [OceanBioME.Models.PISCESModel.Phosphates]
+```
+
+```@autodocs
+Modules = [OceanBioME.Models.PISCESModel.Silicates]
+```
+
+```@autodocs
+Modules = [OceanBioME.Models.PISCESModel.Nitrogen]
+```
+
 ### Sugar kelp (Saccharina latissima)
 
 ```@autodocs

docs/src/model_components/biogeochemical/PISCES/PISCES.md

@@ -7,7 +7,7 @@ An overview of the model structure is available from the [PISCES community websi
 
 ![PISCES model structure](https://www.pisces-community.org/wp-content/uploads/2021/12/PISCES_Operational-1.png)
 
-More documentation to follow..., for now see our list of key differences from [Aumont2015](@citet) and queries we have about the parameterisations on the [notable differences](@ref PISCES_queries).
+More documentation to follow..., for now see our list of key differences from [Aumont2015](@citet) and queries we have about the parametrisations on the [notable differences](@ref PISCES_queries).
 
 ## Model conservation
 When the permanent scavenging of iron, and nitrogen fixation are turned off, PISCES conserves:

src/BoxModel/timesteppers.jl

@@ -1,5 +1,5 @@
 using Oceananigans.Architectures: device
-using Oceananigans.Biogeochemistry: update_tendencies!, biogeochemical_auxiliary_fields
+using Oceananigans.Biogeochemistry: update_tendencies!, biogeochemical_auxiliary_fields, AbstractBiogeochemistry, AbstractContinuousFormBiogeochemistry
 using Oceananigans.Grids: nodes, Center
 using Oceananigans.TimeSteppers: rk3_substep_field!, store_field_tendencies!, RungeKutta3TimeStepper, QuasiAdamsBashforth2TimeStepper
 using Oceananigans.Utils: work_layout, launch!
@@ -41,7 +41,7 @@ function compute_tendencies!(model::BoxModel, callbacks)
     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, model.grid)
+        @inbounds Gⁿ[tracer][1, 1, 1] = tracer_tendency(Val(tracer), model.biogeochemistry, forcing, model.clock, model.fields, model.field_values, model.grid)
     end
 
     for callback in callbacks
@@ -57,8 +57,11 @@ end
 @inline boxmodel_coordinate(::Nothing, grid) = zero(grid)
 @inline boxmodel_coordinate(nodes, grid) = @inbounds nodes[1]
 
-@inline tracer_tendency(val_name, biogeochemistry, forcing, time, model_fields, grid) =
-    biogeochemistry(val_name, boxmodel_xyz(nodes(grid, Center(), Center(), Center()), grid)..., time, model_fields...) + forcing(time, model_fields...)
+@inline tracer_tendency(val_name, biogeochemistry::AbstractContinuousFormBiogeochemistry, forcing, clock, model_fields, model_field_values, grid) =
+    biogeochemistry(val_name, boxmodel_xyz(nodes(grid, Center(), Center(), Center()), grid)..., clock.time, model_field_values...) + forcing(time, model_fields...)
+
+@inline tracer_tendency(val_name, biogeochemistry::AbstractBiogeochemistry, forcing, clock, model_fields, model_field_values, grid) =
+    biogeochemistry(1, 1, 1, grid, val_name, clock, model_fields) + forcing(clock, model_fields)
 
 function rk3_substep!(model::BoxModel, Δt, γⁿ, ζⁿ)
     model_fields = prognostic_fields(model)

src/Light/multi_band.jl

@@ -71,7 +71,7 @@ function MultiBandPhotosyntheticallyActiveRadiation(; grid,
                                                       base_water_attenuation_coefficient = MOREL_kʷ,
                                                       base_chlorophyll_exponent = MOREL_e,
                                                       base_chlorophyll_attenuation_coefficient = MOREL_χ,
-                                                      field_names = [par_symbol(n) for n in 1:length(bands)],
+                                                      field_names = ntuple(n->par_symbol(n), Val(length(bands))),
                                                       surface_PAR = default_surface_PAR,
                                                       surface_PAR_division = fill(1 / length(bands), length(bands)))
     Nbands = length(bands)
@@ -91,11 +91,17 @@ function MultiBandPhotosyntheticallyActiveRadiation(; grid,
 
     sum(surface_PAR_division) == 1 || throw(ArgumentError("surface_PAR_division does not sum to 1"))
 
-    fields = [CenterField(grid; 
-                boundary_conditions = 
-                    regularize_field_boundary_conditions(
-                        FieldBoundaryConditions(top = ValueBoundaryCondition(ScaledSurfaceFunction(surface_PAR, surface_PAR_division[n]))), grid, name)) 
-              for (n, name) in enumerate(field_names)]
+    n_fields = length(field_names)
+
+    surface_boundary_conditions = 
+        ntuple(n-> ValueBoundaryCondition(ScaledSurfaceFunction(surface_PAR, surface_PAR_division[n])), Val(n_fields))
+
+    field_boundary_conditions =
+        ntuple(n -> regularize_field_boundary_conditions(FieldBoundaryConditions(top = surface_boundary_conditions[n]),
+                                                         grid, field_names[n]),
+               Val(n_fields))
+
+    fields = NamedTuple{field_names}(ntuple(n -> CenterField(grid; boundary_conditions = field_boundary_conditions[n]), Val(n_fields)))
 
     total_PAR = sum(fields)
 
@@ -115,14 +121,10 @@ function numerical_mean(λ, C, idx1, idx2)
     return ∫Cdλ / ∫dλ
 end
 
-function par_symbol(n)
-    subscripts = Symbol[]
-    for digit in reverse(digits(n))
-        push!(subscripts, Symbol(Char('\xe2\x82\x80'+digit)))
-    end
+@inline par_symbol(n) = Symbol(:PAR, Char('\xe2\x82\x80'+n)) #Symbol(:PAR, number_subscript(tuple(reverse(digits(n))...))...)
 
-    return Symbol(:PAR, subscripts...)
-end
+@inline number_subscript(digits::NTuple{N}) where N =
+    ntuple(n->Symbol(Char('\xe2\x82\x80'+digits[n])), Val(N))
 
 @kernel function update_MultiBandPhotosyntheticallyActiveRadiation!(grid, field, kʷ, e, χ,
                                                                     _surface_PAR, surface_PAR_division, 
@@ -146,7 +148,6 @@ end
     end
 end
 
-
 function update_biogeochemical_state!(model, PAR::MultiBandPhotosyntheticallyActiveRadiation)
     grid = model.grid
 
@@ -168,9 +169,9 @@ function update_biogeochemical_state!(model, PAR::MultiBandPhotosyntheticallyAct
                 k′)
     end
 
-    for field in PAR.fields
-        fill_halo_regions!(field, model.clock, fields(model))
-    end
+    #for field in PAR.fields
+    #    fill_halo_regions!(field, model.clock, fields(model))
+    #end
 end
 
 summary(par::MultiBandPhotosyntheticallyActiveRadiation) = 
@@ -179,7 +180,19 @@ summary(par::MultiBandPhotosyntheticallyActiveRadiation) =
 show(io::IO, model::MultiBandPhotosyntheticallyActiveRadiation) = print(io, summary(model))
 
 biogeochemical_auxiliary_fields(par::MultiBandPhotosyntheticallyActiveRadiation) = 
-    merge((PAR = par.total, ), NamedTuple{par.field_names}(par.fields))
+    merge((PAR = par.total, ), par.fields)#NamedTuple{field_names(par.field_names, par.fields)}(par.fields))
+
+@inline field_names(field_names, fields) = field_names
+@inline field_names(::Nothing, fields::NTuple{N}) where N = ntuple(n -> par_symbol(n), Val(N))
 
 # avoid passing this into kernels
-Adapt.adapt_structure(to, par::MultiBandPhotosyntheticallyActiveRadiation) = nothing
+Adapt.adapt_structure(to, par::MultiBandPhotosyntheticallyActiveRadiation) = 
+    MultiBandPhotosyntheticallyActiveRadiation(adapt(to, par.total),
+                                               adapt(to, par.fields),
+                                               nothing,
+                                               nothing, 
+                                               nothing,
+                                               nothing,
+                                               nothing,
+                                               nothing)
+

src/Light/prescribed.jl

@@ -1,13 +1,15 @@
-using Oceananigans.Fields: compute!, AbstractField
+using Oceananigans.Architectures: architecture, GPU
+using Oceananigans.Fields: compute!, AbstractField, ConstantField
 
 function maybe_named_fields(field)
 
     isa(field, AbstractField) || @warn "fields: $field is not an `AbstractField"
 
-    return ((:PAR, ), (field, ))
+    return NamedTuple{(:PAR, )}((field, ))
 end
 
-maybe_named_fields(fields::NamedTuple) = (keys(fields), values(fields))
+maybe_named_fields(fields::NamedTuple) = fields
+
 
 """
     PrescribedPhotosyntheticallyActiveRadiation(fields)
@@ -20,26 +22,20 @@ fields which are user specified, e.g. they may be `FunctionField`s or
 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}
+struct PrescribedPhotosyntheticallyActiveRadiation{F}
     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)
+        fields = maybe_named_fields(fields)
 
-        F  = typeof(values)
-        FN = typeof(names)
-
-        return new{F, FN}(values, names)
+        F = typeof(fields)
+
+        return new{F}(fields)
     end
 end
 
 function update_biogeochemical_state!(model, PAR::PrescribedPhotosyntheticallyActiveRadiation)
-    for field in PAR.fields
+    for field in values(PAR.fields)
         compute!(field)
     end
 
@@ -49,9 +45,12 @@ end
 summary(::PrescribedPhotosyntheticallyActiveRadiation) = string("Prescribed PAR")
 show(io::IO, model::PrescribedPhotosyntheticallyActiveRadiation{F}) where {F} = print(io, summary(model), "\n",
                                                                                             "  Fields:", "\n",
-                                                                                            "    └── $(model.field_names)")
+                                                                                            "    └── $(keys(model.fields))")
+
+biogeochemical_auxiliary_fields(par::PrescribedPhotosyntheticallyActiveRadiation) = par.fields
 
-biogeochemical_auxiliary_fields(par::PrescribedPhotosyntheticallyActiveRadiation) = NamedTuple{par.field_names}(par.fields)
+@inline prescribed_field_names(field_names, fields) = field_names
+@inline prescribed_field_names(::Nothing, fields::NTuple{N}) where N = tuple(:PAR, ntuple(n -> par_symbol(n), Val(N-1))...)
 
-adapt_structure(to, par::PrescribedPhotosyntheticallyActiveRadiation) = PrescribedPhotosyntheticallyActiveRadiation(adapt(to, par.fields), 
-                                                                                                                    adapt(to, par.field_names))
+adapt_structure(to, par::PrescribedPhotosyntheticallyActiveRadiation) = 
+    PrescribedPhotosyntheticallyActiveRadiation(adapt(to, par.fields))