CODEWITHSUNDEEP
pythonnumpysumaddition
Finn Årup Nielsen
Finn Årup Nielsen

Reputation: 6726

"Kronecker addition" for indexing in Numpy

I am seeking a "Kronecker addition" facility for Python Numpy for indexing where elements are added instead of multiplied as in the Kronecker product. Not standard Kronecker sum nor direct sum.

This will do ok: 

def kron_add(A, B):
    C = np.kron(A, np.ones(np.asarray(B).shape)) + \
        np.kron(np.ones(np.asarray(A).shape), B)
    return C

>>> kron_add([0, 40, 80, 120], [0, 1, 2, 3])
array([   0.,    1.,    2.,    3.,   40.,   41.,   42.,   43.,   80.,
         81.,   82.,   83.,  120.,  121.,  122.,  123.])

And for a 1-dimensional problem this will do:

>>> [i+j for i in [0, 40, 80, 120] for j in [0, 1, 2, 3]]
[0, 1, 2, 3, 40, 41, 42, 43, 80, 81, 82, 83, 120, 121, 122, 123]

I suppose this last one is fairly readable and it solves my present issue. I was just wondering whether Numpy has a function or method that does this fairly general operation directly.

Upvotes: 1

Views: 5238

Answers (1)

hunse
hunse

Reputation: 3255

You can do this easily with Numpy broadcasting:

a = np.array([0, 40, 80, 120])
b = np.array([0, 1, 2, 3])
c = a[:, None] + b[None, :]

The statement a[:, None] is equivalent to a.reshape(1, -1), making a a column vector. Similarly, b[None, :] makes b a row vector. Numpy's broadcasting then takes over, and makes the resulting matrix c the required sum. If you want the result as a vector (as in your question) instead of a matrix, just call c.ravel().

You can also check out np.einsum, a really powerful tool that lets you express these kinds of computations in Einstein summation notation.

Upvotes: 3

Related Questions