Vectorize leaf dynamics

[fnattino]

Closes #19

[fnattino]

@SarahAlidoost after the discussion we had today, I realized it was very little changes to add the functionality to work with arbitrarily-shaped tensors as parameters, so I have added it. In the current form the leaf_dynamics should be able to work with parameter tensors that have arbitrary shapes, provided that the shape is the same for all parameters (or one parameter has an arbitrary shape and all other parameters are 0-dimensional).

At the current stage, we should be able to reproduce simulations with scalar parameters (all tests pass), I am now working to add some tests on the use of parameter tensors!

[SarahAlidoost]

@fnattino PR #35 is merged. There, I moved test helper functions to utils. So, you may want to merge main into your branch.

[fnattino]

@SarahAlidoost, I left one question below, but this is otherwise ready to be reviewed.

I have added few tests that show that the module can now work with arbitrarily shaped arrays as parameters (provided all parameters have the same shape). There is quite some repeated code in the tests, happy to hear whether you have suggestions for improvements there. Otherwise, I think that things will simplify significantly when we will modify the engine class and make it easier to pass input and get output.

[SarahAlidoost]

@fnattino great job, thanks 👍 The tests of class TestLeafDynamics were passedd with test data test_leafdynamics_wofost72_17.yaml too, well done!

But those tests only check the correctness of the model and not the gradient graphs. The classes TestDiffLeafDynamicsTDWI and TestDiffLeafDynamicsSPAN check the gradient computation graph, but right now they don’t have any tests for the vector case. For example, similar to test_gradients_tdwi_lai_leaf_dynamics, we should add a test where
tdwi = torch.nn.Parameter(torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32)). So each class could have tests such as:

• test_gradients_tdwi_lai_leaf_dynamics_vector
• test_gradients_tdwi_twlv_leaf_dynamics_vector

I added an example. Can you please add those tests? if something is unclear, let me know. Thanks.

[fnattino]

@SarahAlidoost I should have addressed your comments above. When adding tests, I have slightly modified the structure of tests for leaf dynamics to be able to run them with parameters, so it would be easier to add more of these tests.

Now, one of the thing I did to make this implementation easier, is to have DiffLeafDynamics return a dictionary of tensors instead of stacking all the tensors in a larger tensor. I have then ported the same structure to the tests on root dynamics, which unveiled an issue there: the RD variable seems not to have the gradient defined (see failing tests, e.g. here). Earlier, this issue was not caught by the tests, because when stacking RD and TWRT in a single tensor, a gradient would somehow become available - but now this seems to have emerged. Do you agree with the assesment?

[SarahAlidoost]

@SarahAlidoost I should have addressed your comments above. When adding tests, I have slightly modified the structure of tests for leaf dynamics to be able to run them with parameters, so it would be easier to add more of these tests.

Thanks for refactoring the tests using pytest.mark.parametrize looks great! Right now you merged all the parameters in one class. In future, more parameters will be added to the list. This makes debugging more difficult. If we could keep one class per each parameter that would be great. It is okay if we repeat the codes a bit in tests. Another thing is that before there was an assert statement for gradient of TWLV wrt SPAN where assert numerical_grad == 0.0. This will show that changes in this parameter doesnot change the output. see issue #26. Can you please add this assert back to the tests of class SPAN?

Now, one of the thing I did to make this implementation easier, is to have DiffLeafDynamics return a dictionary of tensors instead of stacking all the tensors in a larger tensor. I have then ported the same structure to the tests on root dynamics, which unveiled an issue there: the RD variable seems not to have the gradient defined (see failing tests, e.g. here). Earlier, this issue was not caught by the tests, because when stacking RD and TWRT in a single tensor, a gradient would somehow become available - but now this seems to have emerged. Do you agree with the assesment?

That's correct. Actually variable TDWI doesnot contribute to RD mathematically, I explained this in the section 2.2 of the notebook. So there should not be any gradient function. I just noticed that pytorch returns runtime error. Two tests test_gradients_tdwi_rd_root_dynamics and test_gradients_tdwi_rd_root_dynamics_numerical are not correct and should be replaced with one tests where we assert loss.grad_fn is None. like:

    def test_gradients_tdwi_rd_root_dynamics(self):
        model = get_test_diff_root_model()
        tdwi = torch.nn.Parameter(torch.tensor(0.2, dtype=torch.float32))
        output = model({"TDWI": tdwi})
        rd = output["RD"]
        loss = rd.sum()
        assert loss.grad_fn is None # tdwi doesnot contribute to rd

[fnattino]

Thanks for all the clarifications (and for bearing with me in this PR :) ). I have followed your suggestion and:

• "unpacked" some of the parametrizations in the tests, putting back one class per parameter as it was originally set up;
• put back the check on the gradient being zero for TWLV wrt SPAN. Note, however, that I have setup the comparison on the gradient calculated from autograd (not the numerical one), since the numerical gradient turned out to be small but not exacly zero (~1.e-5). Not sure why this pops up only now.
• fixed the root dynamics tests to correct the behaviour for checking on RD wrt TDWI.

[SarahAlidoost]

@fnattino looks great! Thank you for implementing the vectorization and addressing the comments and suggestions. Ready to merge with only one small change, could you please remove the test test_gradients_tdwi_rd_root_dynamics_numerical before merging? 🚀

[fnattino]

Done @SarahAlidoost - the lasts checks are running, feel free to merge when you think it's ready!

[sonarqubecloud]

Quality Gate passed

Issues
5 New issues
0 Accepted issues

Measures
0 Security Hotspots
97.6% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud