Run multiple crops - first work on engine

src/diffwofost/physical_models/engine.py

@@ -6,6 +6,7 @@
 """
 
 import gc
+from collections.abc import MutableMapping
 from pathlib import Path
 import torch
 from pcse import signals
@@ -27,6 +28,7 @@ class Engine(PcseEngine):
     """
 
     mconf = Instance(Configuration)
+    parameterprovider = Instance(MutableMapping)
 
     def __init__(
         self,
@@ -80,7 +82,6 @@ def _reset_runtime_state(self):
         self.flag_terminate = False
         self.flag_crop_finish = False
         self.flag_crop_start = False
-        self.flag_crop_delete = False
         self.flag_output = False
         self.flag_summary_output = False
 
@@ -138,12 +139,15 @@ def setup(
         self.weatherdataprovider = weatherdataprovider
         self.drv = self._get_driving_variables(self.day)
 
+        # Call AgroManagement module for management actions at initialization
+        self.agromanager(self.day, None)
+
         # Component for simulation of soil processes
         if self.mconf.SOIL is not None:
             self.soil = self.mconf.SOIL(self.day, self.kiosk, parameterprovider)
-
-        # Call AgroManagement module for management actions at initialization
-        self.agromanager(self.day, self.drv)
+        # Component for crop simulation
+        if self.mconf.CROP is not None:
+            self._create_crop(self.day)
 
         # Calculate initial rates
         self.calc_rates(self.day, self.drv)
@@ -152,7 +156,7 @@ def setup(
     def _on_CROP_START(
         self, day, crop_name=None, variety_name=None, crop_start_type=None, crop_end_type=None
     ):
-        """Instantiate the crop component after a crop-start signal.
+        """Set active crop parameters for providers that support that.
 
         Args:
             day: Current simulation day.
@@ -162,35 +166,46 @@ def _on_CROP_START(
                 signal.
             crop_start_type: Crop start mode used by the parameter provider.
             crop_end_type: Crop end mode used by the parameter provider.
-
-        Raises:
-            RuntimeError: If a crop component is already active.
         """
         self.logger.debug(f"Received signal 'CROP_START' on day {day}")
 
-        if self.crop is not None:
-            raise RuntimeError(
-                "A CROP_START signal was received while self.cropsimulation still holds a valid "
-                "cropsimulation object. It looks like you forgot to send a CROP_FINISH signal with "
-                "option crop_delete=True"
+        if hasattr(self.parameterprovider, "set_active_crop"):
+            self.parameterprovider.set_active_crop(
+                crop_name, variety_name, crop_start_type, crop_end_type
             )
 
-        self.parameterprovider.set_active_crop(
-            crop_name, variety_name, crop_start_type, crop_end_type
-        )
+    def _create_crop(self, day):
+        """Setup crop model instance.
 
+        Args:
+            day: Current simulation day
+        """
         crop_args = [day, self.kiosk, self.parameterprovider]
         crop_kwargs = {"shape": self._shape}
 
         if self.mconf.CROP_NN_MODEL is not None:
-            # crop_nn_model initialize doesnot accpet parameterprovider
+            # crop_nn_model initialize does not accept parameterprovider
             crop_args = [day, self.kiosk, self.mconf.CROP_NN_MODEL]
 
         if self.mconf.CROP_COMPONENTS:
             crop_kwargs["component_overrides"] = self._components_overrides
 
         self.crop = self.mconf.CROP(*crop_args, **crop_kwargs)
 
+    def _on_CROP_FINISH(self, day):
+        """Flag finishing of the crop simulation.
+
+        The flag is needed because finishing the crop simulation is deferred to
+        the correct place in the processing loop and is done by the routine
+        _finish_cropsimulation().
+
+        Differently from PCSE, diffWOFOST does not delete the crop simulation instance.
+
+        Args:
+            day: Current simulation day.
+        """
+        self.flag_crop_finish = True
+
     def _finish_cropsimulation(self, day):
         """Finalize and optionally delete the active crop simulation.
 
@@ -202,12 +217,6 @@ def _finish_cropsimulation(self, day):
         self.crop.finalize(day)
         self._save_summary_output()
 
-        if self.flag_crop_delete:
-            self.flag_crop_delete = False
-            self.crop._delete()
-            self.crop = None
-            gc.collect()
-
 
 def _get_params_shape(parameterprovider):
     """Infer the common tensor batch shape from a parameter provider.
@@ -227,7 +236,7 @@ def _get_params_shape(parameterprovider):
         ValueError: If tensor parameters do not share a common shape.
     """
     shape = ()
-    for paramname in parameterprovider._unique_parameters:
+    for paramname in parameterprovider.keys():
         param = parameterprovider[paramname]
         if isinstance(param, torch.Tensor):
             # We need to drop the last dimension from the Afgen table parameters

src/diffwofost/physical_models/parameter_providers.py

@@ -0,0 +1,12 @@
+from pcse.base.parameter_providers import ParameterProvider as PcseParameterProvider
+
+
+class ParameterProvider(PcseParameterProvider):
+    """Temporary implementation of PCSE's ParameterProvider.
+
+    Fixes the `__iter__` method in order to allow for access via dict-like `.items()`, `.keys()`,
+    and `.values()`. Could be dropped when https://github.com/ajwdewit/pcse/pull/121 is merged.
+    """
+
+    def __iter__(self):
+        return iter(self._unique_parameters)

src/diffwofost/physical_models/test.py

@@ -1,12 +1,12 @@
 import torch
 import yaml
 from pcse import signals
-from pcse.base.parameter_providers import ParameterProvider
 from pcse.base.weather import WeatherDataContainer
 from pcse.base.weather import WeatherDataProvider
 from pcse.settings import settings
 from diffwofost.physical_models.config import ComputeConfig
 from diffwofost.physical_models.engine import Engine
+from diffwofost.physical_models.parameter_providers import ParameterProvider
 
 
 class EngineTestHelper(Engine):

tests/physical_models/test_engine.py

@@ -4,7 +4,9 @@
 import pytest
 import torch
 from pcse.base import ParameterProvider
+from pcse.base import SimulationObject
 from pcse.base.variablekiosk import VariableKiosk
+from diffwofost.physical_models.base import TensorParamTemplate
 from diffwofost.physical_models.config import Configuration
 from diffwofost.physical_models.crop.phenology import DVS_Phenology
 from diffwofost.physical_models.crop.wofost72 import Wofost72
@@ -13,6 +15,7 @@
 from diffwofost.physical_models.soil.classic_waterbalance import WaterbalancePP
 from diffwofost.physical_models.test import get_test_data
 from diffwofost.physical_models.test import prepare_engine_input
+from diffwofost.physical_models.traitlets import Tensor
 from . import phy_data_folder
 
 config = Configuration(
@@ -52,6 +55,14 @@ def set_active_crop(self, *args):
         self.active_crop_args = args
 
 
+class DummyCropModel(SimulationObject):
+    class Parameters(TensorParamTemplate):
+        A = Tensor(-99, dtype=int)
+
+    def initialize(self, day, kiosk, parvalues, shape):
+        self.params = self.Parameters(parvalues, shape=shape)
+
+
 @pytest.mark.usefixtures("fast_mode")
 class TestEngine:
     def test_engine_requires_config(self):
@@ -151,39 +162,7 @@ def test_setup_initializes_soil_component_when_configured(self):
         assert returned_engine is engine
         assert engine.soil is not None
 
-    def test_on_crop_start_raises_when_crop_is_already_active(self):
-        _, (crop_model_params_provider, weather_data_provider, agro_management_inputs, _) = (
-            _get_engine_inputs()
-        )
-        engine = Engine(config=config)
-        engine.setup(crop_model_params_provider, weather_data_provider, agro_management_inputs)
-
-        with pytest.raises(RuntimeError, match="A CROP_START signal was received"):
-            engine._on_CROP_START(engine.day)
-
-    def test_finish_cropsimulation_deletes_crop_when_requested(self):
-        _, (crop_model_params_provider, weather_data_provider, agro_management_inputs, _) = (
-            _get_engine_inputs()
-        )
-        engine = Engine(config=config)
-        engine.setup(crop_model_params_provider, weather_data_provider, agro_management_inputs)
-        crop = engine.crop
-        crop.finalize = Mock()
-        crop._delete = Mock()
-        engine.flag_crop_finish = True
-        engine.flag_crop_delete = True
-        engine._save_summary_output = Mock()
-
-        engine._finish_cropsimulation(date(2000, 1, 1))
-
-        assert engine.flag_crop_finish is False
-        assert engine.flag_crop_delete is False
-        crop.finalize.assert_called_once_with(date(2000, 1, 1))
-        engine._save_summary_output.assert_called_once_with()
-        crop._delete.assert_called_once_with()
-        assert engine.crop is None
-
-    def test_finish_cropsimulation_keeps_crop_when_not_deleting(self):
+    def test_finish_cropsimulation_does_not_delete_crop(self):
         _, (crop_model_params_provider, weather_data_provider, agro_management_inputs, _) = (
             _get_engine_inputs()
         )
@@ -193,13 +172,11 @@ def test_finish_cropsimulation_keeps_crop_when_not_deleting(self):
         crop.finalize = Mock()
         crop._delete = Mock()
         engine.flag_crop_finish = True
-        engine.flag_crop_delete = False
         engine._save_summary_output = Mock()
 
         engine._finish_cropsimulation(date(2000, 1, 1))
 
         assert engine.flag_crop_finish is False
-        assert engine.flag_crop_delete is False
         crop.finalize.assert_called_once_with(date(2000, 1, 1))
         engine._save_summary_output.assert_called_once_with()
         crop._delete.assert_not_called()
@@ -258,6 +235,15 @@ def initialize(self, day, kiosk, nn_model, shape=None):
         engine.parameterprovider = ParameterProvider()
         engine.kiosk = VariableKiosk()
         engine._shape = ()
-        engine._on_CROP_START(date(2000, 1, 1))
+        engine._create_crop(date(2000, 1, 1))
         assert engine.mconf.CROP_NN_MODEL == nn_model
         assert engine.crop._initialized is True
+
+    def test_engine_accept_dict_as_parameter_provider(self):
+        engine = Engine(config=Configuration(CROP=DummyCropModel))
+        engine.parameterprovider = {"A": 10}
+        engine.kiosk = VariableKiosk()
+        engine._shape = ()
+        engine._create_crop(date(2000, 1, 1))
+        assert isinstance(engine.crop.params.A, torch.Tensor)
+        assert engine.crop.params.A == 10

tests/physical_models/test_parameter_providers.py

@@ -0,0 +1,12 @@
+from diffwofost.physical_models.parameter_providers import ParameterProvider
+
+
+class TestParameterProvider:
+    def test_parameter_provider_supports_dict_methods(self):
+        p = ParameterProvider(
+            sitedata={"A": 0}, timerdata={"B": 1}, soildata={"C": 2}, cropdata={"D": 3}
+        )
+        p.set_override("E", 4, check=False)
+        assert len(p.items()) == len(p.keys()) == len(p.values()) == 5
+        assert set(p.keys()) == set("ABCDE")
+        assert set(p.values()) == set(range(5))