时间戳

首页 > 解决方案 > 在扁平数组上实现 numpy.roll

python - 在扁平数组上实现 numpy.roll

问题描述

我正在尝试将我的 2d Ising 模型（具有周期性边界条件）模拟器概括为 Nd 作为个人项目。

作为对内容的快速回顾，请参阅此 wiki 页面（因为 Stack Overflow 不支持 Latex 渲染）Ising-Wiki。但是，对于我的问题，了解这方面的细节并不是绝对必要的。

无论如何，为了使一切尽可能通用，我认为使用扁平数组最容易

import numpy as np

# create nd-spin config
Nx, Ny, Nz = 32, 32, 32  # size of each dimension of 3d lattice
spin_config = np.random.choice([-1, 1], (Nx, Ny, Nz)).astype(np.int8)
spin_iter = np.asarray(spin_config.strides) / spin_config.itemsize

对于站点 (i, j, k)，每个站点的最近邻居将是

#  a[(i-1) % Nz, j, k],  a[(i+1) % Nz, j, k],
#  a[i, (j-1) % Ny, k],  a[i, (j+1) % Ny, k],
#  a[i, j, (k-1) % Nx],  a[i, j, (k+1) % Nk]

上面的问题是

解决这个问题的一种方法是使用 numpy.roll

for j in range(spin_config.shape):
    np.roll(spin_config, -1, axes=j)
    # do stuff
    np.roll(spin_config, 1, axes=j)
    # do stuff

问题在于，如果我想利用 numba jit 编译器，jit 编译器不支持带有轴参数的 np.roll参见此处

如果我展平我的数组，那么可以通过以下方式构造最近的邻居

# for a spin s at location niter in spin_config.ravel()
s_nn = np.zeros(2*len(spin_iter))       # sites for nearest spin on array
for j, sval in enumerate(spin_iter):
    s_nn[2*j: 2*(j+1)] = [(niter-sval) % spin_config.size, (niter+sval) % spin_config.size]
spin_nn = spin_config.ravel()[s_nn]     # sub array containing nearest neighbors

不幸的是，我不确定如何处理我的任何站点索引（i、j、k）分别为 Nz-1、Ny-1 或 Nx-1 的边缘情况。

你们有一个聪明的方法在一个未解开的 numpy 数组中绕轴滚动吗？

标签： pythonnumpy-ndarraynumbanumpy-slicing

解决方案

我想我有一个不错的答案。对于初学者，展平的 nd 数组具有与以下索引公式相关联的索引：

# index = (Ny*Nx)*nz + (Nx)*ny + nx
# spin_config.strides / spin_config.itemsize = (Ny*Nx, Nx, 1)

nj = Nj-1每当j=x, y, z时，边界附近都会出现问题。这很有用，因为现在，测试我是否在边界附近是相邻索引是否除以 Nj-1。所以这是一个新的实现

spin_shape = spin_config.shape
s_nn = np.zeros(2*len(spin_iter))
for j, (sval, s_shape) in enumerate(zip(spin_iter, spin_shape)):
    shift_down = s_shape if (niter-sval) % s_shape-1 else (niter-sval)
    shift_up = s_shape if (niter+sval) % s_shape-1 else (niter+sval)
    s_nn[2*j: 2*(j+1)] = [shift_down, shift_up]
spin_nn = spin_config.ravel()[s_nn]

如果我在这里犯了某种可怕的错误，请告诉我。

推荐阅读