Add a pyccel.lambdify function with similar functionality to sympy.lambdify (pyccel#1829)

Given a SymPy expression `f` and type annotations, `lambdify` returns
a "pyccelised" function `f_fast` that can be used in the same Python
session. Fixes pyccel#1827.

Author: EmilyBourne

.dict_custom.txt

@@ -5,6 +5,8 @@ NumPy
 NumPy's
 CuPy
 CuPy's
+SymPy
+SymPy's
 BLAS
 LAPACK
 MPI

.github/actions/pytest_run/action.yml

@@ -28,7 +28,7 @@ runs:
       shell: ${{ inputs.shell_cmd }}
       working-directory: ./tests
     - name: Test C translation
-      run: python -m pytest -n auto -rX ${FLAGS} -m "not (parallel or xdist_incompatible) and c ${{ inputs.pytest_mark }}" --ignore=symbolic --ignore=ndarrays 2>&1 | tee s1_outfile.out
+      run: python -m pytest -n auto -rX ${FLAGS} -m "not (parallel or xdist_incompatible) and c ${{ inputs.pytest_mark }}" --ignore=ndarrays 2>&1 | tee s1_outfile.out
       shell: ${{ inputs.shell_cmd }}
       working-directory: ./tests
       id: pytest_1
@@ -45,19 +45,19 @@ runs:
       id: pytest_2
     - name: Test multi-file C translations
       run: |
-         python -m pytest -rX ${FLAGS} -m "xdist_incompatible and not parallel and c ${{ inputs.pytest_mark }}" --ignore=symbolic --ignore=ndarrays 2>&1 | tee s3_outfile.out
+         python -m pytest -rX ${FLAGS} -m "xdist_incompatible and not parallel and c ${{ inputs.pytest_mark }}" --ignore=ndarrays 2>&1 | tee s3_outfile.out
          pyccel-clean
       shell: ${{ inputs.shell_cmd }}
       working-directory: ./tests
       id: pytest_3
     - name: Test Fortran translations
-      run: python -m pytest -n auto -rX ${FLAGS} -m "not (parallel or xdist_incompatible) and not (c or python) ${{ inputs.pytest_mark }}" --ignore=symbolic --ignore=ndarrays 2>&1 | tee s4_outfile.out
+      run: python -m pytest -n auto -rX ${FLAGS} -m "not (parallel or xdist_incompatible) and not (c or python) ${{ inputs.pytest_mark }}" --ignore=ndarrays 2>&1 | tee s4_outfile.out
       shell: ${{ inputs.shell_cmd }}
       working-directory: ./tests
       id: pytest_4
     - name: Test multi-file Fortran translations
       run: |
-        python -m pytest -rX ${FLAGS} -m "xdist_incompatible and not parallel and not (c or python) ${{ inputs.pytest_mark }}" --ignore=symbolic --ignore=ndarrays 2>&1 | tee s5_outfile.out
+        python -m pytest -rX ${FLAGS} -m "xdist_incompatible and not parallel and not (c or python) ${{ inputs.pytest_mark }}" --ignore=ndarrays 2>&1 | tee s5_outfile.out
         pyccel-clean
       shell: ${{ inputs.shell_cmd }}
       working-directory: ./tests

CHANGELOG.md

@@ -14,6 +14,7 @@ All notable changes to this project will be documented in this file.
 -   #1750 : Add Python support for set method `remove()`.
 -   #1787 : Ensure `STC` is installed with Pyccel.
 -   #1743 : Add Python support for set method `discard()`.
+-   #1830 : Add a `pyccel.lambdify.lambdify` function to accelerate SymPy expressions.
 -   \[INTERNALS\] Added `container_rank` property to `ast.datatypes.PyccelType` objects.
 -   \[DEVELOPER\] Added an improved traceback to the developer-mode errors for errors in function calls.
 

docs/quickstart.md

@@ -453,6 +453,35 @@ Out[9]: 210.99245283018868
 ```
 After subtracting the amount of time required to create an array copy from the given times, we can conclude that the pyccelised function is approximately 210 times faster than the original Python function.
 
+### Interactive Usage with `lambdify`
+
+While Pyccel is usually used to accelerate Python code, it is also possible to accelerate other expressions. The Pyccel library provides the `lambdify` Python function. This function is similar to SymPy's [`lambdify`](https://docs.sympy.org/latest/modules/utilities/lambdify.html) function, given a SymPy expression `f` and type annotations, `lambdify` returns a "pyccelised" function `f_fast` that can be used in the same Python session.
+For example:
+```python
+import numpy as np
+import sympy as sp
+from pyccel import lambdify
+
+x = sp.Symbol('x')
+expr = x**2 + x*5
+f = lambdify(expr, {x : 'float'})
+print(f(3.0))
+
+expr2 = x-x
+f2 = lambdify(expr, {x : 'float'}, result_type = 'float')
+print(f2(3.0))
+
+expr = x**2 + x*5 + 4.5
+f3 = lambdify(expr, {x : 'T'}, templates = {'T': ['float[:]', 'float[:,:]']})
+x_1d = np.ones(4)
+x_2d = np.ones((4,2))
+print(f3(x_1d))
+print(f3(x_2d))
+```
+In practice `lambdify` uses SymPy's `NumPyPrinter` to generate code which is passed to the `epyccel` function. The `epyccel` function copies the code into a temporary Python file in the `__epyccel__` directory.
+Once the file has been copied, `epyccel` calls the `pyccel` command to generate a Python C extension module that contains a single pyccelised function.
+Then finally, it imports this function and returns it to the caller.
+
 ## Other Features
 
 Pyccel's generated code can use parallel multi-threading through [OpenMP](https://en.wikipedia.org/wiki/OpenMP); please read [our documentation](https://github.com/pyccel/pyccel/blob/master/tutorial/openmp.md) for more details.

pyccel/__init__.py

@@ -1 +1,2 @@
 from .version import __version__
+from .commands.lambdify import lambdify

pyccel/commands/lambdify.py

@@ -0,0 +1,89 @@
+"""
+File describing commands associated with the lambdify function which converts a
+SymPy expression into a Pyccel-accelerated function.
+"""
+import sympy as sp
+from packaging import version
+
+from pyccel.epyccel import epyccel
+from pyccel.utilities.strings import random_string
+
+if version.parse(sp.__version__) >= version.parse('1.8'):
+    from sympy.printing.numpy import NumPyPrinter
+else:
+    from sympy.printing.pycode import NumPyPrinter
+
+def lambdify(expr : sp.Expr, args : 'dict[sp.Symbol, str]', *, result_type : str = None,
+             templates : 'dict[str,list[str]]' = None, **kwargs):
+    """
+    Convert a SymPy expression into a Pyccel-accelerated function.
+
+    Convert a SymPy expression into a function that allows for fast
+    numeric evaluation. This is done using SymPy's NumPyPrinter to
+    generate code that can be accelerated by Pyccel.
+
+    Parameters
+    ----------
+    expr : sp.Expr
+        The SymPy expression that should be returned from the function.
+    args : dict[sp.Symbol, str]
+        A dictionary of the arguments of the function being created.
+        The keys are variables representing the arguments that will be
+        passed to the function. The values are the the type annotations
+        for those functions.
+    result_type : str, optional
+        The type annotation for the result of the function. This argument
+        is optional but it is recommended to provide it as SymPy
+        expressions do not always evaluate to the expected type. For
+        example if the SymPy expression simplifies to 0 then the default
+        type will be int even if the arguments are floats.
+    templates : dict[str,list[str]], optional
+        A description of any templates that should be added to the
+        function. The keys are the symbols which can be used as type
+        specifiers, the values are a list of the type annotations which
+        are valid types for the symbol. See
+        <https://github.com/pyccel/pyccel/blob/devel/docs/templates.md>
+        for more details.
+    **kwargs : dict
+        Additional arguments that are passed to epyccel.
+
+    Returns
+    -------
+    func
+        A Pyccel-accelerated function which allows the evaluation of
+        the SymPy expression.
+
+    See Also
+    --------
+    sympy.lambdify
+        <https://docs.sympy.org/latest/modules/utilities/lambdify.html>.
+    epyccel
+        The function that accelerates the generated code.
+    """
+    if not (isinstance(args, dict) and all(isinstance(k, sp.Symbol) and isinstance(v, str) for k,v in args.items())):
+        raise TypeError("Argument 'args': Expected a dictionary mapping SymPy symbols to string type annotations.")
+
+    expr = NumPyPrinter().doprint(expr)
+    args = ', '.join(f'{a} : "{annot}"' for a, annot in args.items())
+    func_name = 'func_'+random_string(8)
+    if result_type:
+        if not isinstance(result_type, str):
+            raise TypeError("Argument 'result_type': Expected a string type annotation.")
+        signature = f'def {func_name}({args}) -> "{result_type}":'
+    else:
+        signature = f'def {func_name}({args}):'
+    if templates:
+        if not (isinstance(templates, dict) and all(isinstance(k, str) and hasattr(v, '__iter__') for k,v in templates.items()) \
+                and all(all(isinstance(type_annot, str) for type_annot in v) for v in templates.values())):
+            raise TypeError("Argument 'templates': Expected a dictionary mapping strings describing type specifiers to lists of string type annotations.")
+
+        decorators = '\n'.join(f'@template("{key}", ['+', '.join(f'"{annot}"' for annot in annotations)+'])' \
+                for key, annotations in templates.items())
+    else:
+        decorators = ''
+    code = f'    return {expr}'
+    numpy_import = 'import numpy\n'
+    func = '\n'.join((numpy_import, decorators, signature, code))
+    package = epyccel(func, **kwargs)
+    return getattr(package, func_name)
+

pyccel/epyccel.py

@@ -118,8 +118,11 @@ def epyccel_seq(function_or_module, *,
 
     Parameters
     ----------
-    function_or_module : function | module
+    function_or_module : function | module | str
         Python function or module to be accelerated.
+        If a string is passed then it is assumed to be the code from a module which
+        should be accelerated. The module must be capable of running as a standalone
+        file so it must include any necessary import statements.
     language : {'fortran', 'c', 'python'}
         Language of generated code (default: 'fortran').
     compiler : str, optional
@@ -167,7 +170,7 @@ def epyccel_seq(function_or_module, *,
     # Store current directory
     base_dirpath = os.getcwd()
 
-    if isinstance(function_or_module, (FunctionType, type)):
+    if isinstance(function_or_module, (FunctionType, type, str)):
         dirpath = os.getcwd()
 
     elif isinstance(function_or_module, ModuleType):
@@ -197,6 +200,11 @@ def epyccel_seq(function_or_module, *,
 
         module_name, module_lock = get_unique_name(pymod.__name__, epyccel_dirpath)
 
+    elif isinstance(function_or_module, str):
+        code = function_or_module
+
+        module_name, module_lock = get_unique_name('mod', epyccel_dirpath)
+
     else:
         raise TypeError('> Expecting a FunctionType, type or a ModuleType')
 
@@ -277,8 +285,10 @@ def epyccel( python_function_or_module, **kwargs ):
 
     Parameters
     ----------
-    python_function_or_module : function | module
+    python_function_or_module : function | module | str
         Python function or module to be accelerated.
+        If a string is passed then it is assumed to be the code from a module which
+        should be accelerated..
     **kwargs :
         Additional keyword arguments for configuring the compilation and acceleration process.
         Available options are defined in epyccel_seq.
@@ -300,7 +310,7 @@ def epyccel( python_function_or_module, **kwargs ):
     >>> one_f = epyccel(one, language='fortran')
     >>> one_c = epyccel(one, language='c')
     """
-    assert isinstance( python_function_or_module, (FunctionType, type, ModuleType) )
+    assert isinstance( python_function_or_module, (FunctionType, type, ModuleType, str) )
 
     comm  = kwargs.pop('comm', None)
     root  = kwargs.pop('root', 0)

pyproject.toml

@@ -21,6 +21,7 @@ dependencies = [
     "termcolor >= 1.0.0",
     "textx >= 2.2",
     "astunparse >= 1.6.0", # astunparse is only needed for Python3.8, we should use ast.unparse when we drop Python3.8.
+    "packaging",
 ]
 
 [project.optional-dependencies]

tests/symbolic/mappings.py

@@ -0,0 +1,31 @@
+""" A file containing 2D analytical mappings to test Pyccel lambdify function.
+"""
+
+class PolarMapping:
+    """
+    Represents a Polar 2D Mapping object (Annulus).
+    """
+    expressions = {'x': 'c1 + (rmin*(1-x1)+rmax*x1)*cos(x2)',
+                    'y': 'c2 + (rmin*(1-x1)+rmax*x1)*sin(x2)'}
+    constants = {'rmin': 0.0, 'rmax': 1.0, 'c1' : 0.0, 'c2' : 0.0}
+
+#==============================================================================
+class TargetMapping:
+    """
+    Represents a Target 2D Mapping object.
+    """
+    expressions = {'x': 'c1 + (1-k)*x1*cos(x2) - D*x1**2',
+                    'y': 'c2 + (1+k)*x1*sin(x2)'}
+    constants = {'c1': 0.0, 'c2': 0.0, 'k' : 0.5, 'D' : '1.0'}
+
+#==============================================================================
+class CzarnyMapping:
+    """
+    Represents a Czarny 2D Mapping object.
+    """
+    expressions = {'x': '(1 - sqrt( 1 + eps*(eps + 2*x1*cos(x2)) )) / eps',
+                    'y': 'c2 + (b / sqrt(1-eps**2/4) * x1 * sin(x2)) /'
+                        '(2 - sqrt( 1 + eps*(eps + 2*x1*cos(x2)) ))'}
+    constants = {'eps' : 0.1, 'c2' : 0.0, 'b' : 1.0}
+
+

tests/symbolic/test_symbolic.py

@@ -4,18 +4,27 @@
 import os
 import pytest
 
+import numpy as np
+import sympy as sp
+
+import mappings
+
+from pyccel import lambdify as pyc_lambdify
 from pyccel.parser.parser   import Parser
 from pyccel.codegen.codegen import Codegen
 from pyccel.errors.errors   import Errors
 
+RTOL = 1e-12
+ATOL = 1e-16
+
 base_dir = os.path.dirname(os.path.realpath(__file__))
 path_dir = os.path.join(base_dir, 'scripts')
 
 files = sorted(os.listdir(path_dir))
-files = [f for f in files if (f.endswith(".py"))]
+files = [f for f in files if f.endswith(".py")]
 
 @pytest.mark.parametrize( "f", files )
-@pytest.mark.xfail(reason="Broken symbolic function support, see issue #330")
+@pytest.mark.skip(reason="Broken symbolic function support, see issue #330")
 def test_symbolic(f):
 
     pyccel = Parser(f)
@@ -34,6 +43,35 @@ def test_symbolic(f):
     errors = Errors()
     errors.reset()
 
+def test_lambdify(language):
+    r1 = np.linspace(0.0, 1.0, 100)
+    p1 = np.linspace(0.0, 2*np.pi, 100)
+    r,p = np.meshgrid(r1, p1)
+    x,y = sp.symbols('x1,x2')
+    for m in (mappings.PolarMapping, mappings.TargetMapping, mappings.CzarnyMapping):
+        expr_x = sp.sympify(m.expressions['x']).subs(m.constants)
+        expr_y = sp.sympify(m.expressions['y']).subs(m.constants)
+        sp_x = sp.lambdify([x, y], expr_x)
+        sp_y = sp.lambdify([x, y], expr_y)
+        pyc_x = pyc_lambdify(expr_x, {x : 'float[:,:]', y : 'float[:,:]'}, result_type = 'float[:,:]',
+                    language = language)
+        pyc_y = pyc_lambdify(expr_y, {x : 'float[:,:]', y : 'float[:,:]'}, result_type = 'float[:,:]',
+                    language = language)
+
+        assert np.allclose(sp_x(r, p), pyc_x(r, p), rtol=RTOL, atol=ATOL)
+        assert np.allclose(sp_y(r, p), pyc_y(r, p), rtol=RTOL, atol=ATOL)
+
+        pyc_x = pyc_lambdify(expr_x, {x : 'T', y : 'T'}, templates = {'T': ['float[:]', 'float[:,:]']},
+                    language = language)
+        pyc_y = pyc_lambdify(expr_y, {x : 'T', y : 'T'}, templates = {'T': ['float[:]', 'float[:,:]']},
+                    language = language)
+
+        assert np.allclose(sp_x(r, p), pyc_x(r, p), rtol=RTOL, atol=ATOL)
+        assert np.allclose(sp_y(r, p), pyc_y(r, p), rtol=RTOL, atol=ATOL)
+
+        assert np.allclose(sp_x(r[0,:], p[0,:]), pyc_x(r[0,:], p[0,:]), rtol=RTOL, atol=ATOL)
+        assert np.allclose(sp_y(r[0,:], p[0,:]), pyc_y(r[0,:], p[0,:]), rtol=RTOL, atol=ATOL)
+
 ######################
 if __name__ == '__main__':
     print('*********************************')
@@ -43,6 +81,6 @@ def test_symbolic(f):
     print('*********************************')
 
     for f in files:
-        print('> testing {0}'.format(str(f)))
+        print(f'> testing {f}')
         test_symbolic(f)
         print('\n')