Allow lambdify to generate functions with out arguments (pyccel#1867)

Add a `use_out` parameter to `pyccel.lambdify` to avoid unnecessary
memory allocation. Auto-generate a docstring for functions generated via
calls to `pyccel.lambdify`.

Author: EmilyBourne

CHANGELOG.md

@@ -20,6 +20,8 @@ All notable changes to this project will be documented in this file.
 -   #1656 : Ensure `gFTL` is installed with Pyccel.
 -   #1830 : Add a `pyccel.lambdify.lambdify` function to accelerate SymPy expressions.
 -   #1844 : Add line numbers and code to errors from built-in function calls.
+-   #1867 : Add a `use_out` parameter to `pyccel.lambdify` to avoid unnecessary memory allocation.
+-   #1867 : Auto-generate a docstring for functions generated via calls to `pyccel.lambdify`.
 -   \[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

@@ -482,6 +482,22 @@ In practice `lambdify` uses SymPy's `NumPyPrinter` to generate code which is pas
 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.
 
+In order to make functions even faster it may be desirable to avoid unnecessary allocations inside the function. This functionality is similar to using an `out` argument in NumPy. Pyccel makes this functionality possible through the use of the `use_out` argument.
+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[:,:]'}, result_type = 'float[:,:]')
+x_2d = np.ones((4,2))
+y_2d = np.empty_like(x_2d)
+f(x_2d, y_2d)
+print(y_2d)
+```
+
 ## 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/devel/docs/openmp.md) for more details.

pyccel/commands/lambdify.py

@@ -14,7 +14,8 @@
     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):
+             templates : 'dict[str, list[str]]' = None, use_out = False,
+             **kwargs):
     """
     Convert a SymPy expression into a Pyccel-accelerated function.
 
@@ -37,13 +38,18 @@ def lambdify(expr : sp.Expr, args : 'dict[sp.Symbol, str]', *, result_type : str
         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
+    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.
+    use_out : bool, default=False
+        If true the function will modify an argument called 'out' instead
+        of returning a newly allocated array. If this argument is set then
+        result_type must be provided. This only works if the result is an
+        array type.
     **kwargs : dict
         Additional arguments that are passed to epyccel.
 
@@ -62,16 +68,37 @@ def lambdify(expr : sp.Expr, args : 'dict[sp.Symbol, str]', *, result_type : str
     """
     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.")
+    if result_type is not None and not isinstance(result_type, str):
+        raise TypeError("Argument 'result_type': Expected a string type annotation.")
 
     expr = NumPyPrinter().doprint(expr)
-    args = ', '.join(f'{a} : "{annot}"' for a, annot in args.items())
+    args_code = ', '.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}":'
+
+    docstring = " \n".join(('    """',
+            "    Expression evaluation created with `pyccel.lambdify`.",
+            "",
+            "    Function evaluating the expression:",
+           f"    {expr}",
+            "",
+            "    Parameters",
+            "    ----------\n"))
+    docstring += '\n'.join(f"    {a} : {type_annot}" for a, type_annot in args.items())
+
+    if use_out:
+        if not result_type:
+            raise TypeError("The result_type must be provided if use_out is true.")
+        else:
+            signature = f'def {func_name}({args_code}, out : "{result_type}"):'
+            docstring += f"\n    out : {result_type}"
+    elif result_type:
+        signature = f'def {func_name}({args_code}) -> "{result_type}":'
+        docstring += "\n".join(("\n",
+                    "     Returns",
+                    "     -------",
+                   f"     {result_type}"))
     else:
-        signature = f'def {func_name}({args}):'
+        signature = f'def {func_name}({args_code}):'
     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())):
@@ -81,9 +108,15 @@ def lambdify(expr : sp.Expr, args : 'dict[sp.Symbol, str]', *, result_type : str
                 for key, annotations in templates.items())
     else:
         decorators = ''
-    code = f'    return {expr}'
+    if use_out:
+        code = f'    out[:] = {expr}'
+    else:
+        code = f'    return {expr}'
     numpy_import = 'import numpy\n'
-    func = '\n'.join((numpy_import, decorators, signature, code))
+
+    docstring += '\n    """'
+
+    func = '\n'.join((numpy_import, decorators, signature, docstring, code))
     package = epyccel(func, **kwargs)
     return getattr(package, func_name)
 

tests/symbolic/test_symbolic.py

@@ -72,6 +72,35 @@ def test_lambdify(language):
         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)
 
+def test_lambdify_out_arg(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[:,:]',
+                    use_out = True, language = language)
+        pyc_y = pyc_lambdify(expr_y, {x : 'float[:,:]', y : 'float[:,:]'}, result_type = 'float[:,:]',
+                    use_out = True, language = language)
+
+        print(pyc_x.__doc__)
+
+        sp_out_x = np.empty_like(r)
+        sp_out_y = np.empty_like(r)
+        pyc_out_x = np.empty_like(r)
+        pyc_out_y = np.empty_like(r)
+        sp_out_x = sp_x(r, p)
+        sp_out_y = sp_y(r, p)
+        pyc_x(r, p, pyc_out_x)
+        pyc_y(r, p, out = pyc_out_y)
+
+        assert np.allclose(sp_out_x, pyc_out_x, rtol=RTOL, atol=ATOL)
+        assert np.allclose(sp_out_y, pyc_out_y, rtol=RTOL, atol=ATOL)
+
 ######################
 if __name__ == '__main__':
     print('*********************************')