Vectorizing Tensor Products from Python to Matlab-CodePudding

I am in the process of currently trying to convert some code from Python into Matlab. I have code working that produces the same results but I am wondering if there may be a way to vectorize some of my for loops in the Matlab code as it take a long time to run. X in an Nxd matrix, diff is an NxNxd tensor, kxy is an NxN matrix, gradK is an NxNx2 tensor, and sumkxy, dxkxy, and obj are all Nxd matrices.

Here is the original Python Code:

diff = x[:, None, :] - x[None, :, :]  # D_{ij, s}
kxy = np.exp(-np.sum(diff ** 2, axis=-1) / (2 * h ** 2)) / np.power(np.pi * 2.0 * h * h, d / 2) # -1 last dimension K_{ij]
gradK = -diff * kxy[:, :, None] / h ** 2  # N * N * 2
sumkxy = np.sum(kxy, axis=1)
dxkxy = np.sum(gradK, axis=1)  # N * 2    sum_{i} d_i K_{ij, s}
obj = np.sum(gradK / sumkxy[None, :, None], axis=1)  # N * 2

and here is my initial Matlab Code with all the for loops:

diff = zeros([n,n,d]);
for i = 1:n
    for j = 1:n
        for k = 1:d
            diff(i,j,k) = x(i,k) - x(j,k);
        end
    end
end

kxy = exp(-sum(dif.^2, 3)/(2*h^2))/((2*pi*h^2)^(d/2));
sumkxy = sum(kxy,2);
gradK = zeros([n,n,d]);
for i = 1:n
    for j = 1:n
        for k = 1:d
            gradK(i,j,k) = -diff(i,j,k)*kxy(i, j)/h^2;
        end
    end
end

dxkxy = squeeze(sum(gradK,2));
a = zeros([n,n,d]);
for i =1:n
    for j = 1:n
        for k = 1:d
            a(i,j,k) = gradK(i,j,k)/sumkxy(i);
        end
    end
end
obj = squeeze(sum(a, 2));

I know a way a faster way to calculate the kxy term is to use the following code:

XY = x*x';
x2= sum(x.^2, 2);
X2e = repmat(x2, 1, n);
 
H = (X2e   X2e' - 2*XY); % calculate pairwise distance
Kxy = exp(-H/(2*h^2))/((2*pi*h*h)^(d/2));

But then I struggle on a way to then calculate gradK efficiently without diff. Any help or suggestions would be greatly appreciated!

CodePudding user response：

I like to try approaching problems like these by breaking it down into "subcomponents" with some bogus data that will execute quickly and that you can use to test the code functionality. The first subcomponent you might start with is your first nested loop calculating diff:

n = 100;
d = 50;
x = round(100*rand(n,d));

tic
diff = zeros([n,n,d]);
for i = 1:n
    for j = 1:n
        for k = 1:d
            diff(i,j,k) = x(i,k) - x(j,k);
        end
    end
end
toc

First, consider the innermost loop on its own:

...
for k = 1:d
    diff(i,j,k) = x(i,k) - x(j,k);
end
...

Looking at this loop (at least for me!) simplifies things greatly. To vectorize just this "subcomponent" we could write something like:

diff(i,j,:) = x(i,:) - x(j,:);

Now that the low hanging fruit is out of the way, lets consider the next layer of loop. Does doing the same trick as before work?

diff(i,:,:) = x(i,:) - x; % where x(:,:) can just be written as x.

If you aren't sure, you can check this by running both the nested loop version and the one above with the same (emphasis on same) bogus data and checking if they are equal using isequal(). To cut to the chase, it should come out the same and now your original loop is down to:

tic
diff = zeros([n,n,d]);
for i = 1:n
    diff(i,:,:) = x(i,:) - x;
end
toc

For this final bit, you can exploit matlab's matrix/array reshaping/permuting functions. Look up the documentation for reshape() or permute() for more details. In brief, if you reshape or change the order of the dimensions of one copy of x from Nxd to 1xNxd, subtracting x from another, regularly sized matrix will perform the operations elementwise in matlab. So for example:

diff = permute(x,[1,3,2]) - permute(x,[3,1,2]); % this is Nx1xd - 1xNxd

should effectively compute the tensor difference you were looking for in the first loop!

I can expand this answer to show how the other loops might be worked out if you want, but give the other ones a try first with this same logic. Hopefully, you can keep diff and then calculate kxy much faster. Without knowing how big your original matrices are, I can't say how much speedup you should expect though.

Update:

I should add, to ensure that you are doing elementwise multiplication, division and transpose operations make sure to add a '.' before each command. E.g.

gradK(i,j,:) = -diff(i,j,:).*kxy(i, j)/h^2;

For more information, look up elementwise operations in Matlab

CodePudding user response：

If your goal is computation of obj you don't need even to compute gradK and a:

sx = sum(x.^2, 2);
H = sx - 2*x*x.'   sx.';

kxy = exp(-H/(2*h^2))/((2*pi*h^2)^(d/2));
kh = kxy / h^2; 
sumkxy = sum(kxy, 2);
khs = kh ./ sumkxy;
obj = khs * x - sum(khs, 2) .* x;