Creating this block matrix in numpy

Question

I have two sets of 3D points in numpy and I would like to create a matrix and vector representations of the points as follows:

| X1 Y1 Z1 0  0  0  0  0  0  1 0 0|     | X1 |
| 0  0  0  X1 Y1 Z1 0  0  0  0 1 0|     | Y1 |
| 0  0  0  0  0  0  X1 Y1 Z1 0 0 1|     | Z1 |
| X2 Y2 Z2 0  0  0  0  0  0  1 0 0|     | X2 |
| 0  0  0  X2 Y2 Z2 0  0  0  0 1 0|     | Y2 |
| 0  0  0  0  0  0  X2 Y2 Z2 0 0 1|     | Z2 |

A usage would be something like:

import numpy as np
pts = np.random.rand(10, 3)

So the matrix would now have the shape (30, 12). 30 rows (3 per point) and 12 columns. The matrix would be 30 elements long in this case. Is there a way to achieve this in python without writing an explicit for loop?

Answer 1

The Kronecker product (np.kron) is very useful for building block matrices like this:

import numpy as np

pts = np.arange(1, 31).reshape(10, 3)

n, d = pts.shape
I = np.eye(d, dtype=pts.dtype)

# the first d**2 columns of xyz values
xyzcols = np.kron(I, pts[:, None]).reshape(-1, d * d)

# the final d columns of ones
eyecols = np.tile(I, n).T

# concatenate
out = np.hstack([xyzcols, eyecols])

print(repr(out[:6]))
# array([[1, 2, 3, 0, 0, 0, 0, 0, 0, 1, 0, 0],
#        [0, 0, 0, 1, 2, 3, 0, 0, 0, 0, 1, 0],
#        [0, 0, 0, 0, 0, 0, 1, 2, 3, 0, 0, 1],
#        [4, 5, 6, 0, 0, 0, 0, 0, 0, 1, 0, 0],
#        [0, 0, 0, 4, 5, 6, 0, 0, 0, 0, 1, 0],
#        [0, 0, 0, 0, 0, 0, 4, 5, 6, 0, 0, 1]])

Answer 2

One vectorized approach abusing linear indexing and broadcasted indexing -

m,n = pts.shape

idx1 = np.arange(n)[:,None] + np.arange(n)*(n*(n+2))
idx2  = idx1 + np.arange(m)[:,None,None]*(n*n*(n+1))
idx3 = (n*n + np.arange(n)*(n*(n+1)+1)) + np.arange(m)[:,None]*(n*n*(n+1))

out = np.zeros((m*n,n*(n+1)),dtype=pts.dtype)
out.ravel()[idx2] = pts[:,:,None]
out.ravel()[idx3] = 1

Sample run -

In [550]: pts
Out[550]: 
array([[47, 34],
       [36, 25],
       [29, 38],
       [35, 20],
       [37, 48]])

In [551]: m,n = pts.shape
     ...: 
     ...: idx1 = np.arange(n)[:,None] + np.arange(n)*(n*(n+2))
     ...: idx2  = idx1 + np.arange(m)[:,None,None]*(n*n*(n+1))
     ...: idx3=(n*n + np.arange(n)*(n*(n+1)+1)) + np.arange(m)[:,None]*(n*n*(n+1))
     ...: 
     ...: out = np.zeros((m*n,n*(n+1)),dtype=pts.dtype)
     ...: out.ravel()[idx2] = pts[:,:,None]
     ...: out.ravel()[idx3] = 1
     ...: 

In [552]: out
Out[552]: 
array([[47, 34,  0,  0,  1,  0],
       [ 0,  0, 47, 34,  0,  1],
       [36, 25,  0,  0,  1,  0],
       [ 0,  0, 36, 25,  0,  1],
       [29, 38,  0,  0,  1,  0],
       [ 0,  0, 29, 38,  0,  1],
       [35, 20,  0,  0,  1,  0],
       [ 0,  0, 35, 20,  0,  1],
       [37, 48,  0,  0,  1,  0],
       [ 0,  0, 37, 48,  0,  1]])