如何向量化适配quadpy的Python函数？含矩阵运算场景

使用quadpy对函数积分时，发现quadpy会传入numpy数组作为参数（比如给f = lambda x: x**2传入x = [0, 0.3, 0.7, 1]这样的向量）。但要积分的函数长达300行，编写时只考虑了传入单个实数的情况，现在需要让函数支持向量输入，numpy.vectorize在部分场景无效。

典型问题示例：

def f(t):
   U = np.array([[1, 0,          0        ],
                 [0, np.cos(t),  np.sin(t)],
                 [0, -np.sin(t), np.cos(t)]])
   V = np.array([[np.cos(t),  0,  np.sin(t)],
                 [0,          1,  0        ],
                 [-np.sin(t), 0, np.cos(t)]])

   return U @ V

运行时quadpy传入数组后，np.cos(t)变成数组，系统报错需要指定dtype=object，但后续执行U @ V会崩溃。另外还有类似constant(t)*vector(l)的运算问题：常数变为向量后，vector(l)变成数组的数组，导致报错。想知道怎么解决这类问题，是否必须从一开始就用向量化思维编写函数？

不需要完全重写为向量化代码，有几种高效的改造方式：

1. 利用numpy广播机制，给数组增加维度 #

针对矩阵运算的示例，只需要让t的运算结果保持合适维度，就能让矩阵运算自动适配向量输入：

import numpy as np

def f(t):
    # 给t增加两个维度，让cos/sin结果从(n,)变为(n,1,1)，适配3x3矩阵的形状
    t = np.asarray(t)[..., np.newaxis, np.newaxis]
    U = np.array([[1, 0,          0        ],
                  [0, np.cos(t),  np.sin(t)],
                  [0, -np.sin(t), np.cos(t)]])
    V = np.array([[np.cos(t),  0,  np.sin(t)],
                  [0,          1,  0        ],
                  [-np.sin(t), 0, np.cos(t)]])
    # U和V现在形状为(n,3,3)，矩阵乘法会自动在每个n对应的矩阵上执行
    return U @ V

当t是数组时，每个元素都会对应一组3x3矩阵的运算，最终返回(n,3,3)形状的结果，完全适配quadpy的输入要求。

2. 保留原标量函数，逐元素处理 #

如果函数逻辑复杂，广播不好适配，可以保留原有的标量逻辑，通过循环逐元素处理后拼接结果：

def f_scalar(t):
    # 原有的标量版本函数
    U = np.array([[1, 0,          0        ],
                  [0, np.cos(t),  np.sin(t)],
                  [0, -np.sin(t), np.cos(t)]])
    V = np.array([[np.cos(t),  0,  np.sin(t)],
                  [0,          1,  0        ],
                  [-np.sin(t), 0, np.cos(t)]])
    return U @ V

def f(t):
    t_arr = np.asarray(t)
    # 遍历数组元素，调用原函数后拼接成结果数组
    return np.array([f_scalar(ti) for ti in t_arr])

这种方式不需要修改原有300行的核心逻辑，只需要套一层简单的循环即可适配向量输入。

3. 处理constant(t)*vector(l)类运算 #

如果constant(t)是标量函数，vector(l)是固定向量，只需要让constant(t)的结果形状和向量匹配，利用广播完成运算：

def constant(t):
    return np.cos(t)

vector_l = np.array([1,2,3])

def g(t):
    c = np.asarray(t)
    # 把c从(n,)变为(n,1)，和vector_l的(3,)广播成(n,3)形状
    return c[..., np.newaxis] * vector_l

这样当t是数组时，每个元素都会和整个向量相乘，得到正确的数组结果。

总结 #

不需要彻底重写300行的函数。优先尝试广播机制（性能最优），如果逻辑复杂就用逐元素处理的方式，两种方法都能快速适配向量输入，不用改变原有代码的核心逻辑结构。

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