特殊函数 (scipy.special)#

scipy.special 包的主要功能是定义了众多数学物理中的特殊函数。可用的函数包括 Airy、椭圆、贝塞尔、伽马、贝塔、超几何、抛物柱面、Mathieu、扁球波、Struve 和 Kelvin 函数。还有一些低级统计函数,不建议一般使用,因为 stats 模块提供了这些函数的更易用接口。大多数这些函数可以接受数组参数,并返回遵循与 Numerical Python 中其他数学函数相同广播规则的数组结果。许多这些函数也接受复数作为输入。有关可用函数的完整列表及其一行描述,请输入 >>> help(special)。每个函数都有自己的文档,可通过 help 访问。如果你没有找到所需的函数,可以考虑编写它并将其贡献给库。你可以用 C、Fortran 或 Python 编写函数。查看库的源代码,了解这些类型函数的示例。

实数阶贝塞尔函数(jv, jn_zeros)#

贝塞尔函数是贝塞尔微分方程的一族解,具有实数或复数阶 alpha:

\[x^2 \frac{d^2 y}{dx^2} + x \frac{dy}{dx} + (x^2 - \alpha^2)y = 0\]

除了其他用途外,这些函数出现在波传播问题中,例如薄鼓头的振动模式。以下是一个边缘固定的圆形鼓头的示例:

alt:

“此代码生成一个鼓面振动模式的3-D表示,从四分之三角度观察。在X-Y平面上定义了一个圆形区域,边缘的Z值为0。在圆内,-X侧存在一个平滑的谷,+X侧存在一个平滑的峰。在这个角度下,图像类似于一个阴阳图。”

>>> from scipy import special
>>> import numpy as np
>>> def drumhead_height(n, k, distance, angle, t):
...    kth_zero = special.jn_zeros(n, k)[-1]
...    return np.cos(t) * np.cos(n*angle) * special.jn(n, distance*kth_zero)
>>> theta = np.r_[0:2*np.pi:50j]
>>> radius = np.r_[0:1:50j]
>>> x = np.array([r * np.cos(theta) for r in radius])
>>> y = np.array([r * np.sin(theta) for r in radius])
>>> z = np.array([drumhead_height(1, 1, r, theta, 0.5) for r in radius])
>>> import matplotlib.pyplot as plt
>>> fig = plt.figure()
>>> ax = fig.add_axes(rect=(0, 0.05, 0.95, 0.95), projection='3d')
>>> ax.plot_surface(x, y, z, rstride=1, cstride=1, cmap='RdBu_r', vmin=-0.5, vmax=0.5)
>>> ax.set_xlabel('X')
>>> ax.set_ylabel('Y')
>>> ax.set_xticks(np.arange(-1, 1.1, 0.5))
>>> ax.set_yticks(np.arange(-1, 1.1, 0.5))
>>> ax.set_zlabel('Z')
>>> plt.show()
../_images/special-1.png

特殊函数的Cython绑定 (scipy.special.cython_special)#

SciPy还提供了许多特殊函数的标量、类型化版本的Cython绑定。以下Cython代码给出了一个简单的示例,说明如何使用这些函数:

cimport scipy.special.cython_special as csc

cdef:
    double x = 1
    double complex z = 1 + 1j
    double si, ci, rgam
    double complex cgam

rgam = csc.gamma(x)
print(rgam)
cgam = csc.gamma(z)
print(cgam)
csc.sici(x, &si, &ci)
print(si, ci)

(参见 Cython 文档 以获取编译 Cython 的帮助。) 在示例中,函数 csc.gamma 本质上与它的 ufunc 对应物 gamma 类似,尽管它接受 C 类型作为参数而不是 NumPy 数组。特别注意,该函数被重载以支持实数和复数参数;在编译时会选择正确的变体。函数 csc.sicisici 的工作方式略有不同;对于 ufunc,我们可以写 ai, bi = sici(x),而在 Cython 版本中,多个返回值通过指针传递。可以将其视为类似于使用输出数组调用 ufunc:sici(x, out=(si, ci))

使用 Cython 绑定的两个潜在优势是:

  • 它们避免了 Python 函数开销

  • 它们不需要 Python 全局解释器锁(GIL)

以下部分讨论如何利用这些优势来潜在地加速代码,当然,首先应该对代码进行性能分析,以确保投入额外努力是值得的。

避免 Python 函数开销#

对于特殊函数中的 ufunc,通过向量化(即传递数组给函数)来避免 Python 函数开销。通常,这种方法效果很好,但有时在循环内部对标量输入调用特殊函数会更方便,例如,在实现自己的 ufunc 时。在这种情况下,Python 函数开销可能会变得显著。考虑以下示例:

import scipy.special as sc
cimport scipy.special.cython_special as csc

def python_tight_loop():
    cdef:
        int n
        double x = 1

    for n in range(100):
        sc.jv(n, x)

def cython_tight_loop():
    cdef:
        int n
        double x = 1

    for n in range(100):
        csc.jv(n, x)

在一台计算机上,python_tight_loop 运行大约需要 131 微秒,而 cython_tight_loop 运行大约需要 18.2 微秒。 运行。显然这个例子是人为设计的:人们可以直接调用 special.jv(np.arange(100), 1) 并获得与 cython_tight_loop 中一样快的速度。重点是,如果 Python 函数的开销在你的代码中变得显著,那么 Cython 绑定可能会派上用场。

释放 GIL#

人们经常需要在许多点上评估特殊函数,并且通常这些评估是微不足道的并行化。由于 Cython 绑定不需要 GIL,因此使用 Cython 的 prange 函数很容易并行运行它们。例如,假设我们想要计算亥姆霍兹方程的基本解:

\[\Delta_x G(x, y) + k^2G(x, y) = \delta(x - y),\]

其中 \(k\) 是波数,\(\delta\) 是狄拉克 delta 函数。已知在二维情况下,唯一的(辐射)解是

\[G(x, y) = \frac{i}{4}H_0^{(1)}(k|x - y|),\]

其中 \(H_0^{(1)}\) 是第一类汉克尔函数,即函数 hankel1。以下示例展示了我们如何并行计算此函数:

from libc.math cimport fabs
cimport cython
from cython.parallel cimport prange

import numpy as np
import scipy.special as sc
cimport scipy.special.cython_special as csc

def serial_G(k, x, y):
    return 0.25j*sc.hankel1(0, k*np.abs(x - y))

@cython.boundscheck(False)
@cython.wraparound(False)
cdef void _parallel_G(double k, double[:,:] x, double[:,:] y,
                      double complex[:,:] out) nogil:
    cdef int i, j

    for i in prange(x.shape[0]):
        for j in range(y.shape[0]):
            out[i,j] = 0.25j*csc.hankel1(0, k*fabs(x[i,j] - y[i,j]))

def parallel_G(k, x, y):
    out = np.empty_like(x, dtype='complex128')
    _parallel_G(k, x, y, out)
    return out

(有关在 Cython 中编译并行代码的帮助,请参见 这里。)如果上述 Cython 代码位于文件 test.pyx 中,那么我们可以编写一个 非正式基准测试,比较了函数的并行和串行版本:

import timeit

import numpy as np

from test import serial_G, parallel_G

def main():
    k = 1
    x, y = np.linspace(-100, 100, 1000), np.linspace(-100, 100, 1000)
    x, y = np.meshgrid(x, y)

    def serial():
        serial_G(k, x, y)

    def parallel():
        parallel_G(k, x, y)

    time_serial = timeit.timeit(serial, number=3)
    time_parallel = timeit.timeit(parallel, number=3)
    print("串行方法耗时 {:.3} 秒".format(time_serial))
    print("并行方法耗时 {:.3} 秒".format(time_parallel))

if __name__ == "__main__":
    main()

在一台四核计算机上,串行方法耗时 1.29 秒,而并行方法耗时 0.29 秒。

不在 scipy.special 中的函数#

一些函数未包含在 special 中,因为它们可以通过 NumPy 和 SciPy 中的现有函数直接实现。为了避免重复造轮子,本节提供了几个此类函数的实现,希望这些示例能说明如何处理类似的函数。在所有示例中,NumPy 被导入为 np,special 被导入为 sc

二元熵函数:

def binary_entropy(x):
    return -(sc.xlogy(x, x) + sc.xlog1py(1 - x, -x))/np.log(2)

在 [0, 1] 上的矩形阶跃函数:

def step(x):
    return 0.5*(np.sign(x) + np.sign(1 - x))

通过平移和缩放可以得到任意阶跃函数。

斜坡函数:

def ramp(x):
    return np.maximum(0, x)