仅在遮罩区域中计算渐变

Question

我有一个非常大的阵列，只有几个感兴趣的小区域。 我需要计算这个数组的梯度，但出于性能原因，我需要将这个计算限制在这些感兴趣的领域。

我做不到这样的事情：

phi_grad0[mask] = np.gradient(phi[mask], axis=0)

由于花式索引的工作原理， phi[mask]只是成为被遮罩像素的一维数组，丢失了空间信息，使得渐变计算变得毫无价值。

np.gradient会处理np.ma.masked_array ，但性能会np.ma.masked_array一个数量级：

import numpy as np
from timeit_context import timeit_context

phi = np.random.randint(low=-100, high=100, size=[100, 100])
phi_mask = np.random.randint(low=0, high=2, size=phi.shape, dtype=np.bool)

with timeit_context('full array'):
    for i2 in range(1000):
        phi_masked_grad1 = np.gradient(phi)

with timeit_context('masked_array'):
    phi_masked = np.ma.masked_array(phi, ~phi_mask)
    for i1 in range(1000):
        phi_masked_grad2 = np.gradient(phi_masked)

这会产生以下输出：

[full array] finished in 143 ms
[masked_array] finished in 1961 ms

我认为它是因为在masked_array运行的操作没有矢量化，但我不确定。

有没有办法限制np.gradient以达到更好的性能？

这个timeit_context是一个方便的计时器，如果有人感兴趣的话：

from contextlib import contextmanager
import time

@contextmanager
def timeit_context(name):
    """
    Use it to time a specific code snippet
    Usage: 'with timeit_context('Testcase1'):'
    :param name: Name of the context
    """
    start_time = time.time()
    yield
    elapsed_time = time.time() - start_time
    print('[{}] finished in {} ms'.format(name, int(elapsed_time * 1000)))

Answer 1

不完全是答案，但这是我为我的情况设法修补的，这非常有效：

我得到条件为真的像素的1D索引（在这种情况下，条件< 5 ）：

def get_indices_1d(image, band_thickness):
    return np.where(image.reshape(-1) < 5)[0]

这给了我一个带有这些索引的数组。

然后我以不同的方式手动计算这些位置的梯度：

def gradient_at_points1(image, indices_1d):
    width = image.shape[1]
    size = image.size

    # Using this instead of ravel() is more likely to produce a view instead of a copy
    raveled_image = image.reshape(-1)

    res_x = 0.5 * (raveled_image[(indices_1d + 1) % size] - raveled_image[(indices_1d - 1) % size])
    res_y = 0.5 * (raveled_image[(indices_1d + width) % size] - raveled_image[(indices_1d - width) % size])

    return [res_y, res_x]


def gradient_at_points2(image, indices_1d):
    indices_2d = np.unravel_index(indices_1d, dims=image.shape)

    # Even without doing the actual deltas this is already slower, and we'll have to check boundary conditions, etc
    res_x = 0.5 * (image[indices_2d] - image[indices_2d])
    res_y = 0.5 * (image[indices_2d] - image[indices_2d])

    return [res_y, res_x]


def gradient_at_points3(image, indices_1d):
    width = image.shape[1]

    raveled_image = image.reshape(-1)

    res_x = 0.5 * (raveled_image.take(indices_1d + 1, mode='wrap') - raveled_image.take(indices_1d - 1, mode='wrap'))
    res_y = 0.5 * (raveled_image.take(indices_1d + width, mode='wrap') - raveled_image.take(indices_1d - width, mode='wrap'))

    return [res_y, res_x]


def gradient_at_points4(image, indices_1d):
    width = image.shape[1]

    raveled_image = image.ravel()

    res_x = 0.5 * (raveled_image.take(indices_1d + 1, mode='wrap') - raveled_image.take(indices_1d - 1, mode='wrap'))
    res_y = 0.5 * (raveled_image.take(indices_1d + width, mode='wrap') - raveled_image.take(indices_1d - width, mode='wrap'))

    return [res_y, res_x]

我的测试数组看起来像这样：

a = np.random.randint(-10, 10, size=[512, 512])

# Force edges to not pass the condition
a[:, 0] = 99
a[:, -1] = 99
a[0, :] = 99
a[-1, :] = 99

indices = get_indices_1d(a, 5)

mask = a < 5

然后我可以运行这些测试：

with timeit_context('full gradient'):
    for i in range(100):
        grad1 = np.gradient(a)

with timeit_context('With masked_array'):
    for im in range(100):
        ma = np.ma.masked_array(a, mask)
        grad6 = np.gradient(ma)

with timeit_context('gradient at points 1'):
    for i1 in range(100):
        grad2 = gradient_at_points1(image=a, indices_1d=indices)

with timeit_context('gradient at points 2'):
    for i2 in range(100):
        grad3 = gradient_at_points2(image=a, indices_1d=indices)

with timeit_context('gradient at points 3'):
    for i3 in range(100):
        grad4 = gradient_at_points3(image=a, indices_1d=indices)

with timeit_context('gradient at points 4'):
    for i4 in range(100):
        grad5 = gradient_at_points4(image=a, indices_1d=indices)

这给出了以下结果：

[full gradient] finished in 576 ms
[With masked_array] finished in 3455 ms
[gradient at points 1] finished in 421 ms
[gradient at points 2] finished in 451 ms
[gradient at points 3] finished in 112 ms
[gradient at points 4] finished in 102 ms

正如您所看到的，方法4是迄今为止最好的（但不关心它消耗多少内存）。

这可能只是因为我的2D数组相对较小（512x512）。 也许对于更大的阵列，这不是真的。

另一个需要注意的是， ndarray.take(indices, mode='wrap')会在图像边缘周围做一些奇怪的事情（一行将'循环'到下一个等等）以保持良好的性能，所以如果边缘对于您的应用程序可能需要在边缘周围填充1个像素的输入数组。

仍然非常有趣的是masked_array的速度有多慢。 在循环外部拉构造函数ma = np.ma.masked_array(a, mask)不会影响时间，因为masked_array本身只保留对数组及其掩码的引用