Numpy добавить меньшую матрицу к большей

Question

У меня есть большие 3D-матрицы, указывающие положение агентов в 3D-пространстве. Значения матрицы равны 0, если на нем нет агента, и 1, если на нем есть агент. Тогда моя проблема заключается в том, что я хочу, чтобы агенты «росли» в том смысле, что я хочу, чтобы их определяли, скажем, куб (3x3x3) из них. Если уже есть способ сделать это, но у меня проблемы, когда агент находится близко к границам. Например, у меня есть матрица позиций 100x100x100, если я знаю, что мой агент находится в позиции (x, y, z), я сделаю:

positions_matrix = numpy.zeros((100, 100, 100))
positions_matrix[x - 1: x + 2, y - 1: y + 2, z - 1: z + 2] +=  numpy.ones((3, 3, 3))

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

Любая идея, как ее решить или если numpy или любой другой пакет имеет реализацию для этого? Мне не удалось найти его, хотя я уверен, что я не первый, кто сталкивается с этим.

Answer 1

Несколько более программный c способ решения проблемы:

import numpy as np


m = np.zeros((100, 100, 100))

slicing = tuple(
    slice(max(0, x_i - 1), min(x_i + 2, d - 1))
    for x_i, d in zip((x, y, z), m.shape))
ones_shape = tuple(s.stop - s.start for s in slicing)

m[slicing] += np.ones(ones_shape)

Но в остальном он совпадает с принятым ответом .

Answer 2

Существует решение, использующее np.put и его опцию 'clip'. Это просто требует небольшой гимнастики, потому что функция требует индексов в сплюснутой матрице; к счастью, функция np.ravel_multi_index выполняет свою работу:

import itertools
import numpy as np
x, y, z = 2, 0, 4
positions_matrix = np.zeros((100,100,100))

indices = np.array( list( itertools.product( (x-1, x, x+1), (y-1, y, y+1), (z-1, z, z+1)) ))
flat_indices = np.ravel_multi_index(indices.T, positions_matrix.shape, mode='clip')
positions_matrix.put(flat_indices, 1+positions_matrix.take(flat_indices))
# positions_matrix[2,1,4] is now 1.0

Приятно, что в этом решении вы можете играть с другими режимами, например 'wrap' (если ваши агенты живут на пончике; - ) или в промежутке c).

Я объясню, как это работает на двухмерной матрице меньшего размера:

import itertools
import numpy as np

positions_matrix = np.zeros((8,8))
ones = np.ones((3,3))
x, y = 0, 4

indices = np.array( list( itertools.product( (x-1, x, x+1), (y-1, y, y+1) )))
# array([[-1,  3],
#        [-1,  4],
#        [-1,  5],
#        [ 0,  3],
#        [ 0,  4],
#        [ 0,  5],
#        [ 1,  3],
#        [ 1,  4],
#        [ 1,  5]])

flat_indices = np.ravel_multi_index(indices.T, positions_matrix.shape, mode='clip')
# array([ 3,  4,  5,  3,  4,  5, 11, 12, 13])

positions_matrix.put(flat_indices, ones, mode='clip')
# positions_matrix is now:
# array([[0., 0., 0., 1., 1., 1., 0., 0.],
#        [0., 0., 0., 1., 1., 1., 0., 0.],
#        [0., 0., 0., 0., 0., 0., 0., 0.],
#        [ ...

Кстати, в данном случае mode='clip' было избыточно для put.

Ну, я просто обманул put выполняет задание. +=1 требует как take, так и put:

positions_matrix.put(flat_indices, ones.flat + positions_matrix.take(flat_indices))
# notice that ones has to be flattened, or alternatively the result of take could be reshaped (3,3)
# positions_matrix is now: 
# array([[0., 0., 0., 2., 2., 2., 0., 0.],
#        [0., 0., 0., 2., 2., 2., 0., 0.],
#        [0., 0., 0., 0., 0., 0., 0., 0.],
#        [ ...

В этом решении есть одно важное отличие по сравнению с другими: матрица ones всегда равна (3,3), что может или не может быть преимуществом. Трюк находится в этом списке flat_indices, в котором есть повторяющиеся записи (результат клипа).

Таким образом, могут потребоваться некоторые меры предосторожности, если вы добавите непостоянную подматрицу с максимальными индексами:

x, y = 1, 7
values = 1 + np.arange(9)
indices = np.array( list( itertools.product( (x-1, x, x+1), (y-1, y, y+1) )))
flat_indices = np.ravel_multi_index(indices.T, positions_matrix.shape, mode='clip')
positions_matrix.put(flat_indices, values, mode='clip')
# positions_matrix is now:
# array([[0., 0., 0., 2., 2., 2., 1., 3.],
#        [0., 0., 0., 2., 2., 2., 4., 6.],
#        [0., 0., 0., 0., 0., 0., 7., 9.],

... вы, вероятно, ожидали, что последний столбец будет 2 5 8. В настоящее время вы можете работать с flat_indices, например, поместив -1 в недопустимые места. Но было бы проще, если бы np.put принимал неплоские индексы или был клип mode='ignore'.

Answer 3

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

import numpy as np
m = np.zeros((100, 100, 100))

x_min, x_max = np.max([0, x-1]), np.min([x+2, m.shape[0]-1])
y_min, y_max = np.max([0, y-1]), np.min([y+2, m.shape[1]-1])
z_min, z_max = np.max([0, z-1]), np.min([z+2, m.shape[2]-1])

m[x_min:x_max, y_min:y_max, z_min:z_max] += np.ones((x_max-x_min, y_max-y_min, z_max-z_min))