如何在numpy中将傅里叶级数部分和向量化

Question

给定函数周期为T，傅里叶级数系数a[n]和b[n]（分别对应余弦和正弦），以及等间隔区间t，以下代码将计算区间t中所有点的部分和（其中a、b和t都是numpy数组）。需要注意的是，t和a长度不相等。

yn=ones(len(t))*a[0]
for n in range(1,len(a)):
    yn=yn+(a[n]*cos(2*pi*n*t/T)-b[n]*sin(2*pi*n*t/T))

我的问题是：这个for循环能被向量化吗？

- NameOfTheRose

---

你的例子中有个错误拼写，肯定是指 len(a) 吧？ - Ahmed Fasih

@Ahmed Fasih，当然正确，已更正，谢谢。 - NameOfTheRose

为什么不直接使用傅里叶反变换呢？如果你想要插值，可以用零填充高频部分。 - Dirklinux

@Dirklinux，我更喜欢自己做，这样可以让我更好地理解。 - NameOfTheRose

Answer 1

这里有一种矢量化的方法，利用broadcasting来创建余弦/正弦输入的2D数组版本：2*pi*n*t/T，然后使用np.dot进行sum-reduction的matrix-multiplication。

r = np.arange(1,len(a))
S = 2*np.pi*r[:,None]*t/T
cS = np.cos(S)
sS = np.sin(S)
out = a[1:].dot(cS) - b[1:].dot(sS) + a[0]

进一步提高性能

为了进一步提高性能，我们可以利用numexpr模块来计算这些三角函数步骤。

import numexpr as ne
cS = ne.evaluate('cos(S)')
sS = ne.evaluate('sin(S)')

运行时测试 -

方法 -

def original_app(t,a,b,T):
    yn=np.ones(len(t))*a[0]
    for n in range(1,len(a)):
        yn=yn+(a[n]*np.cos(2*np.pi*n*t/T)-b[n]*np.sin(2*np.pi*n*t/T))
    return yn

def vectorized_app(t,a,b,T):    
    r = np.arange(1,len(a))
    S = (2*np.pi/T)*r[:,None]*t
    cS = np.cos(S)
    sS = np.sin(S)
    return a[1:].dot(cS) - b[1:].dot(sS) + a[0]

def vectorized_app_v2(t,a,b,T):    
    r = np.arange(1,len(a))
    S = (2*np.pi/T)*r[:,None]*t
    cS = ne.evaluate('cos(S)')
    sS = ne.evaluate('sin(S)')
    return a[1:].dot(cS) - b[1:].dot(sS) + a[0]

此外，还包括@Paul Panzer帖子中的PP函数。

时间 -

In [22]: # Setup inputs
    ...: n = 10000
    ...: t = np.random.randint(0,9,(n))
    ...: a = np.random.randint(0,9,(n))
    ...: b = np.random.randint(0,9,(n))
    ...: T = 3.45
    ...: 

In [23]: print np.allclose(original_app(t,a,b,T), vectorized_app(t,a,b,T))
    ...: print np.allclose(original_app(t,a,b,T), vectorized_app_v2(t,a,b,T))
    ...: print np.allclose(original_app(t,a,b,T), PP(t,a,b,T))
    ...: 
True
True
True

In [25]: %timeit original_app(t,a,b,T)
    ...: %timeit vectorized_app(t,a,b,T)
    ...: %timeit vectorized_app_v2(t,a,b,T)
    ...: %timeit PP(t,a,b,T)
    ...: 
1 loops, best of 3: 6.49 s per loop
1 loops, best of 3: 6.24 s per loop
1 loops, best of 3: 1.54 s per loop
1 loops, best of 3: 1.96 s per loop

- Divakar

---

谢谢。你假设a和t的长度相同，这不是我想要的；但是很容易解决。我会发布修改后的代码。我还没有研究过优化。我也注意到你把标签从傅里叶变成了FFT。我不同意。 - NameOfTheRose

@NameOfTheRose 原始代码中使用了 for n in range(1,len(t)):，后来更改为使用 len(a)。因此，我需要编辑 r 部分以使用 len(a)。已做出相应修改。 - Divakar

那是我的错误，我道歉。既然您已经编辑了您的帖子，我就没有必要再发布修改后的代码了，这只是一个非常简单的修改。 - NameOfTheRose

Answer 2

无法与numexpr相比，但如果不可用，我们可以节省超越函数（测试和基准测试代码大量基于@Divakar的代码，如果您没有注意到;-)）：

import numpy as np
from timeit import timeit

def PP(t,a,b,T):
    CS = np.empty((len(t), len(a)-1), np.complex)
    CS[...] = np.exp(2j*np.pi*(t[:, None])/T)
    np.cumprod(CS, axis=-1, out=CS)
    return a[1:].dot(CS.T.real) - b[1:].dot(CS.T.imag) + a[0]

def original_app(t,a,b,T):
    yn=np.ones(len(t))*a[0]
    for n in range(1,len(a)):
        yn=yn+(a[n]*np.cos(2*np.pi*n*t/T)-b[n]*np.sin(2*np.pi*n*t/T))
    return yn

def vectorized_app(t,a,b,T):    
    r = np.arange(1,len(a))
    S = 2*np.pi*r[:,None]*t/T
    cS = np.cos(S)
    sS = np.sin(S)
    return a[1:].dot(cS) - b[1:].dot(sS) + a[0]

n = 1000
t = 2000
t = np.random.randint(0,9,(t))
a = np.random.randint(0,9,(n))
b = np.random.randint(0,9,(n))
T = 3.45


print(np.allclose(original_app(t,a,b,T), vectorized_app(t,a,b,T)))
print(np.allclose(original_app(t,a,b,T), PP(t,a,b,T)))

print('{:18s} {:9.6f}'.format('orig', timeit(lambda: original_app(t,a,b,T), number=10)/10))
print('{:18s} {:9.6f}'.format('Divakar no numexpr', timeit(lambda: vectorized_app(t,a,b,T), number=10)/10))
print('{:18s} {:9.6f}'.format('PP', timeit(lambda: PP(t,a,b,T), number=10)/10))

输出：

True
True
orig                0.166903
Divakar no numexpr  0.179617
PP                  0.060817

顺便提一下，如果 delta t 能够整除 T，那么可以节省更多的时间，或者运行完整个 fft 算法并且丢弃多余部分。

- Paul Panzer

---

你关于超越数的评论让我希望Numpy很快能够得到sincos函数！ - Ahmed Fasih

@AhmedFasih 很有趣。感谢分享链接。我甚至不知道 sin 和 cos 是机器指令... - Paul Panzer

那个复杂数组的创建思路很聪明！ - Divakar

@Divakar 谢谢。显然，如果 cumprod 变得非常长，累积误差就会成为一个问题... - Paul Panzer

两个答案都是有效的，可以实现很好的加速。这个答案的性能非常接近于numexpr版本。如果没有numexpr可用，那么应该接受这个答案。 - NameOfTheRose

Answer 3

这不是另一个答案，而是对@Paul Panzer的评论，写成答案是因为我需要发布一些代码。如果有办法在评论中发布格式正确的代码，请告知。

受@Paul Panzer cumprod想法的启发，我想出了以下内容：

an = ones((len(a)-1,len(te)))*2j*pi*te/T
CS = exp(cumsum(an,axis=0))
out = (a[1:].dot(CS.real) - b[1:].dot(CS.imag)) + a[0]

尽管它似乎被正确地向量化并产生了正确的结果，但其性能却很糟糕。它不仅比预期的cumprod慢得多，因为要进行len(a)-1次指数运算，而且比原始的非向量化版本慢50%。这种差劲的表现是什么原因造成的呢？

- NameOfTheRose

---

我的猜测是1）如您所建议，不同数量的exp调用。请注意，差异的顺序为len（a）x len（t），因为an的元素数是其维度的乘积，而不是总和。2）内存布局：在您最初的半向量化代码中，所有变量都可能在内存中连续，并因此有利于受益于优化代码，而交错的.real和.imag数组则无法使用。 - Paul Panzer

@Paul Panzer，谢谢你，我的（天真的）直觉是性能类似于Divakar的第一个版本。你是否仍然对numpy的性能感到惊讶？ - NameOfTheRose

比惊讶更印象深刻。他们每次发布都会使它显着改进。但如果你的意思是我有时会感到“这应该比那个更快”的直觉是错误的吗？是的，实际上经常如此。 - Paul Panzer