Make `Wofost72` differentiable and vectorized by SCiarella · Pull Request #94 · WUR-AI/diffWOFOST

SCiarella · 2026-03-03T11:27:49Z

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.

use torch.where in phenology add config flag for carbon balance

SCiarella · 2026-03-03T13:37:50Z

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

~~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]
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.
phenology uses torch.where instead of if statements to handle mixed states
removed the pcse.decorators that are redundant with Tensor

SCiarella · 2026-03-04T08:50:03Z

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

--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.
--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 · 2026-03-04T12:33:00Z

Quality Gate passed

Issues
15 New issues
0 Accepted issues

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

See analysis details on SonarQube Cloud

SarahAlidoost

@SCiarella Awesome! 👍 The changes look very good, nice to see that all differentiable crop modules are now integrated in wofost72 and gradients are valid 😄 In #88, we will show the optimization of wofost72_pp. Thanks for implementing the soil module too.

I left a few suggestions:

I added some inline comment about issue #60. In wofost72.py, we have conditions like if torch.all(self.pheno.states.STAGE == 0). Right now STAGE is either an integer or an array where all values are the same. But actually, with different crops and agro-management, this condition might not work. This should be fixed when engine can handle different crops and agro-management.
I suggested choosing different data for fast test.

After addressing these, we can merge this 🚀

SarahAlidoost · 2026-03-06T13:29:28Z

src/diffwofost/physical_models/crop/phenology.py

        self.rates = self.RateVariables(kiosk, shape=shape)
        self.kiosk = kiosk

        self._connect_signal(self._on_CROP_FINISH, signal=signals.crop_finish)


Suggested change

# see issue #60

self._connect_signal(self._on_CROP_FINISH, signal=signals.crop_finish)

SarahAlidoost · 2026-03-06T13:30:56Z

src/diffwofost/physical_models/crop/phenology.py


        # Send crop_finish signal if maturity reached for one.
        # assumption is that all elements mature simultaneously
        # TODO: revisit this when fixing engine for agromanager


Suggested change

# TODO: revisit this when fixing engine for agromanager, see issue #60

SarahAlidoost · 2026-03-06T13:31:13Z

src/diffwofost/physical_models/crop/phenology.py

        msg = "Finished state integration for %s"
        self.logger.debug(msg % day)

    def _on_CROP_FINISH(self, day, finish_type=None):


Suggested change

# TODO: revisit this when fixing engine for agromanager, see issue #60

def _on_CROP_FINISH(self, day, finish_type=None):

SarahAlidoost · 2026-03-06T13:34:15Z

src/diffwofost/physical_models/crop/phenology.py

                self.params.shape, day.toordinal(), dtype=self.dtype, device=self.device
            )

    def get_variable(self, varname):


We can remove this function. I couldn't do it in the review because this part is not changed in this PR.

SarahAlidoost · 2026-03-06T13:41:33Z

src/diffwofost/physical_models/soil/classic_waterbalance.py

+        EVS = Tensor(-99.0)
+        WTRA = Tensor(-99.0)


Suggested change

EVS = Tensor(-99.0)

WTRA = Tensor(-99.0)

EVS = Tensor(0.0)

WTRA = Tensor(0.0)

SarahAlidoost · 2026-03-06T15:01:45Z

tests/conftest.py

+    test_data_url = params.get("test_data_url")
+    if test_data_url is not None:
+        match = re.search(r"_(\d+)\.yaml$", str(test_data_url))
+        if match and int(match.group(1)) > 5:


Good idea! Files 1:5 are all the same crop and similar agro management data. It is better to choose files 05, 06, 10, 20, 30, 44. These cover different crops and agro management info, and perhaps edge cases, if any. Is it possible to use those files instead? we can hard code the file names here.

SarahAlidoost · 2026-03-06T15:02:50Z

src/diffwofost/physical_models/crop/wofost72.py

+        self.pheno.calc_rates(day, drv)
+
+        # if before emergence there is no need to continue
+        # because only the phenology is running.


Suggested change

# because only the phenology is running.

# because only the phenology is running.

# TODO: revisit this when fixing #60

SarahAlidoost · 2026-03-06T15:03:17Z

src/diffwofost/physical_models/crop/wofost72.py

+        # if before emergence there is no need to continue
+        # because only the phenology is running.
+        # Just run a touch() to to ensure that all state variables are available
+        # in the kiosk


Suggested change

# in the kiosk

# in the kiosk

# TODO: revisit this when fixing #60

SarahAlidoost · 2026-03-06T15:39:14Z

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

--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.

--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.

Good idea! This has been done in pcse as well, quick and full test. I see that the action still takes 20 minutes. Later we can think about optimizing this further.

@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?

That's a good solution. Also, we can remove windows and mac and python 3.11 and 3.12 from --full_wofost72_test that runs once before PR merge. In a separate github action, we can define a workflow with all OS and pythons for --full_wofost72_test with a workflow_dispatch that we can run the test once before a release.

SCiarella added 4 commits February 27, 2026 10:54

core

bb2ff9e

test

a51620d

Fix conflicts

fc4b8f3

Add soil module

08b7dd5

SCiarella requested a review from SarahAlidoost March 3, 2026 11:33

SCiarella marked this pull request as ready for review March 3, 2026 11:33

SCiarella added 2 commits March 3, 2026 13:12

Fix grad combinations

65b2bb8

Add cache to assimilation

0b6cbc5

use torch.where in phenology add config flag for carbon balance

SCiarella added 6 commits March 3, 2026 15:43

Remove decorators

1db27ed

Fix device

fa2b351

Clean comments

772cc64

Fix tests

d8ae1b6

Fix

5081de6

Add pytest options

fd9ebe8

SCiarella added 2 commits March 4, 2026 10:43

Restore carbon_check

31c5f53

Optimize assimilation

1934448

SarahAlidoost approved these changes Mar 6, 2026

View reviewed changes


	# see issue #60
	self._connect_signal(self._on_CROP_FINISH, signal=signals.crop_finish)


	# TODO: revisit this when fixing engine for agromanager, see issue #60


	# TODO: revisit this when fixing engine for agromanager, see issue #60
	def _on_CROP_FINISH(self, day, finish_type=None):

	# because only the phenology is running.
	# because only the phenology is running.
	# TODO: revisit this when fixing #60

	# in the kiosk
	# in the kiosk
	# TODO: revisit this when fixing #60

Conversation

SCiarella commented Mar 3, 2026

Uh oh!

SCiarella commented Mar 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

SCiarella commented Mar 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sonarqubecloud bot commented Mar 4, 2026

Quality Gate passed

Uh oh!

SarahAlidoost left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SarahAlidoost commented Mar 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

SCiarella commented Mar 3, 2026 •

edited

Loading

SCiarella commented Mar 4, 2026 •

edited

Loading