CODEWITHSUNDEEP
pythonnumpy
nneonneo
nneonneo

Reputation: 179392

numpy broadcast from first dimension

In NumPy, is there an easy way to broadcast two arrays of dimensions e.g. (x,y) and (x,y,z)? NumPy broadcasting typically matches dimensions from the last dimension, so usual broadcasting will not work (it would require the first array to have dimension (y,z)).

Background: I'm working with images, some of which are RGB (shape (h,w,3)) and some of which are grayscale (shape (h,w)). I generate alpha masks of shape (h,w), and I want to apply the mask to the image via mask * im. This doesn't work because of the above-mentioned problem, so I end up having to do e.g. 

mask = mask.reshape(mask.shape + (1,) * (len(im.shape) - len(mask.shape)))

which is ugly. Other parts of the code do operations with vectors and matrices, which also run into the same issue: it fails trying to execute m + v where m has shape (x,y) and v has shape (x,). It's possible to use e.g. atleast_3d, but then I have to remember how many dimensions I actually wanted.

Upvotes: 19

Views: 4704

Answers (6)

user3208430
user3208430

Reputation: 676

Easy, you just add a singleton dimension at the end of the smaller array. For example, if xyz_array has shape (x,y,z) and xy_array has shape (x,y), you can do

xyz_array + np.expand_dims(xy_array, xy_array.ndim)

Upvotes: 1

hpaulj
hpaulj

Reputation: 231335

numpy functions often have blocks of code that check dimensions, reshape arrays into compatible shapes, all before getting down to the core business of adding or multiplying. They may reshape the output to match the inputs. So there is nothing wrong with rolling your own that do similar manipulations.

Don't offhand dismiss the idea of rotating the variable 3 dimension to the start of the dimensions. Doing so takes advantage of the fact that numpy automatically adds dimensions at the start.

For element by element multiplication, einsum is quite powerful.

np.einsum('ij...,ij...->ij...',im,mask)

will handle cases where im and mask are any mix of 2 or 3 dimensions (assuming the 1st 2 are always compatible. Unfortunately this does not generalize to addition or other operations.

A while back I simulated einsum with a pure Python version. For that I used np.lib.stride_tricks.as_strided and np.nditer. Look into those functions if you want more power in mixing and matching dimensions.

Upvotes: 5

ev-br
ev-br

Reputation: 26030

Why not just decorate-process-undecorate:

def flipflop(func):
    def wrapper(a, mask):
        if len(a.shape) == 3:
            mask = mask[..., None]
        b = func(a, mask)
        return np.squeeze(b)
    return wrapper

@flipflop
def f(x, mask):
    return x * mask

Then

>>> N = 12
>>> gs = np.random.random((N, N))
>>> rgb = np.random.random((N, N, 3))
>>> 
>>> mask = np.ones((N, N))
>>> 
>>> f(gs, mask).shape
(12, 12)
>>> f(rgb, mask).shape
(12, 12, 3)

Upvotes: 1

HYRY
HYRY

Reputation: 97261

how about use transpose:

(a.T + c.T).T

Upvotes: 15

Eelco Hoogendoorn
Eelco Hoogendoorn

Reputation: 10759

as another angle: if you encounter this pattern frequently, it may be useful to create a utility function to enforce right-broadcasting:

def right_broadcasting(arr, target):
    return arr.reshape(arr.shape + (1,) * (target.ndim - arr.ndim))

Although if there are only two types of input (already having 3 dims or having only 2), id say the single if statement is preferable.

Upvotes: 3

Eelco Hoogendoorn
Eelco Hoogendoorn

Reputation: 10759

Indexing with np.newaxis creates a new axis in that place. Ie

xyz = #some 3d array
xy = #some 2d array
xyz_sum = xyz + xy[:,:,np.newaxis]
or
xyz_sum = xyz + xy[:,:,None]

Indexing in this way creates an axis with shape 1 and stride 0 in this location.

Upvotes: 1

Related Questions