Extended matrix generation

Question

I am extending a matrix with some calculated values. This code runs very slowly.

Is there a way to vectorize this with slice and hstack?

import numpy as np

def expand(X):
    X_expanded = np.zeros((X.shape[0], 6))
    for i in range(X.shape[0]):
        X_expanded[i,0]=X[i,0]
        X_expanded[i,1]=X[i,1]
        X_expanded[i,2]=X[i,0]**2
        X_expanded[i,3]=X[i,1]**2
        X_expanded[i,4]=X[i,0]*X[i,1]
        X_expanded[i,5]=1
    return X_expanded

Answer 1

As you hinted in your question, when working with numpy, anytime you use a Python for loop when manipulating your matrices, there is a good chance there will be an opportunity to "vectorize" the operation, and speed it up. Your code was actually a good start as it allowed you to get the right result with more familiar coding.

To convert it to a vectorized operation you then need to describe the loop in a vector style. In this case it was as simple as removing the for and replacing the i loop variable with :

def expand(X):
    X_expanded = np.zeros((X.shape[0], 6))
    X_expanded[:, 0] = X[:, 0]
    X_expanded[:, 1] = X[:, 1]
    X_expanded[:, 2] = X[:, 0] ** 2
    X_expanded[:, 3] = X[:, 1] ** 2
    X_expanded[:, 4] = X[:, 0] * X[:, 1]
    X_expanded[:, 5] = 1
    return X_expanded
    

The above will almost certaily be much faster than the original. And depending on your needs you might stop there. But since you also asked about stacking, you can remove the intermediate np.zeros like:

def expand(X):
    return np.stack((
        X[:, 0],
        X[:, 1],
        X[:, 0] ** 2,
        X[:, 1] ** 2,
        X[:, 0] * X[:, 1],
        np.ones(X[:, 0].shape)
    ), axis=1)

Test Code:

def original_expand(X):
    X_expanded = np.zeros((X.shape[0], 6))
    for i in range(X.shape[0]):
        X_expanded[i, 0] = X[i, 0]
        X_expanded[i, 1] = X[i, 1]
        X_expanded[i, 2] = X[i, 0] ** 2
        X_expanded[i, 3] = X[i, 1] ** 2
        X_expanded[i, 4] = X[i, 0] * X[i, 1]
        X_expanded[i, 5] = 1
    return X_expanded

data = np.array(list(zip(range(10), range(1, 11))))
assert not (expand(data) - original_expand(data)).any()