Combining outer subtraction with element wise subtraction on the last axis?

Question

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

Consider the above code. I have an array A of shape (4,4,4,16,3), which represents 3D position vectors in a lattice (the last axis of dim 3 are the x,y,z coordinates). The first 4 dimensions index the site in the lattice.

What I want

I would like to generate from A, an array containing all possible separation vectors between sites in the lattice. This means an output array B, of shape (4,4,4,16,4,4,4,16,3). The first 4 dimensions being of site i, next 4 dimensions of site j, then the last dimension of the (x,y,z) coordinate of the position vector difference.

i.e., A[a,b,c,d]: shape (3,) is the (x,y,z) of first site; A[r,s,t,u]: shape (3,) is the (x,y,z) of second site; Then I want B[a,b,c,d,r,s,t,u] to be (x,y,z) difference between the first two.

My attempt

I know about the ufunc.outer function, as you can see in my attempt in code. But I'm stuck at applying it together with performing element-wise subtraction on the last axis (the (x,y,z)) of each A.

In my attempt, B has the correct dimensions I want, but it is obviously wrong. Any hints? (barring the use of any for-loops)

score 1 · Accepted Answer · answered Nov 01 '18 at 14:46

I think you just need to do:

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

In your code:

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

yup I think that's right. I should have just fall back on normal broadcasting.. Thanks (+1). I'll leave the question a while more before accepting, hope that's okay with you — Troy, Nov 01 '18 at 14:57

Combining outer subtraction with element wise subtraction on the last axis?

What I want

My attempt

1 Answers1