Outer minimum vectorization in numpy-CodePudding

Given an NxM matrix A, I want to efficiently obtain the NxMxN tensor whose ith layer is the application of np.minimum between A and the ith row of A. Using a for loop:

> A = np.array([[1, 2], [3, 4], [5,6]])
> output = np.zeros(shape=(A.shape[0], A.shape[1], A.shape[0]))
> for i in range(a.shape[0]):
      output[:, :, i] = np.minimum(A, A[i])
> output
array([[[1., 1., 1.],
        [2., 2., 2.]],

       [[1., 3., 3.],
        [2., 4., 4.]],

       [[1., 3., 5.],
        [2., 4., 6.]]])

This is very slow so I would like to get rid of the for loop and vectorize it. I feel like there should be a general method that works with any function of a matrix and a vector not just, minimum. Using np.minimum.outer does not work as it outputs an order 4 tensor.

CodePudding user response：

With broadcasting we can make a (3,3,2) result:

In [153]: np.minimum(A[:,None,:],A[None,:,:])
Out[153]: 
array([[[1, 2],
        [1, 2],
        [1, 2]],

       [[1, 2],
        [3, 4],
        [3, 4]],

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

and then switch the last 2 dimensions to get the (3,2,3) you want:

In [154]: np.minimum(A[:,None,:],A[None,:,:]).transpose(0,2,1)
Out[154]: 
array([[[1, 1, 1],
        [2, 2, 2]],

       [[1, 3, 3],
        [2, 4, 4]],

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

Or do the transpose first

In [155]: np.minimum(A[:,:,None],A.T[None,:,:])  # (3,2,1) (1,2,3)=>(3,2,3)
Out[155]: 
array([[[1, 1, 1],
        [2, 2, 2]],

       [[1, 3, 3],
        [2, 4, 4]],

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

edit

Adding the sum is trivial:

In [157]: np.minimum(A[:,:,None],A.T[None,:,:]).sum(axis=1)
Out[157]: 
array([[ 3,  3,  3],
       [ 3,  7,  7],
       [ 3,  7, 11]])

CodePudding user response：

To efficiently obtain the tensor you described without using a for loop, you can use the np.einsum function. This function allows you to specify a summation operation using a simple string notation, which can be much faster than using a for loop.

import numpy as np


A = np.array([[1, 2], [3, 4], [5,6]])


output = np.einsum('ij,ik->ijk', A, A)


print(output)

This code will produce the same output as your original code, but without using a for loop.

The key to using np.einsum in this case is the string 'ij,ik->ijk'. This string specifies the summation operation that np.einsum will perform. The first two letters 'ij' specify the indices of the first operand (in this case, A), and the second two letters 'ik' specify the indices of the second operand (also A). Finally, the last three letters 'ijk' specify the indices of the output tensor.

In this case, the np.einsum function will compute the tensor by applying the minimum operator element-wise to the two operands. This is equivalent to your original code, but it is much faster because it uses a vectorized operation instead of a for loop.