Fix issue 22 by defining new VariableKiosk

src/diffwofost/physical_models/base/__init__.py

@@ -1,5 +1,5 @@
-from .states_rates import TensorParamTemplate
-from .states_rates import TensorRatesTemplate
-from .states_rates import TensorStatesTemplate
+from diffwofost.physical_models.base.states_rates import TensorParamTemplate
+from diffwofost.physical_models.base.states_rates import TensorRatesTemplate
+from diffwofost.physical_models.base.states_rates import TensorStatesTemplate
 
 __all__ = ["TensorParamTemplate", "TensorRatesTemplate", "TensorStatesTemplate"]

src/diffwofost/physical_models/base/states_rates.py

@@ -2,8 +2,8 @@
 from pcse.base import RatesTemplate
 from pcse.base import StatesTemplate
 from pcse.traitlets import HasTraits
-from ..traitlets import Tensor
-from ..utils import AfgenTrait
+from diffwofost.physical_models.traitlets import Tensor
+from diffwofost.physical_models.utils import AfgenTrait
 
 
 class TensorContainer(HasTraits):

src/diffwofost/physical_models/crop/wofost72.py

@@ -12,16 +12,22 @@
 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.crop.assimilation import WOFOST72_Assimilation as Assimilation
+from diffwofost.physical_models.crop.evapotranspiration import (
+    EvapotranspirationWrapper as Evapotranspiration,
+)
+from diffwofost.physical_models.crop.leaf_dynamics import WOFOST_Leaf_Dynamics as Leaf_Dynamics
+from diffwofost.physical_models.crop.partitioning import DVS_Partitioning as Partitioning
+from diffwofost.physical_models.crop.phenology import DVS_Phenology as Phenology
+from diffwofost.physical_models.crop.respiration import (
+    WOFOST_Maintenance_Respiration as MaintenanceRespiration,
+)
+from diffwofost.physical_models.crop.root_dynamics import WOFOST_Root_Dynamics as Root_Dynamics
+from diffwofost.physical_models.crop.stem_dynamics import WOFOST_Stem_Dynamics as Stem_Dynamics
+from diffwofost.physical_models.crop.storage_organ_dynamics import (
+    WOFOST_Storage_Organ_Dynamics as Storage_Organ_Dynamics,
+)
 from diffwofost.physical_models.traitlets import Tensor
-from .assimilation import WOFOST72_Assimilation as Assimilation
-from .evapotranspiration import EvapotranspirationWrapper as Evapotranspiration
-from .leaf_dynamics import WOFOST_Leaf_Dynamics as Leaf_Dynamics
-from .partitioning import DVS_Partitioning as Partitioning
-from .phenology import DVS_Phenology as Phenology
-from .respiration import WOFOST_Maintenance_Respiration as MaintenanceRespiration
-from .root_dynamics import WOFOST_Root_Dynamics as Root_Dynamics
-from .stem_dynamics import WOFOST_Stem_Dynamics as Stem_Dynamics
-from .storage_organ_dynamics import WOFOST_Storage_Organ_Dynamics as Storage_Organ_Dynamics
 
 
 class Wofost72(SimulationObject):

src/diffwofost/physical_models/engine.py

@@ -6,7 +6,7 @@
 from pcse.engine import Engine
 from pcse.timer import Timer
 from pcse.traitlets import Instance
-from .config import Configuration
+from diffwofost.physical_models.config import Configuration
 
 
 class Engine(Engine):

src/diffwofost/physical_models/traitlets.py

@@ -1,7 +1,7 @@
 import torch
 from traitlets_pcse import TraitType
 from traitlets_pcse import Undefined
-from .config import ComputeConfig
+from diffwofost.physical_models.config import ComputeConfig
 
 
 class Tensor(TraitType):

src/diffwofost/physical_models/utils.py

@@ -1,8 +1,6 @@
 """This file contains code that is required to run the YAML unit tests.
 
 It contains:
-    - VariableKioskTestHelper: A subclass of the VariableKiosk that can use externally
-      forced states/rates
     - EngineTestHelper: engine specifically for running the YAML tests.
     - WeatherDataProviderTestHelper: a weatherdata provides that takes the weather
       inputs from the YAML file.
@@ -19,77 +17,22 @@
 import yaml
 from pcse import signals
 from pcse.base.parameter_providers import ParameterProvider
-from pcse.base.variablekiosk import VariableKiosk
 from pcse.base.weather import WeatherDataContainer
 from pcse.base.weather import WeatherDataProvider
 from pcse.engine import BaseEngine
 from pcse.settings import settings
 from pcse.timer import Timer
 from pcse.traitlets import TraitType
 from pcse.util import doy
-from .config import ComputeConfig
-from .config import Configuration
-from .engine import Engine
-from .engine import _get_params_shape
+from diffwofost.physical_models.config import ComputeConfig
+from diffwofost.physical_models.config import Configuration
+from diffwofost.physical_models.engine import Engine
+from diffwofost.physical_models.engine import _get_params_shape
+from diffwofost.physical_models.variablekiosk import VariableKiosk
 
 logging.disable(logging.CRITICAL)
 
 
-class VariableKioskTestHelper(VariableKiosk):
-    """Variable Kiosk for testing purposes which allows to use external states."""
-
-    external_state_list = None
-
-    def __init__(self, external_state_list=None):
-        super().__init__()
-        self.current_externals = {}
-        if external_state_list:
-            self.external_state_list = external_state_list
-
-    def __call__(self, day):
-        """Sets the external state/rate variables for the current day.
-
-        Returns True if the list of external state/rate variables is exhausted,
-        otherwise False.
-        """
-        if self.external_state_list:
-            current_externals = self.external_state_list.pop(0)
-            forcing_day = current_externals.pop("DAY")
-            msg = "Failure updating VariableKiosk with external states: days are not matching!"
-            assert forcing_day == day, msg
-            self.current_externals.clear()
-            self.current_externals.update(current_externals)
-            if len(self.external_state_list) == 0:
-                return True
-
-        return False
-
-    def is_external_state(self, item):
-        """Returns True if the item is an external state."""
-        return item in self.current_externals
-
-    def __getattr__(self, item):
-        """Allow use of attribute notation.
-
-        eg "kiosk.LAI" on published rates or states.
-        """
-        if item in self.current_externals:
-            return self.current_externals[item]
-        else:
-            return dict.__getitem__(self, item)
-
-    def __getitem__(self, item):
-        """Override __getitem__ to first look in external states."""
-        if item in self.current_externals:
-            return self.current_externals[item]
-        else:
-            return dict.__getitem__(self, item)
-
-    def __contains__(self, key):
-        """Override __contains__ to first look in external states."""
-        return key in self.current_externals or dict.__contains__(self, key)
-
-
 class EngineTestHelper(Engine):
     """An engine which is purely for running the YAML unit tests."""
 
@@ -113,7 +56,7 @@ def __init__(
         self._shape = _get_params_shape(self.parameterprovider)
 
         # Variable kiosk for registering and publishing variables
-        self.kiosk = VariableKioskTestHelper(external_states)
+        self.kiosk = VariableKiosk(external_states)
 
         # Placeholder for variables to be saved during a model run
         self._saved_output = list()
@@ -157,10 +100,10 @@ def _run(self):
         # Update timer
         self.day, delt = self.timer()
 
-        # When the list of external states is exhausted the VariableKioskTestHelper will
-        # return True signalling the end of the test
-        stop_test = self.kiosk(self.day)
-        if stop_test:
+        self.kiosk(self.day)
+        # When the list of external states is exhausted, send crop_finish to
+        # end the test run
+        if self.kiosk.external_states_exhausted:
             self._send_signal(
                 signal=signals.crop_finish, day=self.day, finish_type="maturity", crop_delete=False
             )

src/diffwofost/physical_models/variablekiosk.py

@@ -0,0 +1,73 @@
+from pcse.base.variablekiosk import VariableKiosk as _PcseVariableKiosk
+
+
+class VariableKiosk(_PcseVariableKiosk):
+    """Extends pcse's VariableKiosk with support for external dependencies.
+
+    The external_state_list parameter accepts a list of per-day dicts, each
+    containing a ``"DAY"`` key and the variable values to inject for that day.
+    Calling the kiosk with a day (``kiosk(day)``) advances to the next entry
+    and makes those variables available via normal attribute/item access.
+
+    All original VariableKiosk behaviour (registering, publishing, flushing)
+    is inherited unchanged from pcse.
+    """
+
+    def __init__(self, external_state_list=None):
+        super().__init__()
+        self.current_externals = {}
+        self._last_called_day = None
+        # Build a day-keyed dict for O(1) lookup
+        self._external_states = {}
+        if external_state_list is not None:
+            self._external_states = {
+                item["DAY"]: {k: v for k, v in item.items() if k != "DAY"}
+                for item in list(external_state_list)
+            }
+
+    def __call__(self, day):
+        """Set the external state/rate variables for the current day.
+
+        If the day has an entry in the external state list, its values are
+        injected into ``current_externals``. If the day has no entry,
+        ``current_externals`` is cleared so the module falls back to normally
+        registered kiosk variables. Does nothing when no list was provided.
+        Always returns False; use ``external_states_exhausted`` to check whether
+        the last entry has been passed.
+        """
+        self._last_called_day = day
+        if self._external_states:
+            self.current_externals.clear()
+            if day in self._external_states:
+                self.current_externals.update(self._external_states[day])
+        return False
+
+    @property
+    def external_states_exhausted(self):
+        """True when the simulation has advanced past the last external state entry."""
+        if not self._external_states or self._last_called_day is None:
+            return False
+        return self._last_called_day >= max(self._external_states.keys())
+
+    def is_external_state(self, item):
+        """Returns True if the item is an external state."""
+        return item in self.current_externals
+
+    def __contains__(self, item):
+        """Checks external states first, then the published kiosk variables."""
+        current_externals = self.__dict__.get("current_externals", {})
+        return item in current_externals or dict.__contains__(self, item)
+
+    def __getitem__(self, item):
+        """Look in external states before falling back to published variables."""
+        current_externals = self.__dict__.get("current_externals", {})
+        if item in current_externals:
+            return current_externals[item]
+        return dict.__getitem__(self, item)
+
+    def __getattr__(self, item):
+        """Allow attribute notation (e.g. ``kiosk.LAI``), checking externals first."""
+        current_externals = self.__dict__.get("current_externals", {})
+        if item in current_externals:
+            return current_externals[item]
+        return dict.__getitem__(self, item)