Make Root module fully differentiable

[SCiarella]

This PR addresses the differentiability requirement of issue #47.

All the parameters except RDRRTB were already differentiable since we have defined them as torch.tensor.
To make RDRRTB differentiable, I had to redefine pcse.utils.Afgen to use tensors instead of list. For this reason I have added src/physical_models/crop/afgen.py which is now differentiable.

I have also restructured test_root_dynamics such that all the combinations of parameter-output are checked, and code duplication is reduced.

---

As a quality of life improvement, I have installed pre-commit that automatically runs ruff checks before each git commit.

[SCiarella]

For the efficiency requirement of issue #47, I believe that the current state already implements vectorized operation, so I think that #47 can be closed once this PR is finalized.

[SarahAlidoost]

@SCiarella thanks. The changes looks good. The PR needs a few modifications specially for the vecorization which its tests currently missing. Please see my comments, if something unclear, let me know.

[SCiarella]

👋 @SarahAlidoost Thanks for the careful review!

I have implemented all the suggestions, except for a few points that remain open:

• [1] I have cleaned the test_root_dynamics as you suggested, but I have kept the single-testclass structure because I think it is easier to extend and maintain and it is not hard to interpret, but let me know what you think
• [2] I do not fully understand the vectorization requirement. To me it is already vectorized (parameter can be passed as tensors). Can you give me a practical example?
• [3] should we ask someone to confirm that the gradient TWRT->RDRRTB is numerically 0?

[SCiarella]

👋 @SarahAlidoost

With the last commit I have implemented the vectorization for the root dynamics.
I have used _get_params_shape and _broadcast_to which I have now moved from the leaf module to utils.py.

Finally, I have copied all the vectorization tests that we use for the leaf dynamics and implemented them in test_root_dynamics.py.
Let me know what you think! 💪

[SarahAlidoost]

@SCiarella thanks for addressing the comments 👍 We need a few more tests to make sure everything works with vectorization:

1. Tests for gradients when parameters are vectors, see my code suggestions in class TestDiffRootDynamicsGradients.
2. Tests for parameter "RDRRTB", see the list in test_root_dynamics_with_one_parameter_vector.
3. Tests for other parameters in test_root_dynamics_with_multiple_parameter_vectors, and test_root_dynamics_with_multiple_parameter_arrays. Currently, only two parameters ("RDI", "RRI") are there.

Please have a look and let me know if something is unclear. Thanks.

[SCiarella]

Thanks @SarahAlidoost!

Now I have fully vectorized also RDRRTB by extending the Afgen class to work with tensors of tables. I have extended the unit tests in test_utils to also cover the tensor case and I have expanded the test_root_dynamics to cover the case when all the parameters (including RDRRTB) are tensors, as you pointed out in the code review.

Notice that I have also added a branch to calculate_numerical_grad in order to return the correct gradient for RDRRTB.

[SarahAlidoost]

@SCiarella thanks for addressing the comments, nice work 💪 I have one suggestion on refactoring of numerical function making it more generic, see my suggestion on the code. If you agree with it, please go ahead with committing it and merge this PR 🚀

[sonarqubecloud]

Quality Gate passed

Issues
2 New issues
0 Accepted issues

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

See analysis details on SonarQube Cloud