Используйте itertools.product
, чтобы получить список точек, который эквивалентен декартову произведению диапазонов каждой координаты:
>>> list(itertools.product(range(3), range(3), range(3)))
[(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 1, 0), (0, 1, 1), (0, 1, 2), (0, 2, 0),
(0, 2, 1), (0, 2, 2), (1, 0, 0), (1, 0, 1), (1, 0, 2), (1, 1, 0), (1, 1, 1),
(1, 1, 2), (1, 2, 0), (1, 2, 1), (1, 2, 2), (2, 0, 0), (2, 0, 1), (2, 0, 2),
(2, 1, 0), (2, 1, 1), (2, 1, 2), (2, 2, 0), (2, 2, 1), (2, 2, 2)]
Вы также можете объединить все 6 координатв один генератор:
coords = itertools.product(range(Nk), range(Nk), range(Nk), range(Ns), range(Ns), range(Ns))
for kz, ky, kx, rz, ry, rx in coords:
k, r = Nk * (Nk * kz + ky) + kx, Ns * (Ns * rz + ry) + rx
dot_product = kx * rx + ky * ry + kz * rz
k_box[k], R_box[r] = perform_some_calculation(dot_product)
В качестве альтернативы, чтобы избежать вычислений индекса, используйте enumerate
:
for k, (kz, ky, kx) in enumerate(itertools.product(range(Nk), range(Nk), range(Nk))):
for r, (rz, ry, rx) in enumerate(itertools.product(range(Ns), range(Ns), range(Ns))):
dot_product = kx * rx + ky * ry + kz * rz
k_box[k], R_box[r] = perform_some_calculation(dot_product)