Improve PISCES composability with additional layers of abstraction

[nanophyto]

@jagoosw I’m very interested in re-using some of the model components of PISCES (with a long term goal of building DARWIN-PISCES 😈), but the current setup makes this quite difficult. What are your thoughts on the suggested changes below to make the model more modular?

The idea would be to add two layers of abstraction to make the model more composable on both the sub-module (e.g. simple iron) and kernel (e.g. ligand_aggregation) side.

Sub-Module (e.g. Iron) contract layer.

The idea here would be to give each process a small concrete struct with parameter and tracer values, such that they can be reused independently without requiring a callback to the bgc PISCES struct.

Iron example:

# same as before 
@kwdef struct SimpleIron{FT} 
   excess_scavenging_enhancement :: FT = 1000.0 # unitless 
    maximum_ligand_concentration :: FT = 0.6    # μmol Fe / m³ 
          dissolved_ligand_ratio :: FT = 0.09   # μmol Fe / mmol C 
end 
 
const SimpleIronPISCES = PISCES{<:Any, <:Any, <:Any, <:Any, <:Any, <:SimpleIron} 
 
# new struct specific to iron, such that requirements are easily inferred: 
 
struct IronInputs{FT} 
   small_particle_iron_remineralisation::FT 
   grazing_waste::FT 
   upper_trophic_waste::FT 
   phytoplankton_iron_uptake::FT 
   ligand_aggregation_loss::FT 
   colloidal_aggregation::FT 
   scavenging::FT 
   bacterial_uptake::FT 
end 
 
# maths only depend on value inputs, not bgc or grid: 
 
@inline function iron_tendency(::SimpleIron, inputs::IronInputs) 
   return ( 
       inputs.small_particle_iron_remineralisation + 
       inputs.grazing_waste + 
       inputs.upper_trophic_waste - 
       inputs.phytoplankton_iron_uptake - 
       inputs.ligand_aggregation_loss - 
       inputs.colloidal_aggregation - 
       inputs.scavenging - 
       inputs.bacterial_uptake 
   ) 
  
# separate function looks up values:

@inline function iron_inputs(bgc::SimpleIronPISCES, i, j, k, grid, clock, fields, auxiliary_fields)
    Fe = @inbounds fields.Fe[i, j, k]
    DOC = @inbounds fields.DOC[i, j, k]
    T = @inbounds fields.T[i, j, k]

    scavenging_rate = iron_scavenging_rate(bgc.particulate_organic_matter, i, j, k, grid, bgc, clock, fields, auxiliary_fields)

    colloidal_aggregation, = aggregation_of_colloidal_iron(bgc.dissolved_organic_matter, i, j, k, grid, bgc, clock, fields, auxiliary_fields)

    scavenging = iron_scavenging(bgc.particulate_organic_matter, i, j, k, grid, bgc, clock, fields, auxiliary_fields)

    bacterial_uptake = bacterial_iron_uptake(bgc.particulate_organic_matter, i, j, k, grid, bgc, clock, fields, auxiliary_fields)

    small_particle_iron_remineralisation = degradation(bgc.particulate_organic_matter, Val(:SFe), i, j, k, grid, bgc, clock, fields, auxiliary_fields)

    phytoplankton_iron_uptake = uptake(bgc.phytoplankton, Val(:Fe), i, j, k, grid, bgc, clock, fields, auxiliary_fields)

    grazing_waste = non_assimilated_iron(bgc.zooplankton, i, j, k, grid, bgc, clock, fields, auxiliary_fields)

    upper_trophic_waste = upper_trophic_dissolved_iron(bgc.zooplankton, i, j, k, grid, bgc, clock, fields, auxiliary_fields)

    ligand_aggregation_loss = ligand_aggregation(bgc.iron, Fe, DOC, T, scavenging_rate)

    return IronInputs(
        small_particle_iron_remineralisation,
        grazing_waste,
        upper_trophic_waste,
        phytoplankton_iron_uptake,
        ligand_aggregation_loss,
        colloidal_aggregation,
        scavenging,
        bacterial_uptake,
    )
end
    
# this is now a light wrapper with value lookup: 
 
@inline function (bgc::SimpleIronPISCES)(i, j, k, grid, ::Val{:Fe}, clock, fields, auxiliary_fields) 
   inputs = iron_inputs(bgc, i, j, k, grid, clock, fields, auxiliary_fields) 
 
   return iron_tendency(bgc.iron, inputs) 
end 
 

Pure kernel layer

For this layer the idea would be to split scalar maths from value look-ups, such that the math does not depend on the grid or bgc objects. This would make each component also re-usable:

For example, split:

# scalar maths plus value lookup: 
 
@inline function aggregation_of_colloidal_iron(dom::DissolvedOrganicCarbon, i, j, k, grid, bgc, clock, fields, auxiliary_fields) 
   _, Φ₁, Φ₂, Φ₃ = aggregation(dom, i, j, k, grid, bgc, clock, fields, auxiliary_fields) 
 
   Fe = @inbounds fields.Fe[i, j, k] 
   DOC = @inbounds fields.DOC[i, j, k] 
 
   Fe′ = free_iron(bgc.iron, i, j, k, grid, bgc, clock, fields, auxiliary_fields) 
   ligand_iron = Fe - Fe′ 
   colloidal_iron = 0.5 * ligand_iron 
 
   CgFe1 = (Φ₁ + Φ₃) * colloidal_iron / (DOC + eps(0.0)) 
   CgFe2 = Φ₂ * colloidal_iron / (DOC + eps(0.0)) 
 
   return CgFe1 + CgFe2, CgFe1, CgFe2 
end 

Into:

# scalar maths: 
@inline function aggregation_of_colloidal_iron(Φ₁, Φ₂, Φ₃, Fe, Fe′, DOC) 
    
   colloidal_iron = 0.5 * (Fe - Fe′) 
   CgFe1 = (Φ₁ + Φ₃) * colloidal_iron / (DOC + eps(0.0)) 
   CgFe2 = Φ₂ * colloidal_iron / (DOC + eps(0.0)) 
 
   return CgFe1 + CgFe2, CgFe1, CgFe2 
end 
 
# value lookup: 
@inline function aggregation_of_colloidal_iron(dom::DissolvedOrganicCarbon, i, j, k, grid, bgc, clock, fields, auxiliary_fields) 
   _, Φ₁, Φ₂, Φ₃ = aggregation(dom, i, j, k, grid, bgc, clock, fields, auxiliary_fields) 
 
   Fe = @inbounds fields.Fe[i, j, k] 
   DOC = @inbounds fields.DOC[i, j, k] 
 
   Fe′ = free_iron(bgc.iron, i, j, k, grid, bgc, clock, fields, auxiliary_fields) 
 
   return aggregation_of_colloidal_iron(Φ₁, Φ₂, Φ₃, Fe, Fe′, DOC) 
end 

Let me know what you think, I’ve started on a proof of concept fork for iron already but would be good to hear your thoughts:

• Do you think submodule re-use would be valuable?
• Do you think re-usable kernel functions would be valuable?
• I’ve started with Iron on a OceanBioME fork, if you like this idea, would you prefer one PR that implements changes for all the components in PISCES, or one component at a time?
• I named the iron sub-module struct *Inputs, is that syntax you are happy with, or should I name it something else? (e.g. IronInputs, PhosphateInputs, NanoPhytoGrowthInputs, GrazingInputs)

A downside of this approach would be additional abstraction and code surface, so something worth mentioning...

[johnryantaylor]

Hi Joost, thank you very much for this idea. In general, it would be great to improve the composability. @jagoosw and I are starting to think about how to implement MARBL in OceanBioME, and it would be good to use the same approach there too (and re-use functionality where possible). It would be good to use the same approach for both of the full complexity models (PISCES and MARBL). We could then think about how/whether to incorporate this into the low and medium complexity models (NPZD and LOBSTER).

A small point: I don't quite understand the use of "input" in this context. To me input sounds like something that the user needs to specify, whereas I think here they are the individual terms in the iron budget (correct me if I'm wrong). What about using "rate" or "source" instead?

[nanophyto]

yes, I agree "Input" is not the best naming convention. Jago suggested using "args" or "arguments" which would work well. Or we could drop this component entirely since it's not strictly necessary (see response to Jago below)

[nanophyto]

And to your other point, I agree it would make sense to re-use KFs between models where appropriate, although then some more re-organization of the full repo would be required. Perhaps a central folder called something like "Library" (this is what we do in Agate.jl) would make sense?

[jagoosw]

Hi Joost,

This seems like a really interesting idea and I think working out some kind of abstraction like this would be really helpful.

I'm not sure I understand why we need the type IronInputs? Is the idea so we could dispatch on it like iron_tendency(::SimpleIron, inputs::IronInputs{Nothing, <:Any, ...}) (which would require changing the types of the inputs), otherwise can't we just pass the arguments as arguments like:

@inline iron_tendency(::SimpleIron, 
                      small_particle_iron_remineralisation,
                      grazing_waste,
                      upper_trophic_waste,
                      phytoplankton_iron_uptake,
                      ligand_aggregation_loss,
                      colloidal_aggregation,
                      scavenging,
                      bacterial_uptake) =
   (small_particle_iron_remineralisation + 
    grazing_waste + 
    upper_trophic_waste - 
    phytoplankton_iron_uptake - 
    ligand_aggregation_loss - 
    colloidal_aggregation - 
    scavenging - 
    bacterial_uptake) 
  
# separate function looks up values:

@inline function iron_inputs(bgc::SimpleIronPISCES, i, j, k, grid, clock, fields, auxiliary_fields)
...

    return (small_particle_iron_remineralisation,
            grazing_waste,
            upper_trophic_waste,
            phytoplankton_iron_uptake,
            ligand_aggregation_loss,
            colloidal_aggregation,
            scavenging,
            bacterial_uptake)
end

Maybe I'm missing something?

To John's comment: if I'm understanding correctly perhaps a better name would be arguments or args not inputs

[nanophyto]

Yes this is a good question. The only reason to include a *Inputs struct would be to reduce the function signature surface in instances that the number of parameters and functions is very large. But you are right this adds abstraction that is not strictly necessary.

[jagoosw]

• Do you think submodule re-use would be valuable?
• Do you think re-usable kernel functions would be valuable?
• I’ve started with Iron on a OceanBioME fork, if you like this idea, would you prefer one PR that implements changes for all the components in PISCES, or one component at a time?
• I named the iron sub-module struct *Inputs, is that syntax you are happy with, or should I name it something else? (e.g. IronInputs, PhosphateInputs, NanoPhytoGrowthInputs, GrazingInputs)

Anyway, I think this is a good idea and my origional attempt to abstract PISCES wasn't the most thought through. It seems like it would make sense to abstract away some of the getting values etc. by having reusable KFs, but I'm not sure exactly how todo that (and on originally writing PISCES and LOBSTER as "continuous" models where we just get the tracer values, I found it a lot harder to abstract, although possibly the XInputs idea solves this problem).

If its possible todo piecewise maybe it is the best approach todo this a component at a time and we can use PISCES as a testing ground since no one really uses it.

[nanophyto]

OK sounds good, let's do it piece-wise then, the surface of the model wouldn't change, so it is more so about internal consistency.

[nanophyto]

Hi both,

Thanks for your enthusiasm and comment already. It sounds like the way forward is to

1. agree if:

• An e.g. *args struct would add value, or is unnecessary abstraction
• If so, what to name it *args, *arguments, or something else

2. Prepare a PR that focused on iron. In which we can further discuss syntax.

3. Prepare PRs for the other PISCES components.

4. Once PISCES is more modular consider a wider repo restructure such that components are shared between PISCES, NPZD, LOBSTER and eventually MARBL (e.g. a central \Libary)

[nanophyto]

I started a bit on an self contained iron kernel which looked something like this:

# self contained kernel that contains all scalar maths:
inline function iron_tendency(iron, dom, pom, Fe, DOC, T, [...])
	ligand_aggregation_loss = ligand_aggregation(iron, Fe, DOC, T, scavenging_rate)
	
	[...]
	
    return (small_particle_iron_remineralisation + grazing_waste + upper_trophic_waste -
            phytoplankton_iron_uptake - ligand_aggregation_loss - colloidal_aggregation -
            scavenging - bacterial_uptake)

# thin wrapper with relevant value lookup
@inline function iron_tendency(bgc::SimpleIronPISCES, i, j, k, grid, clock, fields, auxiliary_fields)
    Fe = @inbounds fields.Fe[i, j, k]
    DOC = @inbounds fields.DOC[i, j, k]
    T = @inbounds fields.T[i, j, k]
    [...]
    
  
    return iron_tendency(bgc.iron,
        bgc.particulate_organic_matter,
        bgc.dissolved_organic_matter,
        Fe, DOC, T, [...])
    
    
# still keep existing high level call-back:
@inline function (bgc::SimpleIronPISCES)(i, j, k, grid, ::Val{:Fe}, clock, fields, auxiliary_fields)
    return iron_tendency(bgc, i, j, k, grid, clock, fields, auxiliary_fields)
end

However, while this is great for re-use since the kernel is self contained, for processes other than iron this would lead to a lot of repeated calculations...

This particularly true for the plankton side where different modules re-use things like prey availability, grazing pressure, and growth efficiency.

I'm thus now thinking about effectively splitting rates and the tendency. This is less relevant for iron, but for plankton it could look be something like this:

# single scalar function which is used to estimate all rates once:

@inline function phytoplankton_rates(
    phyto,
    zoo,
    val_name,
    I,
    IChl,
    IFe,
    T,
    PAR,
    [...]
)
    production = total_production(
        phyto, val_name,
        I, IChl, IFe, T, PAR, [...]
    )

    carbon_growth = [...]
    chlorophyll_growth = [...]
    iron_growth = [...]
    silicate_growth = [...]
    linear_mortality = [...]
    quadratic_mortality = [...]
    grazing_loss = [...]

    # return named tuple, tuple or struct
    return (  
        production
        carbon_growth
        chlorophyll_growth 
	[...]
    )
end

# thin wrapper used for value lookup
@inline function phytoplankton_rates(
    bgc,
    phyto,
    val_name,
    i, j, k, grid, clock, fields, auxiliary_fields,
)
    I    = phytoplankton_concentration(val_name, i, j, k, fields)
    IChl = chlorophyll_concentration(val_name, i, j, k, fields)
    IFe  = iron_quota_concentration(val_name, i, j, k, fields)

    T   = @inbounds fields.T[i, j, k]
    PAR = @inbounds auxiliary_fields.PAR[i, j, k]
    [...]

    return phytoplankton_rates(
        phyto,
        bgc.zooplankton,
        val_name,
        I,
        IChl,
        IFe,
        T,
        PAR,
        [...]
    )
end

# now a set of tendency functions take in the already computed rates:
@inline function carbon_tendency(rates)
    return (
        rates.carbon_growth -
        rates.linear_mortality -
        rates.quadratic_mortality -
        rates.grazing_loss
    )
end

# more complicated scalar maths for e.g. chlorophyll:
@inline function chlorophyll_tendency(rates)
    θChl = rates.IChl / (12 * rates.I + eps(0.0))

    return (
        rates.chlorophyll_growth -
        (rates.linear_mortality + rates.quadratic_mortality + rates.grazing_loss) * θChl * 12
    )
end

@inline function iron_tendency(rates)
    [...]
end

[...] # any other relevant tendencies such as iron, silicate etc.

# putting it together:

phyto = parameterisation(Val(:P), bgc.phytoplankton)

# only estimate rates once
phyto_rates = phytoplankton_rates(
    bgc,
    phyto,
    Val(:P),
    i, j, k, grid, clock, fields, auxiliary_fields,
)

# now re-use rates for tendencies
P    = carbon_tendency(phyto_rates)
PChl = chlorophyll_tendency(phyto_rates)
PFe  = iron_tendency(phyto_rates)

More concrete (diatom) example:

# current syntax:

@inline function (bgc::PISCES{<:NanoAndDiatoms})(i, j, k, grid, val_name::Val{:DSi}, clock, fields, auxiliary_fields)
    phyto = parameterisation(val_name, bgc.phytoplankton)
    
    growth = silicate_growth(phyto, val_name, i, j, k, grid, bgc, clock, fields, auxiliary_fields)

    D = @inbounds fields.D[i, j, k]
    DSi = @inbounds fields.DSi[i, j, k]

    θSi = DSi / (D + eps(0.0))

    linear_mortality, quadratic_mortality = mortality(phyto, val_name, i, j, k, grid, bgc, clock, fields, auxiliary_fields)

    death = (linear_mortality + quadratic_mortality)

    getting_grazed = grazing(bgc.zooplankton, Val(:D), i, j, k, grid, bgc, clock, fields, auxiliary_fields) 

    return growth - (death + getting_grazed) * θSi
end

# new syntax:

@inline function (bgc::PISCES{<:NanoAndDiatoms})(i, j, k, grid, val_name::Val{:DSi}, clock, fields, auxiliary_fields)
    phyto = parameterisation(val_name, bgc.phytoplankton)

    phyto_rates = phytoplankton_rates(
        bgc,
        phyto,
        val_name,
        i, j, k, grid, clock, fields, auxiliary_fields,
    )

    D = @inbounds fields.D[i, j, k]
    DSi = @inbounds fields.DSi[i, j, k]

    return silicate_tendency(phyto_rates, D, DSi)
end


@inline function silicate_tendency(phyto_rates, D, DSi)
    θSi = DSi / (D + eps(0.0))
    death = phyto_rates.linear_mortality + phyto_rates.quadratic_mortality

    return phyto_rates.silicate_growth - (death + phyto_rates.getting_grazed) * θSi
end

Mostly thinking out loud at this point.

[nanophyto]

Actually, looking at the call-backs, the rates are still re-computed every time, which seems kind of unavoidable with how it is currently setup? A high-level call-back or caching might need be necessary, which would be a larger/different refactor.

So perhaps the suggestion is not needed at this point, but could be a step in the right direction at least?

[jagoosw]

Do you mean that to prevent recalculation of rates etc. used for different tracers (e.g. D/DChl/DSi) that we compute them once per time step instead?

That would be quite a significant refactor and I think we'd need to workout how much peformance it gains (given that the biggest cost is tracer advection). We've previously thought of recomputing things like this as a trade off between computation and memory (because if we computed all the different rates and shared them between the tracers we'd have to store them), but I could see a different paridime where we e.g. compute grazing terms and apply them to every tendency field they interact with, then compute every growth term, etc.

[nanophyto]

Yes, but good to know you have already considered/discussed this. In that case I will ignore it for now in that case and keep the PR simple.