在numpy中，对数组或矩阵中的每个元素，进行函数计算

在numpy中，对数组或矩阵中的每个元素，进行函数计算

2019年12月6日 数据分析日记

今天学到了一个利用自建函数，对numpy矩阵种元素进行计算的新招

假设有一个数组或矩阵：
$$\left(\begin{array}{cc}0& 1\\ 1& 0\end{array}\right)$$

需对其中每个元素输入函数：

$$F(x)=44000+165000\times [1-e^{\frac{(-0.3)x}{6}}]$$
那么我们可以创建数据集：

import numpy as np
matrix_test = np.array([[0,1],[1,0]])

并且构造函数：

def function_A(x):
    return (440000 + 165000*(1 - np.exp((-0.3) * x / 6)))

然后直接计算：

function_A(matrix_test)

得到结果：

array([[440000.        , 448047.14495738],
       [448047.14495738, 440000.        ]])

然而，当这个函数变为一个非连续函数时：
$$\begin{gathered} F(x)=\{\begin{array}{cc} \\ 0& \text{if x = 0}\\ 44000+165000\times [1-e^{\frac{(-0.3)x}{6}}]& \text{if x ≥ 0}\end{array} \end{gathered}$$
我们可以构造新的函数：

def function_B(x):
    if x == 0:
        result = float(0)
    else:
        result = 440000 + 165000*(1 - np.exp((-0.3) * x / 6))
    return result

运行时会报错：

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-42-0423f1e4de9c> in <module>()
----> 1 function_B(matrix_test)

<ipython-input-41-99e5c180ace1> in function_B(x)
      1 def function_B(x):
----> 2     if x == 0:
      3         result = float(0)
      4     else:
      5         result = 440000 + 165000*(1 - np.exp((-0.3) * x / 6))

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

显然，对于带有if…else这样带条件的函数，并不能直接输入矩阵进行计算。查找numpy文档发现，应通过numpy.vectorize()这个函数函数进行处理（官方文档: https://docs.scipy.org/doc/numpy/reference/generated/numpy.vectorize.html )

因此先对function进行向量化处理：

function_vector = np.vectorize(function_B)
function_vector(matrix_test)

得到正确结果：

array([[     0.        , 448047.14495738],
       [448047.14495738,      0.        ]])