在 Matlab 中优化“掩码”功能

Question

为了进行基准比较，我考虑了简单的函数：

function dealiasing2d(where_dealiased, data)
[n1, n0, nk] = size(data);
for i0=1:n0 
    for i1=1:n1
        if where_dealiased(i1, i0)
            data(i1, i0, :) = 0.;
        end
    end
end

它在伪光谱模拟中很有用（其中data 是复数的 3d 数组），但基本上它将掩码应用于一组图像，将where_dealiased 为真的一些元素置零。

我在这个简单的案例中比较了不同语言（以及实现、编译器等）的性能。对于 Matlab，我使用 timeit 对函数计时。由于我不想在 Matlab 中衡量我的无知，我想用这种语言真正优化这个功能。在 Matlab 中最快的方法是什么？

我现在使用的简单解决方案是：

function dealiasing2d(where_dealiased, data)
[n1, n0, nk] = size(data);
N = n0*n1;
ind_zeros = find(reshape(where_dealiased, 1, []));
for ik=1:nk
    data(ind_zeros + N*(ik-1)) = 0;
end

我怀疑这不是正确的做法，因为等效的 Numpy 解决方案大约快 10 倍。

import numpy as np

def dealiasing(where, data):
    nk = data.shape[0]
    N = reduce(lambda x, y: x*y, data.shape[1:])
    inds, = np.nonzero(where.flat)
    for ik in xrange(nk):
        data.flat[inds + N*ik] = 0.

最后，如果有人告诉我类似“当你想在 Matlab 中快速使用特定函数时，你应该像这样编译它：[...]”，我会在基准测试中包含这样的解决方案。

---

编辑：

在 2 个答案之后，我对这些命题进行了基准测试，似乎没有明显的性能改进。这很奇怪，因为简单的 Python-Numpy 解决方案确实（一个数量级）快得多，所以我仍在寻找使用 Matlab 的更好解决方案......

Answer 1

如果我理解正确，这可以通过bsxfun 轻松快速地完成：

data = bsxfun(@times, data, ~where_dealiased);

这会将where_dealiased 为true 的条目的所有三维组件设置为0（将它们乘以0），并保持其余部分不变（将它们乘以@ 987654329@)。

当然，这是假设[size(data,1) size(data,2]==size(where_dealiased)。

---

您使用linear indexing 的解决方案可能也非常快。为了节省一些时间，您可以删除reshape，因为find 已经返回线性索引：

ind_zeros = find(where_dealiased);

Answer 2

方法 #1：Logical indexing 与 repmat -

data(repmat(where_dealiased,1,1,size(data,3))) = 0;

方法 #2：Linear indexing 和 bsxfun(@plus -

[m,n,r] =  size(data);
idx = bsxfun(@plus,find(where_dealiased),[0:r-1]*m*n); %// linear indices
data(idx) = 0;

如果where_dealiased 中的非零元素很少，这应该很快。

Answer 3

没有基准就没有优化！所以这里有一些建议的解决方案和性能测量。初始化代码为：

N = 2000;
nk = 10;

where = false([N, N]);
where(1:100, 1:100) = 1;
data = (5.+j)*ones([N, N, nk]);

我使用函数timeit 对函数进行计时，如下所示：

timeit(@() dealiasing2d(where, data))

作为比较，当我对问题中给出的 Numpy 函数执行完全相同的操作时，它的运行时间为 0.0167 秒。

具有 2 个循环的初始 Matlab 函数运行时间约为 0.34 秒，等效的 Numpy 函数（具有 2 个循环）运行速度较慢，运行时间为 0.42 秒。可能是因为 Matlab 使用了 JIT 编译。

Luis Mendo mentions 我可以删除reshape，因为find 已经返回线性索引。我喜欢它，因为代码更简洁，但 reshape 无论如何都非常便宜，所以它并没有真正提高函数的性能：

function dealiasing2d(where, data)
[n1, n0, nk] = size(data);
N = n0*n1;
ind_zeros = find(where);
for ik=1:nk
    data(ind_zeros + N*(ik-1)) = 0;
end

这个函数需要 0.23 秒，比使用 2 个循环的解决方案快，但与 Numpy 解决方案相比确实很慢（慢了约 14 倍！）。这就是我写问题的原因。

Luis Mendo also proposes 一个基于函数bsxfun 的解决方案，它给出：

function dealiasing2d_bsxfun(where, data)
data = bsxfun(@times, data, ~where);

这个解决方案涉及N*N*nk 乘法（乘以 1 或 0），这显然是太多的工作，因为我们只需要将数组 data 中的 100*100*nk 值归零。但是，这些乘法可以向量化，因此与其他 Matlab 解决方案相比“相当快”：0.23 秒，即与使用find的第一个解决方案相同！

Both solutions proposed by Divakar 涉及创建一个大小为N*N*nk 的大型数组。没有 Matlab 循环，所以我们可以期待更好的性能，但是......

function dealiasing2d_bsxfun2(where, data)
[n1, n0, nk] = size(data);
idx = bsxfun(@plus, find(where), [0:nk-1]*n1*n0);
data(idx) = 0;

需要 0.23 秒（仍然与其他功能相同！）和

function dealiasing2d(where, data)
data(repmat(where,[1,1,size(data,3)])) = 0;

耗时 0.30 秒（比其他 Matlab 解决方案多出约 20%）。

总而言之，在这种情况下，似乎有一些东西限制了 Matlab 的性能。也可能是 Matlab 中有更好的解决方案，或者我在基准测试中做错了什么......如果有 Matlab 和 Python-Numpy 的人可以提供其他时间，那就太好了。

---

编辑：

更多关于 Divakar 评论的数据：

对于 N = 500 ; nk = 500：

Method          | time (s) | normalized      
----------------|----------|------------
Numpy           |    0.05  |     1.0
Numpy loop      |    0.05  |     1.0
Matlab bsxfun   |    0.70  |    14.0
Matlab find     |    0.75  |    15.0
Matlab bsxfun2  |    0.76  |    15.2
Matlab loop     |    0.77  |    15.4
Matlab repmat   |    0.96  |    19.2

对于 N = 500 ; nk = 100：

Method          | time (s) | normalized      
----------------|----------|------------
Numpy           |    0.01  |     1.0
Numpy loop      |    0.03  |     3.0
Matlab bsxfun   |    0.14  |    12.7
Matlab find     |    0.15  |    13.6
Matlab bsxfun2  |    0.16  |    14.5
Matlab loop     |    0.16  |    14.5
Matlab repmat   |    0.20  |    18.2

对于 N = 2000 ; nk = 10：

Method          | time (s) | normalized |     
----------------|----------|------------|
Numpy           |    0.02  |     1.0    |
Matlab find     |    0.23  |    13.8    |
Matlab bsxfun2  |    0.23  |    13.8    |
Matlab bsxfun   |    0.24  |    14.4    |
Matlab repmat   |    0.30  |    18.0    |
Matlab loop     |    0.34  |    20.4    |
Numpy loop      |    0.42  |    25.1    |

我真的很想知道为什么与 Numpy 相比，Matlab 看起来这么慢......