Support custom atomic operators

cgp/__init__.py

@@ -7,7 +7,7 @@
 __doc__ = f"{__description__} <{__url__}>"
 
 from . import ea, local_search, utils
-from .cartesian_graph import CartesianGraph
+from .cartesian_graph import CartesianGraph, atomic_operator
 from .genome import Genome
 from .hl_api import evolve
 from .individual import IndividualMultiGenome, IndividualSingleGenome

cgp/cartesian_graph.py

@@ -29,6 +29,14 @@
     from .genome import Genome
 
 
+CUSTOM_ATOMIC_OPERATORS = {}
+
+
+def atomic_operator(f: Callable) -> Callable:
+    CUSTOM_ATOMIC_OPERATORS[f.__name__] = f
+    return f
+
+
 class CartesianGraph:
     """Class representing a particular Cartesian graph defined by a
     Genome.
@@ -256,7 +264,7 @@ def _f(x):
     return [{s}]
 """
         func_str = self._fill_parameter_values(func_str)
-        exec(func_str)
+        exec(func_str, {**globals(), **CUSTOM_ATOMIC_OPERATORS}, locals())
         return locals()["_f"]
 
     def _format_output_str_numpy_of_all_nodes(self):
@@ -299,7 +307,7 @@ def _f(x):
     return np.stack([{s}], axis=1)
 """
         func_str = self._fill_parameter_values(func_str)
-        exec(func_str)
+        exec(func_str, {**globals(), **CUSTOM_ATOMIC_OPERATORS}, locals())
 
         return locals()["_f"]
 
@@ -358,7 +366,7 @@ def forward(self, x):
         class_str += func_str
         class_str = self._fill_parameter_values(class_str)
 
-        exec(class_str)
+        exec(class_str, {**globals(), **CUSTOM_ATOMIC_OPERATORS}, locals())
         exec("_c = _C()")
 
         return locals()["_c"]

test/test_node.py

@@ -517,3 +517,137 @@ def test_repr():
     node = cgp.Add(idx, input_nodes)
     node_repr = str(node)
     assert node_repr == "Add(idx: 0, active: False, arity: 2, input_nodes [1, 2]"
+
+
+def test_custom_node():
+    class MyScaledAdd(cgp.node.OperatorNode):
+
+        _arity = 2
+        _def_output = "2.0 * (x_0 + x_1)"
+
+    primitives = (MyScaledAdd,)
+    params = {
+        "n_inputs": 2,
+        "n_outputs": 1,
+        "n_columns": 1,
+        "n_rows": 1,
+        "levels_back": 1,
+        "primitives": primitives,
+    }
+
+    genome = cgp.Genome(**params)
+
+    # f(x) = 2 * (x[0] + x[1])
+    genome.dna = [
+        ID_INPUT_NODE,
+        ID_NON_CODING_GENE,
+        ID_NON_CODING_GENE,
+        ID_INPUT_NODE,
+        ID_NON_CODING_GENE,
+        ID_NON_CODING_GENE,
+        0,
+        0,
+        1,
+        ID_OUTPUT_NODE,
+        2,
+        ID_NON_CODING_GENE,
+    ]
+
+    x = [5.0, 1.5]
+    y_target = [2 * (x[0] + x[1])]
+
+    _test_graph_call_and_to_x_compilations(genome, x, y_target)
+
+
+def test_custom_node_with_custom_atomic_operator():
+    @cgp.atomic_operator
+    def f_scale(x):
+        return 2.0 * x
+
+    class MyScaledAdd(cgp.node.OperatorNode):
+
+        _arity = 2
+        _def_output = "f_scale((x_0 + x_1))"
+
+    primitives = (MyScaledAdd,)
+    params = {
+        "n_inputs": 2,
+        "n_outputs": 1,
+        "n_columns": 1,
+        "n_rows": 1,
+        "levels_back": 1,
+        "primitives": primitives,
+    }
+
+    genome = cgp.Genome(**params)
+
+    # f(x) = f_scale(x[0] + x[1])
+    genome.dna = [
+        ID_INPUT_NODE,
+        ID_NON_CODING_GENE,
+        ID_NON_CODING_GENE,
+        ID_INPUT_NODE,
+        ID_NON_CODING_GENE,
+        ID_NON_CODING_GENE,
+        0,
+        0,
+        1,
+        ID_OUTPUT_NODE,
+        2,
+        ID_NON_CODING_GENE,
+    ]
+
+    x = [5.0, 1.5]
+    y_target = [2 * (x[0] + x[1])]
+
+    _test_graph_call_and_to_x_compilations(
+        genome, x, y_target, test_graph_call=False, test_to_sympy=False
+    )
+
+
+def test_custom_node_with_custom_atomic_operator_with_external_library():
+    scipy_const = pytest.importorskip("scipy.constants")
+
+    @cgp.atomic_operator
+    def f_scale(x):
+        return scipy_const.golden_ratio * x
+
+    class MyScaledAdd(cgp.node.OperatorNode):
+
+        _arity = 2
+        _def_output = "f_scale((x_0 + x_1))"
+
+    primitives = (MyScaledAdd,)
+    params = {
+        "n_inputs": 2,
+        "n_outputs": 1,
+        "n_columns": 1,
+        "n_rows": 1,
+        "levels_back": 1,
+        "primitives": primitives,
+    }
+
+    genome = cgp.Genome(**params)
+
+    # f(x) = f_scale(x[0] + x[1])
+    genome.dna = [
+        ID_INPUT_NODE,
+        ID_NON_CODING_GENE,
+        ID_NON_CODING_GENE,
+        ID_INPUT_NODE,
+        ID_NON_CODING_GENE,
+        ID_NON_CODING_GENE,
+        0,
+        0,
+        1,
+        ID_OUTPUT_NODE,
+        2,
+        ID_NON_CODING_GENE,
+    ]
+
+    x = [5.0, 1.5]
+    y_target = [scipy_const.golden_ratio * (x[0] + x[1])]
+
+    _test_graph_call_and_to_x_compilations(
+        genome, x, y_target, test_graph_call=False, test_to_sympy=False
+    )