Optimizing parameters in a WOFOST crop model using diffWOFOST

This Jupyter notebook demonstrates the optimization of parameters in a differentiable model using the diffwofost package. The package provides differentiable implementations of the WOFOST model and its associated sub-models. As diffwofost is under active development, this notebook focuses on root_dynamics.

To enable these models to operate independently, certain state variables required by the model are supplied as "external states" derived from the test data. Also, at this stage, only a limited subset of model parameters has been made differentiable.

1. Root dynamics¶

In this section, we will demonstrate how to optimize two parameters TWDI in root_dynamics model using a differentiable version of root_dynamics. The optimization will be done using the Adam optimizer from torch.optim.

1.1 software requirements¶

To run this notebook, we need to install the diffwofost; the differentiable version of WOFOST models. Since the package is constantly under development, make sure you have the latest version of diffwofost installed in your python environment. You can install it using pip:

In [ ]:

# install diffwofost
!pip install diffwofost

# install diffwofost !pip install diffwofost

install diffwofost¶

!pip install diffwofost

In [2]:

# ---- import libraries ----
import copy
import torch
import numpy
import yaml
from pathlib import Path
from diffwofost.physical_models.config import Configuration
from diffwofost.physical_models.crop.root_dynamics import WOFOST_Root_Dynamics
from diffwofost.physical_models.utils import EngineTestHelper
from diffwofost.physical_models.utils import prepare_engine_input
from diffwofost.physical_models.utils import get_test_data

# ---- import libraries ---- import copy import torch import numpy import yaml from pathlib import Path from diffwofost.physical_models.config import Configuration from diffwofost.physical_models.crop.root_dynamics import WOFOST_Root_Dynamics from diffwofost.physical_models.utils import EngineTestHelper from diffwofost.physical_models.utils import prepare_engine_input from diffwofost.physical_models.utils import get_test_data

In [ ]:

# --- run on CPU ------
from diffwofost.physical_models.config import ComputeConfig
ComputeConfig.set_device('cpu')

# --- run on CPU ------ from diffwofost.physical_models.config import ComputeConfig ComputeConfig.set_device('cpu')

In [3]:

# ---- disable a warning: this will be fixed in the future ----
import warnings
warnings.filterwarnings("ignore", message="To copy construct from a tensor.*")

# ---- disable a warning: this will be fixed in the future ---- import warnings warnings.filterwarnings("ignore", message="To copy construct from a tensor.*")

1.2. Data¶

A test dataset of TWRT (Total weight of roots) will be used to optimize parametesr TWDI (total initial dry weight). Note that in root_dynamic, changes in TWDI dont affect RD (Current rooting depth).

The data is stored in PCSE tests folder, and can be doewnloded from PCSE repsository. You can select any of the files related to root_dynamics model with a file name that follwos the pattern test_rootdynamics_wofost72_*.yaml. Each file contains different data depending on the locatin and crop type. For example, you can download the file "test_rootdynamics_wofost72_01.yaml" as:

In [4]:

import urllib.request

url = "https://raw.githubusercontent.com/ajwdewit/pcse/refs/heads/master/tests/test_data/test_rootdynamics_wofost72_01.yaml"
filename = "test_rootdynamics_wofost72_01.yaml"

urllib.request.urlretrieve(url, filename)
print(f"Downloaded: {filename}")

import urllib.request url = "https://raw.githubusercontent.com/ajwdewit/pcse/refs/heads/master/tests/test_data/test_rootdynamics_wofost72_01.yaml" filename = "test_rootdynamics_wofost72_01.yaml" urllib.request.urlretrieve(url, filename) print(f"Downloaded: {filename}")

Downloaded: test_rootdynamics_wofost72_01.yaml

We also need to download a config file to be able to run each crop module. This will change in the future versions. To donwload the config file, you can use the following command:

In [5]:

# ---- Check the path to the files that are downloaded as explained above ----
test_data_path = "test_rootdynamics_wofost72_01.yaml"

# ---- Check the path to the files that are downloaded as explained above ---- test_data_path = "test_rootdynamics_wofost72_01.yaml"

In [6]:

# ---- Here we read the test data and set some variables ----
test_data = get_test_data(test_data_path)
(crop_model_params_provider, weather_data_provider, agro_management_inputs, external_states) = (
    prepare_engine_input(test_data, ["RDI", "RRI", "RDMCR", "RDMSOL", "TDWI", "IAIRDU"])
)

expected_results = test_data["ModelResults"]
expected_twrt = torch.tensor(
    [float(item["TWRT"]) for item in expected_results], dtype=torch.float32
) # shape: [1, time_steps]

# ---- dont change this: in this config file we specified the diffrentiable version of root_dynamics ----
root_dynamics_config = Configuration(
    CROP=WOFOST_Root_Dynamics,
    OUTPUT_VARS=["RD", "TWRT"],
)

# ---- Here we read the test data and set some variables ---- test_data = get_test_data(test_data_path) (crop_model_params_provider, weather_data_provider, agro_management_inputs, external_states) = ( prepare_engine_input(test_data, ["RDI", "RRI", "RDMCR", "RDMSOL", "TDWI", "IAIRDU"]) ) expected_results = test_data["ModelResults"] expected_twrt = torch.tensor( [float(item["TWRT"]) for item in expected_results], dtype=torch.float32 ) # shape: [1, time_steps] # ---- dont change this: in this config file we specified the diffrentiable version of root_dynamics ---- root_dynamics_config = Configuration( CROP=WOFOST_Root_Dynamics, OUTPUT_VARS=["RD", "TWRT"], )

1.3. Helper classes/functions¶

The model parameters shoudl stay in a valid range. To ensure this, we will use BoundedParameter class with (min, max) and initial values for each parameter. You might change these values depending on the crop type and location. But dont use a very small range, otherwise gradiants will be very small and the optimization will be very slow.

In [7]:

# ---- Adjust the values if needed  ----
TDWI_MIN, TDWI_MAX, TDWI_INIT = (0.0, 1.0, 0.30)

# ---- Helper for bounded parameters ----
class BoundedParameter(torch.nn.Module):
    def __init__(self, low, high, init_value):
        super().__init__()
        self.low = low
        self.high = high

        # Normalize to [0, 1]
        init_norm = (init_value - low) / (high - low)

        # Parameter in raw logit space
        self.raw = torch.nn.Parameter(torch.logit(torch.tensor(init_norm, dtype=torch.float32), eps=1e-6))

    def forward(self):
        return self.low + (self.high - self.low) * torch.sigmoid(self.raw)

# ---- Adjust the values if needed ---- TDWI_MIN, TDWI_MAX, TDWI_INIT = (0.0, 1.0, 0.30) # ---- Helper for bounded parameters ---- class BoundedParameter(torch.nn.Module): def __init__(self, low, high, init_value): super().__init__() self.low = low self.high = high # Normalize to [0, 1] init_norm = (init_value - low) / (high - low) # Parameter in raw logit space self.raw = torch.nn.Parameter(torch.logit(torch.tensor(init_norm, dtype=torch.float32), eps=1e-6)) def forward(self): return self.low + (self.high - self.low) * torch.sigmoid(self.raw)

Another helper class is OptRootDynamics which is a subclass of torch.nn.Module. We use this class to wrap the EngineTestHelper function and make it easier to run the model root_dynamic.

In [8]:

# ---- Wrap the model with torch.nn.Module----
class OptDiffRootDynamics(torch.nn.Module):
    def __init__(self, crop_model_params_provider, weather_data_provider, agro_management_inputs, root_dynamics_config, 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 = root_dynamics_config
        self.external_states = external_states

        # bounded parameters
        self.tdwi = BoundedParameter(TDWI_MIN, TDWI_MAX, init_value=TDWI_INIT)
        
    def forward(self):
        # currently, copying is needed due to an internal issue in engine
        crop_model_params_provider_ = copy.deepcopy(self.crop_model_params_provider)
        external_states_ = copy.deepcopy(self.external_states)
        
        tdwi_val = self.tdwi()
        
        # pass new value of parameters to the model
        crop_model_params_provider_.set_override("TDWI", tdwi_val, check=False)

        engine = EngineTestHelper(
            crop_model_params_provider_,
            self.weather_data_provider,
            self.agro_management_inputs,
            self.config,
            external_states_,
        )
        engine.run_till_terminate()
        results = engine.get_output()
        
        return torch.stack([item["TWRT"] for item in results]) # shape: [1, time_steps]

# ---- Wrap the model with torch.nn.Module---- class OptDiffRootDynamics(torch.nn.Module): def __init__(self, crop_model_params_provider, weather_data_provider, agro_management_inputs, root_dynamics_config, 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 = root_dynamics_config self.external_states = external_states # bounded parameters self.tdwi = BoundedParameter(TDWI_MIN, TDWI_MAX, init_value=TDWI_INIT) def forward(self): # currently, copying is needed due to an internal issue in engine crop_model_params_provider_ = copy.deepcopy(self.crop_model_params_provider) external_states_ = copy.deepcopy(self.external_states) tdwi_val = self.tdwi() # pass new value of parameters to the model crop_model_params_provider_.set_override("TDWI", tdwi_val, check=False) engine = EngineTestHelper( crop_model_params_provider_, self.weather_data_provider, self.agro_management_inputs, self.config, external_states_, ) engine.run_till_terminate() results = engine.get_output() return torch.stack([item["TWRT"] for item in results]) # shape: [1, time_steps]

In [9]:

# ----  Create model ---- 
opt_model = OptDiffRootDynamics(
    crop_model_params_provider,
    weather_data_provider,
    agro_management_inputs,
    root_dynamics_config,
    external_states,
)

# ---- Create model ---- opt_model = OptDiffRootDynamics( crop_model_params_provider, weather_data_provider, agro_management_inputs, root_dynamics_config, external_states, )

In [10]:

# ----  Early stopping ---- 
best_loss = float("inf")
patience = 10  # Number of steps to wait for improvement
patience_counter = 0
min_delta = 1e-4 

# ----  Optimizer ---- 
optimizer = torch.optim.Adam(opt_model.parameters(), lr=0.1)

# ----  We use relative MAE as loss because there are two outputs with different untis ----  
denom = torch.mean(torch.abs(expected_twrt)) 

# Training loop (example)
for step in range(101):
    optimizer.zero_grad()
    results = opt_model() 
    mae = torch.mean(torch.abs(results - expected_twrt))
    loss = mae / denom # example: relative mean absolute error
    loss.backward()
    optimizer.step()

    print(f"Step {step}, Loss {loss.item():.8f}, TDWI {opt_model.tdwi().item():.4f}")
    # Early stopping logic
    if loss.item() < best_loss - min_delta:
        best_loss = loss.item()
        patience_counter = 0
    else:
        patience_counter += 1
        if patience_counter >= patience:
            print(f"Early stopping at step {step}")
            break

# ---- Early stopping ---- best_loss = float("inf") patience = 10 # Number of steps to wait for improvement patience_counter = 0 min_delta = 1e-4 # ---- Optimizer ---- optimizer = torch.optim.Adam(opt_model.parameters(), lr=0.1) # ---- We use relative MAE as loss because there are two outputs with different untis ---- denom = torch.mean(torch.abs(expected_twrt)) # Training loop (example) for step in range(101): optimizer.zero_grad() results = opt_model() mae = torch.mean(torch.abs(results - expected_twrt)) loss = mae / denom # example: relative mean absolute error loss.backward() optimizer.step() print(f"Step {step}, Loss {loss.item():.8f}, TDWI {opt_model.tdwi().item():.4f}") # Early stopping logic if loss.item() < best_loss - min_delta: best_loss = loss.item() patience_counter = 0 else: patience_counter += 1 if patience_counter >= patience: print(f"Early stopping at step {step}") break

Step 0, Loss 0.00004644, TDWI 0.3214
Step 1, Loss 0.00004170, TDWI 0.3436
Step 2, Loss 0.00003679, TDWI 0.3665
Step 3, Loss 0.00003172, TDWI 0.3901
Step 4, Loss 0.00002650, TDWI 0.4143
Step 5, Loss 0.00002116, TDWI 0.4389
Step 6, Loss 0.00001571, TDWI 0.4639
Step 7, Loss 0.00001019, TDWI 0.4891
Step 8, Loss 0.00000461, TDWI 0.5144
Step 9, Loss 0.00000099, TDWI 0.5316
Step 10, Loss 0.00000479, TDWI 0.5424
Early stopping at step 10

In [11]:

# ---- validate the results using test data ---- 
print(f"Actual TDWI {crop_model_params_provider["TDWI"].item():.4f}")

# ---- validate the results using test data ---- print(f"Actual TDWI {crop_model_params_provider["TDWI"].item():.4f}")

Actual TDWI 0.5100

In [ ]: