如何折叠/累积一个 numpy 矩阵乘积（点）？

Question

通过使用 python 库 numpy，可以使用函数cumprod来评估累积乘积，例如

a = np.array([1,2,3,4,2])
np.cumprod(a)

给

array([ 1,  2,  6, 24, 48])

确实可以仅沿一个轴应用此功能。

我想对矩阵（表示为 numpy 数组）做同样的事情，例如，如果我有

S0 = np.array([[1, 0], [0, 1]])
Sx = np.array([[0, 1], [1, 0]])
Sy = np.array([[0, -1j], [1j, 0]])
Sz = np.array([[1, 0], [0, -1]])

和

b = np.array([S0, Sx, Sy, Sz])

那么我想要一个类似cumprod的函数，它给出

np.array([S0, S0.dot(Sx), S0.dot(Sx).dot(Sy), S0.dot(Sx).dot(Sy).dot(Sz)])

（这是一个简单的例子，在现实中我有潜在的大矩阵评估了n维meshgrids，所以我寻求最简单和评估这个东西有效的方式。）

在例如 Mathematica 中，我会使用

FoldList[Dot, IdentityMatrix[2], {S0, Sx, Sy, Sz}]

所以我搜索了一个 fold 函数，我找到的只是saccumulate上numpy.ufunc的一个方法。老实说，我知道我可能注定要失败，因为我试图

np.core.umath_tests.matrix_multiply.accumulate(np.array([pauli_0, pauli_x, pauli_y, pauli_z]))

正如在numpy 邮件列表中提到的那样产生错误

Reduction not defined on ufunc with signature

您知道如何（有效地）进行这种计算吗？

提前致谢。

Answer 1

作为深思熟虑，以下是评估 3 个连续点积的 3 种方法：

使用普通的Pythonreduce（也可以写成循环）

In [118]: reduce(np.dot,[S0,Sx,Sy,Sz])
array([[ 0.+1.j,  0.+0.j],
       [ 0.+0.j,  0.+1.j]])

相当于einsum​

In [119]: np.einsum('ij,jk,kl,lm',S0,Sx,Sy,Sz)

索引einsum表达式看起来像是一个操作序列，但实际上它被评估为在 3 个轴上求和的 5d 乘积。在C代码中这是通过nditerand strides来完成的，但是效果如下：

In [120]: np.sum(S0[:,:,None,None,None] * Sx[None,:,:,None,None] *
    Sy[None,None,:,:,None] * Sz[None,None,None,:,:],(1,2,3))

In [127]: np.prod([S0[:,:,None,None,None], Sx[None,:,:,None,None],
    Sy[None,None,:,:,None], Sz[None,None,None,:,:]]).sum((1,2,3))

不久前，在创建补丁时，np.einsum我将该C代码翻译为Python，并且还编写了一个Cython乘积和函数。该代码位于 github 上

https://github.com/hpaulj/numpy-einsum

einsum_py.py是 Python einsum，带有一些有用的调试输出

sop.pyx是 Cython 代码，它被编译为sop.so.

以下是如何使用它来解决您的部分问题。我正在跳过该Sy数组，因为我的sop数组没有针对复数进行编码（但这可以更改）。

import numpy as np
import sop
import einsum_py    

S0 = np.array([[1., 0], [0, 1]])
Sx = np.array([[0., 1], [1, 0]])
Sz = np.array([[1., 0], [0, -1]])

print np.einsum('ij,jk,kl', S0, Sx, Sz)
# [[ 0. -1.] [ 1.  0.]]
# same thing, but with parsing information
einsum_py.myeinsum('ij,jk,kl', S0, Sx, Sz, debug=True)
"""
{'max_label': 108, 'min_label': 105, 'nop': 3, 
 'shapes': [(2, 2), (2, 2), (2, 2)], 
 'strides': [(16, 8), (16, 8), (16, 8)], 
 'ndim_broadcast': 0, 'ndims': [2, 2, 2], 'num_labels': 4,
 ....
 op_axes [[0, -1, 1, -1], [-1, -1, 0, 1], [-1, 1, -1, 0], [0, 1, -1, -1]]
"""    

# take op_axes (for np.nditer) from this debug output
op_axes = [[0, -1, 1, -1], [-1, -1, 0, 1], [-1, 1, -1, 0], [0, 1, -1, -1]]
w = sop.sum_product_cy3([S0,Sx,Sz], op_axes)
print w

正如所写的，sum_product_cy3不能采用任意数量的ops。另外，迭代空间随着每个操作和索引的增加而增加。但我可以想象在 Cython 级别或从 Python 重复调用它。我认为它有可能比repeat(dot...)许多小型阵列更快。

Cython 代码的精简版本是：

def sum_product_cy3(ops, op_axes, order='K'):
    #(arr, axis=None, out=None):
    cdef np.ndarray[double] x, y, z, w
    cdef int size, nop
    nop = len(ops)
    ops.append(None)
    flags = ['reduce_ok','buffered', 'external_loop'...]
    op_flags = [['readonly']]*nop + [['allocate','readwrite']]

    it = np.nditer(ops, flags, op_flags, op_axes=op_axes, order=order)
    it.operands[nop][...] = 0
    it.reset()
    for x, y, z, w in it:
        for i in range(x.shape[0]):
           w[i] = w[i] + x[i] * y[i] * z[i]
    return it.operands[nop]