SymPy autowrap仅返回实部，如何使其同时返回实部和虚部？

The issue you're seeing happens because autowrap defaults to generating code that returns a real number when working with real input variables, even if the expression is complex. Here are a few straightforward ways to fix this and get the full complex result:

Solution 1: Split into Real/Imaginary Parts and Return a Tuple #

You can explicitly split your complex expression into its real and imaginary components, then autowrap a function that returns both as a tuple. You can then combine them into a complex number after calling the function:

import sympy as sym
from sympy.utilities.autowrap import autowrap

x, y = sym.symbols('x y', real=True)
expr = sym.I * x + y**2

# Extract real and imaginary parts of the expression
real_part = sym.re(expr)
imag_part = sym.im(expr)

# Autowrap a function that returns both components
f = autowrap((real_part, imag_part), args=(x, y))

# Call the function and combine into a complex number
real, imag = f(1, 2)
result = complex(real, imag)
print(result)  # Output: 4+1j

Solution 2: Use the Cython Backend #

The default f2py backend sometimes has trouble handling complex return types with real inputs. Switching to the cython backend natively supports complex outputs and requires minimal changes to your code:

import sympy as sym
from sympy.utilities.autowrap import autowrap

x, y = sym.symbols('x y', real=True)
expr = sym.I * x + y**2

# Use cython backend to preserve complex output
f = autowrap(expr, args=(x, y), backend='cython')

print(f(1, 2))  # Output: 4+1j

Solution 3: Explicitly Specify Complex Return Type for Fortran (f2py) #

If you prefer to stick with the f2py backend, you can use codegen_kwargs to tell SymPy to generate Fortran code with a complex return type:

import sympy as sym
from sympy.utilities.autowrap import autowrap

x, y = sym.symbols('x y', real=True)
expr = sym.I * x + y**2

# Force the output to be complex*16 (double-precision complex)
f = autowrap(expr, args=(x, y),
             codegen_kwargs={'type_map': {expr: 'complex*16'}})

print(f(1, 2))  # Output: 4+1j

Each of these methods will give you the full complex result instead of just the real part. The Cython backend is often the most seamless option if you have Cython installed.

内容的提问来源于stack exchange，提问作者Apostolis