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>

<pre><code>In [91]: A=np.array([[1,3],[2,4]])
In [92]: Q=np.array([[-.5,1],[1,.5]])
In [93]: B=np.dot(Q,np.dot(A,Q.T))
In [94]: B
Out[94]: 
array([[ 1.75,  2.75],
       [ 4.  ,  4.5 ]])
</code></pre>

<p>The same calculation with <code>einsum</code>:</p>

<pre><code>In [95]: np.einsum('ij,jk,kl',Q,A,Q)
Out[95]: 
array([[ 1.75,  2.75],
       [ 4.  ,  4.5 ]])
</code></pre>

<p>If I make several copies of <code>A</code> - on a new 1st dimension:</p>

<pre><code>In [96]: AA = np.array([A,A,A])
In [97]: AA.shape
Out[97]: (3, 2, 2)
...
In [99]: BB=np.einsum('ij,pjk,kl-&gt;pil',Q,AA,Q)
In [100]: BB
Out[100]: 
array([[[ 1.75,  2.75],
        [ 4.  ,  4.5 ]],

       [[ 1.75,  2.75],
        [ 4.  ,  4.5 ]],

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

<p><code>BB</code> has a (3,2,2) shape.</p>

<p>The newish <code>matmul</code> (@ operator) lets me do the same thing</p>

<pre><code>In [102]: Q@A@Q.T
Out[102]: 
array([[ 1.75,  2.75],
       [ 4.  ,  4.5 ]])
In [103]: Q@AA@Q.T
Out[103]: 
array([[[ 1.75,  2.75],
        [ 4.  ,  4.5 ]],

       [[ 1.75,  2.75],
        [ 4.  ,  4.5 ]],

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

<p>With <code>einsum</code> it is just as easy to work with the last dimension:</p>

<pre><code>In [104]: AA3=np.stack([A,A,A],-1)    # newish np.stack
In [105]: AA3.shape
Out[105]: (2, 2, 3)
In [106]: np.einsum('ij,jkp,kl-&gt;ilp',Q,AA3,Q)
Out[106]: 
array([[[ 1.75,  1.75,  1.75],
        [ 2.75,  2.75,  2.75]],

       [[ 4.  ,  4.  ,  4.  ],
        [ 4.5 ,  4.5 ,  4.5 ]]])
In [107]: _.shape
Out[107]: (2, 2, 3)
</code></pre>

python numpy matrix-multiplication