Numpy from an array, for each element create a matrix N*M with all values set as the element without for loops

Viewed 291

I have a numpy array like np.array([1, 2, 3])

Without using for loops but just use numpy or pytorch methods, i want a matrix with dimension len(array) * N * M composed by matrices N*M that the first matrix is composed by all ones, the second only by two values and the third only by 3 values.

For istance

N = 4 M = 3
[[[1,1,1,1],[1,1,1,1],[1,1,1,1]], 
  [[2,2,2,2],[2,2,2,2],[2,2,2,2]], 
  [[3,3,3,3],[3,3,3,3],[3,3,3,3]]]

I tried different methods to achive this matrix like unsqueeze and repeat but i wasn't able to find a solution, any suggestion?

2 Answers

Here you go:

np.ones((len(a),M,N)) * a[:,None,None]

Or without multiplication:

np.full((len(a),M,N), a[:,None,None])

Output:

array([[[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]],

       [[2., 2., 2., 2.],
        [2., 2., 2., 2.],
        [2., 2., 2., 2.]],

       [[3., 3., 3., 3.],
        [3., 3., 3., 3.],
        [3., 3., 3., 3.]]])

Note the shape here and given in your expected output is len(a) * M * N, not len(a) * N * M. Swap M,N in np.ones if needed.

In PyTorch you can use expand

a_t = torch.from_numpy(a)

a_t[:,None,None].expand(len(a_t), M, N)
tensor([[[1, 1, 1, 1],
         [1, 1, 1, 1],
         [1, 1, 1, 1]],

        [[2, 2, 2, 2],
         [2, 2, 2, 2],
         [2, 2, 2, 2]],

        [[3, 3, 3, 3],
         [3, 3, 3, 3],
         [3, 3, 3, 3]]], dtype=torch.int32)

Another way would be using repeat

a_t[:,None,None].repeat(1, M, N)

this similar to numpy tile - np.tile(a[:,None,None], (1,M,N))

Related