ShanZhengYang
ShanZhengYang

Reputation: 17641

Solve broadcasting error without for loop, speed up code

I may be misunderstanding how broadcasting works in Python, but I am still running into errors.

scipy offers a number of "special functions" which take in two arguments, in particular the eval_XX(n, x[,out]) functions. See http://docs.scipy.org/doc/scipy/reference/special.html

My program uses many orthogonal polynomials, so I must evaluate these polynomials at distinct points. Let's take the concrete example scipy.special.eval_hermite(n, x, out=None).

I would like the x argument to be a matrix shape (50, 50). Then, I would like to evaluate each entry of this matrix at a number of points. Let's define n to be an a numpy array narr = np.arange(10) (where we have imported numpy as np, i.e. import numpy as np).

So, calling

scipy.special.eval_hermite(narr, matrix)

should return Hermitian polynomials H_0(matrix), H_1(matrix), H_2(matrix), etc. Each H_X(matrix) is of the shape (50,50), the shape of the original input matrix.

Then, I would like to sum these values. So, I call

matrix1 = np.sum( [scipy.eval_hermite(narr, matrix)], axis=0 )

but I get a broadcasting error!

ValueError: operands could not be broadcast together with shapes (10,) (50,50)

I can solve this with a for loop, i.e.

matrix2 = np.sum( [scipy.eval_hermite(i, matrix) for i in narr], axis=0)

This gives me the correct answer, and the output matrix2.shape = (50,50). But using this for loop slows down my code, big time. Remember, we are working with entries of matrices.

Is there a way to do this without a for loop?

Upvotes: 1

Views: 289

Answers (2)

hpaulj
hpaulj

Reputation: 231540

The documentation for these functions is skimpy, and a lot of the code is compiled, so this is just based on experimentation:

special.eval_hermite(n, x, out=None)

n apparently is a scalar or array of integers. x can be an array of floats.

special.eval_hermite(np.ones(5,int)[:,None],np.ones(6)) gives me a (5,6) result. This is the same shape as what I'd get from np.ones(5,int)[:,None] * np.ones(6).

The np.ones(5,int)[:,None] is a (5,1) array, np.ones(6) a (6,), which for this purpose is equivalent of (1,6). Both can be expanded to (5,6).

So as best I can tell, broadcasting rules in these special functions is the same as for operators like *.


Since special.eval_hermite(nar[:,None,None], x) produces a (10,50,50), you just apply sum to axis 0 of that to produce the (50,50).

special.eval_hermite(nar[:,Nar,Nar], x).sum(axis=0)

Like I wrote before, the same broadcasting (and summing) rules apply for this hermite as they do for a basic operation like *.

Upvotes: 1

nneonneo
nneonneo

Reputation: 179552

eval_hermite broadcasts n with x, then evaluates Hn(x) at each point. Thus, the output shape will be the result of broadcasting n with x. So, if you want to make this work, you'll have to make n and x have compatible shapes:

import scipy.special as ss
import numpy as np
matrix = np.ones([100,100]) # example
narr = np.arange(10) # example
ss.eval_hermite(narr[:,None,None], matrix).shape # => (10, 100, 100)

But note that this might actually be faster:

out = np.zeros_like(matrix)
for n in narr:
    out += ss.eval_hermite(n, matrix)

In testing, it appears to be between 5-10% faster than np.sum(...) of above.

Upvotes: 1

Related Questions