Intuitive way to use 3D numpy arrays
<p>Lets take a simplistic example where I have the data array</p>

<pre><code>A = np.asarray([[1,3], [2,4]])
</code></pre>

<p>And this data is to be transformed into another form following a simple transformation: </p>

<pre><code>Q = np.asarray([[-0.5,1], [1,0.5]])
B = np.dot(Q,np.dot(A,Q.T))
print B
</code></pre>

<p>Now assume that I have a set of data that takes the form of a 2d array for several time steps. For simplicity again assume that this data is just <code>A</code> copied for 3 time steps. We can represent this data as a 3d array with dimensions <code>(2,2,N)</code> where <code>N =3</code> in this case. The third dimension then represents the time index of the data. Now it would be natural to demand to have a simple way of transforming the data as above but for each time step, by an intuitive multiplication of 3d arrays, however I have only been able to make the following work which is non-intuitive:</p>

<pre><code># Create the 3d data array
AA = np.tile(A,(3,1,1)) # shape (3,2,2)
BB = np.dot(Q,np.dot(AA,Q.T))

print np.all( BB[:,0,:] == B ) # Returns true 
</code></pre>

<p>So with this method I don't have to recast the <code>Q</code> array to make it work, but now the second dimension acts as the "time" index which is a bit counter intuitive since in <code>AA</code> it was the first dimension that denoted the time... Ideally I would like a solution in which both <code>AA</code> and <code>BB</code> have the time index in the third dimension!</p>

<p>Edit:</p>

<p>Since <code>dot(a, b)[i,j,k,m] = sum(a[i,j,:] * b[k,:,m])</code> from docs I am wondering if what I am trying to achieve is not possible? It seems strange as this should be a relatively common thing one may desire...</p>

<p>I'm not sure I agree with your definition of "intuitive" - to me it seems more natural to represent the time index by the first dimension of the array. Since numpy arrays are <a href="https://en.wikipedia.org/wiki/Row-major_order" rel="nofollow">row-major</a> by default, this ordering of the dimensions will give you better <a href="https://en.wikipedia.org/wiki/Locality_of_reference" rel="nofollow">locality of reference</a> for each 2x2 submatrix, since all of the elements will be in adjacent memory addresses.</p>

<p>Nonetheless, it's possible to adapt your example to work for a <code>(2, 2, N)</code> array by using <a href="http://docs.scipy.org/doc/numpy/reference/generated/numpy.matmul.html" rel="nofollow"><code>np.matmul</code></a> and transposing each of the intermediate arrays:</p>

<pre><code>CC = np.repeat(A[..., None], 3, -1)   # shape (2, 2, 3)
DD = np.matmul(Q.T, np.matmul(CC.T, Q)).T

print(DD.shape)
# (2, 2, 3)

print(repr(DD))
# array([[[ 1.75,  1.75,  1.75],
#         [ 2.75,  2.75,  2.75]],

#        [[ 4.  ,  4.  ,  4.  ],
#         [ 4.5 ,  4.5 ,  4.5 ]]])
</code></pre>

<p>In Python 3.5+ you can make this even more compact by using the <a href="http://legacy.python.org/dev/peps/pep-0465/" rel="nofollow"><code>@</code> operator</a> as shorthand for <code>np.matmul</code>:</p>

<pre><code>DD = (Q.T @ (CC.T @ Q)).T
</code></pre>

python numpy matrix-multiplication