Эффективный поиск по окрестностям в numpy ndarray вместо вложенных условных циклов for

Постановка задачи

У вас есть твердая форма в большом массиве. Из него вырезаешь мяч. Ваша цель — найти индексы поверхности твердого тела внутри шара. Поверхности определяются как любая точка, примыкающая к внешней стороне твердого тела с 6-точечной связностью. Ребра массива также считаются поверхностями.

Решение для более быстрого цикла

Вы уже вычислили маску, представляющую пересечение твердого тела и шара. Вы можете вычислить маску немного более элегантно и вместо этого преобразовать ее в индексы. Я предлагаю сохранить порядок ваших размеров постоянным, вместо того, чтобы переключаться между различными обозначениями. Например, это влияет на порядок RN, и вы рискуете не соответствовать ограничениям оси.

RN = np.array([4, 4, 3])
rad = 3

im = ...

cutout = ((np.indices(im.shape) - RN.reshape(-1, 1, 1, 1))**2).sum(axis=0) <= rad**2
solid = im > 0
mask = solid & cutout
indices = np.argwhere(mask)

Вы также можете получить вырез без изменения формы RN, выполнив

cutout = ((np.rollaxis(np.indices(im.shape, sparse=False), 0, 4) - RN)**2).sum(axis=-1) <= rad**2

Хорошая вещь в вычислении индексов заключается в том, что ваши циклы больше не должны быть огромными. Используя argwhere, вы в основном удаляете три внешних цикла, оставляя только оператор if для цикла. Вы также можете векторизовать проверку подключения. У этого есть приятный побочный эффект: вы можете определить произвольное соединение для каждого пикселя.

limit = np.array(im.shape) - 1  # Edge of `im`
connectivity = np.array([[ 1,  0,  0],  # Add rows to determine connectivity
                         [-1,  0,  0],
                         [ 0,  1,  0],
                         [ 0, -1,  0],
                         [ 0,  0,  1],
                         [ 0,  0, -1]], dtype=indices.dtype)
index_mask = np.ones(len(indices), dtype=bool)

for n, ind in enumerate(indices):
    if ind.all() and (ind < limit).all() and im[tuple((ind + connectivity).T)].all():
        index_mask[n] = False

LA = indices[index_mask, :]

Обратите внимание, что в imtemp нет никакого смысла. Даже в исходном цикле вы могли просто манипулировать mask напрямую. Вместо того, чтобы устанавливать для элементов значение -2, когда они соответствуют вашему критерию, вы можете установить для элементов значение False, если они этого не сделали.

Я делаю что-то подобное здесь. Мы проверяем каждый из фактически выбранных индексов и определяем, находится ли какой-либо из них внутри тела. Эти индексы исключаются из маски. Затем список индексов обновляется на основе маски.

Проверка ind.all() and (ind < limit).all() and im[tuple((ind + connectivity).T)].all() — это сокращение того, что вы делали с условиями or, но в обратном порядке (тестирование не на поверхности, а на поверхности).

• ind.all() проверяет, что ни один из индексов не равен нулю: т. е. не на верхней/передней/левой поверхности.
• (ind < limit).all() проверяет, что ни один из индексов не равен соответствующему размеру изображения минус один.
• im[tuple((ind + connectivity).T)].all() проверяет, что ни один из подключенных пикселей не равен нулю. (ind + connectivity).T — это массив (3, 6) из шести точек, к которым мы подключены (в настоящее время определяется как +/-1 по каждой оси массивом (6, 3) connectivity). Когда вы превращаете его в кортеж, он просто становится причудливым индексом, как если бы вы сделали что-то вроде im[x + connectivity[:, 0], y + connectivity[:, 1], z + connectivity[:, 2]]. Запятые в индексе просто превращают его в кортеж. Способ, который я показываю, лучше подходит для произвольного количества измерений.

Пиксели, прошедшие все три теста, находятся внутри твердого тела и удаляются. Конечно, вы можете написать цикл для проверки по-другому, но тогда вам придется изменить свою маску:

index_mask = np.zeros(len(indices), dtype=bool)

for n, ind in enumerate(indices):
    if (ind == 0).any() or (ind == limit).any() or (im[tuple((ind + connectivity).T)] == 0).any():
        index_mask[n] = True

LA = indices[index_mask, :]

Зацикливание вручную ни в коем случае не является идеальным. Тем не менее, он показывает вам, как сократить цикл (вероятно, на пару порядков) и как определить произвольное соединение с помощью векторизации и широковещания, не увязая в жестком кодировании.

Полностью векторизованное решение

Циклы выше можно полностью векторизовать, используя магию вещания. Вместо того, чтобы перебирать каждую строку в indices, мы можем добавить к ней connectivity массово и фильтровать результаты массово. Хитрость заключается в том, чтобы добавить достаточно измерений, чтобы добавить все connectivity к каждому элементу indices.

Вы все равно захотите опустить пиксели по краям:

edges = (indices == 0).any(axis=-1) | (indices == limit).any(axis=-1)
conn_index = indices[~edges, None, :] + connectivity[None, ...]

index_mask = np.empty(len(indices), dtype=bool)
index_mask[edges] = True
index_mask[~edges] = (im[tuple(conn_index.T)] == 0).any(axis=0)

LA = indices[index_mask, :]

Я ожидаю, что правильно написанный цикл, скомпилированный с помощью numba, будет значительно быстрее, чем это решение, потому что это позволит избежать большей части накладных расходов при конвейерной обработке операций. Это не потребует больших временных буферов или специальной обработки.

TL;DR

# Parameters
RN = np.array([4, 4, 3])
rad = 3

im = ...

# Find subset of interest
cutout = ((np.indices(im.shape) - RN.reshape(-1, 1, 1, 1))**2).sum(axis=0) <= rad**2
solid = im > 0

# Convert mask to indices
indices = np.argwhere(solid & cutout)

# Find image edges among indices
edges = (indices == 0).any(axis=-1) | (indices == limit).any(axis=-1)

# Connectivity elements for non-edge pixels
conn_index = indices[~edges, None, :] + connectivity[None, ...]

# Mask the valid surface pixels
index_mask = np.empty(len(indices), dtype=bool)
index_mask[edges] = True
index_mask[~edges] = (im[tuple(conn_index.T)] == 0).any(axis=0)

# Final result
LA = indices[index_mask, :]

Поскольку ваши циклы используют только прямое индексирование Numpy, вы можете использовать Numba @njit для выполнения этого гораздо более эффективным способом.

@njit
def compute_imtemp(imtemp, X, Y, Z):
    for i in range (Z):
        for j in range (Y-1):
            for k in range (X-1):
                if imtemp[i,j,k]==-1:
                    if i==(Z-1): 
                        imtemp[i,j,k]=-2
                    elif imtemp[i+1,j,k] == 0 or imtemp[i-1,j,k] == 0 or imtemp[i,j+1,k] == 0 or imtemp[i,j-1,k] == 0 or imtemp[i,j,k+1] == 0 or imtemp[i,j,k-1] == 0:
                        imtemp[i,j,k]=-2

[...]
imtemp = im.copy()
imtemp[mask] = -1
compute_imtemp(imtemp, X, Y, Z)
LA = np.argwhere(imtemp==-2)

Вот результаты производительности на моей машине:

281 µs ± 1.43 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
776 ns ± 16.4 ns per loop (mean ± std. dev. of 7 runs, 100 loops each)

Реализация Numba в 362 раза быстрее.

Обратите внимание, что первый вызов compute_imtemp будет медленным из-за компиляции. Один из способов преодолеть это — вызвать compute_imtemp для пустого массива Numpy. Другой способ — вручную скомпилировать функцию с помощью Numba API и явно указать типы для Numba.