2
import numpy as np
import itertools as it

SPIN_POS = np.array([[0, 0, 0], [1, 1, 0], [1, 0, 1], [0, 1, 1],  
                     [2, 2, 0], [3, 3, 0], [3, 2, 1], [2, 3, 1], 
                     [2, 0, 2], [3, 1, 2], [3, 0, 3], [2, 1, 3],  
                     [0, 2, 2], [1, 3, 2], [1, 2, 3], [0, 3, 3] 
                     ]) / 4

def gen_posvecs(xdim:int, ydim:int, zdim:int):
    """
    Generates position vectors of site pairs in the lattice of size xdim,ydim,zdim

    :param x,y,z is the number of unit cells in the x,y,z directions;
    :returns array containing the position vectors
    """
    poss = np.zeros((xdim,ydim,zdim,16,3))
    for x,y,z,s in it.product(range(xdim), range(ydim), range(zdim), range(16)):
        poss[x,y,z,s] = np.array([x,y,z]) + SPIN_POS[s]
    return poss

A = gen_sepvecs(4,4,4)  # A.shape = (4,4,4,16,3)
B = np.subtract.outer(A[...,-1], A)  # my attempt at a soln
assert all(A[1,2,0,12] - A[0,1,3,11] == B[1,2,0,12,0,1,3,11])  # should give true

考虑上面的代码。我有一个形状数组A(4,4,4,16,3),它表示格子中的 3D 位置向量(dim 3 的最后一个轴是 x、y、z 坐标)。前 4 个维度索引点阵中的站点。

我想要的是

我想从 中生成A一个数组,该数组包含晶格中站点之间所有可能的分离向量。B这意味着一个形状为 的输出数组(4,4,4,16,4,4,4,16,3)。前 4 个维度是站点 i,接下来是站点 j 的 4 个维度,然后是位置向量差的 (x,y,z) 坐标的最后一个维度。

A[a,b,c,d]:形状(3,)是第一个站点的(x,y,z);A[r,s,t,u]:形状(3,)是第二个站点的(x,y,z);然后我想B[a,b,c,d,r,s,t,u]成为前两者之间的 (x,y,z) 差异。

我的尝试

我知道该ufunc.outer功能,正如您在我的代码尝试中看到的那样。但我坚持将它与在每个A.

在我的尝试中,B具有我想要的正确尺寸,但这显然是错误的。有什么提示吗?(除非使用任何 for 循环)

4

1 回答 1

1

我认为你只需要这样做:

B = (A[:, :, :, :, np.newaxis, np.newaxis, np.newaxis, np.newaxis] -
     A[np.newaxis, np.newaxis, np.newaxis, np.newaxis])

在您的代码中:

import numpy as np
import itertools as it

SPIN_POS = np.array([[0, 0, 0], [1, 1, 0], [1, 0, 1], [0, 1, 1],  
                     [2, 2, 0], [3, 3, 0], [3, 2, 1], [2, 3, 1], 
                     [2, 0, 2], [3, 1, 2], [3, 0, 3], [2, 1, 3],  
                     [0, 2, 2], [1, 3, 2], [1, 2, 3], [0, 3, 3] 
                     ]) / 4

def gen_posvecs(xdim:int, ydim:int, zdim:int):
    """
    Generates position vectors of site pairs in the lattice of size xdim,ydim,zdim

    :param x,y,z is the number of unit cells in the x,y,z directions;
    :returns array containing the position vectors
    """
    poss = np.zeros((xdim,ydim,zdim,16,3))
    for x,y,z,s in it.product(range(xdim), range(ydim), range(zdim), range(16)):
        poss[x,y,z,s] = np.array([x,y,z]) + SPIN_POS[s]
    return poss

A = gen_posvecs(4,4,4)  # A.shape = (4,4,4,16,3)
B = A[:, :, :, :, np.newaxis, np.newaxis, np.newaxis, np.newaxis] - A[np.newaxis, np.newaxis, np.newaxis, np.newaxis]

assert all(A[1,2,0,12] - A[0,1,3,11] == B[1,2,0,12,0,1,3,11])
# Does not fail
于 2018-11-01T14:46:25.570 回答