Search code examples
pythonscipycomplex-numbersnumpy-ufunc

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


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)


Solution

  • 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