Vectorizing scalar-array multiplication in Python (Numpy)

Question

I've created a 2D mesh grid w = np.meshgrid(x,y) which I'm trying to input into the following function:

def f(w):
    eigs = np.linalg.eigvals(A0 + w[0]*A1 + w[1]*A2)
    return abs(eigs[0] - eigs[-1])

where A0,A1,A2 are square arrays. But I'm getting an error which tells me that the operands cannot be broadcast together. Basically what's happening is that the w[0]*A1 is not vectorized and so w[0] is the entire block of x-values of the mesh instead of each individual x-value, same thing for w[1]*A2 but with the y-values.

I've tried doing np.vectorize(func) but that doesn't work and I get an IndexError.

hpaulj · Accepted Answer · 2015-01-16 04:22:28Z

1

Since w[0] is a (n,m) array, you can 'vectorize' by

def f(w):
    eigs = np.linalg.eigvals(A0+w[0][...,None,None]*A1 +
        w[1[...,None,None]*A2)
    return abs(eigs[...,0] - eigs[...,-1]

The result is the same shape as w[0].

This works because the inner function is linear in x and y, and eigvals accepts a (..., M, M) array_like input. In this case its input will be (n,m,M,M) shape.

answered Jan 16, 2015 at 4:22

hpaulj

233k14 gold badges260 silver badges392 bronze badges

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

1 Comment

Set Over a year ago

Thanks, this works, but damn I would have never figured this out on my own.

Collectives™ on Stack Overflow

Vectorizing scalar-array multiplication in Python (Numpy)

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related