CODEWITHSUNDEEP
pythonnumpy
aarcas
aarcas

Reputation: 309

Numpy add smaller matrix to a bigger one

I have big 3D matrices indicating the position of agents in a 3D space. The values of the matrix are 0 if there is not agent on it and 1 if there is an agent on it. Then, my problem is that I want the agents to 'grow' in the sense that I want them to be determined by lets say a cube (3x3x3) of ones. If already gotten a way to do it but I'm having trouble when the agent is close to the borders. For example, I have a matrix of positions 100x100x100, if I know my agent is at position (x, y, z) I will do:

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))

Of course in my real code I'm looping over more positions but this is basically it. This works but the problem comes when the agent is to close to the border in which the sum can't be made because the resultant matrix from slicing would be smaller than the ones matrix.

Any idea how to solve it or if numpy or any other package have an implementation for this? I couldn't manage to find it although I'm pretty sure I'm not the first one to face against this.

Upvotes: 1

Views: 570

Answers (3)

Demi-Lune
Demi-Lune

Reputation: 1977

There is a solution using np.put, and its 'clip' option. It just requires a little gymnastics because the function requires indices in the flattened matrix; fortunately, the function np.ravel_multi_index does the job:

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

The nice thing about this solution is that you can play with other modes, for instance 'wrap' (if your agents live on a donut ;-) or in a periodic space).

I'll explain how it works on a smaller 2D matrix:

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.],
#        [ ...

By the way, in this case mode='clip' was redundant for put.

Well, I just cheated put does an assignment. The +=1 requires both take and 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.],
#        [ ...

There is one important difference in this solution compared to the others: the ones matrix is always (3,3), which may or may not be an advantage. The trick is in this flat_indices list, that has repeating entries (result of clip).

It may thus require some precautions, if you add a non constant sub-matrix at max 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.],

... you were probably expecting the last column to be 2 5 8. Currently, you could work on flat_indices, for example by putting -1 in the out-of-bounds locations. But it'd all be easier if np.put accepted non-flat indices, or if there was a clip mode='ignore'.

Upvotes: 0

Idea O.
Idea O.

Reputation: 480

A slightly more programmatic way of solving the problem:

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)

But it is otherwise the same as the accepted answer.

Upvotes: 1

FBruzzesi
FBruzzesi

Reputation: 6495

You should cut at the lower and upper bounds, using something like:

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))

Upvotes: 0

Related Questions