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
考虑上面的代码。我有一个形状为{4,4,4,16,3)的数组A
,它表示晶格中的3D位置矢量(昏暗3的最后一个轴是x,y,z坐标)。前四个维度在晶格中索引该位置。
我想从A
生成一个数组,其中包含晶格中各个位点之间的所有可能的分离向量。这意味着形状为B
的输出数组(4,4,4,16,4,4,4,16,3)
。位置向量差的(x,y,z)坐标的最后一个维度是站点i的前四个维度,然后是站点j的下四个维度。
即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
的最后一个轴((x,y,z))上进行逐元素减法。
在我的尝试中,B
具有我想要的正确尺寸,但这显然是错误的。有什么提示吗? (除非使用任何for循环)
答案 0 :(得分: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