Can this loop containing different indices be vectorized using implicit expansion (or otherwise)?

I have a code which is doing some processing over every point in a 3D matrix. Until recently, I was doing this by reshaping the array to 1D first, then doing the processing, then reshaping back to 3D afterwards, like this:

Nx = 8; Ny = 6; Nz = 4; Ntot = Nx*Ny*Nz;
xvals = rand(1,Nx); yvals = rand(1,Ny); zvals = rand(1,Nz);

input_vec_3D = rand(Ny,Nx,Nz);

factor1 = 3.6*xvals;
factor2 = 1.2*yvals;
factor3 = 8.5*zvals;

%%% NON-VECTORIZED METHOD %%%
input_vec_1D = reshape( permute(input_vec_3D,[3,1,2]) , [Ntot 1]); % Reshape to 1D for loop
output_vec1 = zeros(Ntot,1);
for ind = 1:Ntot   % loop over every point in the lattice
    j1 = floor( floor((ind-1) / Nz) / Ny) + 1;
    j2 = mod( floor((ind-1)/Nz) , Ny ) + 1;
    j3 = mod( (ind-1), Nz ) + 1;
    output_vec1(ind) = input_vec_1D(ind) * factor1(j1)*factor2(j2)*factor3(j3); % Do processing
end
output_vec1 = permute( reshape( output_vec1, [Nz,Ny,Nx] ) , [2,3,1] );  % Reshape back to 3D

However, I have since become aware that the above code can be dramatically simplified (and speeded up) by using implicit expansion , which is great, like this:

%%% VECTORIZED METHOD %%%
factor1 = 3.6*xvals;                    % Size is  (1 x Nx x 1) for implicit expansion
factor2 = (1.2*yvals).';                % Size is (Ny x 1 x 1) for implicit expansion
factor3 = permute( 8.5*zvals ,[1,3,2]); % Size is (1 x 1 x Nz) for implicit expansion
output_vec1 = input_vec_3D .* factor1 .* factor2 .* factor3;

In this example, it was crucial for my application that I did not store any temporary extra large matrices due to memory requirements (notice that factor1, factor2, factor3 are only 1D vectors, so memory usage for them is small), which precludes the use of meshgrid() to generate the factors.

I would now like to know is something similar can be done to eliminate the following loop as well:

output_vec2 = zeros(Ntot,1);
for ind = 1:Ntot
    j1 = floor( floor( (ind-1)/Nz ) /Ny ) + 1;   
    j2 = mod( floor( (ind-1)/Nz ) , Ny ) + 1; 
    j3 = mod( (ind-1) , Nz ) + 1;
    n1 = mod( 5*(j1-1) ,Nx);
    n2 = mod( 3*(j2-1) ,Ny);
    n3 = mod( 2*(j3-1) ,Nz);
    ind_prime = mod( ( n3 + Nz*(n2 + Ny*n1) ) , Ntot ) + 1; % a different index for input_vec
    output_vec2(ind) = output_vec2(ind) + input_vec_1D(ind_prime) * factor1(j1)*factor2(j2)*factor3(j3);
end

This is a little more complicated, because in the processing line different indices are used. Again, I do not want to precalculate a large 3D array of repeated indices, because that would use significantly more memory than 1D implicit expansion. How can this be best done?

[Thomas Barrett]

Can anyone show me how to solve this? Thank you.

Answer 1

[darova]

Simply remove keywords for and end. Don't forget about element wise operator (dot)

[Thomas Barrett]

@darova No, this would then generate many large 1 x Ntot arrays of indices, for each of ind, j1,j2,...n1,n2,...ind_prime, etc. This would be a hugh additional RAM overhead.