CODEWITHSUNDEEP
pythonarraysnumpytensorflowvectorization
linello
linello

Reputation: 8704

Vectorized creation of an array of diagonal square arrays from a liner array in Numpy or Tensorflow

I have an array of shape [batch_size, N], for example:

[[1  2]
 [3  4]
 [5  6]]

and I need to create a 3 indices array with shape [batch_size, N, N] where for every batch I have a N x N diagonal matrix, where diagonals are taken by the corresponding batch element, for example in this case, In this simple case, the result I am looking for is:

[
  [[1,0],[0,2]],
  [[3,0],[0,4]],
  [[5,0],[0,6]],
]

How can I make this operation without for loops and exploting vectorization? I guess it is an extension of dimension, but I cannot find the correct function to do this. (I need it as I am working with tensorflow and prototyping with numpy).

Upvotes: 3

Views: 1353

Answers (5)

user9413641
user9413641

Reputation:

You basically want a function that does the opposite of/reverses np.block(..)

I needed the same thing, so I wrote this little function:

def split_blocks(x, m=2, n=2):
    """
    Reverse the action of np.block(..)

    >>> x = np.random.uniform(-1, 1, (2, 18, 20))
    >>> assert (np.block(split_blocks(x, 3, 4)) == x).all()

    :param x: (.., M, N) input matrix to split into blocks
    :param m: number of row splits
    :param n: number of column, splits
    :return:
    """
    x = np.array(x, copy=False)
    nd = x.ndim
    *shape, nr, nc = x.shape
    return list(map(list, x.reshape((*shape, m, nr//m, n, nc//n)).transpose(nd-2, nd, *range(nd-2), nd-1, nd+1)))

Upvotes: 0

caverac
caverac

Reputation: 1637

Yu can use numpy.diag

m = [[1, 2],
 [3, 4],
 [5, 6]]

[np.diag(b) for b in m]

EDIT The following plot shows the average execution time for the solution above (solid line), and compared it against @Divakar's (dashed line) for different batch-sizes and different matrix sizes

enter image description here

I don't believe you get much of an improvement, but this is just based on this simple metric

Upvotes: 1

yatu
yatu

Reputation: 88236

Using np.expand_dims with an element-wise product with np.eye

a = np.array([[1,  2],
              [3,  4],
              [5, 6]])
N = a.shape[1]
a = np.expand_dims(a, axis=1)
a*np.eye(N)

array([[[1., 0.],
       [0., 2.]],

      [[3., 0.],
       [0., 4.]],

      [[5., 0.],
       [0., 6.]]])

Explanation

np.expand_dims(a, axis=1) adds a new axis to a, which will now be a (3, 1, 2) ndarray:

array([[[1, 2]],

       [[3, 4]],

       [[5, 6]]])

You can now multiply this array with a size N identity matrix, which you can generate with np.eye:

np.eye(N)
array([[1., 0.],
       [0., 1.]])

Which will yield the desired output:

a*np.eye(N)

array([[[1., 0.],
        [0., 2.]],

       [[3., 0.],
        [0., 4.]],

       [[5., 0.],
        [0., 6.]]])

Upvotes: 2

Divakar
Divakar

Reputation: 221584

Approach #1

Here's a vectorized one with np.einsum for input array, a -

# Initialize o/p array
out = np.zeros(a.shape + (a.shape[1],),dtype=a.dtype)

# Get diagonal view and assign into it input array values
diag = np.einsum('ijj->ij',out)
diag[:] = a

Approach #2

Another based on slicing for assignment -

m,n = a.shape
out = np.zeros((m,n,n),dtype=a.dtype)
out.reshape(-1,n**2)[...,::n+1] = a

Upvotes: 2

giser_yugang
giser_yugang

Reputation: 6166

Try it in tensorflow:

import tensorflow as tf
A = [[1,2],[3 ,4],[5,6]]
B = tf.matrix_diag(A)
print(B.eval(session=tf.Session()))
[[[1 0]
  [0 2]]

 [[3 0]
  [0 4]]

 [[5 0]
  [0 6]]]

Upvotes: 2

Related Questions