Add deepcopy_method argument (#54)

* Add deepcopy_method argument

* Fix policy state mutation issue

CHANGELOG.md

@@ -4,6 +4,13 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.10.0] - 2022-09-01
+### Added
+- Add a `deepcopy_method` Engine argument to allow setting a custom deepcopy method e.g. `copy.deepcopy` instead of default Pickle methods
+
+### Changed
+- Fix edge case of unintended mutation of state passed to state update function from policy function (see issue #53)
+
 ## [0.9.1] - 2022-07-26
 ### Added
 - Thanks to @vmeylan for a developer experience contribution in `radcad/core.py`: make `_update_state()` error messages more verbose for easier debugging

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "radcad"
-version = "0.9.1"
+version = "0.10.0"
 description = "A Python package for dynamical systems modelling & simulation, inspired by and compatible with cadCAD"
 authors = ["CADLabs <[email protected]>"]
 packages = [

radcad/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.9.1"
+__version__ = "0.10.0"
 
 from radcad.wrappers import Context, Model, Simulation, Experiment
 from radcad.engine import Engine

radcad/core.py

@@ -2,7 +2,13 @@
 import logging
 import pickle
 import traceback
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Callable
+
+
+# Define the default method used for deepcopy operations
+# Must be a function and not a lambda function to ensure multiprocessing can Pickle the object
+def default_deepcopy_method(obj):
+    return pickle.loads(pickle.dumps(obj=obj, protocol=-1))
 
 
 def _update_state(initial_state, params, substep, result, substate, signals, state_update_tuple):
@@ -34,6 +40,7 @@ def _single_run(
     state_update_blocks: list,
     params: dict,
     deepcopy: bool,
+    deepcopy_method: Callable,
     drop_substeps: bool,
 ):
     logging.info(f"Starting simulation {simulation} / run {run} / subset {subset}")
@@ -61,15 +68,20 @@ def _single_run(
             substate: dict = (
                 previous_state.copy() if substep == 0 else substeps[substep - 1].copy()
             )
-            substate_copy = pickle.loads(pickle.dumps(substate, -1)) if deepcopy else substate.copy()
+
+            # Create two independent deepcopies to ensure a policy function
+            # can't mutate the state passed to the state update functions
+            policy_substate_copy = deepcopy_method(substate) if deepcopy else substate.copy()
+            state_update_substate_copy = deepcopy_method(substate) if deepcopy else substate.copy()
+
             substate["substep"] = substep + 1
 
             signals: dict = reduce_signals(
-                params, substep, result, substate_copy, psu, deepcopy
+                params, substep, result, policy_substate_copy, psu, deepcopy
             )
 
             updated_state = map(
-                partial(_update_state, initial_state, params, substep, result, substate_copy, signals),
+                partial(_update_state, initial_state, params, substep, result, state_update_substate_copy, signals),
                 psu["variables"].items()
             )
             substate.update(updated_state)
@@ -90,6 +102,7 @@ def single_run(
     state_update_blocks=[],
     params={},
     deepcopy: bool=True,
+    deepcopy_method: Callable=default_deepcopy_method,
     drop_substeps: bool=False,
 ) -> Tuple[list, Exception, str]:
     result = []
@@ -106,6 +119,7 @@ def single_run(
                 state_update_blocks,
                 params,
                 deepcopy,
+                deepcopy_method,
                 drop_substeps,
             ),
             None, # Error
@@ -174,7 +188,7 @@ def _add_signals(acc, a: Dict[str, any]):
     return acc
 
 
-def reduce_signals(params: dict, substep: int, result: list, substate: dict, psu: dict, deepcopy: bool=True):
+def reduce_signals(params: dict, substep: int, result: list, substate: dict, psu: dict, deepcopy: bool=True, deepcopy_method: Callable=default_deepcopy_method):
     policy_results: List[Dict[str, any]] = list(
         map(lambda function: function(params, substep, result, substate), psu["policies"].values())
     )
@@ -184,6 +198,6 @@ def reduce_signals(params: dict, substep: int, result: list, substate: dict, psu
     if result_length == 0:
         return result
     elif result_length == 1:
-        return pickle.loads(pickle.dumps(policy_results[0], -1)) if deepcopy else policy_results[0].copy()
+        return deepcopy_method(policy_results[0]) if deepcopy else policy_results[0].copy()
     else:
         return reduce(_add_signals, policy_results, result)

radcad/engine.py

@@ -19,7 +19,8 @@ def __init__(self, **kwargs):
             **backend (Backend): Which execution backend to use (e.g. Pathos, Multiprocessing, etc.). Defaults to `Backend.DEFAULT` / `Backend.PATHOS`.
             **processes (int, optional): Number of system CPU processes to spawn. Defaults to `multiprocessing.cpu_count() - 1 or 1`
             **raise_exceptions (bool): Whether to raise exceptions, or catch them and return exceptions along with partial results. Default to `True`.
-            **deepcopy (bool): Whether to enable deepcopy of State Variables, alternatively leaves safety up to user with improved performance. Defaults to `True`.
+            **deepcopy (bool): Whether to enable deepcopy of State Variables to avoid unintended state mutation. Defaults to `True`.
+            **deepcopy_method (Callable): Method to use for deepcopy of State Variables. By default uses Pickle for improved performance, use `copy.deepcopy` for an alternative to Pickle.
             **drop_substeps (bool): Whether to drop simulation result substeps during runtime to save memory and improve performance. Defaults to `False`.
             **_run_generator (tuple_iterator): Generator to generate simulation runs, used to implement custom execution backends. Defaults to  `iter(())`.
         """
@@ -28,6 +29,7 @@ def __init__(self, **kwargs):
         self.backend = kwargs.pop("backend", Backend.DEFAULT)
         self.raise_exceptions = kwargs.pop("raise_exceptions", True)
         self.deepcopy = kwargs.pop("deepcopy", True)
+        self.deepcopy_method = kwargs.pop("deepcopy_method", core.default_deepcopy_method)
         self.drop_substeps = kwargs.pop("drop_substeps", False)
         self._run_generator = iter(())
 
@@ -137,6 +139,7 @@ def _run_stream(self, configs):
                         state_update_blocks,
                         copy.deepcopy(param_set),
                         self.deepcopy,
+                        self.deepcopy_method,
                         self.drop_substeps,
                     )
                     self.executable._after_subset(context=context)

radcad/wrappers.py

@@ -1,4 +1,4 @@
-from radcad.core import _single_run_wrapper, generate_parameter_sweep
+from radcad.core import _single_run_wrapper, generate_parameter_sweep, default_deepcopy_method
 from radcad.engine import Engine
 from collections import namedtuple
 import copy
@@ -13,6 +13,7 @@
     "state_update_blocks",
     "parameters",
     "deepcopy",
+    "deepcopy_method",
     "drop_substeps",
 ])
 Context = namedtuple("Context", "simulation run subset timesteps initial_state parameters")
@@ -35,6 +36,7 @@ def __init__(self, initial_state={}, state_update_blocks=[], params={}):
         self.exceptions = []
         self._raise_exceptions = True
         self._deepcopy = True
+        self._deepcopy_method = default_deepcopy_method
         self._drop_substeps = False
 
     def __iter__(self):
@@ -50,6 +52,7 @@ def __iter__(self):
                 state_update_blocks = self.state_update_blocks,
                 parameters = _params,
                 deepcopy = self._deepcopy,
+                deepcopy_method = self._deepcopy_method,
                 drop_substeps = self._drop_substeps,
             )
             result, exception = _single_run_wrapper((run_args, self._raise_exceptions))
@@ -61,6 +64,7 @@ def __iter__(self):
     def __call__(self, **kwargs):
         self._raise_exceptions = kwargs.pop("raise_exceptions", True)
         self._deepcopy = kwargs.pop("deepcopy", True)
+        self._deepcopy_method = kwargs.pop("deepcopy_method", default_deepcopy_method)
         self._drop_substeps = kwargs.pop("drop_substeps", False)
         return self