CODEWITHSUNDEEP
pythonnumpy
Ofer Sadan
Ofer Sadan

Reputation: 11922

numpy broadcasting to all dimensions

I have a 3d numpy array build like this:

a = np.ones((3,3,3))

And I would like to broadcast values on all dimensions starting from a certain point with given coordinates, but the number of dimensions may vary.

For example if i'm given the coordinates (1,1,1) I can do these 3 functions:

a[1,1,:] = 0
a[1,:,1] = 0
a[:,1,1] = 0

And the result will be my desired output which is: 

array([[[1., 1., 1.],
        [1., 0., 1.],
        [1., 1., 1.]],

       [[1., 0., 1.],
        [0., 0., 0.],
        [1., 0., 1.]],

       [[1., 1., 1.],
        [1., 0., 1.],
        [1., 1., 1.]]])

Or if i'm given the coordinates (0,1,0) the corresponding broadcast will be:

a[0,1,:] = 0
a[0,:,0] = 0
a[:,1,0] = 0

Is there any way to do this in a single action instead of 3? I'm asking because the actual arrays i'm working with have even more dimensions which makes the code seem long and redundant. Also if the number of dimensions change I would have to rewrite the code.

EDIT: It doesn't have to be a single action, I just need to do it in all dimensions programatically such that if the number of dimensions change the code will stay the same.

EDIT 2: About the logic of this, i'm not sure if that's relevant, but i'm being given the value of a point (by coordinates) on a map and based on that I know the values of the entire row, column and height on the same map (that's why i'm updating all 3 with 0 as an example). In other cases the map is 2-dimensions and I still know the same thing about the row and column, but can't figure out a function that works for a varied numbers of dimensions.

Upvotes: 3

Views: 388

Answers (3)

gboffi
gboffi

Reputation: 25023

What you want to do can be done programmatically using a slice object that is instantiated using the slice function — e.g., see " Dealing with variable numbers of indices within programs".

def broadcast_val(mat, val, indices):
    shape = mat.shape
    for i in len(indices):
        len_axis = shape[i]
        i_axis = slice(len_axis)
        index_tuple = *indices[:i], i_axis, *indices[i+1:]
        mat[index_tuple] = val

I wonder if it is possible to do what you want using a single assignment...

Upvotes: 2

Tai
Tai

Reputation: 7984

This method builds the slices with np.repeat and np.fill_diagonal.

a = np.ones((3,3,3))
target = (0, 1, 0)
slices = np.repeat([target], a.ndim, axis=0).astype(object)
np.fill_diagonal(slices, slice(None))

# slices
# array([[slice(None, None, None), 1, 0],
#       [0, slice(None, None, None), 0],
#       [0, 1, slice(None, None, None)]], dtype=object)

for t in slices:
    a[tuple(t)] = 0

array([[[0., 1., 1.],
        [0., 0., 0.],
        [0., 1., 1.]],

       [[1., 1., 1.],
        [0., 1., 1.],
        [1., 1., 1.]],

       [[1., 1., 1.],
        [0., 1., 1.],
        [1., 1., 1.]]])

np.repeat([target], a.ndim, axis=0) repeats (0, 1, 0) for a.ndim times and it gives 

array([[0, 1, 0],
       [0, 1, 0],
       [0, 1, 0]])

Then, we want to change part of each row to slice(None). We cannot do it directly as the type of slice object is different from int. Thus, we use astype(object) to change the type of the array.

Upvotes: 2

Ben Schmidt
Ben Schmidt

Reputation: 401

Here's a way to generate string of exactly the 3 lines of code you're currently using, and then execute them:

import numpy as np

a = np.ones([3,3,3])
coord = [1, 1, 1]

for i in range(len(coord)):
   temp = coord[:]
   temp[i] = ':'
   slice_str = ','.join(map(str, temp))
   exec("a[%s] = 0"%slice_str)

print a

This may not be the best approach, but at least it's amusing. Now that we know that it works, we can go out and find the appropriate syntax to do it without actually generating the string and execing it. For example, you could use slice:

import numpy as np

a = np.ones([3,3,3])
coord = [1, 1, 1]

for i, length in enumerate(a.shape):
   temp = coord[:]
   temp[i] = slice(length)
   a[temp] = 0
print a

Upvotes: 1

Related Questions