Deathcrush
Deathcrush

Reputation: 77

Tensor expansion in numpy

I am trying to find an elegant way to solve the following problem.

I have a tensor y containing n dxd matrices, which I obtain by choosing combinations of rows of a matrix X. I have a second numpy array, of size k x d. My intention is to expand the tensor so that I add each of the rows of the k x d matrix to each of the elements in the y tensor, to obtain a y' tensor with k x n d x (d+1) matrices.

I cannot see how to do it without a for loop. My simple code example is as follows:

#Array x
X = np.arange(27).reshape(9,3)
# Create tensor y
combs = [(0,1,2),(0,1,3),(0,1,4),(0,1,5)]
y = X[combs,:]
# Add a dummy column of 1.0s to each element of the y tensor
b = np.array([1.0,1.0,1.0]).reshape(1,3)
b = b.repeat(y.shape[0],axis=0).reshape(y.shape[0],y.shape[1],1)
# Concatenate the column with the tensor
y_new = np.concatenate((y,b),axis=2)`

This solution is far from ideal because I would have to keep a copy of the original array, loop through all the rows, obtain k tensors and then merge them at the end. In the general problem I am trying to solve, the tensors y are large and multiple processes execute in parallel, so this tensor expansion should ideally be as efficient as possible. Any suggestions appreciated!

Upvotes: 0

Views: 867

Answers (2)

lfog
lfog

Reputation: 11

I believe the following code does what you are requiring without loops. The idea is to expand the tensors to the needed dimension and then perform concatenation with numpy in one line:

X = np.arange(27).reshape(9,3)
# Create tensor y
combs = [(0,1,2),(0,1,3),(0,1,4),(0,1,5)]
y = X[combs,:]
# Create tensor b, (the k x d matrix in your question)
b = np.arange(100,121).reshape(7,3,1)
# expand the tensors 
b = np.stack([b]*4,axis=1)
y = np.stack([y]*7)
# concatenate
y_new = np.concatenate([y,b],axis=3) 

Upvotes: 1

hpaulj
hpaulj

Reputation: 231385

Your repeat and concatenate code can be simplified, and probably sped up with:

In [50]: z = np.zeros((y.shape[:-1]+(y.shape[-1]+1,)))
In [51]: z.shape
Out[51]: (4, 3, 4)
In [52]: z[:,:,:-1]=y
In [53]: z[:,:,-1]=np.array([.1,.2,.3])

That is create target array, and fill with values from y and b. With broadcasting b doesn't need the reshape and repeat.

It sounds as though you are embedding this y_new in some sort of loop, but I haven't followed those details.

Upvotes: 1

Related Questions