diff --git a/docs/notebooks/optimization.ipynb b/docs/notebooks/optimization.ipynb
deleted file mode 100644
index 48fe890..0000000
--- a/docs/notebooks/optimization.ipynb
+++ /dev/null
@@ -1,738 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "id": "c784f1c2-d477-464d-a35f-f2c64e96fb10",
-   "metadata": {},
-   "source": [
-    "<div style=\"display: flex; align-items: center;\">\n",
-    "    <h1>Optimizing parameters in a WOFOST crop model using <code>diffWOFOST</code></h1>\n",
-    "    <img src=\"https://raw.githubusercontent.com/WUR-AI/diffWOFOST/refs/heads/main/docs/logo/diffwofost.png\" width=\"150\" style=\"margin-left: 20px;\">\n",
-    "</div>\n",
-    "\n",
-    "\n",
-    "This Jupyter notebook demonstrates the optimization of parameters in a\n",
-    "differentiable model using the `diffwofost` package. The package provides\n",
-    "differentiable implementations of the WOFOST model and its associated\n",
-    "sub-models. As `diffwofost` is under active development, this notebook focuses on\n",
-    "two sub-models: `leaf_dynamics` and `root_dynamics`. \n",
-    "\n",
-    "To enable these models to operate independently, certain state variables\n",
-    "required by the model are supplied as \"external states\" derived from the test\n",
-    "data. Also, at this stage, only a limited subset of model parameters has been made\n",
-    "differentiable.\n",
-    "\n",
-    "The notebook is organized into two standalone sections that can\n",
-    "be executed independently:\n",
-    "\n",
-    " 1. Leaf Dynamics\n",
-    " 2. Root Dynamics"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "41262fbd-270b-4616-91ad-09ee82451604",
-   "metadata": {},
-   "source": [
-    "## 1. Leaf dynamics\n",
-    "\n",
-    "In this section, we will demonstrate how to optimize two parameters `TWDI` and `SPAN` in\n",
-    "leaf_dynamics model using a differentiable version of leaf_dynamics.\n",
-    "The optimization will be done using the Adam optimizer from `torch.optim`."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "1b6c3f53-6fab-4537-9177-7b16e0a1ccec",
-   "metadata": {},
-   "source": [
-    "### 1.1 software requirements\n",
-    "\n",
-    "To run this notebook, we need to install the `diffwofost`; the differentiable\n",
-    "version of WOFOST models. Since the package is constantly under development, make\n",
-    "sure you have the latest version of `diffwofost` installed in your\n",
-    "python environment. You can install it using pip:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "e4049fea-1d05-41f1-bf9d-f030ae83a324",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# install diffwofost\n",
-    "!pip install diffwofost"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "21731653-3976-4bb9-b83b-b11d78211700",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ---- import libraries ----\n",
-    "import copy\n",
-    "import torch\n",
-    "import numpy\n",
-    "import yaml\n",
-    "from pathlib import Path\n",
-    "from diffwofost.physical_models.utils import EngineTestHelper\n",
-    "from diffwofost.physical_models.utils import prepare_engine_input"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "82a1ef6b-336e-4902-8bd1-2a1ed2020f9d",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ---- disable a warning: this will be fixed in the future ----\n",
-    "import warnings\n",
-    "warnings.filterwarnings(\"ignore\", message=\"To copy construct from a tensor.*\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "47def7fc-f2dd-4aaf-a572-41cc9d1e4679",
-   "metadata": {},
-   "source": [
-    "### 1.2. Data\n",
-    "\n",
-    "A test dataset of `LAI` (Leaf area index, including stem and pod area) and\n",
-    "`TWLV` (Dry weight of total leaves (living + dead)) will be used to optimize\n",
-    "parametesr `TWDI` (total initial dry weight) and `SPAN` (life span of leaves).\n",
-    "Note that in leaf_dynamic, changes in `SPAN` dont affect `TWLV`. \n",
-    "\n",
-    "The data is stored in PCSE tests folder, and can be doewnloded from PCSE repsository.\n",
-    "You can select any of the files related to `leaf_dynamics` model with a file name that follwos the pattern\n",
-    "`test_leafdynamics_wofost72_*.yaml`. Each file contains different data depending on the locatin and crop type.\n",
-    "For example, you can download the file \"test_leafdynamics_wofost72_01.yaml\" as:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "id": "0233a048-e5a2-4249-887d-35a37284769c",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Downloaded: test_leafdynamics_wofost72_01.yaml\n"
-     ]
-    }
-   ],
-   "source": [
-    "import urllib.request\n",
-    "\n",
-    "url = \"https://raw.githubusercontent.com/ajwdewit/pcse/refs/heads/master/tests/test_data/test_leafdynamics_wofost72_01.yaml\"\n",
-    "filename = \"test_leafdynamics_wofost72_01.yaml\"\n",
-    "\n",
-    "urllib.request.urlretrieve(url, filename)\n",
-    "print(f\"Downloaded: {filename}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "e4565b6b-523c-49c4-934e-500248317461",
-   "metadata": {},
-   "source": [
-    "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:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "id": "b4a24f1c-77e4-4b05-bde9-229dd497f09e",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Downloaded: WOFOST_Leaf_Dynamics.conf\n"
-     ]
-    }
-   ],
-   "source": [
-    "url = \"https://raw.githubusercontent.com/WUR-AI/diffWOFOST/refs/heads/main/tests/physical_models/test_data/WOFOST_Leaf_Dynamics.conf\"\n",
-    "filename = \"WOFOST_Leaf_Dynamics.conf\"\n",
-    "\n",
-    "urllib.request.urlretrieve(url, filename)\n",
-    "print(f\"Downloaded: {filename}\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "5a459489-bfcb-4ad6-9102-1b6be5edeb52",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ---- Check the path to the files that are downloaded as explained above ----\n",
-    "test_data_path = \"test_leafdynamics_wofost72_01.yaml\"\n",
-    "config_path = \"WOFOST_Leaf_Dynamics.conf\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "id": "9f3105fb-4fbe-4405-9fd4-e8255b4b119e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ---- Here we read the test data and set some variables ----\n",
-    "(crop_model_params_provider, weather_data_provider, agro_management_inputs, external_states) = (\n",
-    "    prepare_engine_input(test_data_path, [\"SPAN\", \"TDWI\", \"TBASE\", \"PERDL\", \"RGRLAI\"])\n",
-    ")\n",
-    "\n",
-    "expected_results = yaml.safe_load(open(test_data_path))[\"ModelResults\"]\n",
-    "expected_lai_twlv = torch.tensor(\n",
-    "    [[float(item[\"LAI\"]), float(item[\"TWLV\"])] for item in expected_results], dtype=torch.float32\n",
-    ").unsqueeze(0) # shape: [1, time_steps, 2]\n",
-    "\n",
-    "# ---- dont change this: in this config file we specified the diffrentiable version of leaf_dynamics ----\n",
-    "config_path = str(Path(config_path).resolve()) "
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "52b19ae2-3fe6-4b3f-95a7-aea07a2c0958",
-   "metadata": {},
-   "source": [
-    "### 1.3. Helper classes/functions\n",
-    "\n",
-    "The model parameters shoudl stay in a valid range. To ensure this, we will use\n",
-    "`BoundedParameter` class with (min, max) and initial values for each\n",
-    "parameter. You might change these values depending on the crop type and\n",
-    "location. But dont use a very small range, otherwise gradiants will be very\n",
-    "small and the optimization will be very slow."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "id": "e4610238-de0d-42cf-9689-3c074eb2cc0e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ---- Adjust the values if needed  ----\n",
-    "TDWI_MIN, TDWI_MAX, TDWI_INIT = (0.0, 1.0, 0.40)\n",
-    "SPAN_MIN, SPAN_MAX, SPAN_INIT = (10.0, 60.0, 30.0)\n",
-    "\n",
-    "# ---- Helper for bounded parameters ----\n",
-    "class BoundedParameter(torch.nn.Module):\n",
-    "    def __init__(self, low, high, init_value):\n",
-    "        super().__init__()\n",
-    "        self.low = low\n",
-    "        self.high = high\n",
-    "\n",
-    "        # Normalize to [0, 1]\n",
-    "        init_norm = (init_value - low) / (high - low)\n",
-    "\n",
-    "        # Parameter in raw logit space\n",
-    "        self.raw = torch.nn.Parameter(torch.logit(torch.tensor(init_norm, dtype=torch.float32), eps=1e-6))\n",
-    "\n",
-    "    def forward(self):\n",
-    "        return self.low + (self.high - self.low) * torch.sigmoid(self.raw)\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "153e4306-77ed-4278-8797-a04e637e12c8",
-   "metadata": {},
-   "source": [
-    "Another helper class is `OptDiffLeafDynamics` which is a subclass of `torch.nn.Module`. \n",
-    "We use this class to wrap the `EngineTestHelper` function and make it easier to run the model `leaf_dynamic`."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "id": "36dd6463-4812-41c0-b2bf-d4769df1136f",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ---- Wrap the model with torch.nn.Module----\n",
-    "class OptDiffLeafDynamics(torch.nn.Module):\n",
-    "    def __init__(self, crop_model_params_provider, weather_data_provider, agro_management_inputs, config_path, external_states):\n",
-    "        super().__init__()\n",
-    "        self.crop_model_params_provider = crop_model_params_provider\n",
-    "        self.weather_data_provider = weather_data_provider\n",
-    "        self.agro_management_inputs = agro_management_inputs\n",
-    "        self.config_path = config_path\n",
-    "        self.external_states = external_states\n",
-    "\n",
-    "        # bounded parameters\n",
-    "        self.tdwi = BoundedParameter(TDWI_MIN, TDWI_MAX, init_value=TDWI_INIT)\n",
-    "        self.span = BoundedParameter(SPAN_MIN, SPAN_MAX, init_value=SPAN_INIT)\n",
-    "\n",
-    "    def forward(self):\n",
-    "        # currently, copying is needed due to an internal issue in engine\n",
-    "        crop_model_params_provider_ = copy.deepcopy(self.crop_model_params_provider)\n",
-    "        external_states_ = copy.deepcopy(self.external_states)\n",
-    "        \n",
-    "        tdwi_val = self.tdwi()\n",
-    "        span_val = self.span()\n",
-    "        \n",
-    "        # pass new value of parameters to the model\n",
-    "        crop_model_params_provider_.set_override(\"TDWI\", tdwi_val, check=False)\n",
-    "        crop_model_params_provider_.set_override(\"SPAN\", span_val, check=False)\n",
-    "\n",
-    "        engine = EngineTestHelper(\n",
-    "            crop_model_params_provider_,\n",
-    "            self.weather_data_provider,\n",
-    "            self.agro_management_inputs,\n",
-    "            self.config_path,\n",
-    "            external_states_,\n",
-    "        )\n",
-    "        engine.run_till_terminate()\n",
-    "        results = engine.get_output()\n",
-    "        \n",
-    "        return torch.stack(\n",
-    "            [torch.stack([item[\"LAI\"], item[\"TWLV\"]]) for item in results]\n",
-    "        ).unsqueeze(0) # shape: [1, time_steps, 2]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "id": "0dcd85c8-624c-4b58-a7c9-893316110d98",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Step 0, Loss 0.1348, TDWI 0.4242, SPAN 31.2111\n",
-      "Step 10, Loss 0.0589, TDWI 0.4890, SPAN 36.7883\n",
-      "Step 20, Loss 0.0174, TDWI 0.5191, SPAN 34.8344\n",
-      "Step 30, Loss 0.0113, TDWI 0.5056, SPAN 35.0567\n",
-      "Step 40, Loss 0.0067, TDWI 0.5100, SPAN 35.4831\n",
-      "Step 50, Loss 0.0121, TDWI 0.5019, SPAN 34.6534\n",
-      "Step 60, Loss 0.0005, TDWI 0.5038, SPAN 34.6887\n",
-      "Step 70, Loss 0.0079, TDWI 0.5015, SPAN 35.7478\n",
-      "Step 80, Loss 0.0145, TDWI 0.5061, SPAN 34.5574\n",
-      "Step 90, Loss 0.0110, TDWI 0.4985, SPAN 34.3927\n",
-      "Step 100, Loss 0.0271, TDWI 0.5064, SPAN 36.3026\n"
-     ]
-    }
-   ],
-   "source": [
-    "# ----  Create model ---- \n",
-    "opt_model = OptDiffLeafDynamics(\n",
-    "    crop_model_params_provider,\n",
-    "    weather_data_provider,\n",
-    "    agro_management_inputs,\n",
-    "    config_path,\n",
-    "    external_states,\n",
-    ")\n",
-    "\n",
-    "# ----  Optimizer ---- \n",
-    "optimizer = torch.optim.Adam(opt_model.parameters(), lr=0.1)\n",
-    "\n",
-    "# ----  We use relative MAE as loss because there are two outputs with different untis ----  \n",
-    "denom = torch.mean(torch.abs(expected_lai_twlv), dim=1) \n",
-    "\n",
-    "# Training loop (example)\n",
-    "for step in range(101):\n",
-    "    optimizer.zero_grad()\n",
-    "    results = opt_model() \n",
-    "    mae = torch.mean(torch.abs(results - expected_lai_twlv), dim=1)\n",
-    "    rmae = mae / denom\n",
-    "    loss = rmae.sum()  # example: relative mean absolute error\n",
-    "    loss.backward()\n",
-    "    optimizer.step()\n",
-    "\n",
-    "    if step % 10 == 0:\n",
-    "        print(f\"Step {step}, Loss {loss.item():.4f}, TDWI {opt_model.tdwi().item():.4f}, SPAN {opt_model.span().item():.4f}\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "id": "c2d3a463-43a4-4b29-a71f-696c019343d3",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Actual TDWI 0.5100, SPAN 35.0000\n"
-     ]
-    }
-   ],
-   "source": [
-    "# ---- validate the results using test data ---- \n",
-    "print(f\"Actual TDWI {crop_model_params_provider[\"TDWI\"].item():.4f}, SPAN {crop_model_params_provider[\"SPAN\"].item():.4f}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "fedda7e9-02a6-40e5-8c61-08d7886c9519",
-   "metadata": {},
-   "source": [
-    "## 2. Root dynamics \n",
-    "\n",
-    "In this section, we will demonstrate how to optimize two parameters `TWDI` in\n",
-    "root_dynamics model using a differentiable version of root_dynamics.\n",
-    "The optimization will be done using the Adam optimizer from `torch.optim`."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "152ec514-baf1-4213-b98b-76f453d49538",
-   "metadata": {},
-   "source": [
-    "### 2.1 software requirements\n",
-    "\n",
-    "To run this notebook, we need to install the `diffwofost`; the differentiable\n",
-    "version of WOFOST models. Since the package is constantly under development, make\n",
-    "sure you have the latest version of `diffwofost` installed in your\n",
-    "python environment. You can install it using pip:\n",
-    "\n",
-    "```bash\n",
-    "pip install diffwofost\n",
-    "```"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "eaa4c172-9719-4a79-b2f2-d37ea5b6f11d",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ---- import libraries ----\n",
-    "import copy\n",
-    "import torch\n",
-    "import numpy\n",
-    "import yaml\n",
-    "from pathlib import Path\n",
-    "from diffwofost.physical_models.utils import EngineTestHelper\n",
-    "from diffwofost.physical_models.utils import prepare_engine_input"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "26437aed-755d-4ee3-b7b6-82caf8c30ec5",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ---- disable a warning: this will be fixed in the future ----\n",
-    "import warnings\n",
-    "warnings.filterwarnings(\"ignore\", message=\"To copy construct from a tensor.*\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "e245765e-f360-4693-b097-b595360471e3",
-   "metadata": {},
-   "source": [
-    "### 2.2. Data\n",
-    "\n",
-    "A test dataset of `TWRT` (Total weight of roots) will be used to optimize\n",
-    "parametesr `TWDI` (total initial dry weight). Note that in root_dynamic, changes in `TWDI` dont affect `RD` (Current rooting depth). \n",
-    "\n",
-    "The data is stored in PCSE tests folder, and can be doewnloded from PCSE repsository.\n",
-    "You can select any of the files related to `root_dynamics` model with a file name that follwos the pattern\n",
-    "`test_rootdynamics_wofost72_*.yaml`. Each file contains different data depending on the locatin and crop type.\n",
-    "For example, you can download the file \"test_rootdynamics_wofost72_01.yaml\" as:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "id": "22000922-68be-47ed-8afa-2e97c56bb502",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Downloaded: test_rootdynamics_wofost72_01.yaml\n"
-     ]
-    }
-   ],
-   "source": [
-    "import urllib.request\n",
-    "\n",
-    "url = \"https://raw.githubusercontent.com/ajwdewit/pcse/refs/heads/master/tests/test_data/test_rootdynamics_wofost72_01.yaml\"\n",
-    "filename = \"test_rootdynamics_wofost72_01.yaml\"\n",
-    "\n",
-    "urllib.request.urlretrieve(url, filename)\n",
-    "print(f\"Downloaded: {filename}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "c097f6cb-0ded-4b62-844c-795de83df130",
-   "metadata": {},
-   "source": [
-    "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:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "id": "ae2de6dc-1294-4a5a-91d0-3e83f28dc892",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Downloaded: WOFOST_Root_Dynamics.conf\n"
-     ]
-    }
-   ],
-   "source": [
-    "url = \"https://raw.githubusercontent.com/WUR-AI/diffWOFOST/refs/heads/main/tests/physical_models/test_data/WOFOST_Root_Dynamics.conf\"\n",
-    "filename = \"WOFOST_Root_Dynamics.conf\"\n",
-    "\n",
-    "urllib.request.urlretrieve(url, filename)\n",
-    "print(f\"Downloaded: {filename}\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "id": "a69e9279-49eb-4136-8f15-7cff8bb4af52",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ---- Check the path to the files that are downloaded as explained above ----\n",
-    "test_data_path = \"test_rootdynamics_wofost72_01.yaml\"\n",
-    "config_path = \"WOFOST_Root_Dynamics.conf\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "id": "d560476b-64f9-422c-9722-8d0778cfc574",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ---- Here we read the test data and set some variables ----\n",
-    "(crop_model_params_provider, weather_data_provider, agro_management_inputs, external_states) = (\n",
-    "    prepare_engine_input(test_data_path, [\"RDI\", \"RRI\", \"RDMCR\", \"RDMSOL\", \"TDWI\", \"IAIRDU\"])\n",
-    ")\n",
-    "\n",
-    "expected_results = yaml.safe_load(open(test_data_path))[\"ModelResults\"]\n",
-    "expected_twrt = torch.tensor(\n",
-    "    [float(item[\"TWRT\"]) for item in expected_results], dtype=torch.float32\n",
-    ") # shape: [1, time_steps]\n",
-    "\n",
-    "# ---- dont change this: in this config file we specified the diffrentiable version of root_dynamics ----\n",
-    "config_path = str(Path(config_path).resolve()) "
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "d1cef2bd-e42b-4edb-b9b5-b367f776336b",
-   "metadata": {},
-   "source": [
-    "### 2.3. Helper classes/functions\n",
-    "\n",
-    "The model parameters shoudl stay in a valid range. To ensure this, we will use\n",
-    "`BoundedParameter` class with (min, max) and initial values for each\n",
-    "parameter. You might change these values depending on the crop type and\n",
-    "location. But dont use a very small range, otherwise gradiants will be very\n",
-    "small and the optimization will be very slow."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "id": "cbfc9e28-3233-4808-b658-36d6f7d18d75",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ---- Adjust the values if needed  ----\n",
-    "TDWI_MIN, TDWI_MAX, TDWI_INIT = (0.0, 1.0, 0.30)\n",
-    "\n",
-    "# ---- Helper for bounded parameters ----\n",
-    "class BoundedParameter(torch.nn.Module):\n",
-    "    def __init__(self, low, high, init_value):\n",
-    "        super().__init__()\n",
-    "        self.low = low\n",
-    "        self.high = high\n",
-    "\n",
-    "        # Normalize to [0, 1]\n",
-    "        init_norm = (init_value - low) / (high - low)\n",
-    "\n",
-    "        # Parameter in raw logit space\n",
-    "        self.raw = torch.nn.Parameter(torch.logit(torch.tensor(init_norm, dtype=torch.float32), eps=1e-6))\n",
-    "\n",
-    "    def forward(self):\n",
-    "        return self.low + (self.high - self.low) * torch.sigmoid(self.raw)\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "e978cc84-42b4-479e-9817-e5b4ac9fafc8",
-   "metadata": {},
-   "source": [
-    "Another helper class is `OptRootLeafDynamics` which is a subclass of `torch.nn.Module`. \n",
-    "We use this class to wrap the `EngineTestHelper` function and make it easier to run the model `root_dynamic`."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "id": "90296352-4817-4586-842c-12ac6d97d779",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ---- Wrap the model with torch.nn.Module----\n",
-    "class OptDiffRootDynamics(torch.nn.Module):\n",
-    "    def __init__(self, crop_model_params_provider, weather_data_provider, agro_management_inputs, config_path, external_states):\n",
-    "        super().__init__()\n",
-    "        self.crop_model_params_provider = crop_model_params_provider\n",
-    "        self.weather_data_provider = weather_data_provider\n",
-    "        self.agro_management_inputs = agro_management_inputs\n",
-    "        self.config_path = config_path\n",
-    "        self.external_states = external_states\n",
-    "\n",
-    "        # bounded parameters\n",
-    "        self.tdwi = BoundedParameter(TDWI_MIN, TDWI_MAX, init_value=TDWI_INIT)\n",
-    "        \n",
-    "    def forward(self):\n",
-    "        # currently, copying is needed due to an internal issue in engine\n",
-    "        crop_model_params_provider_ = copy.deepcopy(self.crop_model_params_provider)\n",
-    "        external_states_ = copy.deepcopy(self.external_states)\n",
-    "        \n",
-    "        tdwi_val = self.tdwi()\n",
-    "        \n",
-    "        # pass new value of parameters to the model\n",
-    "        crop_model_params_provider_.set_override(\"TDWI\", tdwi_val, check=False)\n",
-    "\n",
-    "        engine = EngineTestHelper(\n",
-    "            crop_model_params_provider_,\n",
-    "            self.weather_data_provider,\n",
-    "            self.agro_management_inputs,\n",
-    "            self.config_path,\n",
-    "            external_states_,\n",
-    "        )\n",
-    "        engine.run_till_terminate()\n",
-    "        results = engine.get_output()\n",
-    "        \n",
-    "        return torch.stack([item[\"TWRT\"] for item in results]) # shape: [1, time_steps]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "id": "f2a1a612-f33c-48a0-a3d4-31318d78b9f4",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ----  Create model ---- \n",
-    "opt_model = OptDiffRootDynamics(\n",
-    "    crop_model_params_provider,\n",
-    "    weather_data_provider,\n",
-    "    agro_management_inputs,\n",
-    "    config_path,\n",
-    "    external_states,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "id": "6c09d1ce-db66-46f3-93d2-1b16f8196e49",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Step 0, Loss 0.00000021, TDWI 0.5340\n",
-      "Step 10, Loss 0.00000015, TDWI 0.5137\n",
-      "Step 20, Loss 0.00000164, TDWI 0.5144\n",
-      "Step 30, Loss 0.00000068, TDWI 0.5107\n",
-      "Step 40, Loss 0.00000019, TDWI 0.5072\n",
-      "Step 50, Loss 0.00000080, TDWI 0.5084\n",
-      "Step 60, Loss 0.00000007, TDWI 0.5071\n",
-      "Step 70, Loss 0.00000100, TDWI 0.5091\n",
-      "Step 80, Loss 0.00000131, TDWI 0.5089\n",
-      "Step 90, Loss 0.00000068, TDWI 0.5115\n",
-      "Step 100, Loss 0.00000071, TDWI 0.5053\n"
-     ]
-    }
-   ],
-   "source": [
-    "# ----  Optimizer ---- \n",
-    "optimizer = torch.optim.Adam(opt_model.parameters(), lr=0.1)\n",
-    "\n",
-    "# ----  We use relative MAE as loss because there are two outputs with different untis ----  \n",
-    "denom = torch.mean(torch.abs(expected_twrt)) \n",
-    "\n",
-    "# Training loop (example)\n",
-    "for step in range(101):\n",
-    "    optimizer.zero_grad()\n",
-    "    results = opt_model() \n",
-    "    mae = torch.mean(torch.abs(results - expected_twrt))\n",
-    "    loss = mae / denom # example: relative mean absolute error\n",
-    "    loss.backward()\n",
-    "    optimizer.step()\n",
-    "\n",
-    "    if step % 10 == 0:\n",
-    "        print(f\"Step {step}, Loss {loss.item():.8f}, TDWI {opt_model.tdwi().item():.4f}\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "id": "2029e2dd-7460-4c71-85a9-b1030a43167b",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Actual TDWI 0.5100\n"
-     ]
-    }
-   ],
-   "source": [
-    "# ---- validate the results using test data ---- \n",
-    "print(f\"Actual TDWI {crop_model_params_provider[\"TDWI\"].item():.4f}\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "a6a511a4-f269-4af4-9f51-2dafa9ba38c0",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.12.11"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}
diff --git a/docs/notebooks/optimization_leaf_dynamics.ipynb b/docs/notebooks/optimization_leaf_dynamics.ipynb
new file mode 100644
index 0000000..d146cf2
--- /dev/null
+++ b/docs/notebooks/optimization_leaf_dynamics.ipynb
@@ -0,0 +1,409 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "c784f1c2-d477-464d-a35f-f2c64e96fb10",
+   "metadata": {},
+   "source": [
+    "<div style=\"display: flex; align-items: center;\">\n",
+    "    <h1>Optimizing parameters in a WOFOST crop model using <code>diffWOFOST</code></h1>\n",
+    "    <img src=\"https://raw.githubusercontent.com/WUR-AI/diffWOFOST/refs/heads/main/docs/logo/diffwofost.png\" width=\"150\" style=\"margin-left: 20px;\">\n",
+    "</div>\n",
+    "\n",
+    "\n",
+    "This Jupyter notebook demonstrates the optimization of parameters in a\n",
+    "differentiable model using the `diffwofost` package. The package provides\n",
+    "differentiable implementations of the WOFOST model and its associated\n",
+    "sub-models. As `diffwofost` is under active development, this notebook focuses on\n",
+    "`leaf_dynamics`. \n",
+    "\n",
+    "To enable these models to operate independently, certain state variables\n",
+    "required by the model are supplied as \"external states\" derived from the test\n",
+    "data. Also, at this stage, only a limited subset of model parameters has been made\n",
+    "differentiable."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "41262fbd-270b-4616-91ad-09ee82451604",
+   "metadata": {},
+   "source": [
+    "## 1. Leaf dynamics\n",
+    "\n",
+    "In this section, we will demonstrate how to optimize two parameters `TWDI` and `SPAN` in\n",
+    "leaf_dynamics model using a differentiable version of leaf_dynamics.\n",
+    "The optimization will be done using the Adam optimizer from `torch.optim`."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1b6c3f53-6fab-4537-9177-7b16e0a1ccec",
+   "metadata": {},
+   "source": [
+    "### 1.1 software requirements\n",
+    "\n",
+    "To run this notebook, we need to install the `diffwofost`; the differentiable\n",
+    "version of WOFOST models. Since the package is constantly under development, make\n",
+    "sure you have the latest version of `diffwofost` installed in your\n",
+    "python environment. You can install it using pip:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e4049fea-1d05-41f1-bf9d-f030ae83a324",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# install diffwofost\n",
+    "!pip install diffwofost"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "21731653-3976-4bb9-b83b-b11d78211700",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- import libraries ----\n",
+    "import copy\n",
+    "import torch\n",
+    "import numpy\n",
+    "import yaml\n",
+    "from pathlib import Path\n",
+    "from diffwofost.physical_models.config import Configuration\n",
+    "from diffwofost.physical_models.crop.leaf_dynamics import WOFOST_Leaf_Dynamics\n",
+    "from diffwofost.physical_models.utils import EngineTestHelper\n",
+    "from diffwofost.physical_models.utils import prepare_engine_input\n",
+    "from diffwofost.physical_models.utils import get_test_data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "82a1ef6b-336e-4902-8bd1-2a1ed2020f9d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- disable a warning: this will be fixed in the future ----\n",
+    "import warnings\n",
+    "warnings.filterwarnings(\"ignore\", message=\"To copy construct from a tensor.*\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "47def7fc-f2dd-4aaf-a572-41cc9d1e4679",
+   "metadata": {},
+   "source": [
+    "### 1.2. Data\n",
+    "\n",
+    "A test dataset of `LAI` (Leaf area index, including stem and pod area) and\n",
+    "`TWLV` (Dry weight of total leaves (living + dead)) will be used to optimize\n",
+    "parametesr `TWDI` (total initial dry weight) and `SPAN` (life span of leaves).\n",
+    "Note that in leaf_dynamic, changes in `SPAN` dont affect `TWLV`. \n",
+    "\n",
+    "The data is stored in PCSE tests folder, and can be doewnloded from PCSE repsository.\n",
+    "You can select any of the files related to `leaf_dynamics` model with a file name that follwos the pattern\n",
+    "`test_leafdynamics_wofost72_*.yaml`. Each file contains different data depending on the locatin and crop type.\n",
+    "For example, you can download the file \"test_leafdynamics_wofost72_01.yaml\" as:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "0233a048-e5a2-4249-887d-35a37284769c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Downloaded: test_leafdynamics_wofost72_01.yaml\n"
+     ]
+    }
+   ],
+   "source": [
+    "import urllib.request\n",
+    "\n",
+    "url = \"https://raw.githubusercontent.com/ajwdewit/pcse/refs/heads/master/tests/test_data/test_leafdynamics_wofost72_01.yaml\"\n",
+    "filename = \"test_leafdynamics_wofost72_01.yaml\"\n",
+    "\n",
+    "urllib.request.urlretrieve(url, filename)\n",
+    "print(f\"Downloaded: {filename}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "5a459489-bfcb-4ad6-9102-1b6be5edeb52",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- Check the path to the files that are downloaded as explained above ----\n",
+    "test_data_path = \"test_leafdynamics_wofost72_01.yaml\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "9f3105fb-4fbe-4405-9fd4-e8255b4b119e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- Here we read the test data and set some variables ----\n",
+    "test_data = get_test_data(test_data_path)\n",
+    "(crop_model_params_provider, weather_data_provider, agro_management_inputs, external_states) = (\n",
+    "    prepare_engine_input(test_data, [\"SPAN\", \"TDWI\", \"TBASE\", \"PERDL\", \"RGRLAI\"])\n",
+    ")\n",
+    "\n",
+    "expected_results = test_data[\"ModelResults\"]\n",
+    "expected_lai_twlv = torch.tensor(\n",
+    "    [[float(item[\"LAI\"]), float(item[\"TWLV\"])] for item in expected_results], dtype=torch.float32\n",
+    ").unsqueeze(0) # shape: [1, time_steps, 2]\n",
+    "\n",
+    "# ---- dont change this: in this config file we specified the diffrentiable version of leaf_dynamics ----\n",
+    "leaf_dynamics_config = Configuration(\n",
+    "    CROP=WOFOST_Leaf_Dynamics,\n",
+    "    OUTPUT_VARS=[\"LAI\", \"TWLV\"],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "52b19ae2-3fe6-4b3f-95a7-aea07a2c0958",
+   "metadata": {},
+   "source": [
+    "### 1.3. Helper classes/functions\n",
+    "\n",
+    "The model parameters shoudl stay in a valid range. To ensure this, we will use\n",
+    "`BoundedParameter` class with (min, max) and initial values for each\n",
+    "parameter. You might change these values depending on the crop type and\n",
+    "location. But dont use a very small range, otherwise gradiants will be very\n",
+    "small and the optimization will be very slow."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "e4610238-de0d-42cf-9689-3c074eb2cc0e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- Adjust the values if needed  ----\n",
+    "TDWI_MIN, TDWI_MAX, TDWI_INIT = (0.0, 1.0, 0.40)\n",
+    "SPAN_MIN, SPAN_MAX, SPAN_INIT = (10.0, 60.0, 25.0)\n",
+    "\n",
+    "# ---- Helper for bounded parameters ----\n",
+    "class BoundedParameter(torch.nn.Module):\n",
+    "    def __init__(self, low, high, init_value):\n",
+    "        super().__init__()\n",
+    "        self.low = low\n",
+    "        self.high = high\n",
+    "\n",
+    "        # Normalize to [0, 1]\n",
+    "        init_norm = (init_value - low) / (high - low)\n",
+    "\n",
+    "        # Parameter in raw logit space\n",
+    "        self.raw = torch.nn.Parameter(torch.logit(torch.tensor(init_norm, dtype=torch.float32), eps=1e-6))\n",
+    "\n",
+    "    def forward(self):\n",
+    "        return self.low + (self.high - self.low) * torch.sigmoid(self.raw)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "153e4306-77ed-4278-8797-a04e637e12c8",
+   "metadata": {},
+   "source": [
+    "Another helper class is `OptDiffLeafDynamics` which is a subclass of `torch.nn.Module`. \n",
+    "We use this class to wrap the `EngineTestHelper` function and make it easier to run the model `leaf_dynamic`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "36dd6463-4812-41c0-b2bf-d4769df1136f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- Wrap the model with torch.nn.Module----\n",
+    "class OptDiffLeafDynamics(torch.nn.Module):\n",
+    "    def __init__(self, crop_model_params_provider, weather_data_provider, agro_management_inputs, leaf_dynamics_config, external_states):\n",
+    "        super().__init__()\n",
+    "        self.crop_model_params_provider = crop_model_params_provider\n",
+    "        self.weather_data_provider = weather_data_provider\n",
+    "        self.agro_management_inputs = agro_management_inputs\n",
+    "        self.config = leaf_dynamics_config\n",
+    "        self.external_states = external_states\n",
+    "\n",
+    "        # bounded parameters\n",
+    "        self.tdwi = BoundedParameter(TDWI_MIN, TDWI_MAX, init_value=TDWI_INIT)\n",
+    "        self.span = BoundedParameter(SPAN_MIN, SPAN_MAX, init_value=SPAN_INIT)\n",
+    "\n",
+    "    def forward(self):\n",
+    "        # currently, copying is needed due to an internal issue in engine\n",
+    "        crop_model_params_provider_ = copy.deepcopy(self.crop_model_params_provider)\n",
+    "        external_states_ = copy.deepcopy(self.external_states)\n",
+    "        \n",
+    "        tdwi_val = self.tdwi()\n",
+    "        span_val = self.span()\n",
+    "        \n",
+    "        # pass new value of parameters to the model\n",
+    "        crop_model_params_provider_.set_override(\"TDWI\", tdwi_val, check=False)\n",
+    "        crop_model_params_provider_.set_override(\"SPAN\", span_val, check=False)\n",
+    "\n",
+    "        engine = EngineTestHelper(\n",
+    "            crop_model_params_provider_,\n",
+    "            self.weather_data_provider,\n",
+    "            self.agro_management_inputs,\n",
+    "            self.config,\n",
+    "            external_states_,\n",
+    "        )\n",
+    "        engine.run_till_terminate()\n",
+    "        results = engine.get_output()\n",
+    "        \n",
+    "        return torch.stack(\n",
+    "            [torch.stack([item[\"LAI\"], item[\"TWLV\"]]) for item in results]\n",
+    "        ).unsqueeze(0) # shape: [1, time_steps, 2]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "3d34c3e8-a8d7-4bc9-94ed-bd2e0234e95c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ----  Create model ---- \n",
+    "opt_model = OptDiffLeafDynamics(\n",
+    "    crop_model_params_provider,\n",
+    "    weather_data_provider,\n",
+    "    agro_management_inputs,\n",
+    "    leaf_dynamics_config,\n",
+    "    external_states,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "78d797f5-4ac4-4380-85f3-6622a7b0f7fb",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 0, Loss 0.0070, TDWI 0.5341, SPAN 36.0670\n",
+      "Step 1, Loss 0.0271, TDWI 0.5325, SPAN 35.8151\n",
+      "Step 2, Loss 0.0216, TDWI 0.5223, SPAN 35.1632\n",
+      "Step 3, Loss 0.0040, TDWI 0.5078, SPAN 34.3827\n",
+      "Step 4, Loss 0.0169, TDWI 0.5028, SPAN 34.0915\n",
+      "Step 5, Loss 0.0262, TDWI 0.5041, SPAN 34.1150\n",
+      "Step 6, Loss 0.0261, TDWI 0.5098, SPAN 34.3638\n",
+      "Step 7, Loss 0.0174, TDWI 0.5189, SPAN 34.7887\n",
+      "Step 8, Loss 0.0083, TDWI 0.5229, SPAN 35.1052\n",
+      "Step 9, Loss 0.0022, TDWI 0.5224, SPAN 35.2232\n",
+      "Step 10, Loss 0.0048, TDWI 0.5183, SPAN 35.1611\n",
+      "Step 11, Loss 0.0040, TDWI 0.5112, SPAN 34.9521\n",
+      "Step 12, Loss 0.0012, TDWI 0.5019, SPAN 34.6612\n",
+      "Step 13, Loss 0.0114, TDWI 0.4970, SPAN 34.5845\n",
+      "Step 14, Loss 0.0126, TDWI 0.4959, SPAN 34.6862\n",
+      "Step 15, Loss 0.0104, TDWI 0.4981, SPAN 34.9379\n",
+      "Step 16, Loss 0.0016, TDWI 0.5030, SPAN 35.2729\n",
+      "Step 17, Loss 0.0059, TDWI 0.5101, SPAN 35.6397\n",
+      "Step 18, Loss 0.0173, TDWI 0.5137, SPAN 35.7822\n",
+      "Step 19, Loss 0.0207, TDWI 0.5140, SPAN 35.7163\n",
+      "Step 20, Loss 0.0194, TDWI 0.5114, SPAN 35.4857\n",
+      "Step 21, Loss 0.0111, TDWI 0.5062, SPAN 35.1246\n",
+      "Step 22, Loss 0.0026, TDWI 0.5043, SPAN 34.8779\n",
+      "Early stopping at step 22\n"
+     ]
+    }
+   ],
+   "source": [
+    "# ----  Early stopping ---- \n",
+    "best_loss = float(\"inf\")\n",
+    "patience = 10  # Number of steps to wait for improvement\n",
+    "patience_counter = 0\n",
+    "min_delta = 1e-4 \n",
+    "\n",
+    "# ----  Optimizer ---- \n",
+    "optimizer = torch.optim.Adam(opt_model.parameters(), lr=0.1)\n",
+    "\n",
+    "# ----  We use relative MAE as loss because there are two outputs with different untis ----  \n",
+    "denom = torch.mean(torch.abs(expected_lai_twlv), dim=1) \n",
+    "\n",
+    "# Training loop (example)\n",
+    "for step in range(101):\n",
+    "    optimizer.zero_grad()\n",
+    "    results = opt_model() \n",
+    "    mae = torch.mean(torch.abs(results - expected_lai_twlv), dim=1)\n",
+    "    rmae = mae / denom\n",
+    "    loss = rmae.sum()  # example: relative mean absolute error\n",
+    "    loss.backward()\n",
+    "    optimizer.step()\n",
+    "\n",
+    "    print(f\"Step {step}, Loss {loss.item():.4f}, TDWI {opt_model.tdwi().item():.4f}, SPAN {opt_model.span().item():.4f}\")\n",
+    "    # Early stopping logic\n",
+    "    if loss.item() < best_loss - min_delta:\n",
+    "        best_loss = loss.item()\n",
+    "        patience_counter = 0\n",
+    "    else:\n",
+    "        patience_counter += 1\n",
+    "        if patience_counter >= patience:\n",
+    "            print(f\"Early stopping at step {step}\")\n",
+    "            break"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "c2d3a463-43a4-4b29-a71f-696c019343d3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Actual TDWI 0.5100, SPAN 35.0000\n"
+     ]
+    }
+   ],
+   "source": [
+    "# ---- validate the results using test data ---- \n",
+    "print(f\"Actual TDWI {crop_model_params_provider[\"TDWI\"].item():.4f}, SPAN {crop_model_params_provider[\"SPAN\"].item():.4f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a6a511a4-f269-4af4-9f51-2dafa9ba38c0",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/docs/notebooks/optimization_root_dynamics.ipynb b/docs/notebooks/optimization_root_dynamics.ipynb
new file mode 100644
index 0000000..90a0d01
--- /dev/null
+++ b/docs/notebooks/optimization_root_dynamics.ipynb
@@ -0,0 +1,389 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "c784f1c2-d477-464d-a35f-f2c64e96fb10",
+   "metadata": {},
+   "source": [
+    "<div style=\"display: flex; align-items: center;\">\n",
+    "    <h1>Optimizing parameters in a WOFOST crop model using <code>diffWOFOST</code></h1>\n",
+    "    <img src=\"https://raw.githubusercontent.com/WUR-AI/diffWOFOST/refs/heads/main/docs/logo/diffwofost.png\" width=\"150\" style=\"margin-left: 20px;\">\n",
+    "</div>\n",
+    "\n",
+    "\n",
+    "This Jupyter notebook demonstrates the optimization of parameters in a\n",
+    "differentiable model using the `diffwofost` package. The package provides\n",
+    "differentiable implementations of the WOFOST model and its associated\n",
+    "sub-models. As `diffwofost` is under active development, this notebook focuses on\n",
+    "`root_dynamics`. \n",
+    "\n",
+    "To enable these models to operate independently, certain state variables\n",
+    "required by the model are supplied as \"external states\" derived from the test\n",
+    "data. Also, at this stage, only a limited subset of model parameters has been made\n",
+    "differentiable."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fedda7e9-02a6-40e5-8c61-08d7886c9519",
+   "metadata": {},
+   "source": [
+    "## 1. Root dynamics \n",
+    "\n",
+    "In this section, we will demonstrate how to optimize two parameters `TWDI` in\n",
+    "root_dynamics model using a differentiable version of root_dynamics.\n",
+    "The optimization will be done using the Adam optimizer from `torch.optim`."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "152ec514-baf1-4213-b98b-76f453d49538",
+   "metadata": {},
+   "source": [
+    "### 1.1 software requirements\n",
+    "\n",
+    "To run this notebook, we need to install the `diffwofost`; the differentiable\n",
+    "version of WOFOST models. Since the package is constantly under development, make\n",
+    "sure you have the latest version of `diffwofost` installed in your\n",
+    "python environment. You can install it using pip:\n",
+    "\n",
+    "```bash\n",
+    "pip install diffwofost\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "eaa4c172-9719-4a79-b2f2-d37ea5b6f11d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- import libraries ----\n",
+    "import copy\n",
+    "import torch\n",
+    "import numpy\n",
+    "import yaml\n",
+    "from pathlib import Path\n",
+    "from diffwofost.physical_models.config import Configuration\n",
+    "from diffwofost.physical_models.crop.root_dynamics import WOFOST_Root_Dynamics\n",
+    "from diffwofost.physical_models.utils import EngineTestHelper\n",
+    "from diffwofost.physical_models.utils import prepare_engine_input\n",
+    "from diffwofost.physical_models.utils import get_test_data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "26437aed-755d-4ee3-b7b6-82caf8c30ec5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- disable a warning: this will be fixed in the future ----\n",
+    "import warnings\n",
+    "warnings.filterwarnings(\"ignore\", message=\"To copy construct from a tensor.*\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e245765e-f360-4693-b097-b595360471e3",
+   "metadata": {},
+   "source": [
+    "### 1.2. Data\n",
+    "\n",
+    "A test dataset of `TWRT` (Total weight of roots) will be used to optimize\n",
+    "parametesr `TWDI` (total initial dry weight). Note that in root_dynamic, changes in `TWDI` dont affect `RD` (Current rooting depth). \n",
+    "\n",
+    "The data is stored in PCSE tests folder, and can be doewnloded from PCSE repsository.\n",
+    "You can select any of the files related to `root_dynamics` model with a file name that follwos the pattern\n",
+    "`test_rootdynamics_wofost72_*.yaml`. Each file contains different data depending on the locatin and crop type.\n",
+    "For example, you can download the file \"test_rootdynamics_wofost72_01.yaml\" as:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "22000922-68be-47ed-8afa-2e97c56bb502",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Downloaded: test_rootdynamics_wofost72_01.yaml\n"
+     ]
+    }
+   ],
+   "source": [
+    "import urllib.request\n",
+    "\n",
+    "url = \"https://raw.githubusercontent.com/ajwdewit/pcse/refs/heads/master/tests/test_data/test_rootdynamics_wofost72_01.yaml\"\n",
+    "filename = \"test_rootdynamics_wofost72_01.yaml\"\n",
+    "\n",
+    "urllib.request.urlretrieve(url, filename)\n",
+    "print(f\"Downloaded: {filename}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c097f6cb-0ded-4b62-844c-795de83df130",
+   "metadata": {},
+   "source": [
+    "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:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "a69e9279-49eb-4136-8f15-7cff8bb4af52",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- Check the path to the files that are downloaded as explained above ----\n",
+    "test_data_path = \"test_rootdynamics_wofost72_01.yaml\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "d560476b-64f9-422c-9722-8d0778cfc574",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- Here we read the test data and set some variables ----\n",
+    "test_data = get_test_data(test_data_path)\n",
+    "(crop_model_params_provider, weather_data_provider, agro_management_inputs, external_states) = (\n",
+    "    prepare_engine_input(test_data, [\"RDI\", \"RRI\", \"RDMCR\", \"RDMSOL\", \"TDWI\", \"IAIRDU\"])\n",
+    ")\n",
+    "\n",
+    "expected_results = test_data[\"ModelResults\"]\n",
+    "expected_twrt = torch.tensor(\n",
+    "    [float(item[\"TWRT\"]) for item in expected_results], dtype=torch.float32\n",
+    ") # shape: [1, time_steps]\n",
+    "\n",
+    "# ---- dont change this: in this config file we specified the diffrentiable version of root_dynamics ----\n",
+    "root_dynamics_config = Configuration(\n",
+    "    CROP=WOFOST_Root_Dynamics,\n",
+    "    OUTPUT_VARS=[\"RD\", \"TWRT\"],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d1cef2bd-e42b-4edb-b9b5-b367f776336b",
+   "metadata": {},
+   "source": [
+    "### 1.3. Helper classes/functions\n",
+    "\n",
+    "The model parameters shoudl stay in a valid range. To ensure this, we will use\n",
+    "`BoundedParameter` class with (min, max) and initial values for each\n",
+    "parameter. You might change these values depending on the crop type and\n",
+    "location. But dont use a very small range, otherwise gradiants will be very\n",
+    "small and the optimization will be very slow."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "cbfc9e28-3233-4808-b658-36d6f7d18d75",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- Adjust the values if needed  ----\n",
+    "TDWI_MIN, TDWI_MAX, TDWI_INIT = (0.0, 1.0, 0.30)\n",
+    "\n",
+    "# ---- Helper for bounded parameters ----\n",
+    "class BoundedParameter(torch.nn.Module):\n",
+    "    def __init__(self, low, high, init_value):\n",
+    "        super().__init__()\n",
+    "        self.low = low\n",
+    "        self.high = high\n",
+    "\n",
+    "        # Normalize to [0, 1]\n",
+    "        init_norm = (init_value - low) / (high - low)\n",
+    "\n",
+    "        # Parameter in raw logit space\n",
+    "        self.raw = torch.nn.Parameter(torch.logit(torch.tensor(init_norm, dtype=torch.float32), eps=1e-6))\n",
+    "\n",
+    "    def forward(self):\n",
+    "        return self.low + (self.high - self.low) * torch.sigmoid(self.raw)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e978cc84-42b4-479e-9817-e5b4ac9fafc8",
+   "metadata": {},
+   "source": [
+    "Another helper class is `OptRootDynamics` which is a subclass of `torch.nn.Module`. \n",
+    "We use this class to wrap the `EngineTestHelper` function and make it easier to run the model `root_dynamic`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "90296352-4817-4586-842c-12ac6d97d779",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- Wrap the model with torch.nn.Module----\n",
+    "class OptDiffRootDynamics(torch.nn.Module):\n",
+    "    def __init__(self, crop_model_params_provider, weather_data_provider, agro_management_inputs, root_dynamics_config, external_states):\n",
+    "        super().__init__()\n",
+    "        self.crop_model_params_provider = crop_model_params_provider\n",
+    "        self.weather_data_provider = weather_data_provider\n",
+    "        self.agro_management_inputs = agro_management_inputs\n",
+    "        self.config = root_dynamics_config\n",
+    "        self.external_states = external_states\n",
+    "\n",
+    "        # bounded parameters\n",
+    "        self.tdwi = BoundedParameter(TDWI_MIN, TDWI_MAX, init_value=TDWI_INIT)\n",
+    "        \n",
+    "    def forward(self):\n",
+    "        # currently, copying is needed due to an internal issue in engine\n",
+    "        crop_model_params_provider_ = copy.deepcopy(self.crop_model_params_provider)\n",
+    "        external_states_ = copy.deepcopy(self.external_states)\n",
+    "        \n",
+    "        tdwi_val = self.tdwi()\n",
+    "        \n",
+    "        # pass new value of parameters to the model\n",
+    "        crop_model_params_provider_.set_override(\"TDWI\", tdwi_val, check=False)\n",
+    "\n",
+    "        engine = EngineTestHelper(\n",
+    "            crop_model_params_provider_,\n",
+    "            self.weather_data_provider,\n",
+    "            self.agro_management_inputs,\n",
+    "            self.config,\n",
+    "            external_states_,\n",
+    "        )\n",
+    "        engine.run_till_terminate()\n",
+    "        results = engine.get_output()\n",
+    "        \n",
+    "        return torch.stack([item[\"TWRT\"] for item in results]) # shape: [1, time_steps]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "f2a1a612-f33c-48a0-a3d4-31318d78b9f4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ----  Create model ---- \n",
+    "opt_model = OptDiffRootDynamics(\n",
+    "    crop_model_params_provider,\n",
+    "    weather_data_provider,\n",
+    "    agro_management_inputs,\n",
+    "    root_dynamics_config,\n",
+    "    external_states,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "6c09d1ce-db66-46f3-93d2-1b16f8196e49",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 0, Loss 0.00000787, TDWI 0.5207\n",
+      "Step 1, Loss 0.00000237, TDWI 0.4957\n",
+      "Step 2, Loss 0.00000316, TDWI 0.4892\n",
+      "Step 3, Loss 0.00000459, TDWI 0.4919\n",
+      "Step 4, Loss 0.00000400, TDWI 0.5000\n",
+      "Step 5, Loss 0.00000220, TDWI 0.5118\n",
+      "Step 6, Loss 0.00000040, TDWI 0.5165\n",
+      "Step 7, Loss 0.00000143, TDWI 0.5159\n",
+      "Step 8, Loss 0.00000132, TDWI 0.5114\n",
+      "Step 9, Loss 0.00000032, TDWI 0.5038\n",
+      "Step 10, Loss 0.00000137, TDWI 0.5009\n",
+      "Early stopping at step 10\n"
+     ]
+    }
+   ],
+   "source": [
+    "# ----  Early stopping ---- \n",
+    "best_loss = float(\"inf\")\n",
+    "patience = 10  # Number of steps to wait for improvement\n",
+    "patience_counter = 0\n",
+    "min_delta = 1e-4 \n",
+    "\n",
+    "# ----  Optimizer ---- \n",
+    "optimizer = torch.optim.Adam(opt_model.parameters(), lr=0.1)\n",
+    "\n",
+    "# ----  We use relative MAE as loss because there are two outputs with different untis ----  \n",
+    "denom = torch.mean(torch.abs(expected_twrt)) \n",
+    "\n",
+    "# Training loop (example)\n",
+    "for step in range(101):\n",
+    "    optimizer.zero_grad()\n",
+    "    results = opt_model() \n",
+    "    mae = torch.mean(torch.abs(results - expected_twrt))\n",
+    "    loss = mae / denom # example: relative mean absolute error\n",
+    "    loss.backward()\n",
+    "    optimizer.step()\n",
+    "\n",
+    "    print(f\"Step {step}, Loss {loss.item():.8f}, TDWI {opt_model.tdwi().item():.4f}\")\n",
+    "    # Early stopping logic\n",
+    "    if loss.item() < best_loss - min_delta:\n",
+    "        best_loss = loss.item()\n",
+    "        patience_counter = 0\n",
+    "    else:\n",
+    "        patience_counter += 1\n",
+    "        if patience_counter >= patience:\n",
+    "            print(f\"Early stopping at step {step}\")\n",
+    "            break"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "2029e2dd-7460-4c71-85a9-b1030a43167b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Actual TDWI 0.5100\n"
+     ]
+    }
+   ],
+   "source": [
+    "# ---- validate the results using test data ---- \n",
+    "print(f\"Actual TDWI {crop_model_params_provider[\"TDWI\"].item():.4f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a6a511a4-f269-4af4-9f51-2dafa9ba38c0",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/src/diffwofost/physical_models/crop/leaf_dynamics.py b/src/diffwofost/physical_models/crop/leaf_dynamics.py
index 4363670..99b2182 100644
--- a/src/diffwofost/physical_models/crop/leaf_dynamics.py
+++ b/src/diffwofost/physical_models/crop/leaf_dynamics.py
@@ -111,7 +111,7 @@ class WOFOST_Leaf_Dynamics(SimulationObject):
 
     [!NOTE]
     Notice that the following gradients are zero:
-        - ∂SPAN/∂LAI
+        - ∂SPAN/∂TWLV
         - ∂PERDL/∂TWLV
         - ∂KDIFTB/∂LAI
     """  # noqa: E501
@@ -288,16 +288,27 @@ def calc_rates(self, day: datetime.date, drv: WeatherDataContainer) -> None:
         # SPAN because the latter would not allow for the gradient to be tracked.
         # the if statement `p.SPAN.requires_grad` to avoid unnecessary
         # approximation when SPAN is not a learnable parameter.
+        # here we use STE (straight through estimator) method.
         # TODO: sharpness can be exposed as a parameter
         if p.SPAN.requires_grad:
-            # 1e-16 is chosen empirically for cases when s.LVAGE - tSPAN is very
-            # small and mask should be 1
-            sharpness = torch.tensor(1e-16, dtype=DTYPE)
+            # 1000 is chosen empirically to approximate a step function
+            sharpness = torch.tensor(1000, dtype=DTYPE)
 
             # 1e-14 is chosen empirically for cases when s.LVAGE - tSPAN is
             # equal to zero and mask should be 0.0
             epsilon = 1e-14
-            span_mask = torch.sigmoid((s.LVAGE - tSPAN - epsilon) / sharpness).to(dtype=DTYPE)
+
+            # soft mask using sigmoid
+            soft_mask = torch.sigmoid((s.LVAGE - tSPAN - epsilon) / sharpness)
+
+            # originial hard mask
+            hard_mask = (s.LVAGE > tSPAN).to(DTYPE)
+
+            # STE method. Here detach is used to stop the gradient flow. This
+            # way, during backpropagation, the gradient is computed only through
+            # the `soft_mask``, while during the forward pass, the `hard_mask``
+            # is used.
+            span_mask = hard_mask.detach() + soft_mask - soft_mask.detach()
         else:
             span_mask = (s.LVAGE > tSPAN).to(dtype=DTYPE)
 
diff --git a/tests/physical_models/crop/test_leaf_dynamics.py b/tests/physical_models/crop/test_leaf_dynamics.py
index d1dd78b..4535e9d 100644
--- a/tests/physical_models/crop/test_leaf_dynamics.py
+++ b/tests/physical_models/crop/test_leaf_dynamics.py
@@ -3,7 +3,6 @@
 from unittest.mock import patch
 import pytest
 import torch
-from numpy.testing import assert_array_almost_equal
 from pcse.models import Wofost72_PP
 from diffwofost.physical_models.config import Configuration
 from diffwofost.physical_models.crop.leaf_dynamics import WOFOST_Leaf_Dynamics
@@ -419,7 +418,11 @@ def test_wofost_pp_with_leaf_dynamics(self, test_data_url):
 
     @pytest.mark.parametrize("test_data_url", leafdynamics_data_urls)
     def test_leaf_dynamics_with_sigmoid_approx(self, test_data_url):
-        """Test if sigmoid approximation gives same results as leaf dynamics."""
+        """Test if sigmoid approximation gives same results as leaf dynamics.
+
+        This should be the case since WOFOST_Leaf_Dynamics uses STE method.
+        In practice, no approximation is done when not insterested in gradients.
+        """
         # prepare model input
         test_data = get_test_data(test_data_url)
         crop_model_params = ["SPAN", "TDWI", "TBASE", "PERDL", "RGRLAI", "KDIFTB", "SLATB"]
@@ -565,8 +568,11 @@ def test_gradients_numerical(self, param_name, output_name, config_type):
         """Test that analytical gradients match numerical gradients."""
         value, _ = self.param_configs[config_type][param_name]
         param = torch.nn.Parameter(torch.tensor(value, dtype=torch.float64))
+
+        # we pass `param` and not `param.data` because we want `requires_grad=True`
+        # for parameter `SPAN`
         numerical_grad = calculate_numerical_grad(
-            get_test_diff_leaf_model, param_name, param.data, output_name
+            get_test_diff_leaf_model, param_name, param, output_name
         )
 
         model = get_test_diff_leaf_model()
@@ -576,7 +582,16 @@ def test_gradients_numerical(self, param_name, output_name, config_type):
         # this is ∂loss/∂param, for comparison with numerical gradient
         grads = torch.autograd.grad(loss, param, retain_graph=True)[0]
 
-        assert_array_almost_equal(numerical_grad, grads.data, decimal=3)
+        # for span, the numerical gradient can't be equal to the pytorch one
+        # because we are using STE method
+        if (param_name, output_name) not in {("SPAN", "LAI")}:
+            # assert_array_almost_equal(numerical_grad, grads.data, decimal=3)
+            torch.testing.assert_close(
+                numerical_grad,
+                grads,
+                rtol=1e-3,
+                atol=1e-3,
+            )
 
         # Warn if gradient is zero (but this shouldn't happen for gradient_params)
         if torch.all(grads == 0):