Allow users to define custom atomic operators

[jakobj]

As discussed in #234, currently the compilation of computational graphs, for example to a numpy-compatible function, relies on the presence of the necessary imports (i.e., numpy) in cartesian_graph.py. We should try to avoid such dependency.
Users should be able to define custom atomic operators.

One possibility is to allow the user to set up a "compilation context" in which eval is called for the respective compilation targets. This could maybe also enable users to make use of custom functions in OperatorNode, e.g.,:

def myifelse(x_0, x_1, x_2):
    if x_0 < 0:
        return x_1
    else:
        return x_2

class MyIfElseNode(cgp.OperatorNode):
    _arity = 3
    _def_output = "myifelse(x_0, x_1, x_2)"

[jakobj]

some progress: the whole trick is to make the custom functions available when exec is called on the function strings (e.g., https://github.com/Happy-Algorithms-League/hal-cgp/blob/master/cgp/cartesian_graph.py#L259). one way to achieve this is to pass an appropriate mapping as the globals argument to exec. here's a quick and dirty implementation outside of hal-cgp as a proof of concept using context managers, but it's still too messy for my taste. custom functions are defined in the __init__ method of a custom context manager and then can be used in the function strings describing node operations just like the builtin functions.

file myeval.py

class EvalContext:

    def __init__(self, subclass_locals=None):
        self.funcs = {}
        for key in subclass_locals:
            if callable(subclass_locals[key]):
                self.funcs[key] = subclass_locals[key]

    def __enter__(self):
        return self.funcs

    def __exit__(self, type, value, traceback):
        pass


def myeval(expr, context=EvalContext):
    func_str = f"""\
def _f(x):
    return {expr}
"""
    with context() as c:
        exec(func_str, c, locals())
    return locals()["_f"]

file main.py

from myeval import myeval, EvalContext


class MyContext(EvalContext):

    def __init__(self):

        def myadd(x):
            return x[0] + x[1]

        super().__init__(locals())


if __name__ == '__main__':

    expr = "x[0] * myadd(x)"
    f = myeval(expr, MyContext)

    x = [3, 7]
    assert f(x) == x[0] * (x[0] + x[1])

[jakobj]

ok, i think i've found an implementation that allows for a nice API: we use a decorator to add user-defined atomic operations to a global (module-local) dict that is used during exec. it also supports using functions from other libraries without, as shown for numpy in the example below.

file myeval_decorator.py

CUSTOM_ATOMIC_OPERATIONS = {}

def atomic_operation(f):
    CUSTOM_ATOMIC_OPERATIONS[f.__name__] = f
    return f

def myeval(expr):
    func_str = f"""\
def _f(x):
    return {expr}
"""
    exec(func_str, CUSTOM_ATOMIC_OPERATIONS, locals())
    return locals()["_f"]

file main.py

import numpy as np

from myeval_decorator import myeval, atomic_operation

@atomic_operation
def myadd(x):
    return x[0] + x[1]

@atomic_operation
def mynumpyadd(x):
    return np.sum(x)

if __name__ == '__main__':

    expr = "x[0] * mynumpyadd(x)"
    f = myeval(expr)

    x = [3, 7]
    assert f(x) == x[0] * (x[0] + x[1])

@mschmidt87 @HenrikMettler what do you think?

[HenrikMettler]

Wow this fancy stuff! 😍
Am I seeing it right that with this the user could compile the cartesian graph into "arbitrary" functions from other libraries? Or am I missing a step?

[jakobj]

Am I seeing it right that with this the user could compile the cartesian graph into "arbitrary" functions from other libraries? Or am I missing a step?

partly. the general idea is that one can use arbitrary functions in the definition of node operations, see the start of this issue. however, the compilation will still have to be made for one specific target we support, i.e., python function, numpy-compatible, torch-compatible, or sympy compatible. does that answer your question? if not could you provide a code sample of what you have in mind?

[HenrikMettler]

partly. the general idea is that one can use arbitrary functions in the definition of node operations, see the start of this issue. however, the compilation will still have to be made for one specific target we support, i.e., python function, numpy-compatible, torch-compatible, or sympy compatible. does that answer your question? if not could you provide a code sample of what you have in mind?

Ah yes I understand, thanks. There is much more to library specifics in the compilation calls, than in the simple myeval above.