Skip to content

Use a single device/dtype #86

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
SCiarella wants to merge 5 commits into
base: main
Choose a base branch
from
+337 −250
Open

Use a single device/dtype #86

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
54a7a04
Use a single device/dtype
SCiarella Feb 3, 2026
c1f2d47
fix
SCiarella Feb 3, 2026
d5237ae
Update notebook devices
SCiarella Feb 3, 2026
367912b
Update ComputeConfig docs
SCiarella Feb 3, 2026
c72de2d
Default dtype/device to pytorch
SCiarella Feb 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
144 changes: 98 additions & 46 deletions docs/notebooks/optimization_leaf_dynamics.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,18 +50,73 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 1,
"id": "e4049fea-1d05-41f1-bf9d-f030ae83a324",
"metadata": {},
"outputs": [],
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Requirement already satisfied: diffwofost in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (0.2.0)\n",
"Requirement already satisfied: torch in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from diffwofost) (2.9.0)\n",
"Requirement already satisfied: pcse in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from diffwofost) (6.0.9)\n",
"Requirement already satisfied: SQLAlchemy<2.0,>=1.3.0 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from pcse->diffwofost) (1.4.54)\n",
"Requirement already satisfied: PyYAML>=5.1 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from pcse->diffwofost) (6.0.3)\n",
"Requirement already satisfied: openpyxl>=3.0.0 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from pcse->diffwofost) (3.1.5)\n",
"Requirement already satisfied: requests>=2.0.0 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from pcse->diffwofost) (2.32.5)\n",
"Requirement already satisfied: pandas>=0.25 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from pcse->diffwofost) (2.3.3)\n",
"Requirement already satisfied: traitlets-pcse==5.0.0.dev in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from pcse->diffwofost) (5.0.0.dev0)\n",
"Requirement already satisfied: dotmap>=1.3 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from pcse->diffwofost) (1.3.30)\n",
"Requirement already satisfied: ipython_genutils in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from traitlets-pcse==5.0.0.dev->pcse->diffwofost) (0.2.0)\n",
"Requirement already satisfied: six in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from traitlets-pcse==5.0.0.dev->pcse->diffwofost) (1.17.0)\n",
"Requirement already satisfied: decorator in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from traitlets-pcse==5.0.0.dev->pcse->diffwofost) (5.2.1)\n",
"Requirement already satisfied: filelock in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (3.20.0)\n",
"Requirement already satisfied: typing-extensions>=4.10.0 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (4.15.0)\n",
"Requirement already satisfied: setuptools in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (80.9.0)\n",
"Requirement already satisfied: sympy>=1.13.3 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (1.14.0)\n",
"Requirement already satisfied: networkx>=2.5.1 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (3.5)\n",
"Requirement already satisfied: jinja2 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (3.1.6)\n",
"Requirement already satisfied: fsspec>=0.8.5 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (2025.9.0)\n",
"Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.8.93 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (12.8.93)\n",
"Requirement already satisfied: nvidia-cuda-runtime-cu12==12.8.90 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (12.8.90)\n",
"Requirement already satisfied: nvidia-cuda-cupti-cu12==12.8.90 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (12.8.90)\n",
"Requirement already satisfied: nvidia-cudnn-cu12==9.10.2.21 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (9.10.2.21)\n",
"Requirement already satisfied: nvidia-cublas-cu12==12.8.4.1 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (12.8.4.1)\n",
"Requirement already satisfied: nvidia-cufft-cu12==11.3.3.83 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (11.3.3.83)\n",
"Requirement already satisfied: nvidia-curand-cu12==10.3.9.90 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (10.3.9.90)\n",
"Requirement already satisfied: nvidia-cusolver-cu12==11.7.3.90 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (11.7.3.90)\n",
"Requirement already satisfied: nvidia-cusparse-cu12==12.5.8.93 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (12.5.8.93)\n",
"Requirement already satisfied: nvidia-cusparselt-cu12==0.7.1 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (0.7.1)\n",
"Requirement already satisfied: nvidia-nccl-cu12==2.27.5 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (2.27.5)\n",
"Requirement already satisfied: nvidia-nvshmem-cu12==3.3.20 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (3.3.20)\n",
"Requirement already satisfied: nvidia-nvtx-cu12==12.8.90 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (12.8.90)\n",
"Requirement already satisfied: nvidia-nvjitlink-cu12==12.8.93 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (12.8.93)\n",
"Requirement already satisfied: nvidia-cufile-cu12==1.13.1.3 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (1.13.1.3)\n",
"Requirement already satisfied: triton==3.5.0 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from torch->diffwofost) (3.5.0)\n",
"Requirement already satisfied: et-xmlfile in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from openpyxl>=3.0.0->pcse->diffwofost) (2.0.0)\n",
"Requirement already satisfied: numpy>=1.26.0 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from pandas>=0.25->pcse->diffwofost) (2.3.4)\n",
"Requirement already satisfied: python-dateutil>=2.8.2 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from pandas>=0.25->pcse->diffwofost) (2.9.0.post0)\n",
"Requirement already satisfied: pytz>=2020.1 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from pandas>=0.25->pcse->diffwofost) (2025.2)\n",
"Requirement already satisfied: tzdata>=2022.7 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from pandas>=0.25->pcse->diffwofost) (2025.2)\n",
"Requirement already satisfied: charset_normalizer<4,>=2 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from requests>=2.0.0->pcse->diffwofost) (3.4.4)\n",
"Requirement already satisfied: idna<4,>=2.5 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from requests>=2.0.0->pcse->diffwofost) (3.11)\n",
"Requirement already satisfied: urllib3<3,>=1.21.1 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from requests>=2.0.0->pcse->diffwofost) (2.5.0)\n",
"Requirement already satisfied: certifi>=2017.4.17 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from requests>=2.0.0->pcse->diffwofost) (2025.10.5)\n",
"Requirement already satisfied: greenlet!=0.4.17 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from SQLAlchemy<2.0,>=1.3.0->pcse->diffwofost) (3.2.4)\n",
"Requirement already satisfied: mpmath<1.4,>=1.1.0 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from sympy>=1.13.3->torch->diffwofost) (1.3.0)\n",
"Requirement already satisfied: MarkupSafe>=2.0 in /home/simone/.virtualenvs/dwof/lib/python3.12/site-packages (from jinja2->torch->diffwofost) (3.0.3)\n"
]
}
],
"source": [
"# install diffwofost\n",
"!pip install diffwofost"
]
},
{
"cell_type": "code",
"execution_count": 1,
"execution_count": 2,
"id": "21731653-3976-4bb9-b83b-b11d78211700",
"metadata": {},
"outputs": [],
Expand All @@ -72,7 +127,7 @@
"import numpy\n",
"import yaml\n",
"from pathlib import Path\n",
"from diffwofost.physical_models.config import Configuration\n",
"from diffwofost.physical_models.config import Configuration, ComputeConfig\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",
Expand All @@ -81,19 +136,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"id": "9c2ca761-22b0-4707-9f55-619f62617d12",
"metadata": {},
"outputs": [],
"source": [
"# --- run on CPU ------\n",
"from diffwofost.physical_models.config import ComputeConfig\n",
"ComputeConfig.set_device('cpu')"
]
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": 3,
"id": "82a1ef6b-336e-4902-8bd1-2a1ed2020f9d",
"metadata": {},
"outputs": [],
Expand Down Expand Up @@ -123,7 +166,7 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": 4,
"id": "0233a048-e5a2-4249-887d-35a37284769c",
"metadata": {},
"outputs": [
Expand All @@ -147,7 +190,7 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": 5,
"id": "5a459489-bfcb-4ad6-9102-1b6be5edeb52",
"metadata": {},
"outputs": [],
Expand All @@ -158,7 +201,7 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": 6,
"id": "9f3105fb-4fbe-4405-9fd4-e8255b4b119e",
"metadata": {},
"outputs": [],
Expand All @@ -171,7 +214,9 @@
"\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",
" [[float(item[\"LAI\"]), float(item[\"TWLV\"])] for item in expected_results],\n",
" dtype=ComputeConfig.get_dtype(),\n",
" device=ComputeConfig.get_device(),\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",
Expand All @@ -197,7 +242,7 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": 7,
"id": "e4610238-de0d-42cf-9689-3c074eb2cc0e",
"metadata": {},
"outputs": [],
Expand All @@ -217,10 +262,17 @@
" 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",
" self.raw = torch.nn.Parameter(\n",
" torch.logit(\n",
" torch.tensor(\n",
" init_norm, dtype=ComputeConfig.get_dtype(), device=ComputeConfig.get_device()\n",
" ),\n",
" eps=1e-6,\n",
" )\n",
" )\n",
"\n",
" def forward(self):\n",
" return self.low + (self.high - self.low) * torch.sigmoid(self.raw)\n"
" return self.low + (self.high - self.low) * torch.sigmoid(self.raw)"
]
},
{
Expand All @@ -234,7 +286,7 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": 8,
"id": "36dd6463-4812-41c0-b2bf-d4769df1136f",
"metadata": {},
"outputs": [],
Expand Down Expand Up @@ -282,7 +334,7 @@
},
{
"cell_type": "code",
"execution_count": 8,
"execution_count": 9,
"id": "3d34c3e8-a8d7-4bc9-94ed-bd2e0234e95c",
"metadata": {},
"outputs": [],
Expand All @@ -299,7 +351,7 @@
},
{
"cell_type": "code",
"execution_count": 9,
"execution_count": 10,
"id": "78d797f5-4ac4-4380-85f3-6622a7b0f7fb",
"metadata": {},
"outputs": [
Expand All @@ -312,22 +364,22 @@
"Step 2, Loss 0.2117, TDWI 0.4727, SPAN 28.3303\n",
"Step 3, Loss 0.1790, TDWI 0.4962, SPAN 29.5140\n",
"Step 4, Loss 0.1502, TDWI 0.5190, SPAN 30.7290\n",
"Step 5, Loss 0.1148, TDWI 0.5351, SPAN 31.8314\n",
"Step 6, Loss 0.0873, TDWI 0.5459, SPAN 32.8227\n",
"Step 7, Loss 0.0633, TDWI 0.5525, SPAN 33.6916\n",
"Step 8, Loss 0.0380, TDWI 0.5559, SPAN 34.4434\n",
"Step 9, Loss 0.0164, TDWI 0.5564, SPAN 35.0537\n",
"Step 10, Loss 0.0016, TDWI 0.5543, SPAN 35.2340\n",
"Step 11, Loss 0.0050, TDWI 0.5500, SPAN 35.0621\n",
"Step 12, Loss 0.0018, TDWI 0.5436, SPAN 34.6755\n",
"Step 13, Loss 0.0109, TDWI 0.5356, SPAN 34.2452\n",
"Step 14, Loss 0.0230, TDWI 0.5261, SPAN 33.8167\n",
"Step 15, Loss 0.0341, TDWI 0.5152, SPAN 33.4017\n",
"Step 16, Loss 0.0472, TDWI 0.5031, SPAN 33.0102\n",
"Step 17, Loss 0.0559, TDWI 0.4950, SPAN 32.9639\n",
"Step 18, Loss 0.0571, TDWI 0.4903, SPAN 33.1766\n",
"Step 19, Loss 0.0519, TDWI 0.4888, SPAN 33.5989\n",
"Step 20, Loss 0.0400, TDWI 0.4899, SPAN 34.1872\n",
"Step 5, Loss 0.1148, TDWI 0.5351, SPAN 31.8327\n",
"Step 6, Loss 0.0873, TDWI 0.5459, SPAN 32.8254\n",
"Step 7, Loss 0.0633, TDWI 0.5525, SPAN 33.6958\n",
"Step 8, Loss 0.0375, TDWI 0.5559, SPAN 34.4362\n",
"Step 9, Loss 0.0164, TDWI 0.5564, SPAN 35.0038\n",
"Step 10, Loss 0.0001, TDWI 0.5543, SPAN 35.2185\n",
"Step 11, Loss 0.0048, TDWI 0.5500, SPAN 35.0892\n",
"Step 12, Loss 0.0019, TDWI 0.5436, SPAN 34.7440\n",
"Step 13, Loss 0.0092, TDWI 0.5356, SPAN 34.3535\n",
"Step 14, Loss 0.0175, TDWI 0.5260, SPAN 33.9622\n",
"Step 15, Loss 0.0296, TDWI 0.5151, SPAN 33.5733\n",
"Step 16, Loss 0.0409, TDWI 0.5030, SPAN 33.2059\n",
"Step 17, Loss 0.0512, TDWI 0.4949, SPAN 33.2068\n",
"Step 18, Loss 0.0512, TDWI 0.4903, SPAN 33.4762\n",
"Step 19, Loss 0.0422, TDWI 0.4888, SPAN 33.9525\n",
"Step 20, Loss 0.0299, TDWI 0.4900, SPAN 34.5985\n",
"Early stopping at step 20\n"
]
}
Expand Down Expand Up @@ -369,7 +421,7 @@
},
{
"cell_type": "code",
"execution_count": 10,
"execution_count": 11,
"id": "c2d3a463-43a4-4b29-a71f-696c019343d3",
"metadata": {},
"outputs": [
Expand Down Expand Up @@ -397,7 +449,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"display_name": "dwof",
"language": "python",
"name": "python3"
},
Expand All @@ -411,7 +463,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.11"
"version": "3.12.3"
}
},
"nbformat": 4,
Expand Down
Loading
Loading