How can I avoid a loop if I need to do a matrix multiplication?

Question

I have the following code:

import numpy as np
import torch

y = torch.ones((1000,10)) #This is the output of a neural network which does not matter here
theta_min = 0; theta_max = np.pi; K = 10; sigma = 10;
z = torch.zeros(y.shape)
for i in range(0,y.shape[0]):
    theta = np.random.uniform(theta_min, theta_max)
    vector = np.reshape(np.exp(-1j * np.arange(0,K) * np.pi * np.sin(theta)),(-1,1))
    vector = torch.tensor(vector)
    alpha = sigma * np.random.randn()
    z[i,:] = alpha * vector @ vector.T @ y[i,:].T

How can I avoid the loop to make the code faster?

Answer 1

I believe this will work:

import numpy as np
import torch

y = torch.ones((1000,10))
theta_min = 0; theta_max = np.pi; K = 10; sigma = 10;
z = torch.zeros(y.shape)

theta = np.random.uniform(theta_min, theta_max, y.shape)
vector = np.exp(-1j * np.arange(0,K) * np.pi * np.sin(theta))
vector = torch.FloatTensor(vector)

temp = (vector.T @ vector)
alpha = sigma * torch.rand(temp.shape)
z = (alpha*temp @ y.T).T

to speed it up further, you can use torch for the whole process:

import numpy as np
import torch

y = torch.ones((1000,10)).float()
theta_min = 0; theta_max = np.pi; K = 10; sigma = 10;
alpha = sigma * torch.rand(1)

theta = torch.rand(y.shape)*(np.pi)
vector = torch.exp(-1j * torch.arange(0,K) * np.pi * torch.sin(theta)).float()

temp = (vector.T @ vector)
alpha = sigma * torch.rand(temp.shape)
z = (alpha*temp @ y.T).T

Answer 2

Following the solution of Megan Hardy, I have tried instead of make use of 3d arrays. We will make some useless operations, but it is better than the for loop. The code would look like this:

y = torch.ones((1000,10)).type(torch.complex64) 
theta_min = 0; theta_max = np.pi; K = 10; sigma = 10;

theta = torch.rand(y.shape[0],1,1) * np.pi  #Each slice in the 0 dimension will be treated separately
temp = torch.reshape(torch.arange(0,K),(-1,1))    #Necessary so that it doesn't have dimensions (10,), but (10,1)
vector = torch.exp(-1j * temp * np.pi * torch.sin(theta))
matrix = vector @ torch.transpose(vector,1,2)   #Maintain the 0 dimension but exchange dimension 1 and 2

alpha = sigma * torch.rand(1)

z = alpha * matrix @ y.T    #This makes more operations that necessary
temp = range(0,y.shape[0])  
z = alpha * z[temp,:,temp]  #Take only the desired columns