[SOLVED] How to make a function work with different sized meshgrid/linspace inputs

How to make a function work with different sized meshgrid/linspace inputs

I have a a function called [U,V,W] = my_function(x,y,z) which takes grid vectors as inputs and calculates some outputs evaluated at those grid points. There are N grid points in total.

The problem is that I want my function to be tolerant to different sizes of inputs. For example,

THIS:
[Y,Z] = meshgrid(0:10,-10:20); X = 6.3;

OR THIS:
[X,Y,Z] = meshgrid(0:10,-10:20,-2:2);

OR THIS:
Z = linspace(-1,1,10); X = 4.2; Y = 1;
etc

But inside the function I need to convert the x y z variables to to a 3 x N array like this:

r = [X(:) Y(:) Z(:)]

(this is because I need to apply some rotation matrices, and then do some further manipulation on the rotated coordinates). This does not work if X,Y, and Z are different sizes.

I have tried to make a helper function which pre-processes the coordinates, as shown in the MWE below, but it is quite messy. I wonder if there is a more succinct / cleaner / more standard way to achieve this?

x = linspace(-10,10,100);
Y = 3;
z = linspace(-12,12,80);
[X,Z] = meshgrid(x,z);

[U,V,W] = my_function(X,Y,Z);




%%% MAIN FUNCTION %%%
function [U,V,W] = my_function(X,Y,Z)

[X,Y,Z] = preprocess_gridvectors(X,Y,Z);  % Make sure grid vectors are all same size
N1 = size(X,1); N2 = size(X,2); N3 = size(X,3);
X = X(:); Y = Y(:); Z = Z(:);

%%% Do some transformation and processing here (rotate coordinates and use them)
R = [cos(pi) -sin(pi) 0;sin(pi) cos(pi) 0;0 0 1]; rotXYZ = R*[X Y Z].';
U = rotXYZ(1,:); V = rotXYZ(2,:); W = rotXYZ(3,:);
%%%

% Put the outputs back into the same shape as the original inputs
U = reshape(U, N1,N2,N3);
V = reshape(V, N1,N2,N3);
W = reshape(W, N1,N2,N3);
end



%%% HELPER FUNCTION %%%
function [X,Y,Z] = preprocess_gridvectors(X,Y,Z)

unique_array_sizes = unique([size(X) size(Y) size(Z)]);  % Determine dimensions to use
unique_array_sizes(unique_array_sizes==1) = [];          % Disregard scalar dimension
N_expected = prod(unique_array_sizes);                   % Total number of unique grid points expected

sizes = zeros(3,3);
if numel(X)==N_expected     % X already has the right dimensions
    sizes(1,1) = size(X,1); % Get the dimensions of X
    sizes(1,2) = size(X,2);
    sizes(1,3) = size(X,3);
end

if numel(Y)==N_expected    % Y already has the right dimensions
    sizes(2,1) = size(Y,1);
    sizes(2,2) = size(Y,2);
    sizes(2,3) = size(Y,3);
end

if numel(Z)==N_expected    % Z already has the right dimensions
    sizes(3,1) = size(Z,1);
    sizes(3,2) = size(Z,2);
    sizes(3,3) = size(Z,3);
end

inds = sizes~=0; % Anything which is non-zero has correct dimensions already
ind_dir = find(inds(:,1),1); % Choose the first axis that has good dimensions

% These should be the correct final dimensions now
N1 = sizes(ind_dir,1);
N2 = sizes(ind_dir,2);
N3 = sizes(ind_dir,3);

if ~any(sizes(1,:)) % X direction needs replicating
    X = repmat(X,[N1 N2 N3]);
end
if ~any(sizes(2,:)) % Y direction needs replicating
    Y = repmat(Y,[N1 N2 N3]);
end
if ~any(sizes(3,:)) % Z direction needs replicating
    Z = repmat(Z,[N1 N2 N3]);
end

end

Solution

You can get implicit expansion to do the work for you:

X = X + 0*Y + 0*Z;
Y = 0*X + Y + 0*Z;
Z = 0*X + 0*Y + Z;

This uses some unnecessary arithmetical operations (i.e. multiplying by 0 and then adding), but the computation time required for that is probably negligible.