使用生成器代替嵌套循环

Question

I have the following nested loop. But it is inefficient time wise. So using a generator would be much better. Do you know how to do that?

x_sph[:] = [r*sin_t*cos_p for cos_p in cos_phi for sin_t in sin_theta for r in p]     

It seems like some of you are of the opinion (looking at comments) that using a generator was not helpful in this case. I am under the impression that using generators will avoid assigning variables to memory, and thus save memory and time. Am I wrong?

Answer 1

Judging from your code snippet you want to do something numerical and you want to do it fast. A generator won't help much in this respect. But using the numpy module will. Do it like so:

import numpy
# Change your p into an array, you'll see why.
r = numpy.array(p) # If p is a list this will change it into 1 dimensional vector.
sin_theta = numpy.array(sin_theta) # Same with the rest.
cos_phi = numpy.array(cos_phi)

x_sph = r.dot(sin_theta).dot(cos_phi)

In fact I'd use numpy even earlier, by doing:

phi = numpy.array(phi) # I don't know how you calculate this but you can start here with a phi list.
theta = numpy.array(theta)

sin_theta  =numpy.sin(theta)
cos_phi = numpy.cos(phi)

You could even skip the intermediate sin_theta and cos_phi assignments and just put all the stuff in one line. It'll be long and complicated so I'll omit it but I do numpy -maths like that sometimes.

And numpy is fast, it'll make a huge difference. At least a noticeable one.

Answer 2

[...] creates a list and (...) a generator :

generator = (r*sin_t*cos_p for cos_p in cos_phi for sin_t in sin_theta for r in p)
for value in generator:
    # Do something

Answer 3

To turn a loop into a generator, you can make it a function and yield :

def x_sph(p, cos_phi, sin_theta):
    for r in p:
        for sin_t in sin_theta:
            for cos_p in cos_phi:
                yield r * sin_t * cos_p

However, note that the advantages of generators are generally only realised if you don't need to calculate all values and can break at some point, or if you don't want to store all the values (the latter is a space rather than time advantage). If you end up calling this:

lst = list(x_sph(p, cos_phi, sin_theta))

then you won't see any gain.