Как периодически суммировать элементы строк решетки

Question

Предположим, у меня есть решетка

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

Я хотел бы создать функцию func(lattice, start, end), которая принимает 3 входа, где начало и конец - это индекс строк, для которых функция суммирует элементы. Например, для func(a,1,3) он будет суммировать все элементы этих строк, так что func(a,1,3) = 2+2+2+2+3+3+3+3+4+4+4+4 = 36.

Теперь я знаю, что это можно легко сделать с помощью нарезки и np.sum() независимо от того, что. Но главное, что я хочу сделать func, это также иметь возможность обернуться вокруг. А именно func(a,2,4) должен вернуть 3+3+3+3+4+4+4+4+1+1+1+1. Еще пара примеров:

func(a,3,4) = 4+4+4+4+1+1+1+1
func(a,3,5) = 4+4+4+4+1+1+1+1+2+2+2+2
func(a,0,1) = 1+1+1+1+2+2+2+2

В моей ситуации я никогда не доберусь до точки, где все снова будет суммироваться, т.е.

func(a,3,6) = sum of all elements

Обновление : для моего алгоритма

for i in range(MC_STEPS_NODE):
    sweep(lattice, prob, start_index_master, end_index_master,
                      rows_for_master) 
    # calculate the energy
    Ene = subhamiltonian(lattice, start_index_master, end_index_master)  
    # calculate the magnetisation
    Mag = mag(lattice, start_index_master, end_index_master)
    E1 += Ene
    M1 += Mag
    E2 += Ene*Ene
    M2 += Mag*Mag

        if i % sites_for_master == 0:
            comm.Send([lattice[start_index_master:start_index_master+1], L, MPI.INT],
                              dest=(rank-1)%size, tag=4)
            comm.Recv([lattice[end_index_master:end_index_master+1], L, MPI.INT],
                              source = (rank+1)%size, tag=4)
            start_index_master = (start_index_master + 1)
            end_index_master = (end_index_master + 1)

            if start_index_master > 100:
                start_index_master = start_index_master % L

            if end_index_master > 100:
                end_index_master = end_index_master % L

Функция, которую я хочу, - это функция mag(), которая вычисляет намагниченность подрешетки, которая является просто суммой всех ее элементов. Представьте себе решетку LxL, разделенную на две подрешетки, одна принадлежит мастеру, а другая рабочему. Каждый sweep охватывает соответствующую подрешетку lattice, где start_index_master и end_index_master определяют начальный и конечный ряд подрешетки. Для каждого i%sites_for_master = 0 индексы уменьшаются, добавляя 1, в конечном итоге изменяя на 100, чтобы предотвратить переполнение памяти в mpi4py. Таким образом, вы можете представить, находится ли подрешетка в центре основной решетки, тогда start_index_master < end_index_master. В конечном итоге подрешетка будет продолжать двигаться вниз до точки, где start_index_master < end_index_master, где end_index_master > L, поэтому в этом случае, если start_index_master = 10 для решетки L=10, самый нижний ряд подрешетки является первым рядом ([0]) главной решетки.

Энергетическая функция :

def subhamiltonian(lattice: np.ndarray, col_len_start: int,
                   col_len_end: int) -> float:

    energy = 0
    for i in range(col_len_start, col_len_end+1):
        for j in range(len(lattice)):
            spin = lattice[i%L, j]
            nb_sum = lattice[(i%L+1) % L, j] + lattice[i%L, (j+1) % L] + \
                     lattice[(i%L-1) % L, j] + lattice[i%L, (j-1) % L]
            energy += -nb_sum*spin

    return energy/4.

Это моя функция для вычисления энергии подрешетки.

Answer 1

Вы можете использовать np.arange для создания индексов для суммирования.

>>> def func(lattice, start, end):
...     rows = lattice.shape[0]
...     return lattice[np.arange(start, end+1) % rows].sum()
... 
>>> func(a,3,4) 
20
>>> func(a,3,5)
28
>>> func(a,0,1)
12

Answer 2

Вы можете проверить, оборачивается ли индекс останова, и добавляет ли он сумму от начала массива к результату. Это эффективно, потому что оно основано на индексации слайсов и выполняет дополнительную работу только при необходимости.

def func(a, start, stop):
    stop += 1
    result = np.sum(a[start:stop])
    if stop > len(a):
        result += np.sum(a[:stop % len(a)])
    return result

Вышеприведенная версия работает для stop - start < len(a), то есть не более одного полного обхода. Для произвольного числа циклических циклов (т. Е. Произвольных значений для start и stop) может использоваться следующая версия:

def multi_wraps(a, start, stop):
    result = 0
    # Adjust both indices in case the start index wrapped around.
    stop -= (start // len(a)) * len(a)
    start %= len(a)
    stop += 1  # Include the element pointed to by the stop index.
    n_wraps = (stop - start) // len(a)
    if n_wraps > 0:
        result += n_wraps * a.sum()
    stop = start + (stop - start) % len(a)
    result += np.sum(a[start:stop])
    if stop > len(a):
        result += np.sum(a[:stop % len(a)])
    return result

В случае n_wraps > 0 некоторые части массива будут суммироваться удвоение, что излишне неэффективно, поэтому мы можем при необходимости вычислить сумму различных частей массива. Следующая версия суммирует каждый элемент массива не более одного раза:

def multi_wraps_efficient(a, start, stop):
    # Adjust both indices in case the start index wrapped around.
    stop -= (start // len(a)) * len(a)
    start %= len(a)
    stop += 1  # Include the element pointed to by the stop index.
    n_wraps = (stop - start) // len(a)
    stop = start + (stop - start) % len(a)  # Eliminate the wraps since they will be accounted for separately.
    tail_sum = a[start:stop].sum()
    if stop > len(a):
        head_sum = a[:stop % len(a)].sum()
        if n_wraps > 0:
            remaining_sum = a[stop % len(a):start].sum()
    elif n_wraps > 0:
        head_sum = a[:start].sum()
        remaining_sum = a[stop:].sum()
    result = tail_sum
    if stop > len(a):
        result += head_sum
    if n_wraps > 0:
        result += n_wraps * (head_sum + tail_sum + remaining_sum)
    return result

На следующем графике показано сравнение производительности между с использованием индексных массивов и двумя методами многократного переноса, представленными выше. Испытания проводятся на решетке (1_000, 1_000). Можно заметить, что для метода multi_wraps увеличивается время выполнения при переходе от 1 до 2 циклических циклов, поскольку он излишне суммирует массив дважды. Метод multi_wraps_efficient имеет одинаковую производительность независимо от количества циклов, поскольку он суммирует каждый элемент массива не более одного раза.

график производительности был создан с использованием пакета perfplot :

perfplot.show(
    setup=lambda n: (np.ones(shape=(1_000, 1_000), dtype=int), 400, n*1_000 + 200),
    kernels=[
        lambda x: index_arrays(*x),
        lambda x: multi_wraps(*x),
        lambda x: multi_wraps_efficient(*x),
    ],
    labels=['index_arrays', 'multi_wraps', 'multi_wraps_efficient'],
    n_range=range(1, 11),
    xlabel="Number of wrap-around",
    equality_check=lambda x, y: x == y,
)