CODEWITHSUNDEEP
pythonarraysnumpy
ty.
ty.

Reputation: 231

How to distribute a Numpy array along the diagonal of an array of higher dimension?

I have three two dimensional Numpy arrays x, w, d and want to create a fourth one called a. w and d define only the shape of a with d.shape + w.shape. I want to have x in the entries of a with a zeros elsewhere. Specifically, I want a loop-free version of this code:

a = np.zeros(d.shape + w.shape)

for j in range(d.shape[1]):
    a[:,j,:,j] = x

For example, given:

x = np.array([
    [2, 3],
    [1, 1],
    [8,10],
    [0, 1]
])

w = np.array([
    [ 0, 1, 1],
    [-1,-2, 1]
])

d = np.matmul(x,w)

I want a to be 

array([[[[ 2.,  0.,  0.],
         [ 3.,  0.,  0.]],

        [[ 0.,  2.,  0.],
         [ 0.,  3.,  0.]],

        [[ 0.,  0.,  2.],
         [ 0.,  0.,  3.]]],


       [[[ 1.,  0.,  0.],
         [ 1.,  0.,  0.]],

        [[ 0.,  1.,  0.],
         [ 0.,  1.,  0.]],

        [[ 0.,  0.,  1.],
         [ 0.,  0.,  1.]]],


       [[[ 8.,  0.,  0.],
         [10.,  0.,  0.]],

        [[ 0.,  8.,  0.],
         [ 0., 10.,  0.]],

        [[ 0.,  0.,  8.],
         [ 0.,  0., 10.]]],


       [[[ 0.,  0.,  0.],
         [ 1.,  0.,  0.]],

        [[ 0.,  0.,  0.],
         [ 0.,  1.,  0.]],

        [[ 0.,  0.,  0.],
         [ 0.,  0.,  1.]]]])

Upvotes: 1

Views: 194

Answers (2)

amzon-ex
amzon-ex

Reputation: 1744

I happen to have an even simpler solution: a = np.tensordot(x, np.identity(3), axes = 0).swapaxes(1,2)
The size of the identity matrix will be decided by the number of times you wish to repeat the elements of x.

Upvotes: 0

Joe
Joe

Reputation: 7131

This answer inspired the following solution:

# shape a: (4, 3, 2, 3)
# shape x: (4, 2)

a = np.zeros(d.shape + w.shape)
a[:, np.arange(a.shape[1]), :, np.arange(a.shape[3])] = x

It uses Numpy's broadcasting (see here or here) im combination with Advanced Indexing to enlarge x to fit the slicing.

Upvotes: 1

Related Questions