Function closure performance

Question

I thought that I improve performance when I replace this code:

def f(a, b):
  return math.sqrt(a) * b
result = []
a = 100
for b in range(1000000):
  result.append(f(a, b))

with:

def g(a):
  def f(b):
    return math.sqrt(a) * b
  return f
result = []
a = 100
func = g(a)
for b in range(1000000):
  result.append(func(b))

I assumed that since a is fixed when the closure is performed, the interpreter would precompute everything that involves a, and so math.sqrt(a) would be repeated just once instead of 1000000 times.

Is my understanding always correct, or always incorrect, or correct/incorrect depending on the implementation?

I noticed that the code object for func is built (at least in CPython) before runtime, and is immutable. The code object then seems to use global environment to achieve the closure. This seems to suggest that the optimization I hoped for does not happen.

Answer 1

I assumed that since a is fixed when the closure is performed, the interpreter would precompute everything that involves a, and so math.sqrt(a) would be repeated just once instead of 1000000 times.

That assumption is wrong, I don't know where it came from. A closure just captures variable bindings, in your case it captures the value of a, but that doesn't mean that any more magic is going on: The expression math.sqrt(a) is still evaluated every time f is called.

After all, it has to be computed every time because the interpreter doesn't know that sqrt is "pure" (the return value is only dependent on the argument and no side-effects are performed). Optimizations like the ones you expect are practical in functional languages (referential transparency and static typing help a lot here), but would be very hard to implement in Python, which is an imperative and dynamically typed language.

That said, if you want to precompute the value of math.sqrt(a), you need to do that explicitly:

def g(a):
  s = math.sqrt(a)
  def f(b):
    return s * b
  return f

Or using lambda:

def g(a): 
  s = math.sqrt(a)
  return lambda b: s * b

Now that g really returns a function with 1 parameter, you have to call the result with only one argument.

Answer 2

The code is not evaluated statically; the code inside the function is still calculated each time. The function object contains all the byte code which expresses the code in the function; it doesn't evaluate any of it. You could improve matters by calculating the expensive value once:

def g(a):
    root_a = math.sqrt(a)
    def f(b):
        return root_a * b
    return f
result = []
a = 100
func = g(a)
for b in range(1000000):
    result.append(func(b))

Naturally, in this trivial example, you could improve performance much more:

a = 100
root_a = math.sqrt(a)
result = [root_a * b for b in range(1000000)]

But I presume you're working with a more complex example than that where that doesn't scale?