Python cartesian product and conditions?

Question

In Python I am using the itertools.product() function to generate input parameters for a simulation.

I have a test function that requires 4 input parameters a1, a2, b1 and b2. I use the following code generate the parameters. Example:

params = itertools.product(range(10,41,2), range(10,41,2), range(0, 2), range(5, 31, 5))

… which gives me 3072 combinations. Unfortunately some combinations logically make no sense. E. g. if a2 is larger than a1 the test results are useless, also when b1 equals 0 the value of b2 is completely irrelevant – so it wouldn’t make sense to test such combinations.

Is there a possibility to restrict or filter the cartesian product beside doing it manually and nesting for-loops? Because my real use case has way more than 4 parameters, that’s why I like the convenience of the cartesian product function from itertools.

Any ideas or alternatives? Any help appreciated, thanks.

Answer 1

Python 3

In Python 3 you can filter out undesirable combinations using itertools.filterfalse:

# predicate is true when need to skip the combination
predicate = (lambda (a1, a2, b1, b2): a1 <= a2 and (b1 != 0 or b2 == 5), params)
filtered_params = itertools.filterfalse(predicate, params)

Python 2

You can use list comprehension or itertools.ifilter:

filtered_params = itertools.ifilter
    (lambda (a1, a2, b1, b2): a1 <= a2 and (b1 != 0 or b2 == 5), params)

Note that both of these versions loop and filter out under the hood. If you want to avoid that, you'll need to construct an improved algorithm that creates the tuples without the undesirables.

Answer 2

If you have many parameters, a constraint-based approach using a module like python-constraint may be easier to work with - let it do the hard work of figuring out which combinations are valid.

This would look something like

from constraint import Problem

prob = Problem()
prob.addVariables(["a1", "a2"], range(10,41,2))
prob.addVariable("b1", [0, 2])
prob.addVariable("b2", range(5, 31, 5))
prob.addConstraint(lambda a1, a2: a2 <= a1, ["a1", "a2"])
prob.addConstraint(lambda b1, b2: b1 != 0 or b2 == 5, ["b1", "b2"])

for params in prob.getSolutionIter():
    run_sim(**params)