python - YALV(又一个循环矢量化)

Question

我可以按照这些思路编写一些代码，以针对不同的频率计算符合给定初始条件的调和函数。

from numpy import *
...
for i in range(n_freqs):
    # A, B so that X(t) = A cos(w t) + B sin(w t)
    # and X(t0) = x, dX/dt(t0) = v
    w = ws[i] # use a frequency
    solver = array(((+cos(w*t0), -sin(w*t0)),
                    (+sin(w*t0), +cos(w*t0)))) 
    AB = solver @ array((x[i], v[i]/w)) # and store somewhere the result

但我想写点更像的东西

Solver = array(((+cos(ws*t0), -sin(ws*t0)),
                (+sin(ws*t0), +cos(ws*t0))))
AB = Solver @ vstack((x,v/ws)

我(不是)我们

from numpy import *

ws = array((1., 2., 3., 4.))
x = array((3., 6., 2., 1.))
v = x
t0 = 10.0
Solver = array(((+cos(ws*t0), -sin(ws*t0)),
                (+sin(ws*t0), +cos(ws*t0))))
AB = Solver @ vstack((x,v/ws)

这给了我以下痕迹

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: shapes (2,2,4) and (2,4) not aligned: 4 (dim 2) != 2 (dim 0)

同时

(Solver @ vstack((x,v/ws).T).shape # -> (2, 2, 2)

当我想要一个(2, 4)(或(4,2)，我不太挑剔)形状时...

是否可以编写一个无循环表达式来同时计算不同三角函数的系数？

Answer 1

w 是 (n,)，所以 cos(w*t0) 也有那个形状。 Solver 布局为 (2,2)，但具有 (n,) 个元素，因此它是 (2,2,n)。您正在用 (2,n) 进行“打点”。但是在哪个维度上，n 还是 2 之一？

solver(ws[i]) @ array((x[i], v[i]/ws[i]))

表示您希望最后一个维度“随心所欲”，并点在 Solver 的最后 2 个维度上。在 Einsum 表示法中:

np.einsum('ijk,jk->ik', Solver, arr)

In [99]: Solver = np.array(((np.cos(wst),-np.sin(wst)),(np.sin(wst),np.cos(wst))))

In [101]: b = np.vstack((x,v/ws))
In [102]: b.shape
Out[102]: (2, 4)
In [103]: for i in range(4):
     ...:     print(Solver[:,:,i]@b[:,i])
     ...:     
[-0.88515125 -4.14927792]
[-0.29034338  6.70191769]
[ 0.96719065 -1.87322895]
[-0.85321635  0.57837865]
In [104]: np.einsum('ijk,jk->ik',Solver,b)
Out[104]: 
array([[-0.88515125, -0.29034338,  0.96719065, -0.85321635],
       [-4.14927792,  6.70191769, -1.87322895,  0.57837865]])

对于 @ 来说，这不是一个简单的案例，因为它假定数组在第一个维度上堆叠。例如 Solver 应该是 (n,2,2)，而 b (n,2,1)`

In [106]: Solver.transpose(2,0,1)@(b.T[...,None])
Out[106]: 
array([[[-0.88515125],
        [-4.14927792]],

       [[-0.29034338],
        [ 6.70191769]],

       [[ 0.96719065],
        [-1.87322895]],

       [[-0.85321635],
        [ 0.57837865]]])
In [107]: _.shape   # need to squeeze out the last dim
Out[107]: (4, 2, 1)

---

如果我的推论是正确的，即您的“点”维度是大小 2，那么在该维度上迭代将几乎一样快，尤其是对于大 n:

res = np.zeros((2,n))
for i in range(2):
    res += Solver[:,i,:] * b[i,:]