numpy.einsum new feature request: repeated output subscripts as diagonal

[PierreAndreNoel]

I think that the following new feature would make numpy.einsum even more powerful/useful/awesome than it already is. Moreover, the change should not interfere with existing code, it would preserve the "minimalistic" spirit of numpy.einsum, and the new functionality would integrate in a seamless/intuitive manner for the users.

In short, the new feature would allow for repeated subscripts to appear in the "output" part of the subscripts parameter (i.e., on the right-hand side of ->). The corresponding dimensions in the resulting ndarray would only be filled along their diagonal, leaving the off diagonal entries to the default value for this dtype (typically zero). Note that the current behavior is to raise an exception when repeated output subscripts are being used.

This is simplest to describe with an example involving the dual behavior of numpy.diag.

# Extracting the diagonal of a 2-D array.
A = arange(16).reshape(4,4)
print(diag(A)) # Output: [ 0 5 10 15 ]
print(einsum('ii->i', A)) # Same as previous line (current behavior).

# Constructing a diagonal 2-D array.
v = arange(4)
print(diag(v)) # Output: [[0 0 0 0] [0 1 0 0] [0 0 2 0] [0 0 0 3]]
print(einsum('i->ii', v)) # New behavior would be same as previous line.
# The current behavior of the previous line is to raise an exception.

By opposition to numpy.diag, the approach generalizes to higher dimensions: einsum('iii->i', A) extracts the diagonal of a 3-D array, and einsum('i->iii', v) would build a diagonal 3-D array.

The proposed behavior really starts to shine in more intricate cases.

# Dummy values, these should be probabilities to make sense below.
P_w_ab = arange(24).reshape(3,2,4)
P_y_wxab = arange(144).reshape(3,3,2,2,4)

# With the proposed behavior, the following two lines should be equivalent.
P_xyz_ab = einsum('wab,xa,ywxab,zy->xyzab', P_w_ab, eye(2), P_y_wxab, eye(3))
also_P_xyz_ab = einsum('wab,ywaab->ayyab', P_w_ab, P_y_wxab)

If this is not convincing enough, replace eye(2) by eye(P_w_ab.shape[1]) and replace eye(3) by eye(P_y_wxab.shape[0]), then imagine more dimensions and repeated indices... The new notation would allow for crisper codes and reduce the opportunities for dumb mistakes.

For those who wonder, the above computation amounts to $P(X=x,Y=y,Z=z|A=a,B=b) = \sum_w P(W=w|A=a,B=b) P(X=x|A=a) P(Y=y|W=w,X=x,A=a,B=b) P(Z=z|Y=y)$ with $P(X=x|A=a)=\delta_{xa}$ and $P(Z=z|Y=y)=\delta_{zy}$ (using LaTeX notation, and $\delta_{ij}$ is Kronecker's delta).

[jaimefrio]

The best way of getting a new feature into numpy is putting it in yourself. And what you describe would be much easier to implement in a Python wrapper: string manipulations are not C's strongest point. It also doesn't seem like a good idea, performance wise, to add large unused chunks of memory into the output array while you are doing the computations.

[njsmith]

Well, the first question is whether it's a good idea, not how to implement
it :-). (And einsum already has code to manipulate these kinds of strings
so it might or might not be easier to do at the python level. Doing it in
one pass in C would certainly be more memory efficient though.)

To the original poster: can you send this proposal to the mailing list?
That's generally where we discuss "is this even a good idea" type questions.
On 14 Aug 2014 17:47, "Jaime" notifications@github.com wrote:

The best way of getting a new feature into numpy is putting it in
yourself. And what you describe would be much easier to implement in a
Python wrapper: string manipulations are not C's strongest point. It also
doesn't seem like a good idea, performance wise, to add large unused chunks
of memory into the output array while you are doing the computations.

—
Reply to this email directly or view it on GitHub
#4965 (comment).

[PierreAndreNoel]

@jaimefrio

The best way of getting a new feature into numpy is putting it in yourself.

Thanks. I may consider that at some point, but not in the near future.

It also doesn't seem like a good idea, performance wise, to add large unused chunks of memory into the output array while you are doing the computations.

I'm not sure that I understand what you mean by "while you are doing the computations". In the above example, the purpose of the code is to obtain P_xyz_ab. The zeros are not "unused" memory, they are part of the desired outcome. Now, it would be much better if there were support for sparce ndarrays with more than 2D, but that is a different story...

@njsmith

can you send this proposal to the mailing list? That's generally where we discuss "is this even a good idea" type questions.

Thanks, and sorry for my ignorance of the etiquette. Will do so right now.

[jaimefrio]

Lets say that you have an axis a of length 100, that you want to have three times in your output. If you do the math, that means that only 100 in 1,000,000, or 1 in 10,000 entries is going to be different from 0. Even in the simplest case of an axis of length 2, repeated only twice, half of the entries are going to be 0. You do not want to step over every item of the output array, check if it is in a diagonal, and put a zero in there if it is not: you memset everything to 0 at the beginning, and then fill the right positions in. If I am not mistaken, it turns out that, for diagonals, you can access them without any checking with some striding magic. If you, e.g. wanted to do an operation like `ijk,ik->ijiji', which is a madman's storage for the products of a stack of matrices with a stack of vectors, you could simply do the following from your Python interpreter:

i,j,k = 10, 20, 30
a = np.random.rand(i, j, k)
b = np.random.rand(j,k)
out = np.zeros((i, j, i, j, i))
from numpy.lib.stride_tricks import as_strided
out_view = as_strided(out, shape=(i, j),
                       strides=(out.strides[0]+out.strides[2]+out.strides[4],
                                out.strides[1]+out.strides[3]))
np.einsum('ijk,jk->ij', a, b, out=out_view)

You can probably get a general case function in about 50 lines of Python code. I have just looked at einsum.c.src and there are probably something like 500 lines of painful to read code, spread in three functions, to do the subscript parsing... It is probably better, and infinitely simpler, to write a new diagonals method that uses einsum's subscript notation, perhaps inside numpy.lib.stride_tricks, to build or extract diagonals from arbitrary dimensional arrays. So your example above would turn into:

from numpy.lib.stride_tricks import diagonals
also_P_xyz_ab = diagonals('ayb->ayyab', einsum('wab,ywaab->ayb', P_w_ab, P_y_wxab))

[PierreAndreNoel]

Umh. Consider the following hack: I implement a diagonals such as the one you propose. I then edit einsum at the place where it would usually raise an exception to do the following:

1. Call einsum without repeated indices in the outputs (such as your 'wab,ywaab->ayb').
2. Call diagonals to add in the missing dimensions (such as your 'ayb->ayyab').

This way, I can't break nor slow down the current einsum, unless the user's call would have raised an error anyway. Thoughts?

(BTW, I sent this to the mailing list: http://mail.scipy.org/pipermail/numpy-discussion/2014-August/070969.html )