Make the Leaf dynamics differentiable

mkdocs.yml

@@ -8,7 +8,6 @@ nav:
 
 theme:
   name: material
-  custom_dir: docs/assets
   features:
     - navigation.instant
     - navigation.tabs

pyproject.toml

@@ -24,6 +24,7 @@ classifiers = [
 ]
 dependencies = [
     "torch",
+    "pcse@git+https://github.com/ajwdewit/pcse.git@master",
 ]
 description = "Differentiable WOFOST"
 keywords = ["pytorch"," differentiable"," crop"," optimization"]
@@ -66,7 +67,7 @@ command_line = "-m pytest"
 
 
 [tool.ruff]
-line-length = 88
+line-length = 100
 output-format = "concise"
 extend-exclude = ["docs", "build"]
 
@@ -104,6 +105,10 @@ ignore = [
     "D203", # 1 blank line required before class docstring
     "D213", # Multi-line docstring summary should start at the second line
     "D413", # Missing blank line after last section
+    "N806", # Variable name should be lowercase
+    "N803", # Argument name should be lowercase
+    "N802", # Function name should be lowercase
+    "N801", # Class name should use CapWords convention
 ]
 pydocstyle.convention = "google"
 

tests/physical_models/__init__.py

@@ -0,0 +1,4 @@
+from pathlib import Path
+
+test_folder = Path(__file__).resolve().parent
+phy_data_folder = test_folder / "test_data"

tests/physical_models/crop/test_leaf_dynamics.py

@@ -0,0 +1,345 @@
+import copy
+from unittest.mock import patch
+import pytest
+import torch
+import torch.testing
+import yaml
+from numpy.testing import assert_almost_equal
+from pcse.base.parameter_providers import ParameterProvider
+from pcse.engine import Engine
+from pcse.models import Wofost72_PP
+from diffwofost.physical_models.crop.leaf_dynamics import WOFOST_Leaf_Dynamics
+from tests.physical_models.pcse_test_code import EngineTestHelper
+from tests.physical_models.pcse_test_code import WeatherDataProviderTestHelper
+from .. import phy_data_folder
+
+
+def prepare_engine_input(file_path):
+    inputs = yaml.safe_load(open(file_path))
+    agro_management_inputs = inputs["AgroManagement"]
+    cropd = inputs["ModelParameters"]
+
+    weather_data_provider = WeatherDataProviderTestHelper(inputs["WeatherVariables"])
+    crop_model_params_provider = ParameterProvider(cropdata=cropd)
+    external_states = inputs["ExternalStates"]
+
+    # convert parameters to tensors
+    crop_model_params_provider.clear_override()
+    for name in ["SPAN", "TDWI", "TBASE", "PERDL", "RGRLAI"]:
+        value = torch.tensor(crop_model_params_provider[name], dtype=torch.float32)
+        crop_model_params_provider.set_override(name, value, check=False)
+
+    # convert external states to tensors
+    tensor_external_states = [
+        {k: v if k == "DAY" else torch.tensor(v, dtype=torch.float32) for k, v in item.items()}
+        for item in external_states
+    ]
+    return (
+        crop_model_params_provider,
+        weather_data_provider,
+        agro_management_inputs,
+        tensor_external_states,
+    )
+
+
+def get_test_data(file_path):
+    inputs = yaml.safe_load(open(file_path))
+    return inputs["ModelResults"], inputs["Precision"]
+
+
+def get_test_diff_leaf_model():
+    test_data_path = phy_data_folder / "test_leafdynamics_wofost72_01.yaml"
+    (crop_model_params_provider, weather_data_provider, agro_management_inputs, external_states) = (
+        prepare_engine_input(test_data_path)
+    )
+    config_path = str(phy_data_folder / "WOFOST_Leaf_Dynamics.conf")
+    return DiffLeafDynamics(
+        copy.deepcopy(crop_model_params_provider),
+        weather_data_provider,
+        agro_management_inputs,
+        config_path,
+        copy.deepcopy(external_states),
+    )
+
+
+class DiffLeafDynamics(torch.nn.Module):
+    def __init__(
+        self,
+        crop_model_params_provider,
+        weather_data_provider,
+        agro_management_inputs,
+        config_path,
+        external_states,
+    ):
+        super().__init__()
+        self.crop_model_params_provider = crop_model_params_provider
+        self.weather_data_provider = weather_data_provider
+        self.agro_management_inputs = agro_management_inputs
+        self.config_path = config_path
+        self.external_states = external_states
+
+    def forward(self, params_dict):
+        # pass new value of parameters to the model
+        for name, value in params_dict.items():
+            self.crop_model_params_provider.set_override(name, value, check=False)
+
+        engine = EngineTestHelper(
+            self.crop_model_params_provider,
+            self.weather_data_provider,
+            self.agro_management_inputs,
+            self.config_path,
+            self.external_states,
+        )
+        engine.run_till_terminate()
+        results = engine.get_output()
+
+        return torch.stack(
+            [torch.stack([item["LAI"], item["TWLV"]]) for item in results]
+        ).unsqueeze(0)  # shape: [1, time_steps, 2]
+
+
+def calculate_numerical_grad(param_name, param_value, output_name):
+    delta = 1e-6
+    p_plus = param_value.item() + delta
+    p_minus = param_value.item() - delta
+
+    model = get_test_diff_leaf_model()
+    output = model({param_name: torch.nn.Parameter(torch.tensor(p_plus, dtype=torch.float64))})
+    if output_name == "LAI":
+        loss_plus = output[0, :, 0].sum()
+    elif output_name == "TWLV":
+        loss_plus = output[0, :, 1].sum()
+
+    model = get_test_diff_leaf_model()
+    output = model({param_name: torch.nn.Parameter(torch.tensor(p_minus, dtype=torch.float64))})
+    if output_name == "LAI":
+        loss_minus = output[0, :, 0].sum()
+    elif output_name == "TWLV":
+        loss_minus = output[0, :, 1].sum()
+
+    return (loss_plus.item() - loss_minus.item()) / (2 * delta)
+
+
+class TestLeafDynamics:
+    def test_leaf_dynamics_with_testengine(self):
+        """EngineTestHelper and not Engine because it allows to specify `external_states`."""
+
+        # prepare model input
+        test_data_path = phy_data_folder / "test_leafdynamics_wofost72_01.yaml"
+        (
+            crop_model_params_provider,
+            weather_data_provider,
+            agro_management_inputs,
+            external_states,
+        ) = prepare_engine_input(test_data_path)
+        config_path = str(phy_data_folder / "WOFOST_Leaf_Dynamics.conf")
+
+        engine = EngineTestHelper(
+            crop_model_params_provider,
+            weather_data_provider,
+            agro_management_inputs,
+            config_path,
+            external_states,
+        )
+        engine.run_till_terminate()
+        actual_results = engine.get_output()
+
+        # get expected results from YAML test data
+        expected_results, expected_precision = get_test_data(test_data_path)
+
+        assert len(actual_results) == len(expected_results)
+
+        for reference, model in zip(expected_results, actual_results, strict=False):
+            assert reference["DAY"] == model["day"]
+            assert all(
+                abs(reference[var] - model[var]) < precision
+                for var, precision in expected_precision.items()
+            )
+
+    def test_leaf_dynamics_with_engine(self):
+        # prepare model input
+        test_data_path = phy_data_folder / "test_leafdynamics_wofost72_01.yaml"
+        (crop_model_params_provider, weather_data_provider, agro_management_inputs, _) = (
+            prepare_engine_input(test_data_path)
+        )
+
+        config_path = str(phy_data_folder / "WOFOST_Leaf_Dynamics.conf")
+
+        # Engine does not allows to specify `external_states`
+        with pytest.raises(ValueError):
+            Engine(
+                crop_model_params_provider,
+                weather_data_provider,
+                agro_management_inputs,
+                config_path,
+            )
+
+    def test_wofost_pp_with_leaf_dynamics(self):
+        # prepare model input
+        test_data_path = phy_data_folder / "test_potentialproduction_wofost72_01.yaml"
+        (crop_model_params_provider, weather_data_provider, agro_management_inputs, _) = (
+            prepare_engine_input(test_data_path)
+        )
+
+        # get expected results from YAML test data
+        expected_results, expected_precision = get_test_data(test_data_path)
+
+        with patch("pcse.crop.leaf_dynamics.WOFOST_Leaf_Dynamics", WOFOST_Leaf_Dynamics):
+            model = Wofost72_PP(
+                crop_model_params_provider, weather_data_provider, agro_management_inputs
+            )
+            model.run_till_terminate()
+            actual_results = model.get_output()
+
+            assert len(actual_results) == len(expected_results)
+
+            for reference, model in zip(expected_results, actual_results, strict=False):
+                assert reference["DAY"] == model["day"]
+                assert all(
+                    abs(reference[var] - model[var]) < precision
+                    for var, precision in expected_precision.items()
+                )
+
+
+class TestDiffLeafDynamicsTDWI:
+    def test_gradients_tdwi_lai_leaf_dynamics(self):
+        model = get_test_diff_leaf_model()
+        tdwi = torch.nn.Parameter(torch.tensor(0.2, dtype=torch.float32))
+        output = model({"TDWI": tdwi})
+        lai = output[0, :, 0]
+        loss = lai.sum()
+
+        # this is ∂loss/∂tdwi without calling loss.backward().
+        # this is called forward gradient here because it is calculated without backpropagation.
+        grads = torch.autograd.grad(loss, tdwi, retain_graph=True)[0]
+        assert grads is not None, "Gradients for TDWI should not be None"
+
+        tdwi.grad = None  # clear any existing gradient
+        loss.backward()
+        # this is ∂loss/∂tdwi calculated using backpropagation
+        grad_backward = tdwi.grad
+
+        assert grad_backward is not None, "Backward gradients for TDWI should not be None"
+        assert grad_backward == grads, "Forward and backward gradients for TDWI should match"
+
+    def test_gradients_tdwi_lai_leaf_dynamics_numerical(self):
+        # first check if the numerical gradient isnot zero i.e. the parameter has an effect
+        tdwi = torch.nn.Parameter(torch.tensor(0.2, dtype=torch.float64))
+        numerical_grad = calculate_numerical_grad("TDWI", tdwi, "LAI")  # this is Δloss/Δtdwi
+
+        model = get_test_diff_leaf_model()
+        output = model({"TDWI": tdwi})
+        lai = output[0, :, 0]
+        loss = lai.sum()
+        # this is ∂loss/∂tdwi, for comparison with numerical gradient
+        grads = torch.autograd.grad(loss, tdwi, retain_graph=True)[0]
+
+        assert_almost_equal(numerical_grad, grads.item(), decimal=3)
+
+    def test_gradients_tdwi_twlv_leaf_dynamics(self):
+        # prepare model input
+        model = get_test_diff_leaf_model()
+        tdwi = torch.nn.Parameter(torch.tensor(0.2, dtype=torch.float32))
+        output = model({"TDWI": tdwi})
+        twlv = output[0, :, 1]
+        loss = twlv.sum()
+
+        # this is ∂loss/∂tdwi
+        # this is called forward gradient here because it is calculated without backpropagation.
+        grads = torch.autograd.grad(loss, tdwi, retain_graph=True)[0]
+        assert grads is not None, "Gradients for TDWI should not be None"
+
+        tdwi.grad = None  # clear any existing gradient
+        loss.backward()
+        # this is ∂loss/∂tdwi calculated using backpropagation
+        grad_backward = tdwi.grad
+
+        assert grad_backward is not None, "Backward gradients for TDWI should not be None"
+        assert grad_backward == grads, "Forward and backward gradients for TDWI should match"
+
+    def test_gradients_tdwi_twlv_leaf_dynamics_numerical(self):
+        # first check if the numerical gradient isnot zero i.e. the parameter has an effect
+        tdwi = torch.nn.Parameter(torch.tensor(0.2, dtype=torch.float64))
+        numerical_grad = calculate_numerical_grad("TDWI", tdwi, "TWLV")  # this is Δloss/Δtdwi
+
+        model = get_test_diff_leaf_model()
+        output = model({"TDWI": tdwi})
+        twlv = output[0, :, 1]
+        loss = twlv.sum()
+        # this is ∂loss/∂tdwi, for comparison with numerical gradient
+        grads = torch.autograd.grad(loss, tdwi, retain_graph=True)[0]
+
+        assert_almost_equal(numerical_grad, grads.item(), decimal=3)
+
+
+class TestDiffLeafDynamicsSPAN:
+    def test_gradients_span_lai_leaf_dynamics(self):
+        # prepare model input
+        model = get_test_diff_leaf_model()
+        span = torch.nn.Parameter(torch.tensor(30, dtype=torch.float32))
+        output = model({"SPAN": span})
+        lai = output[0, :, 0]
+        loss = lai.sum()
+
+        # this is ∂loss/∂span
+        # this is called forward gradient here because it is calculated without backpropagation.
+        grads = torch.autograd.grad(loss, span, retain_graph=True)[0]
+        assert grads is not None, "Gradients for SPAN should not be None"
+
+        span.grad = None  # clear any existing gradient
+        loss.backward()
+        # this is ∂loss/∂span calculated using backpropagation
+        grad_backward = span.grad
+
+        assert grad_backward is not None, "Backward gradients for SPAN should not be None"
+        assert grad_backward == grads, "Forward and backward gradients for SPAN should match"
+
+    def test_gradients_span_lai_leaf_dynamics_numerical(self):
+        # first check if the numerical gradient isnot zero i.e. the parameter has an effect
+        span = torch.nn.Parameter(torch.tensor(30, dtype=torch.float64))
+        numerical_grad = calculate_numerical_grad("SPAN", span, "LAI")  # this is Δloss/Δspan
+
+        model = get_test_diff_leaf_model()
+        output = model({"SPAN": span})
+        lai = output[0, :, 0]
+        loss = lai.sum()
+        # this is ∂loss/∂tdwi, for comparison with numerical gradient
+        grads = torch.autograd.grad(loss, span, retain_graph=True)[0]
+
+        assert_almost_equal(numerical_grad, grads.item(), decimal=3)
+
+    def test_gradients_span_twlv_leaf_dynamics(self):
+        # prepare model input
+        model = get_test_diff_leaf_model()
+        span = torch.nn.Parameter(torch.tensor(30, dtype=torch.float32))
+        output = model({"SPAN": span})
+        twlv = output[0, :, 1]
+        loss = twlv.sum()
+
+        # this is ∂loss/∂span
+        # this is called forward gradient here because it is calculated without backpropagation.
+        grads = torch.autograd.grad(loss, span, retain_graph=True)[0]
+        assert grads is not None, "Gradients for SPAN should not be None"
+
+        span.grad = None  # clear any existing gradient
+        loss.backward()
+        # this is ∂loss/∂span calculated using backpropagation
+        grad_backward = span.grad
+
+        assert grad_backward is not None, "Backward gradients for SPAN should not be None"
+        assert grad_backward == grads, "Forward and backward gradients for SPAN should match"
+
+    def test_gradients_span_twlv_leaf_dynamics_numerical(self):
+        # first check if the numerical gradient isnot zero i.e. the parameter has an effect
+        span = torch.nn.Parameter(torch.tensor(30, dtype=torch.float64))
+        numerical_grad = calculate_numerical_grad("SPAN", span, "TWLV")  # this is Δloss/Δspan
+
+        model = get_test_diff_leaf_model()
+        output = model({"SPAN": span})
+        twlv = output[0, :, 1]
+        loss = twlv.sum()
+        # this is ∂loss/∂tdwi, for comparison with numerical gradient
+        grads = torch.autograd.grad(loss, span, retain_graph=True)[0]
+
+        assert numerical_grad == 0.0
+        assert_almost_equal(numerical_grad, grads.item(), decimal=3)