Facilitating support for tensor weather data

docs/api_reference.md

@@ -40,6 +40,8 @@ hide:
 
 ::: diffwofost.physical_models.test.EngineTestHelper
 
+::: diffwofost.physical_models.weather.to_weather_data_iterator
+
 ::: diffwofost.io.save_model
 
 ::: diffwofost.io.load_model

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     "safetensors",
     "torch",
     "pcse",
+    "xarray",
 ]
 description = "Differentiable WOFOST"
 keywords = ["wofost", "pytorch", "differentiable", "crop", "optimization"]

src/diffwofost/io.py

@@ -27,6 +27,10 @@
 from diffwofost.physical_models.config import Configuration
 from diffwofost.physical_models.config import load_class
 
+# Use the temporary ParameterProvider. The original PCSE ParameterProvider could be used when
+# https://github.com/ajwdewit/pcse/pull/121 is merged an a new version of PCSE released.
+from diffwofost.physical_models.parameter_providers import ParameterProvider
+
 
 def _default_model_filename(model):
     """Stable filename from class name + init_kwargs hash."""
@@ -224,8 +228,6 @@ def load_model(path, *, model_class=None, device=None, dtype=None):
                     param_data[group_name]["nontensor_params"] = params
 
         # Reconstruct ParameterProvider
-        from pcse.base.parameter_providers import ParameterProvider
-
         init_kwargs = {}
         for group_name in ("sitedata", "timerdata", "soildata", "cropdata"):
             key = f"_{group_name}"

src/diffwofost/physical_models/crop/assimilation.py

@@ -6,14 +6,12 @@
 from pcse.base import SimulationObject
 from pcse.base.parameter_providers import ParameterProvider
 from pcse.base.variablekiosk import VariableKiosk
-from pcse.base.weather import WeatherDataContainer
 from diffwofost.physical_models.base import TensorParamTemplate
 from diffwofost.physical_models.base import TensorRatesTemplate
 from diffwofost.physical_models.config import ComputeConfig
 from diffwofost.physical_models.traitlets import Tensor
 from diffwofost.physical_models.utils import AfgenTrait
 from diffwofost.physical_models.utils import _broadcast_to
-from diffwofost.physical_models.utils import _get_drv
 from diffwofost.physical_models.utils import astro
 
 # ---------------------------------------------------------------------------
@@ -343,7 +341,7 @@ def initialize(
         # elements (which share the same weather driver).
         self._astro_cache: dict = {}
 
-    def calc_rates(self, day: datetime.date = None, drv: WeatherDataContainer = None) -> None:
+    def calc_rates(self, day: datetime.date, drv: dict) -> torch.Tensor:
         """Compute the potential gross assimilation rate (PGASS)."""
         p = self.params
         r = self.rates
@@ -356,9 +354,9 @@ def calc_rates(self, day: datetime.date = None, drv: WeatherDataContainer = None
         lai = _broadcast_to(k["LAI"], self.params.shape, dtype=self.dtype, device=self.device)
 
         # Weather drivers
-        irrad = _get_drv(drv.IRRAD, self.params.shape, dtype=self.dtype, device=self.device)
-        dtemp = _get_drv(drv.DTEMP, self.params.shape, dtype=self.dtype, device=self.device)
-        tmin = _get_drv(drv.TMIN, self.params.shape, dtype=self.dtype, device=self.device)
+        irrad = drv["IRRAD"]
+        dtemp = drv["DTEMP"]
+        tmin = drv["TMIN"]
 
         # Assimilation is zero before crop emergence (DVS < 0)
         dvs_mask = dvs >= 0
@@ -373,32 +371,26 @@ def calc_rates(self, day: datetime.date = None, drv: WeatherDataContainer = None
         # latitude and radiation are passed directly – they may be scalars or
         # tensors; the function returns torch.Tensor results in all cases.
         dayl, _daylp, sinld, cosld, difpp, _atmtr, dsinbe, _angot = astro(
-            day, drv.LAT, drv.IRRAD, dtype=self.dtype, device=self.device
+            day, drv["LAT"], drv["IRRAD"], dtype=self.dtype, device=self.device
         )
 
-        dayl_t = _broadcast_to(dayl, self.params.shape, dtype=self.dtype, device=self.device)
-        sinld_t = _broadcast_to(sinld, self.params.shape, dtype=self.dtype, device=self.device)
-        cosld_t = _broadcast_to(cosld, self.params.shape, dtype=self.dtype, device=self.device)
-        difpp_t = _broadcast_to(difpp, self.params.shape, dtype=self.dtype, device=self.device)
-        dsinbe_t = _broadcast_to(dsinbe, self.params.shape, dtype=self.dtype, device=self.device)
-
         # Parameter tables
         amax = p.AMAXTB(dvs)
         amax = amax * p.TMPFTB(dtemp)
         kdif = p.KDIFTB(dvs)
         eff = p.EFFTB(dtemp)
 
         dtga = totass7(
-            dayl_t,
+            dayl,
             amax,
             eff,
             lai,
             kdif,
             irrad,
-            difpp_t,
-            dsinbe_t,
-            sinld_t,
-            cosld_t,
+            difpp,
+            dsinbe,
+            sinld,
+            cosld,
             epsilon=self._epsilon,
             dtype=self.dtype,
             device=self.device,
@@ -414,11 +406,11 @@ def calc_rates(self, day: datetime.date = None, drv: WeatherDataContainer = None
         r.PGASS = pgass * dvs_mask
         return r.PGASS
 
-    def __call__(self, day: datetime.date = None, drv: WeatherDataContainer = None) -> torch.Tensor:
+    def __call__(self, day: datetime.date, drv: dict) -> torch.Tensor:
         """Calculate and return the potential gross assimilation rate (PGASS)."""
         return self.calc_rates(day, drv)
 
-    def integrate(self, day: datetime.date = None, delt=1.0) -> None:
+    def integrate(self, day: datetime.date, delt: float = 1.0) -> None:
         """No state variables to integrate for this module."""
         return
 

src/diffwofost/physical_models/crop/evapotranspiration.py

@@ -3,7 +3,6 @@
 from pcse.base import SimulationObject
 from pcse.base.parameter_providers import ParameterProvider
 from pcse.base.variablekiosk import VariableKiosk
-from pcse.base.weather import WeatherDataContainer
 from pcse.traitlets import Any
 from pcse.traitlets import Bool
 from pcse.traitlets import Instance
@@ -14,7 +13,6 @@
 from diffwofost.physical_models.traitlets import Tensor
 from diffwofost.physical_models.utils import AfgenTrait
 from diffwofost.physical_models.utils import _broadcast_to
-from diffwofost.physical_models.utils import _get_drv
 
 
 def SWEAF(ET0: torch.Tensor, DEPNR: torch.Tensor) -> torch.Tensor:
@@ -102,22 +100,21 @@ def initialize(
         else:
             self.etmodule = Evapotranspiration(day, kiosk, parvalues, shape=shape)
 
-    def calc_rates(self, day: datetime.date = None, drv: WeatherDataContainer = None):
+    def calc_rates(self, day: datetime.date, drv: dict):
         """Delegate rate calculation to the selected evapotranspiration module.
 
         Args:
-            day (datetime.date, optional): The current date of the simulation.
-            drv (WeatherDataContainer, optional): A dictionary-like container holding
-                weather data elements as key/value. The values are
-                arrays or scalars. See PCSE documentation for details.
+            day (datetime.date): The current date of the simulation.
+            drv (dict): A container holding weather data elements as key/value. The values are
+                arrays or scalars.
         """
         return self.etmodule.calc_rates(day, drv)
 
-    def __call__(self, day: datetime.date = None, drv: WeatherDataContainer = None):
+    def __call__(self, day: datetime.date, drv: dict):
         """Callable interface for rate calculation."""
         return self.calc_rates(day, drv)
 
-    def integrate(self, day: datetime.date = None, delt=1.0) -> None:
+    def integrate(self, day: datetime.date, delt: float = 1.0) -> None:
         """Delegate state integration to the selected evapotranspiration module.
 
         Args:
@@ -190,11 +187,11 @@ def _initialize_base(
         self._IDWST = torch.zeros(shape, dtype=self.dtype, device=self.device)
         self._IDOST = torch.zeros(shape, dtype=self.dtype, device=self.device)
 
-    def __call__(self, day: datetime.date = None, drv: WeatherDataContainer = None):
+    def __call__(self, day: datetime.date, drv: dict):
         """Callable interface for rate calculation."""
         return self.calc_rates(day, drv)
 
-    def integrate(self, day: datetime.date = None, delt=1.0) -> None:
+    def integrate(self, day: datetime.date, delt: float = 1.0) -> None:
         """Accumulate stress-day counters for water and oxygen stress."""
         rfws_stress = (self.rates.RFWS < 1.0).to(dtype=self.dtype)
         rfos_stress = (self.rates.RFOS < 1.0).to(dtype=self.dtype)
@@ -211,11 +208,11 @@ def finalize(self, day: datetime.date) -> None:
 class _BaseEvapotranspirationNonLayered(_BaseEvapotranspiration):
     """Shared implementation for non-layered evapotranspiration."""
 
-    def _rf_tramx_co2(self, drv: WeatherDataContainer, et0: torch.Tensor) -> torch.Tensor:
+    def _rf_tramx_co2(self, drv: dict, et0: torch.Tensor) -> torch.Tensor:
         """Return CO2 reduction factor for TRAMX (no CO2 effect in base implementation)."""
         return torch.ones_like(et0)
 
-    def calc_rates(self, day: datetime.date = None, drv: WeatherDataContainer = None):
+    def calc_rates(self, day: datetime.date, drv: dict):
         p = self.params
         r = self.rates
         k = self.kiosk
@@ -227,9 +224,9 @@ def calc_rates(self, day: datetime.date = None, drv: WeatherDataContainer = None
         # TODO see #22
         dvs = _broadcast_to(k["DVS"], self.params_shape, dtype=self.dtype, device=self.device)
 
-        et0 = _get_drv(drv.ET0, self.params_shape, dtype=self.dtype, device=self.device)
-        e0 = _get_drv(drv.E0, self.params_shape, dtype=self.dtype, device=self.device)
-        es0 = _get_drv(drv.ES0, self.params_shape, dtype=self.dtype, device=self.device)
+        et0 = torch.as_tensor(drv["ET0"])
+        e0 = torch.as_tensor(drv["E0"])
+        es0 = torch.as_tensor(drv["ES0"])
         rf_tramx_co2 = self._rf_tramx_co2(drv, et0)
 
         # If DVS < 0, the crop has not yet emerged, so we zero the rates using a mask
@@ -483,10 +480,10 @@ def initialize(
             shape=shape,
         )
 
-    def _rf_tramx_co2(self, drv: WeatherDataContainer, et0: torch.Tensor) -> torch.Tensor:
+    def _rf_tramx_co2(self, drv: dict, et0: torch.Tensor) -> torch.Tensor:
         """Calculate CO2 reduction factor for TRAMX based on atmospheric CO2 concentration."""
-        if hasattr(drv, "CO2") and drv.CO2 is not None:
-            co2 = _get_drv(drv.CO2, self.params_shape, dtype=self.dtype, device=self.device)
+        if "CO2" in drv and drv["CO2"] is not None:
+            co2 = drv["CO2"]
         else:
             co2 = self.params.CO2
         return self.params.CO2TRATB(co2)
@@ -634,15 +631,15 @@ def initialize(
         # Internal DSOS tracker for layered oxygen-stress response
         self._dsos = torch.zeros(self.params_shape, dtype=self.dtype, device=self.device)
 
-    def _rf_tramx_co2(self, drv: WeatherDataContainer, et0: torch.Tensor) -> torch.Tensor:
+    def _rf_tramx_co2(self, drv: dict, et0: torch.Tensor) -> torch.Tensor:
         """Calculate CO2 reduction factor for TRAMX using CO2 from driver or parameters."""
-        if hasattr(drv, "CO2") and drv.CO2 is not None:
-            co2 = _get_drv(drv.CO2, self.params_shape, dtype=self.dtype, device=self.device)
+        if "CO2" in drv and drv["CO2"] is not None:
+            co2 = drv["CO2"]
         else:
             co2 = self.params.CO2
         return self.params.CO2TRATB(co2)
 
-    def calc_rates(self, day: datetime.date = None, drv: WeatherDataContainer = None):
+    def calc_rates(self, day: datetime.date, drv: dict):
         """Calculate daily evapotranspiration rates per soil layer with CO2 effects.
 
         Computes transpiration and stress factors for each soil layer based on root
@@ -658,9 +655,9 @@ def calc_rates(self, day: datetime.date = None, drv: WeatherDataContainer = None
 
         n_layers = self._n_layers
 
-        et0 = _get_drv(drv.ET0, self.params_shape, dtype=self.dtype, device=self.device)
-        e0 = _get_drv(drv.E0, self.params_shape, dtype=self.dtype, device=self.device)
-        es0 = _get_drv(drv.ES0, self.params_shape, dtype=self.dtype, device=self.device)
+        et0 = torch.as_tensor(drv["ET0"])
+        e0 = torch.as_tensor(drv["E0"])
+        es0 = torch.as_tensor(drv["ES0"])
 
         # reduction factor for CO2 on TRAMX
         rf_tramx_co2 = self._rf_tramx_co2(drv, et0)
@@ -786,11 +783,11 @@ def calc_rates(self, day: datetime.date = None, drv: WeatherDataContainer = None
         r.IDOS = bool(torch.any(r.RFOS < 1.0))
         return r.TRA, r.TRAMX
 
-    def __call__(self, day: datetime.date = None, drv: WeatherDataContainer = None):
+    def __call__(self, day: datetime.date, drv: dict):
         """Callable interface for rate calculation."""
         return self.calc_rates(day, drv)
 
-    def integrate(self, day: datetime.date = None, delt=1.0) -> None:
+    def integrate(self, day: datetime.date, delt: float = 1.0) -> None:
         """Accumulate stress-day counters based on any layer experiencing stress."""
         rfws_stress = (self.rates.RFWS < 1.0).any(dim=0).to(dtype=self.dtype)
         rfos_stress = (self.rates.RFOS < 1.0).any(dim=0).to(dtype=self.dtype)

src/diffwofost/physical_models/crop/leaf_dynamics.py

@@ -5,14 +5,12 @@
 from pcse.base import SimulationObject
 from pcse.base.parameter_providers import ParameterProvider
 from pcse.base.variablekiosk import VariableKiosk
-from pcse.base.weather import WeatherDataContainer
 from diffwofost.physical_models.base import TensorParamTemplate
 from diffwofost.physical_models.base import TensorRatesTemplate
 from diffwofost.physical_models.base import TensorStatesTemplate
 from diffwofost.physical_models.config import ComputeConfig
 from diffwofost.physical_models.traitlets import Tensor
 from diffwofost.physical_models.utils import AfgenTrait
-from diffwofost.physical_models.utils import _get_drv
 
 
 class WOFOST_Leaf_Dynamics(SimulationObject):
@@ -248,14 +246,13 @@ def _calc_LAI(self):
         total_LAI = self.states.LASUM + SAI + PAI
         return total_LAI
 
-    def calc_rates(self, day: datetime.date, drv: WeatherDataContainer) -> None:
+    def calc_rates(self, day: datetime.date, drv: dict) -> None:
         """Calculate the rates of change for the leaf dynamics.
 
         Args:
-            day (datetime.date, optional): The current date of the simulation.
-            drv (WeatherDataContainer, optional): A dictionary-like container holding
-                weather data elements as key/value. The values are
-                arrays or scalars. See PCSE documentation for details.
+            day (datetime.date): The current date of the simulation.
+            drv (dict): A container holding weather data elements as key/value. The values are
+                arrays or scalars.
         """
         r = self.rates
         s = self.states
@@ -326,7 +323,7 @@ def calc_rates(self, day: datetime.date, drv: WeatherDataContainer) -> None:
         r.DRLV = torch.maximum(r.DSLV, r.DALV)
 
         # Get the temperature from the drv
-        TEMP = _get_drv(drv.TEMP, p.shape, self.dtype, self.device)
+        TEMP = drv["TEMP"]
 
         # physiologic ageing of leaves per time step
         FYSAGE = (TEMP - p.TBASE) / (35.0 - p.TBASE)

src/diffwofost/physical_models/crop/phenology.py

@@ -19,8 +19,6 @@
 from diffwofost.physical_models.config import ComputeConfig
 from diffwofost.physical_models.traitlets import Tensor
 from diffwofost.physical_models.utils import AfgenTrait
-from diffwofost.physical_models.utils import _broadcast_to
-from diffwofost.physical_models.utils import _get_drv
 from diffwofost.physical_models.utils import _restore_state
 from diffwofost.physical_models.utils import _snapshot_state
 from diffwofost.physical_models.utils import daylength
@@ -182,7 +180,7 @@ def calc_rates(self, day, drv):
         VERNBASE = params.VERNBASE
         DVS = self.kiosk["DVS"]
 
-        TEMP = _get_drv(drv.TEMP, self.params.shape, self.dtype, self.device)
+        TEMP = drv["TEMP"]
 
         # Operate elementwise only on elements not yet vernalised
         not_vernalised = ~self.states.ISVERNALISED
@@ -505,15 +503,13 @@ def calc_rates(self, day, drv):
         p = self.params
         r = self.rates
         s = self.states
-
         # Day length sensitivity
         # daylength returns a Tensor directly; broadcast to parameter shape.
-        DAYLP = daylength(day, drv.LAT, dtype=self.dtype, device=self.device)
-        DAYLP_t = _broadcast_to(DAYLP, p.shape, dtype=self.dtype, device=self.device)
+        DAYLP = daylength(day, drv["LAT"], dtype=self.dtype, device=self.device)
         # Compute DVRED conditionally based on IDSL >= 1
         safe_den = p.DLO - p.DLC
         safe_den = safe_den.sign() * torch.maximum(torch.abs(safe_den), self._epsilon)
-        dvred_active = torch.clamp((DAYLP_t - p.DLC) / safe_den, 0.0, 1.0)
+        dvred_active = torch.clamp((DAYLP - p.DLC) / safe_den, 0.0, 1.0)
         DVRED = torch.where(p.IDSL >= 1, dvred_active, self._ones)
 
         # Vernalisation factor - always compute if module exists
@@ -529,7 +525,7 @@ def calc_rates(self, day, drv):
                 self._ones,
             )
 
-        TEMP = _get_drv(drv.TEMP, p.shape, self.dtype, self.device)
+        TEMP = drv["TEMP"]
 
         # Initialize all rate variables
         r.DTSUME = self._zeros