Инициализировать массив NumPy на основе его индекса

Я создаю пару многомерных массивов с использованием NumPy и инициализируя их на основе индекса следующим образом:

pos_data = []
# Some typical values
d = 2  # Could also be 3
vol_ext = (1000, 500)  # If d = 3, this will have another entry
ratio = [5.0, 8.0]  # Again, if d = 3, it will have another entry

for i in range(d):
    pos_data.append(np.zeros(vol_ext))

if d == 2:
    for y in range(vol_ext[1]):
        for x in range(vol_ext[0]):
            pos_data[0][x, y] = (x - 1.0) * ratio[0]
            pos_data[1][x, y] = (y - 1.0) * ratio[1]
elif d == 3:
    for z in range(vol_ext[2]):
        for y in range(vol_ext[1]):
            for x in range(vol_ext[0]):
                pos_data[0][x, y, z] = (x - 1.0) * ratio[0]
                pos_data[1][x, y, z] = (y - 1.0) * ratio[1]
                pos_data[2][x, y, z] = (z - 1.0) * ratio[2]

Цикл тоже немного уродливый и медленный. Кроме того, если у меня есть трехмерный объект, тогда у меня должен быть другой вложенный цикл.

Мне было интересно, есть ли Pythonic способ генерировать эти значения, поскольку они основаны только на индексах x, y и z. Я попытался использовать бит combinatorics от itertools, но я не мог заставить его работать.

ОТВЕТЫ

Ответ 1

Это легко с np.meshgrid:

pos_data = np.meshgrid(*(r * (np.arange(s) - 1.0)
                         for s, r in zip(vol_ext, ratio)), indexing='ij')

Ответ 2

Я бы сгенерировал двумерный или трехмерный numpy.meshgrid данных, а затем масштабировал каждую запись по отношению к каждому фрагменту.

Для двумерного случая:

(X, Y) = np.meshgrid(np.arange(vol_ext[1]), np.arange(vol_ext[0]))
pos_data = [(Y - 1) * ratio[0], (X - 1) * ratio[1]]

Для 3D-случая:

(X, Y, Z) = np.meshgrid(np.arange(vol_ext[2]), np.arange(vol_ext[1]), np.arange(vol_ext[0]))
pos_data = [(Z - 1) * ratio[0], (Y - 1) * ratio[1], (X - 1) * ratio[2]]

Пример использования ваших 2D-данных

Ваш код генерируется pos_data. Я создал новый список pos_data2 который хранит эквивалентный список, используя указанное выше решение:

In [40]: vol_ext = (1000, 500)

In [41]: (X, Y) = np.meshgrid(np.arange(vol_ext[1]), np.arange(vol_ext[0]))

In [42]: pos_data2 = [(Y - 1) * ratio[0], (X - 1) * ratio[1]]

In [43]: np.allclose(pos_data[0], pos_data2[0])
Out[43]: True

In [44]: np.allclose(pos_data[1], pos_data2[1])
Out[44]: True

Создание этой адаптивной основываясь на vol_ext

Мы можем комбинировать это со списком, в котором мы можем воспользоваться тем, что вывод numpy.meshgrid является кортежем:

pts = [np.arange(v) for v in reversed(vol_ext)]
pos_data = [(D - 1) * R for (D, R) in zip(reversed(np.meshgrid(*pts)), ratio)]

Первая строка кода генерирует диапазон точек на каждый желаемый размер в виде списка. Затем мы используем понимание списка для вычисления желаемых вычислений на срез путем итерации по каждой желаемой сетке точек в желаемом размере в сочетании с правильным соотношением для применения.

Пример Run

In [49]: pts = [np.arange(v) for v in reversed(vol_ext)]

In [50]:  pos_data2 = [(D - 1) * R for (D, R) in zip(reversed(np.meshgrid(*pts)), ratio)]

In [51]: np.all([np.allclose(p, p2) for (p, p2) in zip(pos_data, pos_data2)])
Out[51]: True

Последняя строка проходит через каждый фрагмент и обеспечивает выравнивание обоих списков.

Ответ 3

Я думаю, что есть несколько вещей, которые следует учитывать:

  • есть ли причина, по которой pos_data должен быть списком?
  • не имеют другой переменной (d), которой вы должны обладать жестким кодом, когда она всегда должна быть длиной некоторого другого кортежа.

Имея это в виду, вы можете решить свою проблему с переменными числами для циклов, используя itertools.product, который в основном является сокращением для вложенных циклов. Позиционные аргументы для product - это диапазоны петель.

Моя реализация:

from itertools import product

vol_ext = (1000, 500)  # If d = 3, this will have another entry
ratio = [5.0, 8.0]  # Again, if d = 3, it will have another entry

pos_data_new = np.zeros((len(ratio), *vol_ext))

# now loop over each dimension in 'vol_ext'. Since 'product' expects
# positional arguments, we have to unpack a tuple of 'range(vol)'s.
for inds in product(*(range(vol) for vol in vol_ext)):
    # inds is now a tuple, and we have to combine it with a slice in 
    # in the first dimension, and use it as an array on the right hand 
    # side to do the computation. 
    pos_data_new[(slice(None),) + inds] = (np.array(inds) - 1) * ratio

Я не думаю, что это будет быстрее, но это, конечно, выглядит лучше.

Обратите внимание, что pos_data_new теперь является массивом, чтобы получить его как список в первом измерении, как в исходном примере, достаточно просто.