Use IntFlag for cost functions (#2133)

Add test.
PrincetonUniversity · Oct 1, 2021 · e20a163 · e20a163
2 parents 5405ac9 + 9545538
commit e20a163
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 7 deletions.
diff --git a/psyneulink/core/components/functions/nonstateful/transferfunctions.py b/psyneulink/core/components/functions/nonstateful/transferfunctions.py
@@ -44,7 +44,7 @@
 import numbers
 import types
 import warnings
-from enum import IntEnum
+from enum import IntFlag
 from math import e, pi, sqrt
 
 import numpy as np
@@ -3364,7 +3364,7 @@ def _is_identity(self, context=None):
                      COMBINE_COSTS_FUNCTION]
 
 
-class CostFunctions(IntEnum):
+class CostFunctions(IntFlag):
     """Options for selecting constituent cost functions to be used by a `TransferWithCosts` Function.
 
     These can be used alone or in combination with one another, by enabling or disabling each using the

diff --git a/psyneulink/core/components/ports/modulatorysignals/controlsignal.py b/psyneulink/core/components/ports/modulatorysignals/controlsignal.py
@@ -1075,8 +1075,6 @@ def compute_costs(self, intensity, context=None):
             duration_cost = self.duration_cost_function(self.parameters.cost._get(context), context=context)
             self.parameters.duration_cost._set(duration_cost, context)
 
-        return max(0.0,
-                   self.combine_costs_function([intensity_cost,
-                                                adjustment_cost,
-                                                duration_cost],
-                                               context=context))
+        all_costs = [intensity_cost, adjustment_cost, duration_cost]
+        combined_cost = self.combine_costs_function(all_costs, context=context)
+        return max(0.0, combined_cost)
diff --git a/tests/composition/test_control.py b/tests/composition/test_control.py
@@ -1048,6 +1048,43 @@ def test_control_of_mech_port(self, comp_mode):
         results = comp.run(inputs=inputs, num_trials=1, execution_mode=comp_mode)
         assert np.allclose(comp.results, [[[0.375]]])
 
+    @pytest.mark.control
+    @pytest.mark.composition
+    @pytest.mark.parametrize("cost, expected, exp_values", [
+        (pnl.CostFunctions.NONE, 7.0, [1, 2, 3, 4, 5]),
+        (pnl.CostFunctions.INTENSITY, 3, [-1.71828183, -5.3890561, -17.08553692, -50.59815003, -143.4131591]),
+        (pnl.CostFunctions.ADJUSTMENT, 3, [1, 1, 1, 1, 1] ),
+        (pnl.CostFunctions.INTENSITY | pnl.CostFunctions.ADJUSTMENT, 3, [-1.71828183, -6.3890561, -19.08553692, -53.59815003, -147.4131591]),
+        (pnl.CostFunctions.DURATION, 3, [-19, -22., -25., -28., -31]),
+        # FIXME: combinations with DURATION are broken
+        # (pnl.CostFunctions.DURATION | pnl.CostFunctions.ADJUSTMENT, ,),
+        # (pnl.CostFunctions.ALL, ,),
+        pytest.param(pnl.CostFunctions.DEFAULTS, 3, [-1.71828183, -5.3890561, -17.08553692, -50.59815003, -143.4131591], id="CostFunctions.DEFAULT")],
+        ids=lambda x: x if isinstance(x, pnl.CostFunctions) else "")
+    def test_modulation_simple(self, cost, expected, exp_values):
+        obj = pnl.ObjectiveMechanism()
+        mech = pnl.ProcessingMechanism()
+
+        comp = pnl.Composition(controller_mode=pnl.BEFORE)
+        comp.add_node(mech, required_roles=pnl.NodeRole.INPUT)
+        comp.add_linear_processing_pathway([mech, obj])
+
+        comp.add_controller(
+            pnl.OptimizationControlMechanism(
+                objective_mechanism=obj,
+                control_signals=pnl.ControlSignal(
+                    modulates=('intercept', mech),
+                    modulation=pnl.OVERRIDE,
+                    allocation_samples=pnl.SampleSpec(start=1, stop=5, step=1),
+                    cost_options=cost,
+                )
+            )
+        )
+
+        ret = comp.run(inputs={mech: [2]}, num_trials=1)
+        assert np.allclose(ret, expected)
+        assert np.allclose([float(x) for x in comp.controller.function.saved_values], exp_values)
+
     @pytest.mark.control
     @pytest.mark.composition
     @pytest.mark.parametrize("mode", [pnl.ExecutionMode.Python])