Learnable component - Nutrition modifier

[gmgivi]

In this PR, I plan to implement training of a simple learnable component using monthly depositions of sulphur and nitrogen.

The learnable components is implemented using the function poly_nm. poly_nm evaluates a cross-term polynomial over the selected inputs, then squashes it into a growth multiplier centered at 1:

$$\text{poly}(N, S) = \sum_i \sum_j w_{i,j} , N^i S^j$$

$$f_{\text{poly}} = 1 + \tanh(\text{poly}) \in (0, 2)$$

$f_{\text{poly}}$ scales 3PG's canopy quantum efficiency (alpha_c). Weights $w$ are fit by minimizing scaled squared error against observed/derived variables via Adam/SGD (train_nutrition_modifier), with plots comparing fitted vs. default 3PG output against observations.

@gmermoud @CedricTravelletti

[gmermoud]

question: why do we need this at all? I suppose we should rather always set defaults to the mode of the prior, no?

[gmgivi]

Defaults are set from literature values, not the prior's mode. They're already the best starting guess we have. The issue is that some literature defaults sit exactly at their prior's min/max, crashing NUTS at initialization. clip_defaults_to_priors just put the values at the boundary to just inside it, so it doesn't fail.

[gmermoud]

chore: explain what is the role of extended_params, deposition, etc. and type hint them.

[gmermoud]

question: I am familiar with the Einstein notation, but can you provide an example of what are operands and letters in the specific context of our case here?

[gmgivi]

If input_vars = ("N", "S") and poly_params.shape = (2, 2):

letters = "ab" # one letter per variable
powers = [N_powers, S_powers], each shape (..., 2)
operands = "...a,...b" , i.e. one term per power array, batch dims kept as ...
einsum string becomes "...a,...b,ab->..."

So the call is jnp.einsum("...a,...b,ab->...", N_powers, S_powers, poly_params),
which computes:

poly = Σ_a Σ_b N_powers[..., a] * S_powers[..., b] * poly_params[a, b]

matching the docstring's example.

I wanted to generalize this, as in case we add other metric later, for e.g. T_avg (average temperature)

With a third variable
(input_vars = ("N", "S", "T_avg"), poly_params.shape = (2,2,2)), it'd extend to
letters = "abc", operands = "...a,...b,...c",
einsum string "...a,...b,...c,abc->...".

poly = Σ_a Σ_b Σ_c N_powers[..., a] * S_powers[..., b] * T_avg_powers[..., b] * poly_params[a, b, c]

[gmermoud]

@gmgivi but N_powers contains powers of N, up to which degree?

[gmgivi]

For the moment, set to powers of 2.

[gmermoud]

Thanks for this @gmgivi. Can you clarify:

• Which target do we use for training? As discussed last week, we should soon use only DBH and initialize the stem biomass using allometric relations, and learn the breakdown into foliage/root biomass.
• Can you provide some motivation for the formulation as a cross-term polynomial? I am unclear about the exact form we use and why.
• Do we compare here two fitted models, or the default parameters wherein only the nutriment modifier has been fitted?

[gmgivi]

Thanks for this @gmgivi. Can you clarify:

This is just a first step towards learnable components (LC).

The aim here is to implement a small LC, run a training loop and see if it has some effect. From the plots, we could see that it's doing 'something'. Now, the question is
1. Whether this 'something' makes sense?
2. What should be the input variables that we should take for LC?
3. Whether polynomial function is a right choice?
4. Add other MLP components as LC?
5. Also add these LC with other optimization algorithms (Bayesian calibration, gradient descent and so on).

• Which target do we use for training? As discussed last week, we should soon use only DBH and initialize the stem biomass using allometric relations, and learn the breakdown into foliage/root biomass.

Yes, I will make these adjustments.

• Can you provide some motivation for the formulation as a cross-term polynomial? I am unclear about the exact form we use and why.

The cross polynomial terms takes the intra and inter-dependence of input variable, for the plot shown above we took nitrogen and sulphur depositions as input variables. We can add other terms, for e.g. other depositions or T_avg to it. I wanted to check the impact of these variables, if any. Now we see that there is some impact, and need to verify whether it makes sense.

From here we can progress towards introducing MLP. I wanted to get this approach working first.

• Do we compare here two fitted models, or the default parameters wherein only the nutriment modifier has been fitted?

We compare default with nutriment modifier fitted.

@gmermoud