Блог

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

 ix_array = [0, 2, 5]
iy_array = [0, 1, 2]
iz_array = [0, 0, 4]
w_array = [0.1, 2.1, 3.9]
 

и я хочу использовать первые три массива в качестве индексов для привязки w_array к 3D-сетке. Это можно сделать с помощью чего-то вроде

 for i, (ix, iy, iz) in enumerate((ix_array, iy_array, iz_array)):
   A[ix][iy][iz]  = w_array[i]
 

Есть ли что-нибудь, чтобы ускорить это?

Не совсем ясно, как вы хотите распределять свои w-данные. Ваш пример кода работает только для списков длиной ровно три, поэтому сгенерированный результат может быть не тем, чего вы на самом деле хотите достичь (или я вообще не понимаю вашу задачу). Исходя из ваших имен переменных, я предполагаю, что это что-то в этом роде:

 import numpy as np

ix_array = [0, 1, 1, 3]
iy_array = [0, 1, 1, 0]
iz_array = [0, 0, 2, 3]
w_array = [0.1, 2.1, 3.9, 7.1]

A = np.zeros((4, 5, 6))

A[tuple(zip([ix_array, iy_array, iz_array]))]  = w_array

print(A)
 

Пример вывода:

 [[[0.1 0.  0.  0.  0.  0. ]
  [0.  0.  0.  0.  0.  0. ]
  [0.  0.  0.  0.  0.  0. ]
  [0.  0.  0.  0.  0.  0. ]
  [0.  0.  0.  0.  0.  0. ]]

 [[0.  0.  0.  0.  0.  0. ]
  [2.1 0.  3.9 0.  0.  0. ]
  [0.  0.  0.  0.  0.  0. ]
  [0.  0.  0.  0.  0.  0. ]
  [0.  0.  0.  0.  0.  0. ]]

 [[0.  0.  0.  0.  0.  0. ]
  [0.  0.  0.  0.  0.  0. ]
  [0.  0.  0.  0.  0.  0. ]
  [0.  0.  0.  0.  0.  0. ]
  [0.  0.  0.  0.  0.  0. ]]

 [[0.  0.  0.  7.1 0.  0. ]
  [0.  0.  0.  0.  0.  0. ]
  [0.  0.  0.  0.  0.  0. ]
  [0.  0.  0.  0.  0.  0. ]
  [0.  0.  0.  0.  0.  0. ]]]
 

numpy.ravel_multi_index может использоваться для создания массива плоских индексов и numpy.flat задания массива A по этим индексам:

 import numpy as np

A = np.zeros((6, 3, 5))
ix_array = [0, 2, 5]
iy_array = [0, 1, 2]
iz_array = [0, 0, 4]
w_array = [0.1, 2.1, 3.9]

arr = np.array([ix_array, iy_array, iz_array])
indices = np.ravel_multi_index(arr, A.shape)
A.flat[indices]  = w_array
print(A)
 

вывод:

 array([[[0.1, 0. , 0. , 0. , 0. ],
        [0. , 0. , 0. , 0. , 0. ],
        [0. , 0. , 0. , 0. , 0. ]],
       [[0. , 0. , 0. , 0. , 0. ],
        [0. , 0. , 0. , 0. , 0. ],
        [0. , 0. , 0. , 0. , 0. ]],
       [[0. , 0. , 0. , 0. , 0. ],
        [2.1, 0. , 0. , 0. , 0. ],
        [0. , 0. , 0. , 0. , 0. ]],
       [[0. , 0. , 0. , 0. , 0. ],
        [0. , 0. , 0. , 0. , 0. ],
        [0. , 0. , 0. , 0. , 0. ]],
       [[0. , 0. , 0. , 0. , 0. ],
        [0. , 0. , 0. , 0. , 0. ],
        [0. , 0. , 0. , 0. , 0. ]],
       [[0. , 0. , 0. , 0. , 0. ],
        [0. , 0. , 0. , 0. , 0. ],
        [0. , 0. , 0. , 0. , 3.9]]])