A common way to multiply numpy array of arrays or 2D array by a vector

Question

I want to get a result of ax = as_x-aseg_3d*aline_a

where as_x and aline_a are numpy vectors, and aseg_3d is an array of arrays. as_x.shape and aline_a shape are (S,), and aseg_3d shape is also shown as (S,) as it should be for array of arrays. The result ax is also an array of arrays with the same shape as aseg_3d.

The problem is, that aseg_3d is an array of arrays because the arrays inside have different lengths. However, sometimes they don't and than aseg_3d becomes just a 2D array. The mutliplication above then fails on ValueError: operands could not be broadcast together with shapes (S,N).

I can fix it by doing: ax = as_x[:,None]-aseg_3d*aline_a[:,None]

but this does not work in the same way for the previous case when I have an array of arrays.

Is there a way to make it work in both cases: array of arrays and 2D array?

Or perhaps there is a way to keep a numpy array of arrays dtype="object" even if all the inner arrays have the same length?

Btw. I was trying to do the above multiplication with list comprehensions, which would always work, but it is much slower.

Thank you for all advice!

Answer 1

Expression that works in either case:

You can utilize the transpose. Simple case (just multiplication):

(vector * vec_of_vecs_or_2D.T).T

Full OP example:

(vector1 - vector2 * vec_of_vecs_or_2D.T).T

Example:

>>> v = np.array([3, 5])
>>> A = np.array((*map(np.array,[[1,2],[1,2,3]]),))
>>> B = np.array([[1,2,4],[1,2,3]])
>>> 
>>> v
array([3, 5])
>>> A
array([array([1, 2]), array([1, 2, 3])], dtype=object)
>>> B
array([[1, 2, 4],
       [1, 2, 3]])
>>> 
>>> (v*B.T).T
array([[ 3,  6, 12],
       [ 5, 10, 15]])
>>> (v*A.T).T
array([array([3, 6]), array([ 5, 10, 15])], dtype=object)

How it works:

If vec_of_vecs_or_2D is 1D of object dtype then the transpose does nothing.

If vec_of_vecs_or_2D is genuinely 2D then the transpose moves the axis matching vector to the end, so vector is properly broadcast, and back again afterwards.

Why it's cheap:

Please note that the transpose in numpy is lazy. It does not actually move any data, it just "swaps the axis labels".

For example taking the transpose of a 1000x1000 array takes ~200 ns on my laptop:

>>> a = np.arange(1_000_000).reshape(1000, 1000)
>>> timeit(lambda: a.T)
0.22008929261937737

Alternatively, force array of arrays in any case:

Example

>>> L1 = [np.arange(i, 2*i+2) for i in range(3)]
>>> L2 = [np.arange(i, i+3) for i in range(3)]
>>> L1
[array([0, 1]), array([1, 2, 3]), array([2, 3, 4, 5])]
>>> L2
[array([0, 1, 2]), array([1, 2, 3]), array([2, 3, 4])]

Method 1 (ugly but simple):

>>> A1 = np.array([*L1, None])[:-1]
>>> A2 = np.array([*L2, None])[:-1]
>>> A1
array([array([0, 1]), array([1, 2, 3]), array([2, 3, 4, 5])], dtype=object)
>>> A2
array([array([0, 1, 2]), array([1, 2, 3]), array([2, 3, 4])], dtype=object)

Method 2: (not quite as ugly but complicated)

>>> A1 = np.frompyfunc(L1.__getitem__, 1, 1)(range(len(L1)))
>>> A2 = np.frompyfunc(L2.__getitem__, 1, 1)(range(len(L2)))
>>> A1
array([array([0, 1]), array([1, 2, 3]), array([2, 3, 4, 5])], dtype=object)
>>> A2
array([array([0, 1, 2]), array([1, 2, 3]), array([2, 3, 4])], dtype=object)