How to make n-dimensional nested for-loops in Python?

Question

I have the following situation:

for x1 in range(x1, x2):
    for x2 in range(x3, x4):
        for x3 ...
            ...
                f(x1, x2, x3, ...)

How to convert this to a mechanism in which I only tell python to make n nested loops where the variable name is x1, x2, x3, x4, ...? I don't want to write every possibility manually of course, since there might be very many dimensions.

Answer 1

What you want to do is iterate over a product. Use itertools.product.

import itertools

ranges = [range(x1, x2), range(x3, x4), ...]

for xs in itertools.product(*ranges):
    f(*xs)

Example

import itertools

ranges = [range(0, 2), range(1, 3), range(2, 4)]

for xs in itertools.product(*ranges):
    print(*xs)

Output

0 1 2
0 1 3
0 2 2
0 2 3
1 1 2
1 1 3
1 2 2
1 2 3

Answer 2

Recommended: itertools

itertools is an awesome package for everything related to iteration:

from itertools import product

x1 = 3; x2 = 4
x3 = 0; x4 = 2
x5 = 42; x6 = 42
for x, y, z in product(range(x1, x2), range(x3, x4), range(x4, x5)):
    print(x, y, z)

gives

3 0 2
3 0 3
3 0 4
3 0 5
3 0 6
3 0 7
3 0 8
3 0 9
3 0 10
3 0 11
3 0 12
3 0 13
3 0 14
3 0 15
3 0 16
3 0 17
3 0 18
3 0 19
3 0 20
3 0 21
3 0 22
3 0 23
3 0 24
3 0 25
3 0 26
3 0 27
3 0 28
3 0 29
3 0 30
3 0 31
3 0 32
3 0 33
3 0 34
3 0 35
3 0 36
3 0 37
3 0 38
3 0 39
3 0 40
3 0 41
3 1 2
3 1 3
3 1 4
3 1 5
3 1 6
3 1 7
3 1 8
3 1 9
3 1 10
3 1 11
3 1 12
3 1 13
3 1 14
3 1 15
3 1 16
3 1 17
3 1 18
3 1 19
3 1 20
3 1 21
3 1 22
3 1 23
3 1 24
3 1 25
3 1 26
3 1 27
3 1 28
3 1 29
3 1 30
3 1 31
3 1 32
3 1 33
3 1 34
3 1 35
3 1 36
3 1 37
3 1 38
3 1 39
3 1 40
3 1 41

Create a cross-product function yourself

def product(*args):
    pools = map(tuple, args)
    result = [[]]
    for pool in pools:
        result = [x + [y] for x in result for y in pool]
    for prod in result:
        yield tuple(prod)

Answer 3

You could also use numpy.ndindex to achieve what you are looking for:

import numpy
for y1, y2, y3 in numpy.ndindex(x2, x4, ...):
    ...

Thats especially useful when you are already using numpy for something else in your script. You would have to add the 'starting point' (i.e. x1, x2, ...) to each dimension though (if it's not 0).

Answer 4

Use multiprocessing!

import multiprocessing
from itertools import product

results = [] 
def callback(return_value):
    # this stores results
    results.append(return_value)

if __name__=="__main__" :
    pool = multiprocessing.Pool(4)
    args = product(range1, range2, range)
    for x, y, z in args:
        pool.apply_async(f,(x,y,z),call_back=callback)
    pool.close()
    pool.join()

Now you can run your function while also taking advantage of the full capacity of your machine!