Can a "ufuncify" function be set up to take complex input? "ufunc 'wrapper_module_0' not supported for the input types" error

Question

I am just learning how to use sympy for symbolic variable manipulations. I now have an expression that I'd like to evaluate many times in an optimization scheme, so I'd like for it to run very quickly. The sympy documentation on Numeric Computation describes a few methods for doing this: subs/evalf; lambdify; lambdify-numpy; ufuncify; Theano.

So far, I've gotten lambdify to work, but it still doesn't seem fast enough for me. Reading up on this further, it looks like lambdify operates at python speed, whereas ufuncify could be recompiling the code into C code under the hood, so I am now looking into this option.

So far, I have been able to ufuncify my expression, but it throws an error whenever I pass a complex input as an argument.

I modified the ufuncify example from this link to create my MWE. When I run this:

import sympy as sp
from sympy.utilities.autowrap import ufuncify
import numpy as np

# Create an example expression
a, b, c = sp.symbols('a, b, c')
expr = a + b*c

# Create a binary (compiled) function that broadcasts it's arguments
func = ufuncify((a, b, c), expr)
b = func(np.arange(5), 2.0+1j, 3.0)

print(b)

I get this error: TypeError: ufunc 'wrapper_module_0' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe''

If I change the code to just remove the imaginary part of the second argument, it runs fine:

import sympy as sp
from sympy.utilities.autowrap import ufuncify
import numpy as np

# Create an example expression
a, b, c = sp.symbols('a, b, c')
expr = a + b*c

# Create a binary (compiled) function that broadcasts it's arguments
func = ufuncify((a, b, c), expr)
b = func(np.arange(5), 2.0, 3.0)

print(b)

returns: [ 6. 7. 8. 9. 10.]

Ideally I would like to use cython or f2py as the backend, since I expect them to be the fastest... but I get similar errors:

func = ufuncify((a, b, c), expr, backend='cython')

returns TypeError: Argument 'b_5226316' has incorrect type (expected numpy.ndarray, got complex)

Answer 1

Could be that there is another answer, but I normally avoid encapsulating a complex value in sympy, so you can setup each value in a pair of ar and ai (for real and imaginary part).

You can then also define small expressions (a = ar + sp.I*ai) to make writing your code easier. Also choose upfront where you want to allow complex values.

The code would then be:

import sympy as sp
from sympy.utilities.autowrap import ufuncify
import numpy as np

# Create an example expression
ar, ai, br, bi, cr, ci = sp.symbols('ar, ai, br, bi, cr, ci', real=True)
a = ar + sp.I*ai
b = br + sp.I*bi
c = cr + sp.I*ci

expr = a + b*c

# Create a binary (compiled) function that broadcasts it's arguments
func = ufuncify((ar, ai, br, bi, cr, ci), expr)
b = func(np.arange(5), 0, 2.0, 1, 3.0, 0)

print(b)

I did not compile it, because I do not have the stuff installed, but it at least starts compiling so I am pretty sure that should work.

Update Not compiling was a mistake. The ufuncify function does not allow complex input at the moment (It is also mentioned as a milestone in the CodeGenerator). The ufuncify uses the Codegen module to create C-code from the given expression. In the case above that one looks like:

double autofunc0(double ar, double ai, double br, double bi, double cr, double ci) {

   double autofunc0_result;
   autofunc0_result = I*ai + ar + (I*bi + br)*(I*ci + cr);
   return autofunc0_result;

}

and it does not currently use any complex extensions, although those would be available in C99. Would be a nice extension to the code generator of Sympy I guess. In the documentation they also suggest that one can write ones own native code.

1) If you are really concerned about speed or memory optimizations,
  you will probably get better results by working directly with the
  wrapper tools and the low level code.  However, the files generated
  by this utility may provide a useful starting point and reference
  code. Temporary files will be left intact if you supply the keyword
  tempdir="path/to/files/".

I guess that would be an alternative option, therby using the generated files as a starting point and then using the <complex.h> header to implement the complex values. Not a satisfying answer, though.

Answer 2

In [251]: a, b, c = symbols('a, b, c') 
     ...: expr = a + b*c 
     ...:                                                                                            

In [252]: f = lambdify((a,b,c), expr)                                                                

In [254]: print(f.__doc__)                                                                           
Created with lambdify. Signature:

func(a, b, c)

Expression:

a + b*c

Source code:

def _lambdifygenerated(a, b, c):
    return (a + b*c)


Imported modules:

For this simple expr, the lambdified numpy code looks good - taking advantage of the whole array compiled operators:

In [255]: f(np.arange(5), 2.0, 3.0)                                                                  
Out[255]: array([ 6.,  7.,  8.,  9., 10.])

In [256]: timeit f(np.arange(5), 2.0, 3.0)                                                           
5.55 µs ± 16 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [257]: timeit np.arange(5) + 2.0 * 3.0                                                            
5.04 µs ± 16.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

It is faster if I use numba:

In [258]: import numba                                                                               

In [259]: @numba.njit 
     ...: def foo(a,b,c): 
     ...:     return a+b*c 
     ...:                                                                                            

In [260]: foo(np.arange(5),2.0,3.0)                                                                  
Out[260]: array([ 6.,  7.,  8.,  9., 10.])

In [261]: timeit foo(np.arange(5),2.0,3.0)                                                           
2.68 µs ± 69.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Answer 3

It's a little bit hacky and it may only work for f2py, but I think this is what you want to do:

import sympy as sp
from sympy.utilities.autowrap import ufuncify
import numpy as np

# This seems to be False by default and prevents the f2py codegen from
# creating complex variables in fortran
sp.utilities.codegen.COMPLEX_ALLOWED = True

# Create an example expression, mark them as Complex so that the codegen
# tool doesn't assume they're just floats. May not even be necessary.
a, b, c = sp.symbols('a, b, c', complex=True)
expr = a + b*c

# Create a binary (compiled) function that broadcasts its arguments
func = ufuncify((a, b, c), expr, backend="f2py")
# The f2py backend will want identically-sized arrays
b = func(np.arange(5), np.ones(5) * (2.0 + 1.0j), np.ones(5) * 3)

print(b)

You can confirm that the generated fortran code expects complex variables if you set tempdir="." on ufuncify and open the wrapped_code_0.f90 file.

!******************************************************************************
!*                      Code generated with sympy 1.6.2                       *
!*                                                                            *
!*              See http://www.sympy.org/ for more information.               *
!*                                                                            *
!*                      This file is part of 'autowrap'                       *
!******************************************************************************

subroutine autofunc(y_1038050, a_1038054, b_1038055, c_1038056, &
      m_1038051)
implicit none
INTEGER*4, intent(in) :: m_1038051
COMPLEX*16, intent(out), dimension(1:m_1038051) :: y_1038050
COMPLEX*16, intent(in), dimension(1:m_1038051) :: a_1038054
COMPLEX*16, intent(in), dimension(1:m_1038051) :: b_1038055
COMPLEX*16, intent(in), dimension(1:m_1038051) :: c_1038056
INTEGER*4 :: i_1038052

do i_1038052 = 1, m_1038051
   y_1038050(i_1038052) = a_1038054(i_1038052) + b_1038055(i_1038052)* &
         c_1038056(i_1038052)
end do

end subroutine