Search code examples
matlabvectorizationweighted-average

Vectorize weighted sum matlab


I was trying to vectorize a certain weighted sum but couldn't figure out how to do it. I have created a simple minimal working example below. I guess the solution involves either bsxfun or reshape and kronecker products but I still have not managed to get it working.

rng(1);
N = 200;
T1 = 5;
T2 = 7;

A = rand(N,T1,T2);
w1 = rand(T1,1);
w2 = rand(T2,1);

B = zeros(N,1);

for i = 1:N
for j1=1:T1
for j2=1:T2
    B(i) = B(i) + w1(j1) * w2(j2) * A(i,j1,j2);
end
end
end

A = B;

Solution

  • You could use a combination of bsxfun, reshape and permute to accomplish this.

    We first use permute to move the N dimension to the 3rd dimension of A. We then multiply w1 and the transpose of w2 to create a grid of weights. We can then use bsxfun to perform element-wise multiplication (@times) between this grid and each "slice" of A. We can then reshape the 3D result into M x N and sum across the first dimension.

    B = sum(reshape(bsxfun(@times, w1 * w2.', permute(A, [2 3 1])), [], N)).';
    

    Update

    There's actually a simpler approach which would use matrix multiplication to perform the summation for you. It unfortunately has to be broken into

    % Create the grid of weights
    W = w1 * w2.';
    
    % Perform matrix multiplication between a 2D version of A and the weights
    B = reshape(A, N, []) * W(:);
    

    Or you could use kron to create the flattened grid of weights:

    B = reshape(A, N, []) * kron(w2, w1);