python - `numpy.piecewise` 丢弃函数的虚部。为什么以及如何解决？

Question

我定义了一个分段复杂函数

import numpy as np

def foo(x):
    return np.piecewise(x, [x>0], [np.exp(1j*x)])

当我尝试评估它时，它会发出警告。

print(foo(9.99))

输出:

-0.8444696962887724
C:\Users\pedro\Anaconda3\lib\site-packages\numpy\lib\function_base.py:1151: ComplexWarning: Casting complex values to real discards the imaginary part
  y[condlist[k]] = item

这是由于 numpy piecewise 造成的，它丢弃了 np.exp(1j*x) 的虚部...为什么会发生这种情况？通过运行 source code警告在以下几行中提出:

if not isinstance(item, collections.Callable):
    y[condlist[k]] = item
else:
    vals = x[condlist[k]]
    if vals.size > 0:
        y[condlist[k]] = item(vals)

我无法理解。另外，我该如何解决这个问题？

Answer 1

问题是第一个数组是一个 int/float 数组(在本例中是一个标量)。因此，分段将首先根据该值构造一个数组，然后从中调用函数。但此时，数组的类型已经固定。

但是，您可以通过将其变为复数并添加 0j 来解决该问题:

def foo(x):
    return np.piecewise(x<b>+0j</b>, [x>0], [np.exp(1j*x)])
    #                    ^ turn x into a complex number

或者使用complex(..)构造函数:

def foo(x):
    return np.piecewise(<b>complex(</b>x<b>)</b>, [x>0], [np.exp(1j*x)])
    #                   ^ turn x into a complex number

这会生成:

>>> foo(9.99)
array((-0.8444696962887724-0.5356033346142913j))

如果x已经是一个数组，您可以使用x.astype(dtype=np.complex)将其变成一个复杂的数组。