llvm, Mechanism: Update output ports in "reset" execution variant

The "reset" variants of functions now return a value. Add tests for reset behaviour of: * ProcessingMechanism using an integrator function * DDM mechanism using DriftDiffusionIntegrator function * TransferMechanism in integrator mode The reproducer from PrincetonUniversity#3142 was added as a test of mechanism reset in composition. Fixes: PrincetonUniversity#3142 Signed-off-by: Jan Vesely <[email protected]>
jvesely · Feb 10, 2025 · 9b41a73 · 9b41a73
1 parent 879f12a
commit 9b41a73
Show file tree

Hide file tree

Showing 4 changed files with 101 additions and 0 deletions.
diff --git a/psyneulink/core/components/mechanisms/mechanism.py b/psyneulink/core/components/mechanisms/mechanism.py
@@ -3241,6 +3241,9 @@ def _gen_llvm_function_reset(self, ctx, builder, m_base_params, m_state, m_arg_i
 
             builder.call(reinit_func, [reinit_params, reinit_state, reinit_in, reinit_out])
 
+        # update output ports after getting the reinitialized value
+        builder = self._gen_llvm_output_ports(ctx, builder, reinit_out, m_base_params, m_state, m_arg_in, m_arg_out)
+
         return builder
 
     def _gen_llvm_function(self, *, extra_args=[], ctx:pnlvm.LLVMBuilderContext, tags:frozenset):

diff --git a/tests/mechanisms/test_ddm_mechanism.py b/tests/mechanisms/test_ddm_mechanism.py
@@ -617,6 +617,28 @@ def test_WhenFinished_DDM_Analytical():
     c.is_satisfied()
 
 
+@pytest.mark.mechanism
+@pytest.mark.parametrize("initializer_param", [{}, {pnl.INITIALIZER: 5.0}], ids=["default_initializer", "custom_initializer"])
+@pytest.mark.parametrize("non_decision_time_param", [{}, {pnl.NON_DECISION_TIME: 13.0}], ids=["default_non_decision_time", "custom_non_decision_time"])
+def test_DDM_reset(mech_mode, initializer_param, non_decision_time_param):
+    D = pnl.DDM(function=pnl.DriftDiffusionIntegrator(**initializer_param, **non_decision_time_param))
+
+    ex = pytest.helpers.get_mech_execution(D, mech_mode)
+
+    initializer_value = initializer_param.get(pnl.INITIALIZER, 0)
+    non_decision_time_value = non_decision_time_param.get(pnl.NON_DECISION_TIME, 0)
+
+    ex([1])
+    ex([2])
+    result = ex([3])
+    np.testing.assert_array_equal(result, [[100], [102 - initializer_value + non_decision_time_value]])
+
+    reset_ex = pytest.helpers.get_mech_execution(D, mech_mode, tags=frozenset({"reset"}), member="reset")
+
+    reset_result = reset_ex(None if mech_mode == "Python" else [0])
+    np.testing.assert_array_equal(reset_result, [[initializer_value], [non_decision_time_value]])
+
+
 @pytest.mark.composition
 @pytest.mark.ddm_mechanism
 @pytest.mark.mechanism

diff --git a/tests/mechanisms/test_processing_mechanism.py b/tests/mechanisms/test_processing_mechanism.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 
+import psyneulink as pnl
 from psyneulink.core.components.functions.function import FunctionError
 from psyneulink.core.components.functions.nonstateful.learningfunctions import Hebbian, Reinforcement, TDLearning
 from psyneulink.core.components.functions.nonstateful.objectivefunctions import Distance
@@ -343,3 +344,26 @@ def test_output_ports2(self, op, expected):
         var = [1, 2, 4] if op in {MEAN, MEDIAN, STANDARD_DEVIATION, VARIANCE} else [1, 2, -4]
         PM1.execute(var)
         np.testing.assert_allclose(PM1.output_ports[0].value, expected)
+
+@pytest.mark.mechanism
+@pytest.mark.parametrize("initializer_param", [{}, {pnl.INITIALIZER: 5.0}], ids=["default_initializer", "custom_initializer"])
+def test_processing_mechanism_reset(mech_mode, initializer_param):
+    T = pnl.ProcessingMechanism(function=pnl.AdaptiveIntegrator(**initializer_param, rate=0.5))
+
+    ex = pytest.helpers.get_mech_execution(T, mech_mode)
+    initializer_value = initializer_param.get(pnl.INITIALIZER, 0)
+
+    ex([1])
+    ex([2])
+    result = ex([3])
+    np.testing.assert_array_equal(result, [[2.125 + initializer_value / 8]])
+
+    reset_ex = pytest.helpers.get_mech_execution(T, mech_mode, tags=frozenset({"reset"}), member="reset")
+
+    reset_result = reset_ex(None if mech_mode == "Python" else [0])
+
+    # FIXME: Python returns 3d value with default initializer
+    if mech_mode == "Python" and pnl.INITIALIZER not in initializer_param:
+        reset_result = reset_result[0]
+
+    np.testing.assert_array_equal(reset_result, [[initializer_value]])
diff --git a/tests/mechanisms/test_transfer_mechanism.py b/tests/mechanisms/test_transfer_mechanism.py
@@ -1793,3 +1793,55 @@ def test_combine_standard_output_port(self):
         # np.testing.assert_allclose(T.output_ports[1].value, [2.0])
         # np.testing.assert_allclose(T.output_ports[2].value, [3.0])
         # np.testing.assert_allclose(T.output_ports[3].value, [4.0])
+
+@pytest.mark.mechanism
+@pytest.mark.parametrize("initializer_param", [{}, {pnl.INITIALIZER: 5.0}], ids=["default_initializer", "custom_initializer"])
+def test_integrator_mode_reset(mech_mode, initializer_param):
+    T = pnl.TransferMechanism(integrator_mode=True,
+                              integrator_function=pnl.AdaptiveIntegrator(**initializer_param))
+
+    ex = pytest.helpers.get_mech_execution(T, mech_mode)
+    initializer_value = initializer_param.get(pnl.INITIALIZER, 0)
+
+    ex([1])
+    ex([2])
+    result = ex([3])
+    np.testing.assert_array_equal(result, [[2.125 + initializer_value / 8]])
+
+    reset_ex = pytest.helpers.get_mech_execution(T, mech_mode, tags=frozenset({"reset"}), member="reset")
+
+    reset_result = reset_ex(None if mech_mode == "Python" else [0])
+    np.testing.assert_array_equal(reset_result, [[initializer_value]])
+
+@pytest.mark.composition
+@pytest.mark.usefixtures("comp_mode_no_per_node")
+def test_integrator_mode_reset_in_composition(comp_mode):
+
+    # Note, input_mech is scheduled to only execute on pass 5!
+    input_mech = pnl.TransferMechanism(
+        input_shapes=2,
+        integrator_mode=True,
+        integration_rate=1,
+        reset_stateful_function_when=pnl.AtTrialStart()
+    )
+
+    lca = pnl.LCAMechanism(
+        input_shapes=2,
+        termination_threshold=10,
+        termination_measure=pnl.TimeScale.TRIAL,
+        execute_until_finished=False
+    )
+
+    gate = pnl.ProcessingMechanism(input_shapes=2)
+
+    comp = pnl.Composition()
+    comp.add_linear_processing_pathway([input_mech, lca, gate])
+
+    comp.scheduler.add_condition(input_mech, pnl.AtPass(5))
+
+    comp.scheduler.add_condition(lca, pnl.Always())
+
+    comp.scheduler.add_condition(gate, pnl.WhenFinished(lca))
+
+    comp.run([[1, 0], [0, 1]], execution_mode=comp_mode)
+    np.testing.assert_allclose(comp.results, [[[0.52293998, 0.40526519]], [[0.4336115, 0.46026939]]])