Объединение внешнего вычитания с поэлементным вычитанием на последней оси? - PullRequest
Новая
       15

Объединение внешнего вычитания с поэлементным вычитанием на последней оси?

0 голосов
/ 01 ноября 2018 
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), который представляет трехмерные векторы положения в решетке (последняя ось dim 3 - это координаты x, y, z). Первые 4 измерения индекса сайта в решетке.

Что я хочу

Я хотел бы создать из A массив, содержащий все возможные векторы разделения между узлами в решетке. Это означает, что выходной массив B имеет форму (4,4,4,16,4,4,4,16,3). Первые 4 измерения относятся к сайту i, следующие 4 измерения - к сайту j, а затем к последнему измерению (x, y, z) координаты разности вектора положения.

, т.е. 1012 *: форма (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, как вы можете видеть в моей попытке в коде. Но я застрял в применении его вместе с поэлементным вычитанием по последней оси ((x, y, z)) каждого A.

В моей попытке B имеет правильные размеры, которые я хочу, но это, очевидно, неправильно. Есть намеки? (исключая использование любых циклов for)

1 Ответ

0 голосов
/ 01 ноября 2018

Я думаю, вам просто нужно сделать: 

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
...