NF - added __imp__ and implemented_function to lambdify

Signed-off-by: Ondrej Certik <ondrej@certik.cz>

sympy/utilities/lambdify.py

@@ -97,7 +97,7 @@ def _import(module, reload="False"):
     for sympyname, translation in translations.iteritems():
         namespace[sympyname] = namespace[translation]
 
-def lambdify(args, expr, modules=None):
+def lambdify(args, expr, modules=None, use_imps=True):
     """
     Returns a lambda function for fast calculation of numerical values.
 
@@ -154,6 +154,23 @@ def lambdify(args, expr, modules=None):
 
         Now f would look like:
         >> lambda x: my_cool_function(x)
+
+    Functions present in `expr` can also carry their own numerical
+    implementations, in a callable attached to the ``__imp__``
+    attribute.  Usually you attach this using the
+    ``implemented_function`` factory:
+
+    >>> from sympy.abc import x, y, z
+    >>> from sympy.utilities.lambdify import lambdify, implemented_function
+    >>> from sympy import Function
+    >>> f = implemented_function(Function('f'), lambda x : x+1)
+    >>> func = lambdify(x, f(x))
+    >>> func(4)
+    5
+
+    ``lambdify`` always prefers ``__imp__`` implementations to
+    implementations in other namespaces, unless the ``use_imps`` input
+    parameter is False.
     """
     # If the user hasn't specified any modules, use what is available.
     if modules is None:
@@ -165,18 +182,22 @@ def lambdify(args, expr, modules=None):
             modules = ("math", "numpy", "mpmath", "sympy")
         except ImportError:
             modules = ("math", "mpmath", "sympy")
-
+            
     # Get the needed namespaces.
-    if isinstance(modules, dict): # Check for dict before "__iter__"
-        namespace = _get_namespace(modules)
-    elif hasattr(modules, "__iter__"):
-        namespace = {}
-        # fill namespace with first having highest priority
-        for m in list(modules)[::-1]:
+    # First find any function implementations
+    namespaces = []
+    namespace = {}
+    if use_imps:
+        namespaces.append(_imp_namespace(expr))
+    # Check for dict before iterating
+    if isinstance(modules, dict) or not hasattr(modules, '__iter__'):
+        namespaces.append(modules)
+    else:
+        namespaces += list(modules)
+    # fill namespace with first having highest priority
+    for m in namespaces[::-1]:
             buf = _get_namespace(m)
             namespace.update(buf)
-    else:
-        namespace = _get_namespace(modules)
 
     if hasattr(expr, "atoms") :
         #Try if you can extract symbols from the expression.
@@ -229,3 +250,89 @@ def lambdastr(args, expr):
         args = str(args)
 
     return "lambda %s: (%s)" % (args, expr)
+
+def _imp_namespace(expr, namespace=None):
+    """ Return namespace dict with function implementations
+
+    We need to search for functions in anything that can be thrown at
+    us - that is - anything that could be passed as `expr`.  Examples
+    include sympy expressions, as well tuples, lists and dicts that may
+    contain sympy expressions.
+
+    Parameters
+    ----------
+    expr : object
+       Something passed to lambdify, that will generate valid code from
+       ``str(expr)``. 
+    namespace : None or mapping
+       Namespace to fill.  None results in new empty dict
+
+    Returns
+    -------
+    namespace : dict
+       dict with keys of implemented function names within `expr` and
+       corresponding values being the numerical implementation of
+       function
+    """
+    # Delayed import to avoid circular imports
+    from sympy.core.function import FunctionClass
+    if namespace is None:
+        namespace = {}
+    # tuples, lists, dicts are valid expressions
+    if isinstance(expr, (list, tuple)):
+        for arg in expr:
+            _imp_namespace(arg, namespace)
+        return namespace
+    elif isinstance(expr, dict):
+        for key, val in expr.items():
+            # functions can be in dictionary keys
+            _imp_namespace(key, namespace)
+            _imp_namespace(val, namespace)
+        return namespace
+    # sympy expressions may be Functions themselves
+    if hasattr(expr, 'func'):
+        if (isinstance(expr.func, FunctionClass) and
+            hasattr(expr.func, '__imp__')):
+            name = expr.func.__name__
+            imp = expr.func.__imp__
+            if name in namespace and namespace[name] != imp:
+                raise ValueError('We found more than one '
+                                 'implementation with name '
+                                 '"%s"' % name)
+            namespace[name] = imp
+    # and / or they may take Functions as arguments
+    if hasattr(expr, 'args'):
+        for arg in expr.args:
+            _imp_namespace(arg, namespace)
+    return namespace
+
+def implemented_function(symfunc, implementation):
+    """ Add numerical `implementation` to function `symfunc`
+
+    `symfunc` can by a Function, or a name, in which case we make an
+    anonymous function with this name.  The function is anonymous in the
+    sense that the name is not unique in the sympy namespace.
+
+    Parameters
+    ----------
+    symfunc : str or ``sympy.FunctionClass`` instance
+       If str, then create new anonymous sympy function with this as
+       name.  If `symfunc` is a sympy function, attach implementation to
+       function
+    implementation : callable
+       numerical implementation of function for use in ``lambdify``
+
+    Returns
+    -------
+    afunc : sympy.FunctionClass instance
+       function with attached implementation
+    """
+    # Delayed import to avoid circular imports
+    from sympy.core.function import FunctionClass, Function
+    # if name, create anonymous function to hold implementation
+    if isinstance(symfunc, basestring):
+        symfunc = FunctionClass(Function, symfunc)
+    # We need to attach as a method because symfunc will be a class
+    symfunc.__imp__ = staticmethod(implementation)
+    return symfunc
+

sympy/utilities/tests/test_lambdify.py

@@ -2,6 +2,7 @@
 from sympy import (symbols, lambdify, sqrt, sin, cos, pi, atan, Rational, Real,
         Matrix, Lambda, exp, Integral, oo)
 from sympy import mpmath
+from sympy.utilities.lambdify import implemented_function
 import math, sympy
 
 # high precision output of sin(0.2*pi) is used to detect if precision is lost unwanted
@@ -254,3 +255,61 @@ def test_namespace_order():
     if2 = lambdify(x, g(x), modules=(n2, "sympy"))
     # previously gave 'second f'
     assert if1(1) == 'first f'
+
+def test_imps():
+    # Here we check if the default returned functions are anonymous - in
+    # the sense that we can have more than one function with the same name
+    f = implemented_function('f', lambda x: 2*x)
+    g = implemented_function('f', lambda x: math.sqrt(x))
+    l1 = lambdify(x, f(x))
+    l2 = lambdify(x, g(x))
+    assert str(f(x)) == str(g(x))
+    assert l1(3) == 6
+    assert l2(3) == math.sqrt(3)
+    # check that we can pass in a Function as input
+    func = sympy.Function('myfunc')
+    assert not hasattr(func, '__imp__')
+    my_f = implemented_function(func, lambda x: 2*x)
+    assert hasattr(func, '__imp__')
+    # Error for functions with same name and different implementation
+    f2 = implemented_function("f", lambda x : x+101)
+    try:
+        lambdify(x, f(f2(x)))
+    except ValueError:
+        pass
+    else:
+        raise Exception('ValueError not raised')
+
+def test_lambdify_imps():
+    # Test lambdify with implemented functions
+    # first test basic (sympy) lambdify
+    f = sympy.cos
+    assert lambdify(x, f(x))(0) == 1
+    assert lambdify(x, 1 + f(x))(0) == 2
+    assert lambdify((x, y), y + f(x))(0, 1) == 2
+    # make an implemented function and test
+    f = implemented_function("f", lambda x : x+100)
+    assert lambdify(x, f(x))(0) == 100
+    assert lambdify(x, 1 + f(x))(0) == 101
+    assert lambdify((x, y), y + f(x))(0, 1) == 101
+    # Can also handle tuples, lists, dicts as expressions
+    lam = lambdify(x, (f(x), x))
+    assert lam(3) == (103, 3)
+    lam = lambdify(x, [f(x), x])
+    assert lam(3) == [103, 3]
+    lam = lambdify(x, [f(x), (f(x), x)])
+    assert lam(3) == [103, (103, 3)]
+    lam = lambdify(x, {f(x): x})
+    assert lam(3) == {103: 3}
+    lam = lambdify(x, {f(x): x})
+    assert lam(3) == {103: 3}
+    lam = lambdify(x, {x: f(x)})
+    assert lam(3) == {3: 103}
+    # Check that imp preferred to other namespaces by default
+    d = {'f': lambda x : x + 99}
+    lam = lambdify(x, f(x), d)
+    assert lam(3) == 103
+    # Unless flag passed
+    lam = lambdify(x, f(x), d, use_imps=False)
+    assert lam(3) == 102
+