Search code examples
pythonarraysnumpyfor-loopvectorization

Multiplying array of vectors by matrix without for loop


I have a 2 x 2 numpy.array() matrix, and an array N x 2 X, containing N 2-dimensional vectors.

I want to multiply each vector in X by the 2 x 2 matrix. Below I use a for loop, but I am sure there is a faster way. Please, could someone show me what it is? I assume there is a way using a numpy function.

# the matrix I want to multiply X by
matrix = np.array([[0, 1], [-1, 0]]) 

# initialize empty solution
Y = np.empty((N, 2))

# loop over each vector in X and create a new vector Y with the result
for i in range(0, N):
    Y[i] = np.dot(matrix, X[i]) 

For example, these arrays:

matrix = np.array([
    [0,  1],
    [0, -1]
])

X = np.array([
    [0, 0],
    [1, 1],
    [2, 2]
])

Should result in:

Y = np.array([
    [0,  0],
    [1, -1],
    [2, -2]
])

Solution

  • One-liner is (matrix @ X.T).T

    Just transpose your X, to get your vectors in columns. Then matrix @ X.T or (np.dot(matrix, X.T) if you prefer this solution, but now that @ notation exists, why not using it) is a matrix made of columns of matrix times X[i]. Just transpose back the result if you need Y to be made of lines of results

    matrix = np.array([[0, 1], [-1, 0]]) 
    X = np.array([[1,2],[3,4],[5,6]])
    Y = (matrix @ X.T).T
    

    Y is

    array([[ 2, -1],
           [ 4, -3],
           [ 6, -5]])
    

    As expected, I guess.

    In detail:
    X is

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

    so X.T is

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

    So, you can multiply your 2x2 matrix by this 2x3 matrix, and the result will be a 2x3 matrix whose columns are the result of multiplication of matrix by the column of this. matrix @ X.T is

    array([[ 2,  4,  6],
           [-1, -3, -5]])
    

    And transposing back this gives the already given result.

    So, tl;dr: one-liner answer is (matrix @ X.T).T