Search code examples
pythonarraysnumpynumpy-ndarrayarray-broadcasting

In NumPy arrays, is there a syntax to set a value in the last dimension based on the first dimension?


I've Googled and asked ChatGPT and looked through NumPy docs and can't find any way to do this, so thought I'd ask here.

Suppose I have a 4-dimensional array -- in this case, of shape (3, 2, 2, 2):

a = np.array([
  [[[0, 0], [0, 0]], 
    [[0, 0], [0, 0]]],
  [[[0, 0], [0, 0]], 
    [[0, 0], [0, 0]]],
  [[[0, 0], [0, 0]], 
    [[0, 0], [0, 0]]],
  ])

and I want to set the last (second) element of the last dimension to a different value according to each row of the first dimension. In my example I have 3 rows in the first dimension, so let's suppose I wanted to apply the values [1, 2, 3] to result in:

[
  [[[0, 1], [0, 1]], 
    [[0, 1], [0, 1]]],
  [[[0, 2], [0, 2]], 
    [[0, 2], [0, 2]]],
  [[[0, 3], [0, 3]], 
    [[0, 3], [0, 3]]],
]

The closest syntax I've been able to think of would be:

a[:, ..., 1] = [1, 2, 3]

But it produces an error (ValueError: could not broadcast input array from shape (3,) into shape (3,2,2)). No error is produced if I try:

a[:, ..., 1] = [1, 2]

but it produces a different result which isn't what I want:

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

Is there any way to elegantly and compactly do what I want?

For now I've written a loop to cycle over every row of the first dimension and then set the values per-row, but I wonder if there's a more powerful way to do this in a single line.


Solution

  • You need to provide a 3-1-1 array on the right hand side of the assignment for numpy to be able to broadcast it to your slice's 3-2-2 shape. If I got your request right, following expression would work

    a[:, ..., 1] = np.array([1, 2, 3])[:, None, None]