diff --git a/src/diffwofost/physical_models/engine.py b/src/diffwofost/physical_models/engine.py index 5e597e5..a84ef32 100644 --- a/src/diffwofost/physical_models/engine.py +++ b/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,28 +166,25 @@ 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: @@ -191,6 +192,20 @@ def _on_CROP_START( 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 diff --git a/src/diffwofost/physical_models/parameter_providers.py b/src/diffwofost/physical_models/parameter_providers.py new file mode 100644 index 0000000..204acc1 --- /dev/null +++ b/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) diff --git a/src/diffwofost/physical_models/test.py b/src/diffwofost/physical_models/test.py index bd2817a..1dbbd7c 100644 --- a/src/diffwofost/physical_models/test.py +++ b/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): diff --git a/tests/physical_models/test_engine.py b/tests/physical_models/test_engine.py index b7e760d..2cbf151 100644 --- a/tests/physical_models/test_engine.py +++ b/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 diff --git a/tests/physical_models/test_parameter_providers.py b/tests/physical_models/test_parameter_providers.py new file mode 100644 index 0000000..c064895 --- /dev/null +++ b/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))