I have a numpy array of shape (5, 4, 3)
and another numpy array of shape (4,)
and what I want to do is expand the last dimension of the first array
(5, 4, 3) -> (5, 4, 4)
and then broadcast the other array with shape (4,)
such that it fills up the new array cells respectively.
Example:
JavaScript
x
26
26
1
np.ones((5,4,3))
2
array([[[1., 1., 1.],
3
[1., 1., 1.],
4
[1., 1., 1.],
5
[1., 1., 1.]],
6
7
[[1., 1., 1.],
8
[1., 1., 1.],
9
[1., 1., 1.],
10
[1., 1., 1.]],
11
12
[[1., 1., 1.],
13
[1., 1., 1.],
14
[1., 1., 1.],
15
[1., 1., 1.]],
16
17
[[1., 1., 1.],
18
[1., 1., 1.],
19
[1., 1., 1.],
20
[1., 1., 1.]],
21
22
[[1., 1., 1.],
23
[1., 1., 1.],
24
[1., 1., 1.],
25
[1., 1., 1.]]])
26
becomes
JavaScript
1
25
25
1
array([[[1., 1., 1., 0.],
2
[1., 1., 1., 0.],
3
[1., 1., 1., 0.],
4
[1., 1., 1., 0.]],
5
6
[[1., 1., 1., 0.],
7
[1., 1., 1., 0.],
8
[1., 1., 1., 0.],
9
[1., 1., 1., 0.]],
10
11
[[1., 1., 1., 0.],
12
[1., 1., 1., 0.],
13
[1., 1., 1., 0.],
14
[1., 1., 1., 0.]],
15
16
[[1., 1., 1., 0.],
17
[1., 1., 1., 0.],
18
[1., 1., 1., 0.],
19
[1., 1., 1., 0.]],
20
21
[[1., 1., 1., 0.],
22
[1., 1., 1., 0.],
23
[1., 1., 1., 0.],
24
[1., 1., 1., 0.]]])
25
And then I have another array
JavaScript
1
2
1
array([2., 3., 4., 5.])
2
which I somehow broadcast with the first one to fill the zeros:
JavaScript
1
25
25
1
array([[[1., 1., 1., 2.],
2
[1., 1., 1., 3.],
3
[1., 1., 1., 4.],
4
[1., 1., 1., 5.]],
5
6
[[1., 1., 1., 2.],
7
[1., 1., 1., 3.],
8
[1., 1., 1., 4.],
9
[1., 1., 1., 5.]],
10
11
[[1., 1., 1., 2.],
12
[1., 1., 1., 3.],
13
[1., 1., 1., 4.],
14
[1., 1., 1., 5.]],
15
16
[[1., 1., 1., 2.],
17
[1., 1., 1., 3.],
18
[1., 1., 1., 4.],
19
[1., 1., 1., 5.]],
20
21
[[1., 1., 1., 2.],
22
[1., 1., 1., 3.],
23
[1., 1., 1., 4.],
24
[1., 1., 1., 5.]]])
25
How can I accomplish this?
Advertisement
Answer
You can use numpy.c_
and numpy.tile
:
JavaScript
1
5
1
A = np.ones((5,4,3), dtype='int')
2
B = np.array([2, 3, 4, 5])
3
4
np.c_[A, np.tile(B[:,None], (A.shape[0], 1, 1))]
5
output:
JavaScript
1
25
25
1
array([[[1, 1, 1, 2],
2
[1, 1, 1, 3],
3
[1, 1, 1, 4],
4
[1, 1, 1, 5]],
5
6
[[1, 1, 1, 2],
7
[1, 1, 1, 3],
8
[1, 1, 1, 4],
9
[1, 1, 1, 5]],
10
11
[[1, 1, 1, 2],
12
[1, 1, 1, 3],
13
[1, 1, 1, 4],
14
[1, 1, 1, 5]],
15
16
[[1, 1, 1, 2],
17
[1, 1, 1, 3],
18
[1, 1, 1, 4],
19
[1, 1, 1, 5]],
20
21
[[1, 1, 1, 2],
22
[1, 1, 1, 3],
23
[1, 1, 1, 4],
24
[1, 1, 1, 5]]])
25
How it works:
JavaScript
1
34
34
1
# reshape B to add one dimension
2
>>> B[:, None]
3
array([[2],
4
[3],
5
[4],
6
[5]])
7
8
# tile to match A's first dimension
9
>>> np.tile(B[:,None], (A.shape[0], 1, 1))
10
array([[[2],
11
[3],
12
[4],
13
[5]],
14
15
[[2],
16
[3],
17
[4],
18
[5]],
19
20
[[2],
21
[3],
22
[4],
23
[5]],
24
25
[[2],
26
[3],
27
[4],
28
[5]],
29
30
[[2],
31
[3],
32
[4],
33
[5]]])
34