Skip to content

Make Wofost72 differentiable and vectorized#94

Open
SCiarella wants to merge 14 commits intomainfrom
wofost72
Open

Make Wofost72 differentiable and vectorized#94
SCiarella wants to merge 14 commits intomainfrom
wofost72

Conversation

@SCiarella
Copy link
Collaborator

This PR closes #49.

Noticeably, Wofost72 uses kiosk variables such as EVS that come from the soil model. So in order to make wofost72 consistently differentiable and vectorized, I had to implement a differentiable soil model soil/classic_waterbalance.py which is inspired by the corresponding model from pcse.

This PR also adds a lot of tests, making issue #81 more relevant than ever.

@SCiarella SCiarella requested a review from SarahAlidoost March 3, 2026 11:33
@SCiarella SCiarella marked this pull request as ready for review March 3, 2026 11:33
use torch.where in phenology

add config flag for carbon balance
@SCiarella
Copy link
Collaborator Author

SCiarella commented Mar 3, 2026

In order to reduce the execution time, I have implemented the following modifications:

  1. ComputeConfig now stores a flag _check_carbon_balance that is used by wofost72 to skip some of the expensive checks that slow down the integration [Removed]

  2. assimilation now has a cache to store the constants and share them between different submodules. More importantly, the 3-point quadrature methods have been refactored in order to be vectorized.

  3. phenology uses torch.where instead of if statements to handle mixed states

  4. removed the pcse.decorators that are redundant with Tensor

@SCiarella
Copy link
Collaborator Author

SCiarella commented Mar 4, 2026

In order to make the tests more manageable (the test action takes >2h if unchecked), I have introduced two pytest parser addoptions:

  1. --full_wofost72_test if used, runs the gradient test for all the parameter combinations used by wofost72. By default, we now only test the gradients of the parameters that are directly manipulated in wofost72.
  2. --fast when specified, it: (a) makes pytest use test files 0 to 5 instead of all the 45 available, and (b) skips GPU and tensor-input in all the gradient tests.

@SarahAlidoost, we could think about using --fast by default for the GitHub actions to significantly save time, and only run the tests without --fast once before PR merge. This is probably the only real solution to #81. What do you think?

@sonarqubecloud
Copy link

sonarqubecloud bot commented Mar 4, 2026

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Task]: Make wofost72 differentiable and efficient

1 participant