如何定期求和一个晶格的行的元素

Question

假设我有一个格子

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

我想制作一个具有3个输入的函数func(lattice, start, end)，其中start和end是该函数将元素相加的行的索引。例如，对于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
    stop -= (start // len(a)) * len(a)
    start %= len(a)
    stop += 1
    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):
    stop -= (start // len(a)) * len(a)
    start %= len(a)
    stop += 1
    n_wraps = (stop - start) // len(a)
    stop = start + (stop - start) % len(a)
    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

下图显示了using index arrays与上面介绍的两种多重包装方法之间的性能比较。测试在(1_000, 1_000)晶格上运行。可以观察到，对于multi_wraps方法，从1到2环绕时，运行时会增加，因为它不必要地将数组加总两次。 multi_wraps_efficient方法具有相同的性能，而与环绕数无关，因为它对每个数组元素进行的累加不超过一次。

性能图是使用perfplot package生成的：

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,
)