Jack Jack - 2 months ago 12
Python Question

Intuitive way to use 3D numpy arrays

Lets take a simplistic example where I have the data array

A = np.asarray([[1,3], [2,4]])


And this data is to be transformed into another form following a simple transformation:

Q = np.asarray([[-0.5,1], [1,0.5]])
B = np.dot(Q,np.dot(A,Q.T))
print B


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
A
copied for 3 time steps. We can represent this data as a 3d array with dimensions
(2,2,N)
where
N =3
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:

# 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


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

Edit:

Since
dot(a, b)[i,j,k,m] = sum(a[i,j,:] * b[k,:,m])
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...

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

The same calculation with einsum:

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

If I make several copies of A - on a new 1st dimension:

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

BB has a (3,2,2) shape.

The newish matmul (@ operator) lets me do the same thing

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

With einsum it is just as easy to work with the last dimension:

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