How to fold/accumulate a numpy matrix product (dot

With using python library numpy, it is possible to use the function cumprod to evaluate cumulative products, e.g.

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

gives

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

It is indeed possible to apply this function only along one axis.

I would like to do the same with matrices (represented as numpy arrays), e.g. if I have

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]])

and

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

then I would like to have a cumprod-like function which gives

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

(This is a simple example, in reality I have potentially large matrices evaluated over n-dimensional meshgrids, so I seek for the most simple and efficient way to evaluate this thing.)

In e.g. Mathematica I would use

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

so I searched for a fold function, and all I found is an accumulate method on numpy.ufuncs. To be honest, I know that I am probably doomed because an attempt at

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

as mentioned in a numpy mailing list yields the error

Reduction not defined on ufunc with signature

Do you have an idea how to (efficiently) do this kind of calculation ?

Thanks in advance.

As food for thought, here are 3 ways of evaluating the 3 sequential dot products:

With the normal Python reduce (which could also be written as a loop)

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

The einsum equivalent

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

The einsum index expression looks like a sequence of operations, but it is actually evaluated as a 5d product with summation on 3 axes. In the C code this is done with an nditer and strides, but the effect is as follows:

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))

A while back while creating a patch from np.einsum I translated that C code to Python, and also wrote a Cython sum-of-products function(s). This code is on github at

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

einsum_py.py is the Python einsum, with some useful debugging output

sop.pyx is the Cython code, which is compiled to sop.so.

Here's how it could be used for part of your problem. I'm skipping the Sy array since my sop is not coded for complex numbers (but that could be changed).

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

As written sum_product_cy3 cannot take an arbitrary number of ops. Plus the iteration space increases with each op and index. But I can imagine calling it repeatedly, either at the Cython level, or from Python. I think it has potential for being faster than repeat(dot...) for lots of small arrays.

A condensed version of the Cython code is:

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]