Как назначить новые разные значения массиву numpy в одном измерении, используя функцию итерации

Question

У меня есть следующий многомерный массив Numpy:

a = np.arange(16).reshape(2,2,2,2)

Я хочу назначить новые различные значения массива для каждого элемента определенного измерения, например, 4-е измерение Я использовал следующий код:

    for i in range(a.shape[3]):
          if i == 0:
               for t in np.nditer(a[:,:,:,i], op_flags = ['readwrite']):
                    t[...] = t*2
          if i == 1:
               for t in np.nditer(a[:,:,:,i], op_flags = ['readwrite']):
                    t[...] = t*3

print(a)
print(a.shape)

вывод отображается как

[[[[ 0  1]
   [ 4  3]]

   [[ 8  5]
   [12  7]]]


 [[[16  9]
   [20 11]]

  [[24 13]
   [28 15]]]]
[[[[ 0  3]
   [ 4  9]]

  [[ 8 15]
   [12 21]]]


 [[[16 27]
   [20 33]]

  [[24 39]
   [28 45]]]]
(2, 2, 2, 2)
(2, 2, 2, 2)

Что я понимаю, что он выполняет итерацию по массиву, и на первом i он присваивает новые значения, затем на следующем i он назначает новые значения и создает новый массив с обоими новыми значениями, кроме первого массива первого i, поэтому я получил два массива в одной переменной. Я беспокоюсь только о последнем массиве, где все значения были назначены новые значения. Как я мог извлечь только последний массив. Или есть другой код, который проще и экономит время для этой задачи?

Answer 1

Вы можете просто сделать это:

>>> a[:,:,:,0] = a[:,:,:,0]*2
>>> a[:,:,:,1] = a[:,:,:,1]*3
>>> a
array([[[[ 0,  3],
         [ 4,  9]],

        [[ 8, 15],
         [12, 21]]],


       [[[16, 27],
         [20, 33]],

        [[24, 39],
         [28, 45]]]])

Или даже это:

>>> a[:,:,:,0] *= 2
>>> a[:,:,:,1] *= 3