如何让函数兼容不同尺寸的meshgrid/linspace输入？

我实现了一个函数[U,V,W] = my_function(x,y,z)，用于接收网格向量输入并计算对应N个网格点的输出。需要兼容多种不同尺寸的输入场景：

• 场景一：[Y,Z] = meshgrid(0:10,-10:20); X = 6.3;（单个标量+二维网格）
• 场景二：[X,Y,Z] = meshgrid(0:10,-10:20,-2:2);（三维完整网格）
• 场景三：Z = linspace(-1,1,10); X = 4.2; Y = 1;（两个标量+一维向量）

函数内部需要将x、y、z转换为3×N的数组r = [X(:) Y(:) Z(:)]，用于后续旋转矩阵等操作，但输入尺寸不一致时无法直接转换。我编写了一个预处理辅助函数，但实现过于繁琐，希望找到更简洁规范的实现方式。

原实现代码如下：

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);  % 确保网格向量尺寸一致
N1 = size(X,1); N2 = size(X,2); N3 = size(X,3);
X = X(:); Y = Y(:); Z = Z(:);

%%% 坐标变换与处理（示例：旋转坐标）
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,:);
%%%

% 将输出恢复为输入的原始形状
U = reshape(U, N1,N2,N3);
V = reshape(V, N1,N2,N3);
W = reshape(W, N1,N2,N3);
end

%%% 原辅助函数 %%%
function [X,Y,Z] = preprocess_gridvectors(X,Y,Z)

unique_array_sizes = unique([size(X) size(Y) size(Z)]);  % 确定目标维度
unique_array_sizes(unique_array_sizes==1) = [];          % 忽略标量维度
N_expected = prod(unique_array_sizes);                   % 预期的网格点总数

sizes = zeros(3,3);
if numel(X)==N_expected     % X已符合目标维度
    sizes(1,1) = size(X,1); % 获取X的维度
    sizes(1,2) = size(X,2);
    sizes(1,3) = size(X,3);
end

if numel(Y)==N_expected    % Y已符合目标维度
    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已符合目标维度
    sizes(3,1) = size(Z,1);
    sizes(3,2) = size(Z,2);
    sizes(3,3) = size(Z,3);
end

inds = sizes~=0; % 非零值对应正确维度
ind_dir = find(inds(:,1),1); % 选取第一个维度正确的输入

% 确定最终目标维度
N1 = sizes(ind_dir,1);
N2 = sizes(ind_dir,2);
N3 = sizes(ind_dir,3);

if ~any(sizes(1,:)) % 需要扩展X维度
    X = repmat(X,[N1 N2 N3]);
end
if ~any(sizes(2,:)) % 需要扩展Y维度
    Y = repmat(Y,[N1 N2 N3]);
end
if ~any(sizes(3,:)) % 需要扩展Z维度
    Z = repmat(Z,[N1 N2 N3]);
end

end

简化后的实现方案 #

可以利用MATLAB的维度扩展特性，通过更简洁的逻辑实现预处理：

简化版辅助函数

function [X,Y,Z] = preprocess_gridvectors(X,Y,Z)
    % 获取每个输入的尺寸
    sz = {size(X), size(Y), size(Z)};
    % 计算目标尺寸：每个维度取三个输入中的最大值
    targetSz = max([sz{:}], [], 1);
    
    % 对每个变量，按目标尺寸进行扩展
    X = repmat(X, targetSz ./ sz{1});
    Y = repmat(Y, targetSz ./ sz{2});
    Z = repmat(Z, targetSz ./ sz{3});
end

进一步简化主函数

还可以省略单独记录N1/N2/N3的步骤，直接用预处理后的尺寸：

function [U,V,W] = my_function(X,Y,Z)
    [X,Y,Z] = preprocess_gridvectors(X,Y,Z);
    origShape = size(X); % 记录原始网格形状
    
    % 转换为列向量并处理
    coords = [X(:), Y(:), Z(:)];
    R = [cos(pi) -sin(pi) 0;sin(pi) cos(pi) 0;0 0 1];
    rotCoords = R * coords.';
    
    % 恢复形状并输出
    U = reshape(rotCoords(1,:), origShape);
    V = reshape(rotCoords(2,:), origShape);
    W = reshape(rotCoords(3,:), origShape);
end

方案说明 #

1. 目标尺寸计算：通过取三个输入每个维度的最大值，自动适配所有输入场景，无需手动判断哪个输入符合目标维度。
2. 简洁扩展逻辑：利用targetSz ./ sz{1}自动计算每个维度需要重复的次数，无论是标量、一维向量还是多维网格，都能正确扩展到目标尺寸。
3. 主函数精简：直接记录预处理后的形状，避免单独存储三个维度变量，代码更紧凑。

内容的提问来源于stack exchange，提问作者teeeeee