0

I have 2 arrays:

>>> a.shape
(9, 3, 11)
>>> b.shape
(9,)

I would like to compute the equivalent of c[i, j] = f(a[i, j, :], b[i]) where f(a0, b0) is a function that takes 2 parameters, with len(a0) == 11 and len(b0) == 9. Here, i is iterating on range(9) and j is iterating on range(3).

Is there a way to code this using numpy.vectorize? Or is it simpler with some clever broadcasting?

I have been trying for 2 hours and I just don't understand how to make it work... I tried to broadcast or to use signatures but to no avail.

  • 1
    `np.array([f(a[i,j,:], b[i]) for j in range(3)] for i in range(9)])`. `a + b[:,None,None]` will work but the result is (9,3,11). Or maybe `np.sum(a + b[:,None,None]).axis=2)` if you want (9,3), but then you might as well `np.sum(a, axis=2)` first. – hpaulj Nov 30 '21 at 08:04
  • 1
    `np.vectorize` in default mode passes scalar values to the function. There is a `signature` option that makes it pass arrays, but it is trickier to use, and even slower. `apply_along_axis` can iterate on your `i,j` dimensions of `a` (slowly), but it can't, at the same time iterate on `b`. It's just a one array iterator. The nested list comprehension is probably your best option. – hpaulj Nov 30 '21 at 08:10
  • 1
    A sample of the function might help. Also you need to make it clear what the function returns. Is it as scalar, so `c` can be a numeric dtype array. That requires, I assume, some sort of reduction on the size 11 dimension. – hpaulj Nov 30 '21 at 16:45
  • @hpaulj, in the end I could find a solution: `np.vectorize(f, signature="(k),(1)->()")` and then I have to call it like `f(a, b[:, None, None]`. Do you think this would be much slower than a loop? –  Nov 30 '21 at 17:06

2 Answers2

1

In the end, I could make it work like this:

>>> f = np.vectorize(f, signature="(k),(1)->()")
>>> print(a.shape)
(9, 3, 11)
>>> print(b.shape)
(9,)
>>> print(f(a, b[:, None, None]).shape)
(9, 3)

This ensures that f gets called with the correct shapes and iterates properly. It is frankly not straightforward from the Numpy documentation to understand the trick to use a (1) in the signature for this purpose.

0

numpy.apply_along_axis is what you need.

import numpy as np

a = np.ones( (9,3,11) )
b = np.ones( 9 )

def f(a0, b0):
    return sum(a0[:9]+b0)

c = np.apply_along_axis( f, 2, a, b )
print(c)

c's shape is (9,3).

Tim Roberts
  • 48,973
  • 4
  • 21
  • 30
  • 1
    I get a (9, 3, 9) shaped `c`. – hpaulj Nov 30 '21 at 08:11
  • 1
    OK, `f` needs to return one value. not a list of 9. Now `c` is (9,3). – Tim Roberts Nov 30 '21 at 17:52
  • You could do just `return np.sum(a0) + b0`, so we don't need the index in the function definition! The important thing is to return the scalar value. –  Nov 30 '21 at 19:20
  • 1
    That won't work, because `b0` is an array of 9. It will do a scalar addition and return a list, which is how I ended up with (9,3,9) in the first place. HOWEVER, I only made up that dummy function because you didn't provide yours. I presume yours is a bit more complicated. – Tim Roberts Nov 30 '21 at 19:25