值得深思的是,这里有 3 种评估 3 个连续点积的方法:
使用普通的 Python reduce(也可以写成循环)
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 代码中这是用 annditer
和 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]