Iron proof of concept (scalar module)

[nanophyto]

@jagoosw @johnryantaylor

Following our initial discussion (#343), I've started on a PISCES iron proof of concept refactor. Note that is just the first step of the iron refactor, but due to the complexity a good point to touch base.

I've gone through the dissolved iron module (and it's dependencies) and made everything scalar, such that there is no value look-ups inside the module. In addition, I've made sure that the module does not depend on PISCES-specific structs such as dom, pom, zoo etc.

Instead, there are now two functions:

# purely scalar form:
@inline function iron_tendency(iron::SimpleIron, args::IronTendencyArgs)

    small_particle_iron_remineralisation = degradation(SFe, [...])
    grazing_waste = non_assimilated_iron(args, [...])
                                                       
    return (small_particle_iron_remineralisation + grazing_waste + [...])
end

# value look-up form:
@inline function iron_tendency(bgc::SimpleIronPISCES, i, j, k, grid, clock, fields, auxiliary_fields)
    Fe = @inbounds fields.Fe[i, j, k]
    [...]
    
    z = znode(i, j, k, grid, Center(), Center(), Center())
    [...]
    
    sinking_flux = edible_flux_rate(bgc.particulate_organic_matter, i, j, k, grid, fields, auxiliary_fields)

    nano = bgc.phytoplankton.nano
    
    args = IronTendencyArgs(
        Fe,
        [...],
        z,
        [...],
        sinking_flux,
        [...],
        nano.size_ratio,
        [...]
        )
        
    return iron_tendency(bgc.iron, args)
  

For implementation, I have added "scalar overloads" for all the relevant functions (e.g. non_assimilated_iron), and removed the "grid lookup" forms as appopriate. Since some functions are not iron specific, they now have both a scalar and grid-lookup form. The idea is to remove the latter forms, once all the PISCES module have been updated.

I've also included an IronTendencyArgs struct which holds all the args needed for the scalar iron_tendency(). This is optional, but tidier since this arg list is rather long. In this context, I have also used long-form arguments instead of symbols to keep the struct more readable/generalizable.

What is currently missing is that the model still hard codes e.g. microzooplankton(P), mesozooplankton (M), nanophytoplankton (P), and diatoms (D). It is thus not currently composable with e.g. a MARBL model that has 8P-4Z, or DARWIN type models... This next step still requires some more discussion and thought which I will post in (#343), but I'm aware the diff file is already quite large, so it might be worth finalizing this PR before moving on to the next step?

To keep the main branch tidy, I started a new branch called PISCES-refactor while we are still figuring out the final model structure.

Reminder to myself:

• restore full test once PR is ready

[jagoosw]

Hi @nanophyto,

Thank you for sharing this progress. I've started going through to try and understand and will review once I've understood more.

I had a few questions to start with:

• Is an advantage to having the IronTendencyArgs vs just passing a NamedTuple? I don't think you ever dispatch on ::IronTendencyArgs right?
• Am I understanding correctly that the idea is take parameterisations, for example in the small particle Fe tendency we have aggregation_of_colloidal_iron(i, j, k, ...), and instead of passing the discrete form in straight to the parameterisation, in the top level tendency function we get all the scalar values and pass e.g. aggregation_of_colloidal_iron(iron_model, iron_arguments)?
  If I am understanding this correctly then I wonder if instead of taking the discrete form of the top level tendency function e.g. (bgc::IronPISCES)(i, j, k, ...) putting that in a function iron_tendency(i, j, k, ...) then turning that into a scalar form iron_tendency(model, arguments) then calling all of the parameterisations, if it would be more clear and flexible to have (bgc::IronPISCES)(i, j, k, ...) directly call the scalar form of the parameterisations and then put them all together. Something like:

function (bgc::IronPISCES)(i, j, k, ...)
    iron_args = ...
    
    small_particle_reminersation = small_particle_reminersation(bgc, iron_args) # perhaps just a subset of the relivant arguments?
    aggregation_of_colloidal_iron = aggregation_of_colloidal_iron(bgc, iron_args)

    ...

    return small_particle_reminiersation - aggregation_of_colloidal_iron + ...
end

When I have written the models in the past I have found that with scalar forms it can be very hard to put different versions together because you have the big long lists of arguments. If I understand right then I think we would achieve the goal of being able to reuse the parameterisations by still writing discrete form top level tendencies that directly use continuous/scalar form parameterisations, because then in a new model combination we would just have to pull out the correct values/parameters and call the parameterisation.

[jagoosw]

In summary I think I'm not sure why we need the intermediate scalar form iron_tendency

[nanophyto]

@jagoosw

Both great questions. Yes, IronTendencyArgs could definitely just be a NamedTuple...

You are right that having an scalar iron_tendency does not add a lot at the moment and in your example could be avoided. The main purpose at this point is to make the maths (remin - agg - scav - uptake + ...) re-usable. Assembling it for each model is trivial if there are just two models (MARBL and PISCES), but ideally I would like to re-use it in Agate.jl as well, instead of rewriting the scalar maths every time. Having a separate function also allows unit tests for just the iron component, although this is of course optional.

But to make things more concrete, the eventual vision (and point of scalar iron_tendency) is something like this:

iron_tendency(iron, common_args, biological_terms)

# where:

iron = (
    ligand = (
        excess_scavenging_enhancement = ...,
        maximum_ligand_concentration = ...,
        dissolved_ligand_ratio = ...,
    ),

    oxygen = (
        first_anoxia_threshold = ...,
        second_anoxia_threshold = ...,
    ),
    
    [...]
)


common_args = (
    Fe = ...,
    DOC = ...,
    
    [...]
    
    z = ...,
    mixed_layer_shear = ...,
    sinking_iron_flux = ...,
    
    [...]
)


biological_terms = (
    phytoplankton_iron_uptake = ...,
    grazing_waste = ...,
    upper_trophic_waste = ...,
    bacterial_uptake = ...,
)


# PISCES version:

function biological_terms(bgc::SimpleIronPISCES, i, j, k, grid, clock, fields, auxiliary_fields)
        
    phytoplankton_iron_uptake = ...  # existing nano + diatom logic
    grazing_waste = ...  # existing micro + meso logic
    upper_trophic_waste = ...
    bacterial_uptake = ...

	return (
		phytoplankton_iron_uptake = phytoplankton_iron_uptake,
		grazing_waste = grazing_waste,
		upper_trophic_waste = upper_trophic_waste,
		bacterial_uptake = bacterial_uptake,
	)
end
    
# Agate version:    
    
function biological_terms(bgc::AgateIronModel, i, j, k, grid, clock, fields, auxiliary_fields)
        
    phytoplankton_iron_uptake = sum(agate_phyto_iron_uptake(p) for p in 1:nP)
    grazing_waste  = sum(agate_grazer_iron_waste(z) for z in 1:nZ)
    upper_trophic_waste = sum(agate_upper_trophic_iron_waste(z) for z in 1:nZ)
    bacterial_uptake = agate_bacterial_iron_uptake(...)

	return (
		phytoplankton_iron_uptake = phytoplankton_iron_uptake,
		grazing_waste = grazing_waste,
		upper_trophic_waste = upper_trophic_waste,
		bacterial_uptake = bacterial_uptake,
	)
end

# unsure about MARBL but likely with similar loops over e.g. 4P and 4Z?

Hopefully, a scalar version in that context makes more sense?

[jagoosw]

@nanophyto

Sorry it took me a while to reply to this. I'm not sure I understand the new design still and why going iron_tendency(i, j, ...) to iron_tendency(iron, bunch_of_scalars) helps modularity because it locks in what values the scalars have to have? And also locks in the form of the tendencies (e.g. aggregation now has to take particular arguments). I'm also a bit concerned about how maintainable it is and how much extra code it takes.

I think we can achieve the end goal of modularity like your describing by maintaining something more similar to the existing code. Starting with the current version:

OceanBioME.jl/src/Models/AdvectedPopulations/PISCES/iron/simple_iron.jl

Lines 19 to 53 in bee31b4

	@inline function (bgc::SimpleIronPISCES)(i, j, k, grid, val_name::Val{:Fe}, clock, fields, auxiliary_fields)
	λ̄ = bgc.iron.excess_scavenging_enhancement
	Fe = @inbounds fields.Fe[i, j, k]
	λFe = iron_scavenging_rate(bgc.particulate_organic_matter, i, j, k, grid, bgc, clock, fields, auxiliary_fields)
	Fe′ = free_iron(bgc.iron, i, j, k, grid, bgc, clock, fields, auxiliary_fields)
	total_ligand_concentration = ligand_concentration(bgc.iron, i, j, k, grid, bgc, clock, fields, auxiliary_fields)
	# terminal process which removes iron from the ocean
	ligand_aggregation = λ̄ * λFe * max(0, Fe - total_ligand_concentration) * Fe′
	colloidal_aggregation, = aggregation_of_colloidal_iron(bgc.dissolved_organic_matter, i, j, k, grid, bgc, clock, fields, auxiliary_fields)
	# scavenging and bacterial uptake
	scavenging = iron_scavenging(bgc.particulate_organic_matter, i, j, k, grid, bgc, clock, fields, auxiliary_fields)
	BactFe = bacterial_iron_uptake(bgc.particulate_organic_matter, i, j, k, grid, bgc, clock, fields, auxiliary_fields)
	# particle breakdown
	small_particles = degradation(bgc.particulate_organic_matter, Val(:SFe), i, j, k, grid, bgc, clock, fields, auxiliary_fields)
	# consumption
	consumption = uptake(bgc.phytoplankton, val_name, i, j, k, grid, bgc, clock, fields, auxiliary_fields)
	# waste
	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)
	return (small_particles + grazing_waste + upper_trophic_waste
	- consumption - ligand_aggregation - colloidal_aggregation - scavenging - BactFe)
	end

if we were to put this somewhere with more abstraction like I've made in #341 (BiologyNutrientDetritus) then we can have something like (bgc::BiologyNutrientsDetritus{<:Any, NutrientsWithIron})(i, j, k... where NutrientsWithIron points to some nutrient model that specifies that it uses a PISCES like iron scheme.

Then we can abstract the function like:

@inline function (bgc::SimpleIronPISCES)(i, j, k, grid, val_name::Val{:Fe}, clock, fields, auxiliary_fields) 
     λ̄ = bgc.iron.excess_scavenging_enhancement 
  
     Fe = @inbounds fields.Fe[i, j, k] 
  
     λFe = iron_scavenging_rate(bgc.particulate_organic_matter, i, j, k, grid, bgc, clock, fields, auxiliary_fields) 
      
     Fe′ = free_iron(bgc.iron, i, j, k, grid, bgc, clock, fields, auxiliary_fields) 
  
     total_ligand_concentration = ligand_concentration(bgc.iron, i, j, k, grid, bgc, clock, fields, auxiliary_fields) 
  
     # terminal process which removes iron from the ocean 
     ligand_aggregation = λ̄ * λFe * max(0, Fe - total_ligand_concentration) * Fe′ 
  
     biological_uptake = iron_uptake(bgc, i, j, k, ...
     biological_waste = biology_inorganic_iron_loss(bgc, i, j, k, ...
     detritus_waste = detritus_inorganic_iron_loss(bgc, i, j, k, ...
  
     return (biological_waste - biological_uptake - ligand_aggregation) 
end 

where the biology tendencies belong to the biology, so without changing the rest of the PISCES code:

iron_uptake(bgc::BiologyNutrientDetritus{<:PISCESBiology}, i, j, k, ... ) = uptake(bgc.phytoplankton, Val(:Fe), i, j, k, grid, bgc, clock, fields, auxiliary_fields) 

(probably this isn't the best and we should do some refactoring to the PISCES code to maybe just define iron_uptake directly instead of uptake update(bgc, ::Val{:Fe}...)

Then the waste terms are:

biology_inorganic_iron_loss(bgc::BiologyNutrientDetritus{<:PISCESBiology}, i, j, ...) = 
    non_assimilated_iron(bgc.zooplankton, i, j, k, grid, bgc, clock, fields, auxiliary_fields) +
    upper_trophic_dissolved_iron(bgc.zooplankton, i, j, k, grid, bgc, clock, fields, auxiliary_fields) 

and

detritus_inorganic_iron_loss(bgc::BiologyNutrientDetritus{<:Any, <:Any, <:PISCESDetritus}, i, k, ...) = 
    aggregation_of_colloidal_iron(bgc.dissolved_organic_matter, i, j, k, grid, bgc, clock, fields, auxiliary_fields) +
    iron_scavenging(bgc.particulate_organic_matter, i, j, k, grid, bgc, clock, fields, auxiliary_fields) +
    bacterial_iron_uptake(bgc.particulate_organic_matter, i, j, k, grid, bgc, clock, fields, auxiliary_fields) +
    degradation(bgc.particulate_organic_matter, Val(:SFe), i, j, k, grid, bgc, clock, fields, auxiliary_fields) 

(again we probably want to refactor PISCES, and if we were to put PISCES into the BND format we would need to do e.g. bgc.biology.zooplankton).

Agate could then define:

iron_uptake(bgc::BiologyNutrientDetritus{<:AgatePZ}, i, j, ...) =
    sum(...

etc.

I don't know that this is the best abstraction but I think something like this were we encode the minimum constraints while keeping it as simple as possible would be better. Again, I'm not tied to this idea but maybe a good step would be to take the iron formulation from PISCES and port it into the format in #341? Then we could work on interfacing Agate PZ into that format too so it can interchange different nutrient and detritus formulations? For example, as that PR is now we could define it as PISCESNutrients (and NO3 and NH4 are really simple because the only special thing about PISCES' version I think is the oxygen dependent nitrification), then we could use the simple LOBSTER biology and default detritus by defining:

iron_uptake(bgc::BiologyNutrientDetritus{<:PhytoZoo}, args...) = nutrient_uptake(bgc, args...)

biology_inorganic_iron_loss(bgc::...., args...) = bgc.biology.iron_ration * biology_inorganic_nitrogen_waste(bgc, args...)

detritus_inorganic_iron_loss(bgc::BiologyNutrientDetritus{<:Any, <:Any, <:TwoParticleAndDissolved}, args...) = bgc.biology.iron_ratio * detritus_inorganic_nitrogen_waste(bgc, args...)

since its fixed Redfield ratio but can be limited by iron (would probably need to change a couple of other things so that PhytoZoo knows that PISCESNutrients has iron/nitrate/ammonia but hopefully thats minimal changes.