Compute x**k with x, k being arrays of arbitrary dimensionality

Question

I have two numpy arrays: One array x with shape (n, a0, a1, ...) and one array k with shape (n, b0, b1, ...). I would like to compute and array of exponentials such that the output has dimension (a0, a1, ..., b0, b1, ...) and

out[i0, i1, ..., j0, j1, ...] == prod(x[:, i0, i1, ...] ** k[:, j0, j1, ...])

If there is only one a_i and one b_j, broadcasting does the trick via

import numpy

x = numpy.random.rand(2, 31)
k = numpy.random.randint(1, 10, size=(2, 101))

out = numpy.prod(x[..., None]**k[:, None], axis=0)

If x has a few dimensions more, more Nones have to be added:

x = numpy.random.rand(2, 31, 32, 33)
k = numpy.random.randint(1, 10, size=(2, 101))

out = numpy.prod(x[..., None]**k[:, None, None, None], axis=0)

If x has a few dimensions more, more Nones have to be added at other places:

x = numpy.random.rand(2, 31)
k = numpy.random.randint(1, 10, size=(2, 51, 51))

out = numpy.prod(x[..., None, None]**k[:, None], axis=0)

How to make the computation of out generic with respect to the dimensionality of the input arrays?

i want to know this tooooooo.......

nmlq
– nmlq

2017-10-11 13:27:03 +00:00
Commented Oct 11, 2017 at 13:27 — nmlq
– nmlq, Commented Oct 11, 2017 at 13:27

Divakar · Accepted Answer · 2017-10-11 13:50:21Z

5

Here's one using reshaping of the two arrays so that they are broadcastable against each other and then performing those operations and prod reduction along the first axis -

k0_shp = [k.shape[0]] + [1]*(x.ndim-1) + list(k.shape[1:])
x0_shp = list(x.shape) + [1]*(k.ndim-1)
out = (x.reshape(x0_shp) ** k.reshape(k0_shp)).prod(0)

Here's another way to reshape both inputs to 3D allowing one singleton dim per input and such that they are broadcastable against each other, perform prod reduction to get 2D array, then reshape back to multi-dim array -

s = x.shape[1:] + k.shape[1:] # output shape
out = (x.reshape(x.shape[0],-1,1)**k.reshape(k.shape[0],1,-1)).prod(0).reshape(s)

It must be noted that reshaping merely creates a view into the input array and as such is virtually free both memory-wise and performance-wise.

edited Oct 11, 2017 at 13:50

answered Oct 11, 2017 at 13:36

Divakar

222k19 gold badges274 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Aaron · Accepted Answer · 2017-10-11 15:43:23Z

4

Without understanding fully the math of what you're doing, it seems that you need a constant number of None's for the number of dimensions of each x and k.

does something like this work?

out = numpy.prod(x[[...]+[None]*(k.ndim-1)]**k[[slice(None)]+[None]*(x.ndim-1)])

Here are the slices separately so they're a bit easier to read:

x[  [...] + [None]*(k.ndim-1)  ]

k[  [slice(None)] + [None]*(x.ndim-1)  ]

Compatibility Note:

[...] seems to only be valid in python 3.x If you are using 2.7 (I haven't tested lower) substitute [Ellipsis] instead:

x[  [Ellipsis] + [None]*(k.ndim-1)  ]

edited Oct 11, 2017 at 15:43

answered Oct 11, 2017 at 13:37

Aaron

11.2k1 gold badge28 silver badges43 bronze badges

2 Comments

Nico Schlömer Over a year ago

[...] is only valid in Python 3. Perhaps there's an expression to make it compatible with Python 2 as well?

Aaron Over a year ago

@NicoSchlömer good catch, I made the switch to 3.6 a bit ago and tend to forget to test in 2.7 as well... see edit

Collectives™ on Stack Overflow

Compute x**k with x, k being arrays of arbitrary dimensionality

2 Answers 2

Comments

Compatibility Note:

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Compatibility Note:

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related